Commits · 2e0bfbd92bad3e477d7514c80cde0133a289c078 · OpenDAS / TransformerEngine

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

10 Aug, 2023 1 commit

Add mixed use of cuDNN fprop and flash-attn v2 bprop (#349) · ecd4f808

cyanguwa authored Aug 10, 2023



* Add support for cuDNN fprop and FAv2 bprop
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* skip activation recompute tests if FAv2
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* restrict the use of FAv2 bprop to H100 only
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* move use_FAv2_bwd check to init
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove skipifs for FAv2 in test numerics
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix typos and wording for deterministic checks
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* Remove variables related to FAv2 skipifs
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: cyanguwa <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>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

ecd4f808

09 Aug, 2023 1 commit

Disable FAv2 for deterministic use (#366) · 7804d116

Kirthi Shankar Sivamani authored Aug 08, 2023



* Disable FAv2 for deterministic use
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Also disable FusedAttention backend with deterministic
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

7804d116

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

26 Jul, 2023 1 commit
- Make QK layer scaling opt-in (#339) · 06d5fa97
  Kirthi Shankar Sivamani authored Jul 25, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  06d5fa97
20 Jul, 2023 1 commit
- Relax FA 2.0 checks for Ada (#331) · f0ddab82
  Kirthi Shankar Sivamani authored Jul 19, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  f0ddab82
19 Jul, 2023 2 commits

FlashAttention 2.0 support (#329) · faefacb8

Kirthi Shankar Sivamani authored Jul 18, 2023



* FA v2.0 support
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

faefacb8

FA does not support head_dim > 64 on Ada (#328) · 2a81e939

Kirthi Shankar Sivamani authored Jul 18, 2023



* FA does not support head_dim > 64 on Ada
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add cc8.7 to no FA list
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

2a81e939

18 Jul, 2023 1 commit

layer number bug fixes (#326) · 21162d06

Kirthi Shankar Sivamani authored Jul 18, 2023



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

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

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

21162d06

13 Jul, 2023 1 commit
- Fix dtype for KV inference cache (#319) · b172bad8
  Kirthi Shankar Sivamani authored Jul 13, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  b172bad8
26 Jun, 2023 1 commit

Get default dtype from pytorch (#300) · 79a78cae

Kirthi Shankar Sivamani authored Jun 26, 2023



* Get default dtype from pytorch
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

79a78cae

22 Jun, 2023 1 commit

Add long sequence support for fused attention (#237) · 5c58beaa

cyanguwa authored Jun 22, 2023



* add long sequence support and unify three backends for fused attention
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update cudnn-frontend to v0.9.1
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* replace cpu_float2half_rn with __float2half_rn
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

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

* make cudnn plan caches thread_local
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* replace cuDNN throw with NVTE_CHECK
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix replacement of cuDNN throw with NVTE_CHECK
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* force dropout probablity to 0 in inference mode
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* change negInfinity to be consistent with m512 fused attn
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove float2half conversion for scale_dropout
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add back runtime api for sm detection
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add gemm3 to enums FP8Fwd/BwdTensors
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* change dropout from no to yes for fmha_v1
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove output_rng_state in m512 kernels
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix elts_per_thread calculation in kvpacked fwd
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove dropout=0.0 restriction for m512 fused attn
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove output_rng_state completely from m512 kernels
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>

5c58beaa

21 Jun, 2023 1 commit
- Fix BF16 ONNX export for successful ONNX Runtime Verification (#290) · 804f1203
  asfiyab-nvidia authored Jun 20, 2023
```
Signed-off-by: Asfiya Baig <asfiyab@nvidia.com>
```
  804f1203
13 Jun, 2023 1 commit
- Update FA version (#279) · e17c31c3
  Kirthi Shankar Sivamani authored Jun 13, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  e17c31c3
12 Jun, 2023 1 commit
- Revert "Fix BF16 ONNX export for successful ONNX Runtime Verification (#271)" (#275) · 487871e2
  Tim Moon authored Jun 12, 2023
```
This reverts commit 914f3841

.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
```
  487871e2
11 Jun, 2023 1 commit

Fix BF16 ONNX export for successful ONNX Runtime Verification (#271) · 914f3841

asfiyab-nvidia authored Jun 10, 2023



* fix BF16 onnx export for ort verification
Signed-off-by: Asfiya Baig <asfiyab@nvidia.com>

* Update transformer_engine/pytorch/attention.py
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: asfiyab-nvidia <117682710+asfiyab-nvidia@users.noreply.github.com>

---------
Signed-off-by: Asfiya Baig <asfiyab@nvidia.com>
Signed-off-by: asfiyab-nvidia <117682710+asfiyab-nvidia@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

914f3841

25 May, 2023 1 commit

[PyTorch] Rotary Position Embedding (#246) · 156a074a

Kirthi Shankar Sivamani authored May 25, 2023



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

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

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

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

156a074a

22 May, 2023 1 commit

Relax checks for attn_mask_type in FlashAttention (#226) · 122de2cc

cyanguwa authored May 22, 2023



* relax attn mask type checks for FlashAttention
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* disable flash attn if mask tensor is not None
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix the logic for flash attn
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

122de2cc

09 May, 2023 1 commit

PyTorch refactor (#201) · c6a4a4e0

Kirthi Shankar Sivamani authored May 09, 2023



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

* refactor attention out of transformer.py
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

c6a4a4e0

02 May, 2023 1 commit
- Get packaging module from setuptools (#190) · 25bb8647
  Kirthi Shankar Sivamani authored May 01, 2023
```
Use only built-ins for setup
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  25bb8647
30 Apr, 2023 1 commit
- Remove use of distutils (#186) · 36e0ac56
  Kirthi Shankar Sivamani authored Apr 30, 2023
```
Remove distutils
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  36e0ac56
28 Apr, 2023 1 commit
- Re-add support for PyTorch version 1.x (#180) · 522fecc1
  Kirthi Shankar Sivamani authored Apr 27, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  522fecc1
26 Apr, 2023 1 commit

Faster split of QKV for FlashAttention (#166) · 8adb1b1d

Przemyslaw Tredak authored Apr 26, 2023



* Faster split of QKV for FlashAttention
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

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

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

* Message with assert
Co-authored-by: Przemyslaw Tredak <ptredak@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

* make clarifying comment and check strides
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>

8adb1b1d

21 Apr, 2023 1 commit
- Remove userbuf docs (#164) · 68fc78dd
  Kirthi Shankar Sivamani authored Apr 21, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  68fc78dd
19 Apr, 2023 1 commit

TP communication overlap with userbuffers (#147) · 18da4e88

Kirthi Shankar Sivamani authored Apr 19, 2023



* Port initial changes
Co-authored-by: Sangkug Lym <slym@nvidia.com>
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* readd FA include for PyTorch
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Re-enable sm_70 + cleanup
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* LICENSE, cleanup header
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* 5k -> 173 errors
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* license and fixes in userbuffers-host
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

* Turn off default async amax reduction (#148)
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* remove unused code path
Signed-off-by: Sangkug Lym <slym@nvidia.com>

* cleanup Macros
Signed-off-by: Sangkug Lym <slym@nvidia.com>

* fix conflict resolution bug
Signed-off-by: Sangkug Lym <slym@nvidia.com>

* Fix gencode flags in setup (#145)

* Fix gencode flags based on cuda version
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

* Change overlap config dict error message
Signed-off-by: Sangkug Lym <slym@nvidia.com>

* simplify ub initialization
Signed-off-by: Sangkug Lym <slym@nvidia.com>

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

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

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

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

* fix TE macros in public header
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* compiles with and w/o MPI
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fixes for python side annotations for conditional compile
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* link gdrAPI only when MPI found
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

* load MPI before TE
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add Py side argument checks
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* remove unused code and catch silent failures
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Sangkug Lym <slym@nvidia.com>
Co-authored-by: Sangkug Lym <slym@nvidia.com>
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

18da4e88

17 Apr, 2023 1 commit

Switch to upstream flash-attn (#151) · 7396c527

Kirthi Shankar Sivamani authored Apr 16, 2023



* use upstream flash-attn
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* get correct FA for linting
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

7396c527

07 Apr, 2023 1 commit

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

29 Mar, 2023 1 commit

Fix FlashAttention tests (#99) · bcbd4be0

tcherckez-nvidia authored Mar 29, 2023


Signed-off-by: Tal Cherckez <tcherckez@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

bcbd4be0

28 Mar, 2023 2 commits

Fix zombie process when querying TE install path (#121) · 0963b288

Tim Moon authored Mar 28, 2023



* Remove zombie process from querying TE install path
Co-authored-by: Naman Goyal <naman@fb.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

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

* Fix lint warning
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

0963b288

Fix usage of return_bias argument (#114) · d90bf212

Kirthi Shankar Sivamani authored Mar 28, 2023



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

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

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

d90bf212

22 Mar, 2023 1 commit
- Catch FA internal error with compute capability 8.6 (#113) · 5a881a08
  Kirthi Shankar Sivamani authored Mar 22, 2023
```
FA doesn't support compute 8.6 with head_dim>64
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  5a881a08
11 Mar, 2023 2 commits

Move from Sphinx Autodoc to sphinx-autoapi (#92) · c9245c02

Przemyslaw Tredak authored Mar 11, 2023



* Change from AutoDoc to AutoAPI
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Fixes
Signed-off-by: Przemyslaw Tredak <ptredak@nvidia.com>

* WAR for the wrong autosummary generation
Signed-off-by: Przemyslaw Tredak <ptredak@nvidia.com>

* Change common to be in line with pytorch API docs
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Add GitHub Action to build docs
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

* Trying to fix the versions
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

c9245c02

deprecate qk layer scaling and fp32 softmax args (#90) · 81429b80

Kirthi Shankar Sivamani authored Mar 11, 2023



* deprecate qk layer scaling and fp32 softmax args
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* apply QK layer scaling for fp16 training
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

81429b80

07 Mar, 2023 1 commit

Fix flash attention (#84) · 37a12c4e

Kirthi Shankar Sivamani authored Mar 07, 2023



* ignore self attention mask for causal type
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* further relax checks to run FA, update docs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* minimum ampere requirement for fa
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

37a12c4e

02 Mar, 2023 1 commit

Fix unfused QKV params case; stack vs interleave option (#83) · 7d6c1d02

Kirthi Shankar Sivamani authored Mar 02, 2023



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

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

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

7d6c1d02

22 Feb, 2023 1 commit

flash-attn integration (#62) · 79a9fe29

cyanguwa authored Feb 21, 2023



* add flash attention to TransformerLayer
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* Add docs for FP8 calibration (#61)
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* Fix the integer overflow in fused softmax (#60)
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* prefix flash attn env var with NVTE_
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* Address steady memory increase and bloated checkpoints (#63)
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* fix env var logic
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* fix flash attn env var logic again
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* remove d2d copies (#64)

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

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

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

* Increase number of FP8 tensors per GEMM (#22)

* Increase number of FP8 tensors per GEMM
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Enable FP8 output tensor for fp8_gemm
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* [BERT FP8] Initial TE review comments
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Temporary fix for cuda graph non convergence
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Address review comments-2
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Review comments-3
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Cleanup
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Change for New API
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Remove unnecessary clone for D_scale, D_amax
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Avoid Roll for AMAX history size = 1
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Update onnx_te_gemm API
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Fix Lint errors
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

---------
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* Bug fixes from PR 22 (#65)

* 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>
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* replace rearrange with transpose
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* QKV parameters unfused path fixes and optimization (#66)

* 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>

* Better QKV parameter fusion
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* keep original param for unfused case to retain externally set attrs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

* No need to set data pointers
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* Assert memory loc in NoopCat
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Handle case of different memory in param and buffer
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* Reassign params memory to avoid more concats
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* Fix gradients when using AMP (#70)

retain grad related attrs while casting
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* fix pylint violations 

fixed pyline violations such as trailing white spaces and too long lines 
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>

* fix pylint violation on line 264 with R1719
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>

* fix two more pylint violations
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>

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

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

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

* check for correct flash-attn version
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* lint+build fixes, correct settings for default flash-attn
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* fix onnx and disable flash-attn export test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* cleanup internal API; rm duplication
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* only install TE wheel (exclude flash-attn to rm conflicts)
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* forgot to change install wheel path
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

* review comments and fixes to selective checkpointing
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Charlene Yang <charleney@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Co-authored-by: Charlene Yang <charleney@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

79a9fe29

15 Feb, 2023 1 commit

Zero-centered gamma support in LayerNorm (LayerNorm1p) (#67) · 7324fe2b

Przemyslaw Tredak authored Feb 14, 2023



* C++ implementation of LayerNorm1P
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Expose zero centered gamma to pyTorch
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

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

* Fix ONNX export and tests
Signed-off-by: Przemyslaw Tredak <ptredak@nvidia.com>

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

* Fix backward handling - C++ part
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Fix for backward - Python side
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

* Reenable the pylint check
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

* Change in the bwd kernel
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

7324fe2b

10 Feb, 2023 1 commit

QKV parameters unfused path fixes and optimization (#66) · 04490337

Kirthi Shankar Sivamani authored Feb 10, 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>

* Better QKV parameter fusion
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* keep original param for unfused case to retain externally set attrs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

* No need to set data pointers
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* Assert memory loc in NoopCat
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Handle case of different memory in param and buffer
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* Reassign params memory to avoid more concats
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

04490337

05 Jan, 2023 1 commit

Enforce boolean attention mask type (#49) · 0c9c0ba1

Kirthi Shankar Sivamani authored Jan 04, 2023



* Enforce boolean attention mask type
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

0c9c0ba1