Commits · 85928d0887234a64c63b220e3c09d8a7a0d01c7b · OpenDAS / TransformerEngine

18 Jul, 2023 1 commit

[JAX] Fully remove attn_type and set self_attn_mask_type default to 'causal' (#324) · a3e4e611

zlsh80826 authored Jul 18, 2023



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

a3e4e611

20 Jun, 2023 1 commit

[JAX] Add self_attn_mask_type and replace attn_type (#273) · 16208b3b

zlsh80826 authored Jun 20, 2023



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

16208b3b

16 May, 2023 1 commit

The Implementation of Praxis's Modules (#158) · b20c0531

Ming-Xu Huang authored May 17, 2023



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

b20c0531