DeepSeek-V4 Model Support by wenxie-amd · Pull Request #774 · AMD-AGI/Primus

wenxie-amd · 2026-06-22T18:18:05Z

No description provided.

Initial design / planning materials for integrating DeepSeek-V4 training support into Primus. Documentation only; no production code changes. - techblog/: architecture deep dive (CSA / HCA / mHC / Hash routing / sqrtsoftplus / clamped SwiGLU / dual RoPE / Muon / MTP) plus 4 PNG diagrams rendered via Pillow (see render_diagrams.py). - plan/: 8-phase roadmap, full code-landing list, per-phase task breakdown, and testing strategy. - progress/status.md: 64-task checklist tracking phase progress. - develop_deepseek-v4-in-primus.md: top-level goal and development cadence. Made-with: Cursor

Phase 1 of the V4 development plan. Pure config; no Python code paths exercised yet. Subsequent phases (P2..P4) wire dispatch and modules. * primus/configs/models/megatron/deepseek_v4_base.yaml Extends llama_base, sets model_type=deepseek_v4 and registers V4-specific defaults (hc_mult, hybrid_attention_*, q_lora_rank, attn_sink, hash routing, swiglu_limit, dual-RoPE knobs, etc.). * primus/configs/models/megatron/deepseek_v4_flash.yaml Hyperparams from DeepSeek-V4-Flash/config.json. * primus/configs/models/megatron/deepseek_v4_pro.yaml Hyperparams from DeepSeek-V4-Pro/config.json. * examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml Training scaffold; parallelism / perf knobs are conservative and will be retuned during the perf phase. * primus/backends/megatron/training/tokenizer/tokenizer.py Add DeepSeekV4Tokenizer to CUSTOM_TOKENIZER_TYPES so _add_tokenizer_args accepts it. Note: V4 fields do not need to be registered in Megatron's argparse — Primus's merge_namespace mechanism (train_runtime.py:_initialize_trainer) copies yaml-only fields onto backend_args after MegatronArgBuilder.update. Made-with: Cursor

Phase 2 of the V4 development plan. Wires the end-to-end dispatch from yaml.model_type=deepseek_v4 to a primus-owned model_provider + builder, without changing model behaviour yet. The model class is still a thin GPTModel subclass; Phase 3 swaps the decoder for the V4 transformer block. * primus/core/utils/import_utils.py Add a deepseek_v4 branch to get_model_provider() that imports primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_builders and returns partial(model_provider, deepseek_v4_builder). * primus/backends/megatron/megatron_pretrain_trainer.py Add a model_type == "deepseek_v4" branch alongside gpt / mamba. V4 is a causal-LM with the same data shape as GPT, so we reuse pretrain_gpt's forward_step + train_valid_test_datasets_provider; only the model_provider itself is V4-specific. * primus/backends/megatron/core/models/deepseek_v4/__init__.py (new) Re-export DeepseekV4Model + deepseek_v4_builder + model_provider. * primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_model.py (new) DeepseekV4Model: thin subclass of GPTModel. P3 will replace self.decoder with DeepseekV4TransformerBlock. * primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_builders.py (new) deepseek_v4_builder + model_provider. Uses GPT layer specs in P2; P3 will swap them for V4 specs. Made-with: Cursor

Phase 3 of the V4 development plan. Lands the V4 layer-spec helpers and a transparent V4 transformer-block subclass; attention / MLP behaviour still matches GPT. Phase 4 will plug HC + hybrid attention into the block, and Phase 5 will swap in V4 MoE / clamped SwiGLU through the spec-resolution hooks added here. * primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_layer_specs.py (new) Four V4 layer-spec helpers (layer / decoder_block / decoder_layer_specs / mtp_block) that delegate to the GPT helpers in P3, plus two resolution hooks (_resolve_attention_module_spec / _resolve_mlp_module_spec) that return None for now -- P4 / P5 fill these in. * primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_block.py (new) DeepseekV4TransformerBlock: subclasses TransformerBlock and stashes V4 config fields (hc_mult, compress_ratios, attn_sliding_window, attn_sink, q_lora_rank, index_*) onto self so P4 patches don't have to re-walk the config. Forward behaviour unchanged in P3. * primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_model.py Override __init__: after super().__init__() builds the stock decoder, swap self.decoder for DeepseekV4TransformerBlock (same call signature so GPTModel.forward keeps working). * primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_builders.py _resolve_layer_spec / _resolve_mtp_block_spec now route through the V4 layer-spec helpers instead of the GPT helpers directly. * primus/backends/megatron/core/models/deepseek_v4/__init__.py Re-export DeepseekV4TransformerBlock alongside the existing surface. Made-with: Cursor

…dual-RoPE) Phase 4 of the V4 development plan. Lands the full V4 transformer block: mHC multi-stream residual, per-layer hybrid attention dispatch (Dense / HCA / CSA), sliding-window mask, attention sink, dual-RoPE with YaRN. The V4 block becomes a standalone nn.Module that bypasses Megatron's TransformerBlock + ModuleSpec mechanism so the multi-stream HC loop is expressed cleanly. P5 will swap the placeholder SwiGLU MLP for V4's MoE. New modules under primus/backends/megatron/core/transformer/ :: * hyper_connection.py HyperMixer (per-layer mHC mixer), HyperHead (final K->1 collapse), sinkhorn_normalize (doubly-stochastic projection). Linear weights / scales / biases held in fp32 for stability; fp32 sinkhorn iterates. Unit-tested: row/col errors ~1e-6, hc_mult=1 degenerate path exact. * compressor.py V4 compressor for KV downsampling. ratio=4 overlap mode (CSA, coff=2), ratio=128 non-overlap mode (HCA, coff=1). Internal RMSNorm + learnable APE; RoPE applied externally. * indexer.py Sparse top-K position selector for CSA. Internal mini-Compressor builds the score grid; causal mask + top-K (-1 fill for invalid positions); backward propagates to the indexer params. * sliding_window_kv.py Causal SWA mask + per-query KV index helpers. * attn_sink.py Per-head learnable sink scalar; softmax_with_sink ensures probs.sum() <= 1 with the sink absorbing the residual mass. Backward propagates to the sink params. * dual_rope.py Two RoPE bases (main + compress) with optional YaRN scaling. Partial interleaved RoPE: only ``rotary_dim`` of each head's channels rotated; remaining channels passed through unchanged. * deepseek_v4_attention.py Shared base for V4 attention: QKV projection (optional Q LoRA), partial dual-RoPE, SWA mask, attention sink, output projection. ``_extra_kv`` hook lets HCA / CSA augment KV (full pool or sparse top-K). * hca_attention.py Heavily-Compressed Attention. Subclasses DeepseekV4Attention; adds a non-overlap Compressor and concatenates the full compressed pool to the local KV (always visible). * csa_attention.py Compressed-Sparse Attention. Subclasses DeepseekV4Attention; adds an overlap Compressor + Indexer; per-query attention is computed over the local SWA + the indexer's top-K compressed positions. Updated: * primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_block.py Rewritten as a standalone nn.Module. Holds the dual-RoPE for the whole stack, builds DeepseekV4HybridLayer per layer (Dense/HCA/CSA picked from compress_ratios), and runs the K-stream HC loop. Forward shape: [S, B, D] -> [B, S, D] -> [B, S, K, D] -> ... -> [B, S, D] -> [S, B, D]. Smoke-tested: 8-layer mixed dense/CSA/HCA + hc_mult=4 forward / backward / causality OK. * primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_layer_specs.py Cleaned up to a placeholder spec. The V4 block is standalone and bypasses Megatron's spec mechanism; we still hand a valid GPT-shaped spec to GPTModel.__init__ until P6 refactors that allocation away. * primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_model.py Docstring rewritten for the P4 standalone-block layout; pg_collection switched to getattr(self, "pg_collection", None) for safety. * deepseek-v4/develop/progress/status.md, plan/02-phase-details.md Track P1..P4 completion; add the argparse-not-needed note (Primus's merge_namespace covers V4 fields). Made-with: Cursor

…+ MTP Phase 5 of the V4 development plan. Lands the FFN side of the V4 stack: hash-routed and learned top-K MoE, clamped SwiGLU experts, and the V4 MTP head. The V4 block now plugs the V4 MoE in place of P4's placeholder SwiGLU FFN; the V4 model instantiates a separate-HyperHead MTP block when mtp_num_layers > 0. Layer-aware YaRN was already done in P4 (DualRoPE.get_rope picks main_rope vs compress_rope by compress_ratio). New modules: * primus/backends/megatron/core/transformer/clamped_swiglu.py clamped_swiglu(x, alpha=7.0): silu(gate)*up clamped to [-alpha, alpha]. ClampedSwiGLUMLP wraps it as a fused gate_up + down two-linear MLP. Eager (Python) for v1; perf phase will register a fused kernel. * primus/backends/megatron/core/transformer/moe/v4_hash_router.py HashRouter: static [vocab_size, topk] tid2eid table from a fixed seed. Active for the first num_hash_layers V4 layers; gives each token a permanent expert assignment with uniform weight 1/topk. No learnable parameters; deterministic across PP / TP / EP ranks. * primus/backends/megatron/core/transformer/moe/v4_topk_router.py V4TopKRouter: learned gate with score_function in {"sqrtsoftplus", "sigmoid", "softmax"}. Top-K with optional renorm and optional noaux_tc per-expert bias (selection-only; probs are read from the un-biased score). * primus/backends/megatron/core/transformer/moe/v4_moe.py DeepseekV4MoE: per-layer router pick (hash vs learned) + N ClampedSwiGLUMLP routed experts + 1 shared expert. Pure-PyTorch per-expert dispatch; P6 swaps in Megatron's token-dispatcher / grouped-GEMM / EP path. * primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_mtp.py DeepseekV4MTPBlock: mtp_num_layers V4 layers, each owning its own HyperHead (separate from the main decoder's). Shares the dual-RoPE with the main decoder. Loss-side wiring is deferred to P6; P5 just stands the module up so it can be unit-tested standalone. Updated: * primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_block.py DeepseekV4HybridLayer now picks MoE vs dense FFN based on num_routed_experts. forward() threads token_ids through to the MoE for hash-routed layers. The block-level forward picks token_ids up from a model-side stash (_v4_token_ids) so callers don't have to thread it explicitly through every layer of the call stack. * primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_model.py Builds DeepseekV4MTPBlock when mtp_num_layers > 0 (post-process rank only). forward() overridden to stash input_ids onto self.decoder before delegating to GPTModel.forward, so hash-routed MoE layers can consume them. Cross-PP propagation of input_ids is a P6 concern. * primus/backends/megatron/core/models/deepseek_v4/__init__.py Re-export DeepseekV4MTPBlock alongside the existing surface. Smoke-tested on dev-box PyTorch container (CPU, 7-test suite): * clamped_swiglu: clamp tight; MLP forward+backward OK. * HashRouter: per-token top-K distinct, deterministic across re-runs and re-instantiations w/ same seed, probs sum to 1. * V4TopKRouter: top-K honored, renorm OK, backward OK for all three score functions (sqrtsoftplus, sigmoid, softmax). * DeepseekV4MoE (learned & hash modes): forward + backward; same-token determinism for hash routing. * DeepseekV4TransformerBlock with MoE FFN (4 layers, hc_mult=2, mixed dense + CSA): forward + backward; deterministic in eval mode. * DeepseekV4MTPBlock (mtp_num_layers=2, hc_mult=2): forward + backward; per-MTP HyperHead state_dict separation verified. Deferred to P6 (already noted in progress doc): * Real Megatron-MoE / token-dispatcher / EP integration -- replaces the pure-PyTorch dispatch loop in DeepseekV4MoE.forward. * MTP loss path wiring -- DeepseekV4Model.forward currently builds the MTP block but does not yet feed its outputs through lm_head + the auxiliary loss term. * Numerical alignment vs reference inference/model.py (token-0 logits within 1e-2) -- needs reference checkpoint loading. Made-with: Cursor

Wire DeepSeek-V4 through Megatron P6 integration (PP local-layer build, EP expert sharding, and compatibility fixes) and add the P7 single-node launcher plus progress docs after passing PP=2/EP=4 smoke run. Made-with: Cursor

Add the plan-1 roadmap/detail/test documentation plus progress tracker entries, and update the development target doc with TransformerEngine and Primus-Turbo reference pointers. Made-with: Cursor

Remove GPT placeholder/super-init spec coupling so DeepSeek-V4 builds decoder directly from DeepSeek ModuleSpec submodule trees, and update Phase 8 progress records to match the finalized implementation and validation status. Made-with: Cursor

Unify DeepSeek-V4 runtime module selection under DeepSeekV4SpecProvider and migrate attention/MLP/MoE construction to provider-driven ModuleSpec flows with safe local fallbacks. Document and validate the TE CUDA runtime contract, including an explicit fail-fast guard for non-CUDA TE/Turbo inputs and updated Phase 9 progress records in English. Made-with: Cursor

…chema Align phase10 DeepSeek-V4 modules on explicit spec/provider contracts by enforcing SharedExpertMLP-only shared experts and introducing a dedicated DeepSeekV4TransformerConfig for V4-only runtime fields. Update builder/spec/docs so training resolves the new config type and tracks activation clamp through model config. Made-with: Cursor

Fix HC/attention dtype mismatches and tune the DeepSeek-V4 smoke script defaults so the Phase 10 MI355X run completes reliably end-to-end. Add a dedicated Phase 10 convergence report documenting delivered scope, runtime blockers, and remaining tracked items. Made-with: Cursor

… visuals Closes Phase 12 (plan-2 lockdown) on the engineering side. The current modules on this branch run, but a code review against real DeepSeek-V4 (HF reference, NeMo port, official inference) and Megatron's spec/config/provider/submodule/ build_module pattern surfaced 28 findings (10 CRIT / 11 HIGH / 6 MED / 5 LOW), so plan-1 P11 (validation/release) is paused and a new architecture-faithful rewrite (plan-2) is opened. Changes - rename develop/plan/ -> develop/plan-0/ (the original bring-up plan) - add develop/plan-2/ (README + 00 review-findings + 01 roadmap + 02 target-architecture + 03 phase-details + 04 test-strategy) - add develop/techblog/02-plan-1-as-built-and-plan-2-pointer.md (closes plan-0/1 with as-built notes + pointers to plan-2 docs / timeline / PPT) - update develop/techblog/README.md (banner: plan-2 active plan of record) - update develop/progress/status.md (Phase 12-21 tracking section; the only remaining P12 item is the external stakeholder sign-off on plan-2 scope) - add develop/progress/timeline.html (system fonts; daily-column Gantt with May 02-05 holiday band; P13-P21 packed into May 06-09) - add develop/progress/build_roadmap_pptx.py (generator) + deepseek_v4_roadmap_v1.pptx (13 slides; black-bg tech style; slide 7 "07 - 开发计划" is the day-by-day schedule with 3-row layout date / phase A~B / work content + holiday-gap arrow) Status of plan-2 phases - P12 lockdown: this commit. Only stakeholder sign-off pending. - P13-P21: planned for May 06 -> May 09 (4 working days; May 02-05 is holiday). Made-with: Cursor

…hful dense path) Plan-2 P13 first commit: rewrite the dense (`compress_ratio == 0`) path of DeepSeek-V4 attention to be faithful to the released V4-Flash checkpoint and rooted on Megatron's `MLASelfAttention`. What changed in `primus/backends/megatron/core/transformer/deepseek_v4_attention.py`: - New `DeepseekV4Attention(MLASelfAttention)` subclasses MLA for type identity but bypasses the parent `__init__` chain because V4's KV layout differs from MLA's compressed-KV form. - Single-latent KV: one `linear_kv` projection (`hidden -> head_dim`) feeds both K and V, broadcast across all query heads (matches `inference/model.py:Attention.wkv`). - Per-head `q_rms`: parameter-less RMS on `head_dim` after `linear_q_up_proj` and before partial RoPE (no `q_rms.weight` in the released checkpoint). - Grouped low-rank O: einsum-based `linear_o_a` per group + `linear_o_b` when `config.o_lora_rank > 0`. Falls back to MLA-style flat `linear_proj` when `o_lora_rank == 0`. - Learnable `attn_sink`: direct `nn.Parameter` on the attention (matches the released key `layers.{i}.attn.attn_sink` exactly), with inline softmax-with-sink in `_attention_forward`. - New `DeepseekV4AttentionSubmodules` dataclass with MLA-canonical names (`linear_q_down_proj`, `linear_q_up_proj`, `q_layernorm`, `kv_layernorm`) plus V4 extras (`linear_kv`, `linear_o_a`, `linear_o_b`, `attn_sink`). - `_LegacyDeepseekV4Attention` + `_LegacyDeepseekV4AttentionSubmodules` retained for CSA / HCA inheritance until the compressor / indexer fold-in lands in the P13 follow-up commit. Wiring: - `csa_attention.py` / `hca_attention.py`: switched parent to `_LegacyDeepseekV4Attention`. - `deepseek_v4_layer_specs.py`: route `compress_ratio == 0` to the new class with V4-canonical submodules; `{4, 128}` continue on the legacy path with legacy submodules. - `transformer_engine_spec_provider.py`: added `v4_q_layernorm()`, `v4_kv_layernorm()`, `v4_attention_sink()` factory methods. - `deepseek_v4_block._build_attention` (no-spec fallback) now points to `_LegacyDeepseekV4Attention` for compress_ratio == 0 so legacy CPU unit tests / configs that pre-date the rewrite keep working. - `DeepSeekV4TransformerConfig`: added `o_groups: int = 8` and `o_lora_rank: int = 0` (already exposed in the V4-Flash YAML). Tests (`tests/unit_tests/megatron/transformer/deepseek_v4/`): - State-dict-key contract: V4-canonical keys are present, legacy plan-1 keys (`q_a`/`q_b`/`k_proj`/`v_proj`/`o_proj`) are absent. - Forward shape + finiteness check. - Numerical equivalence vs an inline V4 reference forward (single-latent KV, partial interleaved RoPE, attn-sink as virtual key column, grouped low-rank O), with attn_sink enabled and disabled, gates <=1e-3. - Per-head q_rms is parameter-less (no leaked `q_rms.weight`). - `o_lora_rank == 0` fallback path coverage. - Rejection paths for `compress_ratio != 0` and `q_lora_rank == 0`. Deferred to P13 follow-up commit (still in May-06 budget): - Compressor / Indexer folded into `DeepseekV4Attention.forward` as spec submodules, retiring `csa_attention.py` / `hca_attention.py`. - Switch `linear_q_up_proj` from `parallel_mode='duplicated'` to `ColumnParallelLinear`, `linear_o_b` to `RowParallelLinear`. - TP=2 sharding parity test. - HF-reference numerical alignment (waits for P17 state-dict adapter). Made-with: Cursor

…s; retire CSA/HCA legacy Plan-2 P13 follow-up commit (closes P13). Builds on cad0fb3 to land the compressed-branch attention as spec submodules of the V4-faithful class, switch the attention's tensor-parallel-sensitive projections from `parallel_mode="duplicated"` to ColumnParallel / RowParallel, and retire the plan-1 legacy attention classes. What changed in `deepseek_v4_attention.py`: - `DeepseekV4Attention.__init__` now accepts `compress_ratio in {0, 4, 128}` (up from `0` only). When `compress_ratio > 0` the class builds `self.compressor` from `submodules.compressor` (overlap mode for ratio 4, non-overlap for ratio 128); when `compress_ratio == 4` it also builds `self.indexer` from `submodules.indexer`. - `DeepseekV4AttentionSubmodules` extended with `compressor` and `indexer` fields. - `DeepseekV4Attention.forward` dispatches on `self.compress_ratio`: * `0` — dense / SWA over local KV. * `128` — HCA: compute compressed pool with compress-base partial RoPE on indices `[0..P)`, broadcast to `H` heads, concat to local KV with a compressed-causal mask, share the joint softmax-with-sink. * `4` — CSA: per-query top-K from compressed pool via Indexer + Compressor, joint softmax-with-sink across local + sparse keys so the optional `attn_sink` is shared (not separate softmaxes). - `_LegacyDeepseekV4Attention` and `_LegacyDeepseekV4AttentionSubmodules` removed — the new class covers all three layer types. - Helper modules (`_per_head_rms_norm`, `_build_local_rms_norm`, `_append_sink_softmax`) are kept private to this file. `csa_attention.py` and `hca_attention.py` deleted. Wiring: - `deepseek_v4_layer_specs.py`: `_build_legacy_attention_submodules` is gone; `_build_v4_attention_submodules` now also builds `compressor` / `indexer` `ModuleSpec`s for compressed branches and uses: * `provider.column_parallel_linear()` (`gather_output=True`) for `linear_q_up_proj` so the q-up-proj weight is sharded across TP ranks at `tp > 1`. `gather_output=True` keeps downstream attention math TP-agnostic (full `H * head_dim` width). * `provider.row_parallel_linear()` (`input_is_parallel=False`) for `linear_o_b` (grouped) and `linear_proj` (flat-O fallback) so the output projection is sharded across TP ranks. * `parallel_mode="duplicated"` retained for `linear_q_down_proj`, `linear_kv`, `linear_o_a`. Full grouped-O TP plan is tracked in P14. - `deepseek_v4_block._build_attention` (no-spec fallback) constructs `DeepseekV4Attention` for all branches; the new class builds its own Compressor / Indexer locally when no spec is provided. Tests (`tests/unit_tests/megatron/transformer/deepseek_v4/test_deepseek_v4_attention.py`): - `test_unsupported_compress_ratio_rejected` (replaces the previous `compress_ratio != 0` rejection check) — only `{0, 4, 128}` accepted. - `test_hca_forward_shape_and_finite` — HCA forward produces `[B, S, hidden]` finite output, and `attn.compressor` is built while `attn.indexer` is None. - `test_hca_forward_matches_inline_reference` — HCA forward agrees (<=1e-3) with an inline reference that re-implements the same pooled-KV + compress-base partial RoPE + compressed-causal mask + joint softmax-with-sink + grouped-O math. - `test_csa_forward_shape_and_finite` — CSA path with overlap Compressor + Indexer top-K + per-query joint softmax produces a `[B, S, hidden]` finite output (numerical alignment with HF Indexer ranking is deferred to P17 alongside the state-dict adapter). - `test_attention_spec_uses_column_and_row_parallel` — asserts the spec helper sources `linear_q_up_proj` from the provider's ColumnParallel and `linear_o_b` / `linear_proj` from the provider's RowParallel, with the right `gather_output` / `input_is_parallel` flags. This is the contract that lets TP > 1 actually shard the weights at runtime. - `test_attention_spec_includes_compressor_and_indexer` — asserts compressed-branch submodules carry a Compressor `ModuleSpec` (always) and an Indexer `ModuleSpec` (CSA only). - `test_tp2_sharding_parity_scaffold` — skipif scaffold for the `torchrun --nproc_per_node=2` parity check; the bit-equality vs a duplicated baseline implementation is tracked in P19. Status: `deepseek-v4/develop/progress/status.md` updated to mark all P13 items done (including the previously deferred Compressor/Indexer fold-in, TP projections, and TP=2 scaffold) and to point to P14 for the full grouped-O TP plan. Made-with: Cursor

…(G3/G4) Plan-2 P14 phase-1: rewrite the activation, learned router, and hash router so the math, parameter layout, and state-dict keys match the released DeepSeek-V4-Flash reference exactly. Provider helpers and the MoELayer subclassing land in the P14 phase-2 follow-up. Activation (G3): - Replace post-multiplication clamp with the V4 pre-multiplication semantics: SiLU(clamp(gate, max=alpha)) * clamp(up, +/- alpha). New helpers clamped_swiglu_pre_mul (split inputs) and clamped_swiglu_pre_mul_fused ([gate | up] last-dim concat). - ClampedSwiGLUMLP now uses separate w1 / w2 / w3 Linears so the state_dict keys match HF Expert exactly (no fused gate_up.weight leak). Optional fused_gate_up forward fuses GEMMs at run time without changing the saved layout. - _DenseSwiGLUMLP in deepseek_v4_block.py applies the same pre-mul clamp (it previously did vanilla SwiGLU and ignored swiglu_limit). Learned router (G4): - Rename V4TopKRouter -> DeepseekV4LearnedRouter (alias retained). Gate is exposed as `weight` Parameter (matches Megatron's TopKRouter AND HF Gate.weight). - expert_bias is selection-only; routing weights gather from the un-biased scores so probs gradient flows to weight, not bias. - Renormalization is gated on score_function != softmax (matches HF; softmax already sums to 1). - topk_scaling_factor honors moe_router_topk_scaling_factor (= HF route_scale). - v4_score_fn covers softmax / sigmoid / sqrtsoftplus. Hash router (G4): - Rename HashRouter -> DeepseekV4HashRouter (alias retained). - Add learnable `weight` Parameter (same shape as the learned router); previously the hash router emitted uniform 1/topk weights. - tid2eid is a frozen nn.Parameter (requires_grad=False, dtype=int32) matching HF reference layout — preserves the table across state-dict round-trips without polluting the optimizer state. - forward(hidden, token_ids) gathers learned scores at the static expert ids prescribed by tid2eid; renorm + scale parity with the learned router. MoE wiring: - DeepseekV4MoE._route now passes (hidden, token_ids) to the hash router; both routers receive hidden_size / score_function / topk_scaling_factor at init. Tests: - tests/unit_tests/megatron/transformer/deepseek_v4/test_clamped_swiglu.py (7 tests): pre-mul activation vs HF reference (<= 1e-6 fp32, four alpha values), alpha=0 disables clamp, fused-vs-split agreement, one-sided gate clamp behavior, w1/w2/w3 state-dict keys, fused forward equivalence, end-to-end MLP vs HF Expert.forward. - tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_routers.py (13 tests): score-function parity; learned-router HF agreement across (softmax, sigmoid, sqrtsoftplus) x (with/without expert bias) <= 1e-6; back-compat alias; gradient flows to gate weight; expert_bias detached from probs graph; softmax skips renorm; hash router HF agreement across three score functions; tid2eid is a frozen Parameter; state-dict keys; deterministic table across seeds; OOB / shape-mismatch error paths; gradient flows to weight while tid2eid.grad is None. Status / progress: - deepseek-v4/develop/progress/status.md: marks P14 phase-1 tasks complete with this commit hash, lists deferred items for the phase-2 follow-up (MoELayer subclassing, provider helpers, G5 1L MoE forward), and resolves the "HashRouter has no learnable gate weight" / clamped-SwiGLU blockers in the risks log. Out of scope (lands in P14 phase-2): - DeepseekV4MoE(MoELayer) subclassing (load-balance / z-loss / dispatcher lifecycle inheritance). - Provider v4_grouped_mlp_spec(swiglu_limit) / v4_router_spec. - Threading token_ids as a forward kwarg through TransformerBlock -> TransformerLayer -> MoE (co-removed with P15's hybrid-layer refactor). - G5 1L MoE forward agreement vs HF reference within 1e-3 fp32. Made-with: Cursor

…gnment Plan-2 P14 phase-2: structurally bring DeepseekV4MoE in line with Megatron's spec lifecycle, expose CPU-testable forward path so the MoE math can be pinned against the released HF reference, and add provider helpers. DeepseekV4MoE -> MegatronModule: - Parent class switched from nn.Module to MegatronModule so it picks up the standard Megatron config plumbing and integrates cleanly with TransformerLayer.mlp via the spec lifecycle. - BaseMoELayer-compatible public surface: ``set_layer_number`` is now defined on the V4 MoE (mirrors BaseMoELayer.set_layer_number) and ``local_expert_indices`` is exposed as a list attribute. CPU local-experts path: - New ``local_experts: nn.ModuleList[ClampedSwiGLUMLP]`` and a single ``ClampedSwiGLUMLP`` shared expert are constructed when ``pg_collection`` is None. Each expert mirrors a single HF reference ``Expert`` (separate w1 / w2 / w3 Linears + V4 pre-multiplication clamp). - New ``_local_experts_forward`` runs the per-expert dispatch loop matching ``DeepSeek-V4-Flash/inference/model.py:MoE.forward`` exactly: for each routed expert i, gather tokens routed to i, multiply by the per-token routing weight, accumulate. Production path (``pg_collection`` provided) continues to use the Megatron dispatcher + grouped experts unchanged. Provider helpers (plan-2 P14 §5/§6): - ``DeepSeekV4SpecProvider.v4_grouped_mlp_spec(swiglu_limit, ...)`` returns a ready-to-use ``ModuleSpec(grouped_module, MLPSubmodules)`` for the V4 MoE expert path. The pre-mul clamp itself flows via ``config.activation_func_clamp_value`` -- Megatron's eager ``glu()`` already implements ``SiLU(clamp(gate, max=alpha)) * clamp(up, +/- alpha)`` which is bit-equal to the HF reference. - ``DeepSeekV4SpecProvider.v4_router_spec(learned=True/False)`` returns a bare ``ModuleSpec`` for either DeepseekV4LearnedRouter or DeepseekV4HashRouter. Both routers stay standalone nn.Module so they instantiate cleanly on CPU; aux-loss / z-loss / RouterReplay inheritance via TopKRouter subclassing is tracked into P19 (the upstream parent registers CUDA buffers in ``__init__``, which is impractical to use on CPU). G5 numerical alignment (tests/unit_tests/.../test_v4_moe.py, 11 tests): - Construction sanity: parent class is MegatronModule; CPU path builds local_experts (ClampedSwiGLUMLP) + shared_expert; the dispatcher / grouped_experts attributes stay None; set_layer_number propagates. - Learned-router MoE forward: agreement vs inline HF reference on a 1L toy across (sqrtsoftplus / sigmoid / softmax) x (shared expert on / off), <= 1e-3 fp32 CPU. - Hash-router MoE forward: same, with token_ids feeding tid2eid. - ``moe_router_topk_scaling_factor`` (HF ``route_scale``) propagates to the output. - Backward populates grads on router.weight, on the shared expert, and on at least one routed expert's w1 / w2 / w3. - Hash layer raises a clear error when ``token_ids`` is missing. Status: - deepseek-v4/develop/progress/status.md: P14 phase-2 tasks ticked with this commit; the structural row records the MegatronModule-via-CPU-path approach and explicitly defers the TopKRouter-rooted aux-loss path to P19 (with the rationale). Out of scope (lands later): - Threading token_ids as a forward kwarg through TransformerBlock -> TransformerLayer -> MoE (co-removed with P15's hybrid-layer refactor). - DeepseekV4LearnedRouter / DeepseekV4HashRouter subclassing TopKRouter (depends on the parent gating CUDA-buffer registration on a device check upstream; tracked into P19 alongside the distributed re-validation matrix). Made-with: Cursor

…oken-ids forward kwarg + HC x PP packing Plan-2 P15: bring V4's layer / block onto Megatron's standard ``TransformerLayer`` / ``TransformerBlock`` parents (type identity + spec lifecycle), drop the ``decoder._v4_token_ids`` attribute stash in favor of a real forward kwarg, gate ``HyperHead`` to the post_process stage, and extract HC x PP K-stream packing helpers. DeepseekV4HybridLayer -> TransformerLayer: - Parent class switched from GraphableMegatronModule to TransformerLayer. TransformerLayer.__init__ is bypassed (V4's submodule contract differs from upstream -- no cross-attention, no BDA, V4-specific attention signature); MegatronModule.__init__ is called directly. - DeepseekV4HybridLayerSubmodules now extends TransformerLayerSubmodules and uses upstream-canonical field names: input_layernorm / self_attention / pre_mlp_layernorm / mlp. ``attn_hc`` / ``ffn_hc`` are added as V4-specific HC mixer hooks (None when hc_mult == 1). - forward signature is upstream-compatible: (hidden_states, attention_mask=None, *, position_ids=None, token_ids=None, **kwargs). attention_mask is accepted and ignored (V4 manages SWA / sink mask internally); position_ids is consumed from the caller (fallback to arange(S) for tiny smokes); token_ids feeds hash-routed MoE layers. Accepting **kwargs lets the layer plug into MultiTokenPredictionLayer (P16) without bespoke adapters. DeepseekV4TransformerBlock -> TransformerBlock: - Parent class switched from nn.Module to TransformerBlock (init bypass via MegatronModule for CPU instantiability; V4 has its own layer-spec / lift-lower pipeline). Type identity unlocks Megatron isinstance checks + sharded-state-dict integration. - HyperHead is built only on the post_process stage. Earlier PP stages forward the K-stream tensor via _lower_streams_out (no per-stage HyperHead), saving memory and avoiding correctness drift. - forward consumes ``position_ids`` as a kwarg and threads ``token_ids`` through ``decoder.forward -> layer.forward -> mlp.forward -> hash_router.forward``. Removes the legacy ``decoder._v4_token_ids`` getattr fallback. HC x PP K-stream packing helpers: - _lift_streams_in(hidden_states, pre_process, hc_mult): first PP stage expands [S, B, D] -> [B, S, K, D]; non-first stage unfolds [S*K, B, D] -> [B, S, K, D]. Single-stream (hc_mult == 1) just transposes [S, B, D] -> [B, S, D]. - _lower_streams_out(x, post_process, hc_mult): final stage transposes the post-HyperHead [B, S, D] -> [S, B, D]; non-final stage packs [B, S, K, D] -> [S*K, B, D] so PP P2P kernels see a 3D tensor of the expected rank. Both helpers raise clear errors on shape mismatches. - The packing math is intentionally K-folded-into-seq (not the batch axis) so sequence-parallel chunking lines up cleanly; PP P2P doesn't need to know about K. DeepseekV4Model.forward: - Drops the ``decoder._v4_token_ids = input_ids`` stash (and the try/finally cleanup). - Passes ``token_ids=input_ids`` and ``position_ids=position_ids`` directly to ``self.decoder(...)``; the decoder block + each layer consume them as standard forward kwargs. - An AST-level audit (test_v4_block_pp.py) prevents the attribute stash from regressing. Spec wiring: - deepseek_v4_layer_specs.py renames the four core fields on DeepseekV4HybridLayerSubmodules to match the new dataclass: attn_norm -> input_layernorm, attention -> self_attention, ffn_norm -> pre_mlp_layernorm, ffn -> mlp. - DeepseekV4MTPBlock's per-MTP-layer call switches to layer(stream, position_ids=..., token_ids=...) (kwarg, not positional) to match the new layer forward signature. Tests (tests/.../test_v4_block_pp.py, 16 tests): - Subclass identity: DeepseekV4HybridLayer is a TransformerLayer; DeepseekV4TransformerBlock is a TransformerBlock; the hybrid-layer submodules dataclass extends TransformerLayerSubmodules and exposes attn_hc / ffn_hc. - Lift / lower roundtrip: bit-exact across the four PP-stage permutations (pre_process * post_process), for both single-stream (hc_mult=1) and multi-stream (K=3, K=4). - Error paths: misaligned S*K on non-first stage; collapsed input on non-final lower; uncollapsed input on final lower. - Token-ids stash: AST audit confirms decoder._v4_token_ids is gone from the model source; ``token_ids=input_ids`` kwarg is present. - Forward signatures: block.forward exposes position_ids + token_ids kwargs; layer.forward accepts (hidden_states, attention_mask=None, position_ids, token_ids). Status / blockers: - deepseek-v4/develop/progress/status.md: Phase 15 tasks ticked except G6 (PP=1 vs PP=2 vs PP=4 equivalence on a 4L toy), which requires distributed init and is tracked into P19 distributed re-validation. The CPU-only sub-gate -- _lift_streams_in after _lower_streams_out is bit-exact -- is covered by the lift/lower roundtrip tests, which is the math contract a real PP run depends on. - Two blocker rows resolved: "Custom V4 block / layer / MoE bypass upstream parents" (P14 phase-2 + P15) and "Token-IDs propagation via decoder._v4_token_ids attribute" (P15). Out of scope (lands later): - G6 distributed equivalence test -> P19. - mtp_use_repeated_layer / mtp_layer_pattern integration -> P16. Made-with: Cursor

…TokenPredictionBlock + process_mtp_loss Plan-2 P16 wires V4 onto Megatron's upstream MTP pipeline: ``MultiTokenPredictionBlock`` (per-depth eh_proj + V4 hybrid layer + RMSNorm) plus ``process_mtp_loss`` (per-depth shifted-logits aux loss folded into the LM-loss gradient). The legacy ``DeepseekV4MTPBlock`` is preserved behind ``v4_use_custom_mtp_block`` for back-compat with research checkpoints (planned removal: P21) and now emits a ``DeprecationWarning`` on construction. Spec helper (deepseek_v4_mtp_specs.py, new): - get_v4_mtp_block_spec(config, *, transformer_layer_spec, vp_stage) returns ``ModuleSpec(MultiTokenPredictionBlock, submodules= MultiTokenPredictionBlockSubmodules(layer_specs=[...]*mtp_num_layers))``. - Each per-depth ``MultiTokenPredictionLayer`` spec pulls ``enorm`` / ``hnorm`` / ``layer_norm`` from ``DeepSeekV4SpecProvider.v4_norm_module()`` and ``eh_proj`` from ``provider.column_parallel_linear()``. The inner ``mtp_model_layer`` is the V4 hybrid-layer spec passed in by the model -- so each MTP depth shares HC, hash routing, and clamped-SwiGLU with the main decoder exactly. - Rejects ``mtp_num_layers < 1`` with a clear ValueError. DeepseekV4Model updates (deepseek_v4_model.py): - New default path: when ``mtp_num_layers > 0`` and not ``v4_use_custom_mtp_block``, ``__init__`` builds ``self.mtp = MultiTokenPredictionBlock(spec=get_v4_mtp_block_spec(...))`` on stages where ``mtp_on_this_rank()`` is True. ``mtp_on_this_rank`` is wrapped in try/except so CPU smokes (no parallel_state) do not crash; ``self.mtp_process`` is False and ``self.mtp`` is None on those paths, leaving a forward-compatible inert model. - Legacy ``DeepseekV4MTPBlock`` path stays available behind ``v4_use_custom_mtp_block``; ``self.mtp_block`` is the legacy slot, ``self.mtp`` is the new spec-based slot. Both are None when MTP is disabled. - ``forward`` now mirrors ``GPTModel.forward``: runs ``self.mtp(...)`` on stages with MTP layers (passing ``input_ids`` / ``position_ids`` / ``hidden_states`` / ``attention_mask`` / ``embedding`` / ``packed_seq_params``), then on ``post_process`` with ``mtp_num_layers > 0`` calls ``process_mtp_loss(...)`` which chunks the concatenated hidden states, computes the per-depth shifted MTP loss, and folds it into the gradient via ``MTPLossAutoScaler``. - New forward kwargs: ``loss_mask`` (forwarded to ``process_mtp_loss``) and ``packed_seq_params``. Layer / block forward contract: - ``DeepseekV4HybridLayer.forward`` now returns ``(hidden_states, None)`` instead of just ``hidden_states``. This matches upstream ``TransformerLayer`` (which returns ``(hidden_states, context)``) and is required by ``MultiTokenPredictionLayer._proj_and_transformer_layer`` which unpacks ``hidden_states, _ = self.mtp_model_layer(...)``. - ``DeepseekV4TransformerBlock``'s per-layer iteration updates to ``x, _ = layer(...)`` to handle the new tuple return. - Legacy ``DeepseekV4MTPBlock`` likewise updates to unpack the tuple. V4 attention spec (deepseek_v4_layer_specs.py): - The V4 attention spec now declares ``params={"compress_ratio": ..., "attn_mask_type": AttnMaskType.causal}``. ``MultiTokenPredictionLayer.__init__`` validates the inner layer's ``self_attention.params['attn_mask_type']`` against ``{padding, causal, no_mask, padding_causal}``; without this the MTP block fails to construct. The value is functionally inert for V4 (V4 manages its own SWA / sink mask internally). - ``DeepseekV4Attention.__init__`` accepts and ignores ``attn_mask_type`` plus a ``**kwargs`` catch-all so the spec lifecycle keeps working. Legacy DeepseekV4MTPBlock (deepseek_v4_mtp.py): - Module docstring annotated as deprecated (planned removal: P21). - Construction emits a ``DeprecationWarning`` pointing users at ``get_v4_mtp_block_spec``. Code path unchanged otherwise. Package surface (__init__.py): - Exports ``DeepseekV4HybridLayer`` / ``DeepseekV4HybridLayerSubmodules`` / ``DeepseekV4TransformerBlockSubmodules`` / ``get_v4_mtp_block_spec`` alongside the existing surface. Tests (tests/.../test_v4_mtp.py, ~17 tests): - ``get_v4_mtp_block_spec`` structural assertions: outer module is ``MultiTokenPredictionBlock``; ``layer_specs`` length matches ``mtp_num_layers`` (parametrised 1/2/3); each per-depth spec is a ``MultiTokenPredictionLayer``; the V4 inner layer is threaded through unchanged; norm + linear come from the V4 provider. - Rejects ``mtp_num_layers=0`` with a clear ValueError. - ``DeepseekV4HybridLayerSubmodules`` extends ``TransformerLayerSubmodules`` so MTP picks up the GPT path (not Mamba) in its inner-layer-submodules isinstance check. - ``DeepseekV4HybridLayer.forward`` returns ``(hidden_states, None)`` (source-level assertion on ``return x, None``). - V4 attention spec advertises ``AttnMaskType.causal`` (source-level assertion). - Legacy ``DeepseekV4MTPBlock`` emits ``DeprecationWarning`` on construction. - AST audits on ``deepseek_v4_model.py``: ``process_mtp_loss`` is called; upstream MTP machinery is imported; spec helper is invoked; ``v4_use_custom_mtp_block`` flag is preserved; the ``mtp_num_layers > 0`` guard keeps the no-MTP path inert. Status / blockers: - deepseek-v4/develop/progress/status.md: Phase 16 tasks ticked except G7 (MTP loss appears in train log; ``mtp_num_layers=0`` vs ``mtp_num_layers=1`` ablation matches LM loss to 1e-6), which requires distributed init + MultiTokenPredictionBlock runtime (CP / SP plumbing); tracked into P19 distributed re-validation alongside G6. - Two new follow-on rows recorded for the cross-cutting layer-tuple return + attention attn_mask_type declarations (both required by upstream MTP wiring). Out of scope (lands later): - G7 distributed MTP-loss ablation -> P19. - MTP state-dict adapter (HyperHead per-depth weights) -> P17 alongside the V4-Flash safetensors load gate. - mtp_use_repeated_layer / mtp_layer_pattern tuning -> P19 / P20. Made-with: Cursor

…ose P17 for code cleanup Plan-2 reshuffle (user-driven, 2026-05-01): pre-training is the release path; HF-weight loading is not required for the release. Move the HF state-dict adapter + V4-Flash numerical-alignment gate (old P17 / part of old P20) to a deferred Phase 22+ section, and repurpose the P17 slot for the dead-code / hygiene work that previously sat in P21. Phase shape after this commit: P17 Code cleanup (was: state-dict adapter) - retire _RMSNorm duplicates / dual_rope.py / csa_attention.py / hca_attention.py / legacy DeepseekV4MTPBlock - drop EP all_reduce fallback gate + v4_use_custom_mtp_block flag - drop _v4_token_ids residue everywhere (front-loaded from old P18 task list) - fix yaml comments (4=CSA, 128=HCA) - new gate G14 (static dead-code audit) governs exit P18 Spec audit (unchanged; _v4_token_ids item removed) P19 Distributed re-validation (unchanged; G6 / G7 still here) P20 Convergence + perf gates (HF numerical-alignment row removed; convergence baseline switched to Megatron-bridge) P21 Docs + handover (slimmed; cleanup tasks moved to P17) P22+ HF state-dict adapter + V4-Flash checkpoint load (deferred; activate when SFT / evaluation needs HF weights) Files updated: - plan-2/01-roadmap.md phase table, dep graph, milestones, risks, out-of-scope - plan-2/03-phase-details.md P17 rewritten, P18 trimmed, P20 trimmed, P21 trimmed, P22+ section appended - plan-2/04-test-strategy.md G8 / G9 marked deferred to P22+; G12 retargeted at Megatron-bridge baseline; new G14 added; ownership table refreshed - plan-2/README.md highlights reflect deferred adapter - progress/status.md Phase 17 / Phase 20 / Phase 21 sections rewritten; new Phase 22+ section; blocker log records the reshuffle and marks the HF-load entry as deferred No code changes in this commit; docs / plan only. Made-with: Cursor

… dead-code audit (G14) Plan-2 P17 ships the dead-code retirement that was front-loaded from P21 in the 2026-05-01 reshuffle. Pre-training is the release path; the HF state-dict adapter slot moves out (deferred to P22+) and the cleanup work moves up so P18's spec audit walks a clean tree. Retired in this commit: primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_mtp.py - The legacy primus-owned DeepseekV4MTPBlock was deprecation-warned since P16 commit 6c5875d. The spec-based path (get_v4_mtp_block_spec + upstream MultiTokenPredictionBlock + process_mtp_loss) is the only MTP route now. DeepSeekV4TransformerConfig - v4_use_custom_mtp_block (legacy MTP gate) removed. - mtp_compress_ratios (legacy-only field) removed. DeepseekV4Model.__init__ - Single MTP branch on the spec path; the `if v4_use_custom_mtp_block` arm + self.mtp_block field gone. Dedup'd in this commit: primus/backends/megatron/core/transformer/local_rmsnorm.py (new) - One canonical LocalRMSNorm consumed by: * deepseek_v4_block.py (input_layernorm / pre_mlp_layernorm / final_layernorm fallback) * deepseek_v4_attention.py (q_norm / kv_norm fallback closure) * compressor.py (kv_norm) - The three pre-existing `_RMSNorm` definitions are deleted. YAML cleanup: primus/configs/models/megatron/deepseek_v4_flash.yaml - Inverted comment fixed: 4 = CSA (overlap) and 128 = HCA (non- overlap) match DeepseekV4Attention.forward dispatch. deepseek_v4_pro.yaml + deepseek_v4_base.yaml - Same canonical comment block added so all three V4 yamls are self-documenting. Audit (gate G14): tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p17_dead_code.py (new) - Asserts the retired files are gone (deepseek_v4_mtp.py / csa_attention.py / hca_attention.py). - Asserts the legacy import path raises ImportError. - Asserts the V4 config no longer carries v4_use_custom_mtp_block / mtp_compress_ratios. - Asserts the package __all__ no longer exposes DeepseekV4MTPBlock. - AST scans every V4 source for runtime `_v4_token_ids` access (Attribute / Assign / Name) — docstring mentions are exempt. - AST scans every V4 source for `class _RMSNorm` shadow definitions. - Parameterizes over the three V4 yamls and asserts the canonical `4 = CSA` / `128 = HCA` comment is present. Out of scope for P17 (retained, with notes in status.md): primus/backends/megatron/core/transformer/dual_rope.py — load-bearing for V4's CSA / HCA dual-base partial RoPE; Megatron's RotaryEmbedding only supports a single base. Plan-2 was over-eager listing this for retirement; it stays. Other touched files: tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_mtp.py - Imports / fixtures dropped (legacy block is gone). - Replaces deprecation-warning + flag-preservation tests with "module is gone / package surface drops it / config fields gone / model.__init__ no longer references the legacy class" assertions. deepseek-v4/develop/progress/status.md - Phase 17 task table all checked off; blocker log row for the EP all_reduce fallback marked resolved (flag was already gone). Made-with: Cursor

…stency + compress_ratios normalization (G1) Plan-2 P18 closes the spec-system audit findings D1 / D2 / D4 from the plan-2 review (00-review-findings.md): D1 (HIGH) — DeepSeekV4SpecProvider re-instantiated inside the block builder, the layer-spec factory, and the MTP spec helper. D2 (HIGH) — provider.activation_func() returned the TEActivationOp class, which Megatron MLP only honors when config.use_te_activation_func=True; the default V4 yaml path silently dropped it. D4 (MED) — compress_ratios stored as a YAML string and ast.literal_eval'd on every consumer call; normalization belonged in the dataclass. Provider singleton (D1): primus/backends/megatron/core/models/deepseek_v4/build_context.py (new) - resolve_v4_provider(config) caches a single DeepSeekV4SpecProvider on the config object via a private attribute. Different configs get different providers; the cache is GC'd when the config is released. - reset_v4_provider_cache(config) helper for unit tests that need a fresh provider. All three direct DeepSeekV4SpecProvider(config=config) call sites migrated to the helper: - deepseek_v4_block.py (_build_projection + DeepseekV4MoE shared- expert wiring) - deepseek_v4_layer_specs.py - deepseek_v4_mtp_specs.py AST audit (test_v4_p18_spec_audit.py:: test_no_direct_DeepSeekV4SpecProvider_construction_outside_build_context) rejects future regressions; build_context.py is the only allowed instantiation site. Activation-func consistency (D2): DeepSeekV4SpecProvider.v4_mlp_activation_func() - Returns None when config.use_te_activation_func is False (V4 default - needed so Megatron MLP keeps the eager clamped-SwiGLU path that applies activation_func_clamp_value). - Returns the TEActivationOp class when the user opts in. Layer specs + DeepseekV4MoE shared-expert spec switched to the V4 helper. The base provider's activation_func() is unchanged (BackendSpecProvider contract still says "returns a type"). AST audit `test_v4_specs_use_v4_mlp_activation_func_helper` rejects spec builders that fall back to the unconditional provider.activation_func(). compress_ratios normalization (D4): DeepSeekV4TransformerConfig.__post_init__ - Calls _normalize_compress_ratios_field on the raw value once, so downstream consumers see tuple[int, ...] (or None). - Helper handles strings ("[0, 0, 4, 128, ...]") and real lists. Runtime helpers (_parse_int_sequence / _normalize_compress_ratios in deepseek_v4_block.py) keep accepting both forms for back-compat, but now always receive the normalized form on the live path. Schema gate (G1): tests/unit_tests/configs/test_deepseek_v4_yaml.py (new) - Parameterises over deepseek_v4_{base,flash,pro}.yaml. - parse_yaml() succeeds; required fields present. - DeepSeekV4TransformerConfig builds from the parsed dict. - compress_ratios normalized to tuple[int, ...] with no value drift vs the raw schedule. - Every compress_ratios entry is in {0, 4, 128} (canonical V4 branches; matches the deepseek_v4_attention.py dispatch). - Retired P17 fields (v4_use_custom_mtp_block / mtp_compress_ratios) are gone from the dataclass AND from each YAML. - V4-specific runtime fields (HC, sliding-window, sink, o_groups / o_lora_rank, MoE extras, swiglu_limit) all declared on the dataclass - removing one silently breaks the runtime. - Provider singleton: resolve_v4_provider(cfg_a) returns the same instance on repeated calls; different configs get different providers. - v4_mlp_activation_func contract verified for both branches of use_te_activation_func. Spec audit (light-weight, AST-only): tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p18_spec_audit.py (new) - D1 / D2 audits described above. - Package surface: __init__.py __all__ does not re-export DeepseekV4MTPBlock (P17 cross-check). - Spec builders do not eagerly construct TENorm / TE{Column,Row}ParallelLinear / TELinear / TEActivationOp inside __init__ - they emit ModuleSpec(module=...) references that runtime build_module resolves. Made-with: Cursor

…e hashes Replace `(this commit)` / `(working tree)` / `(P17 commit)` / `(P18 commit)` / `(review)` placeholders in `deepseek-v4/develop/progress/status.md` with the concrete commit hashes that landed each row, mapped per phase section: - P5 -> 5e4008d - P6/P7 -> 97b9720 - P8 v2 -> df273a4 - P9 v2 -> e5fec96 - P10 v2 -> b38e83c (with 752b753 for the clamped-SwiGLU follow-up) - P12 -> 636ab3d - P13 phase-2 (Compressor/Indexer fold + TP shard) -> aa9929a - P14 phase-1 (clamped SwiGLU + routers) -> 1a8bf32 - P14 phase-2 (MoE structural + provider helpers + G5) -> 5fe8bc3 - P15 -> 25ccdb5 - P16 -> 6c5875d - P17 -> e591b89 - P18 -> b583267 The legacy DeepseekV4MTPBlock retirement row now spans the deprecation in P16 (6c5875d) -> deletion in P17 (e591b89). The P13 phase-2 HF-reference 1L attention task note is updated to reflect the 2026-05-01 reshuffle that deferred HF numerical alignment to P22+. Made-with: Cursor

… broadcast for 1F1B / VPP (G10) Lands two primus-patches that close the Phase 19 distributed re-validation gates for V4 under 1F1B and interleaved-1F1B / VPP: * megatron.deepseek_v4.pp_tensor_shape - wraps schedules.get_tensor_shapes (1F1B path) so the seq dim reflects mHC's K-stream packing (`[S * hc_mult, B, hidden]`). - additionally wraps forward_backward_pipelining_with_interleaving (VPP path) and scales its `seq_length` kwarg by `hc_mult` so the schedule's inline `tensor_shape` matches what `_lower_streams_out` emits. Without the second wrap, VPP recv buffers are `[S, B, D]` while sender emits `[S*K, B, D]`; PyTorch P2P silently truncates and `_lift_streams_in` reshapes the truncated copy, surfacing as `DeepseekV4HashRouter: hidden=32 vs token_ids=128`. * megatron.deepseek_v4.pp_token_pre_broadcast - V4's hash-routed MoE layers need raw input_ids on every PP stage that owns one, but `pretrain_gpt.get_batch` returns `None` on middle PP stages. Two earlier hooks (in-`forward` and per-call `get_batch`) deadlocked the interleaved schedule because the pre-warmup `recv_forward.wait()` on PP rank > 0 blocks before PP rank 0 ever issues its first send, so the matching broadcast never pairs up. - This patch instead pre-broadcasts all microbatch / chunk tokens upfront from PP rank 0 across the PP group inside a wrapper around `get_forward_backward_func`, before any send/recv runs, and caches the resulting (tokens, ...) tuples per (vp_stage, microbatch). A companion wrapper around `pretrain_gpt.get_batch` consumes the cache when active and falls back otherwise. Cache is reset in a `finally` after each schedule call. Both patches are gated on `model_type == "deepseek_v4"`, `hc_mult > 1` (for shape) / `num_hash_layers > 0` (for tokens), and `PP > 1`. They are strict no-ops for any other model. Model-side cleanup (deepseek_v4_model.py): - Drop the in-`forward` `input_ids` PP broadcast and the VPP fail-fast assert; the pre-broadcast patch handles it cleanly under both 1F1B and VPP. - Stop pre-assigning `self.mtp = None` so Megatron's `set_current_microbatch` only iterates `model.mtp.layers` when MTP is actually live (matches upstream GPTModel). - Use `getattr(self, "mtp", None)` for downstream MTP guards. Layer-specs (deepseek_v4_layer_specs.py): - Import `DeepSeekV4SpecProvider` so the type annotation resolves at module load (NameError surfaced once turbo path was off). Smoke results (mi355-gpu-12, BF16, 10 iters, MBS=1 GBS=16, seq=128, num_layers=8, num_hash_layers=3, hc_mult=4): - Smoke A 1x8 (PP=1 EP=1) : 10/10 iters - Smoke B 1x8 (PP=2 EP=4) : 10/10 iters - Smoke C 1x8 (PP=4 EP=2) : 10/10 iters - Smoke D 1x8 (PP=2 EP=4 VPP=2) : 10/10 iters Status doc (`deepseek-v4/develop/progress/status.md`) updated with the Phase 19 table, patch summaries, and three resolved entries in the Blockers / Risks Log. Smoke runner scripts under `deepseek-v4/develop/progress/p19/run_smoke{C,D}_v2.sh` are also included. Co-authored-by: Cursor <cursoragent@cursor.com>

…n-2 summary Wraps up the plan-2 program of work in the progress tracker now that the Phase 19 distributed re-validation gates are landed and the EP=8 / PP=2 EP=4 profile traces are captured. status.md (Phase 19 close-out): * mark `c10d::allreduce_` autograd warning as gone — verified absent in smoke A/B/C/D + EP=8 / PP=2 EP=4 profile logs on mi355-gpu-12. The EP routed-output reduction now flows entirely through Megatron's MoEAlltoAllTokenDispatcher / MoEFlexTokenDispatcher (P14 phase-2 + P17 dispatcher migration); cite the P19 logs as the runtime audit. * mark G11 (routing-snapshot diff = 0 across PP / EP changes) as deferred: the snapshot dump tooling never landed and is not on the pre-training release path. P19 smokes already cover the runtime stability of the P15 / P19 patches. * drop Phase 20 (perf / convergence gates), Phase 21 (docs / handover), and Phase 22+ (HF state-dict adapter) sections — they live as documented intent in plan-2/03-phase-details.md and re-enter active work when their respective triggers fire. * refresh the Blockers / Risks log entry for c10d::allreduce_ to point at the actual P19 verification (smokes A/B/C/D + EP=8 / PP=2 EP=4 profile runs on mi355-gpu-12) rather than 'still tracked into P19'. deepseek-v4/develop/progress/plan-2-summary.md (new): * stand-alone summary of the plan-2 architecture-faithful rewrite (P12 through P19), with: per-phase outcome + key commits; P19 deep-dive (smokes, profile traces, the two patches landed, c10d verification); test-gate ledger (G1 / G3 / G4 / G5 / G6 / G7 / G11 / G14 + smokes); plan-1 -> plan-2 architectural-shift table; explicit list of deferred / out-of-scope items (G6 distributed, G7 MTP, G11, P20, P21, P22+) and pointers to the live status.md, plan-2/, p19/ logs, and profile traces. deepseek-v4/develop/progress/p19/ (smoke / profile launchers): * run_profile_ep8.sh — torch.profiler trace for TP=1 PP=1 EP=8 (active step iter 6 -> 7) under the same V4 smoke config as P19 (BF16, MBS=1 GBS=16, seq=128, 8 layers, num_hash_layers=3, hc_mult=4). Output trace lives under output/<team>/<user>/<exp>/tensorboard/ as one chrome-trace JSON per rank. * run_profile_pp2_ep4.sh — same configuration with TP=1 PP=2 EP=4 to capture the multi-stage PP wire + Phase 19 patches (pp_tensor_shape, pp_token_pre_broadcast) under torch.profiler. deepseek-v4/download_ref.sh (new): * idempotent helper that ensures git-lfs is installed and clones the V4 reference assets at pinned commits — HuggingFace transformers, ROCm/TransformerEngine, AMD-AGI/Primus-Turbo, NVIDIA-NeMo/Automodel, plus the four DeepSeek-V4 model repos (Pro / Flash / Flash-Base / Pro-Base) with GIT_LFS_SKIP_SMUDGE=1 so large weights are not downloaded by default. No code changes; all five files land under deepseek-v4/.

…ment Plan-3 picks up after plan-2's distributed re-validation and is strictly scoped to two outcomes: 1. Reporting + spec hygiene fixes that came out of the Phase-19 smokes and the first full-Flash-size bring-up attempt: - Megatron's per-iter TFLOPs uses a generic transformer / MLA formula that does not know about V4's mHC K-stream packing, single-latent KV, grouped low-rank O, Compressor / Indexer side-paths, hash routing, or the V4 MTP head — so the reported TFLOPs is misleading on V4 today (P20). - V4 attention + V4 dense-MLP projection helpers warn-and-fall-back to vanilla nn.Linear when build_module(spec) raises, masking real spec bugs (gather_output=True / input_is_parallel=False rejected by TE wrappers) and producing an unsharded model. P21 makes the build strict + root-causes the rejected kwargs by routing through upstream non-TE ColumnParallelLinear / RowParallelLinear via two new provider helpers. 2. Primus-Turbo enablement for V4, mirroring the V2-Lite recipe: - core_attention submodule for V4's compress_ratio==0 (dense + SWA) layers via provider.core_attention(); HCA + CSA stay on the eager-Python path with code comments documenting the analysis (joint-softmax + shared sink for HCA, per-query top-K gather for CSA — neither is a flash-attn pattern today). attn_sink parameter aliased onto Turbo's self.sinks so the V4-Flash checkpoint key is preserved (P22). - Turbo DeepEP dispatcher reaches V4 specs by probing args.use_turbo_deepep directly (the existing turbo monkey-patch never reached V4 specs because deepseek_v4_layer_specs.py captured MoEFlexTokenDispatcher at top-level import) (P23). - run_deepseek_v4.sh smoke gate exercising the four P19 distributed configurations under the full turbo flag set (P24). Adds: - deepseek-v4/develop/plan-3/{README,01-roadmap,02-phase-details,03-test-strategy}.md - Phase 20-24 sections + four risk rows in deepseek-v4/develop/progress/status.md Co-authored-by: Cursor <cursoragent@cursor.com>

…d-form, rank-0 breakdown) Megatron's stock num_floating_point_operations falls on the standard transformer branch for V4 (V4 is single-latent KV, not MLA's compressed KV). The branch counts dense MLP per layer, a flat attention shape, no Compressor / Indexer / HC streams / V4 grouped low-rank O / V4 hash router / V4 MTP — so the per-iter TFLOPs printed today is misleading. Land a Primus patch that wraps the upstream function with a V4 closed form covering: * Q-LoRA + single-latent KV + grouped low-rank O at S * hc_mult * per-layer compress_ratio (dense / CSA / HCA), Compressor + Indexer side-paths gated on per-layer ratio * hash router (zero GEMM cost) on first num_hash_layers, learned router on the rest * V4 MTP runs a full inner V4 transformer layer per depth + eh_proj * one-shot rank-0 breakdown emit on first invocation Two non-obvious binding invariants surfaced during bring-up and are pinned in the patch: 1. Direct-import binding — primus.modules.trainer.megatron.trainer imports num_floating_point_operations at module load, so a monkey patch on megatron.training.training alone is invisible to the trainer's local copy. The patch calls _rebind_trainer_imports() to refresh the trainer's bound name after wrapping. 2. pretrain() enum overwrite — Megatron's pretrain() rewrites args.model_type from the YAML string "deepseek_v4" to a ModelType enum at training.py:1210 *before* train() ever calls num_floating_point_operations. A naive runtime check (args.model_type == "deepseek_v4") silently falls through to the upstream formula. The wrapper instead captures dispatch_v4 at install time via the closure (the @register_patch condition gates install to the YAML-string state). Smoke verification on mi355-gpu-12 (dev_primus_wenx_693, 8 GPUs, bs=16, S=128, hc_mult=4, L=8, mtp=0): closed-form total = 73.43 TFLOP / global-batch. * EP=8 last-iter: 17.9 TFLOP/s/GPU x 8 x 0.5125s = 73.4 TFLOP (Δ 0.04%) * PP=2 EP=4 last-iter: 14.0 TFLOP/s/GPU x 8 x 0.6559s = 73.5 TFLOP (Δ 0.09%) Logs: deepseek-v4/develop/progress/p20/{smoke_ep8_pp1_final.log, smoke_pp2_ep4_final.log}. Tests: 21/21 green (tests/unit_tests/backends/megatron/test_deepseek_v4_flops_patches.py) covering G16 (closed-form match, 7 per-component byte assertions) and G17 (non-V4 byte-for-byte fall-through, 6 parametrised model types + idempotency + post-pretrain ModelType-enum mutation). Plan-3 P20 closed; Plan-3 M1 advances to "P20 done, P21 open". Co-authored-by: Cursor <cursoragent@cursor.com>

…27787) The amend-after-commit flow on the previous SHA (78f0b416) renumbered HEAD to 4c27787. Backfill the six P20 task rows in status.md so the commit cells point to the live SHA on dev/wenx/deepseek-v4. Co-authored-by: Cursor <cursoragent@cursor.com>

Smokes at full Flash dims surfaced hundreds of ``DeepSeek-V4 attention projection submodule init failed (... gather_output = True ...); fallback to nn.Linear.`` warnings. The fallback masked real spec bugs and produced an unsharded model (vanilla nn.Linear instead of the column / row parallel shards the spec asked for) — TP=1 happened to work, TP>1 would have silently diverged. Plan-3 P21 makes the V4 build strict. Spec / provider: * New non-TE provider helpers ``column_parallel_linear_with_gather_output()`` and ``row_parallel_linear_with_scatter_input()`` returning the upstream Megatron ``ColumnParallelLinear`` / ``RowParallelLinear``. TE / Turbo wrappers explicitly raise on ``gather_output=True`` and ``input_is_parallel=False``; the upstream classes accept those flags natively. * ``_build_column_parallel_spec`` and ``_build_row_parallel_spec`` route through the new helpers when the caller asks for the gather / scatter variant; the standard TE path stays for the other cases. Attention surgery (deepseek_v4_attention.py): * ``_build_projection`` now does ``return build_module(submodule)`` unconditionally — failures propagate. ``submodule is None`` (CPU unit-test path) still returns a vanilla ``nn.Linear``. * ``_build_compressor`` and ``_build_indexer`` likewise drop their ``try/except/return local Compressor|Indexer`` blocks. The spec passed the same Python class as the fallback; the handler was dead code that masked real spec bugs. * ``self.attn_sink_module`` build branch retired (along with the ``submodules.attn_sink`` slot, the ``provider.v4_attention_sink()`` method, and ``primus/backends/megatron/core/transformer/attn_sink.py`` whose only consumer was that method). ``self.attn_sink: nn.Parameter`` stays — the inline softmax-with-sink in ``_attention_forward`` is canonical and matches the released V4-Flash checkpoint key ``layers.{i}.attn.attn_sink`` exactly. * Drop the unused ``logging`` import + module-level logger. Block surgery (deepseek_v4_block.py): * ``_build_projection`` drops its ``try/except/return nn.Linear`` block. ``config is None`` (CPU unit-test path) still returns ``nn.Linear``. Tests: * New ``tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p21_strict_build.py`` (13 cases, all green): - **G15** AST audit: walks every ``.py`` under ``primus/backends/megatron/core/`` and asserts no ``try → except → return nn.Linear(...)`` patterns remain; a string-grep gate also bans the retired warning strings (``"submodule init failed"``, ``"fallback to nn.Linear"``, ``"using local Compressor|Indexer"``, ``"attn_sink submodule init failed"``). - Provider-helper contract tests (``column_parallel_linear_with_gather_output`` returns the upstream non-TE class; row variant mirrors; ``provider.v4_attention_sink`` is gone). - Dataclass-surface test (``DeepseekV4AttentionSubmodules.attn_sink`` is gone; live slots intact). - **G15b** TP=1 build smoke: 1-rank gloo group, builds full V4 attention via ``_build_v4_attention_submodules`` + ``DeepseekV4Attention``, asserts every linear is one of ``{ColumnParallelLinear, RowParallelLinear, TELinear, PrimusTurboLinear}`` — never a bare ``nn.Linear``. * Existing ``test_attention_spec_uses_column_and_row_parallel`` updated to assert the new non-TE helper classes (not the TE classes that reject the flags). Smoke verification (mi355-gpu-12, primus-training container): * ``deepseek-v4/develop/progress/p21/smoke_ep8_pp1.log``: TP=1 PP=1 EP=8, 10/10 iters, lm_loss 11.88 → 11.65, ``grep -c`` for the banned strings = 0. * ``deepseek-v4/develop/progress/p21/smoke_pp2_ep4.log``: TP=1 PP=2 EP=4, 10/10 iters, lm_loss 11.89 → 11.62, ``grep -c`` for the banned strings = 0. Status / docs: * ``deepseek-v4/develop/progress/status.md`` — Phase 21 table marked done; M1 milestone closed; blocker entry resolved. * ``deepseek-v4/develop/plan-3/01-roadmap.md`` — P21 row done. * ``deepseek-v4/develop/plan-3/02-phase-details.md`` — new ``Status (2026-05-07)`` section under §P21 with the surgery summary, gate descriptions, and smoke evidence. Co-authored-by: Cursor <cursoragent@cursor.com>

…1 SHA (a4419ac) Co-authored-by: Cursor <cursoragent@cursor.com>

Copilot