examples/jax/encoder/test_multigpu_encoder.py · d770886f02b6c43da1afb65d19487a44cbb8cc88 · OpenDAS / TransformerEngine

"vscode:/vscode.git/clone" did not exist on "2acd76f346efcdff4f6ca1d92fe1575c448e4b70"

[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_multigpu_encoder.py 18.3 KB

Replace test_multigpu_encoder.py