diff --git a/docker/Dockerfile.linting b/docker/Dockerfile.linting index bf27b768374..6c18b0043b3 100644 --- a/docker/Dockerfile.linting +++ b/docker/Dockerfile.linting @@ -20,4 +20,4 @@ FROM main AS jet ARG JET_API_VERSION RUN --mount=type=secret,id=JET_INDEX_URLS \ JET_INDEX_URLS=$(cat /run/secrets/JET_INDEX_URLS) && \ - uv pip install --no-cache-dir "jet-client~=2.0" --upgrade $JET_INDEX_URLS + uv pip install --no-cache-dir "jet-client~=2.0" --upgrade $JET_INDEX_URLS \ No newline at end of file diff --git a/megatron/core/transformer/moe/fused_a2a.py b/megatron/core/transformer/moe/fused_a2a.py index 5e97d7a6786..61901a213f0 100644 --- a/megatron/core/transformer/moe/fused_a2a.py +++ b/megatron/core/transformer/moe/fused_a2a.py @@ -7,16 +7,19 @@ from megatron.core.utils import internal_api +import torch + +from megatron.core.transformer.moe.fused_a2a_config import FusedA2AConfig + try: from deep_ep import Buffer from deep_ep.utils import EventHandle, EventOverlap + from deep_ep_cpp import Config as DeepEPConfig HAVE_DEEP_EP = True except ImportError: HAVE_DEEP_EP = False -import torch - _buffer = None @@ -68,6 +71,56 @@ def get_buffer(group: torch.distributed.ProcessGroup, hidden_bytes: int): return _buffer +if HAVE_DEEP_EP: + + def _build_deepep_config(default_config, fused_a2a_cfg): + """Build a DeepEP Config from a default config and user overrides in FusedA2AConfig. + + Starts from the hardware-appropriate default (e.g. Buffer.get_dispatch_config) and + overrides only the fields the user explicitly set. Returns the default unchanged + when no overrides are needed, avoiding any object allocation on the hot path. + + Args: + default_config: Config returned by Buffer.get_dispatch_config() or + Buffer.get_combine_config() — hardware-tuned baseline. + fused_a2a_cfg: FusedA2AConfig with optional user overrides, or None. + + Returns: + A DeepEP Config (possibly the original default_config if nothing was overridden). + + Note: + chunk_size maps to num_max_nvl_chunked_send_tokens. The DeepEP C++ assertion + requires chunk_size < num_max_nvl_chunked_recv_tokens (default 256), so + values >= 256 will raise inside the DeepEP C++ constructor. + """ + if fused_a2a_cfg is None: + return default_config + if fused_a2a_cfg.num_sms is None and fused_a2a_cfg.chunk_size is None: + return default_config + num_sms = ( + fused_a2a_cfg.num_sms + if fused_a2a_cfg.num_sms is not None + else default_config.num_sms + ) + chunk_size = ( + fused_a2a_cfg.chunk_size + if fused_a2a_cfg.chunk_size is not None + else default_config.num_max_nvl_chunked_send_tokens + ) + # Build a new Config preserving all RDMA / buffer-size params from the default. + # DeepEPConfig == deep_ep_cpp.Config, guaranteed available here (HAVE_DEEP_EP=True). + return DeepEPConfig( + num_sms, + chunk_size, + default_config.num_max_nvl_chunked_recv_tokens, + default_config.num_max_rdma_chunked_send_tokens, + default_config.num_max_rdma_chunked_recv_tokens, + ) + +else: + _build_deepep_config = None + + class FusedDispatch(torch.autograd.Function): """Fused dispatch operation for MoE routing combining computation and communication.""" @@ -81,11 +134,20 @@ def forward( group, async_finish=False, allocate_on_comm_stream=False, + fused_a2a_config=None, ): """Forward pass of fused dispatch.""" previous_event = None if async_finish: previous_event = EventOverlap(EventHandle()) + # Build hardware-appropriate DeepEP Configs, applying user overrides. + # _build_deepep_config reads fields from the C++ Config struct exposed by pybind11. + dispatch_config = _build_deepep_config( + Buffer.get_dispatch_config(group.size()), fused_a2a_config + ) + combine_config = _build_deepep_config( + Buffer.get_combine_config(group.size()), fused_a2a_config + ) # Calculate layout before actual dispatch buffer = get_buffer(group, get_hidden_bytes(x)) ( @@ -120,6 +182,7 @@ def forward( 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, + config=dispatch_config, previous_event=event, # wait in deepep::intra/inter_dispatch async_finish=async_finish, allocate_on_comm_stream=allocate_on_comm_stream, @@ -129,11 +192,12 @@ def forward( if async_finish: after_event_overlap.current_stream_wait() - # Save for backward + # Save for backward (combine uses the combine config) ctx.group = group ctx.handle = handle ctx.async_finish = async_finish ctx.allocate_on_comm_stream = allocate_on_comm_stream + ctx.combine_config = combine_config tokens_per_expert = torch.tensor(num_recv_tokens_per_expert_list) return (recv_x, recv_token_indices, recv_token_probs, tokens_per_expert, handle) @@ -152,6 +216,7 @@ def backward( grad_output.contiguous(), handle, topk_weights=grad_token_probs.float(), + config=ctx.combine_config, previous_event=previous_event, async_finish=ctx.async_finish, allocate_on_comm_stream=ctx.allocate_on_comm_stream, @@ -159,22 +224,33 @@ def backward( # Make sure current stream is synchronized if ctx.async_finish: after_event.current_stream_wait() - return grad_x, None, grad_token_probs, None, None, None, None + return grad_x, None, grad_token_probs, None, None, None, None, None class FusedCombine(torch.autograd.Function): """Fused combine operation for MoE output combining computation and communication.""" @staticmethod - def forward(ctx, x, group, handle, async_finish=False, allocate_on_comm_stream=False): + def forward( + ctx, x, group, handle, async_finish=False, allocate_on_comm_stream=False, + fused_a2a_config=None, + ): """Forward pass of fused combine.""" previous_event = None if async_finish: previous_event = EventOverlap(EventHandle()) + # Build configs; combine uses combine_config, backward dispatch uses dispatch_config. + combine_config = _build_deepep_config( + Buffer.get_combine_config(group.size()), fused_a2a_config + ) + dispatch_config = _build_deepep_config( + Buffer.get_dispatch_config(group.size()), fused_a2a_config + ) buffer = get_buffer(group, get_hidden_bytes(x)) combined_x, _, after_event = buffer.combine( x, handle=handle, + config=combine_config, async_finish=async_finish, previous_event=previous_event, allocate_on_comm_stream=allocate_on_comm_stream, @@ -187,6 +263,7 @@ def forward(ctx, x, group, handle, async_finish=False, allocate_on_comm_stream=F ctx.group = group ctx.async_finish = async_finish ctx.allocate_on_comm_stream = allocate_on_comm_stream + ctx.dispatch_config = dispatch_config return combined_x, None @staticmethod @@ -199,6 +276,7 @@ def backward(ctx, grad_output, previous_event=None): grad_x, _, _, _, _, after_event = buffer.dispatch( grad_output.contiguous(), handle=ctx.handle, + config=ctx.dispatch_config, previous_event=previous_event, async_finish=ctx.async_finish, allocate_on_comm_stream=ctx.allocate_on_comm_stream, @@ -206,7 +284,7 @@ def backward(ctx, grad_output, previous_event=None): # Make sure current stream is synchronized if ctx.async_finish: after_event.current_stream_wait() - return grad_x, None, None, None, None + return grad_x, None, None, None, None, None if HAVE_DEEP_EP: @@ -219,6 +297,7 @@ def fused_dispatch( group, async_finish=False, allocate_on_comm_stream=False, + config=None, ): """Perform fused dispatch operation if deep_ep is available. @@ -241,9 +320,11 @@ def fused_dispatch( group, async_finish, allocate_on_comm_stream, + config, ) - def fused_combine(x, group, handle, async_finish=False, allocate_on_comm_stream=False): + def fused_combine(x, group, handle, async_finish=False, allocate_on_comm_stream=False, + config=None): """Perform fused combine operation if deep_ep is available. Args: @@ -255,7 +336,7 @@ def fused_combine(x, group, handle, async_finish=False, allocate_on_comm_stream= Returns: Result of FusedCombine """ - return FusedCombine.apply(x, group, handle, async_finish, allocate_on_comm_stream) + return FusedCombine.apply(x, group, handle, async_finish, allocate_on_comm_stream, config) def set_deepep_num_sms(num_sms): """Sets the number of SMs to use for DeepEP""" diff --git a/megatron/core/transformer/moe/fused_a2a_config.py b/megatron/core/transformer/moe/fused_a2a_config.py new file mode 100644 index 00000000000..e930c3f6ff2 --- /dev/null +++ b/megatron/core/transformer/moe/fused_a2a_config.py @@ -0,0 +1,53 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +""" +FusedA2AConfig dataclass for user-tunable fused all-to-all MoE parameters. +""" +from dataclasses import dataclass, field +from typing import Optional + + +@dataclass +class FusedA2AConfig: + """ + Configuration for fused all-to-all MoE (FusedDispatch/FusedCombine). + + Fields: + chunk_size: Optional[int] - Chunk size for all-to-all (must be positive if set) + num_sms: Optional[int] - Number of SMs for kernel (must be positive and even if set) + # Future tunables can be added here + + Precedence: CLI > ENV > CONFIG FILE > DEFAULTS + """ + chunk_size: Optional[int] = None + num_sms: Optional[int] = None + # Future tunables can be added here + + def validate(self): + if self.chunk_size is not None: + if not (self.chunk_size > 0): + raise ValueError(f"chunk_size must be positive, got {self.chunk_size}") + if self.num_sms is not None: + if not (self.num_sms > 0): + raise ValueError(f"num_sms must be positive, got {self.num_sms}") + # DeepEP's Buffer.set_num_sms asserts new_num_sms % 2 == 0 + # and the C++ kernel asserts config.num_sms % 2 == 0. An odd value + # would crash deep inside the kernel with an opaque assertion; + # fail fast at validate_args time instead. + if self.num_sms % 2 != 0: + raise ValueError( + f"num_sms must be even (DeepEP requirement), got {self.num_sms}" + ) + + @staticmethod + def from_dict(cfg: dict) -> 'FusedA2AConfig': + allowed = {'chunk_size', 'num_sms'} + unknown = set(cfg.keys()) - allowed + if unknown: + raise ValueError(f"Unknown FusedA2AConfig keys: {unknown}") + return FusedA2AConfig( + chunk_size=cfg.get('chunk_size'), + num_sms=cfg.get('num_sms'), + ) + + def __repr__(self): + return f"FusedA2AConfig(chunk_size={self.chunk_size}, num_sms={self.num_sms})" diff --git a/megatron/core/transformer/moe/fused_a2a_config_loader.py b/megatron/core/transformer/moe/fused_a2a_config_loader.py new file mode 100644 index 00000000000..da1a9128e4a --- /dev/null +++ b/megatron/core/transformer/moe/fused_a2a_config_loader.py @@ -0,0 +1,57 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +""" +FusedA2A config loader and resolver for Megatron-LM MoE fused all-to-all. +""" +import os +import json +from typing import Optional + +try: + import yaml + HAVE_YAML = True +except ImportError: + HAVE_YAML = False + +from .fused_a2a_config import FusedA2AConfig + + +def load_a2a_config_from_file(path: str) -> dict: + """ + Load config from JSON or YAML file. + Raises ValueError on error. + """ + if path.endswith('.json'): + with open(path, 'r') as f: + return json.load(f) + elif path.endswith(('.yaml', '.yml')): + if not HAVE_YAML: + raise ImportError('pyyaml is required for YAML config files') + with open(path, 'r') as f: + return yaml.safe_load(f) + else: + raise ValueError(f"Unsupported config file extension: {path}") + +def resolve_fused_a2a_config_from_sources(cli_args=None, env=os.environ, config_file_path=None) -> FusedA2AConfig: + """ + Resolve FusedA2AConfig from CLI args, environment, and config file. + Precedence: CLI > ENV > CONFIG FILE > DEFAULTS. + Raises ValueError on any invalid or unknown keys. + """ + file_cfg = {} + if config_file_path: + file_cfg = load_a2a_config_from_file(config_file_path) + env_cfg = {} + if env.get('MOE_A2A_CHUNK_SIZE'): + env_cfg['chunk_size'] = int(env['MOE_A2A_CHUNK_SIZE']) + if env.get('MOE_A2A_NUM_SMS'): + env_cfg['num_sms'] = int(env['MOE_A2A_NUM_SMS']) + cli_cfg = {} + if cli_args is not None: + if getattr(cli_args, 'moe_a2a_chunk_size', None) is not None: + cli_cfg['chunk_size'] = cli_args.moe_a2a_chunk_size + if getattr(cli_args, 'moe_a2a_num_sms', None) is not None: + cli_cfg['num_sms'] = cli_args.moe_a2a_num_sms + merged = {**file_cfg, **env_cfg, **cli_cfg} + config = FusedA2AConfig.from_dict(merged) + config.validate() + return config diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index d7a1772b61e..d00af65a025 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -66,6 +66,9 @@ def __init__( pg_collection (ProcessGroupCollection, optional): Process groups for MoE operations. """ self.config = config + # FusedA2AConfig resolved once at startup (validate_args) and stored on TransformerConfig. + # Single source of truth: CLI > ENV > config file > defaults. Never re-resolved here. + self.fused_a2a_config = getattr(config, 'fused_a2a_config', None) self.shared_experts: Optional[SharedExpertMLP] = None # Whether to use NCCL stream for A2A communication, otherwise default stream is used. self.use_nccl_stream = False # Will be set to True when shared_experts is set. @@ -1213,7 +1216,16 @@ def __init__( "DeepEP is not installed. Please install DeepEP package from " "https://github.com/deepseek-ai/deepep." ) - set_deepep_num_sms(config.moe_deepep_num_sms) + # Determine effective SM count. fused_a2a_config.num_sms (set via --moe-a2a-num-sms / + # MOE_A2A_NUM_SMS / config file) takes precedence over moe_deepep_num_sms so there is a + # single, authoritative source of truth after validate_args has resolved everything. + _fused_a2a_config = getattr(config, 'fused_a2a_config', None) + effective_num_sms = ( + _fused_a2a_config.num_sms + if _fused_a2a_config is not None and _fused_a2a_config.num_sms is not None + else config.moe_deepep_num_sms + ) + set_deepep_num_sms(effective_num_sms) def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor): num_tokens = routing_map.shape[0] @@ -1240,6 +1252,8 @@ def dispatch( "DeepEP only supports float32 probs, please set --moe-router-dtype=fp32" ) self.token_probs = self.token_probs.float() # downcast or upcast + # NOTE: FusedA2AConfig must be passed explicitly from the top-level (single source of truth). + # Do not resolve or mutate config here. No Python-side logic in hot path. hidden_states, dispatched_indices, dispatched_probs, num_tokens_per_expert, handle = ( fused_dispatch( hidden_states, @@ -1249,6 +1263,7 @@ def dispatch( self.group, async_finish=async_finish, allocate_on_comm_stream=allocate_on_comm_stream, + config=self.fused_a2a_config, ) ) self.handle = handle @@ -1300,12 +1315,15 @@ def combine( async_finish: bool = False, allocate_on_comm_stream: bool = False, ) -> torch.Tensor: + # NOTE: FusedA2AConfig must be passed explicitly from the top-level (single source of truth). + # Do not resolve or mutate config here. No Python-side logic in hot path. hidden_states, _ = fused_combine( hidden_states, self.group, self.handle, async_finish=async_finish, allocate_on_comm_stream=allocate_on_comm_stream, + config=self.fused_a2a_config, ) # Release the handle after combine operation self.handle = None diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index ab84abd3a17..b3ed374dfbd 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -27,6 +27,9 @@ ) from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout +# FusedA2AConfig is a lightweight dataclass with no heavy dependencies — safe to import directly. +from megatron.core.transformer.moe.fused_a2a_config import FusedA2AConfig + from .._rank_utils import log_single_rank from ..fusions.fused_bias_geglu import quick_gelu from ..model_parallel_config import ModelParallelConfig @@ -885,6 +888,14 @@ class TransformerConfig(ModelParallelConfig): exceeding this budget will be dropped. None means no token will be dropped. The default is None.""" + fused_a2a_config: Optional[FusedA2AConfig] = None + """User-tunable configuration for fused all-to-all MoE communication (DeepEP backend). + Resolved once at training startup by validate_args from CLI flags, environment variables + (MOE_A2A_CHUNK_SIZE, MOE_A2A_NUM_SMS), and an optional JSON/YAML config file. + Precedence: CLI > ENV > config file > built-in defaults. + None means all tunables fall back to their built-in defaults. + See --moe-a2a-chunk-size, --moe-a2a-num-sms, --moe-a2a-config-file.""" + ################## # Context Parallel ################## diff --git a/megatron/training/argument_utils.py b/megatron/training/argument_utils.py index 2cfb3f0f17b..1e47d98de48 100644 --- a/megatron/training/argument_utils.py +++ b/megatron/training/argument_utils.py @@ -295,7 +295,19 @@ def core_transformer_config_from_args(args, config_class=None): kw_args = {} for f in dataclasses.fields(config_class): if hasattr(args, f.name): - kw_args[f.name] = getattr(args, f.name) + value = getattr(args, f.name) + # Skip explicit None values for fields whose dataclass default is non-None. + # Without this, adding a CLI flag with default=None would shadow the field + # default and propagate None into the config. Example: --moe-deepep-num-sms + # is a new optional CLI flag; when the user does not set it, the field must + # keep its historical default (20), not become None. + if value is None: + has_non_none_default = ( + f.default is not dataclasses.MISSING and f.default is not None + ) or f.default_factory is not dataclasses.MISSING + if has_non_none_default: + continue + kw_args[f.name] = value kw_args['persist_layer_norm'] = not args.no_persist_layer_norm kw_args['deallocate_pipeline_outputs'] = True kw_args['pipeline_dtype'] = args.params_dtype diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index d02151066d9..1fcf58cb0f9 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -10,6 +10,7 @@ import types import torch +from packaging.version import Version as PkgVersion from megatron.core.rerun_state_machine import RerunStateMachine from megatron.core.transformer import TransformerConfig @@ -22,6 +23,7 @@ validate_deprecated_cuda_graph_modules_migration_inputs, ) from megatron.core.transformer.enums import AttnBackend, CudaGraphModule, InferenceCudaGraphScope +from megatron.core.transformer.heterogeneous.heterogeneous_config import MLPConfig from megatron.core.utils import ( get_torch_version, is_flashinfer_min_version, @@ -37,6 +39,15 @@ ) from megatron.core.msc_utils import MultiStorageClientFeature +# ArgumentGroupFactory is consumed in this module; core_transformer_config_from_args +# is re-exported for backwards compatibility (external callers do +# `from megatron.training.arguments import core_transformer_config_from_args`). +# The canonical implementation now lives in argument_utils.py. +from megatron.training.argument_utils import ( # noqa: F401 pylint: disable=unused-import + ArgumentGroupFactory, + core_transformer_config_from_args, +) + from megatron.training.argument_utils import ArgumentGroupFactory, core_transformer_config_from_args # noqa: F401 # pylint: disable=unused-import @@ -1775,6 +1786,37 @@ def validate_args(args, defaults={}): assert args.moe_latent_size > 0, "MoE latent projection dimension has to be greater than zero." assert args.num_experts is not None, "MoE latent projections are applicable only for MoE models." + # ------------------------------------------------------------------ + # Resolve FusedA2AConfig: merges CLI > ENV > config file > defaults. + # Done once here so every downstream consumer gets a single, validated + # object instead of re-reading env/files at model-construction time. + # Result is stored on args.fused_a2a_config and automatically picked up + # by core_transformer_config_from_args → TransformerConfig.fused_a2a_config. + # ------------------------------------------------------------------ + from megatron.core.transformer.moe.fused_a2a_config_loader import ( + resolve_fused_a2a_config_from_sources, + ) + args.fused_a2a_config = resolve_fused_a2a_config_from_sources( + cli_args=args, + config_file_path=getattr(args, 'moe_a2a_config_file', None), + ) + # If --moe-a2a-num-sms is not explicitly set, propagate --moe-deepep-num-sms into the + # resolved config so _DeepepManager has a single authoritative source. + if ( + args.fused_a2a_config.num_sms is None + and getattr(args, 'moe_deepep_num_sms', None) is not None + ): + args.fused_a2a_config = type(args.fused_a2a_config)( + chunk_size=args.fused_a2a_config.chunk_size, + num_sms=args.moe_deepep_num_sms, + ) + # Re-validate after the propagation so the even-SM check also covers + # values supplied via --moe-deepep-num-sms (which the resolver above + # does not see). + args.fused_a2a_config.validate() + if args.fused_a2a_config.chunk_size is not None or args.fused_a2a_config.num_sms is not None: + print_rank_0(f'[MoE A2A] Effective FusedA2AConfig: {args.fused_a2a_config}') + # Print arguments. _print_args("arguments", args) @@ -2111,6 +2153,14 @@ def _add_network_size_args(parser): "persist_layer_norm", "bias_dropout_fusion", "apply_rope_fusion", + # registered manually in _add_moe_args to keep the None-sentinel + # for the FusedA2A config resolution flow; auto-generation would + # shadow the dataclass default and break backward compatibility. + "moe_deepep_num_sms", + # not a CLI arg: populated programmatically by validate_args via + # resolve_fused_a2a_config_from_sources. Auto-generating this as a + # CLI flag would expose a dataclass type that argparse cannot parse. + "fused_a2a_config", ] transformer_factory = ArgumentGroupFactory(TransformerConfig, exclude=exclude) transformer_group = transformer_factory.build_group(parser, "transformer configuration") @@ -3175,6 +3225,33 @@ def _add_moe_args(parser): group.add_argument('--moe-upcycling-granularity', type=int, default=1, help='This param sepecifics how many times smaller is the expert hidden size compared with the original dense FFN hidden size. ' 'For using granular upcycling strategy, please set this param as a positive integer. If this param is set to 1, it means using the default upcycling strategy.') + + # ------------------------------------------------------------------ + # DeepEP / Fused All-to-All tuning arguments + # Resolved once at startup in validate_args via resolve_fused_a2a_config_from_sources. + # Precedence: CLI > ENV (MOE_A2A_CHUNK_SIZE / MOE_A2A_NUM_SMS) > CONFIG FILE > DEFAULTS. + # ------------------------------------------------------------------ + group.add_argument( + '--moe-deepep-num-sms', type=int, default=None, + help='Number of SMs to use for DeepEP all-to-all kernels. ' + 'Overrides the TransformerConfig default (20). ' + 'Superseded by --moe-a2a-num-sms when both are set.') + group.add_argument( + '--moe-a2a-chunk-size', type=int, default=None, + help='Chunk size for fused all-to-all MoE communication (DeepEP). ' + 'When not set, the kernel uses its built-in default. ' + 'Precedence: CLI > ENV(MOE_A2A_CHUNK_SIZE) > config file > default.') + group.add_argument( + '--moe-a2a-num-sms', type=int, default=None, + help='Number of SMs for fused all-to-all MoE kernels (DeepEP). ' + 'Takes priority over --moe-deepep-num-sms when both are set. ' + 'Precedence: CLI > ENV(MOE_A2A_NUM_SMS) > config file > --moe-deepep-num-sms > 20.') + group.add_argument( + '--moe-a2a-config-file', type=str, default=None, + help='Path to a JSON or YAML file with fused all-to-all MoE tuning parameters ' + '(chunk_size, num_sms). CLI flags and environment variables take precedence ' + 'over values in this file.') + return parser def _add_mla_args(parser): diff --git a/tests/unit_tests/moe/test_fused_a2a_config.py b/tests/unit_tests/moe/test_fused_a2a_config.py new file mode 100644 index 00000000000..0cf38f32422 --- /dev/null +++ b/tests/unit_tests/moe/test_fused_a2a_config.py @@ -0,0 +1,71 @@ +import unittest +import types +import os +from megatron.core.transformer.moe.fused_a2a_config import FusedA2AConfig +from megatron.core.transformer.moe.fused_a2a_config_loader import resolve_fused_a2a_config_from_sources + +class DummyArgs: + def __init__(self, **kwargs): + for k, v in kwargs.items(): + setattr(self, k, v) + +class TestFusedA2AConfigResolution(unittest.TestCase): + def setUp(self): + # Clean env + for k in ["MOE_A2A_CHUNK_SIZE", "MOE_A2A_NUM_SMS"]: + if k in os.environ: + del os.environ[k] + + def test_cli_priority(self): + # num_sms must be even (DeepEP requirement). + args = DummyArgs(moe_a2a_chunk_size=42, moe_a2a_num_sms=8, moe_a2a_config_file=None) + cfg = resolve_fused_a2a_config_from_sources(cli_args=args) + self.assertEqual(cfg.chunk_size, 42) + self.assertEqual(cfg.num_sms, 8) + + def test_env_priority(self): + os.environ["MOE_A2A_CHUNK_SIZE"] = "99" + args = DummyArgs(moe_a2a_chunk_size=None, moe_a2a_num_sms=None, moe_a2a_config_file=None) + cfg = resolve_fused_a2a_config_from_sources(cli_args=args) + self.assertEqual(cfg.chunk_size, 99) + + def test_file_priority(self): + import tempfile, json + with tempfile.NamedTemporaryFile(mode="w+", suffix=".json", delete=False) as f: + json.dump({"chunk_size": 123, "num_sms": 8}, f) + f.flush() + args = DummyArgs(moe_a2a_chunk_size=None, moe_a2a_num_sms=None, moe_a2a_config_file=f.name) + cfg = resolve_fused_a2a_config_from_sources(cli_args=args, config_file_path=f.name) + self.assertEqual(cfg.chunk_size, 123) + self.assertEqual(cfg.num_sms, 8) + os.unlink(f.name) + + def test_merge_priority(self): + import tempfile, json + os.environ["MOE_A2A_CHUNK_SIZE"] = "77" + with tempfile.NamedTemporaryFile(mode="w+", suffix=".json", delete=False) as f: + json.dump({"chunk_size": 555, "num_sms": 2}, f) + f.flush() + args = DummyArgs(moe_a2a_chunk_size=88, moe_a2a_num_sms=None, moe_a2a_config_file=f.name) + cfg = resolve_fused_a2a_config_from_sources(cli_args=args, config_file_path=f.name) + self.assertEqual(cfg.chunk_size, 88) # CLI wins + self.assertEqual(cfg.num_sms, 2) # file wins + os.unlink(f.name) + + def test_validation(self): + args = DummyArgs(moe_a2a_chunk_size=-1, moe_a2a_num_sms=0, moe_a2a_config_file=None) + with self.assertRaises(ValueError): + resolve_fused_a2a_config_from_sources(cli_args=args) + + def test_unknown_key(self): + import tempfile, json + with tempfile.NamedTemporaryFile(mode="w+", suffix=".json", delete=False) as f: + json.dump({"chunk_size": 123, "num_sms": 8, "bad_field": 1}, f) + f.flush() + args = DummyArgs(moe_a2a_chunk_size=None, moe_a2a_num_sms=None, moe_a2a_config_file=f.name) + with self.assertRaises(ValueError): + resolve_fused_a2a_config_from_sources(cli_args=args, config_file_path=f.name) + os.unlink(f.name) + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit_tests/moe/test_fused_a2a_config_validation.py b/tests/unit_tests/moe/test_fused_a2a_config_validation.py new file mode 100644 index 00000000000..acdae4e615a --- /dev/null +++ b/tests/unit_tests/moe/test_fused_a2a_config_validation.py @@ -0,0 +1,47 @@ +import unittest +from megatron.core.transformer.moe.fused_a2a_config import FusedA2AConfig + +class TestFusedA2AConfigValidation(unittest.TestCase): + def test_valid(self): + cfg = FusedA2AConfig(chunk_size=32, num_sms=8) + cfg.validate() # Should not raise + + def test_valid_none(self): + # All fields None is a valid configuration (use hardware defaults). + cfg = FusedA2AConfig() + cfg.validate() + + def test_valid_even_num_sms(self): + for n in (2, 4, 6, 8, 16, 20, 24, 64): + cfg = FusedA2AConfig(chunk_size=32, num_sms=n) + cfg.validate() # Should not raise + + def test_invalid_chunk_size(self): + cfg = FusedA2AConfig(chunk_size=0, num_sms=8) + with self.assertRaises(ValueError): + cfg.validate() + + def test_invalid_num_sms(self): + cfg = FusedA2AConfig(chunk_size=32, num_sms=-1) + with self.assertRaises(ValueError): + cfg.validate() + + def test_odd_num_sms_rejected(self): + # DeepEP requires num_sms to be even. An odd value would crash + # deep inside the kernel with an opaque assertion; the validator + # must fail fast. + for n in (1, 3, 7, 21, 33): + cfg = FusedA2AConfig(chunk_size=32, num_sms=n) + with self.assertRaises(ValueError) as ctx: + cfg.validate() + self.assertIn("even", str(ctx.exception)) + self.assertIn(str(n), str(ctx.exception)) + + def test_odd_num_sms_alone_rejected(self): + # Even chunk_size doesn't excuse an odd num_sms. + cfg = FusedA2AConfig(chunk_size=64, num_sms=11) + with self.assertRaises(ValueError): + cfg.validate() + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit_tests/moe/test_fused_a2a_integration.py b/tests/unit_tests/moe/test_fused_a2a_integration.py new file mode 100644 index 00000000000..a1b7675da2b --- /dev/null +++ b/tests/unit_tests/moe/test_fused_a2a_integration.py @@ -0,0 +1,393 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +""" +Integration tests: verify the full FusedA2AConfig pipeline from CLI argument resolution +through TransformerConfig to MoETokenDispatcher. + +These tests exercise the complete data-flow that runs in production: + validate_args (resolve_fused_a2a_config_from_sources) + → args.fused_a2a_config + → core_transformer_config_from_args (dataclass field copy) + → TransformerConfig.fused_a2a_config + → MoETokenDispatcher.fused_a2a_config + → fused_dispatch / fused_combine (config= kwarg) + +No GPU / distributed runtime is required. DeepEP is not imported. +""" + +import dataclasses +import json +import os +import tempfile +import unittest + +from megatron.core.transformer.moe.fused_a2a_config import FusedA2AConfig +from megatron.core.transformer.moe.fused_a2a_config_loader import ( + resolve_fused_a2a_config_from_sources, +) +from megatron.core.transformer.transformer_config import TransformerConfig + + +# --------------------------------------------------------------------------- +# Minimal stand-in for a parsed args namespace +# --------------------------------------------------------------------------- +class _Args: + """Lightweight namespace that mimics the argparse output of parse_args().""" + + def __init__(self, **kwargs): + for k, v in kwargs.items(): + setattr(self, k, v) + + +def _make_args(**overrides): + """Return an _Args with all A2A-related attributes set to benign defaults.""" + defaults = dict( + moe_a2a_chunk_size=None, + moe_a2a_num_sms=None, + moe_a2a_config_file=None, + moe_deepep_num_sms=None, + ) + defaults.update(overrides) + return _Args(**defaults) + + +# --------------------------------------------------------------------------- +# Helper: simulate the validate_args resolution block +# --------------------------------------------------------------------------- +def _resolve(args, clean_env=None): + """Mirror the resolution logic in validate_args, for testing without the full stack.""" + env = dict(os.environ) + if clean_env is not None: + for k in ("MOE_A2A_CHUNK_SIZE", "MOE_A2A_NUM_SMS"): + env.pop(k, None) + env.update(clean_env) + + cfg = resolve_fused_a2a_config_from_sources( + cli_args=args, + env=env, + config_file_path=getattr(args, 'moe_a2a_config_file', None), + ) + # Propagate --moe-deepep-num-sms when --moe-a2a-num-sms is absent (mirrors validate_args) + if cfg.num_sms is None and getattr(args, 'moe_deepep_num_sms', None) is not None: + cfg = type(cfg)(chunk_size=cfg.chunk_size, num_sms=args.moe_deepep_num_sms) + # Re-validate after propagation to fail fast on odd values + # supplied via --moe-deepep-num-sms (mirrors validate_args). + cfg.validate() + return cfg + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- +class TestFusedA2AConfigFieldOnTransformerConfig(unittest.TestCase): + """TransformerConfig must carry a fused_a2a_config field.""" + + def test_field_exists_and_defaults_to_none(self): + names = {f.name for f in dataclasses.fields(TransformerConfig)} + self.assertIn( + 'fused_a2a_config', names, + "TransformerConfig must have a fused_a2a_config dataclass field.", + ) + # Inspect the default + field_obj = next(f for f in dataclasses.fields(TransformerConfig) if f.name == 'fused_a2a_config') + self.assertIsNone( + field_obj.default, + "fused_a2a_config field default must be None for backward compatibility.", + ) + + def test_field_accepts_fused_a2a_config_object(self): + cfg = FusedA2AConfig(chunk_size=64, num_sms=8) + # TransformerConfig cannot be constructed without mandatory fields; use a dummy subclass + # that skips __post_init__ to test field presence only. + tc_fields = {f.name: f.default for f in dataclasses.fields(TransformerConfig) + if f.default is not dataclasses.MISSING} + tc_fields.update(dict( + num_layers=2, hidden_size=64, num_attention_heads=2, + fused_a2a_config=cfg, + )) + # We don't instantiate TransformerConfig here (it has complex __post_init__ with TE checks) + # but we verify that the field is declared Optional[FusedA2AConfig] by checking its type. + field_obj = next(f for f in dataclasses.fields(TransformerConfig) if f.name == 'fused_a2a_config') + # The type hint should reference FusedA2AConfig + hint = str(field_obj.type) + self.assertIn('FusedA2AConfig', hint) + + +class TestValidateArgsResolutionLogic(unittest.TestCase): + """Verify the resolution logic that validate_args will execute.""" + + def setUp(self): + for k in ("MOE_A2A_CHUNK_SIZE", "MOE_A2A_NUM_SMS"): + os.environ.pop(k, None) + + def tearDown(self): + for k in ("MOE_A2A_CHUNK_SIZE", "MOE_A2A_NUM_SMS"): + os.environ.pop(k, None) + + # ---- Precedence: CLI wins over everything -------------------------------- + def test_cli_chunk_size_wins_over_env(self): + args = _make_args(moe_a2a_chunk_size=128) + cfg = _resolve(args, clean_env={"MOE_A2A_CHUNK_SIZE": "999"}) + self.assertEqual(cfg.chunk_size, 128) + + def test_cli_num_sms_wins_over_env(self): + args = _make_args(moe_a2a_num_sms=4) + cfg = _resolve(args, clean_env={"MOE_A2A_NUM_SMS": "999"}) + self.assertEqual(cfg.num_sms, 4) + + def test_cli_wins_over_config_file(self): + with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f: + json.dump({"chunk_size": 512, "num_sms": 32}, f) + fname = f.name + try: + args = _make_args(moe_a2a_chunk_size=64, moe_a2a_num_sms=8, moe_a2a_config_file=fname) + cfg = _resolve(args, clean_env={}) + self.assertEqual(cfg.chunk_size, 64) + self.assertEqual(cfg.num_sms, 8) + finally: + os.unlink(fname) + + # ---- Precedence: ENV wins over file ------------------------------------- + def test_env_wins_over_config_file(self): + with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f: + json.dump({"chunk_size": 512}, f) + fname = f.name + try: + args = _make_args(moe_a2a_config_file=fname) + cfg = _resolve(args, clean_env={"MOE_A2A_CHUNK_SIZE": "256"}) + self.assertEqual(cfg.chunk_size, 256) # env wins + finally: + os.unlink(fname) + + # ---- --moe-deepep-num-sms fallback -------------------------------------- + def test_moe_deepep_num_sms_propagates_when_a2a_num_sms_absent(self): + args = _make_args(moe_deepep_num_sms=12) + cfg = _resolve(args, clean_env={}) + self.assertEqual(cfg.num_sms, 12) + + def test_moe_a2a_num_sms_takes_priority_over_moe_deepep_num_sms(self): + args = _make_args(moe_a2a_num_sms=6, moe_deepep_num_sms=30) + cfg = _resolve(args, clean_env={}) + self.assertEqual(cfg.num_sms, 6) + + def test_odd_moe_deepep_num_sms_propagation_raises(self): + # moe_deepep_num_sms is propagated into fused_a2a_config.num_sms + # when --moe-a2a-num-sms is not set; the post-propagation validate() + # call must reject odd values to fail fast at validate_args time. + args = _make_args(moe_deepep_num_sms=21) + with self.assertRaises(ValueError) as ctx: + _resolve(args, clean_env={}) + self.assertIn("even", str(ctx.exception)) + + # ---- Defaults (all None) ------------------------------------------------ + def test_all_defaults_none_when_nothing_set(self): + args = _make_args() + cfg = _resolve(args, clean_env={}) + self.assertIsNone(cfg.chunk_size) + self.assertIsNone(cfg.num_sms) + + # ---- Fail-fast validation ----------------------------------------------- + def test_invalid_chunk_size_raises(self): + args = _make_args(moe_a2a_chunk_size=0) + with self.assertRaises(ValueError): + _resolve(args, clean_env={}) + + def test_invalid_num_sms_raises(self): + args = _make_args(moe_a2a_num_sms=-5) + with self.assertRaises(ValueError): + _resolve(args, clean_env={}) + + def test_unknown_key_in_config_file_raises(self): + with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f: + json.dump({"chunk_size": 64, "bad_field": 1}, f) + fname = f.name + try: + args = _make_args(moe_a2a_config_file=fname) + with self.assertRaises(ValueError): + _resolve(args, clean_env={}) + finally: + os.unlink(fname) + + # ---- YAML support ------------------------------------------------------- + def test_yaml_config_file_loads_correctly(self): + try: + import yaml + except ImportError: + self.skipTest("pyyaml not installed") + with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f: + yaml.dump({"chunk_size": 32, "num_sms": 4}, f) + fname = f.name + try: + args = _make_args(moe_a2a_config_file=fname) + cfg = _resolve(args, clean_env={}) + self.assertEqual(cfg.chunk_size, 32) + self.assertEqual(cfg.num_sms, 4) + finally: + os.unlink(fname) + + +class TestMoETokenDispatcherPicksUpConfig(unittest.TestCase): + """ + MoETokenDispatcher.__init__ must read fused_a2a_config from TransformerConfig. + We verify this without instantiating the full dispatcher (which requires distributed init) + by directly inspecting the attribute-reading logic. + """ + + def test_dispatcher_reads_fused_a2a_config_from_transformer_config(self): + """ + Simulate: TransformerConfig.fused_a2a_config is set → dispatcher reads it via getattr. + This mirrors the exact line in MoETokenDispatcher.__init__: + self.fused_a2a_config = getattr(config, 'fused_a2a_config', None) + """ + expected = FusedA2AConfig(chunk_size=256, num_sms=16) + + class _MockTransformerConfig: + fused_a2a_config = expected + + result = getattr(_MockTransformerConfig(), 'fused_a2a_config', None) + self.assertIs(result, expected) + + def test_dispatcher_falls_back_to_none_when_field_absent(self): + """Backward compatibility: configs without fused_a2a_config must yield None.""" + + class _LegacyConfig: + pass # no fused_a2a_config attribute + + result = getattr(_LegacyConfig(), 'fused_a2a_config', None) + self.assertIsNone(result) + + +class TestDeepepManagerNumSmsPrecedence(unittest.TestCase): + """ + _DeepepManager must prefer fused_a2a_config.num_sms over moe_deepep_num_sms. + Tests simulate the resolution logic without DeepEP installed. + """ + + def _effective_num_sms(self, fused_a2a_config, moe_deepep_num_sms): + """Mirror the SM-count resolution in _DeepepManager.__init__.""" + return ( + fused_a2a_config.num_sms + if fused_a2a_config is not None and fused_a2a_config.num_sms is not None + else moe_deepep_num_sms + ) + + def test_a2a_config_num_sms_wins(self): + cfg = FusedA2AConfig(num_sms=6) + self.assertEqual(self._effective_num_sms(cfg, moe_deepep_num_sms=20), 6) + + def test_falls_back_to_moe_deepep_num_sms_when_a2a_num_sms_is_none(self): + cfg = FusedA2AConfig(num_sms=None) + self.assertEqual(self._effective_num_sms(cfg, moe_deepep_num_sms=20), 20) + + def test_falls_back_when_fused_a2a_config_is_none(self): + self.assertEqual(self._effective_num_sms(None, moe_deepep_num_sms=20), 20) + + +class TestMoeDeepepNumSmsDefaultRegression(unittest.TestCase): + """ + Regression test for the --moe-deepep-num-sms CLI flag regression. + + Before the fix: adding the new --moe-deepep-num-sms flag with default=None caused + core_transformer_config_from_args to copy None into TransformerConfig.moe_deepep_num_sms, + shadowing the dataclass default of 20. This caused _DeepepManager to call + Buffer.set_num_sms(None) which crashes with TypeError. + + After the fix: when args.moe_deepep_num_sms is None (user did not pass the flag), + the field falls back to its dataclass default of 20. When the user does pass the + flag, the value is honored. + """ + + def setUp(self): + for k in ("MOE_A2A_CHUNK_SIZE", "MOE_A2A_NUM_SMS"): + os.environ.pop(k, None) + + def tearDown(self): + for k in ("MOE_A2A_CHUNK_SIZE", "MOE_A2A_NUM_SMS"): + os.environ.pop(k, None) + + def _simulate_core_transformer_config_from_args(self, args): + """Mirror the kw_args build in core_transformer_config_from_args. + + Includes the regression fix: skip explicit None values for fields whose + dataclass default is non-None. + """ + import dataclasses as _dc + from megatron.core.transformer.transformer_config import TransformerConfig + kw_args = {} + for f in _dc.fields(TransformerConfig): + if hasattr(args, f.name): + value = getattr(args, f.name) + # Fix: skip explicit None when dataclass default is non-None. + if value is None: + has_non_none_default = ( + f.default is not _dc.MISSING and f.default is not None + ) or f.default_factory is not _dc.MISSING + if has_non_none_default: + continue + kw_args[f.name] = value + return kw_args + + def test_default_path_preserves_moe_deepep_num_sms_20(self): + """When --moe-deepep-num-sms is not passed, the field must be 20 (not None). + + This is the historical default preserved across the new CLI flag introduction. + """ + args = _make_args() # moe_deepep_num_sms=None (CLI default) + kw_args = self._simulate_core_transformer_config_from_args(args) + # The moe_deepep_num_sms key should be ABSENT (so the dataclass default of 20 is used) + self.assertNotIn( + 'moe_deepep_num_sms', kw_args, + 'moe_deepep_num_sms must not be propagated to TransformerConfig when the ' + 'user did not pass --moe-deepep-num-sms; this would shadow the default of 20.' + ) + + def test_moe_deepep_num_sms_field_default_is_20(self): + """The dataclass field default must be exactly 20 for backward compatibility.""" + import dataclasses as _dc + from megatron.core.transformer.transformer_config import TransformerConfig + field = next(f for f in _dc.fields(TransformerConfig) if f.name == 'moe_deepep_num_sms') + self.assertEqual( + field.default, 20, + 'moe_deepep_num_sms field default must remain 20 for backward compatibility.' + ) + + def test_explicit_user_override_is_honored(self): + """When the user passes --moe-deepep-num-sms=24, the value must be propagated.""" + args = _make_args(moe_deepep_num_sms=24) + kw_args = self._simulate_core_transformer_config_from_args(args) + self.assertEqual( + kw_args.get('moe_deepep_num_sms'), 24, + 'Explicit user override of moe_deepep_num_sms must be honored.' + ) + + def test_fused_a2a_config_num_sms_overrides_still_work(self): + """The fused_a2a_config.num_sms override path must continue to work after the fix.""" + # User sets --moe-a2a-num-sms=8 (overrides everything) + args = _make_args(moe_a2a_num_sms=8, moe_deepep_num_sms=24) + cfg = resolve_fused_a2a_config_from_sources( + cli_args=args, env=dict(os.environ), config_file_path=None, + ) + # In validate_args the propagation also handles moe_deepep_num_sms -> fused_a2a_config + # But in the raw resolver, moe_deepep_num_sms is not part of the merged dict. + # We just verify that the resolver respects moe_a2a_num_sms=8. + self.assertEqual(cfg.num_sms, 8) + + def test_full_validate_args_propagation_does_not_set_args_moe_deepep_num_sms(self): + """The validate_args propagation block mutates fused_a2a_config but must not + touch args.moe_deepep_num_sms (the regression would have been that the new + field was None; the fix is to not copy it to TransformerConfig).""" + # Simulate: user passes nothing related to deepep + args = _make_args() + # Simulate the resolve + propagation block in validate_args + args.fused_a2a_config = resolve_fused_a2a_config_from_sources( + cli_args=args, env=dict(os.environ), config_file_path=None, + ) + # args.moe_deepep_num_sms remains None (CLI default) + self.assertIsNone(args.moe_deepep_num_sms) + # Then core_transformer_config_from_args would NOT copy this to kw_args + # (verified by test_default_path_preserves_moe_deepep_num_sms_20 above). + kw_args = self._simulate_core_transformer_config_from_args(args) + self.assertNotIn('moe_deepep_num_sms', kw_args) + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/build_sequences_per_dataset.py b/tools/build_sequences_per_dataset.py index 9213a0e1f53..1cc22c66f64 100644 --- a/tools/build_sequences_per_dataset.py +++ b/tools/build_sequences_per_dataset.py @@ -21,7 +21,6 @@ def get_paths_from_blend( blend_per_split: Optional[List[Optional[Tuple[List[str], Optional[List[float]]]]]], ) -> List[str]: """Extract all dataset paths from blend and blend_per_split. - Args: blend (Optional[Tuple[List[str], Optional[List[float]]]]): A blend tuple containing a list of dataset paths and optionally a list of weights, e.g., @@ -29,12 +28,10 @@ def get_paths_from_blend( blend_per_split (Optional[List[Optional[Tuple[List[str], Optional[List[float]]]]]]): A list of 3 blend tuples (for train, valid, test splits), where each element has the same structure as blend - Returns: List[str]: A list of all unique dataset paths found in blend and blend_per_split """ paths = [] - # Extract paths from blend if blend is not None: paths_list, _ = blend