Add direct hybrid architecture descriptors - #6295
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
/ok to test a38a1b6 |
|
/claude strict-review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
Strict Review Summary Findings: CRITICAL: 2 · IMPORTANT: 7 · SUGGESTION: 2
Most impactful 1.
2. Per-layer configs carry a stale
Also worth addressing before the stack lands
Risk assessment Moderate for this PR in isolation; high for the direct-descriptor path once wired up. Nothing in No dtype, loss-scaling, or attention-math concerns — this PR contains no numerical kernels. No process-group concerns either: the module takes |
a38a1b6 to
ab40c9f
Compare
|
/ok to test ab40c9f |
Signed-off-by: Philip Petrakian <ppetrakian@nvidia.com>
ab40c9f to
f673468
Compare
|
/ok to test f673468 |
|
/claude strict-review |
| + ", ".join(active_conflicts) | ||
| ) | ||
|
|
||
| config.virtual_pipeline_model_parallel_size = inferred_vp_size if inferred_vp_size > 1 else None |
There was a problem hiding this comment.
[CRITICAL Correctness] A function named _validate_direct_pipeline mutates the caller's TransformerConfig to infer VPP degree, and it does so at a point in the lifecycle where that value can no longer take effect.
Ordering in pretrain(): initialize_megatron → initialize_model_parallel(..., args.virtual_pipeline_model_parallel_size) (creates the PP/VPP process groups, parallel_state.py:827) → setup_model_and_optimizer → get_model → model_provider → HybridModel.__init__ → (in the follow-up PR of this stack) resolve_hybrid_architecture. By the time this line runs:
get_model's chunk loop (megatron/training/training.py:1757-1780) has already branched onargs.virtual_pipeline_model_parallel_size is not None. If the user only expressed VPP viaPipelineSplit()nodes and left VPP unset, exactly one chunk gets built, whileselect_segmentnow believes there arepp_size * 2segments — so half the model's layers are never instantiated andconfig.virtual_pipeline_model_parallel_size = 2is then read byschedules.py:1517/p2p_communication.py:161, producing a scheduler that disagrees with the built model.- The VPP process groups were sized from the pre-mutation value, so they can't be corrected here either.
Contrast the legacy path, which derives VPP at argument-parse time (megatron/training/arguments.py:948, "Only set VPP to None if it wasn't already derived from --hybrid-layer-pattern"). The direct path has no equivalent arg-time derivation.
Also note the mutation runs once per VPP chunk (the same config object is shared by every model_provider_func call), which makes the write order-dependent for no benefit.
Suggested fix: make this a pure validator that requires the config to already agree, and move derivation to argument/config-construction time where the legacy path does it:
inferred_vp_size = num_segments // pp_size
if pp_size == 1 and inferred_vp_size > 1:
raise ValueError("Virtual pipeline parallelism requires pipeline_model_parallel_size > 1.")
expected_vp_size = inferred_vp_size if inferred_vp_size > 1 else None
if config.virtual_pipeline_model_parallel_size != expected_vp_size:
raise ValueError(
f"PipelineSplit() nodes imply virtual_pipeline_model_parallel_size="
f"{expected_vp_size}, but the config specifies "
f"{config.virtual_pipeline_model_parallel_size}. Set it before the process "
f"groups and model chunks are created."
)
# no mutation hereIf deriving-from-the-pattern is the intended UX, it needs a hook that runs before initialize_model_parallel, not inside the layer resolver.
| def _materialize_config( | ||
| layer_config: TransformerConfig, base_config: TransformerConfig | ||
| ) -> TransformerConfig: | ||
| materialized = copy.deepcopy(layer_config) | ||
| for field_name in _TOPOLOGY_CONFIG_FIELDS: | ||
| setattr(materialized, field_name, getattr(base_config, field_name)) | ||
| return materialized |
There was a problem hiding this comment.
[CRITICAL Correctness] _materialize_config produces one deepcopy of TransformerConfig per layer occurrence, which permanently detaches every layer from the runtime mutations the training loop performs on the single shared config object. Several of those mutations are load-bearing:
| Written at | Field | Read from self.config inside layers |
|---|---|---|
training.py:3566-3588, 2152, 4284, 4300 |
timers, grad_scale_func, no_sync_func, grad_sync_func, param_sync_func, finalize_model_grads_func |
via get_model_config(model) — unwrapped model config, so still OK |
transformer_block.py:310,319 |
_cpu_offloading_context |
tensor_parallel/layers.py:1157,1447 reads self.config._cpu_offloading_context on the layer's config |
transformer/utils.py:367 (set_model_to_sequence_parallel, called from text_generation_controller.py:3480) |
sequence_parallel |
~20 sites read self.config.sequence_parallel at forward time (multi_latent_attention.py:654,745,753, shared_experts.py:246,341,550, moe_layer.py:629, dot_product_attention.py:216, …) |
transformer/utils.py:451 (toggle_cuda_graphs, called from rl_utils.py:1605,2243,2256) |
cuda_graph_impl |
attention.py:1449, transformer_layer.py:460, hybrid_block.py:207,308 |
Both toggle helpers write model.config.<field> — the model-level config — and then walk modules setting the attribute of the same name. Per-layer configs are neither the object written nor an attribute they set, so a layer that reads self.config.sequence_parallel / self.config.cuda_graph_impl sees the stale pre-toggle value. Concretely: MoE+TP inference prefill (text_generation_controller.py:3480) intends to disable SP everywhere except BaseMoELayer; with per-layer configs the non-MoE hybrid layers keep sequence_parallel=True and will scatter/gather on an unsharded sequence — silently wrong activations, not a crash. Similarly _cpu_offloading_context is written onto self.config by the block but read by ColumnParallelLinear off its config, so CPU offloading silently no-ops for offload-enabled hybrid runs.
There's also a deepcopy hazard: TransformerConfig holds non-value fields (quant_recipe: RecipeConfig — a plain class with no __eq__, init_method/activation_func callables, _cpu_offloading_context: ContextManager). Deep-copying these per layer both duplicates state that was meant to be shared and makes the != comparison in _materialize_layer (line 441) identity-based for quant_recipe, so two structurally identical recipes will spuriously raise.
Suggestion: don't deepcopy. Materialize by dataclasses.replace-style sharing of the mutable/runtime fields, or — cleaner — keep the base config as the single object and carry only the diff per occurrence (e.g. HybridLayerSpec(module_spec, overrides: dict)) that layers apply at build time. If deepcopy must stay, add an explicit allowlist of runtime-mutated fields that are re-pointed at the base config after copying (_cpu_offloading_context, sequence_parallel, cuda_graph_impl, timers, *_sync_func, grad_scale_func, quant_recipe) and document why.
| _UNIFORM_CONFIG_FIELDS = ( | ||
| "hidden_size", | ||
| "params_dtype", | ||
| "pipeline_dtype", | ||
| "fp16", | ||
| "bf16", | ||
| "fp32_residual_connection", | ||
| "enable_autocast", | ||
| "autocast_dtype", | ||
| "apply_query_key_layer_scaling", | ||
| "attention_softmax_in_fp32", | ||
| "disable_bf16_reduced_precision_matmul", | ||
| "dsa_indexer_k_norm_fp32", | ||
| "fp8", | ||
| "fp8_recipe", | ||
| "fp8_param", | ||
| "fp8_quantizer_factory", | ||
| "fp8_margin", | ||
| "fp8_interval", | ||
| "fp8_amax_history_len", | ||
| "fp8_amax_compute_algo", | ||
| "fp8_wgrad", | ||
| "fp8_output_proj", | ||
| "fp8_dot_product_attention", | ||
| "fp8_multi_head_attention", | ||
| "tp_only_amax_red", | ||
| "activation_func_fp8_input_store", | ||
| "first_last_layers_bf16", | ||
| "num_layers_at_start_in_bf16", | ||
| "num_layers_at_end_in_bf16", | ||
| "moe_router_dtype", | ||
| "mamba_training_ssm_states_dtype", | ||
| "fp4", | ||
| "fp4_recipe", | ||
| "fp4_param", | ||
| "fp4_quantizer_factory", | ||
| "quant_recipe", | ||
| "normalization", | ||
| "layernorm_epsilon", | ||
| "layernorm_zero_centered_gamma", | ||
| "tensor_model_parallel_size", | ||
| "tensor_parallel_num_weight_shards", | ||
| "gtp_weight_remat_size", | ||
| "pipeline_model_parallel_size", | ||
| "pipeline_model_parallel_comm_backend", | ||
| "context_parallel_size", | ||
| "hierarchical_context_parallel_sizes", | ||
| "max_seqlen_per_dp_cp_rank", | ||
| "hybrid_context_parallel", | ||
| "expert_model_parallel_size", | ||
| "expert_tensor_parallel_size", | ||
| "expert_tensor_parallel_num_weight_shards", | ||
| "expert_gtp_weight_remat_size", | ||
| "sequence_parallel", | ||
| ) |
There was a problem hiding this comment.
[CRITICAL Correctness] _UNIFORM_CONFIG_FIELDS is a hand-maintained denylist over a config surface of 338 fields (259 in TransformerConfig + 79 in ModelParallelConfig). Anything not listed is silently accepted as a legal per-occurrence override, and several omissions are not benign:
multi_latent_attention— read athybrid_model.py:247,482off the model config to decide whether to build the standardRotaryEmbeddingat all. Anmlaoccurrence that flips it per-layer gets no decoupled RoPE while the model builds standard RoPE (or vice versa). This is exactly the mix thathybrid_layer_allocation.py:297and:325explicitly reject for the legacy path ("Not supported to have both Attention and MLA/DSA in one model") — the direct path has no equivalent check forattention+mla/dsacoexistence at all.recompute_granularity/recompute_method/recompute_num_layers/recompute_modules—recompute_num_layersis a whole-block count consumed by the block, not the layer; per-layer values are meaningless and will mis-slice recompute.cpu_offloading/cpu_offloading_num_layers— same whole-block semantics;transformer_block.py:303passesself.config.num_layersintoget_cpu_offload_context.cuda_graph_impl/cuda_graph_modules,deterministic_mode,gradient_accumulation_fusion,tp_comm_overlap,perform_initialization/use_cpu_initialization,moe_token_dispatcher_type,moe_grouped_gemm,batch_invariant_mode— all model-global in practice.- MLA/GDN shape fields are neither in the denylist nor in the inference-shape signatures:
kv_lora_rank,qk_head_dim,v_head_dim,qk_pos_emb_head_dim(the last two feeddynamic_context.py:491'skv_reduced_dim), andlinear_key_head_dim/linear_num_key_heads/linear_value_head_dim/linear_num_value_heads/linear_conv_kernel_dimfor GDN state shapes.
Denylists over a 338-field surface will keep drifting as fields are added — nothing in CI notices. Please invert to an allowlist of the fields a per-occurrence override may legally change (mamba shape group, attention head/KV group, ffn_hidden_size, moe_ffn_hidden_size, moe_router_topk, MLA rank/head-dim group, GDN linear-attention group, window_size, softmax_scale, …) and reject everything else with the same error message. That makes a newly added field fail closed instead of open, and it removes the need to keep 54 unrelated names in sync. If an allowlist is too restrictive for the follow-up PRs, at minimum add a unit test that asserts set(dataclasses.fields(TransformerConfig)) - allowed - denied == set() so a new field forces a triage decision.
| @property | ||
| def metric_num_layers(self) -> int: | ||
| """Return the number of explicit layer slots used by MoE metrics.""" | ||
|
|
||
| return len(self.main_layers) + self.mtp_num_layers * len(self.mtp_layers) |
There was a problem hiding this comment.
[IMPORTANT Correctness] metric_num_layers doesn't match what the MoE metrics tracker actually expects, so a direct-architecture run would allocate a mis-sized metric tensor.
Router._record_aux_loss (megatron/core/transformer/moe/router.py:590-593) computes the tensor size as:
num_layers = self.config.num_layers
if self.config.mtp_num_layers is not None:
num_layers += self.config.mtp_num_layersi.e. num_layers + mtp_num_layers — not num_layers + mtp_num_layers * len(mtp_layers). The same formula appears in MoEMetricsTracker.report (moe_logging.py:184: init_size = num_layers + (mtp_num_layers or 0)), and layer_number for MTP layers is self.layer_number + self.config.num_layers (router.py:596).
With mtp_num_layers=2 and a 2-layer MTP block, this property returns num_layers + 4 while the router allocates num_layers + 2 and indexes into it. Since record initializes on first use (moe_logging.py:122) and later .report() compares/reduces across ranks, a size disagreement between what this property reports and what the router allocates produces either an allreduce shape mismatch across PP ranks or a silently wrong per-layer aux-loss attribution.
Either match the router's formula, or — if the intent is genuinely "one slot per materialized MTP layer" — that's a change to the tracker's indexing convention and needs to land together with the router/tracker change, not ahead of it.
Also: this property has no test, and neither do has_heterogeneous_inference_shapes and has_incompatible_dynamic_inference_shapes. Given the 633-line test file covers the rest of the module thoroughly, these three look like an oversight rather than a deliberate gap.
| @property | ||
| def has_heterogeneous_inference_shapes(self) -> bool: | ||
| """Whether dynamic inference would need non-uniform cache/buffer shapes.""" | ||
|
|
||
| layers = self.main_layers + self.mtp_layers | ||
| signatures: dict[str, set[tuple[Any, ...]]] = { | ||
| "mamba": set(), | ||
| "attention": set(), | ||
| "moe": set(), | ||
| } | ||
| for layer in layers: | ||
| config = layer.config | ||
| if layer.layer_type == "mamba": | ||
| signatures["mamba"].add( | ||
| ( | ||
| config.mamba_state_dim, | ||
| config.mamba_head_dim, | ||
| config.mamba_num_heads, | ||
| config.mamba_num_groups, | ||
| ) | ||
| ) | ||
| elif layer.layer_type == "attention": | ||
| signatures["attention"].add( | ||
| (config.num_query_groups or config.num_attention_heads, config.kv_channels) | ||
| ) | ||
| elif layer.layer_type == "moe": | ||
| signatures["moe"].add((config.moe_router_topk,)) | ||
| return any(len(values) > 1 for values in signatures.values()) | ||
|
|
||
| def has_incompatible_dynamic_inference_shapes(self, base_config: TransformerConfig) -> bool: | ||
| """Whether dynamic inference's model-global buffers fit every occurrence. | ||
|
|
||
| Dynamic inference currently allocates attention KV and MoE routing buffers | ||
| from the model-wide config. Even a family-uniform direct override is | ||
| incompatible when it differs from that allocation source. Mamba state | ||
| shapes are discovered from the built layer, so only differing Mamba | ||
| occurrences are rejected. | ||
| """ | ||
|
|
||
| if self.has_heterogeneous_inference_shapes: | ||
| return True | ||
|
|
||
| base_attention_signature = ( | ||
| base_config.num_query_groups or base_config.num_attention_heads, | ||
| base_config.kv_channels, | ||
| ) | ||
| layers = self.main_layers + self.mtp_layers | ||
| for layer in layers: | ||
| if layer.layer_type == "attention": | ||
| occurrence_signature = ( | ||
| layer.config.num_query_groups or layer.config.num_attention_heads, | ||
| layer.config.kv_channels, | ||
| ) | ||
| if occurrence_signature != base_attention_signature: | ||
| return True | ||
| elif ( | ||
| layer.layer_type == "moe" | ||
| and layer.config.moe_router_topk != base_config.moe_router_topk | ||
| ): | ||
| return True | ||
| return False |
There was a problem hiding this comment.
[IMPORTANT Correctness] These two inference-shape predicates are declared but never called anywhere in the tree (verified by grep across megatron/ and tests/ — the only references are the export lists in __init__.py). Per the PR-stack framing that's presumably intentional for a later PR, but they encode assumptions that will be hard to correct once callers exist, and three of them look wrong today:
-
gdn,mla, anddsaoccurrences are invisible to both predicates. Onlymamba,attention, andmoeare inspected. GDN carries its own recurrent-state shape (linear_key_head_dim,linear_num_key_heads,linear_value_head_dim,linear_num_value_heads,linear_conv_kernel_dim), and MLA/DSA KV-cache sizing comes fromkv_lora_rank+qk_pos_emb_head_dim(dynamic_context.py:491). A model that varies any of those across occurrences is reported as homogeneous and will get an undersized cache. -
The attention signature omits
num_attention_heads.(num_query_groups or num_attention_heads, kv_channels)collapses two genuinely different layers:num_attention_heads=16, num_query_groups=4andnum_attention_heads=8, num_query_groups=4produce the identical signature(4, kv_channels). KV cache sizing does only depend onnum_query_groups, so this may be deliberate — if so, say so in the docstring, because the name "inference shapes" reads as covering the whole attention shape. Note also thatnum_query_groupsis resolved tonum_attention_headsinTransformerConfig.__post_init__(transformer_config.py:1335), making theorfallback dead for any real config. -
has_incompatible_dynamic_inference_shapesshort-circuits onhas_heterogeneous_inference_shapes, which folds Mamba back in — contradicting its own docstring. The docstring says "Mamba state shapes are discovered from the built layer, so only differing Mamba occurrences are rejected," but line 149 returnsTruefor any heterogeneity including Mamba, before the per-family base-config comparison runs. If Mamba heterogeneity is meant to be tolerated here, the check needs to be per-family rather than a blanket short-circuit.
Suggest either extending these to cover all six layer families and fixing the Mamba short-circuit, or deferring both properties to the PR that introduces the caller so the semantics can be reviewed against real use.
| if num_segments > 1: | ||
| conflicts = { | ||
| "pipeline_model_parallel_layout": config.pipeline_model_parallel_layout, | ||
| "num_layers_in_first_pipeline_stage": config.num_layers_in_first_pipeline_stage, | ||
| "num_layers_in_last_pipeline_stage": config.num_layers_in_last_pipeline_stage, | ||
| "account_for_embedding_in_pipeline_split": ( | ||
| config.account_for_embedding_in_pipeline_split | ||
| ), | ||
| "account_for_loss_in_pipeline_split": config.account_for_loss_in_pipeline_split, | ||
| } | ||
| active_conflicts = [name for name, value in conflicts.items() if value not in (None, False)] | ||
| if active_conflicts: | ||
| raise ValueError( | ||
| "PipelineSplit() already defines pipeline ownership and cannot be combined with: " | ||
| + ", ".join(active_conflicts) | ||
| ) |
There was a problem hiding this comment.
[IMPORTANT Correctness] Two issues in the conflict detection:
-
num_segments > 1is the wrong guard. Withpp_size=1and a single segment,num_layers_in_first_pipeline_stageetc. are skipped entirely — but a user who sets them along with directlayer_specshas expressed a contradiction that still deserves an error. The natural guard is "directlayer_specswas supplied at all", not "there are splits". As written,pp_size=1 + num_layers_in_first_pipeline_stage=4passes silently. -
value not in (None, False)misclassifies0.0 == Falsein Python, so0 in (None, False)isTrue. Ifnum_layers_in_first_pipeline_stage=0is ever reachable it is silently treated as unset.transformer_config.py:2023-2024rejects<= 0for that field, so this is latent rather than live today — but the idiom is fragile and will bite the next boolean-vs-int field added to the dict. Prefer explicit per-field predicates:
active_conflicts = [
name
for name, value in conflicts.items()
if value is not None and value is not False
]That distinguishes 0 from False and reads as the intended "explicitly set" test.
| def _validate_mtp_placement(architecture: ResolvedHybridArchitecture) -> None: | ||
| """Reject a prediction block on a logical chunk with no decoder layers.""" | ||
|
|
||
| if ( | ||
| architecture.mtp_layers | ||
| and architecture.mtp_num_layers > 0 | ||
| and not architecture.segments[-1] | ||
| ): | ||
| raise ValueError( | ||
| "MTP must share the final logical PP/VPP chunk with at least one decoder layer; " | ||
| "standalone MTP placement is not supported." | ||
| ) |
There was a problem hiding this comment.
[IMPORTANT Correctness] _validate_mtp_placement only inspects segments[-1], but with VPP the last physical pipeline stage is not the last segment — and MTP is placed on the last physical stage.
Segment order is VPP-major (select_segment: vp_stage * pp_size + pp_rank), so for pp_size=2, vp_size=2 the segments are laid out [pp0/vp0, pp1/vp0, pp0/vp1, pp1/vp1]. segments[-1] is pp_rank=1, vp_stage=1, which happens to be the final chunk here — but the emptiness check that matters is on whichever chunk mtp_on_this_rank selects (hybrid_model.py:222-227, which uses is_pp_last_stage ∧ is_vp_last_stage). Those coincide only because the last segment is pp-last ∧ vp-last. The bug is the reverse case: an empty segments[-1] combined with a non-empty earlier chunk on the same physical rank is rejected, and a non-empty segments[-1] with MTP landing elsewhere is accepted. The check would be clearer and future-proof written in terms of the selection it's guarding:
def _validate_mtp_placement(architecture, pp_size, vp_size) -> None:
last_chunk, _ = architecture.select_segment(
pp_rank=pp_size - 1, pp_size=pp_size, vp_stage=None if vp_size == 1 else vp_size - 1
)
if architecture.mtp_layers and architecture.mtp_num_layers > 0 and not last_chunk:
raise ValueError(...)Separately: the architecture.mtp_num_layers > 0 term is unreachable-as-a-guard here — resolve_hybrid_architecture already raised at line 245 if raw_mtp_layers is non-empty with mtp_num_layers <= 0, so mtp_layers non-empty implies mtp_num_layers > 0. Harmless, but it suggests the invariant isn't as clear as it could be.
| raw_mtp_segments = ( | ||
| flatten_hybrid_layer_pattern(mtp_layer_specs, allow_splits=False) | ||
| if mtp_layer_specs is not None | ||
| else ((),) | ||
| ) | ||
| raw_mtp_layers = raw_mtp_segments[0] |
There was a problem hiding this comment.
[SUGGESTION Simplification] flatten_hybrid_layer_pattern with allow_splits=False can only ever return a 1-tuple, so the ((),) sentinel + [0] indexing is a two-step construction where one suffices. It also makes the invariant implicit — a reader has to reason about why raw_mtp_segments[0] is safe.
raw_mtp_layers: tuple[HybridLayerSpec, ...] = ()
if mtp_layer_specs is not None:
(raw_mtp_layers,) = flatten_hybrid_layer_pattern(mtp_layer_specs, allow_splits=False)The unpacking asserts the 1-segment invariant instead of assuming it.
| def _validate_uniform_expert_count( | ||
| architecture: ResolvedHybridArchitecture, base_config: TransformerConfig | ||
| ) -> None: | ||
| expert_counts = { | ||
| layer.config.num_moe_experts | ||
| for layer in architecture.main_layers + architecture.mtp_layers | ||
| if layer.layer_type == "moe" | ||
| } | ||
| if len(expert_counts) > 1: | ||
| raise ValueError( | ||
| "All MoE occurrences must use one uniform num_moe_experts value; got " | ||
| f"{sorted(expert_counts)}." | ||
| ) | ||
| if expert_counts and next(iter(expert_counts)) != base_config.num_moe_experts: | ||
| occurrence_count = next(iter(expert_counts)) | ||
| raise ValueError( | ||
| "MoE occurrence num_moe_experts must match the model-wide config; " | ||
| f"got {occurrence_count!r} != {base_config.num_moe_experts!r}." | ||
| ) |
There was a problem hiding this comment.
[SUGGESTION Simplification] The first branch is subsumed by the second: if every MoE occurrence must equal base_config.num_moe_experts, then a set with more than one element necessarily contains a non-matching value. The two-error structure gives a nicer message for the "layers disagree with each other" case, which is worth keeping — but it can be expressed without the double next(iter(...)):
def _validate_uniform_expert_count(architecture, base_config) -> None:
expert_counts = {
layer.config.num_moe_experts
for layer in architecture.main_layers + architecture.mtp_layers
if layer.layer_type == "moe"
}
if not expert_counts:
return
if expert_counts != {base_config.num_moe_experts}:
raise ValueError(
"Every MoE occurrence must use the model-wide num_moe_experts "
f"({base_config.num_moe_experts!r}); got {sorted(expert_counts, key=repr)}."
)Note sorted(expert_counts) on line 485 will raise TypeError if any occurrence leaves num_moe_experts at its None default (transformer_config.py:232) alongside an int — reachable if someone tags a layer moe without setting experts. key=repr or a plain list avoids that.
| @dataclass(frozen=True, slots=True) | ||
| class ResolvedHybridArchitecture: | ||
| """A validated global hybrid architecture shared by every PP/VPP chunk.""" | ||
|
|
||
| segments: tuple[tuple[HybridLayerSpec, ...], ...] | ||
| mtp_layers: tuple[HybridLayerSpec, ...] = () | ||
| mtp_num_layers: int = 0 | ||
| source: str = "direct" |
There was a problem hiding this comment.
[SUGGESTION Naming] Two naming points on this dataclass:
-
mtp_num_layersis not a layer count — it's the number of MTP prediction depths.len(mtp_layers)is the layer count. The name collides withTransformerConfig.mtp_num_layers(which is also depths, so the confusion is pre-existing) and directly causes the ambiguity inmetric_num_layers.mtp_num_depths— the nameParsedHybridPatternalready uses for exactly this quantity (hybrid_layer_allocation.py:71) — would be consistent with the module it's sourced from and describe what the thing is. -
source: str = "direct"is a two-valued discriminator typed as free-formstr.Literal["direct", "legacy"]documents the domain and lets mypy catch a typo'd comparison; it's also the repo idiom for small closed sets (recompute_granularity,moe_token_dispatcher_type,cuda_graph_implintransformer_config.pyare allLiteral).
| def _resolve_legacy_architecture( | ||
| config: TransformerConfig, hybrid_stack_spec: ModuleSpec, hybrid_layer_pattern: str | None | ||
| ) -> ResolvedHybridArchitecture: | ||
| parsed = parse_hybrid_pattern(hybrid_layer_pattern) | ||
| submodules = hybrid_stack_spec.submodules | ||
| symbol_fields = { | ||
| Symbols.MAMBA: ("mamba", "mamba_layer"), | ||
| Symbols.GDN: ("gdn", "gdn_layer"), | ||
| Symbols.ATTENTION: ("attention", "attention_layer"), | ||
| Symbols.DS_ATTENTION: ("dsa", "dsa_layer"), | ||
| Symbols.MLA: ("mla", "mla_layer"), | ||
| Symbols.MLP: ("mlp", "mlp_layer"), | ||
| Symbols.MOE: ("moe", "moe_layer"), | ||
| } | ||
|
|
||
| def descriptor(symbol: str) -> HybridLayerSpec: | ||
| layer_type, field_name = symbol_fields[symbol] | ||
| return HybridLayerSpec( | ||
| module_spec=_legacy_module_spec(getattr(submodules, field_name), layer_type), | ||
| config=config, | ||
| ) |
There was a problem hiding this comment.
[SUGGESTION Simplification] _resolve_legacy_architecture takes hybrid_stack_spec purely to read .submodules, and the symbol_fields map duplicates the mapping already implied by HybridStackSubmodules' field names (hybrid_block.py:45-51). Two small cleanups:
-
The tuple values
("mamba", "mamba_layer")are(layer_type, field_name)wherefield_nameis alwaysf"{layer_type}_layer". Either derive it (f"{layer_type}_layer") or keep only the layer type and let_legacy_module_specdo thegetattr— the duplicated string pair is the kind of thing that drifts if a field is ever renamed. -
descriptor()re-tags on every call, so a pattern like"MMMM"walks_with_layer_typefour times for the sameModuleSpec. Since the stock specs now carry the tag already (hybrid_layer_specs.py),_with_layer_typehits the fast-path return at line 511 — but for an untagged custom stack it allocates a freshcopy.copy+ dict per occurrence. Hoisting the tagging out of the per-symbol loop (tag each of the 7 fields once, then look up) is both cheaper and clearer.
Neither is behavior-affecting; flagging for the follow-up that makes this path load-bearing.
| get_inference_optimized_moe_spec, | ||
| get_moe_module_spec, | ||
| ) | ||
| from megatron.core.models.hybrid.hybrid_architecture import HYBRID_LAYER_TYPE |
There was a problem hiding this comment.
[SUGGESTION Simplification] This import creates a cycle in intent even though Python tolerates it today: hybrid_architecture imports from hybrid_layer_allocation, and hybrid_layer_specs now imports from hybrid_architecture just to get a string constant. The stock specs don't otherwise depend on the resolver.
HYBRID_LAYER_TYPE is a pure metadata key with no behavior attached. Putting it next to Symbols in hybrid_layer_allocation.py (which both modules already import) — or in a tiny leaf module — keeps the dependency edges one-directional and makes it obvious the constant is shared vocabulary rather than resolver-owned.
Also worth noting: the tag values ("mamba", "gdn", …) are duplicated as bare string literals across two stack specs plus VALID_HYBRID_LAYER_TYPES — 15 sites. A LayerType enum or module-level constants would turn a typo into a NameError instead of a validation failure at recipe-construction time.
Strict Review — PR #6295 "Add direct hybrid architecture descriptors"Reviewed the full diff (4 files, +1200/-0) plus surrounding context in Findings: CRITICAL 3 · IMPORTANT 4 · SUGGESTION 5 Most impactful1. This is the one I would block on. The write lands inside 2. Per-layer
3. Fails open. Notable omissions: Also flagged
Verification notesConfirmed by grep that all 54 Risk assessment: Medium-HighNothing here is wired into a live training path yet, which caps the blast radius of merging as-is — and the legacy path is genuinely left read-only, which the tests verify well ( Findings 1, 2, and 3 should be addressed before this becomes load-bearing. The remaining items are ordinary follow-ups. |
1 participant
Summary
HybridLayerSpec,PipelineSplit, semantic layer tags, and the global resolver.HybridModelexecution on its original parser and builder path.PR 1 of the draft stack extracted from #6235.