GitLab

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

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

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

* Initial commit for fp8 calibration.
Signed-off-by: Sharan Chetlur <schetlur@dlcluster.nvidia.com>

* Fixes to make unit tests pass
Signed-off-by: Sharan Chetlur <schetlur@dlcluster.nvidia.com>

* Added test and finished implementation
Signed-off-by: Sharan Chetlur <schetlur@nvidia.com>

* Cleaning up handling of save_for_backward in Linear
Signed-off-by: Sharan Chetlur <schetlur@nvidia.com>

* Removing commented lines
Signed-off-by: Sharan Chetlur <schetlur@nvidia.com>

* Minor fix to mnist test.
Signed-off-by: Sharan Chetlur <schetlur@nvidia.com>

* Pylint cleanup
Signed-off-by: Sharan Chetlur <schetlur@nvidia.com>

* Moving stats computation to the forward pass instead of pre_forward, and extending to all other layers
Signed-off-by: Sharan Chetlur <schetlur@nvidia.com>

* Pylint cleanup
Signed-off-by: Sharan Chetlur <schetlur@nvidia.com>

* Pylint cleanup.
Signed-off-by: Sharan Chetlur <schetlur@nvidia.com>

* Fixing unit test failures.
Signed-off-by: Sharan Chetlur <schetlur@nvidia.com>

* Misc changes
Signed-off-by: Sharan Chetlur <schetlur@nvidia.com>

* Fixing bad indentation from master merge and moving some code into the needs_stats conditional
Signed-off-by: Sharan Chetlur <schetlur@nvidia.com>
Signed-off-by: Sharan Chetlur <schetlur@dlcluster.nvidia.com>
Signed-off-by: Sharan Chetlur <schetlur@nvidia.com>
Signed-off-by: schetlur-nv <116769508+schetlur-nv@users.noreply.github.com>
Co-authored-by: Sharan Chetlur <schetlur@dlcluster.nvidia.com>

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

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