Skip to content

[Dev][fix] FSDP EP-overlap CUDA-graph guard uses post-refactor API #4834

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

Merged
hxbai merged 1 commit into from
May 18, 2026
+6 −13
Merged

[Dev][fix] FSDP EP-overlap CUDA-graph guard uses post-refactor API #4834

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
19 changes: 6 additions & 13 deletions megatron/core/distributed/fsdp/mcore_fsdp_adapter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,6 @@
from megatron.core.distributed.data_parallel_base import _BaseDataParallel
from megatron.core.distributed.distributed_data_parallel_config import DistributedDataParallelConfig
from megatron.core.process_groups_config import ProcessGroupCollection
from megatron.core.transformer.enums import CudaGraphScope
from megatron.core.transformer.transformer_config import TransformerConfig
from megatron.core.transformer.transformer_layer import TransformerLayer
from megatron.core.utils import is_te_min_version, log_single_rank
Expand Down Expand Up @@ -163,18 +162,12 @@ def __init__(
"1F1B overlap with FSDP does not support double buffer. "
"Please set fsdp_double_buffer=False in the ddp config."
)
if config.cuda_graph_impl not in ["none", "full_iteration"]:
partial_cuda_graph_scopes = [
scope
for scope in config.cuda_graph_scope
if scope
not in (CudaGraphScope.full_iteration, CudaGraphScope.full_iteration_inference)
]
assert not partial_cuda_graph_scopes, (
"1F1B overlap with FSDP does not support partial CUDA graph scopes "
f"({partial_cuda_graph_scopes}). "
"Please use cuda_graph_scope='full' or disable CUDA graphs."
)
assert config.cuda_graph_impl in ("none", "full_iteration"), (
"1F1B overlap with FSDP does not support per-layer CUDA graphs "
f"(cuda_graph_impl={config.cuda_graph_impl!r}). "
"Use cuda_graph_impl='full_iteration' or disable CUDA graphs "
"(cuda_graph_impl='none')."
)

if (
config.overlap_moe_expert_parallel_comm
Expand Down
Loading