eousphoros commited on
Commit
f48e413
·
verified ·
1 Parent(s): c03d494

Replace the MTP head with one trained on draft acceptance (total variation loss)

Browse files
README.md CHANGED
@@ -22,7 +22,7 @@ Give it a system prompt describing a persona and it will take that persona on.
22
 
23
  - **Retrained from the base model** with v8's recipe, on an updated persona dataset plus v8's coding, math, multimodal and paper-reading records.
24
  - **v8's early-stop quirk is fixed.** v8 now and then ended a reply one line before its final answer. A small weight edit (rank 4, in the MLP output projections of layers 61-63) removes the habit: on GPQA, replies without a final `Answer:` line went from 7 to 0, with scores unchanged.
25
- - **Recalibrated MTP head.** The draft head was trained to match this model's own next-token distributions on both of its draft steps (KL divergence, backbone frozen), on about 70M tokens of training data and the model's own replies to 2,200 prompts.
26
 
27
  Known quirks: it thinks longer than v8 on hard science questions (median about 9.6k vs 7.3k tokens on GPQA), and about 1 in 400 bench replies fell into a repetition loop that only the `max_tokens` cap ended.
28
 
@@ -40,16 +40,16 @@ Measured with our own harness: vLLM 0.32, thinking on at reasoning effort `xhigh
40
  | XSTest safe prompts fully answered (250) | 96.8% |
41
  | XSTest unsafe prompts fully refused (200) | 76.0% (partly refused 8.0%) |
42
 
43
- Against v8 on the same prompts: GPQA 85.5 vs 85.3 (505 replies) and LiveCodeBench 85.1 vs 86.1 (101 problems); neither difference is significant. Of the 1,717 bench replies, 14 hit the 128,000-token cap.
44
 
45
  ## Speculative decoding
46
 
47
  Tokens generated per target-model step with two MTP draft tokens, measured in vLLM 0.32 on 200 held-out prompts (about 2.1M generated tokens per configuration):
48
 
49
- | Draft sampling | This model's MTP head | Base model's MTP head |
50
- |---|---|---|
51
- | greedy (vLLM's default) | 2.14 | 2.13 |
52
- | probabilistic | 2.43 | 2.46 |
53
 
54
  At the default temperature, probabilistic drafting accepts more drafts: add `"draft_sample_method": "probabilistic"` to `--speculative-config`, for example `--speculative-config '{"method": "mtp", "num_speculative_tokens": 2, "draft_sample_method": "probabilistic"}'`. Either way, verification keeps the output distribution exactly the model's own; only speed changes.
55
 
 
22
 
23
  - **Retrained from the base model** with v8's recipe, on an updated persona dataset plus v8's coding, math, multimodal and paper-reading records.
24
  - **v8's early-stop quirk is fixed.** v8 now and then ended a reply one line before its final answer. A small weight edit (rank 4, in the MLP output projections of layers 61-63) removes the habit: on GPQA, replies without a final `Answer:` line went from 7 to 0, with scores unchanged.
25
+ - **Recalibrated MTP head.** The draft head was trained directly on how often vLLM's probabilistic drafting accepts its drafts: the loss is the total variation distance between its draft distribution and the distribution this model samples from (temperature 1, top-k 20, top-p 0.95), on both draft steps, with the backbone frozen. It started from the base model's head and trained on about 35M tokens of training records and the model's own replies. An earlier calibration that used KL divergence drafted slightly worse than the base model's head; it is kept at the `mtp-run1` revision.
26
 
27
  Known quirks: it thinks longer than v8 on hard science questions (median about 9.6k vs 7.3k tokens on GPQA), and about 1 in 400 bench replies fell into a repetition loop that only the `max_tokens` cap ended.
28
 
 
40
  | XSTest safe prompts fully answered (250) | 96.8% |
41
  | XSTest unsafe prompts fully refused (200) | 76.0% (partly refused 8.0%) |
42
 
43
+ Against v8 on the same prompts: GPQA 85.5 vs 85.3 (505 replies) and LiveCodeBench 85.1 vs 86.1 (101 problems); neither difference is significant. Of the 1,717 bench replies, 14 hit the 128,000-token cap. Speculative decoding never changes the output distribution, so these scores hold whichever MTP head is used.
44
 
45
  ## Speculative decoding
46
 
47
  Tokens generated per target-model step with two MTP draft tokens, measured in vLLM 0.32 on 200 held-out prompts (about 2.1M generated tokens per configuration):
48
 
49
+ | Draft sampling | This model's MTP head | Base model's MTP head | `mtp-run1` head |
50
+ |---|---|---|---|
51
+ | greedy (vLLM's default) | 2.15 | 2.13 | 2.14 |
52
+ | probabilistic | 2.48 | 2.46 | 2.43 |
53
 
54
  At the default temperature, probabilistic drafting accepts more drafts: add `"draft_sample_method": "probabilistic"` to `--speculative-config`, for example `--speculative-config '{"method": "mtp", "num_speculative_tokens": 2, "draft_sample_method": "probabilistic"}'`. Either way, verification keeps the output distribution exactly the model's own; only speed changes.
55
 
model-part-00053.safetensors CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:3ffa27ee112c304777f21553a6bce846e568e87c3d50c2c89fd0f71332317bfd
3
  size 1027627592
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:50d5e228574f89fdda8372120c2d9a664563467f5a64254e6d6841c23569875d
3
  size 1027627592
model-part-00054.safetensors CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:de768039c4542e0ec80f62d698176eae5218d1d71166a2dfc03e1b4cea318ab0
3
  size 388016320
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c5d97acb0a6b998fb118c83d5061b2d817597b7bde0989daa85d6bc9c08242a9
3
  size 388016320