GitLab

* Fixed minor bug with DistributedConfigsHelper prematurely crashing the test for insufficient GPUs before @pytest.skip condition.
Signed-off-by: Alp Dener <adener@nvidia.com>

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

* Debug PyTest errors when running on single-GPU system
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: Alp Dener <adener@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>

[JAX] Regression tests for custom ops sharding with both xmap and custom_partitioning.

Coverage:
- layernorm: fwd/grad, zero_centered_gamma, DP, TP_COL, DP_TP_COL
- rmsnorm: fwd/grad, DP, TP_COL, DP_TP_COL
- softmax: fwd/grad, SCALED, SCALED_MASKED, SCALED_UPPER_TRIANG_MASKED, DP, TP_COL, TP_ROW, DP_TP_COL, DP_TP_ROW
- self_fused_attn: fwd/grad, NO_BIAS, PRE_SCALE_BIAS, POST_SCALE_BIAS, NO_MASK, CAUSAL_MASK, PADDING_MASK, DP, TP_COL, DP_TP_COL
- cross_fused_attn: fwd/grad, NO_BIAS, NO_MASK, PADDING_MASK, DP, TP_COL, DP_TP_COL
Signed-off-by: Alp Dener <adener@nvidia.com>

* Experimental FP8 tensor
Co-authored-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Sudhakar Singh <sudhakars@nvidia.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add fp8 tensor to ci test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* Default to FP8 usage
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

* Debug transpose caching

Handle case where transpose cache is updated externally.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Rename FP8GlobalStateManager.with_fp8_parameters
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* remove Float8Tensor from import API
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Avoid caching FP8 transposes if not required
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix import error in FP8 tensor tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix tranpose caching and checkpointing bug
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Improve caching and fix distopt case
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

* Store FP8 attributes in dict

Easier for multiple tensors to share, e.g. detached tensors.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Make sure scale_inv is 1D tensor
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Make sure scale_inv is 1D tensor
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Set default fp8 data type
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Sudhakar Singh <sudhakars@nvidia.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>

* Do not include logging macros in installed C headers
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Debug C++ tests

Use Google style for header includes.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Update CUDA driver macros

Incorporating changes from #389.
Co-authored-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Jan Bielak <jbielak@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use core error checking macros in PyTorch extensions

Hack to get around macro redefinition warning.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix missing arg when getting CUDA driver error string
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Reuse logging header in frameworks
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Jan Bielak <jbielak@nvidia.com>

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

* Drop eval for selective checkpointing tests
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Remove redundant recompute for FA
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* CI fix; Decouple fused attention and numerics tests
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* Ability to check cuDNN version from Python
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Modify the fused attention test to not use the CUDNN_VERSION env
variable which is specific to NGC containers
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

* Add class for RNG state tracker.
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

* Debug PyTorch and Paddle tests on Ada
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Only run Paddle layer tests with cuDNN fMHA on supported archs
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug PyTorch fMHA tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Reduce JAX FP8 GEMM sizes

Avoid split-k kernels on Ada.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Disable JAX fused self-attention test on Ada
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Run supported fused attention tests on Ada
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Run supported fused attention JAX tests on Ada
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Enable Paddle fused attention on Ada
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Update reference scale calculation in TensorFlow test
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Restore backend support to reference FP8 attention impl in PyT test

Review suggestion from @cyanguwa
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Debug Paddle tests on Ada
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Loosen tolerances for Paddle attention tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Assume causal mask implies equal seqlens in Paddle attention tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

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

* add support for h2d/2hd in 8.9.6
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* cull unit tests in fused_attn.py and add skipif for layout tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add workopt=1 flag for dpa tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update support table for arbi_seqlen backend
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix rotary position embedding and add unit tests accordingly
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* further cut down unit tests for CI efficiency
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* remove einops dependency
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: Tim Moon <tmoon@nvidia.com>

* [JAX] Enhance Dropout in TransformerLayer.

1. Fixed missing setup of dropout RNG key in TransformerLayer and
   LayerNormMLP.
2. Allowing seperated dropout rate for FC1's output and other hiddens.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix wrong fp8 scale in _update_fp8_metas_impl
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix typo
Signed-off-by: Ming Huang <mingh@nvidia.com>

---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* initial changes [wip]
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add padding mask support for FA
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* rm causal mask from tests and add padding
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

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

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

* [wip] fix API
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add packing and unpacking
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

* fix atomic_add bf16 torch.compile
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Generate non all True masks
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* Fix core attention export and FusedAttn filter
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

* Optimizations and caching fixes in torch.dynamo
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@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>

* wgrad-accumulation patch for legacy megatron-LM
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

* sb3hd/bs3hd working on top of 3xsbhd/bshd/thd
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* remove qkvso_strides on torch side; cover it in generateQKVStrides
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* all 15 layouts pass
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

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

* add note to deprecate some qkv layouts
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix code for unit tests in test_fused_attn.py
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* remove a couple more comments
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

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

* fix onnx for core attn; not fixed
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove nvtx and add env var for workspace opt
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove testing for env var
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* replace zeros/zeros_like with empty/empty_like
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix nvtx marker name for _q_k_v API
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove sm80 when compiling for h100
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add mapping from qkv layout to layout group and qkv format
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* clean up enums mapping and remove trailing spaces
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* simplify workspace opt control logic; only need env var
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix fp8 test, and minor modifications for other tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* avoid overwriting model configs in unit test
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* random fixes/improvements: get_qkv_format/etc, default values, docstrings, comments
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* change workspace opt logic back to FORCE_WORKSPACE_OPT
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* fix get_backend logic for max512/arbitrary
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* clean up unit tests for layouts, and fix minor lint issue
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* minor tweaks for CI testing: onnx string issue and test fused attn first
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove one unsupported layout from max512 and add a check to qkvpacked API
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix te layer test; reduce test time
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* revert compiler option changes; add back sm80 for even h100
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove some unit tests or make them optional to reduce CI time
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove more unit tests temporarily
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove _q_k_v in naming and add NVTE_ERROR for FP8 Aux_CTX_Tensors size checks
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* remove temp tests from last commit
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* remove prints from last commit
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* remove thd->bs3hd user warning to avoid GPU sync
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* adjust fused attn bs in tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* temporary fix for onnx issue; more fixes in PR 437
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove unused variables
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: Charlene Yang
Signed-off-by: cyanguwa <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* [PyTorch] Implement GQA based on fused q, k, v projection. Additionally fixes #392
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>

* [PyTorch] Extend parameters_split option in Linear and LayerNormLinear to support splitting with different sizes as required by unfused GQA.
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>

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

* fix noop cat to bypass torch.cat and support uneven split
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* fix cuda graph due to noop_cat
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* remove the use of enumerate when possible
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

---------
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Co-authored-by: Markus Schnoes <markus.schnoes@gmx.de>

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

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

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

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

* Change scaling factor from E8M0 to E8M23
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* Eliminate amax_and_scale_update bubbles
Signed-off-by: rewang <rewang@nvidia.com>

* Add CUDA check
Signed-off-by: rewang <rewang@nvidia.com>

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

* Enable the model to be change precision between iterations
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

* get default file name with time
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

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

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

* Improve performance of stack backtrace
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

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

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

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

* Add control of attention dropout and hidden dropout RNG state
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

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

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

* [Paddle] Add TP, DP, PP, FSDP
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

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

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

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

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

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

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

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

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

* Squashed commit of the following:

commit 79e2e5fd774e67dcdda9aae01a9f31a6479c5d70
Author: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>
Date:   Sun Aug 20 14:39:16 2023 +0000

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

commit 1d40ad60540490f97ed82ba877cc6eda8902cbf6
Author: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>
Date:   Sun Aug 20 14:22:25 2023 +0000

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

commit 6632f735a0c8251862355fc74622af59fae3a509
Author: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>
Date:   Sun Aug 20 05:52:18 2023 +0000

    Add TP for attention and transformer layer
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

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

* Add FSDP check for stage 1,2,3
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

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

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

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

* Fix CI errors
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>

* API change and some test fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

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

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

* Fused attention kernel only supports sm80 and sm90
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Update transformer_engine/jax/csrc/modules.cpp
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* arbitary fused kernel supports sm86/sm89 after 8.9.3
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Forward is_fused_attn_kernel_available to cpp backend
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove cpp is_fused_attn_available API
Signed-off-by: Reese Wang <rewang@nvidia.com>

---------
Signed-off-by: Reese Wang <rewang@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

fix rng_state issue and minor compiler warning
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* Expose MultiheadAttention for import
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Expand mask type and add no mask numerical test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.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>

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

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

* Save full global state
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

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

* add back pytorch csrc extensions.h
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add unit tests for dpa checkpointing
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove seqlen%32/64 checks for now
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* add tests for changes regarding rng_state in aux_ctx_tensor
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* reuse rng tracker from numerics in fused attn; skip checkpointing if FAv2 in numerics
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* uncomment comments used for testing
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix pre/post scale bias
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* Update transformer_engine/pytorch/attention.py
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: cyanguwa <8636796+cyanguwa@users.noreply.github.com>

* remove skipifs for FAv2 check after PR366
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove checkpointing tests for transformer layer; dpa tests still provide coverage
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* adjust random number range for tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* Add upper bound to FA version
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Check backend only when using FusedAttention
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* remove imports/variables related to FAv2 checks
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* further fix random number ranges for tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix variable referenced before assignment error
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: cyanguwa <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add dynamically shaped input mask in test_export_softmax
* Fix test_softmax_mask_fn - use env. var `NVTE_ONNX_KVCACHE_MAX_SEQ_LEN` to control whether the test uses the default mask generation function or dynamic TRILU mask slicing.
* Change core_attention ONNX export test: use "no_mask" as attn mask type when testing `te.attention.DotProductAttention` w/o masking.
* Use ORT CUDA backend by default.
Signed-off-by: Neta Zmora <nzmora@nvidia.com>

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

Fix for the RMSNorm tests/doc/ONNX export to match the actual
implementation
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Fix flash attention dropout probability with inference
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add output as the fused attention ctx tensor
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add rng_state as the fused attention ctx tensors
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add flash attention supported lengths to the fused attention
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Refactor attention primitive to reuse abstract shaped array
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Detect backend type to allocate appropriate ctx size
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Skip dropout correctness instead of return success
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Use cudaMemsetAsync and enhance the error handling
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add flash attention kernel elts_per_thread update
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove redundant max 512 suffix
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Keep only DType and remove NVTEDType from python
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix a float32_attention_logits bugs
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Re-calculate workspace size for self attention
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enhance bias/dbias shape guard
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enhance the seed/rng_state checker
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Use jax.core.ShapedArray as jax.abstract_arrays is deprecated
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enhance the unittest docs
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Cast Flax collections to FrozenDict as WAR to adapt Flax 0.7.1.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding min version of flax to requirements.txt in examples
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix praxis tests and rename compare_frozen_dict to compare_dict
Signed-off-by: Ming Huang <mingh@nvidia.com>

* [Paddle] Refactor FP8 state (#350)

Refactor fp8 state
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Store FP8 checkpointing data in CPU (#351)
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Make test_layer be able to run on both Flax >=0.7.1 and <=0.7.0
Signed-off-by: Ming Huang <mingh@nvidia.com>

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

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