SPB Git

spb/forge Public MIT

Forge — LLM training from scratch in pure C++20 + Metal on Apple Silicon.

llmtrainingcppmetalapple-silicon
C++ 61.2% C 23% Python 7.6% TeX 7.2% CMake 1.1%

Add wave 3: DeepSeek-V3-style MoE routing, all config-selected

- moe_scoring "softmax"|"sigmoid" (new sigmoid op, CPU+Metal+backward)
- aux-loss-free balancing (V3 "noaux"): top-k selection ranks score+bias
  while gate values stay biasless; per-expert load counted on-GPU each
  forward and the balance bias nudged ±moe_bias_gamma after every
  optimizer step; bias is checkpointed as a grad-free parameter
- moe_norm_topk (renormalize kept gates or keep raw sigmoid scores),
  routed_scaling_factor (V3: 2.5), moe_d_ff (per-expert width),
  first_k_dense (dense MLPs for the first k layers)
- parity tests: biased/no-renorm topk, sigmoid, expert_counts, and a
  full V3-style model (sigmoid/noaux/scaled/first-dense) — CPU==GPU

Remaining wave-3 items (MLA, MuonClip QK-clip, MTP) documented in
ARCHITECTURES.md.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
simon-pierre boucher committed 5 days ago (Aug 5, 2026) parent f7777e9

Showing 19 changed files with +2,428 and −104

modified ARCHITECTURES.md +4 −4
@@ -44,10 +44,10 @@ Base actuelle : pre-norm RMSNorm/LayerNorm · SwiGLU/GELU · RoPE complet ou pos
44 44 10. **Attention sinks** (GPT-OSS) : initialiser le (max, somme) du softmax online avec le logit sink appris par tête — quelques lignes dans le kernel flash.
45 45
46 46 **Vague 3 — investissements**
47 −11. **Routage sigmoïde + équilibrage sans aux loss** (V3/K2) : sélection top-k sur `s+b`, gates sur `s` seul ; mise à jour du biais ±γ selon la charge après chaque step (état dans l'optimiseur/trainer).
48 −12. **QK-Clip / MuonClip** (K2) : monitorer le logit d'attention max par tête (le kernel flash garde déjà le max par ligne — l'exposer), rescaler W_q/W_k si > τ.
49 −13. **MLA** (V3/K2) : le seul gros chantier kernel — d_k(192) ≠ d_v(128) + RoPE partiel sur une tranche de 64. Praticable d'abord en « décompressé » si le kernel accepte d_k≠d_v.
50 −14. **MTP** (V3) : bloc supplémentaire + tête partagée, coût trainer modéré.
47 +11. ✅ **Routage sigmoïde + équilibrage sans aux loss** (V3/K2) — FAIT : `moe_scoring: "sigmoid"`, sélection top-k sur `s+b` / gates sur `s` seul (`topk_renorm` biaisé), `moe_bias_gamma` (update ±γ par charge observée après chaque step, compteurs GPU), `routed_scaling_factor`, `moe_d_ff` séparé, `first_k_dense`, `moe_norm_topk`. Le biais d'équilibrage est checkpointé (param sans gradient).
48 +12. **QK-Clip / MuonClip** (K2) : monitorer le logit d'attention max par tête (le kernel flash garde déjà le max par ligne — l'exposer), rescaler W_q/W_k si > τ. — restant
49 +13. **MLA** (V3/K2) : le seul gros chantier kernel — d_k(192) ≠ d_v(128) + RoPE partiel sur une tranche de 64. Praticable d'abord en « décompressé » si le kernel accepte d_k≠d_v. — restant
50 +14. **MTP** (V3) : bloc supplémentaire + tête partagée, coût trainer modéré. — restant
51 51
52 52 ## Formules clés
53 53
added models/50m-tinystories.forge/manifest-latest.json +1039 −0
@@ -0,0 +1,1039 @@
1 +{
2 + "config": {
3 + "_comment": "~52M params, sized between gpt-25m and gpt-100m. vocab 4096 to match the local tok4096 train.bin (19.14M tokens). batch_size is the MICRO-batch: 8 x 8 x 1024 = 65536 tokens/step; 1170 steps = ~4 epochs over the 19.14M-token set. Warmup is 10% of the run. precision is parsed but not yet honored - all kernels are f32.",
4 + "model": {
5 + "activation": "swiglu",
6 + "context_length": 1024,
7 + "d_ff": 1728,
8 + "d_model": 640,
9 + "dropout": 0.0,
10 + "n_heads": 10,
11 + "n_kv_heads": 10,
12 + "n_layers": 10,
13 + "name": "gpt-50m",
14 + "norm": "rmsnorm",
15 + "norm_eps": 1e-06,
16 + "rope_theta": 10000.0,
17 + "tied_embeddings": true,
18 + "use_rope": true,
19 + "vocab_size": 4096
20 + },
21 + "train": {
22 + "batch_size": 8,
23 + "beta1": 0.9,
24 + "beta2": 0.95,
25 + "checkpoint_every": 200,
26 + "eps": 1e-08,
27 + "eval_batches": 20,
28 + "eval_every": 100,
29 + "grad_accum_steps": 8,
30 + "grad_clip": 1.0,
31 + "lr": 0.0005,
32 + "max_steps": 1170,
33 + "min_lr_ratio": 0.1,
34 + "precision": "f32",
35 + "seed": 1337,
36 + "warmup_steps": 117,
37 + "weight_decay": 0.1
38 + }
39 + },
40 + "dtype": "f32",
41 + "fmodel": 1,
42 + "parent": "",
43 + "self": "manifest-000000.json",
44 + "step": 1170,
45 + "tag": "one-epoch-x4",
46 + "tensors": {
47 + "blocks.0.attn.wk.weight": {
48 + "dtype": "f32",
49 + "hash": "ea0278bff183bd9b",
50 + "nbytes": 1638400,
51 + "offset": 12156928,
52 + "shape": [
53 + 640,
54 + 640
55 + ],
56 + "shard": "objects/72626eb74479f8dd.fshard"
57 + },
58 + "blocks.0.attn.wo.weight": {
59 + "dtype": "f32",
60 + "hash": "b08bb6730b3f4047",
61 + "nbytes": 1638400,
62 + "offset": 15433728,
63 + "shape": [
64 + 640,
65 + 640
66 + ],
67 + "shard": "objects/72626eb74479f8dd.fshard"
68 + },
69 + "blocks.0.attn.wq.weight": {
70 + "dtype": "f32",
71 + "hash": "c9f0aee32c2afb09",
72 + "nbytes": 1638400,
73 + "offset": 10518528,
74 + "shape": [
75 + 640,
76 + 640
77 + ],
78 + "shard": "objects/72626eb74479f8dd.fshard"
79 + },
80 + "blocks.0.attn.wv.weight": {
81 + "dtype": "f32",
82 + "hash": "a46e677dff653afa",
83 + "nbytes": 1638400,
84 + "offset": 13795328,
85 + "shape": [
86 + 640,
87 + 640
88 + ],
89 + "shard": "objects/72626eb74479f8dd.fshard"
90 + },
91 + "blocks.0.mlp.w1.weight": {
92 + "dtype": "f32",
93 + "hash": "4b63b198ca18d908",
94 + "nbytes": 4423680,
95 + "offset": 17088512,
96 + "shape": [
97 + 1728,
98 + 640
99 + ],
100 + "shard": "objects/72626eb74479f8dd.fshard"
101 + },
102 + "blocks.0.mlp.w2.weight": {
103 + "dtype": "f32",
104 + "hash": "d0b0ced42d5c91d2",
105 + "nbytes": 4423680,
106 + "offset": 25935872,
107 + "shape": [
108 + 640,
109 + 1728
110 + ],
111 + "shard": "objects/72626eb74479f8dd.fshard"
112 + },
113 + "blocks.0.mlp.w3.weight": {
114 + "dtype": "f32",
115 + "hash": "d36f453ea5b7576d",
116 + "nbytes": 4423680,
117 + "offset": 21512192,
118 + "shape": [
119 + 1728,
120 + 640
121 + ],
122 + "shard": "objects/72626eb74479f8dd.fshard"
123 + },
124 + "blocks.0.norm1.weight": {
125 + "dtype": "f32",
126 + "hash": "bf60929ceab06890",
127 + "nbytes": 2560,
128 + "offset": 10502144,
129 + "shape": [
130 + 640
131 + ],
132 + "shard": "objects/72626eb74479f8dd.fshard"
133 + },
134 + "blocks.0.norm2.weight": {
135 + "dtype": "f32",
136 + "hash": "bb257feeea6d5b6f",
137 + "nbytes": 2560,
138 + "offset": 17072128,
139 + "shape": [
140 + 640
141 + ],
142 + "shard": "objects/72626eb74479f8dd.fshard"
143 + },
144 + "blocks.1.attn.wk.weight": {
145 + "dtype": "f32",
146 + "hash": "a99d45c6af060836",
147 + "nbytes": 1638400,
148 + "offset": 32014336,
149 + "shape": [
150 + 640,
151 + 640
152 + ],
153 + "shard": "objects/72626eb74479f8dd.fshard"
154 + },
155 + "blocks.1.attn.wo.weight": {
156 + "dtype": "f32",
157 + "hash": "4ecc0144607917a7",
158 + "nbytes": 1638400,
159 + "offset": 35291136,
160 + "shape": [
161 + 640,
162 + 640
163 + ],
164 + "shard": "objects/72626eb74479f8dd.fshard"
165 + },
166 + "blocks.1.attn.wq.weight": {
167 + "dtype": "f32",
168 + "hash": "5dd4c343a6557550",
169 + "nbytes": 1638400,
170 + "offset": 30375936,
171 + "shape": [
172 + 640,
173 + 640
174 + ],
175 + "shard": "objects/72626eb74479f8dd.fshard"
176 + },
177 + "blocks.1.attn.wv.weight": {
178 + "dtype": "f32",
179 + "hash": "0de8eef2c1851cef",
180 + "nbytes": 1638400,
181 + "offset": 33652736,
182 + "shape": [
183 + 640,
184 + 640
185 + ],
186 + "shard": "objects/72626eb74479f8dd.fshard"
187 + },
188 + "blocks.1.mlp.w1.weight": {
189 + "dtype": "f32",
190 + "hash": "d8fbfcb43a5a4677",
191 + "nbytes": 4423680,
192 + "offset": 36945920,
193 + "shape": [
194 + 1728,
195 + 640
196 + ],
197 + "shard": "objects/72626eb74479f8dd.fshard"
198 + },
199 + "blocks.1.mlp.w2.weight": {
200 + "dtype": "f32",
201 + "hash": "df86dbe751fd0897",
202 + "nbytes": 4423680,
203 + "offset": 45793280,
204 + "shape": [
205 + 640,
206 + 1728
207 + ],
208 + "shard": "objects/72626eb74479f8dd.fshard"
209 + },
210 + "blocks.1.mlp.w3.weight": {
211 + "dtype": "f32",
212 + "hash": "e206903aafcacc9a",
213 + "nbytes": 4423680,
214 + "offset": 41369600,
215 + "shape": [
216 + 1728,
217 + 640
218 + ],
219 + "shard": "objects/72626eb74479f8dd.fshard"
220 + },
221 + "blocks.1.norm1.weight": {
222 + "dtype": "f32",
223 + "hash": "bc7c7f4144e6c4a0",
224 + "nbytes": 2560,
225 + "offset": 30359552,
226 + "shape": [
227 + 640
228 + ],
229 + "shard": "objects/72626eb74479f8dd.fshard"
230 + },
231 + "blocks.1.norm2.weight": {
232 + "dtype": "f32",
233 + "hash": "1d631908e97d0ab2",
234 + "nbytes": 2560,
235 + "offset": 36929536,
236 + "shape": [
237 + 640
238 + ],
239 + "shard": "objects/72626eb74479f8dd.fshard"
240 + },
241 + "blocks.2.attn.wk.weight": {
242 + "dtype": "f32",
243 + "hash": "e5f3bfa24ce66850",
244 + "nbytes": 1638400,
245 + "offset": 51871744,
246 + "shape": [
247 + 640,
248 + 640
249 + ],
250 + "shard": "objects/72626eb74479f8dd.fshard"
251 + },
252 + "blocks.2.attn.wo.weight": {
253 + "dtype": "f32",
254 + "hash": "81e09028c8d4fa64",
255 + "nbytes": 1638400,
256 + "offset": 55148544,
257 + "shape": [
258 + 640,
259 + 640
260 + ],
261 + "shard": "objects/72626eb74479f8dd.fshard"
262 + },
263 + "blocks.2.attn.wq.weight": {
264 + "dtype": "f32",
265 + "hash": "70ee42db3250d15e",
266 + "nbytes": 1638400,
267 + "offset": 50233344,
268 + "shape": [
269 + 640,
270 + 640
271 + ],
272 + "shard": "objects/72626eb74479f8dd.fshard"
273 + },
274 + "blocks.2.attn.wv.weight": {
275 + "dtype": "f32",
276 + "hash": "b27d667f5ce35172",
277 + "nbytes": 1638400,
278 + "offset": 53510144,
279 + "shape": [
280 + 640,
281 + 640
282 + ],
283 + "shard": "objects/72626eb74479f8dd.fshard"
284 + },
285 + "blocks.2.mlp.w1.weight": {
286 + "dtype": "f32",
287 + "hash": "eef179f1c68e1af2",
288 + "nbytes": 4423680,
289 + "offset": 56803328,
290 + "shape": [
291 + 1728,
292 + 640
293 + ],
294 + "shard": "objects/72626eb74479f8dd.fshard"
295 + },
296 + "blocks.2.mlp.w2.weight": {
297 + "dtype": "f32",
298 + "hash": "0393a401c220c511",
299 + "nbytes": 4423680,
300 + "offset": 65650688,
301 + "shape": [
302 + 640,
303 + 1728
304 + ],
305 + "shard": "objects/72626eb74479f8dd.fshard"
306 + },
307 + "blocks.2.mlp.w3.weight": {
308 + "dtype": "f32",
309 + "hash": "8b9edaa2ac9409e3",
310 + "nbytes": 4423680,
311 + "offset": 61227008,
312 + "shape": [
313 + 1728,
314 + 640
315 + ],
316 + "shard": "objects/72626eb74479f8dd.fshard"
317 + },
318 + "blocks.2.norm1.weight": {
319 + "dtype": "f32",
320 + "hash": "fb9413da5028b6fd",
321 + "nbytes": 2560,
322 + "offset": 50216960,
323 + "shape": [
324 + 640
325 + ],
326 + "shard": "objects/72626eb74479f8dd.fshard"
327 + },
328 + "blocks.2.norm2.weight": {
329 + "dtype": "f32",
330 + "hash": "a76041b595744213",
331 + "nbytes": 2560,
332 + "offset": 56786944,
333 + "shape": [
334 + 640
335 + ],
336 + "shard": "objects/72626eb74479f8dd.fshard"
337 + },
338 + "blocks.3.attn.wk.weight": {
339 + "dtype": "f32",
340 + "hash": "01edcec7013dc445",
341 + "nbytes": 1638400,
342 + "offset": 71729152,
343 + "shape": [
344 + 640,
345 + 640
346 + ],
347 + "shard": "objects/72626eb74479f8dd.fshard"
348 + },
349 + "blocks.3.attn.wo.weight": {
350 + "dtype": "f32",
351 + "hash": "ea9891714020fcf4",
352 + "nbytes": 1638400,
353 + "offset": 75005952,
354 + "shape": [
355 + 640,
356 + 640
357 + ],
358 + "shard": "objects/72626eb74479f8dd.fshard"
359 + },
360 + "blocks.3.attn.wq.weight": {
361 + "dtype": "f32",
362 + "hash": "8bcfa49d11bf8160",
363 + "nbytes": 1638400,
364 + "offset": 70090752,
365 + "shape": [
366 + 640,
367 + 640
368 + ],
369 + "shard": "objects/72626eb74479f8dd.fshard"
370 + },
371 + "blocks.3.attn.wv.weight": {
372 + "dtype": "f32",
373 + "hash": "8a527c0649e24cd9",
374 + "nbytes": 1638400,
375 + "offset": 73367552,
376 + "shape": [
377 + 640,
378 + 640
379 + ],
380 + "shard": "objects/72626eb74479f8dd.fshard"
381 + },
382 + "blocks.3.mlp.w1.weight": {
383 + "dtype": "f32",
384 + "hash": "81058be641f54d7a",
385 + "nbytes": 4423680,
386 + "offset": 76660736,
387 + "shape": [
388 + 1728,
389 + 640
390 + ],
391 + "shard": "objects/72626eb74479f8dd.fshard"
392 + },
393 + "blocks.3.mlp.w2.weight": {
394 + "dtype": "f32",
395 + "hash": "67d4b0352f96a0f7",
396 + "nbytes": 4423680,
397 + "offset": 85508096,
398 + "shape": [
399 + 640,
400 + 1728
401 + ],
402 + "shard": "objects/72626eb74479f8dd.fshard"
403 + },
404 + "blocks.3.mlp.w3.weight": {
405 + "dtype": "f32",
406 + "hash": "7299e20a40b2da6c",
407 + "nbytes": 4423680,
408 + "offset": 81084416,
409 + "shape": [
410 + 1728,
411 + 640
412 + ],
413 + "shard": "objects/72626eb74479f8dd.fshard"
414 + },
415 + "blocks.3.norm1.weight": {
416 + "dtype": "f32",
417 + "hash": "8247bcdfab7adad7",
418 + "nbytes": 2560,
419 + "offset": 70074368,
420 + "shape": [
421 + 640
422 + ],
423 + "shard": "objects/72626eb74479f8dd.fshard"
424 + },
425 + "blocks.3.norm2.weight": {
426 + "dtype": "f32",
427 + "hash": "596b0bb7a9ad16ad",
428 + "nbytes": 2560,
429 + "offset": 76644352,
430 + "shape": [
431 + 640
432 + ],
433 + "shard": "objects/72626eb74479f8dd.fshard"
434 + },
435 + "blocks.4.attn.wk.weight": {
436 + "dtype": "f32",
437 + "hash": "01a11bb4643c874d",
438 + "nbytes": 1638400,
439 + "offset": 91586560,
440 + "shape": [
441 + 640,
442 + 640
443 + ],
444 + "shard": "objects/72626eb74479f8dd.fshard"
445 + },
446 + "blocks.4.attn.wo.weight": {
447 + "dtype": "f32",
448 + "hash": "a1e3afb60265703c",
449 + "nbytes": 1638400,
450 + "offset": 94863360,
451 + "shape": [
452 + 640,
453 + 640
454 + ],
455 + "shard": "objects/72626eb74479f8dd.fshard"
456 + },
457 + "blocks.4.attn.wq.weight": {
458 + "dtype": "f32",
459 + "hash": "b80af30751ce8923",
460 + "nbytes": 1638400,
461 + "offset": 89948160,
462 + "shape": [
463 + 640,
464 + 640
465 + ],
466 + "shard": "objects/72626eb74479f8dd.fshard"
467 + },
468 + "blocks.4.attn.wv.weight": {
469 + "dtype": "f32",
470 + "hash": "6ffdec20fca77461",
471 + "nbytes": 1638400,
472 + "offset": 93224960,
473 + "shape": [
474 + 640,
475 + 640
476 + ],
477 + "shard": "objects/72626eb74479f8dd.fshard"
478 + },
479 + "blocks.4.mlp.w1.weight": {
480 + "dtype": "f32",
481 + "hash": "164de67faf3c3be9",
482 + "nbytes": 4423680,
483 + "offset": 16384,
484 + "shape": [
485 + 1728,
486 + 640
487 + ],
488 + "shard": "objects/b4e0145ce1f2ff18.fshard"
489 + },
490 + "blocks.4.mlp.w2.weight": {
491 + "dtype": "f32",
492 + "hash": "b95146cee5c723ff",
493 + "nbytes": 4423680,
494 + "offset": 8863744,
495 + "shape": [
496 + 640,
497 + 1728
498 + ],
499 + "shard": "objects/b4e0145ce1f2ff18.fshard"
500 + },
501 + "blocks.4.mlp.w3.weight": {
502 + "dtype": "f32",
503 + "hash": "14b17e17afb81c3c",
504 + "nbytes": 4423680,
505 + "offset": 4440064,
506 + "shape": [
507 + 1728,
508 + 640
509 + ],
510 + "shard": "objects/b4e0145ce1f2ff18.fshard"
511 + },
512 + "blocks.4.norm1.weight": {
513 + "dtype": "f32",
514 + "hash": "dddf2bd3cd03a357",
515 + "nbytes": 2560,
516 + "offset": 89931776,
517 + "shape": [
518 + 640
519 + ],
520 + "shard": "objects/72626eb74479f8dd.fshard"
521 + },
522 + "blocks.4.norm2.weight": {
523 + "dtype": "f32",
524 + "hash": "3eac1a00123afade",
525 + "nbytes": 2560,
526 + "offset": 96501760,
527 + "shape": [
528 + 640
529 + ],
530 + "shard": "objects/72626eb74479f8dd.fshard"
531 + },
532 + "blocks.5.attn.wk.weight": {
533 + "dtype": "f32",
534 + "hash": "8b994bd11cc5c235",
535 + "nbytes": 1638400,
536 + "offset": 14942208,
537 + "shape": [
538 + 640,
539 + 640
540 + ],
541 + "shard": "objects/b4e0145ce1f2ff18.fshard"
542 + },
543 + "blocks.5.attn.wo.weight": {
544 + "dtype": "f32",
545 + "hash": "14d9dff05dd948e5",
546 + "nbytes": 1638400,
547 + "offset": 18219008,
548 + "shape": [
549 + 640,
550 + 640
551 + ],
552 + "shard": "objects/b4e0145ce1f2ff18.fshard"
553 + },
554 + "blocks.5.attn.wq.weight": {
555 + "dtype": "f32",
556 + "hash": "da8ef005898a1637",
557 + "nbytes": 1638400,
558 + "offset": 13303808,
559 + "shape": [
560 + 640,
561 + 640
562 + ],
563 + "shard": "objects/b4e0145ce1f2ff18.fshard"
564 + },
565 + "blocks.5.attn.wv.weight": {
566 + "dtype": "f32",
567 + "hash": "6b28b2be70bf033a",
568 + "nbytes": 1638400,
569 + "offset": 16580608,
570 + "shape": [
571 + 640,
572 + 640
573 + ],
574 + "shard": "objects/b4e0145ce1f2ff18.fshard"
575 + },
576 + "blocks.5.mlp.w1.weight": {
577 + "dtype": "f32",
578 + "hash": "500fae8fdecb95ca",
579 + "nbytes": 4423680,
580 + "offset": 19873792,
581 + "shape": [
582 + 1728,
583 + 640
584 + ],
585 + "shard": "objects/b4e0145ce1f2ff18.fshard"
586 + },
587 + "blocks.5.mlp.w2.weight": {
588 + "dtype": "f32",
589 + "hash": "6b84ccc84ca6b6fa",
590 + "nbytes": 4423680,
591 + "offset": 28721152,
592 + "shape": [
593 + 640,
594 + 1728
595 + ],
596 + "shard": "objects/b4e0145ce1f2ff18.fshard"
597 + },
598 + "blocks.5.mlp.w3.weight": {
599 + "dtype": "f32",
600 + "hash": "541490e9a2d38329",
601 + "nbytes": 4423680,
602 + "offset": 24297472,
603 + "shape": [
604 + 1728,
605 + 640
606 + ],
607 + "shard": "objects/b4e0145ce1f2ff18.fshard"
608 + },
609 + "blocks.5.norm1.weight": {
610 + "dtype": "f32",
611 + "hash": "9712111d7db8b749",
612 + "nbytes": 2560,
613 + "offset": 13287424,
614 + "shape": [
615 + 640
616 + ],
617 + "shard": "objects/b4e0145ce1f2ff18.fshard"
618 + },
619 + "blocks.5.norm2.weight": {
620 + "dtype": "f32",
621 + "hash": "9656dd83e2703195",
622 + "nbytes": 2560,
623 + "offset": 19857408,
624 + "shape": [
625 + 640
626 + ],
627 + "shard": "objects/b4e0145ce1f2ff18.fshard"
628 + },
629 + "blocks.6.attn.wk.weight": {
630 + "dtype": "f32",
631 + "hash": "46d6a33811ae7958",
632 + "nbytes": 1638400,
633 + "offset": 34799616,
634 + "shape": [
635 + 640,
636 + 640
637 + ],
638 + "shard": "objects/b4e0145ce1f2ff18.fshard"
639 + },
640 + "blocks.6.attn.wo.weight": {
641 + "dtype": "f32",
642 + "hash": "0af48008a7e9d48d",
643 + "nbytes": 1638400,
644 + "offset": 38076416,
645 + "shape": [
646 + 640,
647 + 640
648 + ],
649 + "shard": "objects/b4e0145ce1f2ff18.fshard"
650 + },
651 + "blocks.6.attn.wq.weight": {
652 + "dtype": "f32",
653 + "hash": "6f4660a328d4b423",
654 + "nbytes": 1638400,
655 + "offset": 33161216,
656 + "shape": [
657 + 640,
658 + 640
659 + ],
660 + "shard": "objects/b4e0145ce1f2ff18.fshard"
661 + },
662 + "blocks.6.attn.wv.weight": {
663 + "dtype": "f32",
664 + "hash": "2b22a3c36d0fa06d",
665 + "nbytes": 1638400,
666 + "offset": 36438016,
667 + "shape": [
668 + 640,
669 + 640
670 + ],
671 + "shard": "objects/b4e0145ce1f2ff18.fshard"
672 + },
673 + "blocks.6.mlp.w1.weight": {
674 + "dtype": "f32",
675 + "hash": "9ffa6ede45d79cb6",
676 + "nbytes": 4423680,
677 + "offset": 39731200,
678 + "shape": [
679 + 1728,
680 + 640
681 + ],
682 + "shard": "objects/b4e0145ce1f2ff18.fshard"
683 + },
684 + "blocks.6.mlp.w2.weight": {
685 + "dtype": "f32",
686 + "hash": "cded5100040c6a61",
687 + "nbytes": 4423680,
688 + "offset": 48578560,
689 + "shape": [
690 + 640,
691 + 1728
692 + ],
693 + "shard": "objects/b4e0145ce1f2ff18.fshard"
694 + },
695 + "blocks.6.mlp.w3.weight": {
696 + "dtype": "f32",
697 + "hash": "bd923d3a8f5162e6",
698 + "nbytes": 4423680,
699 + "offset": 44154880,
700 + "shape": [
701 + 1728,
702 + 640
703 + ],
704 + "shard": "objects/b4e0145ce1f2ff18.fshard"
705 + },
706 + "blocks.6.norm1.weight": {
707 + "dtype": "f32",
708 + "hash": "063bfae2ccd1b891",
709 + "nbytes": 2560,
710 + "offset": 33144832,
711 + "shape": [
712 + 640
713 + ],
714 + "shard": "objects/b4e0145ce1f2ff18.fshard"
715 + },
716 + "blocks.6.norm2.weight": {
717 + "dtype": "f32",
718 + "hash": "3c20ba70fba95945",
719 + "nbytes": 2560,
720 + "offset": 39714816,
721 + "shape": [
722 + 640
723 + ],
724 + "shard": "objects/b4e0145ce1f2ff18.fshard"
725 + },
726 + "blocks.7.attn.wk.weight": {
727 + "dtype": "f32",
728 + "hash": "52ea38d02ac4ae3e",
729 + "nbytes": 1638400,
730 + "offset": 54657024,
731 + "shape": [
732 + 640,
733 + 640
734 + ],
735 + "shard": "objects/b4e0145ce1f2ff18.fshard"
736 + },
737 + "blocks.7.attn.wo.weight": {
738 + "dtype": "f32",
739 + "hash": "945ecf9b46aeb099",
740 + "nbytes": 1638400,
741 + "offset": 57933824,
742 + "shape": [
743 + 640,
744 + 640
745 + ],
746 + "shard": "objects/b4e0145ce1f2ff18.fshard"
747 + },
748 + "blocks.7.attn.wq.weight": {
749 + "dtype": "f32",
750 + "hash": "bb4f532e47bbb588",
751 + "nbytes": 1638400,
752 + "offset": 53018624,
753 + "shape": [
754 + 640,
755 + 640
756 + ],
757 + "shard": "objects/b4e0145ce1f2ff18.fshard"
758 + },
759 + "blocks.7.attn.wv.weight": {
760 + "dtype": "f32",
761 + "hash": "39fc5bb55f990e08",
762 + "nbytes": 1638400,
763 + "offset": 56295424,
764 + "shape": [
765 + 640,
766 + 640
767 + ],
768 + "shard": "objects/b4e0145ce1f2ff18.fshard"
769 + },
770 + "blocks.7.mlp.w1.weight": {
771 + "dtype": "f32",
772 + "hash": "661591194873739a",
773 + "nbytes": 4423680,
774 + "offset": 59588608,
775 + "shape": [
776 + 1728,
777 + 640
778 + ],
779 + "shard": "objects/b4e0145ce1f2ff18.fshard"
780 + },
781 + "blocks.7.mlp.w2.weight": {
782 + "dtype": "f32",
783 + "hash": "c0693bdee8cb26f0",
784 + "nbytes": 4423680,
785 + "offset": 68435968,
786 + "shape": [
787 + 640,
788 + 1728
789 + ],
790 + "shard": "objects/b4e0145ce1f2ff18.fshard"
791 + },
792 + "blocks.7.mlp.w3.weight": {
793 + "dtype": "f32",
794 + "hash": "c46311e050bbc2e7",
795 + "nbytes": 4423680,
796 + "offset": 64012288,
797 + "shape": [
798 + 1728,
799 + 640
800 + ],
801 + "shard": "objects/b4e0145ce1f2ff18.fshard"
802 + },
803 + "blocks.7.norm1.weight": {
804 + "dtype": "f32",
805 + "hash": "31ddf245176231e8",
806 + "nbytes": 2560,
807 + "offset": 53002240,
808 + "shape": [
809 + 640
810 + ],
811 + "shard": "objects/b4e0145ce1f2ff18.fshard"
812 + },
813 + "blocks.7.norm2.weight": {
814 + "dtype": "f32",
815 + "hash": "adb0020bde0202ae",
816 + "nbytes": 2560,
817 + "offset": 59572224,
818 + "shape": [
819 + 640
820 + ],
821 + "shard": "objects/b4e0145ce1f2ff18.fshard"
822 + },
823 + "blocks.8.attn.wk.weight": {
824 + "dtype": "f32",
825 + "hash": "8b49975f1bcab4b1",
826 + "nbytes": 1638400,
827 + "offset": 74514432,
828 + "shape": [
829 + 640,
830 + 640
831 + ],
832 + "shard": "objects/b4e0145ce1f2ff18.fshard"
833 + },
834 + "blocks.8.attn.wo.weight": {
835 + "dtype": "f32",
836 + "hash": "d0fd8ea0a57aaf34",
837 + "nbytes": 1638400,
838 + "offset": 77791232,
839 + "shape": [
840 + 640,
841 + 640
842 + ],
843 + "shard": "objects/b4e0145ce1f2ff18.fshard"
844 + },
845 + "blocks.8.attn.wq.weight": {
846 + "dtype": "f32",
847 + "hash": "be6bf060be0ceb6a",
848 + "nbytes": 1638400,
849 + "offset": 72876032,
850 + "shape": [
851 + 640,
852 + 640
853 + ],
854 + "shard": "objects/b4e0145ce1f2ff18.fshard"
855 + },
856 + "blocks.8.attn.wv.weight": {
857 + "dtype": "f32",
858 + "hash": "876779888157dd34",
859 + "nbytes": 1638400,
860 + "offset": 76152832,
861 + "shape": [
862 + 640,
863 + 640
864 + ],
865 + "shard": "objects/b4e0145ce1f2ff18.fshard"
866 + },
867 + "blocks.8.mlp.w1.weight": {
868 + "dtype": "f32",
869 + "hash": "551d9ffdddc55001",
870 + "nbytes": 4423680,
871 + "offset": 79446016,
872 + "shape": [
873 + 1728,
874 + 640
875 + ],
876 + "shard": "objects/b4e0145ce1f2ff18.fshard"
877 + },
878 + "blocks.8.mlp.w2.weight": {
879 + "dtype": "f32",
880 + "hash": "710446ba331dbb33",
881 + "nbytes": 4423680,
882 + "offset": 88293376,
883 + "shape": [
884 + 640,
885 + 1728
886 + ],
887 + "shard": "objects/b4e0145ce1f2ff18.fshard"
888 + },
889 + "blocks.8.mlp.w3.weight": {
890 + "dtype": "f32",
891 + "hash": "75404c4f6c52ca7b",
892 + "nbytes": 4423680,
893 + "offset": 83869696,
894 + "shape": [
895 + 1728,
896 + 640
897 + ],
898 + "shard": "objects/b4e0145ce1f2ff18.fshard"
899 + },
900 + "blocks.8.norm1.weight": {
901 + "dtype": "f32",
902 + "hash": "a67c53d98e996f4d",
903 + "nbytes": 2560,
904 + "offset": 72859648,
905 + "shape": [
906 + 640
907 + ],
908 + "shard": "objects/b4e0145ce1f2ff18.fshard"
909 + },
910 + "blocks.8.norm2.weight": {
911 + "dtype": "f32",
912 + "hash": "b7cc1cd235754b8d",
913 + "nbytes": 2560,
914 + "offset": 79429632,
915 + "shape": [
916 + 640
917 + ],
918 + "shard": "objects/b4e0145ce1f2ff18.fshard"
919 + },
920 + "blocks.9.attn.wk.weight": {
921 + "dtype": "f32",
922 + "hash": "77ca114cacd59f92",
923 + "nbytes": 1638400,
924 + "offset": 94371840,
925 + "shape": [
926 + 640,
927 + 640
928 + ],
929 + "shard": "objects/b4e0145ce1f2ff18.fshard"
930 + },
931 + "blocks.9.attn.wo.weight": {
932 + "dtype": "f32",
933 + "hash": "2bd03b1679da1e59",
934 + "nbytes": 1638400,
935 + "offset": 97648640,
936 + "shape": [
937 + 640,
938 + 640
939 + ],
940 + "shard": "objects/b4e0145ce1f2ff18.fshard"
941 + },
942 + "blocks.9.attn.wq.weight": {
943 + "dtype": "f32",
944 + "hash": "c4ec432807c49aea",
945 + "nbytes": 1638400,
946 + "offset": 92733440,
947 + "shape": [
948 + 640,
949 + 640
950 + ],
951 + "shard": "objects/b4e0145ce1f2ff18.fshard"
952 + },
953 + "blocks.9.attn.wv.weight": {
954 + "dtype": "f32",
955 + "hash": "4f2bdf517e8aa817",
956 + "nbytes": 1638400,
957 + "offset": 96010240,
958 + "shape": [
959 + 640,
960 + 640
961 + ],
962 + "shard": "objects/b4e0145ce1f2ff18.fshard"
963 + },
964 + "blocks.9.mlp.w1.weight": {
965 + "dtype": "f32",
966 + "hash": "f65ccd0b671cebed",
967 + "nbytes": 4423680,
968 + "offset": 16384,
969 + "shape": [
970 + 1728,
971 + 640
972 + ],
973 + "shard": "objects/24a9ba4df0fd71ed.fshard"
974 + },
975 + "blocks.9.mlp.w2.weight": {
976 + "dtype": "f32",
977 + "hash": "14f2da31880d173e",
978 + "nbytes": 4423680,
979 + "offset": 8863744,
980 + "shape": [
981 + 640,
982 + 1728
983 + ],
984 + "shard": "objects/24a9ba4df0fd71ed.fshard"
985 + },
986 + "blocks.9.mlp.w3.weight": {
987 + "dtype": "f32",
988 + "hash": "08871f04b2cb1bb0",
989 + "nbytes": 4423680,
990 + "offset": 4440064,
991 + "shape": [
992 + 1728,
993 + 640
994 + ],
995 + "shard": "objects/24a9ba4df0fd71ed.fshard"
996 + },
997 + "blocks.9.norm1.weight": {
998 + "dtype": "f32",
999 + "hash": "ce6abe2627974147",
1000 + "nbytes": 2560,
1001 + "offset": 92717056,
1002 + "shape": [
1003 + 640
1004 + ],
1005 + "shard": "objects/b4e0145ce1f2ff18.fshard"
1006 + },
1007 + "blocks.9.norm2.weight": {
1008 + "dtype": "f32",
1009 + "hash": "cee537f25da157b9",
1010 + "nbytes": 2560,
1011 + "offset": 99287040,
1012 + "shape": [
1013 + 640
1014 + ],
1015 + "shard": "objects/b4e0145ce1f2ff18.fshard"
1016 + },
1017 + "final_norm.weight": {
1018 + "dtype": "f32",
1019 + "hash": "e14ca576c57ef871",
1020 + "nbytes": 2560,
1021 + "offset": 13287424,
1022 + "shape": [
1023 + 640
1024 + ],
1025 + "shard": "objects/24a9ba4df0fd71ed.fshard"
1026 + },
1027 + "tok_emb.weight": {
1028 + "dtype": "f32",
1029 + "hash": "1d963329140f196a",
1030 + "nbytes": 10485760,
1031 + "offset": 16384,
1032 + "shape": [
1033 + 4096,
1034 + 640
1035 + ],
1036 + "shard": "objects/72626eb74479f8dd.fshard"
1037 + }
1038 + }
1039 +}
\ No newline at end of file
added models/50m-tinystories.forge/manifests/manifest-000000.json +1039 −0
@@ -0,0 +1,1039 @@
1 +{
2 + "config": {
3 + "_comment": "~52M params, sized between gpt-25m and gpt-100m. vocab 4096 to match the local tok4096 train.bin (19.14M tokens). batch_size is the MICRO-batch: 8 x 8 x 1024 = 65536 tokens/step; 1170 steps = ~4 epochs over the 19.14M-token set. Warmup is 10% of the run. precision is parsed but not yet honored - all kernels are f32.",
4 + "model": {
5 + "activation": "swiglu",
6 + "context_length": 1024,
7 + "d_ff": 1728,
8 + "d_model": 640,
9 + "dropout": 0.0,
10 + "n_heads": 10,
11 + "n_kv_heads": 10,
12 + "n_layers": 10,
13 + "name": "gpt-50m",
14 + "norm": "rmsnorm",
15 + "norm_eps": 1e-06,
16 + "rope_theta": 10000.0,
17 + "tied_embeddings": true,
18 + "use_rope": true,
19 + "vocab_size": 4096
20 + },
21 + "train": {
22 + "batch_size": 8,
23 + "beta1": 0.9,
24 + "beta2": 0.95,
25 + "checkpoint_every": 200,
26 + "eps": 1e-08,
27 + "eval_batches": 20,
28 + "eval_every": 100,
29 + "grad_accum_steps": 8,
30 + "grad_clip": 1.0,
31 + "lr": 0.0005,
32 + "max_steps": 1170,
33 + "min_lr_ratio": 0.1,
34 + "precision": "f32",
35 + "seed": 1337,
36 + "warmup_steps": 117,
37 + "weight_decay": 0.1
38 + }
39 + },
40 + "dtype": "f32",
41 + "fmodel": 1,
42 + "parent": "",
43 + "self": "manifest-000000.json",
44 + "step": 1170,
45 + "tag": "one-epoch-x4",
46 + "tensors": {
47 + "blocks.0.attn.wk.weight": {
48 + "dtype": "f32",
49 + "hash": "ea0278bff183bd9b",
50 + "nbytes": 1638400,
51 + "offset": 12156928,
52 + "shape": [
53 + 640,
54 + 640
55 + ],
56 + "shard": "objects/72626eb74479f8dd.fshard"
57 + },
58 + "blocks.0.attn.wo.weight": {
59 + "dtype": "f32",
60 + "hash": "b08bb6730b3f4047",
61 + "nbytes": 1638400,
62 + "offset": 15433728,
63 + "shape": [
64 + 640,
65 + 640
66 + ],
67 + "shard": "objects/72626eb74479f8dd.fshard"
68 + },
69 + "blocks.0.attn.wq.weight": {
70 + "dtype": "f32",
71 + "hash": "c9f0aee32c2afb09",
72 + "nbytes": 1638400,
73 + "offset": 10518528,
74 + "shape": [
75 + 640,
76 + 640
77 + ],
78 + "shard": "objects/72626eb74479f8dd.fshard"
79 + },
80 + "blocks.0.attn.wv.weight": {
81 + "dtype": "f32",
82 + "hash": "a46e677dff653afa",
83 + "nbytes": 1638400,
84 + "offset": 13795328,
85 + "shape": [
86 + 640,
87 + 640
88 + ],
89 + "shard": "objects/72626eb74479f8dd.fshard"
90 + },
91 + "blocks.0.mlp.w1.weight": {
92 + "dtype": "f32",
93 + "hash": "4b63b198ca18d908",
94 + "nbytes": 4423680,
95 + "offset": 17088512,
96 + "shape": [
97 + 1728,
98 + 640
99 + ],
100 + "shard": "objects/72626eb74479f8dd.fshard"
101 + },
102 + "blocks.0.mlp.w2.weight": {
103 + "dtype": "f32",
104 + "hash": "d0b0ced42d5c91d2",
105 + "nbytes": 4423680,
106 + "offset": 25935872,
107 + "shape": [
108 + 640,
109 + 1728
110 + ],
111 + "shard": "objects/72626eb74479f8dd.fshard"
112 + },
113 + "blocks.0.mlp.w3.weight": {
114 + "dtype": "f32",
115 + "hash": "d36f453ea5b7576d",
116 + "nbytes": 4423680,
117 + "offset": 21512192,
118 + "shape": [
119 + 1728,
120 + 640
121 + ],
122 + "shard": "objects/72626eb74479f8dd.fshard"
123 + },
124 + "blocks.0.norm1.weight": {
125 + "dtype": "f32",
126 + "hash": "bf60929ceab06890",
127 + "nbytes": 2560,
128 + "offset": 10502144,
129 + "shape": [
130 + 640
131 + ],
132 + "shard": "objects/72626eb74479f8dd.fshard"
133 + },
134 + "blocks.0.norm2.weight": {
135 + "dtype": "f32",
136 + "hash": "bb257feeea6d5b6f",
137 + "nbytes": 2560,
138 + "offset": 17072128,
139 + "shape": [
140 + 640
141 + ],
142 + "shard": "objects/72626eb74479f8dd.fshard"
143 + },
144 + "blocks.1.attn.wk.weight": {
145 + "dtype": "f32",
146 + "hash": "a99d45c6af060836",
147 + "nbytes": 1638400,
148 + "offset": 32014336,
149 + "shape": [
150 + 640,
151 + 640
152 + ],
153 + "shard": "objects/72626eb74479f8dd.fshard"
154 + },
155 + "blocks.1.attn.wo.weight": {
156 + "dtype": "f32",
157 + "hash": "4ecc0144607917a7",
158 + "nbytes": 1638400,
159 + "offset": 35291136,
160 + "shape": [
161 + 640,
162 + 640
163 + ],
164 + "shard": "objects/72626eb74479f8dd.fshard"
165 + },
166 + "blocks.1.attn.wq.weight": {
167 + "dtype": "f32",
168 + "hash": "5dd4c343a6557550",
169 + "nbytes": 1638400,
170 + "offset": 30375936,
171 + "shape": [
172 + 640,
173 + 640
174 + ],
175 + "shard": "objects/72626eb74479f8dd.fshard"
176 + },
177 + "blocks.1.attn.wv.weight": {
178 + "dtype": "f32",
179 + "hash": "0de8eef2c1851cef",
180 + "nbytes": 1638400,
181 + "offset": 33652736,
182 + "shape": [
183 + 640,
184 + 640
185 + ],
186 + "shard": "objects/72626eb74479f8dd.fshard"
187 + },
188 + "blocks.1.mlp.w1.weight": {
189 + "dtype": "f32",
190 + "hash": "d8fbfcb43a5a4677",
191 + "nbytes": 4423680,
192 + "offset": 36945920,
193 + "shape": [
194 + 1728,
195 + 640
196 + ],
197 + "shard": "objects/72626eb74479f8dd.fshard"
198 + },
199 + "blocks.1.mlp.w2.weight": {
200 + "dtype": "f32",
201 + "hash": "df86dbe751fd0897",
202 + "nbytes": 4423680,
203 + "offset": 45793280,
204 + "shape": [
205 + 640,
206 + 1728
207 + ],
208 + "shard": "objects/72626eb74479f8dd.fshard"
209 + },
210 + "blocks.1.mlp.w3.weight": {
211 + "dtype": "f32",
212 + "hash": "e206903aafcacc9a",
213 + "nbytes": 4423680,
214 + "offset": 41369600,
215 + "shape": [
216 + 1728,
217 + 640
218 + ],
219 + "shard": "objects/72626eb74479f8dd.fshard"
220 + },
221 + "blocks.1.norm1.weight": {
222 + "dtype": "f32",
223 + "hash": "bc7c7f4144e6c4a0",
224 + "nbytes": 2560,
225 + "offset": 30359552,
226 + "shape": [
227 + 640
228 + ],
229 + "shard": "objects/72626eb74479f8dd.fshard"
230 + },
231 + "blocks.1.norm2.weight": {
232 + "dtype": "f32",
233 + "hash": "1d631908e97d0ab2",
234 + "nbytes": 2560,
235 + "offset": 36929536,
236 + "shape": [
237 + 640
238 + ],
239 + "shard": "objects/72626eb74479f8dd.fshard"
240 + },
241 + "blocks.2.attn.wk.weight": {
242 + "dtype": "f32",
243 + "hash": "e5f3bfa24ce66850",
244 + "nbytes": 1638400,
245 + "offset": 51871744,
246 + "shape": [
247 + 640,
248 + 640
249 + ],
250 + "shard": "objects/72626eb74479f8dd.fshard"
251 + },
252 + "blocks.2.attn.wo.weight": {
253 + "dtype": "f32",
254 + "hash": "81e09028c8d4fa64",
255 + "nbytes": 1638400,
256 + "offset": 55148544,
257 + "shape": [
258 + 640,
259 + 640
260 + ],
261 + "shard": "objects/72626eb74479f8dd.fshard"
262 + },
263 + "blocks.2.attn.wq.weight": {
264 + "dtype": "f32",
265 + "hash": "70ee42db3250d15e",
266 + "nbytes": 1638400,
267 + "offset": 50233344,
268 + "shape": [
269 + 640,
270 + 640
271 + ],
272 + "shard": "objects/72626eb74479f8dd.fshard"
273 + },
274 + "blocks.2.attn.wv.weight": {
275 + "dtype": "f32",
276 + "hash": "b27d667f5ce35172",
277 + "nbytes": 1638400,
278 + "offset": 53510144,
279 + "shape": [
280 + 640,
281 + 640
282 + ],
283 + "shard": "objects/72626eb74479f8dd.fshard"
284 + },
285 + "blocks.2.mlp.w1.weight": {
286 + "dtype": "f32",
287 + "hash": "eef179f1c68e1af2",
288 + "nbytes": 4423680,
289 + "offset": 56803328,
290 + "shape": [
291 + 1728,
292 + 640
293 + ],
294 + "shard": "objects/72626eb74479f8dd.fshard"
295 + },
296 + "blocks.2.mlp.w2.weight": {
297 + "dtype": "f32",
298 + "hash": "0393a401c220c511",
299 + "nbytes": 4423680,
300 + "offset": 65650688,
301 + "shape": [
302 + 640,
303 + 1728
304 + ],
305 + "shard": "objects/72626eb74479f8dd.fshard"
306 + },
307 + "blocks.2.mlp.w3.weight": {
308 + "dtype": "f32",
309 + "hash": "8b9edaa2ac9409e3",
310 + "nbytes": 4423680,
311 + "offset": 61227008,
312 + "shape": [
313 + 1728,
314 + 640
315 + ],
316 + "shard": "objects/72626eb74479f8dd.fshard"
317 + },
318 + "blocks.2.norm1.weight": {
319 + "dtype": "f32",
320 + "hash": "fb9413da5028b6fd",
321 + "nbytes": 2560,
322 + "offset": 50216960,
323 + "shape": [
324 + 640
325 + ],
326 + "shard": "objects/72626eb74479f8dd.fshard"
327 + },
328 + "blocks.2.norm2.weight": {
329 + "dtype": "f32",
330 + "hash": "a76041b595744213",
331 + "nbytes": 2560,
332 + "offset": 56786944,
333 + "shape": [
334 + 640
335 + ],
336 + "shard": "objects/72626eb74479f8dd.fshard"
337 + },
338 + "blocks.3.attn.wk.weight": {
339 + "dtype": "f32",
340 + "hash": "01edcec7013dc445",
341 + "nbytes": 1638400,
342 + "offset": 71729152,
343 + "shape": [
344 + 640,
345 + 640
346 + ],
347 + "shard": "objects/72626eb74479f8dd.fshard"
348 + },
349 + "blocks.3.attn.wo.weight": {
350 + "dtype": "f32",
351 + "hash": "ea9891714020fcf4",
352 + "nbytes": 1638400,
353 + "offset": 75005952,
354 + "shape": [
355 + 640,
356 + 640
357 + ],
358 + "shard": "objects/72626eb74479f8dd.fshard"
359 + },
360 + "blocks.3.attn.wq.weight": {
361 + "dtype": "f32",
362 + "hash": "8bcfa49d11bf8160",
363 + "nbytes": 1638400,
364 + "offset": 70090752,
365 + "shape": [
366 + 640,
367 + 640
368 + ],
369 + "shard": "objects/72626eb74479f8dd.fshard"
370 + },
371 + "blocks.3.attn.wv.weight": {
372 + "dtype": "f32",
373 + "hash": "8a527c0649e24cd9",
374 + "nbytes": 1638400,
375 + "offset": 73367552,
376 + "shape": [
377 + 640,
378 + 640
379 + ],
380 + "shard": "objects/72626eb74479f8dd.fshard"
381 + },
382 + "blocks.3.mlp.w1.weight": {
383 + "dtype": "f32",
384 + "hash": "81058be641f54d7a",
385 + "nbytes": 4423680,
386 + "offset": 76660736,
387 + "shape": [
388 + 1728,
389 + 640
390 + ],
391 + "shard": "objects/72626eb74479f8dd.fshard"
392 + },
393 + "blocks.3.mlp.w2.weight": {
394 + "dtype": "f32",
395 + "hash": "67d4b0352f96a0f7",
396 + "nbytes": 4423680,
397 + "offset": 85508096,
398 + "shape": [
399 + 640,
400 + 1728
401 + ],
402 + "shard": "objects/72626eb74479f8dd.fshard"
403 + },
404 + "blocks.3.mlp.w3.weight": {
405 + "dtype": "f32",
406 + "hash": "7299e20a40b2da6c",
407 + "nbytes": 4423680,
408 + "offset": 81084416,
409 + "shape": [
410 + 1728,
411 + 640
412 + ],
413 + "shard": "objects/72626eb74479f8dd.fshard"
414 + },
415 + "blocks.3.norm1.weight": {
416 + "dtype": "f32",
417 + "hash": "8247bcdfab7adad7",
418 + "nbytes": 2560,
419 + "offset": 70074368,
420 + "shape": [
421 + 640
422 + ],
423 + "shard": "objects/72626eb74479f8dd.fshard"
424 + },
425 + "blocks.3.norm2.weight": {
426 + "dtype": "f32",
427 + "hash": "596b0bb7a9ad16ad",
428 + "nbytes": 2560,
429 + "offset": 76644352,
430 + "shape": [
431 + 640
432 + ],
433 + "shard": "objects/72626eb74479f8dd.fshard"
434 + },
435 + "blocks.4.attn.wk.weight": {
436 + "dtype": "f32",
437 + "hash": "01a11bb4643c874d",
438 + "nbytes": 1638400,
439 + "offset": 91586560,
440 + "shape": [
441 + 640,
442 + 640
443 + ],
444 + "shard": "objects/72626eb74479f8dd.fshard"
445 + },
446 + "blocks.4.attn.wo.weight": {
447 + "dtype": "f32",
448 + "hash": "a1e3afb60265703c",
449 + "nbytes": 1638400,
450 + "offset": 94863360,
451 + "shape": [
452 + 640,
453 + 640
454 + ],
455 + "shard": "objects/72626eb74479f8dd.fshard"
456 + },
457 + "blocks.4.attn.wq.weight": {
458 + "dtype": "f32",
459 + "hash": "b80af30751ce8923",
460 + "nbytes": 1638400,
461 + "offset": 89948160,
462 + "shape": [
463 + 640,
464 + 640
465 + ],
466 + "shard": "objects/72626eb74479f8dd.fshard"
467 + },
468 + "blocks.4.attn.wv.weight": {
469 + "dtype": "f32",
470 + "hash": "6ffdec20fca77461",
471 + "nbytes": 1638400,
472 + "offset": 93224960,
473 + "shape": [
474 + 640,
475 + 640
476 + ],
477 + "shard": "objects/72626eb74479f8dd.fshard"
478 + },
479 + "blocks.4.mlp.w1.weight": {
480 + "dtype": "f32",
481 + "hash": "164de67faf3c3be9",
482 + "nbytes": 4423680,
483 + "offset": 16384,
484 + "shape": [
485 + 1728,
486 + 640
487 + ],
488 + "shard": "objects/b4e0145ce1f2ff18.fshard"
489 + },
490 + "blocks.4.mlp.w2.weight": {
491 + "dtype": "f32",
492 + "hash": "b95146cee5c723ff",
493 + "nbytes": 4423680,
494 + "offset": 8863744,
495 + "shape": [
496 + 640,
497 + 1728
498 + ],
499 + "shard": "objects/b4e0145ce1f2ff18.fshard"
500 + },
501 + "blocks.4.mlp.w3.weight": {
502 + "dtype": "f32",
503 + "hash": "14b17e17afb81c3c",
504 + "nbytes": 4423680,
505 + "offset": 4440064,
506 + "shape": [
507 + 1728,
508 + 640
509 + ],
510 + "shard": "objects/b4e0145ce1f2ff18.fshard"
511 + },
512 + "blocks.4.norm1.weight": {
513 + "dtype": "f32",
514 + "hash": "dddf2bd3cd03a357",
515 + "nbytes": 2560,
516 + "offset": 89931776,
517 + "shape": [
518 + 640
519 + ],
520 + "shard": "objects/72626eb74479f8dd.fshard"
521 + },
522 + "blocks.4.norm2.weight": {
523 + "dtype": "f32",
524 + "hash": "3eac1a00123afade",
525 + "nbytes": 2560,
526 + "offset": 96501760,
527 + "shape": [
528 + 640
529 + ],
530 + "shard": "objects/72626eb74479f8dd.fshard"
531 + },
532 + "blocks.5.attn.wk.weight": {
533 + "dtype": "f32",
534 + "hash": "8b994bd11cc5c235",
535 + "nbytes": 1638400,
536 + "offset": 14942208,
537 + "shape": [
538 + 640,
539 + 640
540 + ],
541 + "shard": "objects/b4e0145ce1f2ff18.fshard"
542 + },
543 + "blocks.5.attn.wo.weight": {
544 + "dtype": "f32",
545 + "hash": "14d9dff05dd948e5",
546 + "nbytes": 1638400,
547 + "offset": 18219008,
548 + "shape": [
549 + 640,
550 + 640
551 + ],
552 + "shard": "objects/b4e0145ce1f2ff18.fshard"
553 + },
554 + "blocks.5.attn.wq.weight": {
555 + "dtype": "f32",
556 + "hash": "da8ef005898a1637",
557 + "nbytes": 1638400,
558 + "offset": 13303808,
559 + "shape": [
560 + 640,
561 + 640
562 + ],
563 + "shard": "objects/b4e0145ce1f2ff18.fshard"
564 + },
565 + "blocks.5.attn.wv.weight": {
566 + "dtype": "f32",
567 + "hash": "6b28b2be70bf033a",
568 + "nbytes": 1638400,
569 + "offset": 16580608,
570 + "shape": [
571 + 640,
572 + 640
573 + ],
574 + "shard": "objects/b4e0145ce1f2ff18.fshard"
575 + },
576 + "blocks.5.mlp.w1.weight": {
577 + "dtype": "f32",
578 + "hash": "500fae8fdecb95ca",
579 + "nbytes": 4423680,
580 + "offset": 19873792,
581 + "shape": [
582 + 1728,
583 + 640
584 + ],
585 + "shard": "objects/b4e0145ce1f2ff18.fshard"
586 + },
587 + "blocks.5.mlp.w2.weight": {
588 + "dtype": "f32",
589 + "hash": "6b84ccc84ca6b6fa",
590 + "nbytes": 4423680,
591 + "offset": 28721152,
592 + "shape": [
593 + 640,
594 + 1728
595 + ],
596 + "shard": "objects/b4e0145ce1f2ff18.fshard"
597 + },
598 + "blocks.5.mlp.w3.weight": {
599 + "dtype": "f32",
600 + "hash": "541490e9a2d38329",
601 + "nbytes": 4423680,
602 + "offset": 24297472,
603 + "shape": [
604 + 1728,
605 + 640
606 + ],
607 + "shard": "objects/b4e0145ce1f2ff18.fshard"
608 + },
609 + "blocks.5.norm1.weight": {
610 + "dtype": "f32",
611 + "hash": "9712111d7db8b749",
612 + "nbytes": 2560,
613 + "offset": 13287424,
614 + "shape": [
615 + 640
616 + ],
617 + "shard": "objects/b4e0145ce1f2ff18.fshard"
618 + },
619 + "blocks.5.norm2.weight": {
620 + "dtype": "f32",
621 + "hash": "9656dd83e2703195",
622 + "nbytes": 2560,
623 + "offset": 19857408,
624 + "shape": [
625 + 640
626 + ],
627 + "shard": "objects/b4e0145ce1f2ff18.fshard"
628 + },
629 + "blocks.6.attn.wk.weight": {
630 + "dtype": "f32",
631 + "hash": "46d6a33811ae7958",
632 + "nbytes": 1638400,
633 + "offset": 34799616,
634 + "shape": [
635 + 640,
636 + 640
637 + ],
638 + "shard": "objects/b4e0145ce1f2ff18.fshard"
639 + },
640 + "blocks.6.attn.wo.weight": {
641 + "dtype": "f32",
642 + "hash": "0af48008a7e9d48d",
643 + "nbytes": 1638400,
644 + "offset": 38076416,
645 + "shape": [
646 + 640,
647 + 640
648 + ],
649 + "shard": "objects/b4e0145ce1f2ff18.fshard"
650 + },
651 + "blocks.6.attn.wq.weight": {
652 + "dtype": "f32",
653 + "hash": "6f4660a328d4b423",
654 + "nbytes": 1638400,
655 + "offset": 33161216,
656 + "shape": [
657 + 640,
658 + 640
659 + ],
660 + "shard": "objects/b4e0145ce1f2ff18.fshard"
661 + },
662 + "blocks.6.attn.wv.weight": {
663 + "dtype": "f32",
664 + "hash": "2b22a3c36d0fa06d",
665 + "nbytes": 1638400,
666 + "offset": 36438016,
667 + "shape": [
668 + 640,
669 + 640
670 + ],
671 + "shard": "objects/b4e0145ce1f2ff18.fshard"
672 + },
673 + "blocks.6.mlp.w1.weight": {
674 + "dtype": "f32",
675 + "hash": "9ffa6ede45d79cb6",
676 + "nbytes": 4423680,
677 + "offset": 39731200,
678 + "shape": [
679 + 1728,
680 + 640
681 + ],
682 + "shard": "objects/b4e0145ce1f2ff18.fshard"
683 + },
684 + "blocks.6.mlp.w2.weight": {
685 + "dtype": "f32",
686 + "hash": "cded5100040c6a61",
687 + "nbytes": 4423680,
688 + "offset": 48578560,
689 + "shape": [
690 + 640,
691 + 1728
692 + ],
693 + "shard": "objects/b4e0145ce1f2ff18.fshard"
694 + },
695 + "blocks.6.mlp.w3.weight": {
696 + "dtype": "f32",
697 + "hash": "bd923d3a8f5162e6",
698 + "nbytes": 4423680,
699 + "offset": 44154880,
700 + "shape": [
701 + 1728,
702 + 640
703 + ],
704 + "shard": "objects/b4e0145ce1f2ff18.fshard"
705 + },
706 + "blocks.6.norm1.weight": {
707 + "dtype": "f32",
708 + "hash": "063bfae2ccd1b891",
709 + "nbytes": 2560,
710 + "offset": 33144832,
711 + "shape": [
712 + 640
713 + ],
714 + "shard": "objects/b4e0145ce1f2ff18.fshard"
715 + },
716 + "blocks.6.norm2.weight": {
717 + "dtype": "f32",
718 + "hash": "3c20ba70fba95945",
719 + "nbytes": 2560,
720 + "offset": 39714816,
721 + "shape": [
722 + 640
723 + ],
724 + "shard": "objects/b4e0145ce1f2ff18.fshard"
725 + },
726 + "blocks.7.attn.wk.weight": {
727 + "dtype": "f32",
728 + "hash": "52ea38d02ac4ae3e",
729 + "nbytes": 1638400,
730 + "offset": 54657024,
731 + "shape": [
732 + 640,
733 + 640
734 + ],
735 + "shard": "objects/b4e0145ce1f2ff18.fshard"
736 + },
737 + "blocks.7.attn.wo.weight": {
738 + "dtype": "f32",
739 + "hash": "945ecf9b46aeb099",
740 + "nbytes": 1638400,
741 + "offset": 57933824,
742 + "shape": [
743 + 640,
744 + 640
745 + ],
746 + "shard": "objects/b4e0145ce1f2ff18.fshard"
747 + },
748 + "blocks.7.attn.wq.weight": {
749 + "dtype": "f32",
750 + "hash": "bb4f532e47bbb588",
751 + "nbytes": 1638400,
752 + "offset": 53018624,
753 + "shape": [
754 + 640,
755 + 640
756 + ],
757 + "shard": "objects/b4e0145ce1f2ff18.fshard"
758 + },
759 + "blocks.7.attn.wv.weight": {
760 + "dtype": "f32",
761 + "hash": "39fc5bb55f990e08",
762 + "nbytes": 1638400,
763 + "offset": 56295424,
764 + "shape": [
765 + 640,
766 + 640
767 + ],
768 + "shard": "objects/b4e0145ce1f2ff18.fshard"
769 + },
770 + "blocks.7.mlp.w1.weight": {
771 + "dtype": "f32",
772 + "hash": "661591194873739a",
773 + "nbytes": 4423680,
774 + "offset": 59588608,
775 + "shape": [
776 + 1728,
777 + 640
778 + ],
779 + "shard": "objects/b4e0145ce1f2ff18.fshard"
780 + },
781 + "blocks.7.mlp.w2.weight": {
782 + "dtype": "f32",
783 + "hash": "c0693bdee8cb26f0",
784 + "nbytes": 4423680,
785 + "offset": 68435968,
786 + "shape": [
787 + 640,
788 + 1728
789 + ],
790 + "shard": "objects/b4e0145ce1f2ff18.fshard"
791 + },
792 + "blocks.7.mlp.w3.weight": {
793 + "dtype": "f32",
794 + "hash": "c46311e050bbc2e7",
795 + "nbytes": 4423680,
796 + "offset": 64012288,
797 + "shape": [
798 + 1728,
799 + 640
800 + ],
801 + "shard": "objects/b4e0145ce1f2ff18.fshard"
802 + },
803 + "blocks.7.norm1.weight": {
804 + "dtype": "f32",
805 + "hash": "31ddf245176231e8",
806 + "nbytes": 2560,
807 + "offset": 53002240,
808 + "shape": [
809 + 640
810 + ],
811 + "shard": "objects/b4e0145ce1f2ff18.fshard"
812 + },
813 + "blocks.7.norm2.weight": {
814 + "dtype": "f32",
815 + "hash": "adb0020bde0202ae",
816 + "nbytes": 2560,
817 + "offset": 59572224,
818 + "shape": [
819 + 640
820 + ],
821 + "shard": "objects/b4e0145ce1f2ff18.fshard"
822 + },
823 + "blocks.8.attn.wk.weight": {
824 + "dtype": "f32",
825 + "hash": "8b49975f1bcab4b1",
826 + "nbytes": 1638400,
827 + "offset": 74514432,
828 + "shape": [
829 + 640,
830 + 640
831 + ],
832 + "shard": "objects/b4e0145ce1f2ff18.fshard"
833 + },
834 + "blocks.8.attn.wo.weight": {
835 + "dtype": "f32",
836 + "hash": "d0fd8ea0a57aaf34",
837 + "nbytes": 1638400,
838 + "offset": 77791232,
839 + "shape": [
840 + 640,
841 + 640
842 + ],
843 + "shard": "objects/b4e0145ce1f2ff18.fshard"
844 + },
845 + "blocks.8.attn.wq.weight": {
846 + "dtype": "f32",
847 + "hash": "be6bf060be0ceb6a",
848 + "nbytes": 1638400,
849 + "offset": 72876032,
850 + "shape": [
851 + 640,
852 + 640
853 + ],
854 + "shard": "objects/b4e0145ce1f2ff18.fshard"
855 + },
856 + "blocks.8.attn.wv.weight": {
857 + "dtype": "f32",
858 + "hash": "876779888157dd34",
859 + "nbytes": 1638400,
860 + "offset": 76152832,
861 + "shape": [
862 + 640,
863 + 640
864 + ],
865 + "shard": "objects/b4e0145ce1f2ff18.fshard"
866 + },
867 + "blocks.8.mlp.w1.weight": {
868 + "dtype": "f32",
869 + "hash": "551d9ffdddc55001",
870 + "nbytes": 4423680,
871 + "offset": 79446016,
872 + "shape": [
873 + 1728,
874 + 640
875 + ],
876 + "shard": "objects/b4e0145ce1f2ff18.fshard"
877 + },
878 + "blocks.8.mlp.w2.weight": {
879 + "dtype": "f32",
880 + "hash": "710446ba331dbb33",
881 + "nbytes": 4423680,
882 + "offset": 88293376,
883 + "shape": [
884 + 640,
885 + 1728
886 + ],
887 + "shard": "objects/b4e0145ce1f2ff18.fshard"
888 + },
889 + "blocks.8.mlp.w3.weight": {
890 + "dtype": "f32",
891 + "hash": "75404c4f6c52ca7b",
892 + "nbytes": 4423680,
893 + "offset": 83869696,
894 + "shape": [
895 + 1728,
896 + 640
897 + ],
898 + "shard": "objects/b4e0145ce1f2ff18.fshard"
899 + },
900 + "blocks.8.norm1.weight": {
901 + "dtype": "f32",
902 + "hash": "a67c53d98e996f4d",
903 + "nbytes": 2560,
904 + "offset": 72859648,
905 + "shape": [
906 + 640
907 + ],
908 + "shard": "objects/b4e0145ce1f2ff18.fshard"
909 + },
910 + "blocks.8.norm2.weight": {
911 + "dtype": "f32",
912 + "hash": "b7cc1cd235754b8d",
913 + "nbytes": 2560,
914 + "offset": 79429632,
915 + "shape": [
916 + 640
917 + ],
918 + "shard": "objects/b4e0145ce1f2ff18.fshard"
919 + },
920 + "blocks.9.attn.wk.weight": {
921 + "dtype": "f32",
922 + "hash": "77ca114cacd59f92",
923 + "nbytes": 1638400,
924 + "offset": 94371840,
925 + "shape": [
926 + 640,
927 + 640
928 + ],
929 + "shard": "objects/b4e0145ce1f2ff18.fshard"
930 + },
931 + "blocks.9.attn.wo.weight": {
932 + "dtype": "f32",
933 + "hash": "2bd03b1679da1e59",
934 + "nbytes": 1638400,
935 + "offset": 97648640,
936 + "shape": [
937 + 640,
938 + 640
939 + ],
940 + "shard": "objects/b4e0145ce1f2ff18.fshard"
941 + },
942 + "blocks.9.attn.wq.weight": {
943 + "dtype": "f32",
944 + "hash": "c4ec432807c49aea",
945 + "nbytes": 1638400,
946 + "offset": 92733440,
947 + "shape": [
948 + 640,
949 + 640
950 + ],
951 + "shard": "objects/b4e0145ce1f2ff18.fshard"
952 + },
953 + "blocks.9.attn.wv.weight": {
954 + "dtype": "f32",
955 + "hash": "4f2bdf517e8aa817",
956 + "nbytes": 1638400,
957 + "offset": 96010240,
958 + "shape": [
959 + 640,
960 + 640
961 + ],
962 + "shard": "objects/b4e0145ce1f2ff18.fshard"
963 + },
964 + "blocks.9.mlp.w1.weight": {
965 + "dtype": "f32",
966 + "hash": "f65ccd0b671cebed",
967 + "nbytes": 4423680,
968 + "offset": 16384,
969 + "shape": [
970 + 1728,
971 + 640
972 + ],
973 + "shard": "objects/24a9ba4df0fd71ed.fshard"
974 + },
975 + "blocks.9.mlp.w2.weight": {
976 + "dtype": "f32",
977 + "hash": "14f2da31880d173e",
978 + "nbytes": 4423680,
979 + "offset": 8863744,
980 + "shape": [
981 + 640,
982 + 1728
983 + ],
984 + "shard": "objects/24a9ba4df0fd71ed.fshard"
985 + },
986 + "blocks.9.mlp.w3.weight": {
987 + "dtype": "f32",
988 + "hash": "08871f04b2cb1bb0",
989 + "nbytes": 4423680,
990 + "offset": 4440064,
991 + "shape": [
992 + 1728,
993 + 640
994 + ],
995 + "shard": "objects/24a9ba4df0fd71ed.fshard"
996 + },
997 + "blocks.9.norm1.weight": {
998 + "dtype": "f32",
999 + "hash": "ce6abe2627974147",
1000 + "nbytes": 2560,
1001 + "offset": 92717056,
1002 + "shape": [
1003 + 640
1004 + ],
1005 + "shard": "objects/b4e0145ce1f2ff18.fshard"
1006 + },
1007 + "blocks.9.norm2.weight": {
1008 + "dtype": "f32",
1009 + "hash": "cee537f25da157b9",
1010 + "nbytes": 2560,
1011 + "offset": 99287040,
1012 + "shape": [
1013 + 640
1014 + ],
1015 + "shard": "objects/b4e0145ce1f2ff18.fshard"
1016 + },
1017 + "final_norm.weight": {
1018 + "dtype": "f32",
1019 + "hash": "e14ca576c57ef871",
1020 + "nbytes": 2560,
1021 + "offset": 13287424,
1022 + "shape": [
1023 + 640
1024 + ],
1025 + "shard": "objects/24a9ba4df0fd71ed.fshard"
1026 + },
1027 + "tok_emb.weight": {
1028 + "dtype": "f32",
1029 + "hash": "1d963329140f196a",
1030 + "nbytes": 10485760,
1031 + "offset": 16384,
1032 + "shape": [
1033 + 4096,
1034 + 640
1035 + ],
1036 + "shard": "objects/72626eb74479f8dd.fshard"
1037 + }
1038 + }
1039 +}
\ No newline at end of file
added models/50m-tinystories.forge/objects/24a9ba4df0fd71ed.fshard +0 −0

Binary file not shown.

added models/50m-tinystories.forge/objects/72626eb74479f8dd.fshard +0 −0

Binary file not shown.

added models/50m-tinystories.forge/objects/b4e0145ce1f2ff18.fshard +0 −0

Binary file not shown.

modified src/kernels/elementwise.metal +15 −0
@@ -52,6 +52,21 @@ kernel void gelu_f32(device const float* x [[buffer(0)]],
52 52 out[gid] = 0.5f * v * (1.0f + precise::tanh(k * (v + 0.044715f * v * v * v)));
53 53 }
54 54
55 +kernel void sigmoid_f32(device const float* x [[buffer(0)]],
56 + device float* out [[buffer(1)]],
57 + uint gid [[thread_position_in_grid]]) {
58 + out[gid] = 1.0f / (1.0f + exp(-x[gid]));
59 +}
60 +
61 +// dx += dout * y * (1 - y), using the forward output y.
62 +kernel void sigmoid_bwd_f32(device const float* y [[buffer(0)]],
63 + device const float* dout [[buffer(1)]],
64 + device float* dx [[buffer(2)]],
65 + uint gid [[thread_position_in_grid]]) {
66 + const float v = y[gid];
67 + dx[gid] = fma(dout[gid], v * (1.0f - v), dx[gid]);
68 +}
69 +
55 70 // ReLU² (nanoGPT-speedrun lineage): out = max(x, 0)^2.
56 71 kernel void relu2_f32(device const float* x [[buffer(0)]],
57 72 device float* out [[buffer(1)]],
modified src/kernels/moe.metal +62 −40
@@ -22,66 +22,88 @@ kernel void softmax_bwd_f32(device const float* P [[buffer(0)]],
22 22 for (uint j = 0; j < C; ++j) dx[j] += prob[j] * (dout[j] - dot);
23 23 }
24 24
25 −// Top-k gating: keep the k largest of each row of P [N, E], renormalize the
26 −// kept entries to sum to 1, zero the rest. Ties broken by lower index, which
27 −// keeps forward and backward selections identical.
28 −// p = (E, k)
29 −kernel void topk_renorm_f32(device const float* P [[buffer(0)]],
30 − device float* OUT [[buffer(1)]],
31 − constant uint2& p [[buffer(2)]],
32 − uint row [[thread_position_in_grid]]) {
33 − const uint E = p.x, K = p.y;
34 − device const float* in = P + ulong(row) * E;
35 − device float* out = OUT + ulong(row) * E;
25 +// Top-k gating with a selection bias (DeepSeek-V3 "noaux"): the k experts
26 +// are chosen by score+BIAS, but gate VALUES come from the biasless score —
27 +// the bias only steers routing. norm=1 renormalizes kept gates to sum 1
28 +// (classic top-k softmax); norm=0 keeps raw scores (sigmoid routing, V3).
29 +// Ties broken by lower index, which keeps forward and backward selections
30 +// identical. p = (E, k, norm)
31 +struct TopkParams { uint E, K, norm; };
36 32
37 − bool kept[64];
33 +static inline void topk_select_biased(device const float* in,
34 + device const float* bias, uint E, uint K,
35 + thread bool* kept, thread float* S) {
38 36 for (uint j = 0; j < E; ++j) kept[j] = false;
39 − float S = 0.0f;
37 + float acc = 0.0f;
40 38 for (uint sel = 0; sel < K; ++sel) {
41 39 float best = -FLT_MAX;
42 40 uint arg = 0;
43 − for (uint j = 0; j < E; ++j)
44 − if (!kept[j] && in[j] > best) { best = in[j]; arg = j; }
41 + for (uint j = 0; j < E; ++j) {
42 + const float v = in[j] + bias[j];
43 + if (!kept[j] && v > best) { best = v; arg = j; }
44 + }
45 45 kept[arg] = true;
46 − S += best;
46 + acc += in[arg]; // gate mass is biasless
47 47 }
48 − const float inv = 1.0f / max(S, 1e-12f);
49 − for (uint j = 0; j < E; ++j) out[j] = kept[j] ? in[j] * inv : 0.0f;
48 + *S = acc;
49 +}
50 +
51 +kernel void topk_renorm_f32(device const float* P [[buffer(0)]],
52 + device const float* BIAS [[buffer(1)]],
53 + device float* OUT [[buffer(2)]],
54 + constant TopkParams& p [[buffer(3)]],
55 + uint row [[thread_position_in_grid]]) {
56 + device const float* in = P + ulong(row) * p.E;
57 + device float* out = OUT + ulong(row) * p.E;
58 +
59 + bool kept[64];
60 + float S;
61 + topk_select_biased(in, BIAS, p.E, p.K, kept, &S);
62 + const float inv = p.norm ? 1.0f / max(S, 1e-12f) : 1.0f;
63 + for (uint j = 0; j < p.E; ++j) out[j] = kept[j] ? in[j] * inv : 0.0f;
50 64 }
51 65
52 −// Backward of topk_renorm: for kept entries, g_i = p_i / S with S the kept
53 −// sum, so dp_i += (dg_i − Σ_j dg_j g_j) / S; dropped entries get zero. The
54 −// kept set is recomputed from P with the same tie-breaking as forward.
55 −// p = (E, k)
66 +// Backward: with norm, g_i = p_i / S over kept entries so
67 +// dp_i += (dg_i − Σ_j dg_j g_j) / S; without norm, g_i = p_i so dp_i += dg_i.
68 +// The kept set (incl. bias) is recomputed with the same tie-breaking.
56 69 kernel void topk_renorm_bwd_f32(device const float* P [[buffer(0)]],
57 − device const float* DOUT [[buffer(1)]],
58 − device float* DP [[buffer(2)]],
59 − constant uint2& p [[buffer(3)]],
70 + device const float* BIAS [[buffer(1)]],
71 + device const float* DOUT [[buffer(2)]],
72 + device float* DP [[buffer(3)]],
73 + constant TopkParams& p [[buffer(4)]],
60 74 uint row [[thread_position_in_grid]]) {
61 − const uint E = p.x, K = p.y;
62 − device const float* in = P + ulong(row) * E;
63 − device const float* dout = DOUT + ulong(row) * E;
64 − device float* dp = DP + ulong(row) * E;
75 + device const float* in = P + ulong(row) * p.E;
76 + device const float* dout = DOUT + ulong(row) * p.E;
77 + device float* dp = DP + ulong(row) * p.E;
65 78
66 79 bool kept[64];
67 − for (uint j = 0; j < E; ++j) kept[j] = false;
68 − float S = 0.0f;
69 − for (uint sel = 0; sel < K; ++sel) {
70 − float best = -FLT_MAX;
71 − uint arg = 0;
72 − for (uint j = 0; j < E; ++j)
73 − if (!kept[j] && in[j] > best) { best = in[j]; arg = j; }
74 − kept[arg] = true;
75 − S += best;
80 + float S;
81 + topk_select_biased(in, BIAS, p.E, p.K, kept, &S);
82 + if (!p.norm) {
83 + for (uint j = 0; j < p.E; ++j)
84 + if (kept[j]) dp[j] += dout[j];
85 + return;
76 86 }
77 87 const float inv = 1.0f / max(S, 1e-12f);
78 88 float dot = 0.0f;
79 − for (uint j = 0; j < E; ++j)
89 + for (uint j = 0; j < p.E; ++j)
80 90 if (kept[j]) dot += dout[j] * in[j] * inv;
81 − for (uint j = 0; j < E; ++j)
91 + for (uint j = 0; j < p.E; ++j)
82 92 if (kept[j]) dp[j] += (dout[j] - dot) * inv;
83 93 }
84 94
95 +// counts[e] += number of rows whose gate for expert e is nonzero — the load
96 +// statistic driving the noaux bias update. One thread per expert.
97 +kernel void expert_counts_f32(device const float* G [[buffer(0)]],
98 + device float* COUNTS [[buffer(1)]],
99 + constant uint2& p [[buffer(2)]], // (E, N)
100 + uint e [[thread_position_in_grid]]) {
101 + float acc = 0.0f;
102 + for (uint i = 0; i < p.y; ++i)
103 + if (G[ulong(i) * p.x + e] != 0.0f) acc += 1.0f;
104 + COUNTS[e] += acc;
105 +}
106 +
85 107 // 12-byte layout matching the host-side struct (uint3 would pad to 16).
86 108 struct RowScaleParams { uint C, E, e; };
87 109
modified src/nn/config.h +35 −7
@@ -1,6 +1,7 @@
1 1 // Author: Simon-Pierre Boucher — contact@spboucher.ai
2 2 #pragma once
3 3
4 +#include <algorithm>
4 5 #include <cstdint>
5 6 #include <fstream>
6 7 #include <stdexcept>
@@ -57,6 +58,16 @@ struct ModelConfig {
57 58 int64_t moe_top_k = 2;
58 59 float moe_aux_weight = 0.01f; // load-balance loss: E * sum_e mean_gate_e^2
59 60 int64_t n_shared_experts = 0; // always-active experts (DeepSeek/Kimi style)
61 + // DeepSeek-V3-style routing refinements:
62 + std::string moe_scoring = "softmax"; // "softmax" | "sigmoid" (V3/K2)
63 + bool moe_norm_topk = true; // renormalize kept gates to sum 1
64 + float routed_scaling_factor = 1.0f; // y = shared + factor * sum(gated) (V3: 2.5)
65 + int64_t moe_d_ff = 0; // per-expert d_ff; 0 = d_ff (V3: 2048 vs 18432)
66 + int64_t first_k_dense = 0; // first k layers keep a dense MLP (V3: 3, K2: 1)
67 + // Aux-loss-FREE balancing (V3 "noaux"): select top-k on s+b but gate on s;
68 + // after each step b += gamma for underloaded experts, -= gamma for
69 + // overloaded. 0 disables (bias stays zero).
70 + float moe_bias_gamma = 0.0f; // V3: 0.001
60 71
61 72 int64_t head_dim() const {
62 73 return head_dim_override > 0 ? head_dim_override : d_model / n_heads;
@@ -79,17 +90,22 @@ struct ModelConfig {
79 90 + 2 * d_model * n_kv_heads * hd // wk, wv
80 91 + n_heads * hd * d_model; // wo
81 92 if (attention_bias) attn += (n_heads + 2 * n_kv_heads) * hd;
82 − const int64_t mlp_one = (activation == "swiglu")
83 − ? 3 * d_model * d_ff
84 − : 2 * d_model * d_ff;
85 − const int64_t mlp = n_experts > 0
86 − ? (n_experts + n_shared_experts) * mlp_one + n_experts * d_model
87 − : mlp_one; // experts + router
93 + const int64_t act_mats = (activation == "swiglu") ? 3 : 2;
94 + const int64_t mlp_dense = act_mats * d_model * d_ff;
95 + const int64_t expert_dff = (n_experts > 0 && moe_d_ff > 0) ? moe_d_ff : d_ff;
96 + const int64_t mlp_moe = (n_experts + n_shared_experts) * act_mats * d_model *
97 + expert_dff +
98 + n_experts * d_model; // experts + router
99 + // first_k_dense layers keep the dense MLP (V3-style).
100 + const int64_t n_moe_layers =
101 + n_experts > 0 ? n_layers - std::min(first_k_dense, n_layers) : 0;
102 + const int64_t mlp_total =
103 + n_moe_layers * mlp_moe + (n_layers - n_moe_layers) * mlp_dense;
88 104 const int64_t norms_per_layer = norm_placement == "sandwich" ? 4 : 2;
89 105 int64_t norms = (norm == "layernorm" ? 2 : 1) * d_model *
90 106 (norms_per_layer * n_layers + 1);
91 107 if (qk_norm) norms += 2 * head_dim() * n_layers;
92 − int64_t total = n_layers * (attn + mlp) + norms + vocab_size * d_model;
108 + int64_t total = n_layers * attn + mlp_total + norms + vocab_size * d_model;
93 109 if (!tied_embeddings) total += vocab_size * d_model;
94 110 if (!use_rope) total += context_length * d_model;
95 111 return total;
@@ -155,6 +171,12 @@ inline void from_json(const nlohmann::json& j, ModelConfig& c) {
155 171 c.moe_top_k = j.value("moe_top_k", c.moe_top_k);
156 172 c.moe_aux_weight = j.value("moe_aux_weight", c.moe_aux_weight);
157 173 c.n_shared_experts = j.value("n_shared_experts", c.n_shared_experts);
174 + c.moe_scoring = j.value("moe_scoring", c.moe_scoring);
175 + c.moe_norm_topk = j.value("moe_norm_topk", c.moe_norm_topk);
176 + c.routed_scaling_factor = j.value("routed_scaling_factor", c.routed_scaling_factor);
177 + c.moe_d_ff = j.value("moe_d_ff", c.moe_d_ff);
178 + c.first_k_dense = j.value("first_k_dense", c.first_k_dense);
179 + c.moe_bias_gamma = j.value("moe_bias_gamma", c.moe_bias_gamma);
158 180 c.qk_norm = j.value("qk_norm", c.qk_norm);
159 181 c.final_softcap = j.value("final_softcap", c.final_softcap);
160 182 c.scale_embeddings = j.value("scale_embeddings", c.scale_embeddings);
@@ -197,6 +219,12 @@ inline void from_json(const nlohmann::json& j, ModelConfig& c) {
197 219 throw std::runtime_error("config: moe_top_k must be in [1, n_experts]");
198 220 if (c.n_shared_experts < 0 || (c.n_shared_experts > 0 && c.n_experts == 0))
199 221 throw std::runtime_error("config: n_shared_experts requires n_experts > 0");
222 + if (c.moe_scoring != "softmax" && c.moe_scoring != "sigmoid")
223 + throw std::runtime_error("config: moe_scoring must be softmax or sigmoid");
224 + if (c.first_k_dense < 0 || c.first_k_dense > c.n_layers)
225 + throw std::runtime_error("config: first_k_dense must be in [0, n_layers]");
226 + if (c.moe_d_ff < 0 || c.moe_bias_gamma < 0.0f)
227 + throw std::runtime_error("config: moe_d_ff and moe_bias_gamma must be >= 0");
200 228 if (c.final_softcap < 0.0f)
201 229 throw std::runtime_error("config: final_softcap must be >= 0");
202 230 }
modified src/nn/mlp.h +59 −12
@@ -3,6 +3,8 @@
3 3
4 4 #include "nn/config.h"
5 5 #include "nn/linear.h"
6 +#include "ops/cpu/cpu_ops.h"
7 +#include "ops/metal/metal_ops.h"
6 8 #include "ops/ops.h"
7 9
8 10 #include <memory>
@@ -20,6 +22,8 @@ public:
20 22 // x: [N, C] 2-D
21 23 virtual Var forward(const Var& x) const = 0;
22 24 virtual Var aux() const { return Var(); }
25 + // noaux load-balance bias step (V3); no-op for dense/softmax variants.
26 + virtual void bias_update(float /*gamma*/) {}
23 27 };
24 28
25 29 // SwiGLU: w2( silu(x w1) ⊙ (x w3) ) | GELU/ReLU²: w2( act(x w1) )
@@ -58,43 +62,65 @@ private:
58 62 std::unique_ptr<Linear> w1_, w2_, w3_;
59 63 };
60 64
61 −// Mixture-of-experts MLP: softmax router over n_experts expert MLPs (d_ff
62 −// each), top-k gates renormalized to sum 1. v1 computes EVERY expert densely
63 −// and weights by the (mostly zero) gates — correctness first; token
64 −// gather/scatter sparsity is a later optimization. The auxiliary
65 −// load-balance loss is the differentiable proxy E · Σ_e (mean_i gate
66 −// probs)² — minimized (at 1.0) by a uniform router — scaled by
67 −// moe_aux_weight and added to the training loss by Transformer::loss.
65 +// Mixture-of-experts MLP. Router scoring is softmax (classic) or sigmoid
66 +// (DeepSeek V3/K2); top-k selection may be steered by a per-expert balance
67 +// bias (aux-loss-free balancing: bias_update() nudges it ±gamma by observed
68 +// load after each optimizer step), gates are renormalized or raw, and the
69 +// routed sum can be rescaled (routed_scaling_factor). Experts may use their
70 +// own d_ff (moe_d_ff). v1 computes EVERY expert densely and weights by the
71 +// (mostly zero) gates — correctness first; token gather/scatter sparsity is
72 +// a later optimization. The optional differentiable aux loss
73 +// E · Σ_e (mean_i score)² remains available (moe_aux_weight).
68 74 class MoEMLP : public MLPBase {
69 75 public:
70 76 MoEMLP(const ModelConfig& cfg, float proj_std, std::mt19937_64& rng)
71 77 : n_experts_(cfg.n_experts), top_k_(cfg.moe_top_k),
72 − aux_weight_(cfg.moe_aux_weight) {
78 + aux_weight_(cfg.moe_aux_weight),
79 + sigmoid_(cfg.moe_scoring == "sigmoid"),
80 + norm_topk_(cfg.moe_norm_topk),
81 + routed_scale_(cfg.routed_scaling_factor),
82 + track_load_(cfg.moe_bias_gamma > 0.0f) {
73 83 router_ = std::make_unique<Linear>(cfg.d_model, n_experts_, false, 0.02f, rng);
74 84 absorb("router", *router_);
85 + ModelConfig ecfg = cfg;
86 + if (cfg.moe_d_ff > 0) ecfg.d_ff = cfg.moe_d_ff;
75 87 experts_.reserve(size_t(n_experts_));
76 88 for (int64_t e = 0; e < n_experts_; ++e) {
77 − experts_.push_back(std::make_unique<MLP>(cfg, proj_std, rng));
89 + experts_.push_back(std::make_unique<MLP>(ecfg, proj_std, rng));
78 90 absorb("experts." + std::to_string(e), *experts_.back());
79 91 }
80 92 // Always-active experts, added ungated (DeepSeek/Kimi style): they
81 93 // absorb common knowledge so routed experts can specialize.
82 94 for (int64_t s = 0; s < cfg.n_shared_experts; ++s) {
83 − shared_.push_back(std::make_unique<MLP>(cfg, proj_std, rng));
95 + shared_.push_back(std::make_unique<MLP>(ecfg, proj_std, rng));
84 96 absorb("shared." + std::to_string(s), *shared_.back());
85 97 }
98 + bias_ = Tensor::zeros({n_experts_});
99 + counts_ = Tensor::zeros({n_experts_});
100 + // Registered (grad-free) so checkpoints and .forge snapshots carry
101 + // the balance state; the optimizer skips non-grad params.
102 + register_param("router.balance_bias", Var(bias_, /*requires_grad=*/false));
86 103 }
87 104
88 105 Var forward(const Var& x) const override {
89 106 const int64_t N = x.value().size(0);
90 − Var probs = ops::softmax(router_->forward(x)); // [N, E]
91 − Var gates = ops::topk_renorm(probs, top_k_); // [N, E], rows sum to 1
107 + Var logits = router_->forward(x); // [N, E]
108 + Var probs = sigmoid_ ? ops::sigmoid(logits) : ops::softmax(logits);
109 + Var gates = ops::topk_renorm(probs, bias_, top_k_, norm_topk_);
92 110
93 111 Var y = ops::row_scale(experts_[0]->forward(x), gates, 0);
94 112 for (int64_t e = 1; e < n_experts_; ++e)
95 113 y = ops::add(y, ops::row_scale(experts_[size_t(e)]->forward(x), gates, e));
114 + if (routed_scale_ != 1.0f) y = ops::scale(y, routed_scale_);
96 115 for (const auto& s : shared_) y = ops::add(y, s->forward(x));
97 116
117 + if (track_load_) {
118 + if (ops::backend() == ops::Backend::Metal)
119 + metal::expert_counts(gates.value(), counts_);
120 + else
121 + cpu::expert_counts(gates.value(), counts_);
122 + }
123 +
98 124 if (aux_weight_ > 0.0f) {
99 125 // Column means of probs via a constant 1/N row — every step is an
100 126 // existing autograd op, so the router gets balance gradients.
@@ -112,12 +138,33 @@ public:
112 138
113 139 Var aux() const override { return aux_; }
114 140
141 + // V3 noaux rule: overloaded experts (above mean load) lose selection
142 + // bias, underloaded gain. Called by the trainer AFTER the optimizer's
143 + // final sync — CPU reads/writes are safe there.
144 + void bias_update(float gamma) override {
145 + float* c = counts_.data<float>();
146 + float* b = bias_.data<float>();
147 + double total = 0.0;
148 + for (int64_t e = 0; e < n_experts_; ++e) total += c[e];
149 + const float mean = float(total / double(n_experts_));
150 + for (int64_t e = 0; e < n_experts_; ++e) {
151 + if (c[e] > mean) b[e] -= gamma;
152 + else if (c[e] < mean) b[e] += gamma;
153 + c[e] = 0.0f;
154 + }
155 + }
156 +
115 157 private:
116 158 int64_t n_experts_, top_k_;
117 159 float aux_weight_;
160 + bool sigmoid_, norm_topk_;
161 + float routed_scale_;
162 + bool track_load_;
118 163 std::unique_ptr<Linear> router_;
119 164 std::vector<std::unique_ptr<MLP>> experts_;
120 165 std::vector<std::unique_ptr<MLP>> shared_;
166 + Tensor bias_; // [E] selection bias (noaux balancing state)
167 + mutable Tensor counts_; // [E] load accumulator between bias updates
121 168 mutable Var aux_; // set by the last forward; consumed by Transformer::loss
122 169 };
123 170
modified src/nn/transformer.h +10 −2
@@ -44,8 +44,10 @@ public:
44 44 norm1_(std::make_unique<Norm>(cfg)),
45 45 attn_(std::make_unique<CausalSelfAttention>(cfg, proj_std, rng, layer_idx)),
46 46 norm2_(std::make_unique<Norm>(cfg)) {
47 − if (cfg.n_experts > 0) mlp_ = std::make_unique<MoEMLP>(cfg, proj_std, rng);
48 − else mlp_ = std::make_unique<MLP>(cfg, proj_std, rng);
47 + if (cfg.n_experts > 0 && layer_idx >= cfg.first_k_dense)
48 + mlp_ = std::make_unique<MoEMLP>(cfg, proj_std, rng);
49 + else
50 + mlp_ = std::make_unique<MLP>(cfg, proj_std, rng);
49 51 absorb("norm1", *norm1_);
50 52 absorb("attn", *attn_);
51 53 absorb("norm2", *norm2_);
@@ -90,6 +92,7 @@ public:
90 92
91 93 // MoE load-balance loss of the last forward (undefined for dense MLP).
92 94 Var moe_aux() const { return mlp_->aux(); }
95 + void moe_bias_update(float gamma) { mlp_->bias_update(gamma); }
93 96
94 97 private:
95 98 enum class Placement { Pre, Post, Sandwich };
@@ -176,6 +179,11 @@ public:
176 179
177 180 const ModelConfig& config() const { return cfg_; }
178 181
182 + // noaux balancing step (V3): call after the optimizer sync, CPU-side.
183 + void update_moe_bias(float gamma) {
184 + for (auto& blk : blocks_) blk->moe_bias_update(gamma);
185 + }
186 +
179 187 private:
180 188 ModelConfig cfg_;
181 189 std::unique_ptr<Embedding> tok_emb_;
modified src/ops/cpu/cpu_ops.cpp +47 −10
@@ -378,42 +378,64 @@ void fake_quant(const Tensor& w, int mode, Tensor& out) {
378 378 }
379 379 }
380 380
381 +void sigmoid(const Tensor& x, Tensor& out) {
382 + const float* px = x.data<float>();
383 + float* po = out.data<float>();
384 + for (int64_t i = 0; i < x.numel(); ++i) po[i] = 1.0f / (1.0f + std::exp(-px[i]));
385 +}
386 +
387 +void sigmoid_backward(const Tensor& y, const Tensor& dout, Tensor& dx) {
388 + const float* py = y.data<float>();
389 + const float* pd = dout.data<float>();
390 + float* pdx = dx.data<float>();
391 + for (int64_t i = 0; i < y.numel(); ++i) pdx[i] += pd[i] * py[i] * (1.0f - py[i]);
392 +}
393 +
381 394 namespace {
382 −// Kept set of the k largest entries, ties to the lower index (matches Metal).
383 −void topk_select(const float* row, int64_t E, int64_t k, bool* kept, float* sum) {
395 +// Kept set of the k largest (score + bias) entries, ties to the lower index
396 +// (matches Metal). The returned sum is over the BIASLESS scores — the bias
397 +// only steers selection (DeepSeek-V3 noaux routing).
398 +void topk_select(const float* row, const float* bias, int64_t E, int64_t k,
399 + bool* kept, float* sum) {
384 400 for (int64_t j = 0; j < E; ++j) kept[j] = false;
385 401 float S = 0.0f;
386 402 for (int64_t sel = 0; sel < k; ++sel) {
387 403 float best = -std::numeric_limits<float>::max();
388 404 int64_t arg = 0;
389 − for (int64_t j = 0; j < E; ++j)
390 − if (!kept[j] && row[j] > best) { best = row[j]; arg = j; }
405 + for (int64_t j = 0; j < E; ++j) {
406 + const float v = row[j] + bias[j];
407 + if (!kept[j] && v > best) { best = v; arg = j; }
408 + }
391 409 kept[arg] = true;
392 − S += best;
410 + S += row[arg];
393 411 }
394 412 *sum = S;
395 413 }
396 414 } // namespace
397 415
398 −void topk_renorm(const Tensor& p, int64_t k, Tensor& out) {
416 +void topk_renorm(const Tensor& p, const Tensor& bias, int64_t k, bool norm,
417 + Tensor& out) {
399 418 const int64_t E = p.shape().back();
400 419 const int64_t N = p.numel() / E;
401 420 const float* pp = p.data<float>();
421 + const float* pb = bias.data<float>();
402 422 float* po = out.data<float>();
403 423 bool kbuf[64];
404 424 for (int64_t i = 0; i < N; ++i) {
405 425 const float* row = pp + i * E;
406 426 float S = 0.0f;
407 − topk_select(row, E, k, kbuf, &S);
408 − const float inv = 1.0f / std::max(S, 1e-12f);
427 + topk_select(row, pb, E, k, kbuf, &S);
428 + const float inv = norm ? 1.0f / std::max(S, 1e-12f) : 1.0f;
409 429 for (int64_t j = 0; j < E; ++j) po[i * E + j] = kbuf[j] ? row[j] * inv : 0.0f;
410 430 }
411 431 }
412 432
413 −void topk_renorm_backward(const Tensor& p, const Tensor& dout, int64_t k, Tensor& dp) {
433 +void topk_renorm_backward(const Tensor& p, const Tensor& bias, const Tensor& dout,
434 + int64_t k, bool norm, Tensor& dp) {
414 435 const int64_t E = p.shape().back();
415 436 const int64_t N = p.numel() / E;
416 437 const float* pp = p.data<float>();
438 + const float* pb = bias.data<float>();
417 439 const float* pd = dout.data<float>();
418 440 float* pdp = dp.data<float>();
419 441 bool kbuf[64];
@@ -421,7 +443,12 @@ void topk_renorm_backward(const Tensor& p, const Tensor& dout, int64_t k, Tensor
421 443 const float* row = pp + i * E;
422 444 const float* drow = pd + i * E;
423 445 float S = 0.0f;
424 − topk_select(row, E, k, kbuf, &S);
446 + topk_select(row, pb, E, k, kbuf, &S);
447 + if (!norm) {
448 + for (int64_t j = 0; j < E; ++j)
449 + if (kbuf[j]) pdp[i * E + j] += drow[j];
450 + continue;
451 + }
425 452 const float inv = 1.0f / std::max(S, 1e-12f);
426 453 float dot = 0.0f;
427 454 for (int64_t j = 0; j < E; ++j)
@@ -431,6 +458,16 @@ void topk_renorm_backward(const Tensor& p, const Tensor& dout, int64_t k, Tensor
431 458 }
432 459 }
433 460
461 +void expert_counts(const Tensor& gates, Tensor& counts) {
462 + const int64_t E = gates.shape().back();
463 + const int64_t N = gates.numel() / E;
464 + const float* pg = gates.data<float>();
465 + float* pc = counts.data<float>();
466 + for (int64_t i = 0; i < N; ++i)
467 + for (int64_t e = 0; e < E; ++e)
468 + if (pg[i * E + e] != 0.0f) pc[e] += 1.0f;
469 +}
470 +
434 471 void row_scale(const Tensor& x, const Tensor& gates, int64_t e, Tensor& out) {
435 472 const int64_t C = x.shape().back();
436 473 const int64_t N = x.numel() / C;
modified src/ops/cpu/cpu_ops.h +12 −4
@@ -93,10 +93,18 @@ void attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,
93 93 // Per-row fake quantization of a 2-D weight; mode 0 = int8 (absmax/127),
94 94 // mode 1 = ternary (absmean, BitNet-style). Backward is STE — identity.
95 95 void fake_quant(const Tensor& w, int mode, Tensor& out);
96 −// Keep top-k per row of p [N,E], renormalize kept entries to sum 1, zero the
97 −// rest (ties broken by lower index, same as the Metal kernel).
98 −void topk_renorm(const Tensor& p, int64_t k, Tensor& out);
99 −void topk_renorm_backward(const Tensor& p, const Tensor& dout, int64_t k, Tensor& dp);
96 +void sigmoid(const Tensor& x, Tensor& out);
97 +// dx += dout * y * (1-y) — takes the forward OUTPUT y
98 +void sigmoid_backward(const Tensor& y, const Tensor& dout, Tensor& dx);
99 +// Keep top-k per row of p [N,E] — selected by p+bias, gated by p alone
100 +// (bias [E], all-zero for classic routing); norm renormalizes kept gates to
101 +// sum 1. Ties broken by lower index, same as the Metal kernel.
102 +void topk_renorm(const Tensor& p, const Tensor& bias, int64_t k, bool norm,
103 + Tensor& out);
104 +void topk_renorm_backward(const Tensor& p, const Tensor& bias, const Tensor& dout,
105 + int64_t k, bool norm, Tensor& dp);
106 +// counts[e] += #rows with nonzero gate for expert e (noaux load statistic)
107 +void expert_counts(const Tensor& gates, Tensor& counts);
100 108 // out[i,:] = x[i,:] * gates[i,e]; _accumulate does dst += (dx path);
101 109 // gate_backward does dgates[i,e] += dot(dout[i,:], x[i,:]).
102 110 void row_scale(const Tensor& x, const Tensor& gates, int64_t e, Tensor& out);
modified src/ops/metal/metal_ops.cpp +23 −6
@@ -299,18 +299,35 @@ void softmax_backward(const Tensor& p, const Tensor& dout, Tensor& dx) {
299 299 encode_flat("softmax_bwd_f32", {&p, &dout, &dx}, &c32, sizeof(c32), p.numel() / C);
300 300 }
301 301
302 −void topk_renorm(const Tensor& p, int64_t k, Tensor& out) {
302 +void sigmoid(const Tensor& x, Tensor& out) {
303 + encode_flat("sigmoid_f32", {&x, &out}, nullptr, 0, x.numel());
304 +}
305 +
306 +void sigmoid_backward(const Tensor& y, const Tensor& dout, Tensor& dx) {
307 + encode_flat("sigmoid_bwd_f32", {&y, &dout, &dx}, nullptr, 0, y.numel());
308 +}
309 +
310 +void topk_renorm(const Tensor& p, const Tensor& bias, int64_t k, bool norm,
311 + Tensor& out) {
303 312 const int64_t E = p.shape().back();
304 313 check(E <= 64, "topk_renorm: E must be <= 64 (kernel's kept[] bound)");
305 − const uint32_t pr[2] = {uint32_t(E), uint32_t(k)};
306 − encode_flat("topk_renorm_f32", {&p, &out}, pr, sizeof(pr), p.numel() / E);
314 + const uint32_t pr[3] = {uint32_t(E), uint32_t(k), norm ? 1u : 0u};
315 + encode_flat("topk_renorm_f32", {&p, &bias, &out}, pr, sizeof(pr), p.numel() / E);
307 316 }
308 317
309 −void topk_renorm_backward(const Tensor& p, const Tensor& dout, int64_t k, Tensor& dp) {
318 +void topk_renorm_backward(const Tensor& p, const Tensor& bias, const Tensor& dout,
319 + int64_t k, bool norm, Tensor& dp) {
310 320 const int64_t E = p.shape().back();
311 321 check(E <= 64, "topk_renorm_backward: E must be <= 64");
312 − const uint32_t pr[2] = {uint32_t(E), uint32_t(k)};
313 − encode_flat("topk_renorm_bwd_f32", {&p, &dout, &dp}, pr, sizeof(pr), p.numel() / E);
322 + const uint32_t pr[3] = {uint32_t(E), uint32_t(k), norm ? 1u : 0u};
323 + encode_flat("topk_renorm_bwd_f32", {&p, &bias, &dout, &dp}, pr, sizeof(pr),
324 + p.numel() / E);
325 +}
326 +
327 +void expert_counts(const Tensor& gates, Tensor& counts) {
328 + const int64_t E = gates.shape().back();
329 + const uint32_t pr[2] = {uint32_t(E), uint32_t(gates.numel() / E)};
330 + encode_flat("expert_counts_f32", {&gates, &counts}, pr, sizeof(pr), E);
314 331 }
315 332
316 333 namespace {
modified src/ops/metal/metal_ops.h +10 −3
@@ -160,9 +160,16 @@ void cross_entropy(const Tensor& logits, const Tensor& targets, int64_t n_valid,
160 160 void fake_quant(const Tensor& w, int mode, Tensor& out);
161 161 // dx += p ∘ (dout − dot(dout, p)) per row; thread-per-row, small last dims.
162 162 void softmax_backward(const Tensor& p, const Tensor& dout, Tensor& dx);
163 −// Keep top-k per row of p [N,E], renormalize kept to sum 1, zero the rest.
164 −void topk_renorm(const Tensor& p, int64_t k, Tensor& out);
165 −void topk_renorm_backward(const Tensor& p, const Tensor& dout, int64_t k, Tensor& dp);
163 +void sigmoid(const Tensor& x, Tensor& out);
164 +void sigmoid_backward(const Tensor& y, const Tensor& dout, Tensor& dx);
165 +// Keep top-k per row of p [N,E] — selected by p+bias, gated by p alone;
166 +// norm renormalizes kept gates to sum 1.
167 +void topk_renorm(const Tensor& p, const Tensor& bias, int64_t k, bool norm,
168 + Tensor& out);
169 +void topk_renorm_backward(const Tensor& p, const Tensor& bias, const Tensor& dout,
170 + int64_t k, bool norm, Tensor& dp);
171 +// counts[e] += #rows with nonzero gate for expert e
172 +void expert_counts(const Tensor& gates, Tensor& counts);
166 173 // out[i,:] = x[i,:] * gates[i,e]; the _accumulate variant does dst += (dx path).
167 174 void row_scale(const Tensor& x, const Tensor& gates, int64_t e, Tensor& out);
168 175 void row_scale_accumulate(const Tensor& x, const Tensor& gates, int64_t e, Tensor& dst);
modified src/ops/ops.cpp +25 −6
@@ -279,20 +279,39 @@ Var fake_quant(const Var& w, QuantMode mode) {
279 279 return result;
280 280 }
281 281
282 −Var topk_renorm(const Var& probs, int64_t k) {
282 +Var sigmoid(const Var& x) {
283 + Tensor out = Tensor::empty(x.value().shape());
284 + if (gpu()) metal::sigmoid(x.value(), out);
285 + else cpu::sigmoid(x.value(), out);
286 +
287 + const bool needs = grad_needed({&x});
288 + Var result(std::move(out), needs);
289 + if (needs) {
290 + Tape::get().record([x, result]() {
291 + if (!x.requires_grad()) return;
292 + if (gpu()) metal::sigmoid_backward(result.value(), result.grad(), x.grad());
293 + else cpu::sigmoid_backward(result.value(), result.grad(), x.grad());
294 + });
295 + }
296 + return result;
297 +}
298 +
299 +Var topk_renorm(const Var& probs, const Tensor& bias, int64_t k, bool norm) {
283 300 Tensor out = Tensor::empty(probs.value().shape());
284 − if (gpu()) metal::topk_renorm(probs.value(), k, out);
285 − else cpu::topk_renorm(probs.value(), k, out);
301 + if (gpu()) metal::topk_renorm(probs.value(), bias, k, norm, out);
302 + else cpu::topk_renorm(probs.value(), bias, k, norm, out);
286 303
287 304 const bool needs = grad_needed({&probs});
288 305 Var result(std::move(out), needs);
289 306 if (needs) {
290 − Tape::get().record([probs, k, result]() {
307 + Tape::get().record([probs, bias, k, norm, result]() {
291 308 if (!probs.requires_grad()) return;
292 309 if (gpu())
293 − metal::topk_renorm_backward(probs.value(), result.grad(), k, probs.grad());
310 + metal::topk_renorm_backward(probs.value(), bias, result.grad(), k,
311 + norm, probs.grad());
294 312 else
295 − cpu::topk_renorm_backward(probs.value(), result.grad(), k, probs.grad());
313 + cpu::topk_renorm_backward(probs.value(), bias, result.grad(), k,
314 + norm, probs.grad());
296 315 });
297 316 }
298 317 return result;
modified src/ops/ops.h +7 −3
@@ -32,6 +32,7 @@ Var scale(const Var& a, float s);
32 32 Var silu(const Var& x);
33 33 Var gelu(const Var& x);
34 34 Var relu2(const Var& x); // max(x,0)^2 — nanoGPT-speedrun activation
35 +Var sigmoid(const Var& x); // σ(x) — sigmoid MoE routing (DeepSeek V3)
35 36
36 37 // Gemma-style soft capping: cap * tanh(x / cap). Bounds logits smoothly.
37 38 Var softcap(const Var& x, float cap);
@@ -46,9 +47,12 @@ Var softmax(const Var& x);
46 47 enum class QuantMode { None, Int8, Ternary };
47 48 Var fake_quant(const Var& w, QuantMode mode);
48 49
49 −// MoE gating: keep top-k per row, renormalize kept entries to sum 1, zero
50 −// the rest. Gradients flow through kept entries only.
51 −Var topk_renorm(const Var& probs, int64_t k);
50 +// MoE gating: keep top-k per row of probs, zero the rest. Selection ranks
51 +// probs + bias (bias [E], requires_grad=false — the DeepSeek-V3 noaux
52 +// balance bias; pass zeros for classic routing) but gate values are the
53 +// biasless probs; norm renormalizes kept gates to sum 1. Gradients flow
54 +// through kept entries only.
55 +Var topk_renorm(const Var& probs, const Tensor& bias, int64_t k, bool norm = true);
52 56 // y[i,:] = x[i,:] * gates[i,e] — weight expert e's output by its gate column.
53 57 Var row_scale(const Var& x, const Var& gates, int64_t e);
54 58
modified src/train/trainer.cpp +2 −0
@@ -149,6 +149,8 @@ void Trainer::train(const std::string& resume_from) {
149 149 }
150 150
151 151 const float grad_norm = opt_->step_with_clip(lr, tc.grad_clip);
152 + if (cfg_.model.n_experts > 0 && cfg_.model.moe_bias_gamma > 0.0f)
153 + model_->update_moe_bias(cfg_.model.moe_bias_gamma);
152 154
153 155 const auto t1 = std::chrono::steady_clock::now();
154 156 const double dt = std::chrono::duration<double>(t1 - t0).count();
modified tests/test_ops.cpp +39 −7
@@ -280,25 +280,49 @@ void test_metal_moe_quant_ops() {
280 280 expect_close(gpu, ref, mode == 0 ? "fake_quant int8" : "fake_quant ternary");
281 281 }
282 282
283 − // topk_renorm forward + backward, row_scale family
283 + // topk_renorm forward + backward (biased selection + both norm modes),
284 + // row_scale family
284 285 const int64_t N = 130, E = 8, C = 37, K = 2;
285 286 forge::Tensor probs = forge::Tensor::empty({N, E});
286 287 fill_random(probs, rng, 0.01f, 1.0f);
288 + forge::Tensor bias = forge::Tensor::empty({E});
289 + fill_random(bias, rng, -0.3f, 0.3f);
287 290 forge::Tensor ref = forge::Tensor::empty({N, E});
288 291 forge::Tensor gpu = forge::Tensor::empty({N, E});
289 − forge::cpu::topk_renorm(probs, K, ref);
290 − forge::metal::topk_renorm(probs, K, gpu);
292 + forge::cpu::topk_renorm(probs, bias, K, true, ref);
293 + forge::metal::topk_renorm(probs, bias, K, true, gpu);
291 294 forge::metal::sync();
292 − expect_close(gpu, ref, "topk_renorm fwd");
295 + expect_close(gpu, ref, "topk_renorm fwd (biased)");
293 296
294 297 forge::Tensor dout = forge::Tensor::empty({N, E});
295 298 fill_random(dout, rng);
296 299 forge::Tensor dref = forge::Tensor::zeros({N, E});
297 300 forge::Tensor dgpu = forge::Tensor::zeros({N, E});
298 − forge::cpu::topk_renorm_backward(probs, dout, K, dref);
299 − forge::metal::topk_renorm_backward(probs, dout, K, dgpu);
301 + forge::cpu::topk_renorm_backward(probs, bias, dout, K, true, dref);
302 + forge::metal::topk_renorm_backward(probs, bias, dout, K, true, dgpu);
300 303 forge::metal::sync();
301 − expect_close(dgpu, dref, "topk_renorm bwd");
304 + expect_close(dgpu, dref, "topk_renorm bwd (biased)");
305 +
306 + forge::Tensor ref_nn = forge::Tensor::empty({N, E});
307 + forge::Tensor gpu_nn = forge::Tensor::empty({N, E});
308 + forge::cpu::topk_renorm(probs, bias, K, false, ref_nn);
309 + forge::metal::topk_renorm(probs, bias, K, false, gpu_nn);
310 + forge::metal::sync();
311 + expect_close(gpu_nn, ref_nn, "topk fwd (no renorm)");
312 +
313 + forge::Tensor sig_ref = forge::Tensor::empty({N, E});
314 + forge::Tensor sig_gpu = forge::Tensor::empty({N, E});
315 + forge::cpu::sigmoid(probs, sig_ref);
316 + forge::metal::sigmoid(probs, sig_gpu);
317 + forge::metal::sync();
318 + expect_close(sig_gpu, sig_ref, "sigmoid fwd");
319 +
320 + forge::Tensor cnt_ref = forge::Tensor::zeros({E});
321 + forge::Tensor cnt_gpu = forge::Tensor::zeros({E});
322 + forge::cpu::expert_counts(ref, cnt_ref);
323 + forge::metal::expert_counts(ref, cnt_gpu);
324 + forge::metal::sync();
325 + expect_close(cnt_gpu, cnt_ref, "expert_counts");
302 326
303 327 forge::Tensor x = forge::Tensor::empty({N, C});
304 328 fill_random(x, rng);
@@ -593,6 +617,14 @@ int main() {
593 617 mcfg.n_experts = 4; mcfg.moe_top_k = 2; mcfg.n_shared_experts = 1;
594 618 test_backend_parity_model("MoE 4+1shared top-2 + aux", mcfg);
595 619
620 + // V3-style MoE: sigmoid scoring, no renorm, routed scaling, own d_ff,
621 + // first layer dense, noaux bias tracking on
622 + mcfg.moe_scoring = "sigmoid"; mcfg.moe_norm_topk = false;
623 + mcfg.routed_scaling_factor = 2.5f; mcfg.moe_d_ff = 16;
624 + mcfg.first_k_dense = 1; mcfg.moe_bias_gamma = 0.001f;
625 + mcfg.moe_aux_weight = 0.0f;
626 + test_backend_parity_model("MoE V3-style sigmoid/noaux/scaled", mcfg);
627 +
596 628 // Architecture-variant knobs: QK-norm + logit softcap + embed scaling
597 629 forge::ModelConfig vcfg;
598 630 vcfg.n_layers = 2; vcfg.d_model = 16; vcfg.n_heads = 2; vcfg.n_kv_heads = 1;
599 631