[JAX][Draft] Async issuing D2H memcpy for grouped_gemm group_sizes array #2213
Conversation
Signed-off-by: Hua Huang <[email protected]>
Signed-off-by: Hua Huang <[email protected]>
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
I think it is a good improvement for now.
We should probably provide a GroupedLayerNormMLP VJP op, which encloses the grouped_gemm_copy_group_sizes function and the use_async_d2h_group_sizes option so that we don't expose these two to users as they can be pretty bug-prone.
| "supported number ", max_num_gemms, " to be downloaded in advance."); | ||
| host_num_gemms = num_gemms; | ||
| // Wait for current compute stream to finish | ||
| cudaStream_t compute_stream_0 = nvte_get_compute_stream(0); |
There was a problem hiding this comment.
@mingxu1067 could you check if this causes the same stream sync issue as last time when we used the compute_stream(0) instead of the stream given by XLA?
There was a problem hiding this comment.
Just a note: this part follows the logic in https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/common/gemm/cublaslt_gemm.cu#L915
| auto init = [&]() { | ||
| NVTE_CHECK_CUDA(cudaEventCreate(&d2h_event)); | ||
| NVTE_CHECK_CUDA(cudaMallocHost(&host_group_sizes_internal, sizeof(int32_t) * max_num_gemms)); | ||
| }; |
There was a problem hiding this comment.
If this causes any issues, we could consider moving this allocation into the FFI prepare phase.
Description
This is a draft PR for saving some work and discussion.
Recently we used TE/JAX's
grouped_gemm()interface for a MoE model's inference. Nsys shows a GPU bubble whengrouped_gemm()is copying thegroup_sizesarray from device to host. This is a known issue when we were designing thegrouped_gemm()interface. It's performance impact for training / inference prefill stage is relatively small but cannot be ignored in inference decode stage. This draft aims to partially address the bubble issue.Our target model uses MLP-MoE, i.e., each expert is a MLP layer. After fusing GEMMs, each MLP-MoE layer needs two
grouped_gemm()with the samegroup_sizesarray. This PR allows issuing an async D2H copy of thegroup_sizearray before enteringgrouped_gemm(), thengrouped_gemm()can reuse the downloadedgroup_sizes. We have validated the correctness of the implementation in this PR in our target model.This PR does not solve the issue of breaking CUDA graph in
grouped_gemm()since in the async copy mode it still needs to callcudaEventSynchronize(). Furthermore, the D2H memcpy does not overlap with other operations for copying and dispatching tokens to experts in our implementation for the target model, since those JAX-native operations are captured and executed in CUDA graph, while the async D2H copy does not support CUDA graph.@phu0ngng @mingxu1067 Please let me know your comments and suggestions. Much appreciated!
Type of change
Changes
GroupedGemmCopySizesPrimitivefor async copying ofgroup_sizesfrom GPU to hostuse_async_d2h_group_sizesforgrouped_gemm(), the default value isFalseso the original code path will be used by defaultChecklist: