feat(dsv4): add MXFP4 QAT FP8 expert path#6
Draft
flowerfox-scitix wants to merge 27 commits into
Draft
Conversation
…M PR#28 Absorbs verified DSV4 fixes from coworker tar (origin: radixark/Megatron-LM PR#28, yueming-yuan/miles deepseek-v4 branch). Adapts miles.utils.replay_base imports to sirl.utils.routing_replay. Addresses bugs B/C/D/F/G/H: - Bug B (set_is_mtp chain): Router.is_mtp=False base attr; TopKRouter._init_routing_mode() lazy hash-routing init; TopKRouter.set_is_mtp() removes tid2eid for MTP layers; TopKRouter.set_layer_number() override calls _init_routing_mode; BaseMoELayer.set_is_mtp() delegates to router; MultiTokenPredictionLayer calls mlp.set_is_mtp() on construction. - Bug C (is_mtp guard): topk_routing_with_score_function gains is_mtp, tid2eid, input_ids params; MTP layers bypass routing replay; hash-routed layers bypass softmax routing. - Bug D (double-registration): MTP layers skip routing_replay via is_mtp bypass in topk_routing_with_score_function. - Bug F (padding_mask threading): MoELayer.route() and forward() accept padding_mask and input_ids; threaded through custom_forward and recompute paths. - Bug G (tid2eid dtype): torch.int32 for tid2eid parameter (was default torch.long), matches DSV4 hash-routing lookup contract. - Bug H (SwiGLU stride + moe_latent_size): fc1_stride=2 for gated_linear_unit (stride=1 for kitchen); moe_latent_size support for routed expert input/output dimensions. Gate A (py_compile) passed for all 5 files. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…fix mlp fc2 latent dim, restore padding_mask parity Fixes found by static tar parity audit (course-correction-20260508): 1. router.py: Remove duplicate _init_routing_mode (lines 225-262) that was shadowing the correct DSV4-aware version (uses torch.int32, checks is_mtp, handles dsv4_mode). The duplicate used torch.long and lacked is_mtp/dsv4_mode guards. 2. mlp.py: Fix fc2 latent size dimension order. Was: (ffn_hidden_size if not use_latent_size else moe_latent_size, hidden_size) Correct: (ffn_hidden_size, hidden_size if not use_latent_size else moe_latent_size) The input to fc2 is always ffn_hidden_size; moe_latent_size replaces the output dim. 3. moe_utils.py: Restore padding_mask parity from verified tar: - switch_load_balancing_loss_func: add padding_mask param + masking logic - z_loss_func: add padding_mask param + valid token count computation - get_tokens_per_expert_and_token_count: add missing function (was imported but absent) - compute_routing_scores_for_aux_loss: add padding_mask param + valid mask application Contract tests pass: pytest tests/infra/test_megatron_swiglu_mlp_contract.py (3/3) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
MLATransformerConfig (which extends TransformerConfig) lacked the moe_latent_size field introduced in the verified tar for DSV4 latent projections. mlp.py references self.config.moe_latent_size causing AttributeError at runtime. Adding field with default None — DSV4 (slime) does not use latent projections on experts, so use_latent_size always evaluates False. Fixes cluster crash: AttributeError: 'MLATransformerConfig' has no attribute 'moe_latent_size' when instantiating SharedExpertMLP
…it__ _init_routing_mode checks 'self.local_tokens_per_expert is None' but this attribute is only set inside _init_routing_mode itself. When called via set_layer_number (deferred init path for DSV4 mode), the attribute never exists, causing AttributeError. Fix: initialize local_tokens_per_expert=None and expert_bias=None in __init__ before any _init_routing_mode call.
…register_buffer pattern
Previous fix set local_tokens_per_expert=None as plain Python attribute,
causing register_buffer() to reject it ('attribute already exists').
Correct fix:
- Remove plain None assignment from __init__
- In _init_routing_mode: use 'not hasattr' guard for register_buffer call
- In the not-enable_expert_bias branch: use register_buffer(name, None)
so the attribute is a proper buffer, not a plain attribute
This ensures PyTorch module semantics are correct for checkpoint save/load.
The early-return block returned probs with shape [num_tokens, topk] while routing_map had shape [num_tokens, num_experts]. This caused RuntimeError in moe_utils.permute() when masked_select compared [num_experts, N] vs [topk, N] tensors. Root cause: the sqrtsoftplus branch already handles tid2eid inline and falls through to scatter which produces both outputs [N, num_experts]. The early-return was a duplicate that short-circuited before scatter. Fix: remove early-return; let sqrtsoftplus handle tid2eid as in tar. Parity: matches verified tar (yueming-yuan fork, 4a6aef1a4) exactly. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
get_routing_replay_compute_topk() reads ROUTING_REPLAY_STAGE from os.environ, which is never set — the actor uses manager.stage (object attribute). Fix: use routing_replay_manager.get_topk_fn(compute_topk, return_probs=True) which reads manager.stage directly. Triggered by Gate D fp8-no12-canary: ENABLE_ROUTING_REPLAY=1 injected but ROUTING_REPLAY_STAGE env var absent → RuntimeError on first train forward. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Branches
Preflight / validation
Review note
This PR is the Megatron vendor absorption path for the slime-trainer serving-aligned QAT lane. Please review before updating slime-trainer mainline submodule pointers.