Commits · 8e97c8da81082fddae66b9693219780c9b66736f · OpenDAS / TransformerEngine

17 Oct, 2024 2 commits

[Bugfix] Fix bias for 0-dim tensors in gemm (#1246) · 8e97c8da

Xin Yao authored Oct 17, 2024



* fix bias for 0-dim tensor
Signed-off-by: Xin Yao <xiny@nvidia.com>

* add check
Signed-off-by: Xin Yao <xiny@nvidia.com>

* use numel() instead of nullptr
Signed-off-by: Xin Yao <xiny@nvidia.com>

---------
Signed-off-by: Xin Yao <xiny@nvidia.com>

8e97c8da

[PyTorch] Fix wgrads for GroupedLinear when weights don't require grad (#1258) · 2d7020e2

Xin Yao authored Oct 17, 2024



Fix wgrad for GroupedLinear when weights doesn't require grad
Signed-off-by: Xin Yao <xiny@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

2d7020e2

16 Oct, 2024 5 commits

[PyTorch] Fix FP8 activation recompute (#1254) · a5181512
Kirthi Shankar Sivamani authored Oct 16, 2024
```
Fix FP8 activation recompute
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
a5181512

Upgrade pylint to 3.3.1 (#1257) · 6e90fcb7

Kirthi Shankar Sivamani authored Oct 16, 2024



* Upgrade pylint and first round formatting
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

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

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

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

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

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

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

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

6e90fcb7

[PyTorch] Drop FA as an installation requirement (#1226) · 161b1d98

Charlene Yang authored Oct 15, 2024



* WIP: make FA2 optional
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



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

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



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

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



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

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



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

* add L1 test to test all supported FA versions
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update version to 2.1.1 and trim L1 tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



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

* remove onnxruntime from L1 FA versions tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

---------
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

161b1d98

fix assertion bug for SWA API in TE-JAX (#1242) · 43b9e1ee

Md Fahim Faysal Khan authored Oct 15, 2024



fixed assertion bug for SWA
Signed-off-by: Md Fahim Faysal Khan <mdfahimfaysa@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Phuong Nguyen <36155692+phu0ngng@users.noreply.github.com>

43b9e1ee

[PyTorch] Build custom ORT ops before running ONNX export tests (#1252) · f6b766bd

Tim Moon authored Oct 15, 2024



* Build custom ORT ops before running ONNX tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Remove ONNX from context parallelism tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Export ONNX ops that do compute in FP32

Matches internal impl of TE kernels.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add build script for custom ORT ops
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

f6b766bd

15 Oct, 2024 1 commit
- Check for backend support in Jax context parallel fused attention test (#1227) · 20c55e46
  Michael Goldfarb authored Oct 15, 2024
```
Update test to check support for context parallel attention.
Signed-off-by: Michael Goldfarb <mgoldfarb@nvidia.com>
```
  20c55e46
14 Oct, 2024 1 commit
- Do not link against CUDA driver when building (#1240) · 86f07be4
  Tim Moon authored Oct 14, 2024
```
Signed-off-by: Tim Moon <tmoon@nvidia.com>
```
  86f07be4
12 Oct, 2024 1 commit

[PyTorch] Let Fused RoPE support CP with THD format (#1238) · 55dcbb4b

Xin Yao authored Oct 12, 2024



* Let Fused RoPE support THD with CP
Signed-off-by: Xin Yao <xiny@nvidia.com>

* add comment
Signed-off-by: Xin Yao <xiny@nvidia.com>

---------
Signed-off-by: Xin Yao <xiny@nvidia.com>
Co-authored-by: Xiaowei Ren <103958965+xrennvidia@users.noreply.github.com>

55dcbb4b

11 Oct, 2024 2 commits

Add FlashAttention3 to CP implementations (#1232) · b36bd0a4

Xiaowei Ren authored Oct 11, 2024



* fa2 function import renaming
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* refine fa_fwd_kwargs and fa_bwd_kwargs
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* import FA3 fucntions for CP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix output of FA3 fwd
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix rng_state in a2a implementation with FA3
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* hack lse correction for packed lse format
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* make CP thd out correction work with packed lse
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

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

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* change lse_packed to constexpr
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



---------
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

b36bd0a4

Fix bug in torch compile and seqdim is integer (#1217) · 9ee2dbdd

李金梁 authored Oct 12, 2024



* Fix bug in torch compile and seqdim is integer
Signed-off-by: 李金梁 <975761915@qq.com>

* Update attention.py

change the jit_fuser to torch.compile on flash_attn_fwd_out_correction
Signed-off-by: 李金梁 <975761915@qq.com>

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

---------
Signed-off-by: 李金梁 <975761915@qq.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

9ee2dbdd

10 Oct, 2024 2 commits

Small fixes to Float8Tensor (#1225) · 3b89c36f

Przemyslaw Tredak authored Oct 10, 2024



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

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



---------
Signed-off-by: Przemyslaw Tredak <ptredak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

3b89c36f

[JAX] Expose sliding window attn to TE-JAX API (#1205) · 85e60e64

Hua Huang authored Oct 10, 2024



* Expose JAX sliding window attn API
Signed-off-by: Hua Huang <huah@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* No SWA in context parallel; fix RNG seed in test
Signed-off-by: Hua Huang <huah@nvidia.com>

* Handle SAW API discrepancy in cuDNN and Python
Signed-off-by: Hua Huang <huah@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Add SAW API for flax, all tests passed

Will update tests/jax/test_praxis_layers.py next
Signed-off-by: Hua Huang <huah@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Update test_praxis_layers.py for SWA, test passed
Signed-off-by: Hua Huang <huah@nvidia.com>

* Use tuple window_size; update for PR #1212
Signed-off-by: Hua Huang <huah@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Add and adjust some pytest.skip
Signed-off-by: Hua Huang <huah@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Revised following Reese Wang's comments

Still need further debugging:
FAILED test_fused_attn.py::TestFusedAttn::test_backward[NO_SWA-DROP_0.0-4-128-256-16-16-64-BF16-CROSS-KV_PACKED-NO_MASK-NO_BIAS] - AssertionError:
FAILED test_fused_attn.py::TestFusedAttn::test_backward[NO_SWA-DROP_0.0-4-128-256-16-16-64-BF16-CROSS-KV_PACKED-NO_MASK-POST_SCALE_BIAS-1HSS] - AssertionError:
FAILED test_fused_attn.py::TestFusedAttn::test_backward[NO_SWA-DROP_0.0-4-128-256-16-16-64-BF16-CROSS-SEPARATE-NO_MASK-NO_BIAS] - AssertionError:
FAILED test_fused_attn.py::TestFusedAttn::test_backward[NO_SWA-DROP_0.0-4-128-256-16-16-64-BF16-CROSS-SEPARATE-NO_MASK-POST_SCALE_BIAS-1HSS] - AssertionError:

These errors does not exist in the previous commit
Signed-off-by: Hua Huang <huah@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Fix no-SWA test case errors in previous commit
Signed-off-by: Hua Huang <huah@nvidia.com>

* Add Padding mask w/ sliding windows sanity tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Use float32 for the reference code softmax calculation
Signed-off-by: Reese Wang <rewang@nvidia.com>

---------
Signed-off-by: Hua Huang <huah@nvidia.com>
Signed-off-by: Reese Wang <rewang@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Reese Wang <rewang@nvidia.com>

85e60e64

09 Oct, 2024 3 commits

[PyTorch] Improve `get_qkv_layout` (#1214) · 5b6546c8

Charlene Yang authored Oct 09, 2024



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

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



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

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



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

---------
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

5b6546c8

[PyTorch] Add documentation for FP8 attention checkpointing (#1223) · 2d875521

Charlene Yang authored Oct 09, 2024



* add extra_state change description for different TE versions
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

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

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



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

---------
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

2d875521

[PyTorch] Debug dtype casting in operation-based API (#1202) · 5b89f1ad

Tim Moon authored Oct 08, 2024



* Handle Float8Tensor when casting module dtype

Keep data in Float8Tensor and only change nominal dtype. Monkey-patch PyTorch module casting functions to handle Float8Tensor. Add tests.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Respect autocast dtype in linear op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



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

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

* Tweak comments

Review suggestion from @ptrendx
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

5b89f1ad

08 Oct, 2024 1 commit

[PyTorch] Miscellaneous fixes for FA3 attention (#1174) · e762592e

Charlene Yang authored Oct 08, 2024



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

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



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

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* force the same dtype when comparing FA3 and cuDNN FP8
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* Revert "force the same dtype when comparing FA3 and cuDNN FP8"

This reverts commit 19e7f877026a19a32d2f02c6c9de20df4ae2e064.
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* force the same dtype when comparing FA3 and cuDNN FP8
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add try/except for FA3 when custom qkv descales are not supported
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* replace FA3 installation warning with a debug logging message
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



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

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



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

* avoid varlen_func for FP8 and improve messaging
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



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

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

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* change preference reason for FP8 logic
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



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

---------
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

e762592e

07 Oct, 2024 2 commits

Fix cuDNN sliding window size (#1212) · c3b3cd21

Charlene Yang authored Oct 07, 2024



* adjust window size to (i-window_size_left,i] for cuDNN
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* reduce the window to make any errors more pronouced
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

c3b3cd21

Hierarchical CP implementation (Ulysses + Ring) (#1209) · c24a4c41

Xiaowei Ren authored Oct 07, 2024



* change API for hierarchical CP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* move fp8 code before qkv reshape
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* try to insert A2A for hierarchical CP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* make fwd work
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* remove a redundant sync
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* make bwd of hierarchical CP work
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix dout a2a in bwd
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

* assert hierarchical CP implementation does not support THD format
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

* assert hierarchical CP does not support attn bias
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add unit test for hierarchical CP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix cp_comm_type in unit test
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* bug fix and code cleaning
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* minor change
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* an assert info change
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* move function definitions to the front of the first call
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

* refine CP unit test
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



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

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

* save cp_size_a2a and rank_a2a in fwd
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add more explainations of cp_group in doc_string
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

---------
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

c24a4c41

04 Oct, 2024 1 commit

[PyTorch] Minor optimizations to reduce CPU overheads in modules (#1191) · 9d976bcd

Tim Moon authored Oct 03, 2024



* CPU perf optimization in linear autograd function

Avoid enable_grad context when possible in cast function. Cache distributed group properties.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* CPU perf optimization in prepare_forward function

Avoid torch.nn.Module impl of __setattr__.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Avoid module import in TE module forwards
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use fast getter for params
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Reuse tensor dims in linear autograd func
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Apply optimizations to grouped linear
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



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

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

* Avoid deepcopy in tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Move _fast_setattr logic to __setattr__ method
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

9d976bcd

03 Oct, 2024 1 commit

[PyTorch] Move `block_table` argument to FA varlen function (#1222) · 10cceae9

Charlene Yang authored Oct 03, 2024



move block_table arg to varlen_func section
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

10cceae9

01 Oct, 2024 2 commits

Removed the unused options from GroupedLinear docs and fixed the bug with offsets (#1220) · fb749619

Przemyslaw Tredak authored Sep 30, 2024



* Removing the unused options from GroupedLinear docs and fixing the bug
with offsets
Signed-off-by: Przemyslaw Tredak <ptredak@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



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

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



---------
Signed-off-by: Przemyslaw Tredak <ptredak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

fb749619

[PyTorch] Add pool argument to make_graphed_callable (#1218) · 728c558b
Kirthi Shankar Sivamani authored Sep 30, 2024
```
Add pool argument to make_graphed_callable
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
728c558b

27 Sep, 2024 2 commits

[PyTorch] Fix detection of 3 in 3hd/h3d layouts (#1187) · 8a1b7ee2

Charlene Yang authored Sep 27, 2024



* fix detection of 3 in 3hd/h3d layouts
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* error out when invalid layout group is provided
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



---------
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

8a1b7ee2

[PyTorch] Add GroupedLinear to the docs and fix typos (#1206) · c4a5cb85

Paweł Gadziński authored Sep 27, 2024



* Docs fixes
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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

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

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

c4a5cb85

20 Sep, 2024 1 commit

Restore compatibility with Python 3.8 (#1189) · 0c74535e

Przemyslaw Tredak authored Sep 20, 2024



* Restore compatibility with Python 3.8
Signed-off-by: Przemyslaw Tredak <ptredak@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



---------
Signed-off-by: Przemyslaw Tredak <ptredak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

0c74535e

19 Sep, 2024 1 commit

[PyTorch] Relax the contiguous check for flash attention (#1176) · 0ee5ccda

Xin Yao authored Sep 19, 2024



* relax contiguous check for flash attention
Signed-off-by: Xin Yao <xiny@nvidia.com>

* force contiguous for cp
Signed-off-by: Xin Yao <xiny@nvidia.com>

---------
Signed-off-by: Xin Yao <xiny@nvidia.com>

0ee5ccda

18 Sep, 2024 2 commits

Expose `rotary_base` as an arg instead of hardcoding (#944) · c0caadbe

Sudhakar Singh authored Sep 18, 2024



* make rotary_base arg
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* rotary base can be a float
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

---------
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

c0caadbe

[PyTorch] Check network interface name when initializing Userbuffers (#1175) · 841634ca

Alp Dener authored Sep 18, 2024



* Check if network interface name is valid and show useful warning message when initializing Userbuffers
Signed-off-by: Alp Dener <adener@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Fix formatting issue in warning message.
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Alp Dener <adener@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



---------
Signed-off-by: Alp Dener <adener@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

841634ca

17 Sep, 2024 2 commits

[Common] Default CUDA_HOME to /usr/local/cuda when dynamically loading cuDNN and NVRTC (#1183) · 44fd316f

Alp Dener authored Sep 17, 2024



Defaulted CUDA_HOME/CUDA_PATH to /usr/local/cuda when attempting to dynamically load cuDNN and NVRTC
Signed-off-by: Alp Dener <adener@nvidia.com>

44fd316f

[JAX] Context Parallel Attention with All-Gather (#1106) · 9101a78f

Michael Goldfarb authored Sep 17, 2024



Implementation of context parallel fused attention using all-gather.
Signed-off-by: Michael Goldfarb <mgoldfarb@nvidia.com>

9101a78f

16 Sep, 2024 1 commit

Add dtensor support for TE optimizers (#1171) · af5daa09

Jason Wang authored Sep 16, 2024



add dtensor support for te optimizers
Signed-off-by: jasonwan <jasonwan@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

af5daa09

11 Sep, 2024 1 commit

[PyTorch] Proxy class for low-precision tensor (#1127) · 2d57db8b

Tim Moon authored Sep 11, 2024



* Add base class for tensor proxies
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Move tensor detaching logic to tensor proxy base class
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use Python wrappers to PyTorch extensions
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Include transpose caching logic in proxy encode function
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug dimension mismatch with amax history
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Move dequantize logic to proxy_decode func
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Rename to "QuantizedTensor"
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Rename "proxy_detach" to "detach"
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Include transpose cache in detach and clone funcs
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Update FP8 workspaces with QuantizedTensor functions
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Move logic for FP8 transpose cache in FP8 workspaces to base class
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Remove cast-transpose logic from linear op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Remove unnecessary args for Float8Tensor when using FP8 attr dict
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Remove __torch_function__ to QuantizedTensor
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

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

* Debug FP8 transpose test
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

2d57db8b

09 Sep, 2024 2 commits

Add a context parallelism implementation with QKVO all-to-all (#1160) · 40dda924

Xiaowei Ren authored Sep 09, 2024



* clean code for CP function args
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add a placeholder for Ulysses implementation
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* commit code change to CP+A2A
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* finish the draft fwd implementation of Ulysses
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add draft bwd implementation of Ulysses
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* make swa work with ulysses
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* commit FP8 code for Ulysses
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix qkv type in the bwd of FP8+CP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

* fix qkv_dtype of FP8+CP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* code refactoring
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* minor code change
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* config cp correction dtype of FP8+CP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* code style change
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* save chunk_ids
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* try to make Ulysses A2A async
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* make more a2a async
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

* fix chunk_ids generation for A2A
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* avoid code duplication of a2a before attn
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* remove code duplication of a2a after attn
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add cp_stream in A2A implementation
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

* fix qkv of fp8_fwd + bf16_bwd
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix kernel order in cp a2a communication
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* code cleaning for CP a2a
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

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

* adjust sequence chunk reordering for a2a
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add docstring for A2A implementation
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* change an assert info
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add unit tests of A2A implementation
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add more A2A unit test
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

* add more cp unit tests
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix window size of no_mask
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fused attn does not support swa+no_mask
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* change num_gqa_groups to 2 for A2A implementation
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* function and variable renaming
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* code cleaning for CP all-gather implementation
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* some function renaming
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

* commit code change for kv all-gather implementation
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

* add a window size check
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* code cleaning
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add unit test of all_gather+no_mask
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix all-gather cp implementation
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* code cleaning
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



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

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

* fix FP8 with A2A implementation
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* add paper references to CP implementations with all-gather and all-to-all
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* change pdf to abs
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* elaborate cp_comm_type
Signed-off-by: Xiaowei Ren <xren@nvidia.com>

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

---------
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

40dda924

[PyTorch] Propagate fp8 scale-inverse modification to `GroupedLinear` (#1128) · 047a5072

Xin Yao authored Sep 09, 2024



* propagate scale_inv modification to GroupedLinear
Signed-off-by: Xin Yao <xiny@nvidia.com>

* optimization for separate scale_inv of weights and single output
Signed-off-by: Xin Yao <xiny@nvidia.com>

* let grouped gemm support different input combinations
Signed-off-by: Xin Yao <xiny@nvidia.com>

* fix type
Signed-off-by: Xin Yao <xiny@nvidia.com>

* add contiguous check
Signed-off-by: Xin Yao <xiny@nvidia.com>

* use len() instead of isinstance
Signed-off-by: Xin Yao <xiny@nvidia.com>

* fix ut
Signed-off-by: Xin Yao <xiny@nvidia.com>

---------
Signed-off-by: Xin Yao <xiny@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

047a5072

05 Sep, 2024 4 commits

Revert "[C] Suppress 128-D warning from cudnn-frontend" (#1161) · bdea56fc

Kirthi Shankar Sivamani authored Sep 05, 2024

Revert "[C] Suppress 128-D warning from cudnn-frontend (#1158)"

This reverts commit 206c1d92

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

bdea56fc

[C] Suppress 128-D warning from cudnn-frontend (#1158) · 206c1d92

Charlene Yang authored Sep 05, 2024



suppress 128D warning from cudnn-frontend
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

206c1d92

[PyTorch] Implement Fp8 padding and unpadding module (#1129) · 215db88d

Ruibin Cheung authored Sep 06, 2024



* [TE/PyTorch][MoE] Add FP8 padding and unpadding module 

 1. Add multi-tensor padding kernel for FP8 with padding size = 16.
 2. Add FP8Padding and Fp8Unpadding module
 3. Add Padded GroupedLinear unit tests

---------
Signed-off-by: beinggod <zhangruibin@01.ai>
Co-authored-by: Phuong Nguyen <36155692+phu0ngng@users.noreply.github.com>

215db88d

Added offloading support FP8 attention (#1131) · 454e3895

Selvaraj Anandaraj authored Sep 05, 2024



* Added offloading support FP8 attention
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos02.eos.clusters.nvidia.com>

* Update transformer_engine/pytorch/attention.py
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Selvaraj Anandaraj <anandaraj@wisc.edu>

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

---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos02.eos.clusters.nvidia.com>
Signed-off-by: Selvaraj Anandaraj <anandaraj@wisc.edu>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-eos02.eos.clusters.nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

454e3895