transformer_engine/paddle/fp8.py · daad219f87f7929f3cb3b913361fce44fc06e77d · OpenDAS / TransformerEngine

"tests/pytorch/attention/test_attention.py" did not exist on "e762592ebd2ec8b5a73d5a04912a3ba3cf3f3792"

[Paddle] Optimize memory usage when training in pipeline parallel (#580) · daad219f

Tian Zheng authored Jan 13, 2024



* Actively free tensor in bwd
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* - Add inplace support for fp8 casting
- Allow skipping weight update in fp8 meta update
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Support weight caching for Linear
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Add weight caching for LayernormLinear
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Add weight caching for LayerNormMLP
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Add weight caching for Transformer layer
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Add PP unittests
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Fix CI
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

---------
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

daad219f

fp8.py 12 KB

Replace fp8.py