Add HybridModel DeepSeek-V4 layer construction

[FDecaYed]

Add HybridModel DeepSeek-V4 layer construction

Summary

Add HybridModel symbols, specs, ratio handling, and construction for DeepSeek-V4 after the native attention and mHC runtime contracts are available.

Scope

• Extend Hybrid layer allocation with C/H/W for compression ratios 4/128/0 while preserving the existing DSA D and MLA + symbols.
• Add CSA/HCA/window submodule slots and construct the corresponding D2 attention layers.
• Add the Transformer Engine hybrid_dsv4_stack_spec callable and callable-spec handling in training and builder entry points.
• Infer dsv4_hybrid and MLA use from C/H/W patterns.
• Normalize either compact C/H/W-only or legacy full-layer csa_compress_ratios, validating each ratio against its symbol.
• Add focused allocation, argument-normalization, stack-spec, and builder coverage.

Production diff: 6 files, +195/-20. Test diff: 3 files, +153/-136.

Non-goals

• Reimplementing mHC or Hash routing; those remain P2 and P4.
• Concrete fused CSA kernels; D3 is an independent sibling and is not a dependency.
• Packed THD, context parallelism, CUDA-graph lifecycle, FLOP/accounting, functional recipes, or GPTModel support.

Review boundary and dependencies

Requires #6402 (D2) and #6401 (P2); D3 is not a dependency. The local branch is built on D2 and temporarily carries:

• 993fe75090 — patch-equivalent to the current [Split 1/N of #3430] feat(mHC): basic pytorch implementation of manifold hyper connection #4531 prerequisite;
• 2ebb42e43c — patch-equivalent to P2;
• a82126a5621b09cfda5abaa1c16b90640c10bedb — this PR’s own change.

The exact own diff is 9 files, +348/-156. Please review the last commit only. After D2 and P2 merge, the own commit will be replayed directly onto then-current main.

Provenance

Reconstructed from current HybridModel APIs plus the post-#6020 DSv4 residual. #5795 is the historical behavior/test oracle, #6372 defines the current CSA namespace, and #5960 supplies attention-teacher correctness inherited through D1/D2.

Validation

Passed locally:

• clean worktree and no dependency/lock-file changes;
• signed-commit verification and required Signed-off-by trailer;
• isort, Ruff, Python bytecode compilation, git diff --check, and focused API/import scans;
• focused AST/argument behavior checks.

Black’s whole-file checks for megatron/training/argument_utils.py and megatron/training/arguments.py retain the same baseline formatting failures; the D4 hunks introduce no new formatting regression. Runtime pytest is blocked on the bare macOS environment without the Linux CUDA stack. Focused GPU tests and full functional CI remain required after publication.

Current functional constraint

C/H/W layers currently require explicitly selecting the callable hybrid_dsv4_stack_spec with transformer_impl='transformer_engine'. Omitting the DSv4 spec or selecting an inference-optimized/default stack can leave the new submodule slots as IdentityOp; those configurations are unsupported and must not be advertised. A later early-rejection check is optional and does not block this construction slice.

Sibling rebase constraint

D4 can merge without #6404 (D3). Both touch megatron/training/argument_utils.py and megatron/training/arguments.py; if D3 merges first, D4’s rebase must retain D3’s omitted-versus-explicit backend resolution while adding the C/H/W inference and ratio normalization here.

Local artifact

D4-hybrid-dsv4-layer-construction.patch, generated from the own commit only. Verify it against SHA256SUMS.

Before publication

After D2 and P2 merge, replay only D4’s own commit onto current main, preserve any already-merged D3 backend-default logic, confirm the final HybridModel/mHC call signatures and explicit-spec constraint, then rerun focused construction/forward tests and full functional CI.

[copy-pr-bot]

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.