Slayer1B 1.07: converted weights of the final checkpoint (step 22327)

README.md

@@ -20,10 +20,10 @@ Decoder-only language model, 1 073 912 851 parameters, Polish-dominant. **Slayer
 Slayer1B 1.06** from its final checkpoint (step 13848) on 8 891 005 707 training tokens (about 7.0 B new and 1.9 B
 replayed from the 1.06 data), at the same constant learning rate and global batch.
 
-**Status: training complete at step 22327 (6 November 2026).** All 15 evaluated checkpoints of the run, from the 1.06
-starting point (13848) to the final checkpoint (22327), are in the training-time evaluation table below. A 1.07 model release (converted weights) has not been
-published in this repository yet. Native training checkpoints (raw PyTorch training state `ckpt_N.pt`, not a
-ready-to-use model) are published file by file in
+**Status: training complete at step 22327 (6 November 2026). This repository contains the converted weights of the
+final checkpoint (step 22327)**; see Files below. All 15 evaluated checkpoints of the run, from the 1.06 starting point
+(13848) to the final checkpoint (22327), are in the training-time evaluation table below. Native training checkpoints
+(raw PyTorch training state `ckpt_N.pt`, not a ready-to-use model) are published file by file in
 [SlayerLab/Slayer1B-1.07-checkpoints](https://huggingface.co/SlayerLab/Slayer1B-1.07-checkpoints), with sha256 sums.
 
 - Previous version: [SlayerLab/Slayer1B-1.06](https://huggingface.co/SlayerLab/Slayer1B-1.06) (step 13848,
@@ -111,3 +111,21 @@ were scanned for 13-gram overlap with the test items of the evaluation suite bef
 hit were removed. Exact matches of short answers (below 13 words) were counted but not removed, and one evaluation set
 (EWoK, distributed as a password-protected archive) was not part of the scan. The Slayer 1.06 core includes the source
 passage corpora of PolQA and PoQuAD; passages with a 13-gram hit were removed and the rest remain in the replay.
+
+## Files
+
+`model.safetensors` (weights, float32), `config.json`, `tokenizer.json`, `load_slayer.py` (loader),
+`train_gpt_ref_r23.py` and `SHA256SUMS`. The `config.json` records the sha256 of the source checkpoint `ckpt_22327.pt`
+(`48cdb5eb16b69e0ee9a7a2ed72895ff47c0ba57aca0f4516b1f7e9cba286685c`, also published in the checkpoints repository).
+The file layout and the loader are the same as in [SlayerLab/Slayer1B-1.06](https://huggingface.co/SlayerLab/Slayer1B-1.06):
+
+```python
+import load_slayer
+from tokenizers import Tokenizer
+
+model = load_slayer.load(".")            # float32, eval mode, on CPU
+tok = Tokenizer.from_file("tokenizer.json")  # id 0 is end-of-text
+```
+
+`train_gpt_ref_r23.py` is the model definition (`GPT` class) imported by `load_slayer.py`; `trainer_sha256` in
+`config.json` refers to this file. The architecture is unchanged from 1.00–1.06.

SHA256SUMS

@@ -0,0 +1,5 @@
+9e2062501a87d0b87d7ef75ba8bc5484f3b1afdf42081e727a74dc02f5a4afce  config.json
+53a04c726bebbf8d8a25ff1479ab4f70c3003022675d87332d130139cca1d020  load_slayer.py
+d826b4a11bf4eaccee10fe43ec19a4f6d515dfc5f3e5c859adb631f1a0b76ad5  model.safetensors
+0640d3bd3674d7a6e59c540d945a2ad275a12e1bcaeb6007c37a90751f9815dd  tokenizer.json
+db36c150e95c184e82b5f84c3ee783fd5fd2c62e960a438177f543dc3eba1d0b  train_gpt_ref_r23.py

config.json

@@ -0,0 +1,28 @@
+{
+  "attn_bias": 0,
+  "block": 4096,
+  "ckpt_sha256": "48cdb5eb16b69e0ee9a7a2ed72895ff47c0ba57aca0f4516b1f7e9cba286685c",
+  "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": 22327,
+  "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:d826b4a11bf4eaccee10fe43ec19a4f6d515dfc5f3e5c859adb631f1a0b76ad5
+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()