Blame for train_gpt.py

#!/usr/bin/env python
# NOTE (GoLLeM-v5 leaderboard repo): this is a GENERAL nanoGPT-style causal transformer;
# vocab/dtype are CLI-parameterized. The crown/Path-B leaderboard checkpoints were trained in
# BPE-12k mode: --vocab 12288 --dtype uint16 (NOT the byte-level default below). The header
# doc-comment reflects the file origin as a standard-GPT control vs experimental BDH; the
# leaderboard models are the STANDARD transformer in BPE mode and do NOT use BDH.
# -*- coding: utf-8 -*-
"""
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()


# ---- Muon (Keller Jordan) --------------------------------------------------
# Ref: https://github.com/KellerJordan/Muon  (modded-nanogpt). Muon = momentum
# SGD, ale update ortogonalizowany przez ~5 krokow iteracji Newtona-Schulza
# (przyblizona ortogonalizacja macierzy gradientu). Stosowany TYLKO do
# macierzowych ukrytych wag (ndim>=2: qkv/proj/mlp). Embeddingi (tok/pos), head
# (tied), LayerNorm-gains i biasy ida do zwyklego AdamW.
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)
                # scale ~ sqrt(fan_out/fan_in): zrownuje RMS update niezaleznie od ksztaltu
                # r17 match_rms_adamw: 0.2*sqrt(max(A,B)) jak torch.optim.Muon (RMS kroku ~ AdamW)
                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))  # r21: bez zamrozonych (--train-only)
    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."))  # r19: osadzenia TinyLM (tied head) jak tok
        tl_gain = name.startswith("tl.") and name.endswith(".attn.gain")  # r19: skalar QK-norm na glowe [H,1,1] -> AdamW
        # r13: w blokach G do Muon tylko kwadratowe projekcje q/k/v/o; konwolucje (3D), bramki niskiego rzedu,
        # b/w_proj, A_log, dt_bias, normy -> AdamW. Bez blokow G warunek jest pusty (r12 bit w bit).
        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":  # r13: GatedDeltaNet2 w miejsce uwagi (fla 0.5.2); inicjalizacja fla, nie _init GPT
            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)  # r23: GQA (0 albo n_head = MHA, qkv d -> 3d jak r22)
            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:  # r22: bramka i kandydat IMHA (W_a, W_f in R^{2d x d}) + temperatura per warstwa
            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:  # blok 0 nie ma vr_lambda (r21): v0 = v z iteracji t = 0, w t >= 1 v bez mieszania
                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)                                  # i + 1
        logn = torch.log(cnt).clamp_min(1e-12)                             # log(i+1); i = 0 -> H_0 = 0 nizej
        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)  # [B,H,T,T] fp32
            z = (A.to(v.dtype) @ v)                                        # [B,H,T,hd]
            if not plain:
                ent = -(A * torch.log(A.clamp_min(1e-30))).sum(-1)         # 0 log 0 = 0 (A = 0 -> 0 * log(1e-30) = 0)
                hn = torch.where(pos == 0, torch.zeros_like(ent), ent / logn)
                e = torch.exp(-hn / tau)
                w = e / (e.cumsum(-1) / cnt)                               # srednia prefiksu, E[w] ~ 1
                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:  # r23: GQA, glowy KV powielone na grupy glow q (po RoPE / QK-norm / value-residual)
            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):  # r18: bez skladowej wzdluz wlasnego v tokenu (per glowa)
            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"  # r19: Tiny-Chick-10M@60d32c4d tinylm_arch.py


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":  # r19: caly model = TinyLM z pliku Tiny Chick (klucze "tl.*")
            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)  # r15: faktoryzowane osadzenia (0 = jak r14)
        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)  # r23: GQA dla blokow A
        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")  # value-residual: v0 z pierwszego bloku A (r12: blok 0)
        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  # tie
        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):  # r18: wyjscia galezi rezydualnych od zera (bloki G: proj fla bez zmian)
            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):  # r13: moduly GatedDeltaNet2 zostaja z inicjalizacja fla
            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":  # r19: ich forward (logity z tied emb), CE w fp32 jak u nich
            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))  # wiersze globalnego batcha tego ranku
    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.enable_cudnn_sdp(False)  # B300 cuDNN SDPA failed in the prior local smoke test
    torch.backends.cuda.matmul.allow_tf32 = bool(a.tf32)  # r17: --tf32 0 jak Qielbas M2 (domyslnie 1 = jak dotad)
    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  # okno = block+1 tokenow od off + k*block

        def _perm_of(self, e):
            if self._perm_e != e:  # trzymamy tylko biezaca epoke (~70 MB przy 9.4B tok / 1024)
                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]:  # epoki maja rozna dlugosc (offset), wiec liczymy po kolei
                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)  # sloty, ktorych kroki juz sie zaczely
            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]
        # (N-block, block+1) view; jeden fancy-index kopiuje wybrane okna
        windows = np.lib.stride_tricks.sliding_window_view(data, a.block + 1)
        sel = windows[ix.numpy()]                      # (batch, block+1) materialized
        x = torch.from_numpy(sel[:, :-1].astype(np.int64))   # astype -> contiguous copy
        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:  # przekaz blad do konsumenta
                    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()
            # opróżnij kolejke zeby watek nie zawisl na put()
            try:
                self.q.get_nowait()
            except queue.Empty:
                pass

    if a.arch == "tinylm":  # r19: tylko kombinacje, ktore TinyLM odtwarza; reszta = STOP zamiast cichego pominiecia
        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:  # r23: K = n_head to zwykla MHA -> jak 0 (wagi, config ckpt i wyjscia = r22)
        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:  # r21: zamrozenie wszystkiego poza dopasowanymi nazwami (wzor ft_arc11.py)
        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():  # tied head/tok wystepuje raz, jako tok.weight
            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}")
    # DDP: identyczne wagi startowe na kazdym ranku (ten sam seed) + broadcast z rank 0 w konstruktorze DDP.
    # broadcast_buffers=False: forward walidacji (bez backward) nie robi kolektywow.
    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)  # r20
        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:  # r23: ta sama kolejność danych po wznowieniu
            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)  # r20 (r18/r19: False)
        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)):  # r18: klucze bez wag tez musza sie zgadzac; r22: imha_t; r23: kv_heads
            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:  # r14: nauczyciel KD (po budowie/wznowieniu ucznia, wiec jego wagi startowe bez zmian)
        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=[]):  # budowa (init) nauczyciela nie zmienia globalnego RNG ucznia
            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":
            # warmup-stable-decay (JugnuLM/MiniCPM): constant peak, then 1-sqrt decay in last decay-frac.
            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)   # 0->1 over decay phase
            return a.min_lr + (a.lr - a.min_lr) * (1.0 - d ** 0.5)  # 1-sqrt: fast early decay
        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-steps: galaz z ckpt (np. cosine @320k -> wsd) nie skacze od razu na nowy harmonogram,
    # tylko przechodzi liniowo od LR zapisanego w ckpt (param_groups po opt.load_state_dict).
    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 = _to_device(*_make_batch_cpu("val", shard=ddp))
            # Validation must obey the training microbatch memory limit.
            for xj, yj in zip(xb.split(1), yb.split(1)):
                with ctx:
                    _, loss = model(xj, yj)
                ls.append(loss.item())
        mean_loss = torch.tensor(sum(ls) / len(ls), device=device)
        if ddp:
            dist.all_reduce(mean_loss)
            mean_loss /= world
        model.train()
        return mean_loss.item()

    model.train()
    log(f"START gpt-ref: steps={a.steps} (od {start_step}) lr={a.lr}->{a.min_lr}")
    # dedykowany seeded generator dla train-prefetchera (determinizm niezalezny
    # od RNG val-loopa i timingu watku; ta sama dystrybucja co global-RNG)
    train_gen = torch.Generator()
    ds = a.seed if a.data_seed is None else a.data_seed  # r23: kolejność danych osobno od inicjalizacji
    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:
        # --resume: przewin generator o start_step batchy, zeby wznowiony run losowal DOKLADNIE te okna,
        # ktore run ciagly losowalby od kroku start_step (bez tego resume powtarzal okna z krokow 0..N).
        # Kazdy krok zuzywa 1 x torch.randint(len-block-1, (batch,)); 320k krokow ~1 s CPU.
        _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  # ostatni krok strumienia tego biegu
        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:  # r21: poza "config", wiec --init-from i valport bez zmian
                state["train_only"] = {"regex": a.train_only, "n_trainable": n_train}
            torch.save(state, ckpt_path + ".tmp")
            os.replace(ckpt_path + ".tmp", ckpt_path)  # crash w trakcie zapisu nie zostawia ucietego ckpt.pt
            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()  # inne ranki nie ida dalej (ani nie koncza), zanim rank 0 zapisze ckpt

    def finish():
        prefetcher.close()
        if ddp:
            dist.destroy_process_group()

    # r9 (Monter, 06.10): strata sumowana na urządzeniu (float64), .item() tylko przy logu. r8 robił loss.item() co krok,
    # czyli synchronizację GPU->CPU co krok, przez co CPU nie mógł kolejkować następnego kroku. Obliczenia modelu,
    # kolejność batchy i aktualizacje wag są identyczne z r8; zmienia się tylko moment odczytu straty do logu.
    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:  # r16: K mikrobatchy z wierszy ranku, gradient = srednia (strata_j / K); DDP synchronizuje tylko ostatni
            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
        # Fail all ranks together before applying a non-finite update.
        finite = torch.isfinite(loss.detach()).all() & torch.isfinite(torch.as_tensor(gn, device=device)).all()
        failure = (~finite).to(dtype=torch.int32)
        if ddp:
            dist.all_reduce(failure, op=dist.ReduceOp.MAX)
        if failure.item():
            raise FloatingPointError(f"Non-finite loss or gradient norm at step {step + 1}")
        opt.step(); opt.zero_grad(set_to_none=True)
        running += loss.detach().to(torch.float64)
        if (step + 1) % a.log_every == 0:
            if ddp:  # strata logowana = srednia po wszystkich wierszach globalnego batcha (jak w 1-GPU)
                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)
        # Persist the latest training state before the validation pass.
        if ((step + 1) == 25 or (step + 1) % a.ckpt_every == 0) and not a.synthetic:
            save_ckpt(step + 1)
        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 end_step % a.ckpt_every and (a.queue or not a.synthetic):
        save_ckpt(end_step)  # r20: zawsze ckpt po ostatnim kroku poza rytmem ckpt_every (r18/r19: tylko z --queue; k9 13.10: 1140 zamiast 1145)
    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()