GitLab

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

* Support transpose_bs when decoded=True
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix Bugs,

1. Fix missing dropout_dims in LayerNormMLP.
2. Fix broadcast issues in decoded.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix wrong masks in decoded.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fixed wrong assert condition in TransformerLayer
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix amax is not set as 0 in each step.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Enhance rules conflict checking and docs.
Signed-off-by: Ming Huang <mingh@nvidia.com>

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

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

* Add support for fp8 GEMM BIAS AUX GELU fusion
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Fix Lint error
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Fix Lint error
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

---------
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

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

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

* add dp example
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

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

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

* Add model parallel example and refactor
Update readme
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* align code and readme
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* update verification
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* add mask
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* num_gpu is configurable
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* update readme
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* update readme
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* solvepylint issue
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* ignore markdown and txt file from license check
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* Update README.md
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* add flax into requirements.txt
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

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

* Add tensorflow build

Improve build instructions

Fix pybind enum usage

Fix Python_EXECUTABLE cmake var

Move scale_inv calculations to FW
Signed-off-by: Trevor Morris <tmorris@nvidia.com>
Signed-off-by: kaixih <kaixih@nvidia.com>

* Apply clang-format
Signed-off-by: Trevor Morris <tmorris@nvidia.com>
Signed-off-by: kaixih <kaixih@nvidia.com>

* Format python files
Signed-off-by: Trevor Morris <tmorris@nvidia.com>
Signed-off-by: kaixih <kaixih@nvidia.com>

* Add TF build CI
Signed-off-by: Trevor Morris <tmorris@nvidia.com>
Signed-off-by: kaixih <kaixih@nvidia.com>

* Lint checks
Signed-off-by: kaixih <kaixih@nvidia.com>

* Another round of lint checks
Signed-off-by: kaixih <kaixih@nvidia.com>

* Fix TF image tag
Signed-off-by: Trevor Morris <tmorris@nvidia.com>
Signed-off-by: kaixih <kaixih@nvidia.com>

* Use the existing recipe file
Signed-off-by: kaixih <kaixih@nvidia.com>

* Add license claim blocks
Signed-off-by: kaixih <kaixih@nvidia.com>

* Fix a bug about bias dtype conversion
Signed-off-by: kaixih <kaixih@nvidia.com>

* Add mnist example and cleanup old examples
Signed-off-by: kaixih <kaixih@nvidia.com>

* Autopep8 the tests
Signed-off-by: kaixih <kaixih@nvidia.com>

* Autopep8 the examples
Signed-off-by: kaixih <kaixih@nvidia.com>

* Add example in Readme
Signed-off-by: kaixih <kaixih@nvidia.com>

* Add unit tests and linting for TensorFlow
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add causal mask for non-fused case
Signed-off-by: kaixih <kaixih@nvidia.com>

* Fix the mismatched TF vs TE masks
Signed-off-by: kaixih <kaixih@nvidia.com>

* Addressing CI tests
Signed-off-by: kaixih <kaixih@nvidia.com>

* Run lint test
Signed-off-by: kaixih <kaixih@nvidia.com>

* Add missing import
Signed-off-by: kaixih <kaixih@nvidia.com>

* Skip fp8 tests for pre-Hopper GPUs
Signed-off-by: kaixih <kaixih@nvidia.com>

* Remove non-pytest tests
Signed-off-by: kaixih <kaixih@nvidia.com>

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

---------
Signed-off-by: Trevor Morris <tmorris@nvidia.com>
Signed-off-by: kaixih <kaixih@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: kaixih <kaixih@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>

* Remove zombie process from querying TE install path
Co-authored-by: Naman Goyal <naman@fb.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

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

* Fix lint warning
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>
Co-authored-by: Naman Goyal <naman@fb.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

* Fix GELU ONNX export

* Wrap GELU export with cast to/from FP32 to achieve same compute precision as TE.
* Increase GELU export test thresholds.
* Change export to ONNX opset 17 for smaller representation of LN (single node instead of subgraph).
* Remove the need for LN work-around for ORT
Signed-off-by: Neta Zmora <nzmora@nvidia.com>

* Add docstring to te_onnx_extensions.py::compute_in_fp32
Signed-off-by: Neta Zmora <nzmora@nvidia.com>

* Tune threshold for GELU ONNX export

Ran 8K test instances to verify the threshold.
Allow 2 coefficients to escape threshold. Two wrong coefficients
are not a failure.
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>

FA doesn't support compute 8.6 with head_dim>64
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Initial commit for fp8_transpose_dbias kernel
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

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

* Suggestions and fixes
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>

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

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

* Add a temporary workaround to layernorm export

Seems like ORT is performing template-matching for LN and incorrectly concludes
that it doesn't have a kernel for FP32 LN. The work-around adds the addition of
fake_zero which is meant to prevent the template matching while keeping the graph
virtually unchanged. This also requires `do_constant_folding=False` in
`torch.onnx.export`.
Signed-off-by: Neta Zmora <nzmora@nvidia.com>

* Adjust test threshold
Signed-off-by: Neta Zmora <nzmora@nvidia.com>

* Opened an ORT bug and added the link for tracking
Signed-off-by: Neta Zmora <nzmora@nvidia.com>

* Fix Python linter errors
Signed-off-by: Neta Zmora <nzmora@nvidia.com>

* Simplify the LN workaround solution (ONNX export)

After discussing https://github.com/microsoft/onnxruntime/issues/15021


with Microsoft engineers, replaced the LN workaround with a simpler
implementation.

In addition:
* To make test more robust add `allow_cnt_errors` to `validate_result`
* Add more documentation to clarify the purpose and methodology of the
ONNX export tests
Signed-off-by: Neta Zmora <nzmora@nvidia.com>

* Fix unused import
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* Fix unused import
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* Fix unused import
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

---------
Signed-off-by: Neta Zmora <nzmora@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* Adding JAX to README.rst
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Refine README.rst as the suggestion from review.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Refine the API doc of extend_logical_axis_rules.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

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

Use updated comm API PyTorch
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Updated TE/JAX docs
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding TE/JAX docs' rst files
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Set DType as pybind11::module_local() to avoid generic_type errors.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Updating license and exporting more modules
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adopting autoapi and removing enum_tools.
Signed-off-by: Ming Huang <mingh@nvidia.com>

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

* Make jax.rst be style consistent.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Fixing doc statements as the suggestion from review.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Fixing doc statements as the suggestion from code review.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Update the description of Softmax
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Removed categories in catalog as PyTorch
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

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

* Catch FP8 modulo16 error before cublas and fp8 kernels
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* catch incorrect usage of fp8_autocast
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* catch error on first time double execution
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* Change from AutoDoc to AutoAPI
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

* WAR for the wrong autosummary generation
Signed-off-by: Przemyslaw Tredak <ptredak@nvidia.com>

* Change common to be in line with pytorch API docs
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Add GitHub Action to build docs
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

* Trying to fix the versions
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

* deprecate qk layer scaling and fp32 softmax args
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* apply QK layer scaling for fp16 training
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>

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

* add transformer module , unittests and examples
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

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

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

* remove pylint: disable=line-too-long
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* remove pylint: disable=too-many-func-args
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* Fix the wrong broadcasting dim to dropout masks when enable transpose_bs.
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* Enable 2xACC for WGRAD and DGRAD by default
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* rename LayerNormMlpBlock as LayerNormMLP
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* refactor to avoid line-too-long
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* rename amax_history_size to amax_history_len
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* align dropout mask to TE/PyTorch as default
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* enlarge atol for decoder unittests

Two decoder unittests can pass in old JAX container(e.g., 23.02)
but can't in latest container (devel).

1. The actual(-0.020264) and desired(-0.020386) are very close.
2. The TE kernels are not changed, the diff should come from
   new codegen behavior of XLA.

Thus, it is a common floating-point accumulated error.
Enlarge atol to avoid unittest failures.
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* Adding Amax History Support

1. hide amax update in custom_vjp
2. replace amax indexing with roll(using circular buffer)
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* move kernel_init to __post_init__
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

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

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

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

* remove envvar regarding 2xACC
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* remove unused import
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

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

* ignore self attention mask for causal type
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* further relax checks to run FA, update docs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* minimum ampere requirement for fa
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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