Commits · 277b0be27fb2bf2a40572d1768be638af3e515ed · OpenDAS / TransformerEngine

24 Feb, 2023 1 commit

add building workflow for TE/Jax (#53) · a3ec6a54

Jeng Bai-Cheng authored Feb 25, 2023



* add building workflow for jax modules
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* replace bit_cast with reinterpret_cast
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* add nvtx to cmake check list
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* refactor layernorm fwd
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* refactor rmsnorm fwd
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* refactor layernorm_bwd
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* set pytorch as default in setup.py
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* rename extension from *.cc to *.cpp

cpplint cannot recognize *.cc file, so rename the extension
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* refactor style, to align TE/PyTorch
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* add pybinding, unittest and qa
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* fix license
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* disable c-extension-no-member and no-name-in-module
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* add dataclass avoid pylint error
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* Update transformer_engine/__init__.py
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>

* Update tests/jax/test_custom_call_shape.py

fix typo
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>

* Update tests/jax/test_custom_call_shape.py
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>

* add building workflow for jax modules
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* replace bit_cast with reinterpret_cast
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* add nvtx to cmake check list
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* refactor layernorm fwd
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* refactor rmsnorm fwd
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* refactor layernorm_bwd
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* set pytorch as default in setup.py
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* rename extension from *.cc to *.cpp

cpplint cannot recognize *.cc file, so rename the extension
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* refactor style, to align TE/PyTorch
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* add pybinding, unittest and qa
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* fix license
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* disable c-extension-no-member and no-name-in-module
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* add dataclass avoid pylint error
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* Update transformer_engine/__init__.py
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>

* Update tests/jax/test_custom_call_shape.py

fix typo
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>

* Update tests/jax/test_custom_call_shape.py
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>

* fix conflict due to PR62
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* fix c-extension-no-member and no-name-in-module

1. add transformer_engine_jax into extension-pkg-whitelist
2. convert pylintrc from CRLF to LF format
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* Update setup.py
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>

* remove pylint:disable and refactor import order
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

---------
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

a3ec6a54

15 Feb, 2023 1 commit

Zero-centered gamma support in LayerNorm (LayerNorm1p) (#67) · 7324fe2b

Przemyslaw Tredak authored Feb 14, 2023



* C++ implementation of LayerNorm1P
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Expose zero centered gamma to pyTorch
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

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

* Fix ONNX export and tests
Signed-off-by: Przemyslaw Tredak <ptredak@nvidia.com>

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

* Fix backward handling - C++ part
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

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

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

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

* Change in the bwd kernel
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

7324fe2b

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

27 Jan, 2023 1 commit
- Fix the integer overflow in fused softmax (#60) · f63b27e8
  Przemyslaw Tredak authored Jan 27, 2023
```
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
```
  f63b27e8
19 Jan, 2023 1 commit

Fix NVTX name for LN backward (#55) · f22929cc

Kirthi Shankar Sivamani authored Jan 19, 2023


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

f22929cc

17 Jan, 2023 1 commit

Move calculation of scale inverse to framework (#51) · 02a3582c

Kirthi Shankar Sivamani authored Jan 17, 2023



* Move scale inverse calculation to framework
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

02a3582c

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

10 Jan, 2023 1 commit

Add GeGLU and the corresponding gradient kernels (#47) · eed1fa26

zlsh80826 authored Jan 11, 2023



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

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

* Update copyright year
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

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

eed1fa26

09 Jan, 2023 1 commit

Add RMSNorm (#45) · 37cc3625

zlsh80826 authored Jan 10, 2023



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

* Add rmsnorm cpp unit test
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Apply new Tensor struct
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Move scale/scale_inv/amax into the TE Tensor struct
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Separate rmsnorm kernels from the layernorm
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Update rmsnorm test cases
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Update copyright year
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix the support matrix on the document
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Move register macro out of utils.cuh
Signed-off-by: Reese Wang <rewang@nvidia.com>
Signed-off-by: Reese Wang <rewang@nvidia.com>

37cc3625

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

06 Dec, 2022 1 commit

Softmax docstrings and type fixes (#37) · 89f94ba2

Kirthi Shankar Sivamani authored Dec 05, 2022



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

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

* change API, better naming, const fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

89f94ba2

01 Dec, 2022 1 commit

Framework agnostic softmax kernels (#30) · f674d49e

Kirthi Shankar Sivamani authored Dec 01, 2022



* Make fused softmax kernels PyTorch independent
Co-authored-by: Sean Lee <selee@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* move get_batch_per_block to python
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fix license in softmax.h
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Sean Lee <selee@nvidia.com>

f674d49e

30 Nov, 2022 1 commit

Fix illegal memory access in general layer norm backward kernel (#34) · 693d6541

Tim Moon authored Nov 30, 2022



Fix illegal memory access in layernorm backward kernel
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>

693d6541

28 Nov, 2022 1 commit

Multi-tensor cast-transpose (#18) · 2a3916b4

Tim Moon authored Nov 28, 2022



* Add kernel for multi-tensor cast-transpose
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix incorrect test function in multi-tensor cast-transpose unit test
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Remove std::vector from multi-tensor cast-transpose function signature

Makes sure the main header is C-compatible.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Przemyslaw Tredak <ptredak@nvidia.com>

2a3916b4

17 Nov, 2022 1 commit

Make amax reduction optional (#11) · 681bf9ad

Kirthi Shankar Sivamani authored Nov 16, 2022



* Make amax reduction optional
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* remove setup for global amax redux for optional case
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

* Making checkpointing backwards compatible
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add deprecation warning for old checkpoint loading
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fix checkpointing for fp8 recompute case
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* improvements to deprecation warning
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Przemyslaw Tredak <ptredak@nvidia.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>

681bf9ad

08 Nov, 2022 1 commit
- Fix the out-of-bounds access in the C+T+dbias kernel (#28) · acf98b5c
  Przemyslaw Tredak authored Nov 08, 2022
```
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
```
  acf98b5c
12 Oct, 2022 1 commit

Remove fp8_out from the LN API (#8) · 6a2161bf

Przemyslaw Tredak authored Oct 12, 2022



* Remove fp8_out from LN API
Signed-off-by: Przemyslaw Tredak <ptredak@nvidia.com>

* fix LN test
Signed-off-by: Przemyslaw Tredak <ptredak@nvidia.com>

* Fixes
Signed-off-by: Przemyslaw Tredak <ptredak@nvidia.com>
Signed-off-by: Przemyslaw Tredak <ptredak@nvidia.com>
Co-authored-by: ksivamani <ksivamani@nvidia.com>

6a2161bf

05 Oct, 2022 1 commit

fix documentation for default FP8 recipe format (#6) · 1568c6ed

Kirthi Shankar Sivamani authored Oct 05, 2022


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

1568c6ed

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