[C/JAX] Comm+GEMM Overlap API for TE/JAX #1337
Commits on Nov 14, 2024
-
added XLA custom op defs for TE GEMM
Signed-off-by: Alp Dener <[email protected]> Added XLA FFI custom op for TE GEMM Signed-off-by: Alp Dener <[email protected]> finished GEMM custom op primitive and serial unit test Signed-off-by: Alp Dener <[email protected]> fixed GEMM custom op batcher Signed-off-by: Alp Dener <[email protected]> fixed output dtype error and contracting dimensions options Signed-off-by: Alp Dener <[email protected]> AG overlap working but executes scatter to match outer LHS dim Signed-off-by: Alp Dener <[email protected]> both all-gather and all-reduce are now working Signed-off-by: Alp Dener <[email protected]> code style Signed-off-by: Alp Dener <[email protected]> changed kwargs in abstract to be explicit Signed-off-by: Alp Dener <[email protected]> added fwd/bwd implementation for non-fp8 gemm Signed-off-by: Alp Dener <[email protected]>
-
fixed batching rules to accommodated batched RHS operand for GEMM
Signed-off-by: Alp Dener <[email protected]>
-
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Commits on Nov 15, 2024
-
fixed batching for collective GEMM FWD and BWD
Signed-off-by: Alp Dener <[email protected]>
-
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
-
propagated batching fixes to fp8_gemm backward pass
Signed-off-by: Alp Dener <[email protected]>
-
Merge branch 'jax-collective-gemm' of github.com:denera/TransformerEn…
…gine into jax-collective-gemm
-
added XLA custom ops and C++ infrastructure for comm+GEMM overlap in …
…TE/JAX Signed-off-by: Alp Dener <[email protected]>
-
comm+GEMM overlap API for TE/JAX compiles, untested, but did not brea…
…k collective GEMM op Signed-off-by: Alp Dener <[email protected]>
-
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci