GitLab

* Add fused attention unit tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Use NVTE_* enums
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Use NVTE_Mask_Type and remove FMHADescriptor
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Move common functions to utils
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Change namespace to fused_attn
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Move fused_attn_max_512_fwd_qkvpacked under the general APIs
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Move fused_attn_max_512_bwd_qkvpacked under the general APIs
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove redundant blank line
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix a potential bug for cu_seqlen converter
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Refine the unfused attention warning message
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Rename to fused_attn_max_512
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove the deprecated header
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Rename to fused attn
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

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

* Refactor jax primitive API
* Merge q_cu_seqlen and kv_cu_seqlen
* Remove is_causal_masking
* Replace seed with rng_state
* Add is_training argument
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Add None guard for bias and dropout_rng
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Add is_fused_attn_kernel_available() to correctly dispatch the attention impl
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Adjust the code style
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add the missing blank lines
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Change the order of FADescriptor members
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enhance the readability of fused_attn_max_512.cu
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Generalize the input dimension unpacking
Signed-off-by: Reese Wang <rewang@nvidia.com>

* 16 bits fused attention requires 8.9.1
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Update fused attention support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Handle None type when sharding
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Change to the padding ratio
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Performance optimization for non-bias cases
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Revert the cudnn-frontend PRIVATE keyword which was used for debugging
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Revert "Update fused attention support matrix"

This reverts commit 4effe67d0f08f733919a329ce5ab421958740f4a.
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Treat b * s as total_seqs to align ragged cases
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add FP16/BF16 max_seqlen <= 512 fused attention to the support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Refine test_fused_attn.py

* Replace reference code with flax.linen
* Remove unnecessary comments
* Use AttnMaskType
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Unify the cuDNN compile version
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add dropout to the support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Slightly adjust the headers
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Typo fix: remove redundant either
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Consolidating fused attention requirements
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Replace cudnn_frontend::throw_if with NVTE_CHECK for the better error line report
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Rename to fused_attn_fp16_bf16_max_seqlen_512 for the better readability
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Add more annotations to the custom calls
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

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

* refactor attention out of transformer.py
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* ONNX export refactoring

* Remove infer_ort (to enable more testing)
* Add BF16 ORT tests for Q/DQ ops and GELU.
  * Use FP32 i/o instead of BF16 (because ORT doesn't support BF16 i/o) and add casts from FP32 to BF16 (this is only for subgraph inputs and outputs).
  * We'll need to add more BF16 testing.
* GEMM:
  * Add cast after DQ to achieve better performance (matmul at sub-fp32 precisions).
  * Fold bias into Gemm operation (=> smaller graphs)
  * Wrap GEMM-GELU with FP32 (TE implements GELU in FP32)
* Enable tests for cross attention (test_export_multihead_attention)
* Reduce test thresholds for test_export_layernorm_mlp, test_export_layernorm_linear, test_export_layernorm
Signed-off-by: Neta Zmora <nzmora@nvidia.com>

* Loosen MHA export validation thresholds for FP16
Signed-off-by: Neta Zmora <nzmora@nvidia.com>

---------
Signed-off-by: Neta Zmora <nzmora@nvidia.com>

DDP support for no-bias option
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* move dbias from input list to output list for bwd
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* split asserts into three for bias checks
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* another fix for asserts for bias checks
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>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* using different strems for pushsend and pushrecv
Signed-off-by: Sangkug Lym <slym@nvidia.com>

* fix stream dependency
Signed-off-by: Sangkug Lym <slym@nvidia.com>

* add wait from main_stream to memcpy stream
Signed-off-by: Sangkug Lym <slym@nvidia.com>

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

Use only built-ins for setup
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* iFix LN ONNX export

When exporting LayerNorm make sure that the weights and bias
inputs have the same type as the LN input.
Also:
 * Add a regression test.
 * Add environment variable to override directory of generated test artifacts
Signed-off-by: Neta Zmora <nzmora@nvidia.com>

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

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

---------
Signed-off-by: Neta Zmora <nzmora@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Fix bug in Sequence Parallelism
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Adjust Module Structure.

1. Collect Flax related modules to a sub-folder, flax.
2. Add a function to unify scale_init for zero-centered-gamma LN.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Make changes be compatible to previous versions.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adapt jax/examples to the new module structure.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Update jax/docs and Add deprecated warning.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Update README
Signed-off-by: Ming Huang <mingh@nvidia.com>

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

* Adding deprecated warning to flax modules which imported via transformer_engine.jax
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix CI errors and update docs.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Removing unnecessary deprecated warning in docs.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Implementing __iter__ to DeprecatedEnum.
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>

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

* Remove the nonexistent parameter from fused attention documentation
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Remove the second instance
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

* Faster split of QKV for FlashAttention
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

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

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

* Message with assert
Co-authored-by: Przemyslaw Tredak <ptredak@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

* make clarifying comment and check strides
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

remove used function ternary_pw_op_create
Signed-off-by: Charlene Yang <charleney@nvidia.com>
Co-authored-by: Charlene Yang <charleney@nvidia.com>

add cast for BF16 input/output types for Q/DQ ONNX ops
Signed-off-by: Asfiya Baig <asfiyab@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* Add FP8 fused attention to TE for PyTorch
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* add license for cudnn-frontend, modify installation requirements, and refactor some headers for aesthetics
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* add c api docs for fused attention
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* add exception for unsupported precision/sequence length combinations
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* fix installation requirement for non fused attn use cases
Signed-off-by: Charlene Yang <charleney@nvidia.com>

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

* prefix enums with NVTE_ and replace old MHA_Matrix with NVTE_QKV_Matrix
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* minor fixes based on PR comments
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* fix description for kvpacked fwd
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* fix description of Bias in C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* minor fixes for cudnn requirement and description for QKV tensors
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* fix QKV layout description and support matrix for C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* add asserts to cpp_extensions for qkv layout/bias type/attn mask type
Signed-off-by: Charlene Yang <charleney@nvidia.com>

* fix typo precision
Signed-off-by: Charlene Yang <charleney@nvidia.com>

---------
Signed-off-by: Charlene Yang <charleney@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Charlene Yang <charleney@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* Fix linking issue and make mpi conditional
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* Use max SMs at the last RS chunk in pipelined overlap
Co-authored-by: Sangkug Lym <slym@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* Make userbuffers support opt-in

Decouple userbuffers from MPI. Refactor MPI handling in build system. Standardize names to "userbuffers".
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Sangkug Lym <slym@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>

* Allow update_collections and update_fp8_metas to return both Dict and FrozenDict.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix the wrong shape issue of bias when fused QKV or KV.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Reuse tuplized features for bias creating.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Replace get_args to be more readable.
Signed-off-by: Ming Huang <mingh@nvidia.com>

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

* Fix gencode flags in setup (#145)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

* amax reduction internval
Signed-off-by: Sangkug Lym <slym@nvidia.com>

Skip TP-domain only AMAX reduction when TP-group is not initialized
Signed-off-by: Sangkug Lym <slym@nvidia.com>

* Update transformer_engine/pytorch/fp8.py
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Sangkug Lym <slym@nvidia.com>

* check TP group initialized
Signed-off-by: Sangkug Lym <slym@nvidia.com>

fix
Signed-off-by: Sangkug Lym <slym@nvidia.com>

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

* use upstream flash-attn
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* get correct FA for linting
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>

* Fix model load exception when state resides on GPU

- Whenever converting a torch.tensor to numpy, we need to first
migrate the tensor storage to the host CPU.

- Add a warning not to do contant-folding when exporting to ONNX.
This is due to a torch.onnx export bug.

- Refactor compare_outputs
Signed-off-by: Neta Zmora <nzmora@nvidia.com>

* Onnx export: Improve remark text
Signed-off-by: Neta Zmora <nzmora@nvidia.com>

---------
Signed-off-by: Neta Zmora <nzmora@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add zero_center_gamma/functional pass
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

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

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

* Refactored code style for improved readability and consistency
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Docs enhancement for zero_centered_gamma
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add escape for line break and remove some bad if conditions
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Revise scale_init docs
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

Remove the autocast_variable
Signed-off-by: kaixih <kaixih@nvidia.com>

Fix cyclic import error in TF
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* made bias configurable
Signed-off-by: Naman Goyal <naman@fb.com>

* removed commented lines
Signed-off-by: Naman Goyal <naman@fb.com>

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

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

* fixed incorrect call to fused bias dropout add kernel
Signed-off-by: Naman Goyal <naman@fb.com>

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

* Separate FC1 and FC2 use_bias args; solves all ci errors
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

---------
Signed-off-by: Naman Goyal <naman@fb.com>
Signed-off-by: ngoyal2707 <ngoyal2707@users.noreply.github.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Naman Goyal <naman@fb.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Rename enable_fp8 to is_fp8_enabled.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding an API to get an instance of  DelayedScaling which is set via fp8_autocast.
Signed-off-by: Ming Huang <mingh@nvidia.com>

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

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

* conditional dgrad for Linear
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* add tests and small improvements to LNLinear and LNMLP
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* async amax reduction

add env knob to enable async amax reduction
Signed-off-by: slym <slym@login-preos01.a51.clusters.nvidia.com>

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

* remove is_last_model
Signed-off-by: slym <slym@login-preos01.a51.clusters.nvidia.com>

* fix naming
Signed-off-by: slym <slym@login-preos01.a51.clusters.nvidia.com>

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

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

---------
Signed-off-by: slym <slym@login-preos01.a51.clusters.nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: slym <slym@login-preos01.a51.clusters.nvidia.com>

* Add FP8 support for Ada
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

* better message for no fp8
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* same thing for onnx test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

* Bugfix - compute scale_inv when loading checkpoint
Signed-off-by: Gal Hubara Agam <ghubaraagam@nvidia.com>

* Save inverse scale in extra state tensor + minor CR fixes
Signed-off-by: Gal Hubara Agam <ghubaraagam@nvidia.com>

* Fix lint
Co-authored-by: Gal Hubara Agam <ghubaraagam@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Gal Hubara Agam <ghubaraagam@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

* Change FP8 recipe defaults
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Increase default amax history length
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Always check history size
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* no amax history for onnx export
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* revert onnx export test changes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix indices in onnx test
Co-authored-by: Neta Zmora <nzmora@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Neta Zmora <nzmora@nvidia.com>

Signed-off-by: Tal Cherckez <tcherckez@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>