[GH-ISSUE #272] Does Deepep's Dispatch/Combine Support Autograd Like PyTorch's all_to_all?

bsakaguchi commented

2026-10-26 09:56:27 +00:00

Owner

Originally created by @aicodeclub on GitHub (Jul 2, 2025).
Original GitHub issue: https://github.com/deepseek-ai/DeepEP/issues/272

Does the Deepep interface Dispatch/Combine have automatic gradient backpropagation functionality similar to torch.distributed.nn.functional.all_to_all? The built-in Autograd feature of torch.distributed.nn.functional.all_to_all automatically propagates the gradients of tensors exchanged during All-to-All communication in the forward pass back to the corresponding GPUs correctly. If the Deepep interface lacks this functionality, what are the possible solutions?

Originally created by @aicodeclub on GitHub (Jul 2, 2025). Original GitHub issue: https://github.com/deepseek-ai/DeepEP/issues/272 Does the Deepep interface Dispatch/Combine have automatic gradient backpropagation functionality similar to torch.distributed.nn.functional.all_to_all? The built-in Autograd feature of torch.distributed.nn.functional.all_to_all automatically propagates the gradients of tensors exchanged during All-to-All communication in the forward pass back to the corresponding GPUs correctly. If the Deepep interface lacks this functionality, what are the possible solutions?

bsakaguchi commented

2026-10-26 09:56:28 +00:00

Author

Owner

@LyricZhao commented on GitHub (Jul 2, 2025):

We support backward, but you have to write the autograd function by yourself, as the example in README shows.

def dispatch_forward(x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
                     topk_idx: torch.Tensor, topk_weights: torch.Tensor,
                     num_experts: int, previous_event: Optional[EventOverlap] = None) -> \
        Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], torch.Tensor, torch.Tensor, List, Tuple, EventOverlap]:
    # NOTES: an optional `previous_event` means a CUDA event captured that you want to make it as a dependency 
    # of the dispatch kernel, it may be useful with communication-computation overlap. For more information, please
    # refer to the docs of `Buffer.dispatch`
    global _buffer

    # Calculate layout before actual dispatch
    num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, is_token_in_rank, previous_event = \
        _buffer.get_dispatch_layout(topk_idx, num_experts,
                                    previous_event=previous_event, async_finish=True,
                                    allocate_on_comm_stream=previous_event is not None)
    # Do MoE dispatch
    # NOTES: the CPU will wait for GPU's signal to arrive, so this is not compatible with CUDA graph
    # Unless you specify `num_worst_tokens`, but this flag is for intranode only
    # For more advanced usages, please refer to the docs of the `dispatch` function
    recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, event = \
        _buffer.dispatch(x, topk_idx=topk_idx, topk_weights=topk_weights,
                         num_tokens_per_rank=num_tokens_per_rank, num_tokens_per_rdma_rank=num_tokens_per_rdma_rank,
                         is_token_in_rank=is_token_in_rank, num_tokens_per_expert=num_tokens_per_expert,
                         previous_event=previous_event, async_finish=True,
                         allocate_on_comm_stream=True)
    # For event management, please refer to the docs of the `EventOverlap` class
    return recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, event


def dispatch_backward(grad_recv_x: torch.Tensor, grad_recv_topk_weights: torch.Tensor, handle: Tuple) -> \
        Tuple[torch.Tensor, torch.Tensor, EventOverlap]:
    global _buffer

    # The backward process of MoE dispatch is actually a combine
    # For more advanced usages, please refer to the docs of the `combine` function
    combined_grad_x, combined_grad_recv_topk_weights, event = \
        _buffer.combine(grad_recv_x, handle, topk_weights=grad_recv_topk_weights, async_finish=True)

    # For event management, please refer to the docs of the `EventOverlap` class
    return combined_grad_x, combined_grad_recv_topk_weights, event


def combine_forward(x: torch.Tensor, handle: Tuple, previous_event: Optional[EventOverlap] = None) -> \
        Tuple[torch.Tensor, EventOverlap]:
    global _buffer

    # Do MoE combine
    # For more advanced usages, please refer to the docs of the `combine` function
    combined_x, _, event = _buffer.combine(x, handle, async_finish=True, previous_event=previous_event,
                                           allocate_on_comm_stream=previous_event is not None)

    # For event management, please refer to the docs of the `EventOverlap` class
    return combined_x, event


def combine_backward(grad_combined_x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
                     handle: Tuple, previous_event: Optional[EventOverlap] = None) -> \
        Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], EventOverlap]:
    global _buffer

    # The backward process of MoE combine is actually a dispatch
    # For more advanced usages, please refer to the docs of the `dispatch` function
    grad_x, _, _, _, _, event = _buffer.dispatch(grad_combined_x, handle=handle, async_finish=True,
                                                 previous_event=previous_event,
                                                 allocate_on_comm_stream=previous_event is not None)

    # For event management, please refer to the docs of the `EventOverlap` class
    return grad_x, event

You should define a class derived from torch.autograd.Function, and write similar code as above into def forward and def backward.

@LyricZhao commented on GitHub (Jul 2, 2025): We support backward, but you have to write the autograd function by yourself, as the example in README shows. ```python def dispatch_forward(x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], topk_idx: torch.Tensor, topk_weights: torch.Tensor, num_experts: int, previous_event: Optional[EventOverlap] = None) -> \ Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], torch.Tensor, torch.Tensor, List, Tuple, EventOverlap]: # NOTES: an optional `previous_event` means a CUDA event captured that you want to make it as a dependency # of the dispatch kernel, it may be useful with communication-computation overlap. For more information, please # refer to the docs of `Buffer.dispatch` global _buffer # Calculate layout before actual dispatch num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, is_token_in_rank, previous_event = \ _buffer.get_dispatch_layout(topk_idx, num_experts, previous_event=previous_event, async_finish=True, allocate_on_comm_stream=previous_event is not None) # Do MoE dispatch # NOTES: the CPU will wait for GPU's signal to arrive, so this is not compatible with CUDA graph # Unless you specify `num_worst_tokens`, but this flag is for intranode only # For more advanced usages, please refer to the docs of the `dispatch` function recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, event = \ _buffer.dispatch(x, topk_idx=topk_idx, topk_weights=topk_weights, num_tokens_per_rank=num_tokens_per_rank, num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, is_token_in_rank=is_token_in_rank, num_tokens_per_expert=num_tokens_per_expert, previous_event=previous_event, async_finish=True, allocate_on_comm_stream=True) # For event management, please refer to the docs of the `EventOverlap` class return recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, event def dispatch_backward(grad_recv_x: torch.Tensor, grad_recv_topk_weights: torch.Tensor, handle: Tuple) -> \ Tuple[torch.Tensor, torch.Tensor, EventOverlap]: global _buffer # The backward process of MoE dispatch is actually a combine # For more advanced usages, please refer to the docs of the `combine` function combined_grad_x, combined_grad_recv_topk_weights, event = \ _buffer.combine(grad_recv_x, handle, topk_weights=grad_recv_topk_weights, async_finish=True) # For event management, please refer to the docs of the `EventOverlap` class return combined_grad_x, combined_grad_recv_topk_weights, event def combine_forward(x: torch.Tensor, handle: Tuple, previous_event: Optional[EventOverlap] = None) -> \ Tuple[torch.Tensor, EventOverlap]: global _buffer # Do MoE combine # For more advanced usages, please refer to the docs of the `combine` function combined_x, _, event = _buffer.combine(x, handle, async_finish=True, previous_event=previous_event, allocate_on_comm_stream=previous_event is not None) # For event management, please refer to the docs of the `EventOverlap` class return combined_x, event def combine_backward(grad_combined_x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], handle: Tuple, previous_event: Optional[EventOverlap] = None) -> \ Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], EventOverlap]: global _buffer # The backward process of MoE combine is actually a dispatch # For more advanced usages, please refer to the docs of the `dispatch` function grad_x, _, _, _, _, event = _buffer.dispatch(grad_combined_x, handle=handle, async_finish=True, previous_event=previous_event, allocate_on_comm_stream=previous_event is not None) # For event management, please refer to the docs of the `EventOverlap` class return grad_x, event ``` You should define a class derived from `torch.autograd.Function`, and write similar code as above into `def forward` and `def backward`.

bsakaguchi commented

2026-10-26 09:56:29 +00:00

Author

Owner

@aicodeclub commented on GitHub (Jul 2, 2025):

Thanks!

@aicodeclub commented on GitHub (Jul 2, 2025): Thanks!

bsakaguchi commented

2026-10-26 09:56:29 +00:00

Author

Owner

@aicodeclub commented on GitHub (Jul 3, 2025):

@LyricZhao
I define a class derived from torch.autograd.Function, and define forward and backward like this:
forward(): sequentially calls dispatch_forward followed by combine_forward;
backward(): sequentially calls combine_backward followed by dispatch_backward.
Does this implementation have any problem? Thanks

@aicodeclub commented on GitHub (Jul 3, 2025): @LyricZhao I define a class derived from torch.autograd.Function, and define forward and backward like this: forward(): sequentially calls dispatch_forward followed by combine_forward; backward(): sequentially calls combine_backward followed by dispatch_backward. Does this implementation have any problem? Thanks

Rows
Columns

[GH-ISSUE #272] Does Deepep's Dispatch/Combine Support Autograd Like PyTorch's all_to_all? #199