GitLab

* Link attention docs to the main docs and fix errors reported by Sphinx
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Lower the version of nbsphinx
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

* Change the URL of example_attention.py to GitHub
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* More fixes in the attention tutorial
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

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

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

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

* Reformat FP8 Meta

1. Reformat FP8 meta to be one-set-per-tensor.
2. Remove fp8_max and scale_inv.
3. Remove unused functions in fp8.py
Signed-off-by: Ming Huang <mingh@nvidia.com>

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

* Remove ShardingType and MajorShardingType
Signed-off-by: Ming Huang <mingh@nvidia.com>

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

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

* Rename few variables.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Add jit to update_amax_list
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fixed naming error in LayernormMLP
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fixed bugs in test_distributed_layernorm_mlp.py
Signed-off-by: Ming Huang <mingh@nvidia.com>

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

* Remove interval arg from recipe
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Remove usage of interval and use explicit kwarg for testing recipes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* Support FP8 Meta Dtype (FM32) and Align FP8 Scale Update with PyTorch.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Modify with the feedback of code review
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Hiding FlaxFloatMeta32 inside fp8.py
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Make functions to be JAX tracable objects.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Rebased with mian.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Update jax images for github workflow.
Signed-off-by: Ming Huang <mingh@nvidia.com>

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

Signed-off-by: Przemek Tredak <ptredak@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>

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

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

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

* Use the same default in the function to what the class default.
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>

* Assert instead of silently ignoring not supported variation. Small doc correction, amax_compute_algo is partially supported.
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>

* Fix line lenght to fix the CI.
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>

* Apply suggestions from code review
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Frédéric Bastien <frederic.bastien@gmail.com>

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

* Clarify that it is only TE/JAX that don't support that faeture.
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>

* Update transformer_engine/jax/fp8.py
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Frédéric Bastien <frederic.bastien@gmail.com>

* Update the test following the change in default value
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>

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

---------
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>
Signed-off-by: Frédéric Bastien <frederic.bastien@gmail.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>

* Ignore IDE files
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Fix typing errors
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Ignore devcontainer files
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Avoid import from private module
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Apply @timmoon10 's suggestions
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

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

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

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

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

* better FP8 checker
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

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

* replace te.* with te.flax* to remove deprecated warning
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* remove nouse os.environ
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

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

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

* Update examples/jax/encoder/README.md
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* Update examples/jax/encoder/README.md
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* Update examples/jax/encoder/README.md
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* Update examples/jax/encoder/README.md
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* Update examples/jax/encoder/test_multiprocessing_encoder.py
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

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

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

* Update examples/jax/encoder/README.md
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* fix cpp lint

fix issue of "Could not find a newline character at the end of the file."
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* fix AssertionError: 1 GPU per process
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* replace tfds with datasets

The Flax application crash if it use TensorFlow Dataset (tfds) in NVIDIA JAX container.
The tfds is very useful for downloading well-knwon dataset (e.g., MNIST, GLUE) and commonly used by TF/JAX community.
However, it seems like that it is NOT compatible with NVIDIA TensorFlow in NVIDIA JAX container and somehow affects JAX.
It triggers random errors at JAX initialization depending on different versions, and make CI unstable.
Thus, this commit replaces tfds with "huggingface datasets" to download needed datasets.
See "nvbugs 4039266" for more details.
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* fix input sharding

Unlike SPMD mode, in multiprocessing mode, the input tensor must be sharded manually.
Using DP=4, TP=2 as an example, the device mesh looks like:

mesh.device_ids = [[0, 1],
                   [2, 3],
                   [4, 5],
                   [6, 7]]

Assume that the process ID is mapped to GPU ID.
The process 0 and process 1 are grouped for model parallelism,
process 2 and process 3 are grouped together too, and so on.

The process 0 and process 1 need to share the same micro-batch in the training step,
process 0 and process 2, 4, and 6 have different micro-batch.

Thus, `shard_array_wrapper` partitions inputs to 4 parts (and setup
needed arguments for jax.make_array_from_single_device_arrays).
The process 0 and process 1 take the first quarter,
process 2 and process 3 take the second quarter, and so on.
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* refactor UT for multiprocess example

Use Python `multiprocessing` to test the multiprocessing example,
if the system has multiple GPU. 1 GPU per process.

Because `jax.distributed.initialize` must be called before any other JAX or Flax API,
GPU info cannot be queried by calling jax.local_devices() in TestEncoder.
Thus, `unittest_query_gpu()` forks another process to query number of GPUs and
FP8 capability.
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* remove nouse arg `--num-gpu`

JAX doesn't have an API to setup number of GPU used in SPMD mode.
The only way is to use `CUDA_VISIBLE_DEVICES` for now.
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

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

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

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

* increase batch-size for multigpu example

The batch-size 64 is too small to be partitioned for 8xH100.
If batch-size is 64, the GEMM shape is 256x8192x8 per GPU.
The 8 is too small for FP8 GEMM kernel, and
cuBLASLt will throw "Failed to query heuristics".
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* fix downloading mnist error

To download MNIST via `huggingface datasets`, it requires Pillow.
Otherwise, it throws `An error occurred while generating the
dataset`
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>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.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>

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

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

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

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

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

* add building workflow for jax modules
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* replace bit_cast with reinterpret_cast
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* add nvtx to cmake check list
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

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

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

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

* set pytorch as default in setup.py
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* rename extension from *.cc to *.cpp

cpplint cannot recognize *.cc file, so rename the extension
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* refactor style, to align TE/PyTorch
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* add pybinding, unittest and qa
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

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

* disable c-extension-no-member and no-name-in-module
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* add dataclass avoid pylint error
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

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

* Update tests/jax/test_custom_call_shape.py

fix typo
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>

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

* add building workflow for jax modules
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* replace bit_cast with reinterpret_cast
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* add nvtx to cmake check list
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

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

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

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

* set pytorch as default in setup.py
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* rename extension from *.cc to *.cpp

cpplint cannot recognize *.cc file, so rename the extension
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* refactor style, to align TE/PyTorch
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* add pybinding, unittest and qa
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

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

* disable c-extension-no-member and no-name-in-module
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* add dataclass avoid pylint error
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

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

* Update tests/jax/test_custom_call_shape.py

fix typo
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Jeng Bai-Cheng <jeng1220@users.noreply.github.com>

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

* fix conflict due to PR62
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

* fix c-extension-no-member and no-name-in-module

1. add transformer_engine_jax into extension-pkg-whitelist
2. convert pylintrc from CRLF to LF format
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>

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

* remove pylint:disable and refactor import order
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: Kirthi Shankar Sivamani <ksivamani@nvidia.com>