GitLab

* initialize output tensors to 0 for THD while waiting for cuDNN bug fix
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



* move fill_() to F16 loop
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* correct typo in check_set_window_size
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* use nvtx3 instead
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>

* NVTE_OVERRIDE_MAX_SEQ_LEN
Signed-off-by: Frank Lin <eee4017@gmail.com>

* small fix
Signed-off-by: Frank Lin <eee4017@gmail.com>

* preserve old amax_and_scale_update_inplace and new amax_and_scale_update_inplace
Signed-off-by: Frank Lin <eee4017@gmail.com>

* remove useless code path; try to simplify logic within the baseline
Signed-off-by: Frank Lin <eee4017@gmail.com>

* simplify logic
Signed-off-by: Frank Lin <eee4017@gmail.com>

* small fix
Signed-off-by: Frank Lin <eee4017@gmail.com>

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

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



* fix comments from Timmoon
Signed-off-by: Frank Lin <eee4017@gmail.com>

* fix comments from Timmoon
Signed-off-by: Frank Lin <eee4017@gmail.com>

* Update transformer_engine/paddle/distributed.py
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Frank Lin <eee4017@gmail.com>

* diable bw fp8 update
Signed-off-by: Frank Lin <eee4017@gmail.com>

* fix lint
Signed-off-by: Frank Lin <eee4017@gmail.com>

* fix ci error
Signed-off-by: Frank Lin <eee4017@gmail.com>

---------
Signed-off-by: Frank Lin <eee4017@gmail.com>
Co-authored-by: Frank Lin (Engrg-Hardware 1) <fralin@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>

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

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

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

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

* support main_grad
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* update main_grad and fuse_wgrad_accumulation
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* fix ci errors
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

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

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

* use separate qkv
Signed-off-by: jaywan <jaywan@nvidia.com>

add support for GQA
Signed-off-by: jaywan <jaywan@nvidia.com>

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

change rtol
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

fix reshape issue
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

add rmsnorm and rotary position embedding
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

update rmsnorm
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

refactor layernorm and rmsnorm
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

support swiglu
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

add fused rope
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

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

add rope api to __init__
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

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

fix fp8 dtype issue
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* simplify ut cases
Signed-off-by: jaywan <jaywan@nvidia.com>

* Update transformer_engine/paddle/layer/attention.py
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Shijie <505749828@qq.com>

* fix name issue
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

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

* use separate qkv
Signed-off-by: jaywan <jaywan@nvidia.com>

* add support for GQA
Signed-off-by: jaywan <jaywan@nvidia.com>

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

* change rtol
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* fix reshape issue
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

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

* Actively free tensor in bwd
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* - Add inplace support for fp8 casting
- Allow skipping weight update in fp8 meta update
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

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

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

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

* Add weight caching for Transformer layer
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

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

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

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

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

* fix cudnn FA softmax shape
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* set inplace rng_state
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

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

* Deprecate QKV_INTERLEAVED use in JAX
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Deprecate QKV_INTERLEAVED use in Paddle
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enhance qkv enum mappings
Signed-off-by: rewang <rewang@nvidia.com>

* Fix LD_LIBRARY_PATH issue
Signed-off-by: rewang <rewang@nvidia.com>

* Arbitrary seqlen kernels only support self attention currently
Signed-off-by: rewang <rewang@nvidia.com>

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

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

* Add nn.layer: softmax, attention, transformer
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

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

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

* update docs and set dropout=0.1
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

* Update transformer_engine/paddle/layer/attention.py
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Shijie Wang <jaywan@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add FP8 support

- Add FP8 recipe
- Add FP8 path for nn layers
- Add MNIST FP8 example
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

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

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

* Fix FP8 training in float32 accumulation
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Fix FP8 checkpointing for non forward execution cases (same as #323)
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

* Refactors and improvements for better code stype, readability and organization
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

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

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

* add flash attn tests
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

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

* fix random seed
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

---------
Signed-off-by: Shijie Wang <jaywan@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* fix bug
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

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

* minor change on coding style
Signed-off-by: Shijie Wang <jaywan@nvidia.com>

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

* Add cast_transpose

Add gelu, gelu_fp8

Add cast_transpose_bgrad_dgelu

Add layernorm_fwd and layernorm_fwd_fp8

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

* Fix missing header
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>

* First step of PaddlePaddle integration
- Add build option for paddle
- Add basic test framework
- Add 3 basic operators: cast_from_fp8, cast_to_fp8, gemm
Signed-off-by: Tian Zheng <tizheng@nvidia.com>
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

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

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

* Add paddle build support for new building framework
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

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

* Clean up build process for Paddle stub file
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Fix pylint "wrong-import-order" warning
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

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

* Skip BF16 GEMM tests for unsupported arch
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

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