Skip to content

Commit 3fe1b62

Browse files
cursoragentFluffyAIcode
cursoragent
and
FluffyAIcode
committed
K3 Stage 2 evidence: alignment training lifts DFlash acceptance 0.003 -> 0.074
Treated the residual fidelity gap as f_theta alignment (per user direction): froze the verifier + drafter backbone, trained fc+hidden_norm+norm (47.6M params, 600 steps, 8 verifier-greedy prompts, ~12s). Train match-rate 0->0.17. Re-eval on H200 (aligned drafter, 4 prompts, block 16): acceptance_rate 0.003 -> 0.074 acceptance_length 1.04 -> 2.04 (blocks now accept up to 3-4 drafts) lossless_vs_ar True (reference ~0.447 / ~7.7). Clear validation that the native integration + alignment works and improves with training. Larger full-scope run next to push toward the reference. (Aligned weights ~860MB live on the H200, not committed.) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>
1 parent 73dc785 commit 3fe1b62

4 files changed

Lines changed: 231 additions & 0 deletions

File tree

‎results/research/dflash_aligned_fcnorms_report.json‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
{
2+
"kind": "k3_dflash_alignment_train",
3+
"config": {
4+
"verifier_id": "google/gemma-4-26B-A4B-it",
5+
"drafter_id": "z-lab/gemma-4-26B-A4B-it-DFlash",
6+
"steps": 600,
7+
"lr": 0.0001,
8+
"block_size": 16,
9+
"n_prompts": 8,
10+
"gen_len": 160,
11+
"prompt_min_ctx": 8,
12+
"train_scope": "fc_norms",
13+
"seed": 0,
14+
"save": "results/research/dflash_aligned_fcnorms.pt",
15+
"log_every": 25
16+
},
17+
"trainable_params": 47584768,
18+
"n_windows": 1041,
19+
"final_loss": 3.496910614967346,
20+
"final_train_match": 0.1725,
21+
"elapsed_s": 12.202172419056296
22+
}

‎results/research/k3_dflash_specdecode_aligned1.json‎

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,162 @@
1+
{
2+
"schema_version": 1,
3+
"kind": "k3_dflash_specdecode_acceptance",
4+
"config": {
5+
"verifier_id": "google/gemma-4-26B-A4B-it",
6+
"drafter_id": "z-lab/gemma-4-26B-A4B-it-DFlash",
7+
"block_size": 16,
8+
"num_steps": 1,
9+
"max_new_tokens": 48,
10+
"n_prompts": 4,
11+
"aux_layer_ids": [
12+
2,
13+
7,
14+
12,
15+
18,
16+
23,
17+
28
18+
]
19+
},
20+
"aggregate": {
21+
"acceptance_rate": 0.07372881355932204,
22+
"acceptance_length": 2.0357142857142856,
23+
"total_accepted": 87,
24+
"total_drafted": 1180,
25+
"total_blocks": 84,
26+
"lossless_vs_ar": true,
27+
"reference_humaneval": {
28+
"acceptance_length": 7.7,
29+
"acceptance_rate": 0.447
30+
}
31+
},
32+
"per_prompt": [
33+
{
34+
"prompt": "Write a Python function that returns the n-th Fibonacci number.",
35+
"blocks": 22,
36+
"block_accepts": [
37+
1,
38+
1,
39+
0,
40+
2,
41+
0,
42+
1,
43+
0,
44+
1,
45+
0,
46+
1,
47+
1,
48+
0,
49+
2,
50+
4,
51+
4,
52+
0,
53+
3,
54+
1,
55+
2,
56+
0,
57+
2,
58+
0
59+
],
60+
"mean_accepted_per_block": 1.1818181818181819,
61+
"tokens_generated": 48,
62+
"verifier_forwards_spec": 22,
63+
"lossless_vs_ar": true,
64+
"decoded": "There are several ways to implement this depending on whether you prioritize readability, memory, or speed. Below are the three most common approaches.\n\n### 1. The Efficient Approach (Iterative)\nThis "
65+
},
66+
{
67+
"prompt": "Explain in two sentences why the sky is blue.",
68+
"blocks": 30,
69+
"block_accepts": [
70+
1,
71+
1,
72+
0,
73+
0,
74+
1,
75+
1,
76+
1,
77+
0,
78+
0,
79+
2,
80+
0,
81+
0,
82+
1,
83+
4,
84+
0,
85+
0,
86+
1,
87+
1,
88+
0,
89+
1,
90+
1,
91+
0,
92+
1,
93+
0,
94+
0,
95+
0,
96+
0,
97+
1,
98+
0,
99+
1
100+
],
101+
"mean_accepted_per_block": 0.6333333333333333,
102+
"tokens_generated": 48,
103+
"verifier_forwards_spec": 30,
104+
"lossless_vs_ar": true,
105+
"decoded": "The sky appears blue because of a phenomenon called Rayleigh scattering, where sunlight interacts with the gases and particles in Earth's atmosphere. As sunlight reaches the atmosphere, shorter blue w"
106+
},
107+
{
108+
"prompt": "List three prime numbers greater than 100.",
109+
"blocks": 14,
110+
"block_accepts": [
111+
3,
112+
2,
113+
4,
114+
1,
115+
1,
116+
0,
117+
4,
118+
1,
119+
3,
120+
1,
121+
4,
122+
1,
123+
1,
124+
0
125+
],
126+
"mean_accepted_per_block": 1.8571428571428572,
127+
"tokens_generated": 40,
128+
"verifier_forwards_spec": 14,
129+
"lossless_vs_ar": true,
130+
"decoded": "Here are three prime numbers greater than 100:\n\n1. **101**\n2. **103**\n3. **107**"
131+
},
132+
{
133+
"prompt": "Summarize the plot of Romeo and Juliet in one sentence.",
134+
"blocks": 18,
135+
"block_accepts": [
136+
0,
137+
1,
138+
1,
139+
0,
140+
1,
141+
1,
142+
1,
143+
0,
144+
0,
145+
2,
146+
1,
147+
0,
148+
2,
149+
1,
150+
0,
151+
2,
152+
1,
153+
2
154+
],
155+
"mean_accepted_per_block": 0.8888888888888888,
156+
"tokens_generated": 34,
157+
"verifier_forwards_spec": 18,
158+
"lossless_vs_ar": true,
159+
"decoded": "Two star-crossed lovers from feuding noble families take their own lives in a tragic misunderstanding, ultimately transforming their deaths into a catalyst for peace between their households."
160+
}
161+
]
162+
}

‎results/research/logs/k3_dflash_align_fcnorms.log‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
[align] loading verifier google/gemma-4-26B-A4B-it
2+
Loading weights: 0%| | 0/1013 [00:00<?, ?it/s]Loading weights: 0%| | 2/1013 [00:00<00:59, 16.87it/s]Loading weights: 0%| | 4/1013 [00:00<01:42, 9.81it/s]Loading weights: 2%|▏ | 23/1013 [00:00<00:16, 60.45it/s]Loading weights: 3%|▎ | 32/1013 [00:00<00:17, 55.40it/s]Loading weights: 4%|▍ | 41/1013 [00:00<00:23, 42.15it/s]Loading weights: 6%|▌ | 63/1013 [00:01<00:17, 53.48it/s]Loading weights: 9%|▉ | 91/1013 [00:01<00:10, 85.89it/s]Loading weights: 10%|█ | 103/1013 [00:01<00:11, 77.87it/s]Loading weights: 11%|█ | 113/1013 [00:01<00:11, 80.01it/s]Loading weights: 12%|█▏ | 123/1013 [00:01<00:12, 69.57it/s]Loading weights: 13%|█▎ | 132/1013 [00:02<00:16, 52.17it/s]Loading weights: 15%|█▌ | 156/1013 [00:02<00:10, 79.74it/s]Loading weights: 16%|█▋ | 167/1013 [00:02<00:11, 71.35it/s]Loading weights: 18%|█▊ | 178/1013 [00:02<00:10, 76.43it/s]Loading weights: 19%|█▊ | 188/1013 [00:02<00:12, 67.04it/s]Loading weights: 20%|█▉ | 199/1013 [00:03<00:15, 53.90it/s]Loading weights: 21%|██▏ | 217/1013 [00:03<00:12, 64.06it/s]Loading weights: 23%|██▎ | 233/1013 [00:03<00:09, 80.09it/s]Loading weights: 24%|██▍ | 244/1013 [00:03<00:09, 84.55it/s]Loading weights: 25%|██▌ | 254/1013 [00:03<00:10, 74.50it/s]Loading weights: 26%|██▌ | 263/1013 [00:03<00:09, 76.45it/s]Loading weights: 27%|██▋ | 272/1013 [00:04<00:11, 66.75it/s]Loading weights: 28%|██▊ | 287/1013 [00:04<00:08, 83.41it/s]Loading weights: 29%|██▉ | 297/1013 [00:04<00:09, 72.05it/s]Loading weights: 31%|███ | 309/1013 [00:04<00:08, 82.08it/s]Loading weights: 31%|███▏ | 319/1013 [00:04<00:09, 71.82it/s]Loading weights: 33%|███▎ | 331/1013 [00:04<00:08, 81.78it/s]Loading weights: 34%|███▎ | 341/1013 [00:04<00:09, 70.41it/s]Loading weights: 35%|███▍ | 351/1013 [00:05<00:08, 76.20it/s]Loading weights: 36%|███▌ | 360/1013 [00:05<00:09, 66.51it/s]Loading weights: 37%|███▋ | 375/1013 [00:05<00:10, 58.07it/s]Loading weights: 39%|███▉ | 396/1013 [00:05<00:07, 82.15it/s]Loading weights: 40%|████ | 406/1013 [00:05<00:08, 72.66it/s]Loading weights: 41%|████▏ | 418/1013 [00:06<00:07, 79.58it/s]Loading weights: 42%|████▏ | 428/1013 [00:06<00:08, 68.97it/s]Loading weights: 43%|████▎ | 440/1013 [00:06<00:07, 74.96it/s]Loading weights: 44%|████▍ | 449/1013 [00:06<00:08, 64.51it/s]Loading weights: 46%|████▌ | 462/1013 [00:06<00:07, 75.03it/s]Loading weights: 46%|████▋ | 471/1013 [00:06<00:08, 64.22it/s]Loading weights: 48%|████▊ | 484/1013 [00:06<00:06, 75.81it/s]Loading weights: 49%|████▊ | 493/1013 [00:07<00:08, 63.44it/s]Loading weights: 50%|████▉ | 506/1013 [00:07<00:06, 74.50it/s]Loading weights: 51%|█████ | 515/1013 [00:07<00:07, 63.74it/s]Loading weights: 52%|█████▏ | 527/1013 [00:07<00:06, 72.62it/s]Loading weights: 53%|█████▎ | 536/1013 [00:07<00:07, 62.37it/s]Loading weights: 54%|█████▍ | 549/1013 [00:07<00:06, 74.19it/s]Loading weights: 55%|█████▌ | 558/1013 [00:08<00:07, 63.19it/s]Loading weights: 56%|█████▋ | 571/1013 [00:08<00:05, 74.80it/s]Loading weights: 57%|█████▋ | 580/1013 [00:08<00:07, 61.50it/s]Loading weights: 59%|█████▊ | 593/1013 [00:08<00:05, 72.28it/s]Loading weights: 59%|█████▉ | 602/1013 [00:08<00:06, 62.47it/s]Loading weights: 61%|██████ | 615/1013 [00:08<00:05, 73.67it/s]Loading weights: 62%|██████▏ | 624/1013 [00:09<00:06, 63.12it/s]Loading weights: 63%|██████▎ | 637/1013 [00:09<00:05, 74.88it/s]Loading weights: 64%|██████▍ | 646/1013 [00:09<00:05, 63.82it/s]Loading weights: 76%|███████▋ | 774/1013 [00:09<00:00, 309.72it/s]Loading weights: 91%|█████████ | 922/1013 [00:09<00:00, 570.96it/s]Loading weights: 100%|██████████| 1013/1013 [00:09<00:00, 104.50it/s]
3+
[align] loading drafter z-lab/gemma-4-26B-A4B-it-DFlash (fp32 for training)
4+
[align] trainable params (fc_norms): 47,584,768
5+
[align] seq 0: len=186 (prompt 26)
6+
[align] seq 1: len=183 (prompt 23)
7+
[align] seq 2: len=81 (prompt 24)
8+
[align] seq 3: len=185 (prompt 25)
9+
[align] seq 4: len=187 (prompt 27)
10+
[align] seq 5: len=181 (prompt 21)
11+
[align] seq 6: len=183 (prompt 23)
12+
[align] seq 7: len=182 (prompt 22)
13+
[align] 1041 training windows
14+
[align] step=25 loss=8.7057 match=0.000
15+
[align] step=50 loss=6.2933 match=0.000
16+
[align] step=75 loss=5.5856 match=0.000
17+
[align] step=100 loss=5.4675 match=0.000
18+
[align] step=125 loss=5.1793 match=0.000
19+
[align] step=150 loss=4.8446 match=0.000
20+
[align] step=175 loss=4.7523 match=0.040
21+
[align] step=200 loss=5.1121 match=0.033
22+
[align] step=225 loss=4.6334 match=0.068
23+
[align] step=250 loss=4.2954 match=0.102
24+
[align] step=275 loss=4.0659 match=0.115
25+
[align] step=300 loss=4.0129 match=0.133
26+
[align] step=325 loss=3.9662 match=0.150
27+
[align] step=350 loss=3.7828 match=0.133
28+
[align] step=375 loss=3.4494 match=0.172
29+
[align] step=400 loss=3.9362 match=0.115
30+
[align] step=425 loss=3.6597 match=0.160
31+
[align] step=450 loss=3.7880 match=0.158
32+
[align] step=475 loss=3.8354 match=0.117
33+
[align] step=500 loss=3.7113 match=0.152
34+
[align] step=525 loss=3.3665 match=0.172
35+
[align] step=550 loss=3.3897 match=0.168
36+
[align] step=575 loss=3.5502 match=0.177
37+
[align] step=600 loss=3.4969 match=0.172
38+
[align] DONE in 12s; final loss=3.4969 train_match=0.172; saved -> results/research/dflash_aligned_fcnorms.pt

‎results/research/logs/k3_dflash_specdecode_aligned1.log‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
[k3-sd] loading verifier google/gemma-4-26B-A4B-it
2+
Loading weights: 0%| | 0/1013 [00:00<?, ?it/s]Loading weights: 0%| | 2/1013 [00:00<01:50, 9.11it/s]Loading weights: 0%| | 4/1013 [00:00<02:10, 7.72it/s]Loading weights: 2%|▏ | 25/1013 [00:00<00:17, 55.99it/s]Loading weights: 3%|▎ | 35/1013 [00:00<00:17, 55.64it/s]Loading weights: 5%|▍ | 48/1013 [00:01<00:18, 52.66it/s]Loading weights: 7%|▋ | 69/1013 [00:01<00:11, 80.95it/s]Loading weights: 8%|▊ | 80/1013 [00:01<00:12, 72.28it/s]Loading weights: 9%|▉ | 91/1013 [00:01<00:11, 79.67it/s]Loading weights: 10%|▉ | 101/1013 [00:01<00:12, 74.04it/s]Loading weights: 11%|█▏ | 114/1013 [00:01<00:14, 62.89it/s]Loading weights: 13%|█▎ | 134/1013 [00:02<00:10, 86.56it/s]Loading weights: 14%|█▍ | 145/1013 [00:02<00:11, 78.50it/s]Loading weights: 15%|█▌ | 156/1013 [00:02<00:10, 84.85it/s]Loading weights: 16%|█▋ | 166/1013 [00:02<00:10, 77.43it/s]Loading weights: 18%|█▊ | 178/1013 [00:02<00:09, 86.64it/s]Loading weights: 19%|█▊ | 188/1013 [00:02<00:11, 74.76it/s]Loading weights: 20%|█▉ | 200/1013 [00:02<00:09, 83.74it/s]Loading weights: 21%|██ | 210/1013 [00:03<00:11, 72.85it/s]Loading weights: 22%|██▏ | 222/1013 [00:03<00:09, 80.58it/s]Loading weights: 23%|██▎ | 231/1013 [00:03<00:11, 66.46it/s]Loading weights: 24%|██▍ | 244/1013 [00:03<00:09, 77.67it/s]Loading weights: 25%|██▍ | 253/1013 [00:03<00:11, 67.79it/s]Loading weights: 26%|██▌ | 265/1013 [00:03<00:09, 75.91it/s]Loading weights: 27%|██▋ | 274/1013 [00:03<00:11, 66.41it/s]Loading weights: 28%|██▊ | 286/1013 [00:04<00:09, 76.16it/s]Loading weights: 29%|██▉ | 295/1013 [00:04<00:11, 64.60it/s]Loading weights: 31%|███ | 309/1013 [00:04<00:08, 79.86it/s]Loading weights: 31%|███▏ | 319/1013 [00:04<00:10, 67.01it/s]Loading weights: 33%|███▎ | 331/1013 [00:04<00:09, 75.55it/s]Loading weights: 34%|███▎ | 340/1013 [00:04<00:10, 64.10it/s]Loading weights: 35%|███▍ | 353/1013 [00:05<00:08, 77.10it/s]Loading weights: 36%|███▌ | 362/1013 [00:05<00:10, 65.00it/s]Loading weights: 37%|███▋ | 375/1013 [00:05<00:08, 78.16it/s]Loading weights: 38%|███▊ | 385/1013 [00:05<00:08, 69.78it/s]Loading weights: 39%|███▉ | 396/1013 [00:05<00:11, 54.95it/s]Loading weights: 41%|████▏ | 418/1013 [00:05<00:07, 83.09it/s]Loading weights: 42%|████▏ | 429/1013 [00:06<00:07, 78.73it/s]Loading weights: 44%|████▎ | 441/1013 [00:06<00:08, 64.97it/s]Loading weights: 46%|████▌ | 462/1013 [00:06<00:08, 68.71it/s]Loading weights: 48%|████▊ | 484/1013 [00:06<00:05, 91.48it/s]Loading weights: 49%|████▉ | 496/1013 [00:06<00:06, 83.86it/s]Loading weights: 50%|█████ | 507/1013 [00:07<00:07, 66.17it/s]Loading weights: 52%|█████▏ | 527/1013 [00:07<00:05, 86.66it/s]Loading weights: 53%|█████▎ | 539/1013 [00:07<00:05, 80.40it/s]Loading weights: 54%|█████▍ | 550/1013 [00:07<00:07, 65.23it/s]Loading weights: 56%|█████▋ | 571/1013 [00:07<00:04, 88.69it/s]Loading weights: 58%|█████▊ | 583/1013 [00:08<00:05, 81.97it/s]Loading weights: 59%|█████▊ | 594/1013 [00:08<00:06, 67.11it/s]Loading weights: 60%|██████ | 609/1013 [00:08<00:06, 64.98it/s]Loading weights: 62%|██████▏ | 629/1013 [00:08<00:05, 69.61it/s]Loading weights: 76%|███████▌ | 771/1013 [00:08<00:00, 283.62it/s]Loading weights: 93%|█████████▎| 942/1013 [00:08<00:00, 542.64it/s]Loading weights: 100%|██████████| 1013/1013 [00:09<00:00, 112.06it/s]
3+
[k3-sd] loading drafter z-lab/gemma-4-26B-A4B-it-DFlash
4+
[k3-sd] loaded aligned drafter state from results/research/dflash_aligned_fcnorms.pt
5+
[k3-sd] prompt 0: blocks=22 mean_accept=1.18 accepts=[1, 1, 0, 2, 0, 1, 0, 1, 0, 1, 1, 0, 2, 4, 4, 0, 3, 1, 2, 0, 2, 0] lossless=True
6+
[k3-sd] prompt 1: blocks=30 mean_accept=0.63 accepts=[1, 1, 0, 0, 1, 1, 1, 0, 0, 2, 0, 0, 1, 4, 0, 0, 1, 1, 0, 1, 1, 0, 1, 0, 0, 0, 0, 1, 0, 1] lossless=True
7+
[k3-sd] prompt 2: blocks=14 mean_accept=1.86 accepts=[3, 2, 4, 1, 1, 0, 4, 1, 3, 1, 4, 1, 1, 0] lossless=True
8+
[k3-sd] prompt 3: blocks=18 mean_accept=0.89 accepts=[0, 1, 1, 0, 1, 1, 1, 0, 0, 2, 1, 0, 2, 1, 0, 2, 1, 2] lossless=True
9+
[k3-sd] AGGREGATE acceptance_rate=0.074 acceptance_length=2.04 lossless=True (ref ~0.447 / ~7.7) -> results/research/k3_dflash_specdecode_aligned1.json

0 commit comments

Comments
 (0)