From 089a5edf97cbab4dd8c99917abf1c5cc173004ef Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Sat, 16 May 2026 18:49:04 -0700 Subject: [PATCH] Fix FSDP EP-overlap CUDA-graph guard for the post-refactor API PR #3796 ("Support A2A Overlap for Megatron-FSDP") landed on dev with guard logic that iterates the legacy config.cuda_graph_scope list: if config.cuda_graph_impl not in ["none", "full_iteration"]: partial_cuda_graph_scopes = [ scope for scope in config.cuda_graph_scope ... ] After PR #4293 normalized cuda_graph_scope to None in __post_init__, the inner iteration crashes with TypeError: 'NoneType' object is not iterable whenever a user combines overlap_moe_expert_parallel_comm with cuda_graph_impl in {"local", "transformer_engine"}. Drop the legacy iteration; the outer check on cuda_graph_impl is the only signal needed under the new API. Also drop the now-unused CudaGraphScope import. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../distributed/fsdp/mcore_fsdp_adapter.py | 19 ++++++------------- 1 file changed, 6 insertions(+), 13 deletions(-) diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index e80ad8c268c..a852ef226b4 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py @@ -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 @@ -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