Commits · 2574a1ca23f6d7fe9b4748c6cc347f158d232e22 · OpenDAS / TransformerEngine

"rccl/git@developer.sourcefind.cn:wh1225/rccl-auto-tuning.git" did not exist on "7dc4e964ce411652266839f732396d5a5da20c15"

11 Oct, 2023 2 commits

move cp_group setting to DotProductAttention (#468) · 2574a1ca

Xiaowei Ren authored Oct 11, 2023



* rename set_context_parallel_running to set_context_parallel_group
Signed-off-by: xren <xren@nvidia.com>

* bug fix
Signed-off-by: xren <xren@nvidia.com>

---------
Signed-off-by: xren <xren@nvidia.com>

2574a1ca

[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

10 Oct, 2023 2 commits
- Remove deprecated APIs (#464) · daa5e184
  Kirthi Shankar Sivamani authored Oct 10, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  daa5e184
- Update recipe documentation (#465) · 29b4670c
  Kirthi Shankar Sivamani authored Oct 10, 2023
```
Fix docs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  29b4670c
09 Oct, 2023 3 commits

[C/PyTorch] RoPE fixes and minor improvements for fused attention (#453) · 92d1ba0d

cyanguwa authored Oct 09, 2023



* add support for h2d/2hd in 8.9.6
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* cull unit tests in fused_attn.py and add skipif for layout tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add workopt=1 flag for dpa tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update support table for arbi_seqlen backend
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix rotary position embedding and add unit tests accordingly
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* further cut down unit tests for CI efficiency
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

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

92d1ba0d

Disable GitHub CI jobs using NGC PyTorch container (#462) · 79f5fac7
Tim Moon authored Oct 09, 2023
```
Signed-off-by: Tim Moon <tmoon@nvidia.com>
```
79f5fac7
Update reference scale calculation in TensorFlow test (#463) · bf3e1715
Tim Moon authored Oct 09, 2023
```
Signed-off-by: Tim Moon <tmoon@nvidia.com>
```
bf3e1715

06 Oct, 2023 2 commits

Add user to TE CI and enable passing options to the CI (#460) · 38b85c35
Przemyslaw Tredak authored Oct 06, 2023
```
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
```
38b85c35

[JAX] Enhance Dropout in TransformerLayer. (#444) · 80222dc0

Ming-Xu Huang authored Oct 07, 2023



* [JAX] Enhance Dropout in TransformerLayer.

1. Fixed missing setup of dropout RNG key in TransformerLayer and
   LayerNormMLP.
2. Allowing seperated dropout rate for FC1's output and other hiddens.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix wrong fp8 scale in _update_fp8_metas_impl
Signed-off-by: Ming Huang <mingh@nvidia.com>

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

80222dc0

05 Oct, 2023 1 commit

Atomic gemm and FP8 Reduce Scatter (#449) · 958e1889

vasunvidia authored Oct 05, 2023



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

* Repro for RS output mismatch with Single GEMM + Split pipelined RS
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* minor changes for AG->GEMM pipelined overlap
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Add Atomic Gemm cublasApi attributes and initial implementation of AG->Atomic GEMM
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* AtomicGemm+RS functional with workaround
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* add amax update to layernorm_linear for FP8 unit test accuracy
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Enable reducescatter2_userbuff_strided variants
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

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

* AG+AtomicGemm overlap functional but gemm doesnt overlap with comm
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Add userbuffers_sendrecv kernel variants
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* TransformerLayer API changes to enable AtomicGemm+RS overlap
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

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

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

* [UB] AllGather Atomic GEMM overlap using userbuffer_sendrecv kernels
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Code cleanup + bug fix for multiatomic sendrecv kernel
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

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

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

* [UB] Add shuffling for better AG AtomicGEMM overlap
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Bug fix for AG AtomicGemm overlap
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Bug fix for multiAtomicAG and singleAtomicAG
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Use chunk_i+1 as recv_chunk for multiatomic_AG with shuffling
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Launch AtomicGEMM after first-chunk AG
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Rebase to main
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Add FP8 ReduceScatter kernels, AtomicGEMM+FP8 RS not functional
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Revert "Add FP8 ReduceScatter kernels, AtomicGEMM+FP8 RS not functional"

This reverts commit 80a47a76355440cd5fb4314c96fe9fda632d87f9.
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Add support for NVLS-MC and FP8 Reduce Scatter
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

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

* Atomic and Multiatomic FP8 RS functional
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Remove debug print
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* UB comm initialization hang fix
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

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

* Create new GEMM API for Atomic GEMM
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

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

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

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

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

* Revert NVLS-MC
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Check cu* versions for running atomic gemms
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

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

* Add warning to c api
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

958e1889

04 Oct, 2023 2 commits

Pin astroid to 2.15.7 (#457) · be67f219
Przemyslaw Tredak authored Oct 04, 2023
```
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
```
be67f219

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

03 Oct, 2023 2 commits

Make an import optional for a JAX API change that will happen soon. (#454) · d3157e2a
Frédéric Bastien authored Oct 03, 2023
```
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>
```
d3157e2a

[Paddle] Some refactor and optimization on fused attention (#411) · 5f60f82f

Shijie authored Oct 03, 2023



* fix mask conversion and rng_state
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* refactor fused attn
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* use CUB to do prefix sum
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* fuse dropout add
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* minor changes
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* optimize kernel
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

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

---------
Signed-off-by: Shijie Wang <jaywan@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

5f60f82f

02 Oct, 2023 1 commit

Update README.rst - Integrations section (#429) · 30d51226

Santosh Bhavani authored Oct 02, 2023



minor grammatical changes and added "JAX Toolbox" to integrations
Signed-off-by: Santosh Bhavani <santosh@semantic.md>

30d51226

01 Oct, 2023 1 commit

wgrad-accumulation patch for legacy megatron-LM (#451) · aaba1b01

Kirthi Shankar Sivamani authored Sep 30, 2023



* wgrad-accumulation patch for legacy megatron-LM
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

aaba1b01

27 Sep, 2023 4 commits
- Change version to 1.0.0dev · a29e2e11
  Kirthi Shankar Sivamani authored Sep 27, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  a29e2e11
- Keep previous FA version (#450) · c86ac7fb
  Kirthi Shankar Sivamani authored Sep 26, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  c86ac7fb
- Add release to deprecation warnings (#447) · 95dea389
  Kirthi Shankar Sivamani authored Sep 26, 2023
```
Change deprecation warnings
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  95dea389
- [PyTorch] Mcore DDP support (#446) · 06eebf66
  Kirthi Shankar Sivamani authored Sep 26, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  06eebf66
25 Sep, 2023 1 commit

[C/Pytorch] Expand layout support for fused attention (#403) · 76669cdd

cyanguwa authored Sep 24, 2023



* add flexible layout support
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add support for flexible qkv layout
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

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

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

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

* more changes; WIP
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* more changes; WIP
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* sb3hd/bs3hd working on top of 3xsbhd/bshd/thd
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* remove qkvso_strides on torch side; cover it in generateQKVStrides
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* all 15 layouts pass
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* removed most debug info/clean up
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add note to deprecate some qkv layouts
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix code for unit tests in test_fused_attn.py
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* further remove debug info
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove a couple more comments
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

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

* fix onnx for core attn; not fixed
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove nvtx and add env var for workspace opt
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove testing for env var
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* replace zeros/zeros_like with empty/empty_like
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix nvtx marker name for _q_k_v API
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove sm80 when compiling for h100
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add mapping from qkv layout to layout group and qkv format
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* clean up enums mapping and remove trailing spaces
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* simplify workspace opt control logic; only need env var
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix fp8 test, and minor modifications for other tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* avoid overwriting model configs in unit test
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* random fixes/improvements: get_qkv_format/etc, default values, docstrings, comments
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* change workspace opt logic back to FORCE_WORKSPACE_OPT
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix FP8 tests and generateStrides function
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix get_backend logic for max512/arbitrary
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix unit tests; need cleanup
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* clean up unit tests for layouts, and fix minor lint issue
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* minor tweaks for CI testing: onnx string issue and test fused attn first
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove one unsupported layout from max512 and add a check to qkvpacked API
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix te layer test; reduce test time
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* revert compiler option changes; add back sm80 for even h100
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove some unit tests or make them optional to reduce CI time
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove more unit tests temporarily
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove _q_k_v in naming and add NVTE_ERROR for FP8 Aux_CTX_Tensors size checks
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add more deprecation notes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove temp tests from last commit
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* remove prints from last commit
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove redundant contiguous()
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove thd->bs3hd user warning to avoid GPU sync
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* temporary fix for onnx issue; more fixes in PR 437
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove unused variables
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: Charlene Yang
Signed-off-by: cyanguwa <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

76669cdd

23 Sep, 2023 3 commits

Fix layernorm in GQA (#434) · db589510

cyanguwa authored Sep 22, 2023



* [PyTorch] Implement GQA based on fused q, k, v projection. Additionally fixes #392
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>

* [PyTorch] Extend parameters_split option in Linear and LayerNormLinear to support splitting with different sizes as required by unfused GQA.
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>

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

* fix noop cat to bypass torch.cat and support uneven split
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* fix cuda graph due to noop_cat
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* remove the use of enumerate when possible
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

---------
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Co-authored-by: Markus Schnoes <markus.schnoes@gmx.de>

db589510

[PyTorch] Fix ONNX exports (#437) · 903e1f4f

Kirthi Shankar Sivamani authored Sep 22, 2023



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

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

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

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

903e1f4f

Change scaling factor from E8M0 to E8M23 (#427) · 2da34d41

Kirthi Shankar Sivamani authored Sep 22, 2023



* Change scaling factor from E8M0 to E8M23
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

2da34d41

22 Sep, 2023 3 commits

add flash implementation with context parallelism (#362) · 479dbb73

Xiaowei Ren authored Sep 22, 2023



* add flash implementation with context parallelism
Signed-off-by: xren <xren@nvidia.com>

* next more comments
Signed-off-by: xren <xren@nvidia.com>

* code comment fix
Signed-off-by: xren <xren@nvidia.com>

* comment fix
Signed-off-by: xren <xren@nvidia.com>

* add missing space
Signed-off-by: xren <xren@nvidia.com>

* fix docstrings
Signed-off-by: xren <xren@nvidia.com>

* try to add fa v2 api
Signed-off-by: xren <xren@nvidia.com>

* fix a comment
Signed-off-by: xren <xren@nvidia.com>

* fix padded kv return
Signed-off-by: xren <xren@nvidia.com>

* add docstrings of context parallelism
Signed-off-by: xren <xren@nvidia.com>

* minor fix
Signed-off-by: xren <xren@nvidia.com>

* minor docstring fix
Signed-off-by: xren <xren@nvidia.com>

* fix positional arguments
Signed-off-by: xren <xren@nvidia.com>

* make docstring line shorter
Signed-off-by: xren <xren@nvidia.com>

* add fa v2 backward api for flash_attn_with_cp
Signed-off-by: xren <xren@nvidia.com>

* remove redundant code
Signed-off-by: xren <xren@nvidia.com>

* make sure hidden size per attn head is multiple of 8 for FA2
Signed-off-by: xren <xren@nvidia.com>

* remove an unnecessary assert check for FA2
Signed-off-by: xren <xren@nvidia.com>

* indention fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

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

---------
Signed-off-by: xren <xren@nvidia.com>
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

479dbb73

[PyTorch] set SP attr on bias param for reduction (#440) · b95c1818
Kirthi Shankar Sivamani authored Sep 22, 2023
```
Fix for sequence-parallel
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
b95c1818

[Paddle] Eliminate amax update bubbles by using custom_ops (#436) · c32a62cc

zlsh80826 authored Sep 22, 2023



* Eliminate amax_and_scale_update bubbles
Signed-off-by: rewang <rewang@nvidia.com>

* Add CUDA check
Signed-off-by: rewang <rewang@nvidia.com>

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

c32a62cc

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

19 Sep, 2023 1 commit
- Change version to 0.14.0dev · e7eff4a3
  Przemek Tredak authored Sep 19, 2023
```
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
```
  e7eff4a3
18 Sep, 2023 1 commit
- Add an option for TP-only AMAX reduction (#431) · eb64ec2a
  Sangkug Lym authored Sep 18, 2023
```
Signed-off-by: Sangkug Lym <slym@nvidia.com>
```
  eb64ec2a
12 Sep, 2023 1 commit

[C/PyTorch] Add workspace optimization for fused attention arbitrary_seqlen backend (#396) · a150d286

cyanguwa authored Sep 12, 2023



* add workspace optimization for arbitrary_seqlen fused attn
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* add use_workspace_opt to cudnn plan cache and fix workspace estimate
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* modify workspace opt logic; move zero fill to FP8 API only; other minor fixes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* fix std string error when input is nullptr
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* Add = for required vs allowed workspace comparison
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>

a150d286

07 Sep, 2023 1 commit

[PyTorch] Distributed testing (#398) · 4e37499b

Kirthi Shankar Sivamani authored Sep 07, 2023



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

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

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

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

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

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

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

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

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

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

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

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

* get default file name with time
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

4e37499b

06 Sep, 2023 1 commit

[Paddle] Support recompute (#412) · 45a2ac41

Tian Zheng authored Sep 07, 2023



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

* Support recompute core attention
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

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

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

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

* Improve performance of stack backtrace
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

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

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

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

45a2ac41

05 Sep, 2023 1 commit
- [JAX] Use the new API when it is available. (#419) · a5dbf1e2
  Frédéric Bastien authored Sep 05, 2023
```
Use the new API when it is available.
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>
```
  a5dbf1e2
01 Sep, 2023 1 commit

[Paddle] Add control of RNG state (#410) · 805b9872

Tian Zheng authored Sep 02, 2023



* Add control of attention dropout and hidden dropout RNG state
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

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

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

805b9872

30 Aug, 2023 1 commit

[JAX] Fix incorrect sharding when only enable FSDP and Mem Misaligned in LN_BWD. (#379) · 3a63b13d

Ming-Xu Huang authored Aug 31, 2023



* [JAX] Fix incorrect sharding when only enable FSDP.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* [JAX] Add WAR to memory misaligned issues of LN BWD.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* [JAX] Reuse sm_arch for avoiding duplicate code.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* [JAX] Support multiple sizes allocation in WorkspaceManager.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* [JAX] Use template and ariadic arguments to improve multple sizes allocator.
Signed-off-by: Ming Huang <mingh@nvidia.com>

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

3a63b13d

26 Aug, 2023 2 commits

[Paddle] Add parallel support (#357) · b8ba734e

Tian Zheng authored Aug 27, 2023



* [Paddle] Add TP, DP, PP, FSDP
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

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

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

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

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

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

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

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

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

* Squashed commit of the following:

commit 79e2e5fd774e67dcdda9aae01a9f31a6479c5d70
Author: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>
Date:   Sun Aug 20 14:39:16 2023 +0000

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

commit 1d40ad60540490f97ed82ba877cc6eda8902cbf6
Author: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>
Date:   Sun Aug 20 14:22:25 2023 +0000

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

commit 6632f735a0c8251862355fc74622af59fae3a509
Author: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>
Date:   Sun Aug 20 05:52:18 2023 +0000

    Add TP for attention and transformer layer
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

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

* Add FSDP check for stage 1,2,3
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

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

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

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

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

---------
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

b8ba734e

[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

25 Aug, 2023 2 commits

Error handle for non-sm80/sm90 GPUs when using fused attention (#393) · 94c57e4d

zlsh80826 authored Aug 26, 2023



* Fused attention kernel only supports sm80 and sm90
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* arbitary fused kernel supports sm86/sm89 after 8.9.3
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Forward is_fused_attn_kernel_available to cpp backend
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

---------
Signed-off-by: Reese Wang <rewang@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

94c57e4d

Fix rng_state issue and minor compiler warning (#395) · b90a8d3a

cyanguwa authored Aug 25, 2023



fix rng_state issue and minor compiler warning
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

b90a8d3a