Commits · 0db0f4d2d7ca7ae6e761294aedc74b6e30a8aaf4 · OpenDAS / TransformerEngine

06 Oct, 2025 1 commit

[JAX] Fix for GEMM + fuse bias + AllReduce (#2230) · 0db0f4d2

Phuong Nguyen authored Oct 06, 2025



* not fuse bias for output all reduction case + unit tests
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* norm to reduce dgamma along tpsp as well
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

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

* fix test_distributed_layernorm byte counts
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* increase tols for jax_gemm
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

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

0db0f4d2