[C/JAX] Comm+GEMM Overlap API for TE/JAX #1337

denera · 2024-11-15T16:33:29Z

Description

>>> Depends on PR #1307 <<<

This PR implements JAX/XLA custom ops and primitives for comm+GEMM overlap kernels in TE/common, and the pure Python/JAX infrastructure required to bootstrap the functionality.

Current limitations and considerations:

Requires a distributed launch with 1 process per GPU and execution with jax.distributed.initialize(). JAX does not have its own distributed launch utility like torchrun, so this is typically done with mpirun launch + mpi4py in Python.
TE has to be compiled with NVTE_UB_WITH_MPI=1 and Userbuffers has to be bootstrapped with MPI because XLA custom ops cannot execute XLA collectives. Unlike PyTorch, this does not introduce a new dependency because distributed launch with JAX already depends on MPI.
Userbuffers communication buffers are allocated outside of the XLA memory pool. Since XLA has no knowledge of these allocations, its memory allowance as a % of total device memory needs to be decreased to avoid OOM issues.

To do:
[x] Implement XLA custom ops w/ both old API and new FFI interfaces.
[x] Extend JAX CollectiveGemmPrimitive to support comm+GEMM overlap.
[x] Implement bootstrapping and utility functions with PyBind11 bindings.
[x] Verify that comm+GEMM overlap extensions do not break non-overlap collective GEMM functionality.
[ ] Add new unit tests for comm+GEMM overlap.

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refractor

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

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]>

Signed-off-by: Alp Dener <[email protected]>

for more information, see https://pre-commit.ci

… passing test Signed-off-by: Alp Dener <[email protected]>

for more information, see https://pre-commit.ci

Signed-off-by: Alp Dener <[email protected]>

for more information, see https://pre-commit.ci

Signed-off-by: Alp Dener <[email protected]>

for more information, see https://pre-commit.ci

Signed-off-by: Alp Dener <[email protected]>

for more information, see https://pre-commit.ci

…ide the custom op Signed-off-by: Alp Dener <[email protected]>

…xt-parallel LHS operands Signed-off-by: Alp Dener <[email protected]>

for more information, see https://pre-commit.ci

Signed-off-by: Alp Dener <[email protected]>

… and TP-only meshes Signed-off-by: Alp Dener <[email protected]>

for more information, see https://pre-commit.ci

Signed-off-by: Alp Dener <[email protected]>

…TE/JAX Signed-off-by: Alp Dener <[email protected]> comm+GEMM overlap API for TE/JAX compiles, untested, but did not break 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 fixed static args 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

Signed-off-by: Alp Dener <[email protected]>

for more information, see https://pre-commit.ci

Signed-off-by: Alp Dener <[email protected]>

for more information, see https://pre-commit.ci

Signed-off-by: Alp Dener <[email protected]>

for more information, see https://pre-commit.ci

Signed-off-by: Alp Dener <[email protected]>

for more information, see https://pre-commit.ci

Signed-off-by: Alp Dener <[email protected]>

…TransformerEngine into jax-collective-gemm-with-overlap

for more information, see https://pre-commit.ci

abhinavgoel95 · 2024-12-17T23:47:21Z

transformer_engine/jax/cpp_extensions/gemm.py

+
+    (out, out_amax, out_scale, pre_gelu_out, _, extra_out) = (  # bias_grad in non-FP8 GEMM
+        CollectiveGemmPrimitive.outer_primitive.bind(
+            rhs_t,


@denera the order of these ops need to change.

abhinavgoel95 · 2024-12-17T23:47:57Z

transformer_engine/jax/cpp_extensions/gemm.py

+    comm_overlap_config: Optional[dict] = None,
+) -> Tuple[ArrayLike, ...]:
+    """FP8 mat-mul with `nvte_cublas_gemm()` custom op."""
+    out_shape_batched = (*lhs.shape[:-2], lhs.shape[-1], rhs_t.shape[-1])


The output shape looks wrong to me. Should be (*lhs.shape[:-2], lhs.shape[-2], rhs_t.shape[-2]).
@denera

abhinavgoel95 · 2024-12-18T20:55:39Z

transformer_engine/jax/cpp_extensions/gemm.py

+                out_amax_updated_dtype == out_scale_updated_dtype == jnp.float32
+            ), "Invalid output amax or scale dtype."
+        else:
+            assert out_dtype == lhs_dtype, (


@denera this assertion is wrong for FP8, and needs to be guarded

abhinavgoel95 · 2024-12-18T21:42:25Z

transformer_engine/jax/cpp_extensions/gemm.py

+        if lhs_2d_shape is not None and lhs.ndim > 2:
+            lhs = jax.lax.reshape(lhs, lhs_2d_shape, dimensions=lhs_layout)
+            if jax_dtype_is_fp8(lhs.dtype):
+                lhs = jax.lax.transpose(lhs, (1, 0))


@denera do we need this transpose on the LHS? It seems wrong to me

denera added enhancement New feature or request jax labels Nov 15, 2024

denera requested review from nouiz, phu0ngng and ksivaman November 15, 2024 16:33

denera self-assigned this Nov 15, 2024

huanghua1994 self-requested a review November 15, 2024 17:35

denera force-pushed the jax-collective-gemm-with-overlap branch 2 times, most recently from 11ad5ec to e44d5cf Compare November 21, 2024 20:38

denera force-pushed the jax-collective-gemm-with-overlap branch from c35b351 to 616e301 Compare December 3, 2024 14:08

denera mentioned this pull request Dec 5, 2024

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

Open

17 tasks

denera and others added 19 commits December 5, 2024 21:33

fixed batching rules to accommodated batched RHS operand for GEMM

c9774d8

Signed-off-by: Alp Dener <[email protected]>

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

e523018

for more information, see https://pre-commit.ci

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

2c3dbf1

… passing test Signed-off-by: Alp Dener <[email protected]>

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

448eaa9

for more information, see https://pre-commit.ci

batched operands for GEMM custom op seem to be working now

cb6ae3c

Signed-off-by: Alp Dener <[email protected]>

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

6f67355

for more information, see https://pre-commit.ci

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

4b2b2d4

Signed-off-by: Alp Dener <[email protected]>

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

2b2753e

for more information, see https://pre-commit.ci

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

969f597

Signed-off-by: Alp Dener <[email protected]>

fixed logic to remove FSDP sharding

ce86dcb

Signed-off-by: Alp Dener <[email protected]>

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

b215f20

for more information, see https://pre-commit.ci

retained FSDP dims and pushed FSDP all-gather of weight array to outs…

cbab16c

…ide the custom op Signed-off-by: Alp Dener <[email protected]>

Added useful warning about DGRAD sharding not matching sequence/conte…

0ea55c0

…xt-parallel LHS operands Signed-off-by: Alp Dener <[email protected]>

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

2acb92f

for more information, see https://pre-commit.ci

documentation fixes

b07bb2d

Signed-off-by: Alp Dener <[email protected]>

added unit test, both AG+GEMM and GEMM+AR passing with FSDP+TP, DP+TP…

765b844

… and TP-only meshes Signed-off-by: Alp Dener <[email protected]>

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

2ce4377

for more information, see https://pre-commit.ci

restored old test_custom_call_compute.py to remove erroneous changes

f68d71e

Signed-off-by: Alp Dener <[email protected]>

denera and others added 9 commits December 5, 2024 21:34

AG+GEMM overlap working

b306608

Signed-off-by: Alp Dener <[email protected]>

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

aa16726

for more information, see https://pre-commit.ci

added comm+GEMM overlap example script

a569e3b

Signed-off-by: Alp Dener <[email protected]>

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

69db12e

for more information, see https://pre-commit.ci

RS overlap also works

ec2d5ae

Signed-off-by: Alp Dener <[email protected]>

added missing copy of AG+GEMM input into comm buffer

8fe3942

Signed-off-by: Alp Dener <[email protected]>

updated FWD/BWD wrappers for non-FP8 and FP8 gemm

adf4046

Signed-off-by: Alp Dener <[email protected]>

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

c4c608b

for more information, see https://pre-commit.ci

denera force-pushed the jax-collective-gemm-with-overlap branch from 39f0375 to c4c608b Compare December 5, 2024 21:34

denera and others added 5 commits December 5, 2024 22:42

added more documentation to the TE/JAX comm+GEMM overlap example

4707df3

Signed-off-by: Alp Dener <[email protected]>

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

18a6249

for more information, see https://pre-commit.ci

fixed RS overlap in the example

b144941

Signed-off-by: Alp Dener <[email protected]>

Merge branch 'jax-collective-gemm-with-overlap' of github.com:denera/…

76b26bc

…TransformerEngine into jax-collective-gemm-with-overlap

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

6ad5651

for more information, see https://pre-commit.ci

denera mentioned this pull request Dec 14, 2024

[TP comm overlap unit test]CUDA Error: misaligned address error when testing with recent cublas (or pytorch container) #1332

Open

abhinavgoel95 reviewed Dec 17, 2024

View reviewed changes

abhinavgoel95 reviewed Dec 18, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[C/JAX] Comm+GEMM Overlap API for TE/JAX #1337

[C/JAX] Comm+GEMM Overlap API for TE/JAX #1337

denera commented Nov 15, 2024

abhinavgoel95 Dec 17, 2024

abhinavgoel95 Dec 17, 2024

abhinavgoel95 Dec 18, 2024

abhinavgoel95 Dec 18, 2024

[C/JAX] Comm+GEMM Overlap API for TE/JAX #1337

Are you sure you want to change the base?

[C/JAX] Comm+GEMM Overlap API for TE/JAX #1337

Conversation

denera commented Nov 15, 2024

Description

>>> Depends on PR #1307 <<<

Type of change

Checklist:

abhinavgoel95 Dec 17, 2024

Choose a reason for hiding this comment

abhinavgoel95 Dec 17, 2024

Choose a reason for hiding this comment

abhinavgoel95 Dec 18, 2024

Choose a reason for hiding this comment

abhinavgoel95 Dec 18, 2024

Choose a reason for hiding this comment