Commits · 47ca514aa25589f6c9d6f7e9aade29fb5bb370ae · OpenDAS / TransformerEngine

04 Oct, 2023 1 commit

Support packed input for FA (#302) · 47ca514a

Kirthi Shankar Sivamani authored Oct 03, 2023



* initial changes [wip]
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add padding mask support for FA
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Address review comments
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* rm causal mask from tests and add padding
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fix some conflicts
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* conflicts
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add unpadding mask
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix padding mask
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix docs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* [wip] fix API
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add packing and unpacking
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* More fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* docs fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fix atomic_add bf16 torch.compile
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Generate non all True masks
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Lint fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix core attention export and FusedAttn filter
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix all ONNX tests
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Memory optimization
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* More fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Optimizations and caching fixes in torch.dynamo
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Bug fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Review comments
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Padding optimizations
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fixes and reviews
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

47ca514a

27 Sep, 2023 1 commit
- [PyTorch] Mcore DDP support (#446) · 06eebf66
  Kirthi Shankar Sivamani authored Sep 26, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  06eebf66
20 Sep, 2023 1 commit

[pyTorch] Enable the model to change precision between iterations (#414) · 7e759174

Przemyslaw Tredak authored Sep 21, 2023



* Enable the model to be change precision between iterations
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Add test
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Fix for the test
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

---------
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

7e759174

26 Aug, 2023 1 commit

[PyTorch] move mask types to fprop (#402) · 6aa1fcc8

Kirthi Shankar Sivamani authored Aug 25, 2023



* API change and some test fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* more test fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* ONNX fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fixed fused attention tests
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* rm duplicate test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

6aa1fcc8

16 Aug, 2023 1 commit

[PyTorch] FP8 fixes (#380) · 2e0bfbd9

Kirthi Shankar Sivamani authored Aug 16, 2023



* Initial refactor
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Reorder methods by purpose
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Save full global state
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix tests
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* More fixes to test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

2e0bfbd9

10 Aug, 2023 1 commit
- AMP support for LN and RMSNorm (#371) · 3cc2c1d2
  Kirthi Shankar Sivamani authored Aug 10, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  3cc2c1d2
27 Jul, 2023 1 commit

Exposing RMSNorm in pyTorch (#306) · fdc09f42

Przemyslaw Tredak authored Jul 27, 2023



* Exposing RMSNorm in pyTorch extensions
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* First pass at the Python API
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Small fixes
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Added numerics tests and fixed issues
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Lint fixes
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Added RMSNorm to LayerNormMLP
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Added ONNX export and tests for RMSNorm
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Fix python lint
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Fix BERT case
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Added normalization option to the TransformerLayer
Added tests
Fixed test failures
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Fix documentation
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix kwarg bug
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix IMA and invalid type error
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Increase RMSNorm threshold for bf16 case
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix ONNX tests
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

fdc09f42

13 Jun, 2023 1 commit

Adding other activation types to LayerNormMLP (#265) · c67bb2fc

Przemyslaw Tredak authored Jun 13, 2023



* Added ReLU and GLU variants to common
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* pyTorch changes
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* PyTorch C++ lint
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Bug fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* More fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix storage errors
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Compute bgrad
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix numerical tests
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix ONNX export tests
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Review comments
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

c67bb2fc

17 Apr, 2023 1 commit

[PyTorch] Add tests for cuda graph capture (#144) · f126a04f

Kirthi Shankar Sivamani authored Apr 16, 2023



* Add tests for cuda graph capture
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* add sanity test and address reviews
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

f126a04f

07 Apr, 2023 2 commits

make bias configurable (#130) · 82dde778

ngoyal2707 authored Apr 07, 2023



* made bias configurable
Signed-off-by: Naman Goyal <naman@fb.com>

* removed commented lines
Signed-off-by: Naman Goyal <naman@fb.com>

* Update transformer_engine/pytorch/jit.py
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: ngoyal2707 <ngoyal2707@users.noreply.github.com>

* Update transformer_engine/pytorch/jit.py
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: ngoyal2707 <ngoyal2707@users.noreply.github.com>

* fixed incorrect call to fused bias dropout add kernel
Signed-off-by: Naman Goyal <naman@fb.com>

* Update transformer_engine/pytorch/jit.py
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* Separate FC1 and FC2 use_bias args; solves all ci errors
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* jit fusion improvement
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Docs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Naman Goyal <naman@fb.com>
Signed-off-by: ngoyal2707 <ngoyal2707@users.noreply.github.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Naman Goyal <naman@fb.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

82dde778

Conditional dgrad computation for Linear API (#134) · a2e19b7a

Kirthi Shankar Sivamani authored Apr 06, 2023



* small cleanup before starting
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* conditional dgrad for Linear
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* add tests and small improvements to LNLinear and LNMLP
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

a2e19b7a

04 Apr, 2023 1 commit

Add FP8 support for Ada (#129) · 96ad903c

Kirthi Shankar Sivamani authored Apr 04, 2023



* Add FP8 support for Ada
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* better message
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* lint fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Address review comments
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* better message for no fp8
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* same thing for onnx test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix CI and review
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

96ad903c

17 Mar, 2023 1 commit

Improve PyTorch test harness (#102) · 2c996359

Kirthi Shankar Sivamani authored Mar 17, 2023



* add layernorm1p fp8 test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* combine tests for easy maintenance
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* using torch.autocast for AMP and check grad types
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add test for wgrad accumulation fusion
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* rename file
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Setup numerical tests + SAR
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add test for full activation recompute
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add tests for checkpoint load/store
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* TE vs framework numerical tests
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fix ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* relax thresholds
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

2c996359

24 Feb, 2023 1 commit

Move TE/PyTorch UT to tests/pytorch/ (#78) · 97b344cd

Jeng Bai-Cheng authored Feb 24, 2023



* move TE/PyTorch UT to tests/pytorch

1. move tests/* files to tests/pytorch/
2. adjust UT paths in qa/L0_unittest/test.sh
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* update build.yml
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

---------
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

97b344cd

07 Feb, 2023 1 commit

Bug fixes from PR 22 (#65) · 78b4e933

Kirthi Shankar Sivamani authored Feb 06, 2023



* Bug fixes from PR 22
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add FP8 tests to ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* bundle unittests for ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

78b4e933

03 Jan, 2023 1 commit

Update copyright year (#48) · 64a8dc90

Przemyslaw Tredak authored Jan 03, 2023


Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

64a8dc90

03 Nov, 2022 1 commit

Conditional wgrad support (#21) · 78210127

schetlur-nv authored Nov 03, 2022



* Conditional dgrad/wgrad support
Signed-off-by: Sharan Chetlur <schetlur@dlcluster.nvidia.com>

* Fixing the change to depend only on requires_grad. Also updating LayerNorm MLP
Signed-off-by: Sharan Chetlur <schetlur@dlcluster.nvidia.com>

* Minor fixes.
Signed-off-by: Sharan Chetlur <schetlur@dlcluster.nvidia.com>

* Adding conditional wgrad for LayerNormLinear
Signed-off-by: Sharan Chetlur <schetlur@dlcluster.nvidia.com>

* bug fix and remove conditional dgrad

Co-authored-by: schetlur-nv schetlur@nvidia.com
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Adding unit test for wgrad disabled path
Signed-off-by: Sharan Chetlur <schetlur@dlcluster.nvidia.com>

* Adding more unit tests for wgrad disabled path
Signed-off-by: Sharan Chetlur <schetlur@dlcluster.nvidia.com>

* Adding unit tests for fp8 wgrad disabling, and cleaning up the code.
Signed-off-by: Sharan Chetlur <schetlur@dlcluster.nvidia.com>

* fix lint errors
Co-Authored-By: Sharan Chetlur <schetlur@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Sharan Chetlur <schetlur@dlcluster.nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Sharan Chetlur <schetlur@dlcluster.nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

78210127

28 Sep, 2022 1 commit

Inital code drop · 996ea169

Przemek Tredak authored Sep 27, 2022


Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

996ea169