tests/jax/test_distributed_fused_attn.py · d770886f02b6c43da1afb65d19487a44cbb8cc88 · OpenDAS / TransformerEngine

"transformer_engine/pytorch/attention/rope.py" did not exist on "94bff09955baa739818f7675e05693b006a4df47"

[JAX] Add `tpsp_resource` in the `MeshResource` map (#2113) · d770886f

Phuong Nguyen authored Aug 26, 2025



* clean up sharding
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

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

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

* rework test for MeshResource
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* add mesh_resource into fp8_autocast in test_helper.py
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

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

d770886f

test_distributed_fused_attn.py 18.9 KB

Replace test_distributed_fused_attn.py