Skip to content

[main](moe):Refine A2A overlap under CUDA_DEVICE_MAX_CONNECTIONS=1 #2719

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Baidu-AIAK wants to merge 1 commit into from
Closed

[main](moe):Refine A2A overlap under CUDA_DEVICE_MAX_CONNECTIONS=1 #2719

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
196 changes: 153 additions & 43 deletions megatron/core/models/common/model_chunk_schedule_plan.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
from contextlib import nullcontext
from typing import Optional

import os
import torch
from torch import Tensor

Expand Down Expand Up @@ -40,6 +41,7 @@ class TransformerLayerSchedulePlan:
├── moe_dispatch (TransformerLayerNode): dispatch All2All
├── mlp (TransformerLayerNode): mlp module
├── moe_combine (TransformerLayerNode): combine All2All
├── post_combine (TransformerLayerNode): combine post process -> mlp_bda
└── mtp_post_process (PostProcessNode): mtp post process

Note that MTP layer has the same operation and execution order with TransformerLayer regarding
Expand All @@ -55,6 +57,7 @@ class TransformerLayerSchedulePlan:
moe_dispatch = None
mlp = None
moe_combine = None
post_combine = None
mtp_post_process = None

def __init__(self, layer, event, chunk_state, comp_stream, comm_stream, extra_args={}):
Expand Down Expand Up @@ -137,6 +140,7 @@ def create_node(stream, module, name):
moe_dispatch_module,
mlp_module,
moe_combine_module,
post_combine_module,
mtp_post_process_module,
) = fwd_callables

Expand All @@ -148,10 +152,12 @@ def create_node(stream, module, name):
self.post_attn = create_node(comp_stream, post_attn_module, "post_attn")
self.moe_dispatch = create_node(comm_stream, moe_dispatch_module, "moe_dispatch")
self.moe_combine = create_node(comm_stream, moe_combine_module, "moe_combine")
self.post_combine = create_node(comp_stream, post_combine_module, "post_combine")
else:
self.post_attn = NoopScheduleNode()
self.moe_dispatch = NoopScheduleNode()
self.moe_combine = NoopScheduleNode()
self.post_combine = NoopScheduleNode()

if is_mtp:
self.mtp_post_process = create_node(
Expand All @@ -174,68 +180,152 @@ def get_fp8_context(self):
)

@staticmethod
def run(f_layer, b_layer, f_input=None, b_grad=None, is_last_layer_in_bwd=False):
def run(
f_layer,
b_layer,
f_input=None,
b_grad=None,
is_last_layer_in_bwd=False,
fine_grained_overlap=False,
post_forward=None,
post_backward=None,
f_schedule_plan=None,
b_schedule_plan=None,
is_last_layer=False,
):
"""Schedule one-forward-one-backward operations for a single transformer layer.

This function interleaves forward and backward operations, overlapping the communications
(dispatch or combine) of one with the computations (att or mlp) of the other
to maximize parallelism and efficiency.

Two execution modes are supported:

1) Coarse-grained overlap (fine_grained_overlap = False):

When f_layer and b_layer are not None, forward and backward pass are overlapped as follows:
comm_stream: combine_bwd | dispatch_fwd->dispatch_bwd | combine_fwd
comp_stream: attn_fwd->post_attn_fwd| mlp_bwd->mlp_bwd_dw->mlp_fwd| post_attn_bwd->attn_bwd
For MTP, mtp_post_process_fwd is executed after the combine_fwd in the comp_stream,
and mtp_post_process_bwd is executed before the combine_bwd in the comp_stream.

2) Fine-grained overlap (fine_grained_overlap = True):

This mode further decomposes communication and computation into smaller
stages and interleaves them more aggressively, including post-processing
steps (e.g., MTP / post-combine hooks).

The execution timeline becomes:

comm_stream: combine_bwd | dispatch_fwd | dispatch_bwd | combine_fwd | PP_fwd | PP_bwd |
comp_stream: post_combine_bwd → attn_fwd → post_attn_fwd | mlp_bwd | mlp_fwd | mlp_bwd_dw → post_attn_bwd → post_combine_fwd | attn_bwd | attn_bwd_dw |

Args:
f_layer (TransformerLayerSchedulePlan): Forward layer (for current microbatch)
b_layer (TransformerLayerSchedulePlan): Backward layer (for previous microbatch)
f_input (Tensor): Input for forward computation
b_grad (Tensor): Gradient for backward computation
is_last_layer_in_bwd (bool):
Whether the current layer is the last layer in the backward pass.

fine_grained_overlap (bool):
Enable fine-grained communication / computation overlap
post_forward (callable or None):
The function to call after the forward pass
post_backward (callable or None):
The function to call after the backward pass
f_schedule_plan (TransformerModelChunkSchedulePlan):
The forward schedule plan
b_schedule_plan (TransformerModelChunkSchedulePlan):
The backward schedule plan
is_last_layer (bool):
Whether the current layer is the overlap boundary layer of the current chunk
Returns:
Functions or values for next iteration's computation
"""

if b_layer is not None:
b_grad = b_layer.mtp_post_process.backward(b_grad)
b_grad = b_layer.post_combine.backward(b_grad)
b_grad = b_layer.moe_combine.backward(b_grad)

if f_layer is not None:
with f_layer.get_fp8_context():
f_input = f_layer.attn.forward(f_input)
f_input = f_layer.post_attn.forward(f_input)

if b_layer is not None:
b_grad = b_layer.mlp.backward(b_grad)
if fine_grained_overlap:
if f_layer is not None:
with f_layer.get_fp8_context():
f_input = f_layer.moe_dispatch.forward(f_input)

if b_layer is not None:
b_grad = b_layer.mlp.backward(b_grad)

if b_layer is not None:
b_grad = b_layer.moe_dispatch.backward(b_grad)

if f_layer is not None:
with f_layer.get_fp8_context():
f_input = f_layer.mlp.forward(f_input)

if f_layer is not None:
with f_layer.get_fp8_context():
f_input = f_layer.moe_combine.forward(f_input)

if b_layer is not None:
b_layer.mlp.backward_dw()
b_grad = b_layer.post_attn.backward(b_grad)

if f_layer is not None:
with f_layer.get_fp8_context():
f_input = f_layer.post_combine.forward(f_input)
f_input = f_layer.mtp_post_process.forward(f_input)

if is_last_layer:
if f_schedule_plan is not None and post_forward is not None:
f_schedule_plan.wait_current_stream()
post_forward(f_input, f_schedule_plan.vp_stage)

if b_layer is not None:
b_grad = b_layer.attn.backward(b_grad)

if is_last_layer:
if b_schedule_plan is not None and post_backward is not None:
b_schedule_plan.wait_current_stream()
post_backward(b_grad, b_schedule_plan.vp_stage)

if b_layer is not None:
b_layer.attn.backward_dw()
else:
if b_layer is not None:
b_grad = b_layer.mlp.backward(b_grad)

if f_layer is not None:
with f_layer.get_fp8_context():
f_input = f_layer.moe_dispatch.forward(f_input)
if f_layer is not None:
with f_layer.get_fp8_context():
f_input = f_layer.moe_dispatch.forward(f_input)

if b_layer is not None:
b_layer.mlp.backward_dw()
b_grad = b_layer.moe_dispatch.backward(b_grad)
if b_layer is not None:
b_layer.mlp.backward_dw()
b_grad = b_layer.moe_dispatch.backward(b_grad)

if f_layer is not None:
with f_layer.get_fp8_context():
f_input = f_layer.mlp.forward(f_input)
if f_layer is not None:
with f_layer.get_fp8_context():
f_input = f_layer.mlp.forward(f_input)

if f_layer is not None:
with f_layer.get_fp8_context():
f_input = f_layer.moe_combine.forward(f_input)
f_input = f_layer.mtp_post_process.forward(f_input)
if f_layer is not None:
with f_layer.get_fp8_context():
f_input = f_layer.moe_combine.forward(f_input)
f_input = f_layer.post_combine.forward(f_input)
f_input = f_layer.mtp_post_process.forward(f_input)

if b_layer is not None:
b_grad = b_layer.post_attn.backward(b_grad)
b_grad = b_layer.attn.backward(b_grad)
if b_layer is not None:
b_grad = b_layer.post_attn.backward(b_grad)
b_grad = b_layer.attn.backward(b_grad)

# Delay the last attn_dw in backward pass (attn_dw of the first layer)
# for overlapping with the p2p comm
if b_layer is not None and not is_last_layer_in_bwd:
b_layer.attn.backward_dw()
# Delay the last attn_dw in backward pass (attn_dw of the first layer)
# for overlapping with the p2p comm
if b_layer is not None and not is_last_layer_in_bwd:
b_layer.attn.backward_dw()

return f_input, b_grad

Expand Down Expand Up @@ -444,6 +534,9 @@ def run(
f_num_layers = f_schedule_plan.num_layers() if f_schedule_plan is not None else 0
b_num_layers = b_schedule_plan.num_layers() if b_schedule_plan is not None else 0
overlapped_layers = min(f_num_layers, b_num_layers)
equal_layers = (f_num_layers == b_num_layers)
fine_grained_overlap = os.environ.get("CUDA_DEVICE_MAX_CONNECTIONS") == "1"


# combined forward and backward pass for overlapped layers
for i in range(overlapped_layers):
Expand All @@ -456,6 +549,12 @@ def run(
f_input=f_input,
b_grad=b_grad,
is_last_layer_in_bwd=(i == b_num_layers - 1),
fine_grained_overlap = fine_grained_overlap,
post_forward=post_forward if fine_grained_overlap else None,
post_backward=post_backward if fine_grained_overlap else None,
f_schedule_plan=f_schedule_plan if fine_grained_overlap else None,
b_schedule_plan=b_schedule_plan if fine_grained_overlap else None,
is_last_layer=(i == overlapped_layers - 1 and equal_layers) if fine_grained_overlap else False,
)
torch.cuda.nvtx.range_pop()

Expand All @@ -464,34 +563,45 @@ def run(
b_layer = b_schedule_plan.get_layer(b_num_layers - 1 - i)
torch.cuda.nvtx.range_push(f"layer_{b_num_layers - 1 - i}b")
_, b_grad = TransformerLayerSchedulePlan.run(
None, b_layer, b_grad=b_grad, is_last_layer_in_bwd=(i == b_num_layers - 1)
None,
b_layer,
b_grad=b_grad,
is_last_layer_in_bwd=(i == b_num_layers - 1),
fine_grained_overlap = fine_grained_overlap,
)
torch.cuda.nvtx.range_pop()

# forward pass for the remaining layers
for i in range(overlapped_layers, f_num_layers):
f_layer = f_schedule_plan.get_layer(i)
torch.cuda.nvtx.range_push(f"layer_{i}f")
f_input, _ = TransformerLayerSchedulePlan.run(f_layer, None, f_input=f_input)
f_input, _ = TransformerLayerSchedulePlan.run(
f_layer,
None,
f_input=f_input,
fine_grained_overlap = fine_grained_overlap,
)
torch.cuda.nvtx.range_pop()

if f_schedule_plan is not None and post_forward is not None:
# post_forward()/send_forward_recv_forward() is running in the communication stream,
# so the p2p comm could be overlapped with the attn backward
with torch.cuda.stream(get_comm_stream()):
f_schedule_plan.wait_current_stream()
post_forward(f_input, f_schedule_plan.vp_stage)

# post_backward()/send_backward_recv_backward() is running in the computation stream,
# so the p2p comm could be overlapped with the wgrad of attn backward
if b_schedule_plan is not None and post_backward is not None:
b_schedule_plan.wait_current_stream()
post_backward(b_grad, b_schedule_plan.vp_stage)

# Delay the last attn_dw in backward pass (attn_dw of the first layer)
# for overlapping with the p2p comm
if b_num_layers > 0:
b_schedule_plan.get_layer(0).attn.backward_dw()
if (not equal_layers and fine_grained_overlap) or not fine_grained_overlap:
if f_schedule_plan is not None and post_forward is not None:
# post_forward()/send_forward_recv_forward() is running in the communication stream,
# so the p2p comm could be overlapped with the attn backward
with torch.cuda.stream(get_comm_stream()):
f_schedule_plan.wait_current_stream()
post_forward(f_input, f_schedule_plan.vp_stage)

# post_backward()/send_backward_recv_backward() is running in the computation stream,
# so the p2p comm could be overlapped with the wgrad of attn backward
if b_schedule_plan is not None and post_backward is not None:
b_schedule_plan.wait_current_stream()
post_backward(b_grad, b_schedule_plan.vp_stage)

if not fine_grained_overlap:
# Delay the last attn_dw in backward pass (attn_dw of the first layer)
# for overlapping with the p2p comm
if b_num_layers > 0:
b_schedule_plan.get_layer(0).attn.backward_dw()

# post process forward
if f_schedule_plan is not None and f_schedule_plan.post_process is not None:
Expand Down
Loading