transformer_engine/jax/flax/transformer.py · 403ade2f75402e07ad7f09f01bfb9bd623e860f8 · OpenDAS / TransformerEngine

"cacheflow/model_executor/model_loader.py" did not exist on "1ce13335732c8b28d4a76118821b391f6b219b7c"

[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

transformer.py 51 KB

Replace transformer.py