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

[JAX] Debug distributed attention tests (#1038) · e1e83598

Tim Moon authored Jul 24, 2024



* Remove extra args to fused attention func
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add missing arg to fused attention func
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>

e1e83598

test_distributed_fused_attn.py 10.7 KB

Replace test_distributed_fused_attn.py