feat(speculative): add DeepSeek V4 DSpark drafter and V4-Flash training

[khazic]

What

Adds a DSpark speculative-decoding drafter for DeepSeek V4 targets (Qwen3 and Gemma4 are already supported), plus the recipe wiring to train it against a full V4-Flash target, a one-shot smoke runner, and docs.

Components (CPU unit-tested, 9 passed)

• draft_deepseek_v4.py: DeepseekV4DSparkModel with a V4 attention backbone (Q-LoRA, single shared K=V latent, grouped O-LoRA, interleaved partial RoPE with an output inverse-RoPE, attention sink), a plain pre-norm decoder + dense SwiGLU MLP, and an fp32 inv_freq pin. Mirrors the Qwen3 DSpark draft's control flow.
• registry: DeepseekV4ForCausalLM -> DeepseekV4DSparkModel.
• build_deepseek_v4_draft_config: derives the draft config from a V4 target config (shrink depth, disable DSA / hash routing / MTP, add the DSpark fields, SDPA dense-mask path).
• target.py: HFDSparkTargetModel collapses V4's 4D Hyper-Connection streams (intermediate feature layers) to 3D via the stream mean; the final-norm output is already 3D.
• 9 CPU unit tests (forward shapes, three-term loss + backprop with frozen embed/lm_head, Markov head variants, config builder, registry, inv_freq fp32 pin). Re-run after the review cleanups: 9 passed.

Additionally validated on an A800 with a layer-pruned (not full) V4-Flash target: GPU/CPU forward / loss / backward finite, 4-GPU FSDP2 draft training (loss decreases, all-reduce), and real FP8 e4m3 + FP4 expert + e8m0 block-scale weights dequantizing to bf16 with finite hidden states.

Full V4-Flash training recipe (needs DeepEP + 8-GPU validation)

• train_dspark.py: a deepseek_v4 target branch that loads the full V4-Flash model as a frozen target through the V4 finetune recipe's EP/FSDP distributed path (FP8 dequant via the state-dict adapter), seeds the draft from the gathered embed_tokens / lm_head, and shards the draft over the world group. The Qwen3 / Gemma4 branches are unchanged.
• examples/speculative/dspark/deepseek_v4_flash_dspark.yaml: aligned to the official V4-Flash-DSpark hyperparameters (block_size=5, target_layer_ids=[40,41,42], markov_rank=256, mask_token_id=128799), now a multi-node (ep_size=16+) layout; see the file header for why ep_size=8 on one 8x80GB node OOMs at load.

The MoE token dispatcher (hybridep / deepep) requires DeepEP installed; without it setup() raises ImportError: HybridEP is not installed. So this path needs a DeepEP-equipped box.

Smoke test + docs (added)

• examples/speculative/dspark/deepseek_v4_flash_smoke.yaml: the same real EP path (full target, ep_size=8, DeepEP) shrunk (short seq, few anchors, 1 epoch, target_num_hidden_layers reduced) to reach a training step in minutes on one node.
• tests/functional_tests/speculative/run_deepseek_v4_flash_smoke.sh: one shot that runs the CPU unit tests, builds tiny smoke data, runs an 8-GPU training smoke, asserts a finite and non-increasing loss, and prints the EP sanity points.
• examples/speculative/dspark/README_deepseek_v4_flash.md: data prep, the smoke run, full training, and validation status.

Validation status

• Verified: the 9 CPU unit tests (re-run, 9 passed), the draft / recipe code correctness (static review), and, thanks to @HuiyingLi, the full distributed path end to end: an 8x H100 smoke (reduced target layers) and a 4-node / 32-GPU pilot (full 43-layer target, ep_size=32, HybridEP) exercising the distributed target load, hidden-state capture, draft training, and checkpoint save/LATEST restore. See the PR comments below for the detailed findings from those runs (missing chat template on the V4-Flash tokenizer, the full-target OOM at ep_size=8 root-caused to FP4->bf16 expert dequant and fixed by moving to a multi-node layout, and a TileLang JIT failure with target_attn_backend: eager now the default fallback).

• Loss-spike investigation and fix: the 4-node pilot showed unstable early loss (8.94 -> 14.81 -> 100.26 -> 30.55 over the first 4 steps, see comments below). Root-caused to two pre-existing issues in the shared TrainDSparkRecipe scaffolding (also used by the Qwen3 / Gemma4 branches, not specific to this PR's V4 code):

  1. Optimizer construction hardcoded torch.optim.AdamW directly on the model's native (bf16) parameters, so any optimizer._target_ / master_weights / master_weight_dtype / exp_avg_dtype / exp_avg_sq_dtype / store_param_remainders override was silently ignored; training ran with bf16 Adam moments instead of the fp32-master setup the config asked for.
  2. warmup_steps = warmup_ratio * total_optim_steps collapses to a handful of steps on short / small-dataset runs (4 steps for a 100-step run at the default warmup_ratio=0.04), dropping the freshly-initialized draft to near-peak LR almost immediately.

  Fixed: optimizer construction now routes through the shared build_optimizer() / OptimizerConfig system (defaults to plain AdamW when no _target_ is set, so existing Qwen3 / Gemma4 / V4 configs are unaffected), and warmup_steps is floored at 20 optimizer steps unless warmup_ratio<=0 (an explicit opt-out, e.g. the smoke config keeps its intentional no-warmup behavior).

  Re-validated by @HuiyingLi on the same 4-node / 32x H100 layout, with TE FusedAdam + master_weights: true + fp32 master/moments, lr=1e-5, on the 10,240-row regenerated dataset (5 epochs / 100 optimizer steps):

  step	train/loss	train/tv_loss	train/confidence_loss
  1	~9.5	~1.99	~0.7
  5	~5	~1.95	~0.1
  20	~3.3	~1.93	~0.15
  85	~3.0	~1.88	~0.2

  No spike; train/lr holds steady at 1e-5 and train/confidence_loss converges quickly. train/tv_loss stays in the 1.85-2.0 range across this short (100-step, 10k-row) run; likely just insufficient training volume for the block_size=5 / 129280-vocab objective at this scale rather than a correctness issue, worth re-checking once training is scaled to the full corpus / more epochs.

Signed-off-by: khazic khazzz1c@gmail.com

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[khazic]

@HuiyingLi this PR adds DSpark drafter support for DeepSeek V4 plus a full V4-Flash training recipe, and it needs validation on 8 GPUs (I could only verify the components on 1 to 4 GPUs). Thanks in advance for testing. Here is how to run it.

1. Unit tests (CPU, ~30s) confirm the draft model, config builder, registry, and the inv_freq fp32 pin:

python -m pytest tests/unit_tests/speculative/test_dspark_draft_deepseek_v4.py -q

Expect 9 passed. (They passed on the server before the final review cleanup; please confirm they still pass after it.)

2. Full V4-Flash DSpark training (8x80GB):

torchrun --standalone --nproc_per_node=8 -m nemo_automodel.recipes.llm.train_dspark -c examples/speculative/dspark/deepseek_v4_flash_dspark.yaml

The yaml loads the full deepseek-ai/DeepSeek-V4-Flash as a frozen target (EP/FSDP, FP8 dequant via the state-dict adapter, the same path the V4 finetune recipe uses) and trains the draft. As in the DSpark paper, regenerate the training responses with the target model first, then point recipe_args.train_data_path at them.

Key assumptions to validate on 8 GPUs (the ones I could not test):

1. FP8/FP4 dequant triggers: the V4 config carries quantization_config, so the EP-sharded adapter should dequantize the experts to bf16 rather than leaving fp8 scalar params for FSDP2.
2. The hidden states captured by HFDSparkTargetModel's forward hooks are full per rank (FSDP2 shards parameters, not activations), so no all-gather is needed before feeding the draft. Please confirm none come back as DTensors.
3. The frozen V4 forward accepts the 2D padding attention_mask plus use_cache=False that the wrapper passes.
4. ep_size (8 in the yaml) divides n_routed_experts (256), and pp_size stays 1 (the wrapper hooks a single non-pipelined model).
5. Memory: roughly 66 GiB/rank of weights on 80 GB cards; lower seq_length or enable activation_checkpointing if it is tight.

The draft itself runs the dense SDPA mask path (not tilelang/flex). Already verified on fewer GPUs: a single-card real V4-Flash load producing finite hidden states, and 4-GPU FSDP2 draft training with decreasing loss. Happy to fix anything that comes up.

[khazic]

@HuiyingLi this PR now ships a one-shot smoke runner and a README, so you can validate the full V4-Flash path on your DeepEP box without piecing commands together.

Prereqs: DeepEP (or HybridEP) installed, 8x at least 80 GB GPUs, and a local DeepSeek-V4-Flash checkpoint. The V4 MoE token dispatcher needs DeepEP; without it setup() raises ImportError: HybridEP is not installed.

1. Smoke first (a few minutes). This exercises the real EP path (full target, ep_size=8, tilelang, DeepEP) shrunk to reach a training step quickly, and asserts a finite, non-increasing loss:

TARGET=/path/to/DeepSeek-V4-Flash bash tests/functional_tests/speculative/run_deepseek_v4_flash_smoke.sh

It runs the CPU unit tests, builds 64 tiny rows, then an 8-GPU training smoke. If it prints SMOKE OK, the whole pipeline (distributed target load, hidden capture, draft train, loss) works end to end.

2. Full training. Regenerate Open-PerfectBlend responses with V4-Flash first (teacher-forced; see examples/speculative/dspark/README_deepseek_v4_flash.md), set target_model_name_or_path and train_data_path in the yaml, then:

torchrun --standalone --nproc_per_node=8 -m nemo_automodel.recipes.llm.train_dspark -c examples/speculative/dspark/deepseek_v4_flash_dspark.yaml

Per-step metrics (loss, ce_loss, l1_loss, confidence_loss, lr, mem) land in <output_dir>/dspark_train_metrics.jsonl.

What is and is not verified. The 9 CPU unit tests pass and I reviewed the draft + recipe code, but I do not have a DeepEP box, so the full distributed target load and cross-shard hidden capture are not yet validated. When you run the smoke, please eyeball these 5 points:

1. FP8/FP4 dequant on load keeps the target forward finite.
2. Captured hidden states are full per rank (not DTensors needing an all-gather).
3. The frozen V4 forward takes a 2D padding mask with use_cache=False.
4. ep_size=8 divides n_routed_experts=256 (32 experts/rank).
5. About 66 GiB/rank of weights; lower seq_length or set activation_checkpointing: true if it OOMs.

Thanks!

[HuiyingLi]

@khazic I ran the smoke on 8x H100 80 GB with DeepEP/HybridEP installed and a complete local DeepSeek-V4-Flash checkpoint.

Results

• The CPU suite passed: 9 passed.
• The unmodified full-target smoke does not finish loading. It OOMs in DeepseekV4StateDictAdapter._aggregate_experts() -> create_dtensor_from_local() while moving a stacked BF16 expert tensor to CUDA, before the first forward/loss:
  • typical rank: 78.87 GiB in use, 234.94 MiB free, failed a 512 MiB allocation
  • rank 0 failed a 1 GiB allocation with 490.94 MiB free
  • reserved-but-unused memory was only ~12 MiB, so this is capacity pressure rather than allocator fragmentation
  • reducing sequence length or activation checkpointing cannot address this load-time peak

To validate the remainder of the path, I added a local diagnostic option that changes target_config.num_hidden_layers from 43 to 4 before model construction:

recipe_args:
  target_num_hidden_layers: 4
  target_layer_ids: [1, 2, 3]

All target layers 0-3 are loaded and executed. [1,2,3] only identifies the three intermediate hidden states captured for the drafter; it is the four-layer equivalent of capturing the original target's final three layers [40,41,42]. The four-layer checkpoint load succeeded.

Missing tokenizer chat template

The V4-Flash tokenizer in this checkpoint has no chat_template. Because the smoke runner generates OpenAI messages rows, ChatDataset initially stopped with:

ValueError: ChatDataset requires a tokenizer with chat template support.

I resolved this locally by allowing recipe_args.chat_template and applying it immediately after tokenizer construction:

from nemo_automodel.components.datasets.llm.formatting_utils import _resolve_chat_template

self.tokenizer = NeMoAutoTokenizer.from_pretrained(target_path, trust_remote_code=trust_remote_code)
chat_template = recipe_cfg.get("chat_template", None)
if chat_template is not None:
    self.tokenizer.chat_template = _resolve_chat_template(str(chat_template))

For the smoke data I used a minimal template that marks only assistant content as the generation/loss span:

chat_template: |-
  {% for message in messages %}
  {% if message['role'] == 'assistant' %}{% generation %}assistant:
  {{ message['content'] }}{% endgeneration %}
  {% else %}{{ message['role'] }}:
  {{ message['content'] }}
  {% endif %}
  {% endfor %}

This produced 38 non-padding tokens and 25 assistant loss tokens for the first generated row.

Four-layer execution

With the PR's target_attn_backend: tilelang, the run reached the first target forward but failed in TileLang JIT compilation:

AttributeError: '_NestedLoopCheckVisitor' object has no attribute '_inst'

Using the model's Torch fallback via target_attn_backend: eager (while retaining 8-way EP, hybridep, hidden-state capture, and draft training) completed all 8 steps:

steps=8  first=9.6629  last=6.2684  min=5.9278
SMOKE OK

Peak logged memory was 16.06 GiB/rank. So the four-layer EP/hidden-capture/draft-training path works with eager attention; the original full-target load and the TileLang JIT path remain blockers in this environment. These were local diagnostic changes only.

[khazic]

Thanks for the 8x H100 run, @HuiyingLi. Your findings all map to fixes I just pushed. The short version: the full 43-layer target genuinely does not fit on one 8x80GB node, so full training is now a multi-node config, and the two patches you wrote by hand are now supported recipe options.

Why the full-target load OOMs (and why seq_length / activation checkpointing can't help). V4-Flash stores its 256 routed experts in FP4, but this path dequantizes them to bf16 for the grouped-expert GEMM, a 4x expansion. The routed experts alone are ~500 GiB in bf16, so at ep_size=8 that is ~63 GiB/rank of experts plus the replicated dense weights, which is exactly the ~79 GiB resident peak you hit in _aggregate_experts -> create_dtensor_from_local before the first forward. It is a resident-weight peak, so lowering seq_length or enabling activation checkpointing cannot move it, as you observed.

Full training is now multi-node. deepseek_v4_flash_dspark.yaml now defaults to ep_size=16 (2 nodes of 8x80GB), which puts experts at ~32 GiB/rank (~50 GiB/rank total) and fits comfortably; ep_size=32 (4 nodes) gives more headroom. Run the same command on every node, varying only --node-rank:

torchrun --nnodes=2 --node-rank=0 --nproc_per_node=8 --master-addr=<NODE0_IP> --master-port=29500 -m nemo_automodel.recipes.llm.train_dspark -c examples/speculative/dspark/deepseek_v4_flash_dspark.yaml

Your two local patches are now first-class options:

• recipe_args.chat_template: the V4-Flash tokenizer ships no chat template, so both configs now carry one (DeepSeek-style, with a {% generation %} assistant loss span). It is applied to the tokenizer just as in your patch, and the recipe fails fast with a clear message when it is absent. Set it to match the template you regenerate responses with.
• recipe_args.target_num_hidden_layers: the layer reduction you did by hand is now a supported diagnostic knob. The smoke config uses it (4 layers + target_layer_ids: [1, 2, 3]) so the smoke passes on a single 8x80GB node, which is the path you validated.

Attention backend defaults to eager. Since TileLang JIT-failed in your env ('_NestedLoopCheckVisitor' object has no attribute '_inst'), both configs default to target_attn_backend: eager for reliability; tilelang stays an opt-in speedup.

CPU unit tests pass (the original draft/config/registry suite plus new tests for the two recipe helpers). The one thing I still cannot run myself is the full multi-node load and cross-shard hidden capture at ep_size>=16. If you can give it a 2-node run, the 5 EP sanity points are in README_deepseek_v4_flash.md. Thanks again for the thorough testing.

[HuiyingLi]

@khazic Experimental setup/status update — this is not a training-success claim.

I pushed commit 37737a3 directly to the original PR branch khazic/feat/dspark-deepseek-v4. The commit contains only:

• examples/speculative/dspark/deepseek_v4_flash_dspark.yaml
• nemo_automodel/recipes/llm/train_dspark.py

Experimental setup

• Full frozen 43-layer DeepSeek-V4-Flash target
• 4 nodes / 32 H100 GPUs, EP32 with HybridEP
• 2-layer DSpark draft
• Sequence length 4096, 512 anchors, global batch 512
• Direct cleaned Open-PerfectBlend responses, without DSV4 regeneration; this is therefore off-policy training
• The checked-in YAML retains the reliable eager target-attention default. The current experiment overrides it with a local TileLang 0.1.11 environment.
• Added rank-zero W&B metric logging to the DSpark recipe

Current execution status

The distributed load/forward/backward path and checkpoint write/LATEST restore have been exercised. A 100-step pilot is currently running; full training has not been started.

W&B: https://wandb.ai/Nemo-automodel/deepseek-v4-dspark/runs/awf06tzq

The initial pilot loss is unstable: approximately 8.94 -> 14.81 -> 100.26 -> 30.55 in the first four steps. It later moved into roughly the 3.6-4.0 range by around step 40, but this behavior still needs investigation. In particular, the direct off-policy responses, chat/loss masking, objective scaling, and LR warmup should be checked before interpreting the run.

I do not consider this setup validated or successful yet.

[HuiyingLi]

@khazic Follow-up validation on current PR head 9fcfe1e. This is a 10k-row pilot/status report, not a full-training or convergence-success claim.

Integration note on the optimizer fix

Before running, I found one real-ConfigNode issue in the new optimizer path. ConfigNode.to_dict() resolves optimizer._target_ to the Python class, while build_optimizer(model, (name, kwargs)) requires name to remain a string. The new unit test uses a stub and does not expose this. I used this local correction:

target = opt_cfg.get_as_string("_target_", "torch.optim.AdamW")
kwargs.pop("_target_", None)

I also added a runtime optimizer-state log. The exact YAML -> _resolve_dspark_optimizer_spec -> build_optimizer path passed a 2-rank FSDP2 GPU preflight, and the full run logged:

Optimizer=FusedAdam lr=1.000e-05 master_weights=True
master_weight_dtype=torch.float32 store_param_remainders=True
exp_avg_dtype=torch.float32 exp_avg_sq_dtype=torch.float32

Common experiment setup

• Full frozen 43-layer DeepSeek-V4-Flash target; TileLang target attention; HybridEP
• Regenerated DSV4 responses: 10,000 rows plus 240 repeated rows for exact batch alignment (not the original PerfectBlend responses)
• Validation remained the held-out original PerfectBlend train[:4096]
• Sequence length 4096, 512 anchors, microbatch 1, global batch 512
• TE FusedAdam, betas (0.9, 0.95), constant LR 1e-5, FP32 master and moments
• PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
• 5 epochs, 20 optimizer steps/epoch, 100 total steps

2 nodes / EP16 result

Slurm job 13292639, EP16, accumulation 32. The full target loaded and TE optimizer initialized, but the first uncompromised draft forward OOMed in eager_attention_with_sink at the FP32 softmax. Each rank needed another 4.06 GiB with only about 3.1-3.4 GiB free. With expandable segments, reserved-unused memory was only about 0.16-0.24 GiB, so this is capacity pressure, not fragmentation. No optimizer step completed.

W&B: https://wandb.ai/Nemo-automodel/deepseek-v4-dspark/runs/70iomuua

4 nodes / EP32 result

Slurm job 13292649, EP32, accumulation 16: COMPLETED 0:0 in 44:44. Peak allocated memory was 60.83 GiB/GPU. W&B has 100 contiguous train-loss points, five validation points, and constant train/lr=1e-5.

step	train/loss	val/loss
1	9.2867
20	3.5881	3.5286
40	3.1633	3.2306
60	3.0255	3.0572
80	2.9654	2.9784
100	2.8767	2.9575

Minimum train loss was 2.8590 at step 96; mean train loss over steps 91-100 was 2.8932. Final step components were CE 9.2785, TV 1.8548, confidence 0.2095. There was no early loss spike in this setup.

W&B: https://wandb.ai/Nemo-automodel/deepseek-v4-dspark/runs/jv7t32yq

The final step-100 checkpoint includes model/optimizer/RNG state and a consolidated HF safetensors export at:

/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_nemofw/users/huiyingl/sd/test-runs/pr2866-regenerated-draft-10k-4node-latest/run-13292649/checkpoints/epoch_4_step_100

This validates the updated optimizer path and 4-node distributed execution for the short regenerated-data pilot. The objective still needs evaluation at full data scale before interpreting this as drafter quality/convergence.

[HuiyingLi]

LGTM Thank you!

[khazic]

@HuiyingLi CI status update after the loss-spike fix (optimizer/warmup, commit 6e40b02) and the codecov fix:

• codecov/patch is now passing.
• Every other check passed except L2_HF_DCP, which failed with a distributed-checkpoint-save crash unrelated to this PR: tests/functional_tests/checkpoint/test_hf_consolidated_llm.py::test_consolidated_llm_checkpoint (the main LLM finetune recipe's checkpoint test, not DSpark) aborted inside torch.distributed.checkpoint's reduce_scatter/gather_object after one rank hung and was force-killed by the elastic launcher (Fatal Python error: Aborted, exitcode -6). The same test suite passed on gb200_L2_HF_DCP in this same run, and nothing in this PR touches checkpointing code, so this looks like a one-off distributed-collective hang on that runner rather than a real regression.

I don't have permission to re-run the failed job (gh run rerun needs repo admin rights). Could you re-run L2_HF_DCP on this run (https://github.com/NVIDIA-NeMo/Automodel/actions/runs/28495213096) when you get a chance?

[NKNaN]

Hi, this is a great work. We wonder is there any open sourced checkpoint trained with full Opend-PerfectBlend corpus of this dense dspark draft model? And how about the performance of this model compared to the open sourced dsv4 dspark draft model (moe + mhc) in terms of accept rate or other metrics?