[JAX] Collective GEMM custom op with `nvte_cublas_gemm` (no comm. overlap)#1307

Open

denera wants to merge 19 commits intoNVIDIA:mainfrom denera:jax-collective-gemm

+1,944-121

Commits on Dec 5, 2024

added XLA custom op defs for TE GEMM
denera
committed
fixed batching rules to accommodated batched RHS operand for GEMM
denera
committed
[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot]
authored and
denera
committed
re-applied bug fixes to working older version, updated backward pass, passing test
denera
committed
[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot]
authored and
denera
committed
batched operands for GEMM custom op seem to be working now
denera
committed
[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot]
authored and
denera
committed
fixed batch size 1 issue and enabled FSDP sharding for RHS operand
denera
committed
[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot]
authored and
denera
committed
fixed FSDP+TP w/ DP=1 and TP+DP, but FSDP+TP w/ DP>1 still crashes
denera
committed
fixed logic to remove FSDP sharding
denera
committed
[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot]
authored and
denera
committed
retained FSDP dims and pushed FSDP all-gather of weight array to outside the custom op
denera
committed
Added useful warning about DGRAD sharding not matching sequence/context-parallel LHS operands
denera
committed
[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot]
authored and
denera
committed
documentation fixes
denera
committed
added unit test, both AG+GEMM and GEMM+AR passing with FSDP+TP, DP+TP and TP-only meshes
denera
committed
[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot]
authored and
denera
committed
restored old test_custom_call_compute.py to remove erroneous changes
denera
committed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[JAX] Collective GEMM custom op with `nvte_cublas_gemm` (no comm. overlap)#1307

Commits on Dec 5, 2024

added XLA custom op defs for TE GEMM

fixed batching rules to accommodated batched RHS operand for GEMM

[pre-commit.ci] auto fixes from pre-commit.com hooks

re-applied bug fixes to working older version, updated backward pass, passing test

[pre-commit.ci] auto fixes from pre-commit.com hooks

batched operands for GEMM custom op seem to be working now

[pre-commit.ci] auto fixes from pre-commit.com hooks

fixed batch size 1 issue and enabled FSDP sharding for RHS operand

[pre-commit.ci] auto fixes from pre-commit.com hooks

fixed FSDP+TP w/ DP=1 and TP+DP, but FSDP+TP w/ DP>1 still crashes

fixed logic to remove FSDP sharding

[pre-commit.ci] auto fixes from pre-commit.com hooks

retained FSDP dims and pushed FSDP all-gather of weight array to outside the custom op

Added useful warning about DGRAD sharding not matching sequence/context-parallel LHS operands

[pre-commit.ci] auto fixes from pre-commit.com hooks

documentation fixes

added unit test, both AG+GEMM and GEMM+AR passing with FSDP+TP, DP+TP and TP-only meshes

[pre-commit.ci] auto fixes from pre-commit.com hooks

restored old test_custom_call_compute.py to remove erroneous changes