feat(moe): allow user-tuned FusedA2A config (chunk_size, num_sms) for DeepEP#5230
feat(moe): allow user-tuned FusedA2A config (chunk_size, num_sms) for DeepEP#5230DhineshPonnarasan wants to merge 10 commits into
Conversation
…and Dockerfile.linting
…ion tests - Add FusedA2AConfig dataclass (chunk_size, num_sms) with fail-fast validation, unknown-key detection, and __repr__ for observability - Add resolve_fused_a2a_config_from_sources() with precedence CLI > ENV(MOE_A2A_CHUNK_SIZE/MOE_A2A_NUM_SMS) > JSON/YAML file > defaults - Fix _add_moe_args: new args were prepended before group was defined (NameError); move them inside the function in a clearly-labeled DeepEP/A2A tuning section - Add --moe-deepep-num-sms, --moe-a2a-chunk-size, --moe-a2a-num-sms, --moe-a2a-config-file CLI flags to _add_moe_args - Wire resolution into validate_args: resolved once, stored as args.fused_a2a_config, automatically picked up by core_transformer_config_from_args - Add fused_a2a_config: Optional[FusedA2AConfig] = None field to TransformerConfig so config flows from validate_args -> TransformerConfig -> dispatcher without any per-module resolution - Update MoETokenDispatcher.__init__: read fused_a2a_config from config (single source) instead of always initializing to None - Update _DeepepManager.__init__: fused_a2a_config.num_sms takes priority over moe_deepep_num_sms, eliminating the two-parallel-sources inconsistency - Add rank-0 startup logging of effective FusedA2AConfig when any value is set - Add --moe-deepep-num-sms fallback propagation into fused_a2a_config.num_sms so _DeepepManager has one authoritative source regardless of which flag is used - Add 27 unit + integration tests covering all precedence rules, validation, fail-fast, backward compat (None defaults), dispatcher field wiring, and _DeepepManager SM-count resolution; all tests pass Backward compatible: all new fields and args default to None / existing defaults.
Following dadadada-147's comment on issue NVIDIA#4316, deep_ep.Buffer.dispatch() and .combine() accept an optional Config object (deep_ep_cpp.Config) that exposes all kernel-level tunables including num_max_nvl_chunked_send_tokens (chunk_size). Changes: - Import deep_ep_cpp.Config as DeepEPConfig in try/except block - Import FusedA2AConfig at module level - Add _build_deepep_config() helper: starts from Buffer.get_dispatch/combine_config() baseline and overrides only the fields the user explicitly set (num_sms, chunk_size -> num_max_nvl_chunked_send_tokens); all RDMA/buffer-size params preserved from default, ensuring both intranode and internode correctness - Add fused_a2a_config=None to FusedDispatch.forward; build dispatch_config and combine_config once per call; pass dispatch_config to buffer.dispatch() - Save combine_config in ctx for FusedDispatch.backward -> buffer.combine() - Add fused_a2a_config=None to FusedCombine.forward; pass combine_config to buffer.combine(); save dispatch_config in ctx for backward -> buffer.dispatch() - Update backward return tuples (+1 None per added forward param) - Add config= kwarg to fused_dispatch() and fused_combine() wrapper functions, pass through to FusedDispatch.apply() / FusedCombine.apply() This closes the runtime TypeError: token_dispatcher.py already passed config=self.fused_a2a_config but fused_a2a.py did not accept it. chunk_size is now fully functional end-to-end: --moe-a2a-chunk-size / MOE_A2A_CHUNK_SIZE / JSON-YAML config file -> FusedA2AConfig.chunk_size -> _build_deepep_config() -> deep_ep_cpp.Config.num_max_nvl_chunked_send_tokens -> buffer.dispatch() / buffer.combine() All 27 tests pass.
- Move FusedA2AConfig import before the DeepEP try/except (clean import order) - Define _build_deepep_config inside 'if HAVE_DEEP_EP' block so DeepEPConfig (deep_ep_cpp.Config) is never referenced when DeepEP is not installed - Set _build_deepep_config = None in else branch for explicit, inspectable state - Remove '# noqa: F821' workaround that was masking the real fix Verified: importing fused_a2a without DeepEP installed now correctly shows HAVE_DEEP_EP=False, _build_deepep_config=None with no NameError. All 27 tests pass.
The new --moe-deepep-num-sms flag defaults to None. When the user does not pass it, core_transformer_config_from_args was copying the None value into TransformerConfig.moe_deepep_num_sms, shadowing the historical dataclass default of 20 and causing _DeepepManager to call Buffer.set_num_sms(None), which crashes. Skip explicit None values when the corresponding dataclass field has a non-None default. Explicit user overrides (e.g. --moe-deepep-num-sms=24) are still honored, and the 67 Optional-with-default-None fields are unaffected (their default is None, so the skip rule does not trigger). Add 5 regression tests in TestMoeDeepepNumSmsDefaultRegression covering the default path, the field default, explicit user overrides, the fused_a2a_config override path, and the full validate_args propagation flow. All 32 tests pass.
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
|
Hi @dadadada-147, Following up on the discussion in #4316, I've opened this PR implementing configurable DeepEP settings for fused A2A dispatch/combine, including the chunk_size wiring discussed in the issue. I'd appreciate it if you could take a look when you have a chance. Thanks for pointing me to the DeepEP configuration APIs and for the guidance throughout the discussion. |
I believe this precisely addresses the issue where custom configurations could not be used. |
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
Thanks for taking a look. I appreciate the feedback. |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| # FusedA2AConfig is a lightweight dataclass with no heavy dependencies — safe to import directly. | ||
| from megatron.core.transformer.moe.fused_a2a_config import FusedA2AConfig |
There was a problem hiding this comment.
Bug (critical): This change removes from .._rank_utils import log_single_rank, but log_single_rank is still used at line 1194 of this file (inside __post_init__). This will cause a NameError at runtime when fp32_residual_connection is enabled and pipeline_dtype is not torch.float.
The fix is to keep the original import alongside the new one:
from .._rank_utils import log_single_rank
# FusedA2AConfig is a lightweight dataclass with no heavy dependencies — safe to import directly.
from megatron.core.transformer.moe.fused_a2a_config import FusedA2AConfig
| def core_transformer_config_from_args(args, config_class=None): | ||
|
|
||
| # Config class. | ||
| config_class = config_class or TransformerConfig | ||
|
|
||
| if args.multi_latent_attention: | ||
| config_class = MLATransformerConfig | ||
|
|
||
| if args.heterogeneous_layers_config_path is not None: | ||
| assert not args.multi_latent_attention, "Multi latent attention with heterogeneous layers is not supported." | ||
| config_class = HeterogeneousTransformerConfig | ||
|
|
||
| # Translate args to core transformer configuration | ||
| kw_args = {} | ||
| for f in dataclasses.fields(config_class): | ||
| if hasattr(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 | ||
| kw_args['batch_p2p_comm'] = not args.overlap_p2p_comm | ||
| kw_args['num_moe_experts'] = args.num_experts | ||
| kw_args['rotary_interleaved'] = args.rotary_interleaved | ||
| kw_args['num_layers_in_first_pipeline_stage']= args.decoder_first_pipeline_num_layers | ||
| kw_args['num_layers_in_last_pipeline_stage']= args.decoder_last_pipeline_num_layers | ||
| kw_args['fp8_param'] = args.fp8_param_gather | ||
| if args.swiglu: | ||
| kw_args['activation_func'] = F.silu | ||
| kw_args['gated_linear_unit'] = True | ||
| kw_args['bias_activation_fusion'] = args.bias_swiglu_fusion | ||
| else: | ||
| kw_args['bias_activation_fusion'] = args.bias_gelu_fusion | ||
| if args.squared_relu: | ||
| assert not args.swiglu | ||
| kw_args['activation_func'] = squared_relu | ||
| elif args.quick_geglu: | ||
| assert not args.swiglu | ||
| kw_args['gated_linear_unit'] = True | ||
| kw_args['activation_func'] = quick_gelu | ||
| if args.init_method_xavier_uniform: | ||
| kw_args['init_method'] = torch.nn.init.xavier_uniform_ | ||
| kw_args['scaled_init_method'] = torch.nn.init.xavier_uniform_ | ||
| if args.group_query_attention: | ||
| kw_args['num_query_groups'] = args.num_query_groups | ||
| else: | ||
| kw_args['num_query_groups'] = None | ||
| kw_args['config_logger_dir'] = args.config_logger_dir | ||
| if args.rope_type is None: | ||
| # Pop 'rope_type' to let the config class use the default value. | ||
| kw_args.pop('rope_type', None) | ||
| else: | ||
| assert (args.multi_latent_attention or args.rope_type == 'rope'), ( | ||
| f'Common attention only support rope_type="rope", but got {args.rope_type}.' | ||
| ) | ||
|
|
||
| if len(args.cp_comm_type) == 1: | ||
| kw_args['cp_comm_type'] = args.cp_comm_type[0] | ||
| if args.hybrid_layer_pattern is not None: | ||
| kw_args['is_hybrid_model'] = True | ||
|
|
||
| kw_args['inference_sampling_seed'] = args.seed | ||
|
|
||
| # handle quantization config | ||
| # NOTE: Kitchen arguments are only added to the namespace when | ||
| # Kitchen library is available. | ||
| if hasattr(args, "kitchen_config_file") and args.kitchen_config_file is not None: | ||
| kw_args['use_kitchen'] = True | ||
| kw_args['quant_recipe'] = load_quantization_recipe(args.kitchen_config_file) | ||
| elif hasattr(args, 'kitchen_recipe_number') and args.kitchen_recipe_number is not None: | ||
| kw_args['use_kitchen'] = True | ||
| kw_args['quant_recipe'] = kitchen_quantization_recipe_config(args.kitchen_recipe_number) | ||
|
|
||
| kw_args['moe_latent_size'] = args.moe_latent_size | ||
|
|
||
| if args.te_precision_config_file: | ||
| assert not 'quant_recipe' in kw_args, "Quantization recipe already configured." | ||
| # TODO(kwyss): Prohibit fp8_params or fp4_params with this flexibility | ||
| kw_args['quant_recipe'] = load_quantization_recipe(args.te_precision_config_file) | ||
|
|
||
| if hasattr(args, "use_kitchen_attention"): | ||
| kw_args['use_kitchen_attention'] = args.use_kitchen_attention | ||
| if hasattr(args, "kitchen_attention_backend"): | ||
| kw_args['kitchen_attention_backend'] = args.kitchen_attention_backend | ||
|
|
||
| # Return config. | ||
| return config_class(**kw_args) |
There was a problem hiding this comment.
Bug (critical): This PR adds a second core_transformer_config_from_args definition here, but the canonical version is already imported from argument_utils.py at line 55:
from megatron.training.argument_utils import ArgumentGroupFactory, core_transformer_config_from_argsThis local definition shadows the import. The PR's copy diverges from the canonical version in argument_utils.py — it is missing at least:
kw_args['fp4_param'] = args.fp4_param_gather(FP4 quantization will silently break)- The DSA hybrid layer handling (
Symbols.DS_ATTENTION→experimental_attention_variant)
Any caller that does from megatron.training.arguments import core_transformer_config_from_args will now get this incomplete copy.
Recommended fix: Do not duplicate this function. Instead, apply the None-skipping fix to the canonical version in megatron/training/argument_utils.py and remove this entire duplicate definition.
There was a problem hiding this comment.
Two critical bugs found:
-
transformer_config.py: Removed import still in use — The PR deletesfrom .._rank_utils import log_single_rankbutlog_single_rankis still called at line 1194. This will raise aNameErrorat runtime whenfp32_residual_connectionis enabled. -
arguments.py: Duplicatecore_transformer_config_from_args— The PR adds a second definition of this function (line 1822) that shadows the canonical import fromargument_utils.py(line 55). The duplicate is missingfp4_paramhandling and DSA hybrid layer logic, which will silently break those features. The None-skipping fix should be applied to the canonical version inargument_utils.pyinstead.
See inline comments for details.
The FusedA2AConfig import addition accidentally removed the 'from .._rank_utils import log_single_rank' line. log_single_rank is still used at line 1028 in __post_init__ when fp32_residual_connection is enabled and pipeline_dtype is not torch.float — removing the import causes a NameError at runtime. Restores the import. Fixes critical bug flagged by claude[bot] in PR NVIDIA#5230.
|
Thanks for working on this. I agree the However, I do not think this is ready as-is:
|
@guihong-nv Thanks for the review. I've addressed the issues around the duplicated Regarding the DeepEP V2 work, I took a closer look at #5153 and #4632. My understanding is that this PR is focused on exposing configuration for the existing DeepEP V1 dispatch/combine path, while the V2 work introduces a separate backend. There may be some rebase work needed if those land first, but I didn't find a direct functional conflict with the chunk_size configuration introduced here. Please let me know if you'd prefer this functionality to be rebased onto the V2 work instead, and I'm happy to adjust accordingly. |
Move the None-default preservation logic into the canonical core_transformer_config_from_args implementation in argument_utils.py and remove the duplicate implementation from arguments.py. This preserves existing FP4 and DSA handling while fixing the moe_deepep_num_sms default regression introduced by the new optional CLI flag.
|
Thanks for the update. The previous two blockers look fixed ( Main issue:
This should cause an argparse conflict during parser construction. The new A couple of smaller issues:
We strongly suggest the author run some tests before uploading the code. These issues can be quickly found and located. |
Address review comments from guihong-nv on PR NVIDIA#5230. A) Add moe_deepep_num_sms and fused_a2a_config to ArgumentGroupFactory's exclude list in _add_network_size_args. Without this, the factory auto-generates --moe-deepep-num-sms from TransformerConfig (with default=20) and --fused-a2a-config from the FusedA2AConfig field, which would conflict with the manual --moe-deepep-num-sms in _add_moe_args (raising ArgumentError at parser-construction time) and pollute --help with a non-functional --fused-a2a-config flag. B) Re-validate args.fused_a2a_config after the --moe-deepep-num-sms propagation block in validate_args, so the even-SM check also covers values supplied via --moe-deepep-num-sms (which the resolver does not see). C) Add fail-fast even-SM validation in FusedA2AConfig.validate(). DeepEP's Buffer.set_num_sms and the C++ kernel both assert new_num_sms % 2 == 0; odd values would crash deep inside the kernel with an opaque assertion. Validate up front. D) Remove trailing whitespace at line 18 of fused_a2a_config.py (flagged by git diff --check on the PR diff). E) The whitespace-only edits in docker/Dockerfile.linting and tools/build_sequences_per_dataset.py were brought in via merges from main (commits d81b638 and 9ef8a2a), not by the PR's own feature commits. They are not PR-introduced changes and are left alone. Tests: - tests/unit_tests/moe/test_fused_a2a_config_validation.py: add test_valid_none, test_valid_even_num_sms, test_odd_num_sms_rejected, test_odd_num_sms_alone_rejected. - tests/unit_tests/moe/test_fused_a2a_integration.py: add test_odd_moe_deepep_num_sms_propagation_raises, update _resolve helper to call validate() after the propagation block (mirroring the validate_args fix). - tests/unit_tests/moe/test_fused_a2a_config.py: update test_cli_priority to use num_sms=8 (was 7, now rejected as odd). All 37 unit + integration tests pass locally. Signed-off-by: Dhinesh Ponnarasan <dhineshponnarasan@gmail.com>
@guihong-nv Thanks for the review. I have addressed the reported issues in commit
The commit has been pushed and CI has been re-triggered. Please let me know if there are any remaining concerns. |
5 participants
Summary
Adds user-configurable DeepEP settings for fused A2A dispatch/combine.
This addresses #4316 by allowing selected DeepEP configuration values to be overridden instead of always using the default hardware-tuned configuration.
Motivation
DeepEP already supports configurable dispatch/combine settings through its Config object, but FusedDispatch and FusedCombine currently use default values. This change exposes selected tuning parameters while preserving existing behavior when no overrides are provided.
Implementation
Added FusedA2AConfig for user-defined DeepEP settings.
Added support for:
Added JSON/YAML configuration support.
Added _build_deepep_config() which starts from DeepEP hardware-tuned defaults and overrides only user-specified values.
Wired configuration through argument parsing, TransformerConfig, MoETokenDispatcher, and DeepEP dispatch/combine paths.
chunk_size maps to Config.num_max_nvl_chunked_send_tokens.
Backward Compatibility
Testing
Added unit and integration coverage for:
32 tests pass locally.
Closes #4316