Skip to content

Commit 9aa1f51

Browse files
K3 f_theta v1 trained checkpoint (Gemma4 26B-A4B verifier, per-layer KV; loss 50.8->3.70, 13.74x)
Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>
1 parent 844aaac commit 9aa1f51

4 files changed

Lines changed: 175 additions & 0 deletions

File tree

.gitattributes

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,2 @@
11
models/dflash-kakeya-baseline/*.safetensors filter=lfs diff=lfs merge=lfs -text
2+
results/research/f_theta_v1/f_theta_weights.pt filter=lfs diff=lfs merge=lfs -text

results/research/f_theta_v1.json

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
{
2+
"kind": "k3_f_theta_train",
3+
"config": {
4+
"verifier_id": "google/gemma-4-26B-A4B-it",
5+
"drafter_id": "models/dflash-kakeya-baseline",
6+
"steps": 4000,
7+
"lr": 0.001,
8+
"weight_decay": 0.01,
9+
"rank": 256,
10+
"n_prompts": 64,
11+
"gen_len": 128,
12+
"sample_positions": 256,
13+
"save": "results/research/f_theta_v1",
14+
"seed": 0,
15+
"log_every": 50,
16+
"eval_every": 500
17+
},
18+
"f_theta_config": {
19+
"drafter_num_layers": 5,
20+
"drafter_num_kv_heads": 8,
21+
"drafter_head_dim": 128,
22+
"verifier_num_layers": 30,
23+
"verifier_num_kv_heads": 8,
24+
"verifier_head_dim": 256,
25+
"rank": 256,
26+
"verifier_layer_kv_heads": [
27+
8,
28+
8,
29+
8,
30+
8,
31+
8,
32+
2,
33+
8,
34+
8,
35+
8,
36+
8,
37+
8,
38+
2,
39+
8,
40+
8,
41+
8,
42+
8,
43+
8,
44+
2,
45+
8,
46+
8,
47+
8,
48+
8,
49+
8,
50+
2,
51+
8,
52+
8,
53+
8,
54+
8,
55+
8,
56+
2
57+
],
58+
"verifier_layer_head_dims": [
59+
256,
60+
256,
61+
256,
62+
256,
63+
256,
64+
512,
65+
256,
66+
256,
67+
256,
68+
256,
69+
256,
70+
512,
71+
256,
72+
256,
73+
256,
74+
256,
75+
256,
76+
512,
77+
256,
78+
256,
79+
256,
80+
256,
81+
256,
82+
512,
83+
256,
84+
256,
85+
256,
86+
256,
87+
256,
88+
512
89+
]
90+
},
91+
"n_params": 31457280,
92+
"n_sequences": 62,
93+
"collect_seconds": 494.367059551063,
94+
"train_seconds": 59.45806223200634,
95+
"initial_loss": 50.82746124267578,
96+
"final_loss": 3.69950083732605,
97+
"loss_reduction_factor": 13.739005200337568
98+
}
Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
{
2+
"drafter_num_layers": 5,
3+
"drafter_num_kv_heads": 8,
4+
"drafter_head_dim": 128,
5+
"verifier_num_layers": 30,
6+
"verifier_num_kv_heads": 8,
7+
"verifier_head_dim": 256,
8+
"rank": 256,
9+
"verifier_layer_kv_heads": [
10+
8,
11+
8,
12+
8,
13+
8,
14+
8,
15+
2,
16+
8,
17+
8,
18+
8,
19+
8,
20+
8,
21+
2,
22+
8,
23+
8,
24+
8,
25+
8,
26+
8,
27+
2,
28+
8,
29+
8,
30+
8,
31+
8,
32+
8,
33+
2,
34+
8,
35+
8,
36+
8,
37+
8,
38+
8,
39+
2
40+
],
41+
"verifier_layer_head_dims": [
42+
256,
43+
256,
44+
256,
45+
256,
46+
256,
47+
512,
48+
256,
49+
256,
50+
256,
51+
256,
52+
256,
53+
512,
54+
256,
55+
256,
56+
256,
57+
256,
58+
256,
59+
512,
60+
256,
61+
256,
62+
256,
63+
256,
64+
256,
65+
512,
66+
256,
67+
256,
68+
256,
69+
256,
70+
256,
71+
512
72+
]
73+
}
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
version https://git-lfs.github.com/spec/v1
2+
oid sha256:1905f15eeb955b1f08ebdd7b45e752b21cb3a2c1092c606559489439322ec2e5
3+
size 125852105

0 commit comments

Comments
 (0)