GitLab

* Integrate cuDNN frontend v1 to fused attention and miscellaneous fixes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

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

* simplify stride generation
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix and/or logic in get_backend
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* remove v.contiguous() since get_qkv_layout covers it
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

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

* revert mask type to comma-separated list
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* integrate v1/pre-release-5
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* cleanup prerelease5 integration and fix FA2.1 commit
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* force dropout to 0 if not training
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* testing bias/alibi and padding+causal; add alibi to unfused DPA
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* set flag_arb to false when non determinism is not allowed
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* followup on prev commit; remove redundant python env var setting
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* fix determinism logic for fused attn
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

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

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

* remove use of non-deter var for backend selection
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* fix workspace size in bwd and uncomment bias test
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

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

* disable arbitrary backend + post scale bias in Jax; waiting on PR 525
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* swap bias/rng_state order in aux_ctx_tensor and add bias to aux_ctx_tensor in _qkvpacked/_kvpacked API
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove support for padding_causal + cross for max512
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* change alibi bias to float32 for bias_1_4/5 tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* fix thd fwd output shape for FlashAttention and add backend info for DPA
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix definition of workspace limit when dbias is present
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* further tweak DP_WORKSPACE_LIMIT definition
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* disallow alibi+no_mask for sdpa flash and update alibi tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update jax/paddle after PR525 and fix DP_WORKSPACE_LIMIT for dbias Jax tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* disable dbias for non-hopper archs
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* remove build dir in setup.py
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* change selection logic to prefer fused attn on sm90
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* fix h and s order in header
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* remove manual setting of workopt path due to dbias after v1 update
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* add post_scale_bias and alibi to sdpa flash support matrix
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix support matrix in header files
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* move headers back to .cu and change seed/offset to int64
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update Megatron commit in L1 test and remove all prints in fused attn test
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* fix fp8 arg in L1 Megatron script
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* print only when debug flag is on
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove checkpointing loading to avoid loading other tests results
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>

* Float8Tensor uses cached transpose if available
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix bug with non-2D transpose
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Custom pickling for Float8Tensor
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug test for pickling Float8Tensor
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Review suggestions from @sudhakarsingh27

Avoid FP8 casts when copying between Float8Tensors. Make make_like a class function.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add unit test for checkpointing model with FP8 params

Debugged pickling and copy functions.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Add H200 perf non-alpha image
Signed-off-by: Santosh Bhavani <santosh@semantic.md>

* Update README.rst - non-transparent H200 plot
Signed-off-by: Santosh Bhavani <santosh@semantic.md>

---------
Signed-off-by: Santosh Bhavani <santosh@semantic.md>

* Update README.rst - FP8 convergence

- added FP8 convergence section
- removed model support (to be replaced with a feature support table)
Signed-off-by: Santosh Bhavani <santosh@semantic.md>

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

* Update README.rst

Updated Latest News
Signed-off-by: Santosh Bhavani <santosh@semantic.md>

* Update README.rst

Add plot for H200
Signed-off-by: Santosh Bhavani <santosh@semantic.md>

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

Signed-off-by: Zhenhuan Liu <nkulzh16@gmail.com>

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

* [PyTorch] TransformerLayer: add parallel_attention_mlp to support Falcon models
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>

* [PyTorch] add test for parallel_attention_mlp to test_numerics
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>

* [PyTorch] TorchGPT: fix dropout for parallel_attention_mlp

Now uses nn.functional.dropout because depending on the path there are one or two dropouts.
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>

* Apply suggestions from code review
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* [PyTorch] test_gpt_accuracy: fix spelling in construction of TorchGPT
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>

---------
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

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

Fix incorrect variable name in LayerNormMLP backward
Signed-off-by: Tim Moon <tmoon@nvidia.com>

[PyTorch] fix amax calculate during fp8 calibration
Signed-off-by: ladyrick <ladyrick@qq.com>

* Add rng_state output for cross fused attention
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add rng_state and output for the flash attention backward
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add bias for the jax cross attn API
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Add bias in the backward for the arbitrary fused attn backend
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

wgrad should be zero'ed out if a weight parameter is shared among multiple layers
Signed-off-by: Deepak Narayanan <dnarayanan@nvidia.com>

Support layernorm sm_margin through environment variables
Signed-off-by: Reese Wang <rewang@nvidia.com>

Use relative idx to ScaledUpperTriangMaskedSoftmaxFwdPrimitive.abstract to support batching.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* [PyTorch] Linear: fix computation for wgrad if sequence_parallel=True
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>

* Remove buggy gather_along_last_dim
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* [PyTorch] Linear: fix line length
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>

* Simplify logic for saving input tensor for Linear backward
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>

Getting warnings of the following form using ToT TE:

```
/usr/local/lib/python3.10/dist-packages/transformer_engine/pytorch/attention.py:852: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()
   data_ptr = grad_outputs[0].storage().data_ptr()
```
Signed-off-by: Deepak Narayanan <2724038+deepakn94@users.noreply.github.com>

Use unsigned char for byte dtype
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Add set CUDA_DEVICE_MAX_CONNECTIONS warning for tp overlap
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: Tim Moon <4406448+timmoon10@users.noreply.github.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>

* Remove assertion for NO_MASK
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix JAX distributed unit tests name
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

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

* disable FAv2.1 if causal+cross attn
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* include both causal and padding+causal
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* Delay caching of transposes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

add comments
Signed-off-by: Jon Chuang <9093549+jon-chuang@users.noreply.github.com>

* fix condition checks related to FA head_dim
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* force q,k,v contiguous when RoPE is in use
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

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

* Refactor sharding.py for the further custom_partitioning migration
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Migrating both FWD and BWD of LayerNorm/RMSNorm from xmap to custom_partitioning.
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Migrating both FWD and BWD of all kinds of softmax from xmap to custom_partitioning.
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Fix the wrong order of parameters to LN/RMSN bwd in ln_mlp_fp8.
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* WAR to LN/RMSN_fp8 before migrating to CP.
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Fix the wrong order of parameters of bwd of LN/RMSN_fp8.
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Following review feedback to modify
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Force the hidden dim in Norm ops to no sharding and add warning msg.
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Reuse fwd_rule in VJP functions
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Migrating both FWD and BWD of self-fused-attn from xmap to custom_partitioning.
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Migrating both FWD and BWD of cross-fused-attn from xmap to custom_partitioning.
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* add gelu and dgelu.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Reuse fwd_rule in VJP functions for attentions
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Apply native FP8 Dtypes to fp8.py
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Migrating cast_and_transpose from xmap to custom_partitioning
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Migrating transpose from xmap to custom_partitioning
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Apply XLA pattern match to perform FP8 GEMM.
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* migrate layernorm_fp8 to custom_partitioning.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Unify code style
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Extend supported of Transpose with FP8
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Implementing layernorm_fp8_dot based on migrated custom calls.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Renaming variables and publish NVTE_FP8_COLLECTION_NAME
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Replace Q/DQ custom calls with native XLA implementations
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* migrate gelu_fp to custom_partitioning.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Miner fix
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Support custom calls with mutli-dims
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Support gerneral dot indices in _fp8_dot_impl
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Implementing layernrom_geglu_fp8_mlp
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Remove GEMM custom calls
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Remove xmap related code
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Fix typo and add query-function to FP8MetaPackage
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Fix some bugs of custom calls
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Fix CT's bugs
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Update UTs/eaxmaples to adapt to the API changes.
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Unify kernel initilization in MLP.
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Modifing with code review's feedback
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Update README and Add deprecating warning to *ShardingType
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Canonicalize the dtype
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding assertion for non-supported batch dims.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding doc/examples to _multidim_transpose
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Set FP8 meta as WeightHParamsCollection.OVERWRITE_WITH_GRADIENT in Praxis modules.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Set FP8 meta as WeightHParamsCollection.OVERWRITE_WITH_GRADIENT in Praxis modules.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Apply dtype-based rtol/atol to UTs
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Deprecate QKV_INTERLEAVED enum
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Skip test_distributed_custom_ops.py
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix the wrong sharding of bias in SelfAttn
Signed-off-by: Ming Huang <mingh@nvidia.com>

* WAR to fix the wrong cu_seqlen of MHA when DP/FSDP enabled
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding distributed ops unit-tests
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding license to test_distributed_*
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Follow review feedback to modify
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Use total bytes involved in collective ops as criteria.
Signed-off-by: Ming Huang <mingh@nvidia.com>

---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Co-authored-by: Donglin Yang <dongliny@nvidia.com>

* Update README.rst - Installation section

Added pip install instructions and cleaned up pre-reqs and FlashAttention-2 section
Signed-off-by: Santosh Bhavani <santosh@semantic.md>

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

Improve PyTorch memory usage
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

[C/JAX] Support more mask types for the arbitrary seqlen kernels and minor changes of JAX bias (#469)

* Move bias to float32
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Increase neg infinity abs values
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enable varlen tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

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

* Support variable sequence length after cuDNN 8.9.6
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Use unique_ptr instead of shared_ptr
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add a new mask type: PADDING_CAUSAL_MASK
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Support flash padding mask after 8.9.6
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enhance the Max512 handling for causal masking and add the related tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Update the fused attn support lists
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove padding_aware from the caching
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix libtransformer.so issue
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Reduce the pad ratio tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix a bug with cuDNN 8.9.5
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Release backend resource after the module level unit test
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Clean the jax live arrays before running the unit tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix too-few-public-methods lint
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Make user buffer name configurable
Signed-off-by: Sangkug Lym <slym@nvidia.com>

* Apply suggestions from code review
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

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

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

---------
Signed-off-by: Sangkug Lym <slym@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@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>

* Returning an empty tensor of param dtype for wgrad
Signed-off-by: Selvaraj Anandaraj <selvaraja@computelab-frontend-4-ub22.nvidia.com>

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

---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@computelab-frontend-4-ub22.nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@computelab-frontend-4-ub22.nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@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>

* Use FP8 tolerances in JAX FP8 tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Programmatically compute expected floating point error
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Loosen tolerance for MNIST test
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

Disable large test cases for PyTorch Transformer layer
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

fix bwd error with FA v2
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@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>

Inherit the existing XLA_FLAGS
Signed-off-by: Reese Wang <rewang@nvidia.com>

Signed-off-by: Frederic Bastien <fbastien@nvidia.com>