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>

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

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

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

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

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

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

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>

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