| |
| |
| |
| |
| |
| |
| |
| """ |
| Referencyjny ZWYKLY transformer (byte-level nanoGPT-style) ~25M — apples-to-apples vs BDH-25M. |
| Ta sama data (train.bin/val.bin uint8), ten sam scale (~25M), ten sam byte-level (vocab256). |
| Rozni sie TYLKO architektura (standard causal transformer vs BDH fast-weights) -> czysta referencja. |
| |
| CLI mirror train_bdh.py. GPU ROCm/CUDA bf16, cosine+warmup+clip, ckpt/resume, logging. |
| Autor: Hart (N-02). |
| |
| Smoke throughput (bez danych PII, syntetyczny bufor): --synthetic --steps 60 |
| Realny: --data-dir . --run-id gpt25m_run1 --steps 30000 |
| r8 (Monter): wiele GPU przez torchrun (DDP). --batch = batch GLOBALNY; kazdy rank bierze swoje wiersze z tego samego |
| batcha (te same starty okien co przebieg 1-GPU), gradienty usredniane przez DDP, wiec krok = krok 1-GPU z tym samym |
| --batch (do numeryki). Logi, metryki, zdarzenia i ckpt pisze tylko rank 0. Bez torchrun (WORLD_SIZE=1) = r7 bez zmian. |
| torchrun --standalone --nproc-per-node 4 train_gpt_ref_r8.py ... --batch 64 |
| r11 (Monter, 10.10, Wzrost etap 1b: dlugie okna): dwie opcje tylko na pamiec, domyslnie WYLACZONE (bez nich = r10 bit w bit): |
| --act-ckpt przeliczanie aktywacji po bloku w treningu (torch.utils.checkpoint, use_reentrant=False); |
| --loss-chunk N strata liczona kawalkami po N pozycji (glowica + CE w checkpoint), bez pelnych logitow [B,T,V]; |
| forward zwraca wtedy logits=None, a strata = ta sama srednia po B*T (i ten sam z-loss). |
| Okno i baza RoPE z CLI dzialaja z --init-from bez zmian (RoPE bez parametrow, --pos rope nie ma pos-embeddingu). |
| r12 (Monter, 10.10, Slayer9-Eff): --loop K = wspoldzielenie wag miedzy warstwami (looped transformer, pomysl Kacpra 18:30Z): |
| n_layer unikalnych blokow przechodzi sie K razy (glebokosc liczenia n_layer*K, parametry n_layer). Tablica liczy parametry |
| modelu, wiec K nie zmienia rozmiaru. Blok 0 w kazdym przejsciu na nowo ustala v0 (value-residual). K w config ckpt ("loop"); |
| --init-from sprawdza zgodnosc K. Domyslnie K=1 = r11 bit w bit (config ckpt bez klucza "loop"). |
| r13 (Hart, 11.10, kampania 2 Slayer9, Latarnik 11:40Z): hybryda GatedDeltaNet2 (flash-linear-attention 0.5.2) + uwaga. |
| --mixer-layout STR litera na blok: A = uwaga jak r12, G = fla.layers.gdn2.GatedDeltaNet2 (hidden d, head_dim |
| --gdn-head-dim, num_heads d/head_dim, expand_v 1, short conv 4, mode chunk) w miejsce uwagi; FFN, |
| normy i reszta bloku bez zmian. Domyslnie brak = same A = r12 bit w bit (config bez nowych kluczy). |
| Value-residual: v0 ustala PIERWSZY blok A (w r12 blok 0). Wagi G inicjalizuje fla (bez _init GPT). |
| Optymalizator: w blokach G Muon tylko q/k/v/o_proj; reszta G (konwolucje, bramki, A_log, dt_bias, |
| normy) -> AdamW (karta kampanii 2: czesc „budowy Q-ARCH”, ustalona przed biegiem). |
| --attn-bias 0|1 bias w qkv/proj uwagi (domyslnie 1 jak r12; Qielbas M2 = 0). |
| --rope-style interleaved (r12) | half (rotate-half jak Qielbas M2; baza z --rope-theta). |
| Import fla tylko przy G w ukladzie (PYTHONPATH z fla-lib), wiec same A dzialaja bez fla. |
| r14 (Hart, 11.10, kampania 3 Slayer9: KD w pretreningu, Arek 15:56Z): destylacja od nauczyciela na zywo, domyslnie WYLACZONA |
| (bez --kd-teacher = r13 bit w bit). |
| --kd-teacher CKPT --kd-teacher-sha HEX64 ckpt nauczyciela (r12/r13, ten sam trener i vocab co uczen; sha pelne, inaczej |
| STOP). Nauczyciel budowany z config ckpt, eval, bez gradientu, forward w tym samym autocast co uczen |
| (bf16), w kazdym kroku na tym samym batchu. Budowa nauczyciela nie zuzywa globalnego RNG (fork_rng). |
| --kd-alpha A strata = (1-A)*CE + A*T^2*KL(p_naucz || p_ucz), obie srednie po B*T pozycjach; KL po pelnym vocab, |
| w fp32, z logitow /T (uczen po --logit-cap, jak w CE). Z --loss-chunk liczone tymi samymi kawalkami. |
| --kd-temp T temperatura (domyslnie 1.0). |
| Walidacja (VAL loss) = sama CE jak w r13. W logu kroku dodatkowo ce (sama CE). W ckpt klucz "kd" (poza "config", |
| wiec valport/--init-from bez zmian). |
| r15 (Hart, 11.10, Arek 21:33Z „eksperymenty do skutku”, przeglad arXiv #4): faktoryzowane osadzenia, domyslnie WYLACZONE |
| (bez --emb-dim = r14 bit w bit). |
| --emb-dim E tok = Embedding(V, E) -> emb_up Linear(E, d); glowica: emb_down Linear(d, E) -> head Linear(E, V) |
| zwiazana z tok (te same wagi V x E). P osadzen V*E + 2*E*d zamiast V*d. emb_up/emb_down: init |
| N(0, 1/E) i N(0, 1/d) (skala strumienia jak bez faktoryzacji), optymalizator jak macierze ukryte (Muon). |
| W config ckpt "emb_dim" tylko przy E > 0. |
| r16 (Hart, 11.10, kolejka nocna pkt 4 Latarnika: biegi ~V1 po 1 na karte): akumulacja gradientu, domyslnie WYLACZONA |
| (bez --grad-accum = r15 bit w bit). |
| --grad-accum K wiersze ranku (batch / world) dzielone na K rownych, ciaglych mikrobatchy; backward (strata_j / K), |
| pod DDP no_sync poza ostatnim. Te same wiersze co bez akumulacji (sampler globalny), wiec krok = |
| krok bez akumulacji z dokladnoscia do kolejnosci sumowania. Strata logu = srednia mikrobatchy. |
| r17 (Hart, 12.10, karta Q-REPRO Mierniczego a92e717c: replika przepisu Qielbas M2): opcje optymalizatora z ich |
| `pretraining-report.json`, wszystkie domyslnie WYLACZONE (bez nich = r16 bit w bit). |
| --muon-scale match_rms_adamw skala kroku Muon jak torch.optim.Muon(adjust_lr_fn="match_rms_adamw"): |
| 0.2 * sqrt(max(wiersze, kolumny)) zamiast max(1, wiersze/kolumny)^0.5; z --muon-lr = --lr |
| (ich „peak lr=0.001” wspolne dla Muon i AdamW). |
| --muon-wd W decoupled weight decay na parametrach Muon (p *= 1 - lr*W, lr przed skalowaniem, jak torch); |
| domyslnie 0 (jak dotad: WD tylko na AdamW-aux). |
| --tf32 0 TF32 wylaczone w matmul/cudnn (ich „TF32 disabled”); domyslnie 1 (jak dotad). |
| Harmonogram cosine -> min-lr, warmup, swiezy optymalizator przy --init-from: bez zmian (juz byly). |
| r18 (Hart, 13.10, ramie TC-arch po krzywej rozmiaru). Fragmenty XSA, partial RoPE, ReLU^2 i zero-init ADAPTED FROM |
| kurcontko/Tiny-Chick-10M@60d32c4d tinylm_arch.py (sha256 4876824d...), Apache License 2.0 |
| (https://www.apache.org/licenses/LICENSE-2.0); zmiany: wpiecie w Block/GPT r17, flagi CLI, walidacja, config ckpt. |
| wszystkie domyslnie WYLACZONE (bez nich = r17 bit w bit, config ckpt bez nowych kluczy): |
| --xsa 1 po uwadze z wyjscia kazdej glowy usuwana skladowa wzdluz jej (znormalizowanego) wektora value tokenu: |
| y -= (y . v^) v^ (v = to, co weszlo do SDPA, czyli po value-residual). |
| --rope-dims N RoPE tylko na pierwszych N wymiarach glowy (parzyste, 0 < N < head_dim; czestotliwosci liczone z N), |
| reszta wymiarow bez pozycji; 0 = wszystkie (jak r17). Styl interleaved/half bez zmian, w obrebie N. |
| --ffn relu2 MLP down(relu(up(x))^2) bez biasu, hidden = round(--ffn-mult * d) (Tiny Chick: mult 4); gelu/swiglu bez zmian. |
| --zero-init 1 po inicjalizacji zerowane wagi wyjscia uwagi (proj, bloki A) i wyjscia MLP (down / ostatni Linear, kazdy blok); |
| bloki G zostaja z inicjalizacja fla. |
| --init-from i nauczyciel KD: config ckpt bez kluczy r18 = wartosci domyslne; rozjazd xsa/rope_dims/ffn z argumentami = STOP. |
| r19 (Hart, 13.10, karta TC-REPLIKA): backend --arch tinylm = model z ORYGINALNEGO pliku kurcontko/Tiny-Chick-10M@60d32c4d |
| tinylm_arch.py (sha256 4876824de8c697a6..., Apache License 2.0, uzyty bez zmian, ladowany z --tinylm-file i sprawdzany po |
| sha; NOTICE: "derived from kurcontko/Tiny-Chick-10M"). Bez --arch tinylm = r18 bit w bit (config ckpt bez nowych kluczy). |
| Z --arch tinylm: GPT trzyma TinyLM jako self.tl (klucze state_dict "tl.<ich nazwa>"); forward zwraca (logity, CE) jak GPT. |
| TinyConfig z argumentow: vocab, n_layer, n_embd, n_head, --kv-heads, --ffn relu2|swiglu, --ffn-mult, --qk-norm, |
| --rope-dims, --rope-theta, --xsa, --zero-init, --norm-eps, --block. Optymalizator Muon: tl.emb -> AdamW (jak tok/head), |
| skalarny gain QK-norm (ndim 3) -> AdamW, reszta 2D -> Muon. Niezgodne z tinylm (STOP): --loop > 1, --mixer-layout, |
| --emb-dim, --value-residual, --kd-teacher, --act-ckpt, --loss-chunk, --logit-cap, --z-loss, --pos learned, --rope-style |
| interleaved, --norm layernorm, --attn-bias 1. |
| r20 (Hart, 13.10): torch.load ckpt przy --resume i --init-from z weights_only=True (uwaga Wartownika do r18 l. 832: |
| weights_only=False wykonuje dowolny pickle). Ckpt trenera (model / opt / step / config / queue) laduje sie tak bez zmian |
| (sprawdzone na ba223162). Zapis ckpt po ostatnim kroku biegu także bez --queue, gdy end_step % ckpt_every != 0 |
| (r18/r19: tylko z --queue; CPT k9 13.10 skonczyl z ckpt 1140 zamiast 1145). Poza tym r20 = r19 bit w bit. |
| r21 (Monter dla Harta, 14.10, Axolote maggi-uxz6, przepis_axolote_v0 2.2): --train-only REGEX, domyslnie WYLACZONE (bez |
| flagi = r20 bit w bit, ckpt bez nowych kluczy). Parametry, ktorych nazwa nie pasuje do REGEX (re.search), maja |
| requires_grad False i nie trafiaja do optymalizatora (muon juz je pomijal, adamw filtruje), wiec zostaja bit w bit. |
| Tylko z --init-from; regex bez trafien albo chwytajacy osadzenia / tied head = STOP. Wzor: ft_arc11.py --train-only. |
| W ckpt klucz "train_only" {regex, n_trainable} poza "config" (jak "kd"), wiec --init-from i valport bez zmian. |
| r22 (Hart, 17.10, maggi-v38t / karta maggi-jyoy): --imha-t T, domyslnie 0 = WYLACZONE (r21 bit w bit, ckpt bez nowych kluczy). |
| T >= 1: podwarstwa uwagi blokow A = IMHA (Divakar i Priyadarshini, TAI 2(1) 2026, doi 10.53941/tai.2026.100014, Algorytm 1): |
| h0 = ln1(x); T iteracji tymi samymi qkv / proj: A = softmax(QK^T/sqrt(dk) + M) jawnie (fp32), z = AV; entropia wiersza |
| per glowa H_i znormalizowana przez log(i+1) (H_0 = 0); w_i = exp(-H_i/tau) / srednia prefiksu exp(-H_k/tau), k <= i; |
| z_i *= w_i; z = proj([z_h]); g_i = srednia prefiksu z; h += sigmoid(W_a[z;g]) * relu(W_f[z;g]); wyjscie x += h_T - h0. |
| Nowe parametry na blok: imha_a, imha_f (Linear 2d -> d, bez biasu: w artykule W in R^{2d x d}) i imha_log_tau |
| (tau = exp(log_tau) > 0, init tau = 1; artykul: "learnable temperature per layer, initialised at 1"). |
| Decyzje poza artykulem (nasz blok ma wiecej niz MHA z artykulu): qk-norm i RoPE stosowane w kazdej iteracji jak w r21; |
| value residual: v0 ustala blok pierwszy w iteracji 0, kolejne iteracje i bloki mieszaja v z v0 jak w r21. |
| STOP: --imha-t z --xsa, --mixer-layout (bloki G), --arch tinylm. Atrybut testowy cfg.imha_plain (tylko test_r22.py): |
| entropia, kontekst globalny i bramka wylaczone -> przy T = 1 ma byc rowne SDPA r21 (pozytywna kontrola z artykulu, 3.3.7). |
| Koszt: jawna macierz uwagi B x H x N x N zamiast SDPA (wolniej i wiecej pamieci) -> rownowazenie kosztu z tok/s z pomiaru. |
| r23 (Hart, 20.10, karta dzien dowodow 400M, Mierniczy 19:26Z): --kv-heads dla --arch gpt (GQA), domyslnie 0 = r22 bit w bit |
| (te same nazwy/ksztalty wag, ten sam RNG inicjalizacji, te same wyjscia, config ckpt bez nowych kluczy; K = n_head = jak 0). |
| K > 0: w blokach A qkv = jeden Linear d -> d + 2*K*hd (bias wg --attn-bias, Muon jak dotad, bo 2D); q ma n_head glow, |
| k i v po K glow. QK-norm i RoPE na k z K glowami, value residual (v0, vr_lambda) na v z K glowami, potem k, v powielone |
| repeat_interleave na n_head glow (grupa = n_head / K kolejnych glow q) przed SDPA; XSA na v po powieleniu (per glowa q). |
| Bloki G bez zmian. W config ckpt "kv_heads" tylko przy K > 0; --init-from sprawdza kv_heads (brak klucza = 0), nauczyciel |
| KD czyta go z configu. STOP: n_head % K != 0, K > n_head, K < 0, --kv-heads z --imha-t. |
| --data-seed S (karta dowodów v0.1 pkt 2, Mierniczy 20:56Z): ziarno kolejności danych (generator prefetchera i samplery |
| epoch / queue) osobno od --seed (inicjalizacja modelu); domyślnie brak = --seed, czyli r22 bit w bit. Ramiona R × 3 |
| różnią się --seed przy stałym --data-seed. W config ckpt "data_seed" tylko, gdy podane; --resume z innym = STOP. |
| """ |
| import argparse |
| import hashlib |
| import json |
| import math |
| import os |
| import re |
| import time |
| import queue |
| import threading |
| from contextlib import nullcontext |
|
|
| import numpy as np |
| import torch |
| import torch.distributed as dist |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch.utils.checkpoint import checkpoint |
| from torch.nn.parallel import DistributedDataParallel as DDP |
|
|
|
|
| def get_args(): |
| p = argparse.ArgumentParser() |
| p.add_argument("--data-dir", default=".") |
| p.add_argument("--out-dir", default=None) |
| p.add_argument("--run-id", default="gpt25m_run1") |
| p.add_argument("--steps", type=int, default=30000) |
| p.add_argument("--batch", type=int, default=32) |
| p.add_argument("--block", type=int, default=256) |
| p.add_argument("--n-layer", type=int, default=8) |
| p.add_argument("--n-embd", type=int, default=512) |
| p.add_argument("--n-head", type=int, default=8) |
| p.add_argument("--lr", type=float, default=6e-4) |
| p.add_argument("--min-lr", type=float, default=6e-5) |
| p.add_argument("--warmup", type=int, default=200) |
| p.add_argument("--schedule", choices=["cosine", "wsd"], default="cosine", |
| help="cosine (default) | wsd (warmup-stable-decay: constant peak then 1-sqrt decay, JugnuLM ARC-closer)") |
| p.add_argument("--decay-frac", type=float, default=0.2, |
| help="WSD: fraction of steps in the final decay phase (default 0.2)") |
| p.add_argument("--rewarm-steps", type=int, default=0, |
| help="przy --resume: liniowe przejscie LR od wartosci z ckpt do harmonogramu przez N krokow (0=off)") |
| p.add_argument("--wd", type=float, default=0.1) |
| p.add_argument("--grad-clip", type=float, default=1.0) |
| p.add_argument("--log-every", type=int, default=50) |
| p.add_argument("--eval-every", type=int, default=500) |
| p.add_argument("--eval-iters", type=int, default=50) |
| p.add_argument("--ckpt-every", type=int, default=1000) |
| p.add_argument("--seed", type=int, default=1337) |
| p.add_argument("--data-seed", type=int, default=None, |
| help="r23: ziarno kolejności danych (domyślnie = --seed); --seed zostaje ziarnem inicjalizacji modelu") |
| p.add_argument("--sampler", choices=["random", "epoch"], default="random", |
| help="random (r5: okna ze zwracaniem, ~75%% korpusu przy 1.4 ep.) | epoch (r6: kazde okno raz na epoke, " |
| "per-epoka permutacja + losowe przesuniecie siatki; resume dokladny bez fast-forwardu)") |
| p.add_argument("--resume", action="store_true") |
| p.add_argument("--init-from", default=None, |
| help="r10 (Wzrost): start z wag innego ckpt (np. z grow_depth.py), krok 0 i swiezy optymalizator; " |
| "ignorowane, gdy --resume wznawia wlasny ckpt.pt tego biegu") |
| p.add_argument("--mixer-layout", default=None, |
| help="r13: litera na blok (A = uwaga, G = GatedDeltaNet2), dlugosc = --n-layer; brak = same A (r12)") |
| p.add_argument("--gdn-head-dim", type=int, default=64, help="r13: head_dim GatedDeltaNet2") |
| p.add_argument("--attn-bias", type=int, choices=(0, 1), default=1, help="r13: bias w qkv/proj uwagi (1 = r12)") |
| p.add_argument("--rope-style", choices=("interleaved", "half"), default="interleaved", |
| help="r13: interleaved (r12) | half (rotate-half)") |
| p.add_argument("--data-offset-steps", type=int, default=0, |
| help="r10 (Wzrost, tylko --sampler epoch): dane od kroku N strumienia, np. dalszy ciag po ckpt " |
| "mniejszego modelu (ten sam batch i seed), bez powtarzania widzianych okien") |
| p.add_argument("--data-dir-2", default=None, |
| help="r10 (Wzrost, tylko --sampler epoch): od kroku strumienia --switch-step okna z DATA_DIR_2/train.bin " |
| "(np. swieze dokumenty rozlaczne z base), wlasna permutacja epokowa liczona od --switch-step") |
| p.add_argument("--switch-step", type=int, default=0, help="krok strumienia przejscia na --data-dir-2") |
| p.add_argument("--batch-sha", default=None, |
| help="r10 (Wzrost): bez treningu; sha256 paczek (x+y) w podanych krokach biegu, np. 0,24413,24414, " |
| "z tego samego Prefetchera co trening (kontrola strumienia), potem koniec") |
| p.add_argument("--synthetic", action="store_true", help="smoke throughput na losowym uint8 (bez danych)") |
| p.add_argument("--vocab", type=int, default=256) |
| p.add_argument("--dtype", default="uint8", help="bin dtype: uint8 (byte) | uint16 (BPE)") |
| p.add_argument("--optimizer", choices=["adamw", "muon"], default="adamw", |
| help="adamw (default, backward-compat) | muon (Newton-Schulz ortho dla 2D-weights + AdamW dla reszty)") |
| p.add_argument("--muon-lr", type=float, default=0.02, |
| help="peak LR dla Muon (macierzowe params); AdamW-aux uzywa --lr. Muon skalowany ta sama cosine-schedule co AdamW przez lr_mult=muon_lr/lr") |
| p.add_argument("--norm", choices=["layernorm", "rmsnorm"], default="layernorm", |
| help="layernorm (default, backward-compat) | rmsnorm (Qwen3-style, fp32-compute)") |
| p.add_argument("--norm-eps", type=float, default=1e-6) |
| p.add_argument("--pos", choices=["learned", "rope"], default="learned", |
| help="learned (default) | rope (parameter-free RoPE na q,k; usuwa learned pos-embedding)") |
| p.add_argument("--rope-theta", type=float, default=100000.0, |
| help="RoPE theta; top-3-board (JugnuLM/GPT-X2) uzywaja 100000 (nie-default-10K)") |
| p.add_argument("--ffn", choices=["gelu", "swiglu", "relu2"], default="gelu", |
| help="gelu (default 4x MLP) | swiglu (Qwen3 gated-MLP, hidden=--ffn-mult*d) | relu2 (r18: down(relu(up)^2), bez biasu)") |
| p.add_argument("--ffn-mult", type=float, default=2.667, |
| help="mnoznik hidden dla swiglu (~param-parity z 4x-gelu przy 8/3)") |
| p.add_argument("--value-residual", action="store_true", |
| help="ResFormer value-residuals: v_l += lambda_l*v0 (lambda init 0), ARC-targeted") |
| p.add_argument("--qk-norm", action="store_true", |
| help="Qwen3 QK-Norm: RMSNorm per-head na Q,K przed-attention (stabilnosc z Muon/high-LR)") |
| p.add_argument("--events-jsonl", default=None, |
| help="jesli podane: emituj events.jsonl (fabryka-track sidecar-format: update/evaluation/checkpoint/end)") |
| p.add_argument("--compile", action="store_true", |
| help="torch.compile model (2-3x throughput; state_dict zapisywany bez _orig_mod prefix via raw_model)") |
| p.add_argument("--logit-cap", type=float, default=0.0, |
| help="Gemma-style logit soft-cap: logits=cap*tanh(logits/cap); 0=off (prob. 30)") |
| p.add_argument("--z-loss", type=float, default=0.0, |
| help="z-loss coef: +coef*mean(logsumexp(logits)^2), stabilizuje logity; 0=off (prob. 1e-4)") |
| p.add_argument("--act-ckpt", action="store_true", |
| help="r11: przeliczanie aktywacji po bloku w treningu (pamiec ~ wejscia blokow zamiast wszystkich aktywacji)") |
| p.add_argument("--loss-chunk", type=int, default=0, |
| help="r11: strata kawalkami po N pozycji, bez pelnych logitow [B,T,V]; 0 = off (r10)") |
| p.add_argument("--loop", type=int, default=1, |
| help="r12: ile razy przejsc przez n_layer wspoldzielonych blokow (1 = bez wspoldzielenia, jak r11)") |
| p.add_argument("--queue", default=None, |
| help="r7: manifest kolejki shardow (JSON: steps_per_shard, shards{id:{lo,hi,sha256,source}}, rotations[[id..]..]). " |
| "Slot j = kroki [j*S,(j+1)*S) na shardzie schedule[j]; schedule = rotacje po kolei. Bez flagi = r6 bit w bit.") |
| p.add_argument("--queue-one-slot", action="store_true", |
| help="r7: trenuj tylko do konca biezacego slotu i zapisz ckpt (krok nadzorcy kolejki decyduje o nastepnym)") |
| p.add_argument("--kd-teacher", default=None, help="r14: ckpt nauczyciela KD (r12/r13); bez = r13 bit w bit") |
| p.add_argument("--kd-teacher-sha", default=None, help="r14: pelne sha256 ckpt nauczyciela (wymagane z --kd-teacher)") |
| p.add_argument("--kd-alpha", type=float, default=0.0, help="r14: waga KL w (1-A)*CE + A*T^2*KL; 0 < A <= 1 z nauczycielem") |
| p.add_argument("--kd-temp", type=float, default=1.0, help="r14: temperatura KD") |
| p.add_argument("--emb-dim", type=int, default=0, help="r15: faktoryzacja osadzen V x E + E x d (0 = wylaczone, r14)") |
| p.add_argument("--grad-accum", type=int, default=1, help="r16: mikrobatche na krok (1 = wylaczone, r15)") |
| p.add_argument("--muon-scale", choices=["sqrt_ratio", "match_rms_adamw"], default="sqrt_ratio", |
| help="r17: skala kroku Muon (sqrt_ratio = jak dotad; match_rms_adamw = torch.optim.Muon)") |
| p.add_argument("--muon-wd", type=float, default=0.0, help="r17: decoupled WD na parametrach Muon (0 = jak dotad)") |
| p.add_argument("--tf32", type=int, choices=[0, 1], default=1, help="r17: 0 = TF32 wylaczone (jak Qielbas M2)") |
| p.add_argument("--xsa", type=int, choices=[0, 1], default=0, help="r18: 1 = z wyjscia glowy usuwana skladowa wzdluz wlasnego v") |
| p.add_argument("--rope-dims", type=int, default=0, help="r18: RoPE na pierwszych N wymiarach glowy (0 = wszystkie, r17)") |
| p.add_argument("--zero-init", type=int, choices=[0, 1], default=0, help="r18: 1 = zerowe wagi proj uwagi i wyjscia MLP") |
| p.add_argument("--arch", choices=["gpt", "tinylm"], default="gpt", help="r19: tinylm = model z pliku Tiny Chick (--tinylm-file)") |
| p.add_argument("--tinylm-file", default=None, help="r19: tinylm_arch.py z kurcontko/Tiny-Chick-10M@60d32c4d (sha 4876824d...)") |
| p.add_argument("--kv-heads", type=int, default=0, |
| help="r19 tinylm / r23 gpt: liczba glow KV (GQA, dzielnik --n-head); 0 = --n-head (MHA)") |
| p.add_argument("--train-only", default=None, |
| help="r21: regex nazw uczonych parametrow; reszta requires_grad False i poza optymalizatorem; tylko z --init-from") |
| p.add_argument("--imha-t", type=int, default=0, |
| help="r22: IMHA z T iteracjami w blokach uwagi (0 = zwykla uwaga, r21 bit w bit); artykul: T = 2") |
| return p.parse_args() |
|
|
|
|
| |
| |
| |
| |
| |
| |
| def zeropower_via_newtonschulz5(G, steps=5): |
| """Ortogonalizacja macierzy G przez quintic Newton-Schulz (bf16). Zwraca |
| macierz ~ U V^T z SVD(G)=U S V^T. Wspolczynniki (a,b,c) z impl. Kellera.""" |
| assert G.ndim == 2 |
| a, b, c = (3.4445, -4.7750, 2.0315) |
| X = G.bfloat16() |
| transposed = G.size(0) > G.size(1) |
| if transposed: |
| X = X.T |
| X = X / (X.norm() + 1e-7) |
| for _ in range(steps): |
| A = X @ X.T |
| B = b * A + c * (A @ A) |
| X = a * X + B @ X |
| if transposed: |
| X = X.T |
| return X |
|
|
|
|
| class Muon(torch.optim.Optimizer): |
| """Momentum-SGD z ortogonalizowanym update. weight_decay domyslnie 0 (Muon-params |
| czysto; WD trzymamy na AdamW-aux). lr_mult pozwala petli lr-schedule skalowac |
| Muon proporcjonalnie do AdamW.""" |
| def __init__(self, params, lr=0.02, lr_mult=1.0, momentum=0.95, nesterov=True, |
| ns_steps=5, weight_decay=0.0, scale_mode="sqrt_ratio"): |
| defaults = dict(lr=lr, lr_mult=lr_mult, momentum=momentum, nesterov=nesterov, |
| ns_steps=ns_steps, weight_decay=weight_decay, scale_mode=scale_mode) |
| super().__init__(params, defaults) |
|
|
| @torch.no_grad() |
| def step(self, closure=None): |
| loss = None |
| if closure is not None: |
| with torch.enable_grad(): |
| loss = closure() |
| for group in self.param_groups: |
| lr = group["lr"]; momentum = group["momentum"]; wd = group["weight_decay"] |
| for p in group["params"]: |
| g = p.grad |
| if g is None: |
| continue |
| if g.ndim > 2: |
| g = g.reshape(g.size(0), -1) |
| state = self.state[p] |
| if "momentum_buffer" not in state: |
| state["momentum_buffer"] = torch.zeros_like(g) |
| buf = state["momentum_buffer"] |
| buf.mul_(momentum).add_(g) |
| g = g.add(buf, alpha=momentum) if group["nesterov"] else buf |
| u = zeropower_via_newtonschulz5(g, steps=group["ns_steps"]) |
| if wd != 0: |
| p.mul_(1 - lr * wd) |
| |
| |
| if group.get("scale_mode", "sqrt_ratio") == "match_rms_adamw": |
| scale = 0.2 * max(g.size(0), g.size(1)) ** 0.5 |
| else: |
| scale = max(1.0, p.size(0) / p.size(1)) ** 0.5 |
| p.add_(u.reshape(p.shape).to(p.dtype), alpha=-lr * scale) |
| return loss |
|
|
|
|
| class MuonWithAuxAdam: |
| """Kontener: Muon dla macierzowych ukrytych wag + AdamW dla reszty. Wystawia |
| param_groups/step/zero_grad/state_dict tak, by petla treningowa dzialala bez zmian.""" |
| def __init__(self, muon, adamw): |
| self.muon = muon |
| self.adamw = adamw |
|
|
| @property |
| def param_groups(self): |
| return self.muon.param_groups + self.adamw.param_groups |
|
|
| @property |
| def state(self): |
| return {**self.muon.state, **self.adamw.state} |
|
|
| def step(self, closure=None): |
| self.muon.step() |
| self.adamw.step() |
|
|
| def zero_grad(self, set_to_none=True): |
| self.muon.zero_grad(set_to_none=set_to_none) |
| self.adamw.zero_grad(set_to_none=set_to_none) |
|
|
| def state_dict(self): |
| return {"muon": self.muon.state_dict(), "adamw": self.adamw.state_dict()} |
|
|
| def load_state_dict(self, sd): |
| self.muon.load_state_dict(sd["muon"]) |
| self.adamw.load_state_dict(sd["adamw"]) |
|
|
|
|
| def build_optimizer(a, model): |
| """--optimizer adamw -> DOKLADNIE poprzedni AdamW (backward-compat). |
| --optimizer muon -> Muon(2D-hidden) + AdamW(embeddingi/head/norm/bias).""" |
| if a.optimizer == "adamw": |
| return torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=a.lr, weight_decay=a.wd, |
| betas=(0.9, 0.95)) |
| muon_params, adamw_params, seen = [], [], set() |
| for name, p in model.named_parameters(): |
| if not p.requires_grad or id(p) in seen: |
| continue |
| seen.add(id(p)) |
| is_embed_or_head = name.startswith(("tok.", "pos.", "head.", "tl.emb.")) |
| tl_gain = name.startswith("tl.") and name.endswith(".attn.gain") |
| |
| |
| gdn_aux = ".gdn." in name and not name.endswith(("q_proj.weight", "k_proj.weight", "v_proj.weight", "o_proj.weight")) |
| if p.ndim >= 2 and not is_embed_or_head and not gdn_aux and not tl_gain: |
| muon_params.append(p) |
| else: |
| adamw_params.append(p) |
| lr_mult = a.muon_lr / a.lr if a.lr > 0 else 1.0 |
| muon = Muon(muon_params, lr=a.muon_lr, lr_mult=lr_mult, weight_decay=getattr(a, "muon_wd", 0.0), |
| scale_mode=getattr(a, "muon_scale", "sqrt_ratio")) |
| adamw = torch.optim.AdamW(adamw_params, lr=a.lr, weight_decay=a.wd, betas=(0.9, 0.95)) |
| return MuonWithAuxAdam(muon, adamw) |
|
|
|
|
| class RMSNorm(nn.Module): |
| """Qwen3-style RMSNorm (fp32-compute dla stabilnosci). 1D weight -> AdamW w split-Muon.""" |
| def __init__(self, d, eps=1e-6): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(d)) |
| self.eps = eps |
|
|
| def forward(self, x): |
| return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps).to(x.dtype) * self.weight |
|
|
|
|
| def make_norm(d, cfg): |
| return RMSNorm(d, cfg.norm_eps) if cfg.norm == "rmsnorm" else nn.LayerNorm(d) |
|
|
|
|
| def apply_rope(x, base=100000.0, style="interleaved", nrot=0): |
| """Parameter-free RoPE na [B,H,T,D] (interleaved-conv, port z qwen_model.py). Train==eval |
| MUSZA uzywac tej samej konwencji (self-contained eval -> spojne). |
| r13: style="half" = rotate-half (pary i, i+D/2), jak Qielbas M2; "interleaved" = r12 bez zmian. |
| r18: nrot (0 < nrot < D) = obrot tylko pierwszych nrot wymiarow (czestotliwosci z nrot), reszta bez zmian; 0 = calosc.""" |
| if 0 < nrot < x.size(-1): |
| return torch.cat([apply_rope(x[..., :nrot], base, style), x[..., nrot:]], dim=-1) |
| _, _, T, dim = x.shape |
| pos = torch.arange(T, device=x.device, dtype=torch.float32) |
| freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=x.device, dtype=torch.float32) / dim)) |
| ang = torch.outer(pos, freq) |
| if style == "half": |
| ang = torch.cat([ang, ang], dim=-1) |
| cos, sin = ang.cos().to(x.dtype)[None, None], ang.sin().to(x.dtype)[None, None] |
| a, b = x.chunk(2, dim=-1) |
| return x * cos + torch.cat([-b, a], dim=-1) * sin |
| cos, sin = ang.cos().to(x.dtype)[None, None], ang.sin().to(x.dtype)[None, None] |
| even, odd = x[..., ::2], x[..., 1::2] |
| return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2) |
|
|
|
|
| class SwiGLU(nn.Module): |
| """Qwen3 gated-MLP: down(silu(gate(x))*up(x)). 3x 2D bez-bias -> wszystkie do Muon.""" |
| def __init__(self, d, hidden): |
| super().__init__() |
| self.gate = nn.Linear(d, hidden, bias=False) |
| self.up = nn.Linear(d, hidden, bias=False) |
| self.down = nn.Linear(hidden, d, bias=False) |
|
|
| def forward(self, x): |
| return self.down(F.silu(self.gate(x)) * self.up(x)) |
|
|
|
|
| class ReLU2(nn.Module): |
| """r18 (Tiny Chick): down(relu(up(x))^2), 2x 2D bez biasu -> Muon.""" |
| def __init__(self, d, hidden): |
| super().__init__() |
| self.up = nn.Linear(d, hidden, bias=False) |
| self.down = nn.Linear(hidden, d, bias=False) |
|
|
| def forward(self, x): |
| return self.down(F.relu(self.up(x)).square()) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, d, nh, block, cfg, is_first=False, mixer="A", layer_idx=0): |
| super().__init__() |
| self.ln1 = make_norm(d, cfg) |
| self.ln2 = make_norm(d, cfg) |
| self.mixer = mixer |
| if mixer == "G": |
| from fla.layers.gdn2 import GatedDeltaNet2 |
| hd = getattr(cfg, "gdn_head_dim", 64) |
| if d % hd: |
| raise SystemExit(f"STOP: n_embd {d} nie dzieli sie przez --gdn-head-dim {hd}") |
| self.gdn = GatedDeltaNet2(hidden_size=d, head_dim=hd, num_heads=d // hd, expand_v=1.0, |
| use_short_conv=True, mode="chunk", layer_idx=layer_idx) |
| for sm in self.gdn.modules(): |
| sm._r13_fla_init = True |
| else: |
| bias = bool(getattr(cfg, "attn_bias", 1)) |
| kvh = int(getattr(cfg, "kv_heads", 0) or 0) |
| self.nkv = kvh if 0 < kvh < nh else nh |
| self.kvd = self.nkv * (d // nh) |
| self.qkv = nn.Linear(d, d + 2 * self.kvd, bias=bias) |
| self.proj = nn.Linear(d, d, bias=bias) |
| if cfg.ffn == "swiglu": |
| self.mlp = SwiGLU(d, int(round(cfg.ffn_mult * d))) |
| elif cfg.ffn == "relu2": |
| self.mlp = ReLU2(d, int(round(cfg.ffn_mult * d))) |
| else: |
| self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d)) |
| self.nh = nh |
| self.d = d |
| self.cfg = cfg |
| self.is_first = is_first |
| if mixer == "A" and cfg.value_residual and not is_first: |
| self.vr_lambda = nn.Parameter(torch.zeros(1)) |
| if mixer == "A" and cfg.qk_norm: |
| hd = d // nh |
| self.q_norm = RMSNorm(hd, cfg.norm_eps) |
| self.k_norm = RMSNorm(hd, cfg.norm_eps) |
| self.imha_t = int(getattr(cfg, "imha_t", 0) or 0) if mixer == "A" else 0 |
| if self.imha_t and self.nkv != nh: |
| raise SystemExit("STOP --imha-t: niezgodne z --kv-heads (GQA)") |
| if self.imha_t: |
| self.imha_a = nn.Linear(2 * d, d, bias=False) |
| self.imha_f = nn.Linear(2 * d, d, bias=False) |
| self.imha_log_tau = nn.Parameter(torch.zeros(1)) |
|
|
| def _qkv(self, h, v0, set_v0): |
| B, T, D = h.size() |
| q, k, v = self.qkv(h).split(self.d, dim=2) |
| hd = D // self.nh |
| q = q.view(B, T, self.nh, hd).transpose(1, 2) |
| k = k.view(B, T, self.nh, hd).transpose(1, 2) |
| v = v.view(B, T, self.nh, hd).transpose(1, 2) |
| if self.cfg.qk_norm: |
| q = self.q_norm(q) |
| k = self.k_norm(k) |
| if self.cfg.pos == "rope": |
| style = getattr(self.cfg, "rope_style", "interleaved") |
| nrot = int(getattr(self.cfg, "rope_dims", 0) or 0) |
| q = apply_rope(q, self.cfg.rope_theta, style, nrot) |
| k = apply_rope(k, self.cfg.rope_theta, style, nrot) |
| if self.cfg.value_residual: |
| if self.is_first: |
| if set_v0: |
| v0 = v |
| else: |
| v = v + self.vr_lambda * v0 |
| return q, k, v, v0 |
|
|
| def _imha(self, x, v0): |
| """r22: IMHA (Algorytm 1); zwraca (h_T - h_0, v0). Wszystkie sumy po k <= i (prefiks), wiec przyczynowo.""" |
| B, T, D = x.size() |
| hd = D // self.nh |
| h0 = self.ln1(x) |
| h = h0 |
| pos = torch.arange(T, device=x.device) |
| cnt = (pos + 1).to(torch.float32) |
| logn = torch.log(cnt).clamp_min(1e-12) |
| mask = torch.ones(T, T, device=x.device, dtype=torch.bool).tril() |
| plain = bool(getattr(self.cfg, "imha_plain", False)) |
| tau = self.imha_log_tau.exp() |
| for t in range(self.imha_t): |
| q, k, v, v0 = self._qkv(h, v0, set_v0=(t == 0)) |
| s = (q.float() @ k.float().transpose(-2, -1)) * (hd ** -0.5) |
| A = torch.softmax(s.masked_fill(~mask, float("-inf")), dim=-1) |
| z = (A.to(v.dtype) @ v) |
| if not plain: |
| ent = -(A * torch.log(A.clamp_min(1e-30))).sum(-1) |
| hn = torch.where(pos == 0, torch.zeros_like(ent), ent / logn) |
| e = torch.exp(-hn / tau) |
| w = e / (e.cumsum(-1) / cnt) |
| z = z * w.unsqueeze(-1).to(z.dtype) |
| z = self.proj(z.transpose(1, 2).contiguous().view(B, T, D)) |
| if plain: |
| h = h + z |
| continue |
| g = z.cumsum(1) / cnt.view(1, T, 1).to(z.dtype) |
| u = torch.cat([z, g], dim=-1) |
| h = h + torch.sigmoid(self.imha_a(u)) * F.relu(self.imha_f(u)) |
| return h - h0, v0 |
|
|
| def forward(self, x, v0=None): |
| B, T, D = x.size() |
| if self.imha_t: |
| dx, v0 = self._imha(x, v0) |
| x = x + dx |
| x = x + self.mlp(self.ln2(x)) |
| return x, v0 |
| h = self.ln1(x) |
| if self.mixer == "G": |
| x = x + self.gdn(h, use_cache=False)[0] |
| x = x + self.mlp(self.ln2(x)) |
| return x, v0 |
| q, k, v = self.qkv(h).split([self.d, self.kvd, self.kvd], dim=2) |
| hd = D // self.nh |
| q = q.view(B, T, self.nh, hd).transpose(1, 2) |
| k = k.view(B, T, self.nkv, hd).transpose(1, 2) |
| v = v.view(B, T, self.nkv, hd).transpose(1, 2) |
| if self.cfg.qk_norm: |
| q = self.q_norm(q) |
| k = self.k_norm(k) |
| if self.cfg.pos == "rope": |
| style = getattr(self.cfg, "rope_style", "interleaved") |
| nrot = int(getattr(self.cfg, "rope_dims", 0) or 0) |
| q = apply_rope(q, self.cfg.rope_theta, style, nrot) |
| k = apply_rope(k, self.cfg.rope_theta, style, nrot) |
| if self.cfg.value_residual: |
| if self.is_first: |
| v0 = v |
| else: |
| v = v + self.vr_lambda * v0 |
| if self.nkv != self.nh: |
| k = k.repeat_interleave(self.nh // self.nkv, dim=1) |
| v = v.repeat_interleave(self.nh // self.nkv, dim=1) |
| y = F.scaled_dot_product_attention(q, k, v, is_causal=True) |
| if getattr(self.cfg, "xsa", 0): |
| vn = F.normalize(v, dim=-1) |
| y = y - (y * vn).sum(-1, keepdim=True) * vn |
| y = y.transpose(1, 2).contiguous().view(B, T, D) |
| x = x + self.proj(y) |
| x = x + self.mlp(self.ln2(x)) |
| return x, v0 |
|
|
|
|
| TINYLM_SHA = "4876824de8c697a657e5f58351c216ffb84cb90d6a2a53ecd94933f7ddcc4cfa" |
|
|
|
|
| def load_tinylm_module(path): |
| """r19: oryginalny tinylm_arch.py (Apache-2.0, kurcontko/Tiny-Chick-10M@60d32c4d) z pliku przypietego sha. |
| Bez --tinylm-file: tinylm_arch.py obok tego trenera (np. eval, ktory odtwarza model z configu ckpt).""" |
| import importlib.util |
| path = path or os.path.join(os.path.dirname(os.path.abspath(__file__)), "tinylm_arch.py") |
| if not os.path.exists(path): |
| raise SystemExit(f"STOP --arch tinylm: brak {path}") |
| with open(path, "rb") as f: |
| h = hashlib.sha256(f.read()).hexdigest() |
| if h != TINYLM_SHA: |
| raise SystemExit(f"STOP --arch tinylm: sha {path} = {h[:8]} != {TINYLM_SHA[:8]}") |
| spec = importlib.util.spec_from_file_location("tinylm_arch_tc", path) |
| mod = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(mod) |
| return mod |
|
|
|
|
| class GPT(nn.Module): |
| def __init__(self, vocab, n_layer, n_embd, n_head, block, cfg): |
| super().__init__() |
| self.cfg = cfg |
| if getattr(cfg, "arch", "gpt") == "tinylm": |
| tlm = load_tinylm_module(getattr(cfg, "tinylm_file", None)) |
| self.tl = tlm.TinyLM(tlm.TinyConfig( |
| vocab_size=vocab, d=n_embd, layers=n_layer, heads=n_head, |
| kv_heads=int(getattr(cfg, "kv_heads", 0) or 0) or n_head, mlp_mult=cfg.ffn_mult, act=cfg.ffn, |
| qk_norm=bool(cfg.qk_norm), qk_gain=1.0, rope_dims=int(getattr(cfg, "rope_dims", 0) or 0), |
| xsa=bool(getattr(cfg, "xsa", 0)), zero_init=bool(getattr(cfg, "zero_init", 0)), |
| max_position_embeddings=block, rope_theta=cfg.rope_theta, eps=cfg.norm_eps)) |
| self.block = block |
| return |
| E = int(getattr(cfg, "emb_dim", 0) or 0) |
| self.emb_dim = E |
| self.tok = nn.Embedding(vocab, E if E else n_embd) |
| self.use_rope = cfg.pos == "rope" |
| if not self.use_rope: |
| self.pos = nn.Embedding(block, n_embd) |
| layout = getattr(cfg, "mixer_layout", None) or "A" * n_layer |
| if len(layout) != n_layer or set(layout) - {"A", "G"}: |
| raise SystemExit(f"STOP: --mixer-layout {layout!r}: dlugosc {n_layer}, tylko litery A/G") |
| nrot = int(getattr(cfg, "rope_dims", 0) or 0) |
| kvh = int(getattr(cfg, "kv_heads", 0) or 0) |
| if kvh < 0 or kvh > n_head or (kvh and n_head % kvh): |
| raise SystemExit(f"STOP: --kv-heads {kvh}: 0 (MHA) albo dzielnik n_head {n_head}") |
| if nrot and (nrot % 2 or not 0 < nrot < n_embd // n_head or cfg.pos != "rope"): |
| raise SystemExit(f"STOP: --rope-dims {nrot}: parzyste, 0 < N < head_dim {n_embd // n_head}, tylko z --pos rope") |
| first_a = layout.find("A") |
| self.blocks = nn.ModuleList([Block(n_embd, n_head, block, cfg, is_first=(i == first_a), mixer=layout[i], layer_idx=i) |
| for i in range(n_layer)]) |
| self.lnf = make_norm(n_embd, cfg) |
| self.head = nn.Linear(E if E else n_embd, vocab, bias=False) |
| self.head.weight = self.tok.weight |
| if E: |
| self.emb_up = nn.Linear(E, n_embd, bias=False) |
| self.emb_down = nn.Linear(n_embd, E, bias=False) |
| self.block = block |
| self.apply(self._init) |
| if getattr(cfg, "zero_init", 0): |
| for b in self.blocks: |
| if b.mixer == "A": |
| nn.init.zeros_(b.proj.weight) |
| nn.init.zeros_(b.mlp.down.weight if hasattr(b.mlp, "down") else b.mlp[2].weight) |
| if E: |
| nn.init.normal_(self.emb_up.weight, 0.0, E ** -0.5) |
| nn.init.normal_(self.emb_down.weight, 0.0, n_embd ** -0.5) |
|
|
| def _head(self, h): |
| return self.head(self.emb_down(h)) if self.emb_dim else self.head(h) |
|
|
| def _init(self, m): |
| if getattr(m, "_r13_fla_init", False): |
| return |
| if isinstance(m, nn.Linear): |
| nn.init.normal_(m.weight, 0.0, 0.02) |
| if m.bias is not None: |
| nn.init.zeros_(m.bias) |
| elif isinstance(m, nn.Embedding): |
| nn.init.normal_(m.weight, 0.0, 0.02) |
|
|
| def _loss_sum(self, h, y, t=None): |
| """r11: suma CE (+ z-loss) dla kawalka pozycji; h po lnf [B,t,D], y [B,t]. Ta sama arytmetyka co sciezka r10. |
| r14: t = logity nauczyciela tego kawalka [B,t,V] -> (suma straty mieszanej KD, suma samej CE).""" |
| logits = self._head(h) |
| cap = getattr(self.cfg, "logit_cap", 0.0) |
| if cap and cap > 0: |
| logits = cap * torch.tanh(logits / cap) |
| flat = logits.view(-1, logits.size(-1)) |
| s = F.cross_entropy(flat, y.reshape(-1), reduction="sum") |
| zc = getattr(self.cfg, "z_loss", 0.0) |
| if zc and zc > 0: |
| lse = torch.logsumexp(flat, dim=-1) |
| s = s + zc * (lse * lse).sum() |
| if t is None: |
| return s |
| return self._kd_mix(s, flat, t, 1), s |
|
|
| def _kd_mix(self, ce, flat, t, div): |
| """r14: (1-A)*ce + A*T^2*KL(p_naucz || p_ucz); KL sumowane po pozycjach kawalka / div (1 = suma, N = srednia). |
| fp32, pelny vocab; flat = logity ucznia [N,V] (po logit_cap), t = logity nauczyciela [...,V].""" |
| A, T = self.cfg.kd_alpha, self.cfg.kd_temp |
| ls = F.log_softmax(flat.float() / T, dim=-1) |
| lt = F.log_softmax(t.reshape(-1, t.size(-1)).float() / T, dim=-1) |
| kl = (lt.exp() * (lt - ls)).sum() / div |
| return (1.0 - A) * ce + (A * T * T) * kl |
|
|
| def forward(self, idx, targets=None, teacher_logits=None): |
| if getattr(self.cfg, "arch", "gpt") == "tinylm": |
| logits = self.tl(idx).logits |
| if targets is None: |
| return logits, None |
| return logits, F.cross_entropy(logits.reshape(-1, logits.size(-1)).float(), targets.reshape(-1)) |
| B, T = idx.size() |
| x = self.tok(idx) |
| if self.emb_dim: |
| x = self.emb_up(x) |
| if not self.use_rope: |
| pos = torch.arange(T, device=idx.device) |
| x = x + self.pos(pos)[None] |
| v0 = None |
| act_ckpt = self.training and getattr(self.cfg, "act_ckpt", False) |
| for _ in range(getattr(self.cfg, "loop", 1)): |
| for b in self.blocks: |
| x, v0 = checkpoint(b, x, v0, use_reentrant=False) if act_ckpt else b(x, v0) |
| h = self.lnf(x) |
| n = getattr(self.cfg, "loss_chunk", 0) |
| if targets is not None and n and n < T: |
| if teacher_logits is None: |
| s = sum(checkpoint(self._loss_sum, h[:, i:i + n], targets[:, i:i + n], use_reentrant=False) |
| for i in range(0, T, n)) |
| return None, s / (B * T) |
| parts = [checkpoint(self._loss_sum, h[:, i:i + n], targets[:, i:i + n], teacher_logits[:, i:i + n], |
| use_reentrant=False) for i in range(0, T, n)] |
| return None, sum(p[0] for p in parts) / (B * T), sum(p[1] for p in parts).detach() / (B * T) |
| logits = self._head(h) |
| cap = getattr(self.cfg, "logit_cap", 0.0) |
| if cap and cap > 0: |
| logits = cap * torch.tanh(logits / cap) |
| loss = None |
| if targets is not None: |
| flat = logits.view(-1, logits.size(-1)) |
| loss = F.cross_entropy(flat, targets.view(-1)) |
| zc = getattr(self.cfg, "z_loss", 0.0) |
| if zc and zc > 0: |
| lse = torch.logsumexp(flat, dim=-1) |
| loss = loss + zc * (lse * lse).mean() |
| if teacher_logits is not None: |
| return logits, self._kd_mix(loss, flat, teacher_logits, flat.size(0)), loss.detach() |
| return logits, loss |
|
|
|
|
| def main(): |
| a = get_args() |
| world = int(os.environ.get("WORLD_SIZE", "1")) |
| rank = int(os.environ.get("RANK", "0")) |
| local_rank = int(os.environ.get("LOCAL_RANK", "0")) |
| ddp = world > 1 |
| if ddp: |
| if a.batch % world: |
| raise SystemExit(f"STOP: --batch {a.batch} (globalny) nie dzieli sie przez liczbe procesow {world}") |
| if torch.cuda.is_available(): |
| torch.cuda.set_device(local_rank) |
| dist.init_process_group("nccl" if torch.cuda.is_available() else "gloo") |
| if a.grad_accum < 1 or (a.batch // world) % a.grad_accum: |
| raise SystemExit(f"STOP: --grad-accum {a.grad_accum}: wiersze ranku {a.batch // world} nie dziela sie na K mikrobatchy") |
| master = rank == 0 |
| rows = slice(rank * (a.batch // world), (rank + 1) * (a.batch // world)) |
| out_dir = a.out_dir or os.path.join(a.data_dir, "runs", a.run_id) |
| os.makedirs(out_dir, exist_ok=True) |
| log_path = os.path.join(out_dir, "train.log") |
| metrics_path = os.path.join(out_dir, "metrics.jsonl") |
| ckpt_path = os.path.join(out_dir, "ckpt.pt") |
|
|
| def log(msg): |
| if not master: |
| return |
| line = f"[{time.strftime('%H:%M:%S')}] {msg}" |
| print(line, flush=True) |
| with open(log_path, "a", encoding="utf-8") as f: |
| f.write(line + "\n") |
|
|
| events_path = a.events_jsonl |
|
|
| def emit(kind, step, metrics=None, **extra): |
| if not events_path or not master: |
| return |
| rec = {"kind": kind, "updates": int(step), "tokens": int(step) * a.block * a.batch} |
| if metrics: |
| rec["metrics"] = metrics |
| rec.update(extra) |
| with open(events_path, "a", encoding="utf-8") as f: |
| f.write(json.dumps(rec) + "\n") |
|
|
| torch.manual_seed(a.seed) |
| torch.backends.cuda.matmul.allow_tf32 = bool(a.tf32) |
| torch.backends.cudnn.allow_tf32 = bool(a.tf32) |
| device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu") |
| use_bf16 = device.type == "cuda" and torch.cuda.is_bf16_supported() |
| ptdtype = torch.bfloat16 if use_bf16 else torch.float32 |
| ctx = torch.amp.autocast(device_type=device.type, dtype=ptdtype) if device.type == "cuda" else nullcontext() |
| log(f"device={device} bf16={use_bf16} dev={torch.cuda.get_device_name(local_rank) if device.type=='cuda' else 'cpu'}" |
| f" world={world}") |
|
|
| if a.synthetic: |
| rng = np.random.default_rng(a.seed) |
| train_data = rng.integers(0, 256, size=8_000_000, dtype=np.uint8) |
| val_data = train_data[:200_000] |
| log("SYNTHETIC uint8 (smoke throughput, zero danych PII)") |
| else: |
| train_data = np.memmap(os.path.join(a.data_dir, "train.bin"), dtype=np.dtype(a.dtype), mode="r") |
| val_data = np.memmap(os.path.join(a.data_dir, "val.bin"), dtype=np.dtype(a.dtype), mode="r") |
| train2_data = (np.memmap(os.path.join(a.data_dir_2, "train.bin"), dtype=np.dtype(a.dtype), mode="r") |
| if a.data_dir_2 else None) |
| log(f"dane: train={len(train_data):,}B block={a.block} batch={a.batch} tok/step={a.block*a.batch:,}") |
|
|
| class EpochSampler: |
| """Losowanie BEZ zwracania: epoka e = permutacja okien na siatce co `block` tokenow, przesunietej o losowe |
| offset_e w [0, block) (kolejne epoki widza inne granice okien). Krok s bierze okna s*batch..(s+1)*batch-1 |
| z ciaglego strumienia epok, wiec start od dowolnego kroku (resume) daje dokladnie ten sam ciag co run ciagly. |
| Epoka e ma wlasny RNG ziarna (seed, e): najpierw offset, potem permutacja; niezalezna od innych epok.""" |
| def __init__(self, n_tokens, block, batch, seed): |
| self.n, self.block, self.batch, self.seed = n_tokens, block, batch, seed |
| self._perm_e, self._perm = None, None |
|
|
| def _meta(self, e): |
| rng = np.random.default_rng([self.seed, e]) |
| off = int(rng.integers(self.block)) |
| return rng, off, (self.n - 1 - off) // self.block |
|
|
| def _perm_of(self, e): |
| if self._perm_e != e: |
| rng, _, nwin = self._meta(e) |
| self._perm_e, self._perm = e, rng.permutation(nwin) |
| return self._perm |
|
|
| def starts(self, step): |
| g, need, out, e = step * self.batch, self.batch, [], 0 |
| while g >= self._meta(e)[2]: |
| g -= self._meta(e)[2]; e += 1 |
| while need: |
| _, off, nwin = self._meta(e) |
| take = min(need, nwin - g) |
| out.append(off + self._perm_of(e)[g:g + take].astype(np.int64) * self.block) |
| need -= take; g = 0; e += 1 |
| return np.concatenate(out) |
|
|
| class QueueSampler: |
| """r7: pula jako kolejka shardow (zakresy tokenow [lo,hi) w train.bin, wyrownane do dokumentow). |
| schedule = rotacje manifestu po kolei; slot j = kroki [j*S,(j+1)*S) na shardzie schedule[j]. W slocie kazde okno |
| uzyte najwyzej raz: siatka co `block` od lo+off, off i permutacja z RNG (seed, j, id shardu). Starty okien zaleza |
| tylko od (manifest, seed, krok), wiec resume z ckpt na granicy slotu jest dokladny, a usuniecie shardu ze slotu j |
| nie zmienia krokow < j*S. prefix_sha(krok) = hash slotow juz uzytych; zapisywany w ckpt i sprawdzany przy resume.""" |
| def __init__(self, manifest, n_tokens, block, batch, seed): |
| self.S, self.block, self.batch, self.seed = int(manifest["steps_per_shard"]), block, batch, seed |
| self.shards = manifest["shards"] |
| self.schedule = [sid for rot in manifest["rotations"] for sid in rot] |
| for sid in self.schedule: |
| sh = self.shards[sid] |
| if not 0 <= sh["lo"] < sh["hi"] <= n_tokens: |
| raise SystemExit(f"STOP kolejka: shard {sid} [{sh['lo']},{sh['hi']}) poza train.bin ({n_tokens})") |
| if (sh["hi"] - sh["lo"] - block) // block < self.S * batch: |
| raise SystemExit(f"STOP kolejka: shard {sid} ma za malo okien na {self.S} krokow x {batch}") |
| used = sorted((self.shards[s]["lo"], self.shards[s]["hi"], s) for s in dict.fromkeys(self.schedule)) |
| for (_, hi0, s0), (lo1, _, s1) in zip(used, used[1:]): |
| if lo1 < hi0: |
| raise SystemExit(f"STOP kolejka: shardy {s0} i {s1} nakladaja sie w train.bin (dane uzyte dwa razy)") |
| self._cache = (None, None) |
|
|
| def _slot(self, j): |
| if self._cache[0] != j: |
| sid = self.schedule[j] |
| sh = self.shards[sid] |
| idn = int(hashlib.sha256(sid.encode()).hexdigest()[:8], 16) |
| rng = np.random.default_rng([self.seed, j, idn]) |
| off = int(rng.integers(self.block)) |
| nwin = (sh["hi"] - sh["lo"] - 1 - off) // self.block |
| perm = rng.permutation(nwin)[: self.S * self.batch].astype(np.int64) |
| self._cache = (j, (sh["lo"] + off, perm)) |
| return self._cache[1] |
|
|
| def starts(self, step): |
| j, k = divmod(step, self.S) |
| if j >= len(self.schedule): |
| raise SystemExit(f"STOP kolejka: krok {step} poza harmonogramem ({len(self.schedule)} slotow x {self.S})") |
| base, perm = self._slot(j) |
| return base + perm[k * self.batch:(k + 1) * self.batch] * self.block |
|
|
| def prefix_sha(self, step): |
| used = -(-step // self.S) |
| spec = [self.S, self.seed, self.block, self.batch, |
| [[sid, self.shards[sid]["lo"], self.shards[sid]["hi"], self.shards[sid]["sha256"]] for sid in self.schedule[:used]]] |
| return hashlib.sha256(json.dumps(spec).encode()).hexdigest() |
|
|
| def _make_batch_cpu(split, generator=None, ix=None, shard=False): |
| """Wektoryzowane budowanie batcha: JEDEN numpy fancy-index zamiast |
| python-loop per-item. sliding_window_view daje strided-view (N-block, block+1) |
| BEZ kopiowania; windows[ix] materializuje tylko wybrane wiersze naraz. |
| Zwraca (x,y) long CPU (pinned jesli cuda). Rozklad batchy IDENTYCZNY jak |
| stary torch.stack-loop: x=data[i:i+block], y=data[i+1:i+1+block]. ix: gotowe starty okien (sampler epoch). |
| shard=True (trening pod DDP): starty losowane dla CALEGO globalnego batcha jak w 1-GPU, rank bierze swoje wiersze.""" |
| data = train_data if split == "train" else (train2_data if split == "train2" else val_data) |
| if ix is None: |
| ix = torch.randint(len(data) - a.block - 1, (a.batch,), generator=generator) |
| if shard: |
| ix = ix[rows] |
| |
| windows = np.lib.stride_tricks.sliding_window_view(data, a.block + 1) |
| sel = windows[ix.numpy()] |
| x = torch.from_numpy(sel[:, :-1].astype(np.int64)) |
| y = torch.from_numpy(sel[:, 1:].astype(np.int64)) |
| if device.type == "cuda": |
| x = x.pin_memory(); y = y.pin_memory() |
| return x, y |
|
|
| def _to_device(x, y): |
| if device.type == "cuda": |
| return x.to(device, non_blocking=True), y.to(device, non_blocking=True) |
| return x.to(device), y.to(device) |
|
|
| def get_batch(split, generator=None): |
| return _to_device(*_make_batch_cpu(split, generator)) |
|
|
| class Prefetcher: |
| """Async double-buffer: 1 background-thread buduje NASTEPNY batch na CPU |
| (pinned) podczas gdy GPU liczy biezacy. queue depth=2. Konsument robi |
| .next() -> H2D-copy (non_blocking) w watku glownym. Watek uzywa wlasnego |
| torch.Generator (seeded), wiec ciag train-batchy jest deterministyczny i |
| NIEZALEZNY od timingu watku oraz od RNG val-loopa (dystrybucja bez zmian). |
| sampler (EpochSampler) zamiast generatora: starty okien z kroku, liczonego od start_step.""" |
| def __init__(self, split, generator, depth=2, sampler=None, start_step=0): |
| self.split = split |
| self.gen = generator |
| self.sampler, self.step = sampler, start_step |
| self.q = queue.Queue(maxsize=depth) |
| self._stop = threading.Event() |
| self.t = threading.Thread(target=self._worker, daemon=True) |
| self.t.start() |
|
|
| def _worker(self): |
| while not self._stop.is_set(): |
| try: |
| if self.sampler is not None and sampler2 is not None and self.step >= a.switch_step: |
| item = _make_batch_cpu("train2", ix=torch.from_numpy(sampler2.starts(self.step - a.switch_step)), |
| shard=True) |
| self.step += 1 |
| elif self.sampler is not None: |
| item = _make_batch_cpu(self.split, ix=torch.from_numpy(self.sampler.starts(self.step)), shard=True) |
| self.step += 1 |
| else: |
| item = _make_batch_cpu(self.split, self.gen, shard=True) |
| except Exception as e: |
| self.q.put(e) |
| return |
| while not self._stop.is_set(): |
| try: |
| self.q.put(item, timeout=0.5) |
| break |
| except queue.Full: |
| continue |
|
|
| def next(self): |
| item = self.q.get() |
| if isinstance(item, Exception): |
| raise item |
| return _to_device(*item) |
|
|
| def close(self): |
| self._stop.set() |
| |
| try: |
| self.q.get_nowait() |
| except queue.Empty: |
| pass |
|
|
| if a.arch == "tinylm": |
| bad = [n for n, v in (("--loop > 1", a.loop != 1), ("--mixer-layout", a.mixer_layout), ("--emb-dim", a.emb_dim), |
| ("--value-residual", a.value_residual), ("--kd-teacher", a.kd_teacher), |
| ("--act-ckpt", a.act_ckpt), ("--loss-chunk", a.loss_chunk), ("--logit-cap", a.logit_cap), |
| ("--z-loss", a.z_loss), ("--pos != rope", a.pos != "rope"), |
| ("--rope-style != half", a.rope_style != "half"), ("--norm != rmsnorm", a.norm != "rmsnorm"), |
| ("--attn-bias != 0", a.attn_bias != 0), ("--ffn gelu", a.ffn == "gelu"), |
| ("--init-from", a.init_from)) if v] |
| if bad: |
| raise SystemExit(f"STOP --arch tinylm: niezgodne {bad}") |
| if a.n_embd % a.n_head or a.n_head % (a.kv_heads or a.n_head): |
| raise SystemExit(f"STOP --arch tinylm: n_embd {a.n_embd} / n_head {a.n_head} / kv_heads {a.kv_heads}") |
| elif a.tinylm_file: |
| raise SystemExit("STOP: --tinylm-file tylko z --arch tinylm") |
| elif a.kv_heads == a.n_head: |
| a.kv_heads = 0 |
| raw_model = GPT(a.vocab, a.n_layer, a.n_embd, a.n_head, a.block, a).to(device) |
| nparam = sum(p.numel() for p in raw_model.parameters()) |
| log(f"model GPT-ref: {nparam/1e6:.1f}M param (L{a.n_layer} d{a.n_embd} h{a.n_head}" |
| + (f" kv{a.kv_heads}" if a.arch == "gpt" and a.kv_heads else "") + ")") |
| if a.arch == "tinylm": |
| log(f"r19 arch=tinylm (Tiny-Chick-10M@60d32c4d tinylm_arch.py {TINYLM_SHA[:8]}): P={nparam} " |
| f"kv_heads={a.kv_heads or a.n_head} ffn={a.ffn} x{a.ffn_mult} qk_norm={a.qk_norm} rope_dims={a.rope_dims} " |
| f"xsa={a.xsa} zero_init={a.zero_init} eps={a.norm_eps} theta={a.rope_theta}") |
| n_train = nparam |
| if a.train_only: |
| if not a.init_from: |
| raise SystemExit("STOP --train-only: wymaga --init-from") |
| if a.arch == "tinylm" or a.loop != 1 or a.mixer_layout or a.kd_teacher or a.compile or ddp: |
| raise SystemExit("STOP --train-only: kombinacja nietestowana (tinylm / loop / mixer-layout / kd / compile / DDP)") |
| rx = re.compile(a.train_only) |
| for name, p in raw_model.named_parameters(): |
| p.requires_grad_(bool(rx.search(name))) |
| n_train = sum(p.numel() for p in raw_model.parameters() if p.requires_grad) |
| if n_train == 0: |
| raise SystemExit(f"STOP --train-only: {a.train_only!r} nie pasuje do zadnego parametru") |
| if raw_model.tok.weight.requires_grad: |
| raise SystemExit(f"STOP --train-only: {a.train_only!r} chwyta osadzenia / tied head") |
| log(f"r21 train-only {a.train_only!r}: uczone {n_train} z {nparam}") |
| opt = build_optimizer(a, raw_model) |
| log(f"optimizer={a.optimizer}" + (f" muon_lr={a.muon_lr} (mult={a.muon_lr/a.lr:.1f}x)" if a.optimizer == "muon" else "")) |
| if a.act_ckpt or a.loss_chunk: |
| log(f"r11 pamiec: act_ckpt={a.act_ckpt} loss_chunk={a.loss_chunk}") |
| if a.imha_t < 0: |
| raise SystemExit("STOP: --imha-t >= 0") |
| if a.imha_t: |
| bad = [n for n, v in (("--xsa", a.xsa), ("--mixer-layout", a.mixer_layout), ("--arch tinylm", a.arch == "tinylm")) if v] |
| if bad: |
| raise SystemExit(f"STOP --imha-t: kombinacja nietestowana {bad}") |
| log(f"r22 IMHA: T={a.imha_t} w {a.n_layer} blokach (jawna macierz uwagi, entropia + bramka); P={nparam}") |
| if a.loop < 1: |
| raise SystemExit("STOP: --loop >= 1") |
| if a.loop > 1: |
| log(f"r12 loop={a.loop}: {a.n_layer} wspoldzielonych blokow x {a.loop} = glebokosc {a.n_layer * a.loop}") |
| if a.mixer_layout or a.attn_bias != 1 or a.rope_style != "interleaved": |
| log(f"r13: mixer_layout={a.mixer_layout or 'A' * a.n_layer} gdn_head_dim={a.gdn_head_dim} " |
| f"attn_bias={a.attn_bias} rope_style={a.rope_style}") |
| if a.xsa or a.rope_dims or a.zero_init or a.ffn == "relu2": |
| log(f"r18: xsa={a.xsa} rope_dims={a.rope_dims or 'wszystkie'} zero_init={a.zero_init} ffn={a.ffn} ffn_mult={a.ffn_mult}") |
| |
| |
| model = DDP(raw_model, device_ids=[local_rank] if device.type == "cuda" else None, |
| broadcast_buffers=False) if ddp else raw_model |
| model = torch.compile(model) if a.compile else model |
| if a.compile: |
| log("torch.compile enabled") |
|
|
| start_step = 0 |
| if a.resume and os.path.exists(ckpt_path): |
| ck = torch.load(ckpt_path, map_location=device, weights_only=True) |
| raw_model.load_state_dict(ck["model"]); opt.load_state_dict(ck["opt"]); start_step = ck["step"] |
| if ck.get("config", {}).get("data_seed") != a.data_seed: |
| raise SystemExit(f"STOP --resume: data_seed w ckpt {ck.get('config', {}).get('data_seed')} != {a.data_seed}") |
| log(f"RESUME @ {start_step}") |
| if a.init_from and start_step == 0: |
| ck0 = torch.load(a.init_from, map_location=device, weights_only=True) |
| c0 = ck0.get("config", {}) |
| for key in ("vocab", "n_layer", "n_embd", "n_head"): |
| if c0.get(key) is not None and c0[key] != getattr(a, key): |
| raise SystemExit(f"STOP --init-from: {key}={c0[key]} w ckpt != {getattr(a, key)} w argumentach") |
| if c0.get("loop", 1) != a.loop: |
| raise SystemExit(f"STOP --init-from: loop={c0.get('loop', 1)} w ckpt != {a.loop} w argumentach") |
| for key, dflt in (("xsa", 0), ("rope_dims", 0), ("ffn", "gelu"), ("imha_t", 0), ("kv_heads", 0)): |
| if c0.get(key, dflt) != getattr(a, key): |
| raise SystemExit(f"STOP --init-from: {key}={c0.get(key, dflt)} w ckpt != {getattr(a, key)} w argumentach") |
| raw_model.load_state_dict(ck0["model"], strict=True) |
| log(f"INIT-FROM {a.init_from} (wagi, krok 0, swiezy optymalizator; grown_from={ck0.get('grown_from')})") |
| del ck0 |
| teacher = None |
| if a.kd_teacher: |
| if not 0.0 <= a.kd_alpha <= 1.0 or a.kd_temp <= 0: |
| raise SystemExit(f"STOP --kd: alpha {a.kd_alpha} poza [0,1] albo temp {a.kd_temp} <= 0") |
| _h = hashlib.sha256() |
| with open(a.kd_teacher, "rb") as _f: |
| for _chunk in iter(lambda: _f.read(1 << 20), b""): |
| _h.update(_chunk) |
| _tsha = _h.hexdigest() |
| if _tsha != (a.kd_teacher_sha or "").lower(): |
| raise SystemExit(f"STOP --kd-teacher: sha {_tsha} != --kd-teacher-sha") |
| ckt = torch.load(a.kd_teacher, map_location="cpu", weights_only=True) |
| ct = ckt["config"] |
| if ct["vocab"] != a.vocab or ct["block"] < a.block: |
| raise SystemExit(f"STOP --kd-teacher: vocab {ct['vocab']} / block {ct['block']} nauczyciela " |
| f"vs vocab {a.vocab} / block {a.block} ucznia") |
| import types |
| with torch.random.fork_rng(devices=[]): |
| teacher = GPT(ct["vocab"], ct["n_layer"], ct["n_embd"], ct["n_head"], ct["block"], types.SimpleNamespace(**ct)) |
| teacher.load_state_dict({k.removeprefix("_orig_mod."): v for k, v in ckt["model"].items()}, strict=True) |
| teacher = teacher.to(device).eval().requires_grad_(False) |
| log(f"r14 KD: nauczyciel {a.kd_teacher} sha {_tsha[:8]} krok {ckt.get('step')} " |
| f"{sum(p.numel() for p in teacher.parameters())/1e6:.1f}M (L{ct['n_layer']} d{ct['n_embd']} " |
| f"{ct.get('mixer_layout') or 'A' * ct['n_layer']}), alpha {a.kd_alpha} T {a.kd_temp}") |
| del ckt |
| teacher_fwd = torch.compile(teacher) if a.compile else teacher |
| elif a.kd_alpha or a.kd_teacher_sha: |
| raise SystemExit("STOP: --kd-alpha/--kd-teacher-sha bez --kd-teacher") |
| if a.data_offset_steps and a.sampler != "epoch": |
| raise SystemExit("STOP: --data-offset-steps wymaga --sampler epoch") |
|
|
| def _lr_sched(s): |
| if s < a.warmup: |
| return a.lr * (s + 1) / a.warmup |
| if s >= a.steps: |
| return a.min_lr |
| if a.schedule == "wsd": |
| |
| decay_start = int(a.steps * (1.0 - a.decay_frac)) |
| if s < decay_start: |
| return a.lr |
| d = (s - decay_start) / max(1, a.steps - decay_start) |
| return a.min_lr + (a.lr - a.min_lr) * (1.0 - d ** 0.5) |
| r = (s - a.warmup) / max(1, a.steps - a.warmup) |
| return a.min_lr + 0.5 * (a.lr - a.min_lr) * (1 + math.cos(math.pi * r)) |
|
|
| |
| |
| rewarm_from = None |
| if start_step and a.rewarm_steps > 0: |
| _g0 = opt.param_groups[0] |
| rewarm_from = _g0["lr"] / _g0.get("lr_mult", 1.0) |
| log(f"rewarm: LR {rewarm_from:.3e} -> harmonogram przez {a.rewarm_steps} krokow od {start_step}") |
|
|
| def lr_at(s): |
| base = _lr_sched(s) |
| if rewarm_from is not None and s < start_step + a.rewarm_steps: |
| f = (s - start_step + 1) / a.rewarm_steps |
| return rewarm_from + (base - rewarm_from) * f |
| return base |
|
|
| @torch.no_grad() |
| def eval_val(): |
| model.eval() |
| ls = [] |
| for _ in range(a.eval_iters): |
| xb, yb = get_batch("val") |
| with ctx: |
| _, loss = model(xb, yb) |
| ls.append(loss.item()) |
| model.train() |
| return sum(ls) / len(ls) |
|
|
| model.train() |
| log(f"START gpt-ref: steps={a.steps} (od {start_step}) lr={a.lr}->{a.min_lr}") |
| |
| |
| train_gen = torch.Generator() |
| ds = a.seed if a.data_seed is None else a.data_seed |
| train_gen.manual_seed(ds) |
| sampler = None |
| sampler2 = None |
| if a.queue: |
| if a.synthetic: |
| raise SystemExit("STOP: --queue wymaga prawdziwego train.bin (bez --synthetic)") |
| with open(a.queue, "rb") as f: |
| _qraw = f.read() |
| sampler = QueueSampler(json.loads(_qraw), len(train_data), a.block, a.batch, ds) |
| if start_step: |
| _ckq = ck.get("queue") or {} |
| if _ckq.get("prefix_sha") != sampler.prefix_sha(start_step): |
| raise SystemExit(f"STOP kolejka: sloty uzyte do kroku {start_step} w manifescie != zapisane w ckpt " |
| "(resume nie bylby dokladny)") |
| log(f"sampler queue: manifest {hashlib.sha256(_qraw).hexdigest()[:8]}, {len(sampler.schedule)} slotow x " |
| f"{sampler.S} krokow, start od kroku {start_step} = slot {start_step // sampler.S}") |
| elif a.sampler == "epoch" and not a.synthetic: |
| sampler = EpochSampler(len(train_data), a.block, a.batch, ds) |
| _, _off0, _nwin0 = sampler._meta(0) |
| log(f"sampler epoch: {_nwin0:,} okien/epoke (offset0 {_off0}), start od kroku {start_step} " |
| f"= epoka {start_step * a.batch / _nwin0:.3f}") |
| if a.data_dir_2: |
| if a.switch_step <= 0 or _nwin0 // a.batch < a.switch_step: |
| raise SystemExit(f"STOP --data-dir-2: --switch-step {a.switch_step} poza epoka 0 train.bin " |
| f"({_nwin0 // a.batch} krokow)") |
| sampler2 = EpochSampler(len(train2_data), a.block, a.batch, ds) |
| _nwin2 = sampler2._meta(0)[2] |
| log(f"data-dir-2: {a.data_dir_2} {len(train2_data):,} tok, {_nwin2:,} okien/epoke = {_nwin2 // a.batch} krokow, " |
| f"od kroku strumienia {a.switch_step}") |
| elif a.data_dir_2: |
| raise SystemExit("STOP: --data-dir-2 wymaga --sampler epoch") |
| elif start_step and not a.synthetic: |
| |
| |
| |
| _n = len(train_data) - a.block - 1 |
| for _ in range(start_step): |
| torch.randint(_n, (a.batch,), generator=train_gen) |
| log(f"train RNG fast-forward o {start_step} krokow (resume)") |
| if a.data_offset_steps: |
| log(f"data-offset: dane od kroku strumienia {start_step + a.data_offset_steps} (krok biegu {start_step} + {a.data_offset_steps})") |
| prefetcher = Prefetcher("train", train_gen, sampler=sampler, start_step=start_step + a.data_offset_steps) |
| end_step = a.steps |
| if a.queue and a.queue_one_slot: |
| end_step = min(a.steps, (start_step // sampler.S + 1) * sampler.S) |
| if sampler2 is not None: |
| _last = a.data_offset_steps + end_step - 1 |
| if _last - a.switch_step >= sampler2._meta(0)[2] // a.batch: |
| raise SystemExit(f"STOP --data-dir-2: krok strumienia {_last} wychodzi poza 1 epoke DATA_DIR_2 (powtorzenie)") |
|
|
| if a.batch_sha: |
| _want = sorted({int(s) for s in a.batch_sha.split(",")}) |
| if _want[0] < start_step or _want[-1] >= end_step: |
| raise SystemExit(f"STOP --batch-sha: kroki poza [{start_step}, {end_step})") |
| for s in range(start_step, _want[-1] + 1): |
| xb, yb = prefetcher.next() |
| if s in _want: |
| _h = hashlib.sha256(xb.cpu().numpy().tobytes() + yb.cpu().numpy().tobytes()).hexdigest() |
| log(f"BATCH-SHA krok biegu {s} = krok strumienia {s + a.data_offset_steps}: {_h}") |
| prefetcher.close() |
| log("BATCH-SHA DONE") |
| return |
|
|
| def save_ckpt(done): |
| if master: |
| state = {"model": raw_model.state_dict(), "opt": opt.state_dict(), "step": done, |
| "config": {"vocab": a.vocab, "n_layer": a.n_layer, "n_embd": a.n_embd, |
| "n_head": a.n_head, "block": a.block, "norm": a.norm, |
| "norm_eps": a.norm_eps, "pos": a.pos, "rope_theta": a.rope_theta, |
| "ffn": a.ffn, "ffn_mult": a.ffn_mult, "value_residual": a.value_residual, |
| "qk_norm": a.qk_norm, "logit_cap": a.logit_cap, "z_loss": a.z_loss, |
| **({"loop": a.loop} if a.loop > 1 else {}), |
| **({"data_seed": a.data_seed} if a.data_seed is not None else {}), |
| **({"mixer_layout": a.mixer_layout, "gdn_head_dim": a.gdn_head_dim} if a.mixer_layout else {}), |
| **({"attn_bias": a.attn_bias} if a.attn_bias != 1 else {}), |
| **({"rope_style": a.rope_style} if a.rope_style != "interleaved" else {}), |
| **({"emb_dim": a.emb_dim} if a.emb_dim else {}), |
| **({"xsa": a.xsa} if a.xsa else {}), |
| **({"rope_dims": a.rope_dims} if a.rope_dims else {}), |
| **({"zero_init": a.zero_init} if a.zero_init else {}), |
| **({"imha_t": a.imha_t} if a.imha_t else {}), |
| **({"kv_heads": a.kv_heads} if a.arch == "gpt" and a.kv_heads else {}), |
| **({"arch": "tinylm", "kv_heads": a.kv_heads or a.n_head, "tinylm_sha256": TINYLM_SHA} |
| if a.arch == "tinylm" else {})}} |
| if a.queue: |
| state["queue"] = {"prefix_sha": sampler.prefix_sha(done), "slots_used": -(-done // sampler.S)} |
| if teacher is not None: |
| state["kd"] = {"teacher": a.kd_teacher, "teacher_sha256": _tsha, "alpha": a.kd_alpha, "temp": a.kd_temp} |
| if a.train_only: |
| state["train_only"] = {"regex": a.train_only, "n_trainable": n_train} |
| torch.save(state, ckpt_path + ".tmp") |
| os.replace(ckpt_path + ".tmp", ckpt_path) |
| log(f"ckpt @ {done}") |
| if a.events_jsonl: |
| _h = hashlib.sha256() |
| with open(ckpt_path, "rb") as _cf: |
| for _chunk in iter(lambda: _cf.read(1 << 20), b""): |
| _h.update(_chunk) |
| emit("checkpoint", done, sha256=_h.hexdigest()) |
| if ddp: |
| dist.barrier() |
|
|
| def finish(): |
| prefetcher.close() |
| if ddp: |
| dist.destroy_process_group() |
|
|
| |
| |
| |
| t0 = time.time(); running = torch.zeros((), dtype=torch.float64, device=device) |
| running_ce = torch.zeros((), dtype=torch.float64, device=device) |
| for step in range(start_step, end_step): |
| lr = lr_at(step) |
| for g in opt.param_groups: |
| g["lr"] = lr * g.get("lr_mult", 1.0) |
| xb, yb = prefetcher.next() |
| if a.grad_accum == 1: |
| with ctx: |
| if teacher is None: |
| _, loss = model(xb, yb) |
| else: |
| with torch.no_grad(): |
| tl = teacher_fwd(xb)[0] |
| _, loss, ce = model(xb, yb, tl) |
| del tl |
| running_ce += ce.to(torch.float64) |
| loss.backward() |
| else: |
| K = a.grad_accum; m = xb.size(0) // K |
| loss = torch.zeros((), dtype=torch.float32, device=device) |
| for j in range(K): |
| xj, yj = xb[j * m:(j + 1) * m], yb[j * m:(j + 1) * m] |
| with (model.no_sync() if ddp and j < K - 1 else nullcontext()): |
| with ctx: |
| if teacher is None: |
| _, lj = model(xj, yj) |
| else: |
| with torch.no_grad(): |
| tl = teacher_fwd(xj)[0] |
| _, lj, ce = model(xj, yj, tl) |
| del tl |
| running_ce += ce.to(torch.float64) / K |
| (lj / K).backward() |
| loss += lj.detach().float() / K |
| gn = torch.nn.utils.clip_grad_norm_(model.parameters(), a.grad_clip) if a.grad_clip > 0 else 0.0 |
| opt.step(); opt.zero_grad(set_to_none=True) |
| running += loss.detach().to(torch.float64) |
| if (step + 1) % a.log_every == 0: |
| if ddp: |
| dist.all_reduce(running) |
| running = running / world |
| if teacher is not None: |
| dist.all_reduce(running_ce) |
| running_ce = running_ce / world |
| running = running.item() |
| kd_ce = running_ce.item() / a.log_every if teacher is not None else None |
| dt = time.time() - t0 |
| tok_s = a.log_every * a.block * a.batch / dt |
| mem = torch.cuda.max_memory_allocated(device)/1e9 if device.type == "cuda" else 0.0 |
| log(f"step {step+1}/{a.steps} loss {running/a.log_every:.4f} lr {lr:.2e} gnorm {float(gn):.2f} {tok_s:,.0f} tok/s peakVRAM {mem:.1f}GB" |
| + (f" ce {kd_ce:.4f}" if kd_ce is not None else "")) |
| if master: |
| with open(metrics_path, "a", encoding="utf-8") as f: |
| f.write(json.dumps({"step": step+1, "loss": running/a.log_every, "lr": lr, "tok_s": tok_s, |
| **({"ce": kd_ce} if kd_ce is not None else {})}) + "\n") |
| emit("update", step + 1, metrics={"loss": running / a.log_every, "tokens_per_second": tok_s, |
| "gradient_norm": float(gn), "learning_rate": lr}) |
| running = torch.zeros((), dtype=torch.float64, device=device); t0 = time.time() |
| running_ce = torch.zeros((), dtype=torch.float64, device=device) |
| if (step + 1) % a.eval_every == 0: |
| vloss = eval_val() |
| log(f" >> VAL loss {vloss:.4f} @ {step+1}") |
| emit("evaluation", step + 1, metrics={"loss": vloss}) |
| if (step + 1) % a.ckpt_every == 0 and not a.synthetic: |
| save_ckpt(step + 1) |
| if end_step % a.ckpt_every and (a.queue or not a.synthetic): |
| save_ckpt(end_step) |
| if end_step < a.steps: |
| log(f"SLOT DONE @ {end_step}") |
| finish() |
| return |
| log("DONE") |
| emit("end", a.steps, status="completed") |
| finish() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
