Commits · dff113403293c420bd4ca1ef00050fc41737318a · OpenDAS / TransformerEngine

12 Jun, 2024 1 commit

[JAX] Rewrite the Format of FP8 Meta and Remove unused ShardingTypes. (#842) · dff11340

Ming-Xu Huang authored Jun 12, 2024



* Reformat FP8 Meta

1. Reformat FP8 meta to be one-set-per-tensor.
2. Remove fp8_max and scale_inv.
3. Remove unused functions in fp8.py
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix unit-tests
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Remove ShardingType and MajorShardingType
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix lint errors
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fixed unittests.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Rename few variables.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Add jit to update_amax_list
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fixed naming error in LayernormMLP
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fixed bugs in test_distributed_layernorm_mlp.py
Signed-off-by: Ming Huang <mingh@nvidia.com>

---------
Signed-off-by: Ming Huang <mingh@nvidia.com>

dff11340

07 Jun, 2024 1 commit

Remove interval arg from recipe (#892) · d75db5f7

Kirthi Shankar Sivamani authored Jun 07, 2024



* Remove interval arg from recipe
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Remove usage of interval and use explicit kwarg for testing recipes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

d75db5f7

06 Jun, 2024 1 commit
- Build system refactor for wheels (#877) · c1b915ae
  Kirthi Shankar Sivamani authored Jun 06, 2024
```
Cleanup
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  c1b915ae
01 Jun, 2024 1 commit

Added comments about Llama3 weights to Llama tutorial (#830) · 8b210490

Paweł Gadziński authored May 31, 2024



* Llama 3 update
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Times update
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Times update
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* utils.py fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* utils.py fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* utils.py fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* update te llama tutorial to allow running with llama 3 weights
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* small fixes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* small fix
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* small fix
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* add llama 3 vs llama 2 distinctions
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* paraphrasing and corrected facts
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* fix
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* fix
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

---------
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Co-authored-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Sudhakar Singh <sudhakars@nvidia.com>

8b210490

29 May, 2024 1 commit

New NVIDIA footer in documentation (#876) · 4e30bc4b

Przemyslaw Tredak authored May 29, 2024



* Change the documentation footer
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Update docs toolchain versions
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

4e30bc4b

28 May, 2024 1 commit

Use correct FP8 group in multi-GPU docs (#852) · 9ff2c076

Tim Moon authored May 28, 2024



* Use correct FP8 group in multi-GPU docs

FP8 process group should be tensor-parallel group
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Synchronize FP8 scales over world group in multi-GPU docs
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>

9ff2c076

25 May, 2024 1 commit

Different dimension for attention (#833) · 66736890

Paweł Gadziński authored May 24, 2024



* Fixed Llama tutorial. Changed batch size and added fused=True.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: root <root@ipp2-0037.nvidia.com>

* Tutorial updated but not complete yet.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: root <root@ipp2-0037.nvidia.com>

* Tutorial notebook reseted - removed fuse=true
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: root <root@ipp2-0037.nvidia.com>

* Removed fused=true
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: root <root@ipp2-0037.nvidia.com>

* Batch size back to 8
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: root <root@ipp2-0037.nvidia.com>

* Typo and commented out line
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: root <root@ipp2-0037.nvidia.com>

* fixed whitespace
Signed-off-by: root <root@ipp2-0037.nvidia.com>

* fixed whitespace
Signed-off-by: root <root@ipp2-0037.nvidia.com>

* Added comment to attention line. Fixed potential bug with loading weights - now loading works correctly, confirmed by the generation code.
Signed-off-by: root <root@ipp2-1661.nvidia.com>

* Comments
Signed-off-by: root <root@ipp2-1661.nvidia.com>

* Models cast added again
Signed-off-by: root <root@ipp2-1661.nvidia.com>

* Weight download info
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Moved parameter gate_proj_size to config
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* gate_proj_size removed and put immediate_size instead
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Llama 3 added to tutorial
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Typos fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Typos fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Fixed model loading
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Loading fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Different dim for attention
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Reversed other commit
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Changed name to kv_channels
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Fixed typo
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Back to kv_channels in transformer layer
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Back to kv_channels in transformer layer
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Small bug fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Small bug fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Test fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* changed file modes
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* lint fix and resolved conflict
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* lint fix and resolved conflict
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Lint fix, hopefully last
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

---------
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: root <root@ipp2-0037.nvidia.com>
Signed-off-by: root <root@ipp2-1661.nvidia.com>
Co-authored-by: root <root@ipp2-2373.nvidia.com>
Co-authored-by: root <root@ipp2-1588.nvidia.com>
Co-authored-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: root <root@ipp2-0037.nvidia.com>
Co-authored-by: root <root@ipp2-1661.nvidia.com>
Co-authored-by: root <root@ipp2-2371.nvidia.com>
Co-authored-by: root <root@ipp2-1589.nvidia.com>
Co-authored-by: Sudhakar Singh <sudhakars@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

66736890

12 Apr, 2024 1 commit

[PyTorch] cuda graph support (#575) · 73f8d90f

Kirthi Shankar Sivamani authored Apr 12, 2024



* FP8 cuda graphs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: Charlene Yang <charleney@nvidia.com>

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

* exclude torch compile from numerics tests
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

* rm fusion from unfused path
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: Charlene Yang <charleney@nvidia.com>

73f8d90f

31 Mar, 2024 1 commit

Llama tutorial fixes (#730) · 16a469df

Paweł Gadziński authored Mar 31, 2024



Llama tutorial fixes - all
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Pawel Gadzinski <pgadzinski@nvidia.com>

16a469df

20 Mar, 2024 1 commit

Llama accelerate tutorial (#720) · c38779be

Sudhakar Singh authored Mar 20, 2024



* tutorial and doc fixes
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* remove extra code
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* fix typos
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

---------
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

c38779be

01 Mar, 2024 1 commit
- Create a small tutorial on how to accelerate HF Llama models with Transformer-Engine (#615) · 0bd84ed9
  Sudhakar Singh authored Feb 29, 2024
  
  0bd84ed9
22 Feb, 2024 1 commit

[JAX] Refine MHA API and add DPA API (#653) · 9b2fed51

Reese Wang authored Feb 23, 2024



* Refine MHA API
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Reuse func from the flax
Signed-off-by: Reese Wang <rewang@nvidia.com>

* DPA draft
Signed-off-by: Reese Wang <rewang@nvidia.com>

* qkv packed draft
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix test_layer with fused attn
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add attn_bias_type and enhance a few code flow
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Move scale_factor from __call__ to init
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enhance the docs
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add DPA public API and tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Refine docs
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Refine docs
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix conflict
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add qkv separate fused attn
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Apply BSHD_BSHD_BSHD format
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove debug log
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add fused attention layer tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add NVTE_FUSED_ATTN docs
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fine-grained fused attn settings
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove the default value of num_attetnion_head and head_dim
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add teardown for fused attn env
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Unify the Optional notation
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix Pre/Post scale bias comments
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add no_mask tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add checkpoint_name for fused attn
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix the fused attn batcher
Signed-off-by: Reese Wang <rewang@nvidia.com>

---------
Signed-off-by: Reese Wang <rewang@nvidia.com>

9b2fed51

08 Feb, 2024 1 commit
- Update example to use new TE_DType path (#660) · 379c1ee3
  Quentin Anthony authored Feb 08, 2024
```
Signed-off-by: Quentin Anthony <qganthony@yahoo.com>
```
  379c1ee3
02 Feb, 2024 1 commit

[JAX] Support SP + RoPE + GeLU (#602) · ce163f9e

Ming-Xu Huang authored Feb 03, 2024



* Adding support of sequence parallelism
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding RoPE
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix wrong batch_logical_axes
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Rnaming FSDP outer env var
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Poring RoPE to Praxis layers.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Porting GeLU + [FP8 Cast].
Signed-off-by: Ming Huang <mingh@nvidia.com>

* WAR to make XLA successfully match FP8 GEMM on FFN1 with GeLU.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Allowing arbitrary dimension of NVShape for the workspace allocation
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding checkpoint_name to fused functions of mlp.py to get better perf with nn.scan.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Modify with review feedback.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix bugs
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix typo.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fixed for lint
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Follow review feedback to modify code.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix typo.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Port SP to Praxis
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Fix an issue when enabling both GQA and RoPE.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Update docs
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

ce163f9e

21 Jan, 2024 1 commit

Activation offloading to CPU's for the Linear, Layernorm Linear and the... · f196d14b

Selvaraj Anandaraj authored Jan 21, 2024


Activation offloading to CPU's for the Linear, Layernorm Linear and the Layernorm MLP modules (#571)

* Added support activation offloading to CPU's
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>

* Moving CPU offloading library to TE
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>

* Restructured code, added switch to choose between weight/activation offloading
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>

* Removed arg during constructor
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>

* Fix nit-pick errors
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>

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

* Fix to the code block in docs
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Added offloading unit test
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>

* Fixed formatting
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>

* wgrad fusion fix, minor errors and lint
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* Fixed stray PyT tensors in LayernormMLP getting offloaded
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>

* Fixed typi
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>

* Fix offloading for rmsnorm, rm test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>
Co-authored-by: Przemyslaw Tredak <ptredak@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

f196d14b

19 Jan, 2024 1 commit
- chore: Fix multiple typos (#613) · b4b8ae7b
  hugo-syn authored Jan 19, 2024
```
Signed-off-by: hugo-syn <hugo.vincent@synacktiv.com>
```
  b4b8ae7b
03 Jan, 2024 1 commit
- Change the copyright to include 2024 (#583) · cd798c97
  Przemyslaw Tredak authored Jan 02, 2024
```
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
```
  cd798c97
16 Dec, 2023 1 commit
- Update default value in the recipe docs (#550) · 9020344b
  Przemyslaw Tredak authored Dec 16, 2023
```
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
```
  9020344b
06 Dec, 2023 1 commit

Update README.md - Latest News section (#554) · 14c51e62

Santosh Bhavani authored Dec 06, 2023



* Add H200 perf non-alpha image
Signed-off-by: Santosh Bhavani <santosh@semantic.md>

* Update README.rst - non-transparent H200 plot
Signed-off-by: Santosh Bhavani <santosh@semantic.md>

---------
Signed-off-by: Santosh Bhavani <santosh@semantic.md>

14c51e62

04 Dec, 2023 1 commit
- Update doc to match with latest code. (#511) · 5debfdb3
  Zhenhuan Liu authored Dec 05, 2023
```
Signed-off-by: Zhenhuan Liu <nkulzh16@gmail.com>
```
  5debfdb3
31 Oct, 2023 1 commit

[PyTorch] Experimental FP8 tensor class (#452) · b1820c44

Tim Moon authored Oct 31, 2023



* Experimental FP8 tensor
Co-authored-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Sudhakar Singh <sudhakars@nvidia.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add fp8 tensor to ci test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* Default to FP8 usage
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

* Debug transpose caching

Handle case where transpose cache is updated externally.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Rename FP8GlobalStateManager.with_fp8_parameters
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* remove Float8Tensor from import API
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Avoid caching FP8 transposes if not required
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix import error in FP8 tensor tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix tranpose caching and checkpointing bug
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Improve caching and fix distopt case
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

* Store FP8 attributes in dict

Easier for multiple tensors to share, e.g. detached tensors.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Make sure scale_inv is 1D tensor
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Make sure scale_inv is 1D tensor
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Set default fp8 data type
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Sudhakar Singh <sudhakars@nvidia.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>

b1820c44

24 Oct, 2023 2 commits

Document how to install a specific PR. (#488) · 7a025c6b
Frédéric Bastien authored Oct 24, 2023
```
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>
```
7a025c6b

[paddle] add documentation (#489) · 4d1f92df

Kirthi Shankar Sivamani authored Oct 24, 2023



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

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

4d1f92df

19 Oct, 2023 1 commit
- [PyTorch] rm unused docs (#484) · 136acacb
  Kirthi Shankar Sivamani authored Oct 19, 2023
```
RM unused docs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  136acacb
17 Oct, 2023 1 commit
- Improve documentation (#478) · 0963020f
  Kirthi Shankar Sivamani authored Oct 16, 2023
```
Improve docs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  0963020f
12 Oct, 2023 1 commit

[PyTorch] RNG state support for model parallelism (#473) · 8e757a45

Kirthi Shankar Sivamani authored Oct 12, 2023



* Add class for RNG state tracker.
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

8e757a45

11 Oct, 2023 2 commits
- Remove TF support (#467) · 1f4c3979
  Kirthi Shankar Sivamani authored Oct 11, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  1f4c3979
- [PyTorch] Inference params (KV cache) support (#466) · d7511ec4
  Kirthi Shankar Sivamani authored Oct 11, 2023
```
Inference params support
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  d7511ec4
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

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

31 May, 2023 1 commit

Refactor build system (#235) · 37bbfc76

Tim Moon authored May 31, 2023



* Refactor Setuptools build system

Successfully launches CMake install, but installs CMake extensions in temp dir.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug JAX build

Fix pybind11 import. Distinguish between build-time and run-time dependencies.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add helper function to determine dependencies
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add missing license
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug case where system CMake is too old
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add missing license
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Simplify sanity import tests

Just importing modules provides richer error messages.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Properly install submodules
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Install helper library for TensorFlow
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Update documentation
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Do not install Ninja by default
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Include Git commit hash in version string
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Override build_ext.build_extensions instead of build_ext.run
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix incorrect include path

Restore Ninja dependency. Restore overriding build_ext.run func.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Review suggestions from @nouiz
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Disable parallel Ninja jobs in GitHub actions
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Properly install userbuffers lib
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Tweak install docs

Review suggestion from @ksivaman
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add examples for specifying framework in docs
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>

37bbfc76

26 May, 2023 1 commit
- Documentation fixes (#248) · 215dfe7e
  Kirthi Shankar Sivamani authored May 26, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  215dfe7e
19 May, 2023 1 commit

Support for NVRTC kernels (#138) · e9022290

Tim Moon authored May 19, 2023



* Initial implementation of NVRTC infrastructure
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Initial NVRTC impl for transpose

NVRTC gives compilation errors at runtime. Everything else compiles and passes tests as expected.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug NVRTC transpose impl

NVRTC kernel compiles, runs, and passes tests with FP32.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use variadic template for kernel arguments in RTC kernel launch func
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Refactoring

Added utility header for CUDA Runtime API. Optimized concat_strings function.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add helper function for regex substitutions in strings
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add option to disable NVRTC support
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add support for header includes in NVRTC kernels
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Access lazily-initialized CUDA driver lib and add option to specify CUDA header dir
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Configure NVRTC transpose kernel with simple perf model
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Revert change to tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Style fixes
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add prime-valued test cases
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix multiple definition error
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Optimize NVRTC transpose kernel for small data sizes
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Mention NVRTC in docs
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add unit tests for NVRTC and string utils
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add comment in install docs about NVRTC

Review suggestion from @nouiz
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug perf model for RTC transpose kernel
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Remove NVRTC discussion from docs
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Require CUDA headers unless NVRTC is explicitly disabled
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use diagonal coords in transpose kernel to avoid partition camping
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use std::call_once for thread-safety
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Minor fixes
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug CMake error
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Remove unnecessary call_once
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Remove diagonal coordinates from transpose kernel
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use size_t indices instead of int
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Review suggestions from @ptrendx

Check build-time CUDA include path for run-time CUDA headers. Handle case where CUDA context is initially uninitialized.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>

e9022290

09 May, 2023 1 commit

Add FP16/BF16 fused_attention support with max_seqlen=512 (#175) · 73c9f421

zlsh80826 authored May 10, 2023



* Add fused attention unit tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Use NVTE_* enums
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Use NVTE_Mask_Type and remove FMHADescriptor
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Move common functions to utils
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Change namespace to fused_attn
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Move fused_attn_max_512_fwd_qkvpacked under the general APIs
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add fused_attn_max_512_bwd_qkvpacked
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Move fused_attn_max_512_bwd_qkvpacked under the general APIs
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove redundant blank line
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix a potential bug for cu_seqlen converter
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Reformat fused_attn_max_512
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Refine the unfused attention warning message
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Rename to fused_attn_max_512
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove the deprecated header
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix flax import
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Rename to fused attn
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add attention related mask
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add attn_mask_type and attn_bias_type
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Refactor jax primitive API
* Merge q_cu_seqlen and kv_cu_seqlen
* Remove is_causal_masking
* Replace seed with rng_state
* Add is_training argument
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove dsoftmax from the customcall
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add None guard for bias and dropout_rng
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add version guard
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add is_fused_attn_kernel_available() to correctly dispatch the attention impl
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix the merge conflict
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Adjust the code style
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add the missing blank lines
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Change the order of FADescriptor members
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enhance the readability of fused_attn_max_512.cu
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Generalize the input dimension unpacking
Signed-off-by: Reese Wang <rewang@nvidia.com>

* 16 bits fused attention requires 8.9.1
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Update fused attention support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Handle None type when sharding
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Change to the padding ratio
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Performance optimization for non-bias cases
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Revert the cudnn-frontend PRIVATE keyword which was used for debugging
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Revert "Update fused attention support matrix"

This reverts commit 4effe67d0f08f733919a329ce5ab421958740f4a.
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Treat b * s as total_seqs to align ragged cases
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add FP16/BF16 max_seqlen <= 512 fused attention to the support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Refine test_fused_attn.py

* Replace reference code with flax.linen
* Remove unnecessary comments
* Use AttnMaskType
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Unify the cuDNN compile version
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add dropout to the support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Slightly adjust the headers
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Typo fix: remove redundant either
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Consolidating fused attention requirements
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Replace cudnn_frontend::throw_if with NVTE_CHECK for the better error line report
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Rename to fused_attn_fp16_bf16_max_seqlen_512 for the better readability
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove CUDNN_FRONTEND_UNUSED
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add more annotations to the custom calls
Signed-off-by: Reese Wang <rewang@nvidia.com>

---------
Signed-off-by: Reese Wang <rewang@nvidia.com>

73c9f421

03 May, 2023 1 commit

Use NVIDIA Sans font in Transformer Engine documentation (#179) · a2c9c635

Przemyslaw Tredak authored May 03, 2023



* Use NVIDIA Sans font in Transformer Engine documentation
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Use NV font for the names in the API
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

a2c9c635

29 Apr, 2023 1 commit

Correct cuDNN version requirement (#184) · 14c2b719

cyanguwa authored Apr 29, 2023



correct cuDNN version requirement
Signed-off-by: Charlene Yang <charleney@nvidia.com>

14c2b719

28 Apr, 2023 1 commit

[JAX] Adjust Module Structure. (#169) · 0792ded4

Ming-Xu Huang authored Apr 28, 2023



* Adjust Module Structure.

1. Collect Flax related modules to a sub-folder, flax.
2. Add a function to unify scale_init for zero-centered-gamma LN.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Make changes be compatible to previous versions.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adapt jax/examples to the new module structure.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Update jax/docs and Add deprecated warning.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Update README
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding deprecated_wrapper
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding deprecated warning to flax modules which imported via transformer_engine.jax
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix CI errors and update docs.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Removing unnecessary deprecated warning in docs.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Implementing __iter__ to DeprecatedEnum.
Signed-off-by: Ming Huang <mingh@nvidia.com>

---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

0792ded4

26 Apr, 2023 1 commit

Recursive submodule checkout (#176) · e5a69d92

Kirthi Shankar Sivamani authored Apr 26, 2023



* Recursively checkout submodules
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>

e5a69d92

25 Apr, 2023 1 commit

Add guide to build from source (#172) · 4f3d6341

Kirthi Shankar Sivamani authored Apr 25, 2023



* Add guide to build from source
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>

4f3d6341

21 Apr, 2023 1 commit

Add FP8 fused attention (#155) · 989a53a0

cyanguwa authored Apr 21, 2023



* Add FP8 fused attention to TE for PyTorch
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* add license for cudnn-frontend, modify installation requirements, and refactor some headers for aesthetics
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* add c api docs for fused attention
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* add exception for unsupported precision/sequence length combinations
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* fix installation requirement for non fused attn use cases
Signed-off-by: Charlene Yang <charleney@nvidia.com>

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

* prefix enums with NVTE_ and replace old MHA_Matrix with NVTE_QKV_Matrix
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* minor fixes based on PR comments
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* fix description for kvpacked fwd
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* fix description of Bias in C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* minor fixes for cudnn requirement and description for QKV tensors
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* fix QKV layout description and support matrix for C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* add asserts to cpp_extensions for qkv layout/bias type/attn mask type
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* fix typo precision
Signed-off-by: Charlene Yang <charleney@nvidia.com>

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

989a53a0