GitLab

* scaling enum abstract

* rm NVTE_ from ScalingMode names

* rework scaling mode enum in grouped gemm

* fix norm sharding

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* refactor + mxfp8

* added grouped gemm

* rename linear to dense

* added cublas init phase for groupedGemm

* relax the tol of test encoder multiprocessing mxfp8 by 0.001
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Co-authored-by: Hua Huang <huah@nvidia.com>
Co-authored-by: Jeremy Berchtold <jberchtold@nvidia.com>

flax module with compute dtype inferred from the inputs
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* fixes L1 test

* fix test_multigpu_encoder

* fixes for other multi-encoder tests

* jax.extend.ffi to jax.ffi

* initialization with float32

* add init_dtype as an optional arg to all modules

* update use_scan query from xla flags

* relax threshold for test_encoder fp8

* relax the tols

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

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

* Skip encoder tests on V100

* Fix mulitprocessing jax.distributed.init

* Remove XLA xla_gpu_deterministic_ops which causes segfault

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

* Use jit instead of pjit

---------
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>
Co-authored-by: Phuong Nguyen <36155692+phu0ngng@users.noreply.github.com>

* Switch to nltk>3.8.1 and new data
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@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>

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>

* Cast Flax collections to FrozenDict as WAR to adapt Flax 0.7.1.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding min version of flax to requirements.txt in examples
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix praxis tests and rename compare_frozen_dict to compare_dict
Signed-off-by: Ming Huang <mingh@nvidia.com>

* [Paddle] Refactor FP8 state (#350)

Refactor fp8 state
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Store FP8 checkpointing data in CPU (#351)
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Make test_layer be able to run on both Flax >=0.7.1 and <=0.7.0
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Update Flax version
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Tian Zheng <tizheng@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>

* Fully remove attn_type and set self_attn_mask_type default to 'causal'
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix tests with new arguments
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Explicit self_attn_mask_type for examples
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Update transformer_engine/jax/flax/transformer.py
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: zlsh80826 <rewang@nvidia.com>

* Update transformer_engine/jax/flax/transformer.py
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: zlsh80826 <rewang@nvidia.com>

---------
Signed-off-by: Reese Wang <rewang@nvidia.com>
Signed-off-by: zlsh80826 <rewang@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* Cast the uint32 key/counter to int64
Signed-off-by: Reese Wang <rewang@nvidia.com>

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

* Revise devPtrCuSeqlen* to align the naming
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Support different Jax PRNG impls
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Revert CastAsync since it is not used
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Implement is_training for 16-bit fused attn
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add fused attn with dropout sanity unit tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enhance the comments readability and rng_state checker
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Change the attention dropout shape to align other frameworks
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Make encoder tests deterministic
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Change the default seed for the jax encoder tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Maintain offset in TE
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enhance the resource safety
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Revert rng_state type to allow only i64
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Handle the corner case for elts_per_threads calculation
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Populate rng state by kernels
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Rename rng_state as seed in cpp_extensions
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Update the attention dropout comment
Signed-off-by: Reese Wang <rewang@nvidia.com>

---------
Signed-off-by: Reese Wang <rewang@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Adding JAX/Praxis modules and dependencies.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding UTs to JAX/Praxis modules.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Remove praxis as a dependency due to not strictly needed
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Repalce is_fp8_supported to is_fp8_available
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Make Praxis as an optional dependency.

1. Removed 'from . import praxis' in __init__.py.
    1.1 Noted, keep 'from . import flax' for deprecated warning.
2. Changed te.flax to te_flax in examples and README.rst.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding a workaround to FP8 training on Praxis.
Signed-off-by: Ming Huang <mingh@nvidia.com>

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

bugfix for softmax lowering
Signed-off-by: Ryan Jeng <rjeng@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>

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