[micahtyong]

Summary

The DeepSeek-V4.1-Flash recipes train one sequence per micro-batch: the V4.1 attention in Miles takes token positions from the row, so sequences cannot be packed. A step of 128 GRPO samples on 16 H200s at TP4 / EP16 is therefore 32 micro-batches per data-parallel rank, each run through 40 layers forward, recomputed and backward. The step-2 profile of the shipped recipe shows 2.2 million kernel launches of a few microseconds each, the GPU idle 86% of forward_backward, and 4% of the time in communication. The step is bound by launching kernels, so the fix is fewer micro-batches per rank.

Expert parallelism already spans all 16 GPUs, so tensor parallelism only divides the attention and shared-expert work, and every TP rank still launches every kernel. TP1 gives 16 data-parallel ranks with 8 micro-batches each and no tensor-parallel collectives. Peak memory rises from 71 GiB to 84 GiB per H200. The 32K recipe keeps its TP8 override; the 1017 8K recipe inherits the change.

Trainer-only replay of the same 128 sequences (cross-entropy, 5 steps). Parity is the mean |Δ logprob| at step 0 against the shipped recipe; 3.5e-2 is the run-to-run floor, measured by rerunning the shipped recipe with deterministic mode off. The TP1 row was run with data_pad_size_multiplier: 32, which on its own changed nothing (77.7 → 80.8 s at TP4) and is not part of this PR.

Trainer layout (16×H200, EP16)	Micro-batches per rank	forward_backward	Peak memory	Parity
TP4 / DP4 (shipped)	32	77.7 s	71 GiB	—
TP2 / DP8	16	36.8 s	75 GiB	3.6e-2
TP1 / DP16	8	16.0 s	84 GiB	3.6e-2
TP1 / DP16, 16 rows of 8K tokens (the recipe's ceiling)	1	11.8 s	124 GiB (129 GiB reserved)	—

At the 8K ceiling the layout keeps about 11 GiB of headroom on an H200. A micro-batch holds one row, so the peak depends on row length, not on the number of samples per step. Recompute has to stay on: without it a single 8K row runs out of memory at TP1. B200 and B300, which the recipe prefers, have more memory but were not run.

E2E results

Setup: deepseek_v41_flash_1017_lora_8k on H200 (2 trainer nodes, 1 sampler replica), stock tinker-cookbook math_rl GSM8K GRPO, 128 samples of up to 512 tokens per step, 5 steps; steady state is the median of steps 1, 3 and 4. Before / After

Per step	Before	After
Total step time	129.5 s	73.0 s
Training step	80.5 s	26.0 s
forward_backward (trainer side)	78.8 s	22.4 s (16.6 s by step 4)
Checkpoint for the sampler	5.5 s	7.8 s

Peak trainer memory per GPU is 84 GiB instead of 71 GiB; KL between sampler and trainer 0.0025 (0.0032 before). Trainer cold start is unchanged (model loaded at 206 s, serving at 324 s, first step compiled in 258 s). Sampling is untouched by this PR; the sampler stack (#146–#149) covers that side. With both together a step takes 47.9 s (run, 8 steps, same medians).

🤖 Generated with Claude Code

[micahtyong]

Hold for one more check: every run so far used 512–640-token rows. A 16-row × 8K-token probe at TP1 is running to confirm the recipe's full context still fits on H200 with recompute; numbers to follow here.

[micahtyong]

8K check done: 16 rows of ~8,000 tokens (one per rank at DP16) train at TP1 with full recompute, 3/3 steps, peak 124 GiB allocated / 129 GiB reserved on H200, forward_backward 11.8 s per 128K-token step. Without recompute the same rows OOM in the first forward, so the recipe's recompute settings stay. Table in the description updated.

[micahtyong]

One thing to weigh before merging: user configs that subclass deepseek_v41_flash_lora_8k and only raise max_context_length inherit this TP1 setting. At 32K a TP1 row does not fit an H200 (an 8K row already peaks at 124 GiB), so such a config would go from TP4 (which #128 found fits to ~24K) to an OOM. The shipped 32K recipe is unaffected (it overrides TP to 8). I'm measuring TP8/TP4/TP2/TP1 at 32K rows on H200 and Blackwell now and will add the numbers and a note to the recipe docstring, so the right TP per context length and GPU is written down.