Slayer1B 1.00 weights: step 1430 (ckpt sha 7d092d4d), exported 1:1 from the trainer checkpoint

SHA256SUMS

@@ -0,0 +1,5 @@
+55445c12307a780e841b234bc381865a1ec9b301dc9e5a76c85fff1e8130a111  config.json
+53a04c726bebbf8d8a25ff1479ab4f70c3003022675d87332d130139cca1d020  load_slayer.py
+1c97c0527e1cff6baf3a07539a1947469b476769891399c505854824ebd7a0e6  model.safetensors
+0640d3bd3674d7a6e59c540d945a2ad275a12e1bcaeb6007c37a90751f9815dd  tokenizer.json
+db36c150e95c184e82b5f84c3ee783fd5fd2c62e960a438177f543dc3eba1d0b  train_gpt_ref_r23.py

config.json

@@ -0,0 +1,28 @@
+{
+  "attn_bias": 0,
+  "block": 4096,
+  "ckpt_sha256": "7d092d4d0bf838da1115a5369bf0a6729fe4c170e0d5ad303873f7320d6c1866",
+  "data_seed": 20261023,
+  "ffn": "swiglu",
+  "ffn_mult": 2.667,
+  "logit_cap": 0.0,
+  "n_embd": 2048,
+  "n_head": 16,
+  "n_layer": 20,
+  "n_params": 1073912851,
+  "norm": "rmsnorm",
+  "norm_eps": 1e-06,
+  "pos": "rope",
+  "qk_norm": true,
+  "rope_theta": 500000.0,
+  "step": 1430,
+  "tied_weights": {
+    "tok.weight": "head.weight"
+  },
+  "tokenizer_sha256": "0640d3bd3674d7a6e59c540d945a2ad275a12e1bcaeb6007c37a90751f9815dd",
+  "trainer": "train_gpt_ref_r23.py",
+  "trainer_sha256": "db36c150e95c184e82b5f84c3ee783fd5fd2c62e960a438177f543dc3eba1d0b",
+  "value_residual": true,
+  "vocab": 32768,
+  "z_loss": 0.0
+}

load_slayer.py

@@ -0,0 +1,38 @@
+"""Load Slayer weights (model.safetensors + config.json) into the GPT class of train_gpt_ref_r23.py.
+
+Usage:
+    import load_slayer
+    model = load_slayer.load("path/to/Slayer1B-1.00")   # float32, eval mode, on CPU
+    logits, _ = model(token_ids)                          # token_ids: LongTensor [batch, seq], seq <= config["block"]
+Tokenizer: tokenizers.Tokenizer.from_file("tokenizer.json"); id 0 is end-of-text.
+"""
+import importlib.util
+import json
+import os
+from types import SimpleNamespace
+
+from safetensors.torch import load_file
+
+
+def _trainer(path):
+    spec = importlib.util.spec_from_file_location("train_gpt_ref_r23", os.path.join(path, "train_gpt_ref_r23.py"))
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+def build(cfg, path):
+    """Empty model with the architecture from config (no weights loaded)."""
+    c = SimpleNamespace(**cfg)
+    return _trainer(path).GPT(c.vocab, c.n_layer, c.n_embd, c.n_head, c.block, c)
+
+
+def load(path, device="cpu"):
+    with open(os.path.join(path, "config.json"), encoding="utf-8") as f:
+        cfg = json.load(f)
+    model = build(cfg, path)
+    sd = load_file(os.path.join(path, "model.safetensors"))
+    for dst, src in cfg.get("tied_weights", {}).items():
+        sd.setdefault(src, sd[dst])
+    model.load_state_dict(sd, strict=True)
+    return model.to(device).float().eval()

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1c97c0527e1cff6baf3a07539a1947469b476769891399c505854824ebd7a0e6
+size 4295670676

train_gpt_ref_r23.py

@@ -0,0 +1,1321 @@
+#!/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.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 = 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}")
+    # 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
+        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)
+        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)  # 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()