Support dpsk-4.1 sft and rl - #201
Open
tastelikefeet wants to merge 19 commits into
Open
tastelikefeet wants to merge 19 commits into
tastelikefeet wants to merge 19 commits into
Conversation
- Add deepseek_v41 GPT bridge: CSA2 compressor/indexer, single-pass mHC, native trainable Engram modules, Vision/Aligner, DSpark (MTP) draft stack - Add DSpark stack module (markov_head / confidence_head endpoints) - Config parser/model_config support for composite deepseek_v41 config - GPTBridge.export_weights: add skip_unsupported_export flag (RL weight sync skips engram export; save/checkpoint path unaffected) - FP8/FP4 dequant handling in safetensors util - Add engram load/export unit test
…text backbone) Migrate the DeepSeek-V4.1 text backbone from GPTModel to the PP-capable HybridModel path, keeping GPTModel as the golden baseline. - Add DeepseekV41HybridLoader/Bridge: 1 GPT layer -> 2 hybrid layers (D attn + E/- mlp), 2x-derived CSA2 arrays/pattern, block-aligned PP/VPP segmentation, Engram placement at hybrid layer 2e+1, and an HC wrapper forward-override that applies the Engram delta on the n-stream residual to match the golden order. - Route to the hybrid path when pipeline_model_parallel_size > 1, with a deepseek_v41_hybrid config flag (via --megatron_extra_kwargs) to force it at PP1. - Fix DSv4HybridSelfAttention layer_type lookup for the doubled hybrid index space. - Normalize _convert layer count across the load (GPT-space) / export (doubled) config spaces via _num_hybrid_layers, fixing an export crash on out-of-range layers. GPU-validated on a tiny model: hybrid PP1/PP2/PP2xEP2/PP2xDP2 all align with the GPTModel baseline at iter-1 (<0.06% rel loss/grad) with clean weight exports. Unit tests: 62 passed (hybrid 27 + engram 35).
…(B3) Backbone-agnostic DSpark support for the PP-capable HybridModel path: - deepseek_v41.py: extract _attach_dspark() and _convert_dspark_stack() from the GPT build_model/_convert_additional_layers (behavior-preserving). - deepseek_v41_hybrid.py: build_model attaches DSpark on the text-only model itself; _convert_additional_layers maps mtp.* via _lm() with a non-last-stage guard, reusing _convert_dspark_stack. - Inference-time capture hook rehang deferred (documented): hybrid has no speculative-decoding forward, so a standalone hook would be dead code. Verified (tiny SFT): iter-1 loss/grad match the B1 hybrid baseline within noise; mtp.* key-ledger round-trips cleanly (52 DSpark keys, the sole delta vs B1 export, no unexpected keys vs golden GPT baseline). 33 unit tests pass.
Add DeepseekV41MultimodalHybridModel (composition over DeepseekV41Multimodal- GPTModel, swapping language_model_cls to the PP-capable HybridModel backbone) and route the hybrid loader/bridge through it, so the vision tower + image-embed injection reuse the golden GPT multimodal path verbatim. - DeepseekV41HybridStackModel: expose extra_forward_keys=[]; forward() now unpacks the wrapper's extra_block_kwargs container (upstream HybridModel.forward has no such param -- it threads input_ids itself) before stripping visual keys. - HybridLoader.build_model: resolve language_model via getattr and apply the engram/is_expert/dspark fix-ups on the nested backbone. - HybridBridge: _set_word_embeddings via _lm; _convert_pre_process dispatches to the GPT vision path when visual is present, else text-only word emb (guards the visual=None non-first PP stage on export). Tests: 41 passed (+2 forward extra_block_kwargs unpack / pixel_values guard). GPU accept: text-only hybrid_pp1_b4 iter-1 loss=12.89138508/grad=102.55239105 in the aligned family band (wrapper does not perturb text loss); export key ledger 285 keys, 0 missing/0 unexpected vs both source and golden GPT baseline (+42 vision/aligner/image_* over B3's 243).
Switch _deepseek_v41_use_hybrid to return True by default so every layout (including PP1/single-GPU) builds on the PP-capable HybridModel path, now that B1-B4 validated it against the GPTModel golden baseline (iter-1 loss/grad within the bf16/MoE non-determinism band, clean weight key ledger). The GPTModel path is retained as a force-off regression baseline via deepseek_v41_hybrid=False. Update routing docstrings/module header and loader/bridge/use_hybrid unit tests (41 passed). GPU smoke at PP1 with no force flag confirms default routing builds DeepseekV41MultimodalHybridModel and aligns (iter-1 loss 12.8874/grad 102.589). VPP remains an optional follow-up (needs Engram VPP validation coverage).
Enable PP + VPP (interleaved schedule) for the DeepSeek-V4.1 HybridModel path. Layer segmentation was already vp_stage-aware from the B1 migration; this lands the remaining three fixes, all in mcore-bridge: - engram: relax the upstream Engram parallelism guard to also allow VPP (rename _ContextParallelSizeOneView -> _RelaxedParallelismView, hiding both context_parallel_size and virtual_pipeline_model_parallel_size while keeping every other check, e.g. etp!=tp). Engram.forward is self-contained (local n-gram hashing + slicing, no cross-stage state); layer placement is vp_stage-aware and the whole-block guard rejects half-block stages. - deepseek_v41: thread vp_stage into _attach_dspark's build_module so the DSpark draft layers satisfy get_transformer_layer_offset's VPP assertion (vp_stage is not None). Offset is 0 for the tiny draft stack, placement unchanged. - mm_gpt_model: surface the language model's typed-pipeline payload interface (pipeline_payload_factory / pipeline_payload_spec) on MultimodalGPTModel. get_attr_wrapped_model only descends via .module and cannot reach self.language_model, so the CSA2 custom cross-stage payload fell back to the shape-based P2PCommunicator, which calls .size() on the payload and crashes under both 1F1B and interleaved schedules. Backbones without a payload (GPTModel, GLM) expose None and keep the legacy path. Tests: test_deepseek_v41_engram.py + test_deepseek_v41_hybrid.py -> 76 passed. GPU smoke (tiny, multimodal default): PP2xVPP2 iter-1 loss 12.89163 / grad 102.519 aligns with same-model PP2 (1F1B) baseline 12.89124 / 102.558.
…tary and weight fidelity Found while validating PP / CP / baseline loss-grad alignment and the RL scripts on a 4-layer tiny model. * Keep the PP export collectives symmetric: `_convert_pre_process` no longer branches on this rank's `visual`, and non-last stages convert the DSpark stack through an empty structural proxy. A per-rank guard desynchronized the pp-group all-reduces so the last stage read a stale `has_model` and raised StopIteration. * Pre-index the MLA rotary table by `position_ids` under `thd` packing, mirroring what `GPTModel.forward` does. The DSv4 attention consumes per-token frequencies, so without this the packed and CP paths hit `freqs.shape[0] != tokens`. `HybridModel.forward` never threads `position_ids` into the decoder, hence the stash for the rotary pre-hook. * Reject non-contiguous CP partitioning on the non-packed Engram path too: it reads the layout from the transformer config instead of packed_seq_params, so zigzag used to slip through and silently mis-align hashes with the local hidden states. * Mark the DSpark router `expert_bias` keep-in-fp32. The draft stack never runs in the training forward, so mcore's lazy `_maintain_float32_expert_bias` never fired and the checkpoint's fp32 `mtp.*.ffn.gate.bias` round-tripped through bf16. * Guard the hyper-connection `alpha_*` export with `_peft_format` (both the GPT and hybrid copies). These are frozen base weights written outside `_set_state_dict`, so a LoRA run used to write `hc_*_scale` into adapter_model.safetensors and demand it back on load. Validation on the tiny model (iter-1 relative difference, the clean signal; a same-config rerun measures 0.0141% loss / 0.145% grad as the noise floor): PP2 0.0083% / 0.108%, CP2 0.0134% / 0.027%, PP2+packing 0.0068% / 0.031%. Full-param export stays at 285 keys matching the source; the LoRA adapter now holds 155 lora_A + 155 lora_B and nothing else.
…plementation
The GPTModel-based V4.1 path was only kept as the golden baseline while the
HybridModel backbone was being brought up. Now that HybridModel is the default
for every layout (TP/PP/VPP/CP/EP/DP + packing, all verified against the PP1
baseline), the two-file / two-class-layer split is pure overhead: every load or
export walked DeepseekV41Hybrid{Loader,Bridge} -> DeepseekV41{Loader,Bridge},
where the parent half of each pair was mostly dead code.
* merge deepseek_v41_hybrid.py into deepseek_v41.py (registration must live in
an eagerly imported module) and collapse the two inheritance layers into one
DeepseekV41Loader / DeepseekV41Bridge; the only cross-layer super() call
(_convert_pre_process) is merged by hand
* guard register_model with _HYBRID_MODEL_AVAILABLE: V4.1 now requires a
megatron with megatron.core.models.hybrid
* drop the ModelConfig.deepseek_v41_hybrid selector and both __new__ routers,
which no longer select anything
Removes (irreversible, recover from 1461d8b or earlier): DeepseekV41GPTModel and
its DSpark speculative-decoding methods (forward_dspark /
compute_dspark_speculative_tokens / the capture hooks, ~240 lines) plus their 3
unit tests. That code was already unreachable on the HybridModel backbone -- the
wrapper delegated to methods HybridModel does not define, so any call raised
AttributeError. The shared DeepseekV41DSparkAttention (prefill_dspark /
reset_dspark_cache) and the mtp.* draft-stack weight conversion are untouched.
Tests: 70 passed (test_deepseek_v41_engram.py, test_deepseek_v41_hybrid.py).
Four stale expectations updated -- two were already failing before this commit
(they still asserted the pre-PP-fix "skip the vision block / draft stack on
non-last stages" behaviour), and two asserted GPT-path semantics that the merge
replaced (monkeypatching the parent bridge; the non-doubled Engram layer_number).
Tiny 4-layer regression, iteration-1 loss / grad_norm:
PP1 baseline reproduced across the refactor : 0.0125% / 0.0166%
packed CP1 baseline reproduced : 0.0003% / 0.0017%
PP2 vs PP1 (was 0.0487% / 0.0484%) : 0.0147% / 0.0050%
CP2 vs CP1 (was 0.0210% / 0.2688%) : 0.0179% / 0.2429%
Export: full-parameter PP2 round-trips all 285 HF keys (no missing/unexpected);
LoRA exports 310 adapter tensors (155 lora_A + 155 lora_B) and no base weights.
The shared DSv4 attention applied a per-head query RMS norm unconditionally. V4.1 already normalizes the query latent via q_layernorm and feeds wq_b's output straight into RoPE, so re-normalizing every head here rescaled the attention scores. Gate it on dsv4_version == 'v4'.
…fork The bridge's TransformerLayer replaces upstream's __init__ rather than extending it, and had not mirrored upstream's engram composition point, so self.engram was never set. The inherited _forward_attention reads self.engram on every layer, so every model assembled through this fork raised AttributeError on the first forward. Mirror the upstream engram block here, guarding submodules.engram with getattr for Megatron-Core builds that predate the composition point.
num_nextn_predict_layers is the standard MTP key carried by every DeepSeek V3/V4 config, so reading it as the DSpark layer count made plain V3/V4 checkpoints fail DSpark validation. Mark a DSpark checkpoint by dspark_block_size instead (V4.1-Flash and the V4 Vision experiment carry it, plain V4-Flash does not) and clear dspark_num_layers otherwise. The draft stack's own expert counts are optional: a checkpoint omitting them reuses the backbone's MoE shape, so the builder and weight mapping fall back to it. Freeze the draft stack after attaching it: it never runs in the training forward, so its parameters only ever hold a zero gradient, but Adam's decoupled weight decay would still erode the mtp.* weights every step. The bridge reads and writes param.data directly, so the checkpoint still round-trips unchanged.
The mcore_bridge.inference package (DeepseekV41TextGenerationController / DeepseekV41DynamicInferenceEngine) was scaffolding for a Megatron-native DSpark speculative-decoding path that is not wired into either SFT or RL: swift trains via the forward/backward path and serves V4.1 rollout through vLLM, so nothing imported it beyond one isolated unit test. Remove the package and its test; the training-side DSpark stack (model/modules/dspark.py) is untouched.
The Engram module was carrying a full Megatron-native inference machinery that swift never exercises: rollout/serving goes through vLLM and training only runs the forward, so inference_context is always None in production. Remove the inference-only hashing (_static_inference_hashes, _dynamic_inference_hashes and their private _hash_windows helper), collapse _build_hash_ids to the training path, and delete the two dead standalone helpers adapt_deepseek_v41_layer_specs (the loader builds the D-layer spec inline) and allow_engram_inference (bypass for a native-inference guard that is never entered). forward now keeps its inference_context parameter -- the bridge wrapper still passes it positionally -- but fails loud instead of silently running inference, so the training forward is unchanged. Drop the now-unused contextmanager import and the three tests that covered the removed code.
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.
No description provided.