Skip to content

feat(deepseek-v4): enable bf16 RL, bump vLLM to 0.29.0 and prime-envs - #3543

Draft
garrett361 wants to merge 9 commits into
mainfrom
feat/ds-v4-bf16-rl
Draft

garrett361 wants to merge 9 commits into
mainfrom
feat/ds-v4-bf16-rl

Conversation

@garrett361

@garrett361 garrett361 commented Sep 14, 2026

Copy link
Copy Markdown
Contributor

DeepSeek V4 Flash passes the KL mismatch merge bar for bf16 RL. The blocker was a vLLM weight-reload bug that left captured CUDA graphs reading a freed pointer, fixed upstream in vllm#52626 and picked up here by pinning vLLM 0.29.0.

math environment, batch_size = 64, 20 steps, lr = 0, full depth, 4 trainer nodes + 1 inference node:

step replay on replay off
1 0.00037 0.00091
2 0.00030 0.00097
3 0.00040 0.00108
4 0.00025 0.00092
5 0.00028 0.00068
6 0.00031 0.00083
7 0.00037 0.00097
8 0.00030 0.00070
9 0.00040 0.00077
10 0.00030 0.00101
11 0.00031 0.00067
12 0.00035 0.00099
13 0.00028 0.00085
14 0.00034 0.00075
15 0.00035 0.00065
16 0.00030 0.00108
17 0.00037 0.00094
18 0.00037 0.00081
19 0.00029 0.00075
20 0.00039 0.00110

All 40 entries under the 0.015 bar: max 0.00040 with replay on, 0.00110 with replay off.

reverse-text

The same model and topology on reverse-text at batch_size = 32, over 20 steps:

configuration mismatch_kl max
vLLM 0.28, router replay off 0.194 0.490
this PR, router replay off 0.024 0.053
this PR, as shipped 0.0016 0.0027

The vLLM 0.28 row is a 2026-09-14 measurement carried forward. The two this PR rows were
re-measured on the current branch, so rerunning the baseline would compare against a different
main than the one it was originally a baseline for.

  • The math table was measured before this branch bumped prime-envs. At the old pin 1f1e050ab, environments/math/ set network_allow = [], which makes verifiers refuse the subprocess runtime every config in this repo uses, so the run needed a scratch wrapper taskset restoring the wildcard. The bump to a6fc0c66f picks up prime-envs Simple implementation of different loss scaling #795, which reverts that default-deny, so the environment now runs unwrapped. The other two measurement-only settings, constant_trainer_batch_size = false and an inert length-penalty tiebreaker so frozen-policy groups ship, are still not committed.
  • Serving the published repo would make vLLM allocate fp8 linears and MXFP4 experts that bf16 weights cannot fill, since its config.json carries quantization_config and expert_dtype = "fp4". These configs use PrimeIntellect/DeepSeek-V4-Flash-0731-bf16 instead, currently private.

@garrett361 garrett361 changed the title feat(deepseek-v4): enable bf16 RL and bump vLLM to 0.29.0 feat(deepseek-v4): enable bf16 RL, bump vLLM to 0.29.0 and prime-envs Sep 17, 2026
garrett361 added a commit that referenced this pull request Sep 21, 2026
Mirrors PR #3543's layout: per-step math table, per-step reverse-text
table, and a summary. The bf16 columns reproduce #3543's own numbers,
which makes the FP8 columns directly comparable to it.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
garrett361 added a commit that referenced this pull request Sep 21, 2026
Mirrors PR #3543's layout: per-step math table, per-step reverse-text
table, and a summary. The bf16 columns reproduce #3543's own numbers,
which makes the FP8 columns directly comparable to it.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
garrett361 added a commit that referenced this pull request Sep 22, 2026
Mirrors PR #3543's layout: per-step math table, per-step reverse-text
table, and a summary. The bf16 columns reproduce #3543's own numbers,
which makes the FP8 columns directly comparable to it.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
garrett361 and others added 8 commits September 22, 2026 13:23
…okenize-params patch

The patched _text_len_check and get_encode_kwargs are full-body copies of
upstream, and 0.29 added truncation-side overrides to both originals:
encode must skip tokenizer-level truncation (its default side may differ
from the request) and _text_len_check pre-trims characters from the
requested side. Restore both branches, keyed on max_total_tokens as the
copies already are.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…ight loaders

DeepseekV4Model.load_weights writes the attention sinks with a bare
params_dict[name][:n].copy_(narrow_weight) instead of going through
param.weight_loader. Layerwise reload moves a layer's tensors to meta and
wraps each loader to buffer the incoming tensor, so that copy lands in a
meta tensor and is discarded silently: load_numel stays 0, finalize
restores the boot value with only a warning, and the loader still does
loaded_params.add(name), so a named_parameters() - loaded_params diff
cannot see the loss either. Attention sinks are trainable, so every weight
update keeps serving the sinks the server booted with.

Padding this rank's heads back up with -inf, the parameter's own init
value, makes it an ordinary full-parameter load through the existing
loader contract. Still required on vLLM 0.29.0 (model.py:1546).

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Writes a bf16 HF-layout export of a prime-rl conversion cache, restoring
the published key names, splitting the fused expert tensors back per
expert, and emitting a config.json with the quantization fields removed.

DeepSeek V4 Flash cannot otherwise be served against a bf16 weight
broadcast: the published repo's config.json carries quantization_config
and expert_dtype = "fp4", so vLLM allocates fp8 linears and MXFP4 experts
that bf16 weights cannot fill. The routed experts then fail loudly in
_load_w13 while the fp8 linears fail silently, since a bf16 tensor matches
the fp8 parameter's shape and copy_ casts it with no scale applied.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
rl.toml trains deepseek-ai/DeepSeek-V4-Flash-0731 at full depth over a
bf16 NCCL weight broadcast, 4 train nodes plus 1 inference node, with
router replay on. As written it is a fidelity measurement rather than a
training run: lr = 0 holds the policy fixed so the trainer-vs-inference
mismatch KL is measurable without policy drift.

inference_bf16.toml is the serving half on its own, a standalone
pre-flight for whether this cluster's vLLM allocates bf16 parameters for
this checkpoint. Both serve the bf16 export built by
tools/convert_prime_to_hf.py, not the published repo.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The bf16 export is now published as PrimeIntellect/DeepSeek-V4-Flash-0731-bf16,
so the configs name a hub id instead of an absolute path under
/home/hf-cache/exports. That path existed on exactly one cluster, which made
the shipped configs unrunnable anywhere else.

Publishing the artifact also retires tools/convert_prime_to_hf.py from this PR.
It was included so others could rebuild the export, and a published checkpoint
needs no rebuild. It remains in git as ebadd7589 on exp/ds-v4-router-replay.

The trainer keeps loading deepseek-ai/DeepSeek-V4-Flash-0731 and dequantizing
it through its own conversion cache, so only the serving side moves.
validate_shared_model_name requires the orchestrator and vLLM model names to
be identical, so those two move together.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
propagate_shared_fields fans a top-level [model] name out to trainer,
orchestrator and inference, so the config names the bf16 repo once instead of
three times, matching configs/advanced/minimax-m2.5/swe.toml.

The trainer previously loaded the published fp8 checkpoint and dequantized it
on the way in. That yields the same weights by construction, since the bf16
export was written from the trainer's own conversion cache and the index's
metadata.total_size is byte-identical to it, but naming a single source removes
a link from the chain the mismatch-KL evidence rests on.

This costs disk rather than saving it: trainer/model.py:663 builds a prime
conversion cache for any HF-format snapshot, and both repos are HF-format, so
bf16 means a 530 GB download plus a 530 GB cache where fp8 meant 167 GB plus
the same cache.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Restores the math environment's default solver network access. prime-envs #780 had
set `network_allow = []` on `MathData`, which makes verifiers treat every math task as
requiring a framework-aware network policy and refuse the `subprocess` runtime that every
config in this repo uses. prime-envs #795 reverts that default-deny, so the KL merge-bar
measurement on `math` no longer needs a wrapper taskset to restore the wildcard.

The bump spans ten commits, so it also brings fair-use solver prompts to the SWE and
terminal training tasksets, moves MCP Atlas and Enterprise Ops Gym into VMs, and adds
OfficeQA Pro V2, along with a large `registry.json` regeneration.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
garrett361 added a commit that referenced this pull request Sep 22, 2026
Mirrors PR #3543's layout: per-step math table, per-step reverse-text
table, and a summary. The bf16 columns reproduce #3543's own numbers,
which makes the FP8 columns directly comparable to it.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
vllm-project/vllm#54955 is a "[Preview][DO NOT MERGE]" branch, so it can
never be the event that lets us delete this patch. Cite #57798 instead,
which fixes the DeepSeek V4 target models on top of the shared padded-sink
helper in #57797.

Verified against upstream main d90f0eade5: the bare
params_dict[name][:n].copy_() is still there, so the patch is still needed.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
garrett361 added a commit that referenced this pull request Sep 22, 2026
Mirrors PR #3543's layout: per-step math table, per-step reverse-text
table, and a summary. The bf16 columns reproduce #3543's own numbers,
which makes the FP8 columns directly comparable to it.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant