[JAX] cuBlasMp integration for CollectiveGemm custom op #2361
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Signed-off-by: Alp Dener <[email protected]>
for more information, see https://pre-commit.ci
denera
changed the title
Contributor
Greptile OverviewGreptile SummaryThis PR integrates cuBLASMp library bindings into the JAX CollectiveGemm custom operation to enable communication-computation overlap for distributed GEMM operations. Critical Issues Found:
Implementation Approach:
The code will not compile in its current state when Confidence Score: 0/5
Important Files ChangedFile Analysis
Sequence DiagramsequenceDiagram
participant User as JAX User Code
participant Init as CollectiveGemmInitFFI
participant Registry as CollectiveGemmPlanRegistry
participant Handler as CommunicatorHandler
participant GemmFFI as GemmFFI
participant cuBlasMp as cuBLASMp Context
participant UB as Userbuffers (P2P)
Note over User,UB: Initialization Phase
User->>Init: CollectiveGemmInitFFI(buffer_shape, dtype, collective_op)
Init->>Registry: get_context(buffer_shape, dtype, collective_op)
alt NVTE_WITH_CUBLASMP
Registry->>Handler: get_comm_for_current_device()
Registry->>cuBlasMp: nvte_comm_gemm_ctx_create(comm, nranks, rank)
cuBlasMp-->>Registry: ctx (NVTECommGemmCtx)
else Without cuBLASMp
Registry->>UB: new CommOverlapP2PBase(...)
UB-->>Registry: ctx (CommOverlapCore)
end
Registry-->>Init: context cached in plan_map
Note over User,UB: Computation Phase
User->>GemmFFI: GemmFFI(lhs, rhs, collective_op)
GemmFFI->>Registry: get_context(buffer_shape, dtype, collective_op)
Registry-->>GemmFFI: ctx (from cache)
alt REDUCE_SCATTER with cuBLASMp
GemmFFI->>cuBlasMp: nvte_gemm_reduce_scatter(ctx, m, n, k_global, ...)
cuBlasMp->>cuBlasMp: Split GEMM + ReduceScatter overlap
cuBlasMp-->>GemmFFI: result in output buffer
else ALL_GATHER with cuBLASMp
GemmFFI->>cuBlasMp: nvte_all_gather_gemm(ctx, m, n_global, k, ...)
cuBlasMp->>cuBlasMp: AllGather + GEMM overlap
cuBlasMp-->>GemmFFI: result in output buffer
else Without cuBLASMp (Userbuffers)
GemmFFI->>UB: ctx->split_overlap_rs/ag(...)
UB->>UB: P2P-based comm-compute overlap
UB-->>GemmFFI: result in output buffer
end
GemmFFI-->>User: computation complete
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Comment on lines
+160
to
+164
| #ifndef NVTE_WITH_CUBLASMP | ||
| typedef CollectiveGemmCtx CommOverlapCore; | ||
| #else | ||
| typedef CollectiveGemmCtx CommGemmCtx; | ||
| #endif |
Contributor
There was a problem hiding this comment.
syntax: CollectiveGemmCtx is undefined, creating circular typedef
Suggested change
| #ifndef NVTE_WITH_CUBLASMP | |
| typedef CollectiveGemmCtx CommOverlapCore; | |
| #else | |
| typedef CollectiveGemmCtx CommGemmCtx; | |
| #endif | |
| #ifndef NVTE_WITH_CUBLASMP | |
| typedef CommOverlapCore CollectiveGemmCtx; | |
| #else | |
| typedef NVTECommGemmCtx CollectiveGemmCtx; | |
| #endif |
Comment on lines
+205
to
+207
| ctx = nvte_comm_gemm_ctx_create(comm_handler.get_comm_for_current_device(), | ||
| comm_handler.num_total_devices, comm_handler.get_global_rank(), | ||
| te::cuda::current_device()); |
Contributor
There was a problem hiding this comment.
syntax: function signature mismatch - nvte_comm_gemm_ctx_create takes 3 parameters (comm, nranks, rank) but 4 are passed here
Suggested change
| ctx = nvte_comm_gemm_ctx_create(comm_handler.get_comm_for_current_device(), | |
| comm_handler.num_total_devices, comm_handler.get_global_rank(), | |
| te::cuda::current_device()); | |
| ctx = nvte_comm_gemm_ctx_create(comm_handler.get_comm_for_current_device(), | |
| comm_handler.num_total_devices, comm_handler.get_global_rank()); |
| const int n = (transb) ? lhs_shape[1] : lhs_shape[0]; | ||
| const int k = (transa) ? rhs_shape[1] : rhs_shape[0]; | ||
|
|
||
| k *= ctx->nranks; // convert contracting dimension to global size |
Contributor
There was a problem hiding this comment.
syntax: k is declared const int on line 272, cannot be modified
Suggested change
| k *= ctx->nranks; // convert contracting dimension to global size | |
| int k_global = k * ctx->nranks; // convert contracting dimension to global size |
|
|
||
| k *= ctx->nranks; // convert contracting dimension to global size | ||
| NVTE_CHECK_CUBLASMP( | ||
| nvte_gemm_reduce_scatter(ctx, m, n, k, rhs_.data(), lhs_.data(), out_.data(), bias_.data(), |
Contributor
There was a problem hiding this comment.
syntax: out_ variable is undefined in the NVTE_WITH_CUBLASMP code path (only defined in the #ifndef block above)
Suggested change
| nvte_gemm_reduce_scatter(ctx, m, n, k, rhs_.data(), lhs_.data(), out_.data(), bias_.data(), | |
| nvte_gemm_reduce_scatter(ctx, m, n, k_global, rhs_.data(), lhs_.data(), output->untyped_data(), bias_.data(), |
| const int n = (transb) ? lhs_shape[1] : lhs_shape[0]; | ||
| const int k = (transa) ? rhs_shape[1] : rhs_shape[0]; | ||
|
|
||
| n *= ctx->nranks; // convert all-gathered dimension to global size |
Contributor
There was a problem hiding this comment.
syntax: n is declared const int on line 297, cannot be modified
Suggested change
| n *= ctx->nranks; // convert all-gathered dimension to global size | |
| int n_global = n * ctx->nranks; // convert all-gathered dimension to global size |
|
|
||
| n *= ctx->nranks; // convert all-gathered dimension to global size | ||
| NVTE_CHECK_CUBLASMP( | ||
| nvte_all_gather_gemm(_ctx, m, n, k, rhs_.data(), lhs_.data(), out_.data(), bias_.data(), |
Contributor
There was a problem hiding this comment.
syntax: typo: _ctx should be ctx
Suggested change
| nvte_all_gather_gemm(_ctx, m, n, k, rhs_.data(), lhs_.data(), out_.data(), bias_.data(), | |
| nvte_all_gather_gemm(ctx, m, n_global, k, rhs_.data(), lhs_.data(), output->untyped_data(), bias_.data(), |
Comment on lines
+276
to
+278
| nvte_gemm_reduce_scatter(ctx, m, n, k, rhs_.data(), lhs_.data(), out_.data(), bias_.data(), | ||
| pre_gelu_.data(), rhs_transposed, lhs_transposed, grad, | ||
| use_split_accumulator, 0, stream, CUBLASMP_MATMUL_ALGO_SPLIT_P2P)); |
Contributor
There was a problem hiding this comment.
syntax: wrong enum constant - should use kNVTECommGemmAlgoSplitP2P from NVTECommGemmAlgoType enum instead of undefined CUBLASMP_MATMUL_ALGO_SPLIT_P2P
Suggested change
| nvte_gemm_reduce_scatter(ctx, m, n, k, rhs_.data(), lhs_.data(), out_.data(), bias_.data(), | |
| pre_gelu_.data(), rhs_transposed, lhs_transposed, grad, | |
| use_split_accumulator, 0, stream, CUBLASMP_MATMUL_ALGO_SPLIT_P2P)); | |
| nvte_gemm_reduce_scatter(ctx, m, n, k_global, rhs_.data(), lhs_.data(), output->untyped_data(), bias_.data(), | |
| pre_gelu_.data(), rhs_transposed, lhs_transposed, grad, | |
| use_split_accumulator, 0, stream, kNVTECommGemmAlgoSplitP2P)); |
Comment on lines
+302
to
+304
| nvte_all_gather_gemm(_ctx, m, n, k, rhs_.data(), lhs_.data(), out_.data(), bias_.data(), | ||
| pre_gelu_.data(), rhs_transposed, lhs_transposed, grad, | ||
| use_split_accumulator, 0, stream, CUBLASMP_MATMUL_ALGO_SPLIT_P2P)); |
Contributor
There was a problem hiding this comment.
syntax: wrong enum constant - should use kNVTECommGemmAlgoSplitP2P from NVTECommGemmAlgoType enum instead of undefined CUBLASMP_MATMUL_ALGO_SPLIT_P2P
Suggested change
| nvte_all_gather_gemm(_ctx, m, n, k, rhs_.data(), lhs_.data(), out_.data(), bias_.data(), | |
| pre_gelu_.data(), rhs_transposed, lhs_transposed, grad, | |
| use_split_accumulator, 0, stream, CUBLASMP_MATMUL_ALGO_SPLIT_P2P)); | |
| nvte_all_gather_gemm(ctx, m, n_global, k, rhs_.data(), lhs_.data(), output->untyped_data(), bias_.data(), | |
| pre_gelu_.data(), rhs_transposed, lhs_transposed, grad, | |
| use_split_accumulator, 0, stream, kNVTECommGemmAlgoSplitP2P)); |
Collaborator
|
Hi, could you add some unit tests? |
| CommOverlapCore *CollectiveGemmPlanRegistry::get_executor(std::vector<size_t> buffer_shape, | ||
| DType dtype, | ||
| JAXX_Collective_Op collective_op) { | ||
| CollectiveGemmCtx *CollectiveGemmPlanRegistry::get_context(std::vector<size_t> buffer_shape, |
Collaborator
There was a problem hiding this comment.
Let's restructure this to instead of a compile-time flag, we compile both backend implementations, similar to TE's support for NormPlanRegistry that can use cuDNN or TE kernels. This way a user can easily switch between the two without needing to rebuild TE/JAX
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
This PR integrates TE/common cuBlasMp bindings into the TE/JAX CollectiveGemm custom op.
Type of change
Checklist: