Commits · 94c57e4d5cc58655233196934bb3cf4d7a6a1f5f · OpenDAS / TransformerEngine

19 Aug, 2023 1 commit

PyTorch MultiheadAttention API (#387) · 8aa2da17

Kirthi Shankar Sivamani authored Aug 19, 2023



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

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

* Expose MultiheadAttention for import
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Expand mask type and add no mask numerical test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

8aa2da17

11 Aug, 2023 1 commit

Miscellaneous fixes for core attention (#344) · cbfb8c6b

cyanguwa authored Aug 10, 2023



* miscellenous fixes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add back pytorch csrc extensions.h
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add unit tests for dpa checkpointing
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove seqlen%32/64 checks for now
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix tests for core attn bias
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add tests for changes regarding rng_state in aux_ctx_tensor
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* reuse rng tracker from numerics in fused attn; skip checkpointing if FAv2 in numerics
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* uncomment comments used for testing
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix pre/post scale bias
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* remove skipifs for FAv2 check after PR366
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove checkpointing tests for transformer layer; dpa tests still provide coverage
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* adjust random number range for tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* Add upper bound to FA version
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Check backend only when using FusedAttention
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* remove imports/variables related to FAv2 checks
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* further fix random number ranges for tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix variable referenced before assignment error
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

---------
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: cyanguwa <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

cbfb8c6b

08 Aug, 2023 1 commit
- Fix for the RMSNorm tests/doc/ONNX export to match the actual implementation (#364) · a0e1cf99
  Przemyslaw Tredak authored Aug 09, 2023
```
Fix for the RMSNorm tests/doc/ONNX export to match the actual
implementation
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
```
  a0e1cf99
29 Jul, 2023 1 commit

Add support for multi-query and grouped-query attention (#338) · 1cb4b25a

cyanguwa authored Jul 28, 2023



* add support for multi-query/grouped-query attention
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix lint
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* revert to flash-attn 1.0.6 and build 2.0.0.post1 manually in CI
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add keyword name for DPA input
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix fused attn tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix skipif for pytest
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* Update transformer_engine/pytorch/attention.py
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Update tests/pytorch/test_fused_attn.py
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* add skipifs for pytest
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove higher limit for flash-attn version
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

---------
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

1cb4b25a

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

12 May, 2023 1 commit

PyTorch API numeric tests (#215) · ec0d40d6

Kirthi Shankar Sivamani authored May 11, 2023



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

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

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

ec0d40d6

18 Apr, 2023 1 commit

Tighten tolerances for graph capture test (#153) · b2b3fbe7

Tim Moon authored Apr 17, 2023


Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

b2b3fbe7

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

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