Commits · 14c1ecd09c86a2bbab4702e15ea19fc0da0e05b4 · OpenDAS / TransformerEngine

19 Apr, 2024 1 commit

NVRTC kernels for cast-transpose (#258) · 14c1ecd0

Tim Moon authored Apr 19, 2024



* Add NVRTC kernels for cast-transpose
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Add noop flag to NVRTC cast-transpose kernel
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Apply suggestions from code review
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>

14c1ecd0

18 Apr, 2024 2 commits

[JAX] Fixing CI failure due to incorrect use of `static_argnums` in jax.jit (#785) · c63766d4

Alp Dener authored Apr 18, 2024



* fixed static argnums for jax.jit in single gpu encoder test, changed warning filtering for pytest
Signed-off-by: Alp Dener <adener@nvidia.com>

* propagating the fix to the JAX mnist example
Signed-off-by: Alp Dener <adener@nvidia.com>

* fixed missing space ibetween flags i QAA scripts
Signed-off-by: Alp Dener <adener@nvidia.com>

* added TE warnings into the ignore list
Signed-off-by: Alp Dener <adener@nvidia.com>

---------
Signed-off-by: Alp Dener <adener@nvidia.com>

c63766d4

[PyTorch] Fix for type checking failure on custom callables (#790) · fad0e273

Alp Dener authored Apr 17, 2024



fix type checking in checkpointing to assume that there must be TE modules in custom callables
Signed-off-by: Alp Dener <adener@nvidia.com>

fad0e273

17 Apr, 2024 2 commits

[UB] Adding configurable timeout for userbuffer and improving error reporting... · 08e5e4b1

Pavel Shamis (Pasha) authored Apr 17, 2024


[UB] Adding configurable timeout for userbuffer and improving error reporting for potential hangs (#757)

* Improving error reporting and hang detection logic

* Adding verbose error reporting in case of UB hang
* Adding CE hang detector
* Replacing hard-coded timeout with configurable one
Signed-off-by: Pasha (Pavel) Shamis <pasharesearch@gmail.com>

* Cleaning up warnings in the code
Signed-off-by: Pasha (Pavel) Shamis <pasharesearch@gmail.com>

* Removing unused codes
Signed-off-by: Pasha (Pavel) Shamis <pasharesearch@gmail.com>

* Fixing styling issues reported on github
Signed-off-by: Pasha (Pavel) Shamis <pasharesearch@gmail.com>

* Addressing lint new line and casting warnings
Signed-off-by: Pasha (Pavel) Shamis <pasharesearch@gmail.com>

* Addressing lint warning about the usage of `unsigned long long`
Signed-off-by: Pasha (Pavel) Shamis <pasharesearch@gmail.com>

* Removing unused case causing build issues on multi-arch setup
Signed-off-by: Pasha (Pavel) Shamis <pasharesearch@gmail.com>

* Post GRDCOPY removal cleanup

* Remove cmake check
* Remove unused includes
Signed-off-by: Pasha (Pavel) Shamis <pasharesearch@gmail.com>

---------
Signed-off-by: Pasha (Pavel) Shamis <pasharesearch@gmail.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

08e5e4b1

[PyTorch] Misc fixes for release_v1.6 (#784) · cd54a8cd

Kirthi Shankar Sivamani authored Apr 17, 2024



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

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

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

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

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

cd54a8cd

16 Apr, 2024 6 commits

Support Low Rank Adaptation (LoRA). (#745) · 7c1828f8
Ming-Xu Huang authored Apr 16, 2024

7c1828f8
Add new users to TE CI · 1442b47e
Przemek Tredak authored Apr 16, 2024
```
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
```
1442b47e

[PyTorch] TE checkpoint pass-through logic fix (#782) · f96f3407

Alp Dener authored Apr 16, 2024



* changed TE checkpoint passthrough logic to also recursively look for TE submodules
Signed-off-by: Alp Dener <adener@nvidia.com>

* simplified search for TE modules in the checkpointed network
Signed-off-by: Alp Dener <adener@nvidia.com>

---------
Signed-off-by: Alp Dener <adener@nvidia.com>

f96f3407

[PyTorch] Use __torch_function__ as a class method (#783) · d3552ddb
Kirthi Shankar Sivamani authored Apr 16, 2024
```
Use torch function as a class method
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
d3552ddb
Changed VERSION to 1.7.0dev · 63c7a1a3
Kirthi Shankar Sivamani authored Apr 15, 2024
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
63c7a1a3

[C/PyTorch] Add FP8 DPA and MHA (#768) · 83a4c219

cyanguwa authored Apr 15, 2024



* WIP: fp8 v1 fprop integration
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

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

* fprop working for h1; w/ debug info
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* cleanup; bprop running but has mismatches
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* clean up and add back v0.9.2 FE support; fprop/bprop passing with 5e-2 tols
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix after merge; add bias_b/h to caching descriptor
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* distinguish fwd/bwd tensor types for bprop
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* minor fix for F16 cases; include added dqkv_type and d_scale_dp
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* adjust out shape for bwd in test
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add casting from/to FP8 to DPA module
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* WIP: support all sbhd/bshd layouts
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* add qkvpacked and kvpacked support in both FusedAttnFunc and C levels
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove qkvpacked/kvpacked calls in DPA module (used for testing)
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove tp setup; add allow_non_contiguous; update FE; revert to sbh3d in tests; clean up
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add NVTE_FP8_DPA_BWD to control whether to use FP8 bwd or F16 bwd
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* fix MQA/GQA in FP8 v1 API
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update FE to 705d8e3, with API change
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* restrict mha_fill for THD format
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix fused attn with CP and comment out is_alibi code
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* clean up FE0.9 vs FE1.0 FP8 implementations, and related unit tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* change NVTE_FP8_DPA_BWD default to 1, and fix its use in qkvpacked/kvpacked APIs
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix lint and self.tp_size/group in FusedAttention()
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update FE to 6902c94
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* update to FE v1.3.0
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* minor fixes for FP8 MHA with different configs
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* emit stats regardless of is_training
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix linear when input is not Float8Tensor
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix d_out type when f16 bprop
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix user buffer for layernorm_linear/linear and revert two FP8 casts in MHA
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add docstring for fp8_dpa/mha in recipe
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

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

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

* replace two more transpose with transpose_2d
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* Revert "add FP8 initialization to FusedAttention"

This reverts commit 15fffd825d6f23f31ea709b16ba01dfd61efabf8.
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Change order of ctxs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* add back docs and mark as beta
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

83a4c219

15 Apr, 2024 1 commit
- [PyTorch] Don't use autograd hook for bwd reduction (#781) · f69e45be
  Kirthi Shankar Sivamani authored Apr 15, 2024
```
Don't use autograd hook for bwd reduction
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  f69e45be
12 Apr, 2024 3 commits

Add SM margin to LayerNorm in inference (#772) · 5d34b2ac

Sangkug Lym authored Apr 12, 2024



* Add LN margin to inference
Signed-off-by: Sangkug Lym <slym@nvidia.com>

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

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

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

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

5d34b2ac

[PyTorch] Fix kernel_bulk launch config (#775) · b4ef463d
Kirthi Shankar Sivamani authored Apr 12, 2024
```
Fix 0 grid size
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
b4ef463d

[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

10 Apr, 2024 1 commit
- Fix undefined symbol issue for transformer_engine::getenv (#763) · 1b20f2d6
  Jinze Xue authored Apr 10, 2024
```
Signed-off-by: Jinze Xue <jinzex@nvidia.com>
Co-authored-by: Jinze Xue <jinzex@nvidia.com>
```
  1b20f2d6
06 Apr, 2024 3 commits

[JAX] Adapt latest JAX/PAX image (#744) · bfe21c3d

Reese Wang authored Apr 07, 2024



* value_and_grad requires same shape for input and gradients
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Use high precision layernorm
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove local_device_ids as it caused unexpected behaviors
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Revert "Remove local_device_ids as it caused unexpected behaviors"

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

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

bfe21c3d

Fix the default userbuffer communicator init settings (#755) · d541d208
Sangkug Lym authored Apr 05, 2024
```
fix the default userbuffer communicator init settings
Signed-off-by: Sangkug Lym <slym@nvidia.com>
```
d541d208

Enable DGRAD RS overlap (#754) · e3de4037

Jaemin Choi authored Apr 05, 2024



* Enable DGRAD RS overlap
Signed-off-by: Jaemin Choi <jaeminc@nvidia.com>

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

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

e3de4037

04 Apr, 2024 4 commits

userbuffer: support fp8 buffer for individual overlap instance (#750) · 7d8ef9bf

Sangkug Lym authored Apr 04, 2024



* userbuffer fp8 reduction support for individual overlap
Signed-off-by: Sangkug Lym <slym@nvidia.com>

* cleanup dict ub_cfg dict value load
Signed-off-by: Sangkug Lym <slym@nvidia.com>

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

* Remove unnecessary fence from producer

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

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

7d8ef9bf

Compile tuned RMSNorm kernels for hidden size 8192 (#747) · 27cb9950
Tim Moon authored Apr 04, 2024
```
Signed-off-by: Tim Moon <tmoon@nvidia.com>
```
27cb9950

[PyTorch] Fix backward compatibility for checkpoint API (#748) · ce328fac

Kirthi Shankar Sivamani authored Apr 04, 2024



* Args can be None
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

ce328fac

Fixing potential integer overflow on sequence counter (#729) · e1e2b76e

Pavel Shamis (Pasha) authored Apr 03, 2024



* Fixing potential integer overflow on sequence counter

Current implementation may potential cause hangs or data corruption
Signed-off-by: Pasha (Pavel) Shamis <pasharesearch@gmail.com>

* Fixing typo in comments

Addressing reviewers comments
Signed-off-by: Pasha (Pavel) Shamis <pasharesearch@gmail.com>

---------
Signed-off-by: Pasha (Pavel) Shamis <pasharesearch@gmail.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

e1e2b76e

03 Apr, 2024 5 commits

Update README.rst (#733) · 1fa5bf18

Santosh Bhavani authored Apr 03, 2024



* Update README.rst

1. Updated latest news with databricks blog
2. Fixed formatting issues
3. Added GTC 2024 video
Signed-off-by: Santosh Bhavani <santosh@semantic.md>

* Update README.rst

added back overview marker for docs generation
Signed-off-by: Santosh Bhavani <santosh@semantic.md>

* Added MPT-13B convergence result
Signed-off-by: Santosh Bhavani <santosh@semantic.md>

* Added Levanter/JAX to integrations section of README
Signed-off-by: Santosh Bhavani <santosh@semantic.md>

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

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

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

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

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

---------
Signed-off-by: Santosh Bhavani <santosh@semantic.md>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

1fa5bf18

Atomic gemm for TP-AR and TP-RS overlap with P2P exchanges (#732) · 180de056

Sangkug Lym authored Apr 02, 2024



* Atomic gemm for TP-AR and TP-RS overlap with P2P exchanges
Signed-off-by: Sangkug Lym <slym@nvidia.com>

* FP8 reduction for atomic TP-RS with p2p exchange
Signed-off-by: Sangkug Lym <slym@nvidia.com>

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

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

180de056

Do not store input activations when not computing weight gradients (#739) · 580eb52b

Sangkug Lym authored Apr 02, 2024



* Do not store input activations when not computing weight gradients
Signed-off-by: Sangkug Lym <slym@nvidia.com>

* fix userbuffer tp comm overlap case
Signed-off-by: Sangkug Lym <slym@nvidia.com>

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

580eb52b

Add NVLS-MC based UB kernels (#721) · 8e7795e1

vasunvidia authored Apr 02, 2024



Fix license, and sign off everything
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

8e7795e1

Revert "Update FA version to 2.5.6 (#714)" · 47276e1b
Kirthi Shankar Sivamani authored Apr 02, 2024
```
This reverts commit 965803c9.
```
47276e1b

02 Apr, 2024 1 commit
- Set CUDA context before loading NVRTC kernels (#734) · 2dd6b146
  Tim Moon authored Apr 02, 2024
```
Signed-off-by: Tim Moon <tmoon@nvidia.com>
```
  2dd6b146
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

29 Mar, 2024 2 commits

[PyTorch] Fix backward compatibility with checkpoint API (#740) · 12cbd863

Kirthi Shankar Sivamani authored Mar 28, 2024



* Fix backward compatibility with checkpoint API
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

12cbd863

[PyTorch] Fix bug in FP8 cast in LayerNormLinear/LayerNormMLP (#738) · df1b16da
Tim Moon authored Mar 28, 2024
```
Perform FP8 cast on gathered layernorm output in LayerNormLinear
Signed-off-by: Tim Moon <tmoon@nvidia.com>
```
df1b16da

22 Mar, 2024 2 commits

Enable TP-AG overlap with return_layernorm_output (#727) · c1a68f6c

Jaemin Choi authored Mar 22, 2024



* Enable TP-AG overlap with return_layernorm_output
Signed-off-by: Jaemin Choi <jaeminc@nvidia.com>

* Use ub_overlap_ag
Signed-off-by: Jaemin Choi <jaeminc@nvidia.com>

---------
Signed-off-by: Jaemin Choi <jaeminc@nvidia.com>
Co-authored-by: Jaemin Choi <jaeminc@nvidia.com>

c1a68f6c

[JAX] Refactor fused attention (#711) · 8e672ff0

Reese Wang authored Mar 22, 2024



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

* Unify the fused attn workspace size cpp code
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Reduce the skipped cases
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Rename self/cross attention to qkvpacked/kvpacked
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Refine the attn mask implementations
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

8e672ff0

21 Mar, 2024 2 commits

TP-RS overlap with send/recv ring-exchange (#724) · b855656b

Sangkug Lym authored Mar 21, 2024



* TP-RS overlap with send/recv

Atomic GEMM based TP-RS overlap with send/recv
Signed-off-by: Sangkug Lym <slym@nvidia.com>

Specify userbuffer overlap method of each overlap instance
Signed-off-by: Sangkug Lym <slym@nvidia.com>

P2P TP-RS overlap with fp8 GEMM outputs
Signed-off-by: Sangkug Lym <slym@nvidia.com>

Fix TP-RS overlap with send/recv
Signed-off-by: Sangkug Lym <slym@nvidia.com>

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

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

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

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

---------
Signed-off-by: Sangkug Lym <slym@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

b855656b

[Pytorch] Update context parallel softmax lse correction func (#716) · 59bfc17b

Kite0011 authored Mar 21, 2024



[Pytorch] Update context parallel softmax lse correction func.
Signed-off-by: kitefang <kitefang@tencent.com>
Co-authored-by: kitefang <kitefang@tencent.com>

59bfc17b

20 Mar, 2024 2 commits

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

Update FA version to 2.5.6 (#714) · 965803c9
Kirthi Shankar Sivamani authored Mar 20, 2024
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
965803c9

18 Mar, 2024 1 commit
- Changed VERSION to 1.6.0dev · a3ba77b8
  Kirthi Shankar Sivamani authored Mar 18, 2024
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  a3ba77b8
15 Mar, 2024 1 commit

Rachitg/dp carveout (#722) · 1ec33ae1

Rachit Garg authored Mar 15, 2024



* fix the perf regression because of constant property polling of the device
Signed-off-by: Rachit Garg <rachitg@nvidia.com>

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

---------
Signed-off-by: Rachit Garg <rachitg@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Rachit Garg <rachitg@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>

1ec33ae1