Commits · 958e1889de8df3d7fe9c5e20ddbf4de556871c30 · OpenDAS / TransformerEngine

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

26 Jun, 2023 1 commit
- Fix compute type for GEMM (#296) · c56646e4
  Kirthi Shankar Sivamani authored Jun 26, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  c56646e4
19 Apr, 2023 1 commit

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

Kirthi Shankar Sivamani authored Apr 19, 2023



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

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

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

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

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

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

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

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

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

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

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

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

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

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

* Fix gencode flags in setup (#145)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

18da4e88

28 Mar, 2023 1 commit

Add support for fp8 GEMM BIAS AUX GELU fusion (#116) · 2d72c11f

vasunvidia authored Mar 28, 2023



* Add support for fp8 GEMM BIAS AUX GELU fusion
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

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

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

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

2d72c11f

03 Feb, 2023 1 commit

Increase number of FP8 tensors per GEMM (#22) · 14198f20

vasunvidia authored Feb 03, 2023



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

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

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

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

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

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

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

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

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

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

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

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

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

14198f20

12 Jan, 2023 1 commit

Add NVTX to TE modules (#50) · aadd3e7c

Przemyslaw Tredak authored Jan 12, 2023



* Add NVTX to TE modules
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

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

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

* Add NVTX to C API
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

* Add NVTX to GeGlu
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

aadd3e7c

03 Jan, 2023 1 commit

Update copyright year (#48) · 64a8dc90

Przemyslaw Tredak authored Jan 03, 2023


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

64a8dc90

08 Dec, 2022 1 commit

Move the amax/scale/scale_inv into the TE Tensor struct. (#33) · a5ba71f3

Przemyslaw Tredak authored Dec 08, 2022



* Move the amax/scale/scale_inv into the TE Tensor struct.
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

* Changed softmax to new Tensor
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* First pass at the cpp tests
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

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

* Fix cast_to_fp8
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Przemyslaw Tredak <ptrendx@gmail.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

a5ba71f3

28 Sep, 2022 1 commit

Inital code drop · 996ea169

Przemek Tredak authored Sep 27, 2022


Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

996ea169