Add training recipe/script (README.md, train_125m.py)

README.md

@@ -70,7 +70,7 @@ Trzecia iteracja serii GoLLeM. **Główna zmiana vs v2: druga epoka na tym samym
 | Kontekst | 512 tokenów |
 | Tokenizer | polski BPE (dynaword-32k), słownik 32 000, `<\|endoftext\|>`=0 |
 | Dane | korpus v2 ~2,0 mld (58% curated: Wikipedia/Wikisource/Wolne Lektury/1000 Novels/Wiki\*/eltec + 42% HPLT v3 web clean; **zero legalese**), **2 epoki** (~4 mld tok widzianych) |
-| Trening | kontynuacja z ckpt v2 (krok 60 733 → 121 466), bf16, AdamW, cosine LR + warmup, wd 0.1, batch 64 (grad-accum 4) |
+| Trening | kontynuacja z ckpt v2 (krok 60 733 → 121 466), bf16, AdamW, cosine LR (bez ponownej rozgrzewki, patrz „Przepis treningu”), wd 0.1, batch 64 (grad-accum 4) |
 | Sprzęt | 1× AMD Radeon RX 7900 XTX 24GB (gfx1100), ROCm/WSL2, ~9 h |
 
 ## Ewaluacja
@@ -95,6 +95,36 @@ Protokół: OpenPL (`polish4`, 0-shot) + **własna reprodukcja domain-PMI** Oris
 
 **Wniosek:** v3 przewyższa v2 pod spójnym scoringiem (Sygnał-6 +2,55, kompozyt +2,99), rośnie na 7/9. **Regresja:** CBD (mowa nienawiści) −2,0. Pozycja na oficjalnej tablicy OrisTeam — do potwierdzenia przez zgłoszenie modelu do nich (autorytatywny scoring po ich stronie).
 
+## Przepis treningu
+
+Wszystko poniżej pochodzi ze skryptu treningu i logu biegu tego modelu; wartości, których log nie zapisał, są oznaczone jako niezapisane.
+
+**Skrypt:** `train_125m.py` w tym repozytorium (kopia pliku o sha256 `a7fc85ad5352a2e7…` z lokalnymi ścieżkami domyślnymi zamienionymi na `./…` i wczytywaniem checkpointów przy wznowieniu przez `torch.load(..., weights_only=True)`; bez innych zmian). Jeden proces na jednej karcie AMD Radeon RX 7900 XTX (24 GB), ROCm na WSL2.
+
+| | |
+|---|---|
+| Architektura | 12 warstw, d_model 768, 12 głów, kontekst 512, słownik 32 000; bloki pre-LayerNorm, atencja przyczynowa, MLP 4× z GELU, uczone embeddingi pozycji, wspólne wagi embeddingu wejścia i głowy wyjścia; 110 025 216 parametrów |
+| Inicjalizacja | domyślna PyTorch (bez własnej) |
+| Optymalizator | AdamW, beta (0.9, 0.95), eps 1e-8 (domyślne PyTorch), weight decay 0.1 na wszystkich parametrach (jedna grupa), przycinanie gradientu 1.0 |
+| Learning rate | szczyt 3e-4, liniowa rozgrzewka przez 1 % zaplanowanych kroków, potem cosine do 3e-5 (0.1 × szczyt) |
+| Batch | 64 sekwencje × 512 tokenów = 32 768 tokenów na krok |
+| Precyzja | autocast bf16; wagi i stan optymalizatora w fp32 |
+| Próbkowanie danych | losowe okna 512 tokenów z powtórzeniami z jednego pliku tokenów, więc „epoka” = budżet tokenów ÷ rozmiar korpusu, a nie kolejne przejście po danych |
+| Seed | nieustawiony (biegi nie są odtwarzalne bit w bit) |
+| Checkpointy | co 2 000 kroków; wznowienie wczytuje model, optymalizator i krok (bez stanu generatora losowego); wydanie = ostatni krok |
+
+**Ten model (v3):**
+- Kontynuacja v2: model i stan AdamW wczytane z checkpointu v2 (krok 60 733), druga epoka na tym samym korpusie ([SlayerLab/gollem-corpus-2b-pl](https://huggingface.co/datasets/SlayerLab/gollem-corpus-2b-pl)); łącznie ~4,0 mld widzianych tokenów.
+- Harmonogram przeliczono na 2 epoki (121 466 kroków, rozgrzewka 1 214). Ponieważ trening wznowiono od kroku 60 733, rozgrzewka się nie powtórzyła: learning rate skoczył z 3e-5 (koniec v2) do 1,67e-4 i dalej szedł cosine do 3e-5.
+- Gradient accumulation 4 (16 sekwencji w mikro-batchu), według tej karty i jej komendy; log biegu go nie zapisuje. Wydany checkpoint = krok 121 466; ~9,2 h.
+- Komenda (z sekcji Reprodukcja; ścieżki jako placeholdery):
+
+```bash
+python train_125m.py --run-id gollem_v3_e2b --data <gollem_v2_train_32k.bin> \
+  --epochs 2 --batch 64 --accum-steps 4 --lr 3e-4
+# z checkpointem v2 (krok 60 733) w katalogu checkpointów biegu
+```
+
 ## Użycie
 
 ```python
@@ -127,6 +157,6 @@ Benchmark i protokół ewaluacyjny: **OrisTeam** ([Polish-SLM-Benchmark](https:/
 
 ## Reprodukcja
 
-Trening: `train_125m.py --run-id gollem_v3_e2b --data gollem_v2_train_32k.bin --epochs 2 --batch 64 --accum-steps 4 --lr 3e-4` (resume z ckpt v2). Dane treningowe jawnie: [`SlayerLab/gollem-corpus-2b-pl`](https://huggingface.co/datasets/SlayerLab/gollem-corpus-2b-pl) (dokładny korpus v2/v3). Ewaluacja: `board_eval.py` (domain-PMI). Ślad: repo `amd-torch`.
+Trening: `train_125m.py --run-id gollem_v3_e2b --data gollem_v2_train_32k.bin --epochs 2 --batch 64 --accum-steps 4 --lr 3e-4` (resume z ckpt v2). Dane treningowe jawnie: [`SlayerLab/gollem-corpus-2b-pl`](https://huggingface.co/datasets/SlayerLab/gollem-corpus-2b-pl) (dokładny korpus v2/v3). Ewaluacja: `board_eval.py` (domain-PMI). Skrypt treningu: `train_125m.py` w tym repozytorium.
 
 **Autor:** Arkadiusz Słota.

train_125m.py

@@ -0,0 +1,130 @@
+"""125M/110M pretrening LM na ROCm/7900XTX (bf16). Resumable, checkpointowany, logowany.
+Uruchamiaj z `python -u` w środowisku z PyTorch (trenowane na ROCm).
+
+Spec uzgodniony (forum 2026-08-10): d768/12L/12H, vocab=32000, block=512, bf16, wd=0.1, 1 epoka.
+Dane: --data = .bin uint16 tokenow (mmap), najlepiej na szybkim dysku lokalnym.
+Logi: <log-dir>/<run-id>/{train.log, metrics.jsonl}.
+Checkpointy: ./ckpt-<run-id>/ (duze pliki).
+
+Przyklad (Arek odpala live):
+  python -u train_125m.py --data <data-dir>/corpus_32k.bin
+"""
+import argparse, json, math, os, time, glob, datetime
+import numpy as np, torch, torch.nn as nn, torch.nn.functional as F
+
+class Block(nn.Module):
+    def __init__(s, d, nh):
+        super().__init__(); s.nh = nh
+        s.ln1 = nn.LayerNorm(d); s.ln2 = nn.LayerNorm(d)
+        s.qkv = nn.Linear(d, 3*d); s.proj = nn.Linear(d, d)
+        s.fc = nn.Linear(d, 4*d); s.fc2 = nn.Linear(4*d, d)
+    def forward(s, x):
+        B, T, D = x.shape
+        q, k, v = s.qkv(s.ln1(x)).split(D, 2)
+        q = q.view(B, T, s.nh, D//s.nh).transpose(1, 2)
+        k = k.view(B, T, s.nh, D//s.nh).transpose(1, 2)
+        v = v.view(B, T, s.nh, D//s.nh).transpose(1, 2)
+        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        y = y.transpose(1, 2).contiguous().view(B, T, D)
+        x = x + s.proj(y)
+        x = x + s.fc2(F.gelu(s.fc(s.ln2(x))))
+        return x
+
+class GPT(nn.Module):
+    def __init__(s, vocab, d, nl, nh, block):
+        super().__init__()
+        s.tok = nn.Embedding(vocab, d); s.pos = nn.Embedding(block, d)
+        s.blocks = nn.ModuleList([Block(d, nh) for _ in range(nl)])
+        s.lnf = nn.LayerNorm(d); s.head = nn.Linear(d, vocab, bias=False)
+        s.tok.weight = s.head.weight  # weight tying (GPT-2 standard, 110M)
+    def forward(s, idx, tgt=None):
+        B, T = idx.shape
+        x = s.tok(idx) + s.pos(torch.arange(T, device=idx.device))[None]
+        for b in s.blocks: x = b(x)
+        logits = s.head(s.lnf(x))
+        loss = None if tgt is None else F.cross_entropy(logits.view(-1, logits.size(-1)), tgt.view(-1))
+        return logits, loss
+
+def get_batch(data, block, B, dev):
+    ix = torch.randint(len(data) - block - 1, (B,))
+    x = torch.stack([torch.from_numpy(data[i:i+block].astype(np.int64)) for i in ix])
+    y = torch.stack([torch.from_numpy(data[i+1:i+1+block].astype(np.int64)) for i in ix])
+    return x.to(dev, non_blocking=True), y.to(dev, non_blocking=True)
+
+def lr_at(step, warmup, total, lr):
+    if step < warmup: return lr * (step+1) / warmup
+    r = (step - warmup) / max(1, total - warmup)
+    return 0.1*lr + 0.5*(0.9*lr) * (1 + math.cos(math.pi * min(1.0, r)))
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--data", required=True); ap.add_argument("--vocab", type=int, default=32000)
+    ap.add_argument("--d", type=int, default=768); ap.add_argument("--nl", type=int, default=12); ap.add_argument("--nh", type=int, default=12)
+    ap.add_argument("--block", type=int, default=512); ap.add_argument("--batch", type=int, default=64)
+    ap.add_argument("--lr", type=float, default=3e-4); ap.add_argument("--wd", type=float, default=0.1)
+    ap.add_argument("--warmup-frac", type=float, default=0.01); ap.add_argument("--grad-clip", type=float, default=1.0)
+    ap.add_argument("--epochs", type=float, default=1.0); ap.add_argument("--total-tokens", type=float, default=0.0)
+    ap.add_argument("--log-dir", default="./runs")
+    ap.add_argument("--ckpt-every", type=int, default=2000); ap.add_argument("--log-every", type=int, default=20)
+    ap.add_argument("--accum-steps", type=int, default=1)  # grad-accum: microbatch = batch//accum; --batch = EFEKTYWNY batch (LR/total_steps/licznik liczone z batch, nie micro)
+    ap.add_argument("--run-id", default=datetime.datetime.now().strftime("%Y%m%d_%H%M%S"))
+    a = ap.parse_args()
+    dev = "cuda"; assert torch.cuda.is_available(), "ROCm/cuda niedostepny — source env.sh?"
+    rundir = os.path.join(a.log_dir, a.run_id); os.makedirs(rundir, exist_ok=True)
+    ckpt_dir = f"./ckpt-{a.run_id}"; os.makedirs(ckpt_dir, exist_ok=True)
+    logf = open(os.path.join(rundir, "train.log"), "a", buffering=1)
+    metf = open(os.path.join(rundir, "metrics.jsonl"), "a", buffering=1)
+    def log(msg):
+        line = f"[{datetime.datetime.now().strftime('%H:%M:%S')}] {msg}"
+        print(line, flush=True); logf.write(line + "\n")
+
+    data = np.memmap(a.data, dtype=np.uint16, mode="r")
+    tok_per_step = a.batch * a.block
+    total_tokens = a.total_tokens if a.total_tokens > 0 else a.epochs * len(data)
+    total_steps = int(total_tokens / tok_per_step)
+    warmup = max(1, int(a.warmup_frac * total_steps))
+    log(f"run={a.run_id} dev={torch.cuda.get_device_name(0)} data={len(data):,}tok "
+        f"cfg=d{a.d}/{a.nl}L/{a.nh}H/vocab{a.vocab}/block{a.block}/B{a.batch} "
+        f"steps={total_steps:,} tok/step={tok_per_step} wd={a.wd} warmup={warmup} ckpt={ckpt_dir}")
+
+    model = GPT(a.vocab, a.d, a.nl, a.nh, a.block).to(dev)
+    P = sum(p.numel() for p in model.parameters()); log(f"params={P:,}")
+    opt = torch.optim.AdamW(model.parameters(), lr=a.lr, weight_decay=a.wd, betas=(0.9, 0.95), foreach=True)
+    step0 = 0
+    cks = sorted(glob.glob(os.path.join(ckpt_dir, "ckpt_*.pt")))
+    if cks:
+        ck = torch.load(cks[-1], map_location=dev, weights_only=True)
+        model.load_state_dict(ck["model"]); opt.load_state_dict(ck["opt"]); step0 = ck["step"]
+        log(f"RESUME z {cks[-1]} @ step {step0}")
+
+    model.train(); t0 = time.time(); seen = 0
+    for step in range(step0, total_steps):
+        for g in opt.param_groups: g["lr"] = lr_at(step, warmup, total_steps, a.lr)
+        micro = max(1, a.batch // a.accum_steps)
+        opt.zero_grad(set_to_none=True)
+        for _ in range(a.accum_steps):
+            xb, yb = get_batch(data, a.block, micro, dev)
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                _, loss = model(xb, yb)
+            (loss / a.accum_steps).backward()
+        gnorm = nn.utils.clip_grad_norm_(model.parameters(), a.grad_clip); opt.step()
+        seen += tok_per_step
+        if (step+1) % a.log_every == 0:
+            torch.cuda.synchronize(); dt = time.time() - t0; tps = seen / dt
+            eta = (total_steps - step - 1) * tok_per_step / tps / 3600
+            lv, lr_now = loss.item(), opt.param_groups[0]["lr"]
+            log(f"step {step+1:>8,}/{total_steps:,} loss {lv:.4f} lr {lr_now:.2e} "
+                f"grad {float(gnorm):.2f} tok/s {tps:,.0f} elapsed {dt/3600:.2f}h ETA {eta:.1f}h")
+            metf.write(json.dumps({"step": step+1, "loss": lv, "lr": lr_now, "grad_norm": float(gnorm),
+                                   "tok_s": round(tps), "elapsed_s": round(dt), "eta_h": round(eta, 2)}) + "\n")
+        if (step+1) % a.ckpt_every == 0:
+            p = os.path.join(ckpt_dir, f"ckpt_{step+1:08d}.pt")
+            torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "step": step+1, "cfg": vars(a)}, p)
+            log(f"[ckpt] {p}")
+    fp = os.path.join(ckpt_dir, f"ckpt_{total_steps:08d}.pt")
+    if not os.path.exists(fp):  # final ckpt gdy total_steps nie trafia w ckpt_every (np. 41348 % 2000 != 0)
+        torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "step": total_steps, "cfg": vars(a)}, fp)
+        log(f"[ckpt-final] {fp}")
+    log("DONE")
+
+if __name__ == "__main__": main()