GitLab

* Initially commit for FSDP
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding support to fsdp xmap sharding
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Specify WeightHParamsCollection of fp8 meta.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Support partial FP8 custom calls with FSDP.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding amax reduction on the fsdp mesh dim.
Signed-off-by: Ming Huang <mingh@nvidia.com>

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

* Fix the wrong batch axis in logic_axis_rules and add sharding_constraint to BMM1
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Support FSDP in fMHA.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix missing all-reduce of wgrads along FSDP axis.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Change default value of fsdp_axis_name to  for aligning with others
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix RuntimeError: with_sharding_constraint requires a non-empty
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Slightly changes (review feedback)
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Removed unnecessary comments
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Mergeing input_dp_dim into weight_fsdp_dim_map
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Update transformer_engine/jax/sharding.py
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.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>

* Add self_attn_mask_type and replace attn_type
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Refine the keyword style for the better readability
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Replace attn_type with attn_mask_type in praxis transformer
Signed-off-by: Reese Wang <rewang@nvidia.com>

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