Upload 7 files

data_prepare.py

@@ -0,0 +1,177 @@
+# -- coding: utf-8 --
+import os
+from datasets import load_dataset
+from tqdm import tqdm
+import sentencepiece as spm
+import numpy as np
+
+# ===========================================================
+# KONFIGURACE
+# ===========================================================
+TARGET_TOKENS = 1_000_000_000       # 100M pro test, může být 1_000_000_000 a víc
+VOCAB_SIZE = 32_000
+RAW_TEXT_PATH = "dataset.txt"
+
+TOKENIZER_MODEL_PATH = "tokenizer.model"
+
+BIN_TRAIN_PATH = "dataset.bin"
+BIN_VALID_PATH = "valid.bin"
+
+TRAIN_RATIO = 0.98     # 98% trénink, 2% valid
+
+SPECIAL_TOKENS = {
+    "unk_id": 0,
+    "bos_id": 1,
+    "eos_id": 2,
+    "pad_id": 3,
+}
+
+# ===========================================================
+# 1) STREAMOVANÉ STAŽENÍ FINEWEB -> dataset.txt
+# ===========================================================
+def download_fineweb_streaming():
+    if os.path.exists(RAW_TEXT_PATH):
+        print("✔ dataset.txt už existuje, přeskočeno.")
+        return
+
+    print("📥 Stahuji FineWeb-Edu streamovacím způsobem...")
+
+    dataset = load_dataset(
+        "HuggingFaceFW/fineweb-edu",
+        name="sample-10BT",
+        split="train",
+        streaming=True
+    )
+
+    tokens_so_far = 0
+
+    with open(RAW_TEXT_PATH, "w", encoding="utf-8") as f:
+        for example in tqdm(dataset, desc="Stahuji dataset"):
+            text = example["text"].strip() + "\n\n"
+            approx = len(text) // 4  # odhad tokenů
+
+            if tokens_so_far + approx > TARGET_TOKENS:
+                remaining = TARGET_TOKENS - tokens_so_far
+                chars = remaining * 4
+                f.write(text[:chars])
+                print("✔ dataset.txt hotovo.")
+                return
+
+            f.write(text)
+            tokens_so_far += approx
+
+            if tokens_so_far >= TARGET_TOKENS:
+                print("✔ dataset.txt hotovo.")
+                return
+
+# ===========================================================
+# 2) TRÉNINK SENTENCEPIECE TOKENIZERU
+# ===========================================================
+def train_tokenizer():
+    if os.path.exists(TOKENIZER_MODEL_PATH):
+        print("✔ Tokenizer už existuje, přeskakuji.")
+        return
+
+    print("🔧 Trénuji SentencePiece tokenizer...")
+
+    prefix = TOKENIZER_MODEL_PATH.replace(".model", "")
+
+    spm.SentencePieceTrainer.train(
+        input=RAW_TEXT_PATH,
+        model_prefix=prefix,
+        vocab_size=VOCAB_SIZE,
+        model_type="unigram",
+        character_coverage=1.0,
+        byte_fallback=True,
+
+        unk_id=SPECIAL_TOKENS["unk_id"],
+        bos_id=SPECIAL_TOKENS["bos_id"],
+        eos_id=SPECIAL_TOKENS["eos_id"],
+        pad_id=SPECIAL_TOKENS["pad_id"],
+
+        train_extremely_large_corpus=True,
+    )
+
+    print("✔ Tokenizer natrénován.")
+
+# ===========================================================
+# 3) STREAMOVÁ TOKENIZACE → BIN FILE (INT32)
+# ===========================================================
+def tokenize_to_bin_streaming():
+    """
+    Streamovací tokenizace velkého datasetu do binárních souborů (int32),
+    bez držení celého datasetu v paměti.
+    """
+    if os.path.exists(BIN_TRAIN_PATH) and os.path.exists(BIN_VALID_PATH):
+        print("✔ dataset.bin + valid.bin už existují.")
+        return
+
+    print("🔠 Streamuji text → tokeny (int32) → dataset.bin...")
+
+    sp = spm.SentencePieceProcessor(model_file=TOKENIZER_MODEL_PATH)
+    EOS = sp.eos_id()
+
+    # ===========================================================
+    # 1️⃣ ZJIŠTĚNÍ CELKOVÉHO POČTU TOKENŮ
+    # ===========================================================
+    print("🔎 Počítám tokeny...")
+    total_tokens = 0
+    with open(RAW_TEXT_PATH, "r", encoding="utf-8") as f:
+        for line in tqdm(f, desc="Počítám tokeny"):
+            line = line.strip()
+            if not line:
+                continue
+            total_tokens += len(sp.encode(line)) + 1  # +1 pro EOS
+
+    train_tokens = int(total_tokens * TRAIN_RATIO)
+    valid_tokens = total_tokens - train_tokens
+
+    print(f"Celkem tokenů: {total_tokens:,}")
+    print(f"Train: {train_tokens:,}")
+    print(f"Valid: {valid_tokens:,}")
+
+    # ===========================================================
+    # 2️⃣ VYTVOŘENÍ MEMMAP SOUBORŮ
+    # ===========================================================
+    train_mm = np.memmap(BIN_TRAIN_PATH, dtype=np.int32, mode="w+", shape=(train_tokens,))
+    valid_mm = np.memmap(BIN_VALID_PATH, dtype=np.int32, mode="w+", shape=(valid_tokens,))
+
+    # ===========================================================
+    # 3️⃣ STREAMOVÁ TOKENIZACE A ZÁPIS
+    # ===========================================================
+    print("✍ Tokenizuji a zapisují do memmap...")
+    ti, vi = 0, 0  # indexy do train/valid memmap
+
+    with open(RAW_TEXT_PATH, "r", encoding="utf-8") as f:
+        for line in tqdm(f, desc="Tokenizuji dataset"):
+            line = line.strip()
+            if not line:
+                continue
+
+            ids = sp.encode(line) + [EOS]
+
+            for tok in ids:
+                if ti < train_tokens:
+                    train_mm[ti] = tok
+                    ti += 1
+                else:
+                    valid_mm[vi] = tok
+                    vi += 1
+
+    # ===========================================================
+    # 4️⃣ FLUSH MEMMAP
+    # ===========================================================
+    train_mm.flush()
+    valid_mm.flush()
+
+    print("✔ Hotovo — dataset.bin + valid.bin připravené pro trénink!")
+
+# ===========================================================
+# MAIN
+# ===========================================================
+if __name__ == "__main__":
+    download_fineweb_streaming()
+    train_tokenizer()
+    tokenize_to_bin_streaming()
+    print("\n🎉 HOTOVO — dataset.bin + valid.bin připravené pro trénink!")
+

eval-loss.py

@@ -0,0 +1,245 @@
+# -- coding: utf-8 --
+# Compare validation loss of multiple GPT checkpoints
+# Works with old and new checkpoint formats
+# Compatible with Antonín Tomeček Transformer code
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import sentencepiece as spm
+import numpy as np
+from torch.utils.data import Dataset, DataLoader
+from tqdm import tqdm
+
+# =========================
+# CONFIG
+# =========================
+CHECKPOINTS = {
+    "pretrain_900k": "checkpoints/step_900000.pt",
+    "continual_100k": "checkpoints/step_100000.pt",
+    "continual_200k": "checkpoints/step_200000.pt",
+    "continual_300k": "checkpoints/step_300000.pt",
+    "continual_400k": "checkpoints/step_400000.pt",
+    "continual_500k": "checkpoints/step_500000.pt",
+}
+
+TOKENIZER_MODEL_PATH = "tokenizer.model"
+VALID_BIN = "valid.bin"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+BATCH_SIZE = 1  # můžeš zvýšit podle VRAM
+
+# =========================
+# ModelArgs
+# =========================
+from dataclasses import dataclass
+
+@dataclass
+class ModelArgs:
+    dim: int = 768
+    n_layers: int = 12
+    n_heads: int = 12
+    n_kv_heads: int = 4
+    vocab_size: int = 32000
+    multiple_of: int = 256
+    ffn_dim_multiplier: float = 3.0
+    norm_eps: float = 1e-5
+    max_seq_len: int = 1024
+
+# =========================
+# Dataset
+# =========================
+class MemmapDataset(Dataset):
+    def __init__(self, path: str, max_seq_len: int, stride=None):
+        self.tokens = np.memmap(path, dtype=np.int32, mode="r")
+        self.max_seq_len = max_seq_len
+        self.stride = stride or max_seq_len // 2
+
+        max_start = len(self.tokens) - (max_seq_len + 1)
+        if max_start <= 0:
+            raise ValueError("Dataset too small")
+
+        self.starts = list(range(0, max_start, self.stride))
+        if self.starts[-1] != max_start:
+            self.starts.append(max_start)
+
+    def __len__(self):
+        return len(self.starts)
+
+    def __getitem__(self, idx):
+        i = self.starts[idx]
+        seq = torch.from_numpy(
+            self.tokens[i:i + self.max_seq_len + 1].copy()
+        ).long()
+        return seq[:-1], seq[1:]
+
+# =========================
+# Transformer model
+# =========================
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
+
+def precompute_freqs_cis(dim, seq_len, theta=10000.0):
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2) / dim))
+    t = torch.arange(seq_len)
+    freqs = torch.outer(t, freqs)
+    return freqs.cos(), freqs.sin()
+
+def apply_rotary_emb(x, cos, sin):
+    x1, x2 = x[..., 0::2], x[..., 1::2]
+    cos = cos.unsqueeze(0).unsqueeze(2)
+    sin = sin.unsqueeze(0).unsqueeze(2)
+    out = torch.empty_like(x)
+    out[..., 0::2] = x1 * cos - x2 * sin
+    out[..., 1::2] = x1 * sin + x2 * cos
+    return out
+
+class Attention(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.n_heads = args.n_heads
+        self.head_dim = args.dim // args.n_heads
+        self.n_kv_heads = args.n_kv_heads
+        self.repeat_kv = args.n_heads // args.n_kv_heads
+
+        self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
+        self.wk = nn.Linear(args.dim, args.n_kv_heads * self.head_dim, bias=False)
+        self.wv = nn.Linear(args.dim, args.n_kv_heads * self.head_dim, bias=False)
+        self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
+
+    def forward(self, x, cos, sin):
+        B, T, _ = x.shape
+        q = self.wq(x).view(B, T, self.n_heads, self.head_dim)
+        k = self.wk(x).view(B, T, self.n_kv_heads, self.head_dim)
+        v = self.wv(x).view(B, T, self.n_kv_heads, self.head_dim)
+        k = k.repeat_interleave(self.repeat_kv, dim=2)
+        v = v.repeat_interleave(self.repeat_kv, dim=2)
+        q = apply_rotary_emb(q, cos, sin)
+        k = apply_rotary_emb(k, cos, sin)
+        q = q.transpose(1,2)
+        k = k.transpose(1,2)
+        v = v.transpose(1,2)
+        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        out = out.transpose(1,2).contiguous().view(B, T, -1)
+        return self.wo(out)
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, multiple_of, mult):
+        super().__init__()
+        hidden = multiple_of * ((int(dim * mult) + multiple_of -1)//multiple_of)
+        self.w1 = nn.Linear(dim, hidden, bias=False)
+        self.w2 = nn.Linear(hidden, dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden, bias=False)
+    def forward(self,x):
+        return self.w2(F.silu(self.w1(x))*self.w3(x))
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.attn = Attention(args)
+        self.ffn = FeedForward(args.dim, args.multiple_of, args.ffn_dim_multiplier)
+        self.attn_norm = RMSNorm(args.dim, args.norm_eps)
+        self.ffn_norm = RMSNorm(args.dim, args.norm_eps)
+    def forward(self, x, cos, sin):
+        x = x + self.attn(self.attn_norm(x), cos, sin)
+        x = x + self.ffn(self.ffn_norm(x))
+        return x
+
+class Transformer(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.tok_emb = nn.Embedding(args.vocab_size, args.dim)
+        self.layers = nn.ModuleList([TransformerBlock(args) for _ in range(args.n_layers)])
+        self.norm = RMSNorm(args.dim, args.norm_eps)
+        self.out = nn.Linear(args.dim, args.vocab_size, bias=False)
+        cos, sin = precompute_freqs_cis(args.dim//args.n_heads, args.max_seq_len*2)
+        self.register_buffer("cos_cached", cos, persistent=False)
+        self.register_buffer("sin_cached", sin, persistent=False)
+    def forward(self, tokens):
+        B, T = tokens.shape
+        h = self.tok_emb(tokens)
+        cos = self.cos_cached[:T]
+        sin = self.sin_cached[:T]
+        for layer in self.layers:
+            h = layer(h, cos, sin)
+        h = self.norm(h)
+        return self.out(h)
+
+# =========================
+# Eval function
+# =========================
+def evaluate_checkpoint(path, valid_loader, tokenizer, args):
+    ckpt = torch.load(path, map_location="cpu", weights_only=False)
+
+    # Podpora starého i nového formátu checkpointu
+    if isinstance(ckpt, dict) and "model_state_dict" in ckpt:
+        state_dict = ckpt["model_state_dict"]
+    else:
+        state_dict = ckpt
+
+    model = Transformer(args)
+    model.load_state_dict(state_dict)
+    model.to(DEVICE)
+    model.eval()
+
+    total_loss = 0.0
+    total_tokens = 0
+
+    with torch.no_grad():
+        for x, y in valid_loader:
+            x = x.to(DEVICE)
+            y = y.to(DEVICE)
+
+            logits = model(x)
+            loss = F.cross_entropy(
+                logits.view(-1, logits.size(-1)),
+                y.view(-1),
+                ignore_index=tokenizer.pad_id(),
+                reduction="sum",
+            )
+
+            total_loss += loss.item()
+            total_tokens += (y != tokenizer.pad_id()).sum().item()
+
+    return total_loss / total_tokens
+
+# =========================
+# MAIN
+# =========================
+def main():
+    # pevné ModelArgs
+    args = ModelArgs()
+    tokenizer = spm.SentencePieceProcessor(model_file=TOKENIZER_MODEL_PATH)
+    args.vocab_size = tokenizer.vocab_size()
+
+    # dataset
+    valid_ds = MemmapDataset(VALID_BIN, args.max_seq_len)
+    valid_loader = DataLoader(valid_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=2, pin_memory=True)
+
+    print("="*70)
+    print("Checkpoint comparison (validation)")
+    print("="*70)
+
+    results = {}
+    for name, path in CHECKPOINTS.items():
+        print(f"[Eval] {name}")
+        loss = evaluate_checkpoint(path, valid_loader, tokenizer, args)
+        ppl = math.exp(loss)
+        results[name] = (loss, ppl)
+        print(f"  Val loss: {loss:.6f}")
+        print(f"  Perplexity: {ppl:.2f}")
+        print("-"*50)
+
+    print("\nSummary:")
+    for name, (loss, ppl) in results.items():
+        print(f"{name:20s} | loss {loss:.6f} | ppl {ppl:.2f}")
+    print("="*70)
+
+if __name__ == "__main__":
+    main()

test-checkpoints.py

@@ -0,0 +1,86 @@
+# -- coding: utf-8 --
+# Author: Antonín Tomeček
+# Date: 10 Jan 2026
+# Description: Standalone text generation from GPT-style checkpoint 500k
+
+import os
+import torch
+import sentencepiece as spm
+
+# importuj model a třídy z tvého tréninkového souboru
+from train import Transformer, ModelArgs, generate_text  # uprav podle názvu souboru
+
+# =========================
+# CONFIG
+# =========================
+CHECKPOINT_PATH = "checkpoints/step_500000.pt"
+TOKENIZER_MODEL_PATH = "tokenizer.model"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+MAX_NEW_TOKENS = 200
+TEMPERATURE = 0.8
+TOP_P = 0.95
+EOS_ID = 1  # podle tokenizeru, většinou 1 je </s>
+
+# =========================
+# Povolit ModelArgs při odpickle
+# =========================
+torch.serialization.add_safe_globals([ModelArgs])
+
+# =========================
+# LOAD TOKENIZER
+# =========================
+tokenizer = spm.SentencePieceProcessor(model_file=TOKENIZER_MODEL_PATH)
+vocab_size = tokenizer.vocab_size()
+
+# =========================
+# LOAD CHECKPOINT
+# =========================
+if not os.path.exists(CHECKPOINT_PATH):
+    raise FileNotFoundError(f"Checkpoint {CHECKPOINT_PATH} not found")
+
+checkpoint = torch.load(CHECKPOINT_PATH, map_location=DEVICE, weights_only=False)
+
+# načteme model podle uložených args
+model_args = checkpoint.get("model_args", ModelArgs())
+model_args.vocab_size = vocab_size
+model = Transformer(model_args).to(DEVICE)
+
+# načteme váhy
+model.load_state_dict(checkpoint["model_state_dict"])
+model.eval()
+
+print(f"[Info] Loaded checkpoint from step {checkpoint.get('step', 'unknown')}")
+print(f"[Info] Model has {sum(p.numel() for p in model.parameters() if p.requires_grad):,} params")
+
+# =========================
+# PROMPTS
+# =========================
+prompts = [
+    "Once upon a time",
+    "In a distant future",
+    "Artificial intelligence will",
+    "First step to build a rocket",
+    "Capital city of France"
+]
+
+# =========================
+# GENERATE TEXT
+# =========================
+results = generate_text(
+    model,
+    tokenizer,
+    prompts,
+    max_new_tokens=MAX_NEW_TOKENS,
+    temperature=TEMPERATURE,
+    top_p=TOP_P,
+    eos_id=EOS_ID
+)
+
+# =========================
+# PRINT RESULTS
+# =========================
+for prompt, text in results.items():
+    print("="*50)
+    print(f"Prompt: {prompt}")
+    print(f"Generated: {text}")

tokenizer.model

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

train.py

@@ -0,0 +1,519 @@
+# -- coding: utf-8 --
+# Author: Antonín Tomeček
+# Date: 3 Jan. 2026
+# Description: GPT-style Transformer with Flash Attention 2, Memmap dataset,
+#              correct gradient accumulation, and clean English logging.
+
+import os
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from dataclasses import dataclass
+from typing import Optional
+from torch.utils.data import Dataset, DataLoader
+from accelerate import Accelerator
+from tqdm import tqdm
+import sentencepiece as spm
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+# =========================
+# FLASH ATTENTION 2
+# =========================
+try:
+    print(f"[Info] Torch version: {torch.__version__}")
+    print(f"[Info] CUDA available: {torch.cuda.is_available()}")
+    if torch.cuda.is_available():
+        print(f"[Info] CUDA version: {torch.version.cuda}")
+
+    from flash_attn import flash_attn_func
+    FLASH_ATTENTION_2 = True
+    print("[OK] Flash Attention 2 enabled")
+except Exception:
+    FLASH_ATTENTION_2 = False
+    print("[WARN] Flash Attention 2 not available – using PyTorch SDPA")
+
+# =========================
+# CONFIG
+# =========================
+@dataclass
+class ModelArgs:
+    dim: int = 768
+    n_layers: int = 12
+    n_heads: int = 12
+    n_kv_heads: int = 4
+    vocab_size: int = 32000
+    multiple_of: int = 256
+    ffn_dim_multiplier: float = 3.0
+    norm_eps: float = 1e-5
+    max_seq_len: int = 1024
+
+
+SAVE_EVERY_STEPS = 100_000
+TOKENIZER_MODEL_PATH = "tokenizer.model"
+TRAIN_BIN = "dataset.bin"
+VALID_BIN = "valid.bin"
+CHECKPOINT_DIR = "checkpoints"
+os.makedirs(CHECKPOINT_DIR, exist_ok=True)
+
+# =========================
+# MODEL
+# =========================
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
+
+
+def precompute_freqs_cis(dim, seq_len, theta=10000.0):
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2) / dim))
+    t = torch.arange(seq_len)
+    freqs = torch.outer(t, freqs)
+    return freqs.cos(), freqs.sin()
+
+
+def apply_rotary_emb(x, cos, sin):
+    x1, x2 = x[..., 0::2], x[..., 1::2]
+    cos = cos.unsqueeze(0).unsqueeze(2)
+    sin = sin.unsqueeze(0).unsqueeze(2)
+    out = torch.empty_like(x)
+    out[..., 0::2] = x1 * cos - x2 * sin
+    out[..., 1::2] = x1 * sin + x2 * cos
+    return out
+
+
+class Attention(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.n_heads = args.n_heads
+        self.head_dim = args.dim // args.n_heads
+        self.n_kv_heads = args.n_kv_heads
+        self.repeat_kv = args.n_heads // args.n_kv_heads
+
+        self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
+        self.wk = nn.Linear(args.dim, args.n_kv_heads * self.head_dim, bias=False)
+        self.wv = nn.Linear(args.dim, args.n_kv_heads * self.head_dim, bias=False)
+        self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
+
+    def forward(self, x, cos, sin):
+        B, T, _ = x.shape
+
+        q = self.wq(x).view(B, T, self.n_heads, self.head_dim)
+        k = self.wk(x).view(B, T, self.n_kv_heads, self.head_dim)
+        v = self.wv(x).view(B, T, self.n_kv_heads, self.head_dim)
+
+        k = k.repeat_interleave(self.repeat_kv, dim=2)
+        v = v.repeat_interleave(self.repeat_kv, dim=2)
+
+        q = apply_rotary_emb(q, cos, sin)
+        k = apply_rotary_emb(k, cos, sin)
+
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+
+        if FLASH_ATTENTION_2:
+            out = flash_attn_func(q, k, v, causal=True)
+        else:
+            out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+
+        out = out.transpose(1, 2).contiguous().view(B, T, -1)
+        return self.wo(out)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, multiple_of, mult):
+        super().__init__()
+        hidden = multiple_of * ((int(dim * mult) + multiple_of - 1) // multiple_of)
+        self.w1 = nn.Linear(dim, hidden, bias=False)
+        self.w2 = nn.Linear(hidden, dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden, bias=False)
+
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.attn = Attention(args)
+        self.ffn = FeedForward(args.dim, args.multiple_of, args.ffn_dim_multiplier)
+        self.attn_norm = RMSNorm(args.dim, args.norm_eps)
+        self.ffn_norm = RMSNorm(args.dim, args.norm_eps)
+        self.gradient_checkpointing = False
+
+    def forward(self, x, cos, sin):
+        x = x + self.attn(self.attn_norm(x), cos, sin)
+
+        if self.training and self.gradient_checkpointing:
+            x = x + torch.utils.checkpoint.checkpoint(
+                self._ffn, x, use_reentrant=False
+            )
+        else:
+            x = x + self.ffn(self.ffn_norm(x))
+        return x
+
+    def _ffn(self, x):
+        return self.ffn(self.ffn_norm(x))
+
+
+class Transformer(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.tok_emb = nn.Embedding(args.vocab_size, args.dim)
+        self.layers = nn.ModuleList([TransformerBlock(args) for _ in range(args.n_layers)])
+        self.norm = RMSNorm(args.dim, args.norm_eps)
+        self.out = nn.Linear(args.dim, args.vocab_size, bias=False)
+
+        cos, sin = precompute_freqs_cis(args.dim // args.n_heads, args.max_seq_len * 2)
+        self.register_buffer("cos_cached", cos, persistent=False)
+        self.register_buffer("sin_cached", sin, persistent=False)
+
+        self.apply(self._init)
+
+    def gradient_checkpointing_enable(self):
+        for layer in self.layers:
+            layer.gradient_checkpointing = True
+        print("[OK] Gradient checkpointing enabled")
+
+    def _init(self, m):
+        if isinstance(m, (nn.Linear, nn.Embedding)):
+            nn.init.normal_(m.weight, std=0.02)
+
+    def forward(self, tokens):
+        B, T = tokens.shape
+        h = self.tok_emb(tokens)
+        cos = self.cos_cached[:T]
+        sin = self.sin_cached[:T]
+
+        for layer in self.layers:
+            h = layer(h, cos, sin)
+
+        h = self.norm(h)
+        return self.out(h)
+
+    def get_num_params(self):
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+
+# =========================
+# MEMMAP DATASET (FIXED)
+# =========================
+class MemmapDataset(Dataset):
+    def __init__(self, path: str, max_seq_len: int, stride: Optional[int] = None):
+        self.tokens = np.memmap(path, dtype=np.int32, mode="r")
+        self.max_seq_len = max_seq_len
+        self.stride = stride or max_seq_len // 2
+
+        max_start = len(self.tokens) - (max_seq_len + 1)
+        if max_start <= 0:
+            raise ValueError("Dataset too small for the given max_seq_len")
+
+        self.starts = list(range(0, max_start, self.stride))
+        if self.starts[-1] != max_start:
+            self.starts.append(max_start)
+
+    def __len__(self):
+        return len(self.starts)
+
+    def __getitem__(self, idx):
+        i = self.starts[idx]
+        seq = torch.from_numpy(
+            self.tokens[i:i + self.max_seq_len + 1].copy()
+        ).long()
+        return seq[:-1], seq[1:]
+
+# =========================
+# TEXT GENERATION
+# =========================
+@torch.no_grad()
+def generate_text(model, tokenizer, prompts,
+                  max_new_tokens=128, temperature=0.8, top_p=0.95, eos_id=1):
+    model.eval()
+    device = next(model.parameters()).device
+    results = {}
+
+    for prompt in prompts:
+        ids = tokenizer.encode(prompt)
+        x = torch.tensor([ids], device=device)
+
+        for _ in range(max_new_tokens):
+            logits = model(x)[0, -1] / temperature
+            sorted_logits, sorted_idx = torch.sort(logits, descending=True)
+            probs = torch.softmax(sorted_logits, dim=0)
+
+            cum_probs = probs.cumsum(dim=0)
+            mask = cum_probs > top_p
+            mask[1:] = mask[:-1].clone()
+            mask[0] = False
+
+            logits[sorted_idx[mask]] = -float("inf")
+            probs = torch.softmax(logits, dim=0)
+
+            next_tok = torch.multinomial(probs, 1)
+            x = torch.cat([x, next_tok.unsqueeze(0)], dim=1)
+
+            if next_tok.item() == eos_id:
+                break
+
+        results[prompt] = tokenizer.decode(x[0].tolist())
+
+    return results
+
+# =========================
+# TRAINING
+# =========================
+def train(
+    model,
+    train_ds,
+    valid_ds,
+    tokenizer,
+    args,
+    batch_size=1,
+    grad_accum=8,
+    epochs=1,
+    lr=1e-5,
+    warmup_steps=0,
+):
+    accelerator = Accelerator(
+        mixed_precision="bf16" if torch.cuda.is_bf16_supported() else "fp16",
+        gradient_accumulation_steps=grad_accum,
+    )
+
+    model.gradient_checkpointing_enable()
+
+    train_loader = DataLoader(
+        train_ds,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=2,
+        pin_memory=True,
+    )
+
+    valid_loader = DataLoader(
+        valid_ds,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=2,
+        pin_memory=True,
+    )
+
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        lr=lr,
+        betas=(0.9, 0.95),
+        weight_decay=0.01,
+    )
+
+    total_steps = math.ceil(len(train_loader) / grad_accum) * epochs
+
+    def lr_lambda(step):
+        if step < warmup_steps:
+            return step / max(1, warmup_steps)
+        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
+        return 0.5 * (1.0 + math.cos(math.pi * progress))
+
+    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+
+    model, optimizer, train_loader, valid_loader, scheduler = accelerator.prepare(
+        model, optimizer, train_loader, valid_loader, scheduler
+    )
+
+    if accelerator.is_main_process:
+        eff_bs = batch_size * grad_accum * accelerator.num_processes
+        print(f"Model params: {model.get_num_params():,}")
+        print(f"Effective batch size: {eff_bs}")
+        print(f"Total optimizer steps: {total_steps}")
+        print(f"Flash Attention: {FLASH_ATTENTION_2}")
+        print("-" * 60)
+
+    global_step = 0
+    best_val = float("inf")
+
+    for epoch in range(epochs):
+        model.train()
+        running_loss = 0.0
+
+        pbar = tqdm(
+            train_loader,
+            disable=not accelerator.is_local_main_process,
+            desc=f"Epoch {epoch+1}/{epochs}",
+        )
+
+        for step, (x, y) in enumerate(pbar):
+            with accelerator.accumulate(model):
+                logits = model(x)
+                loss = F.cross_entropy(
+                    logits.view(-1, logits.size(-1)),
+                    y.view(-1),
+                    ignore_index=tokenizer.pad_id(),
+                )
+                accelerator.backward(loss)
+
+                if accelerator.sync_gradients:
+                    accelerator.clip_grad_norm_(model.parameters(), 1.0)
+                    optimizer.step()
+                    scheduler.step()
+                    optimizer.zero_grad()
+
+            # ======== global_step podle training steps (batchů) ========
+            global_step += 1
+
+            # ==========================================
+            # PERIODIC CHECKPOINT + TEXT GENERATION
+            # ==========================================
+            if accelerator.is_main_process and global_step % SAVE_EVERY_STEPS == 0:
+                ckpt_path = f"{CHECKPOINT_DIR}/step_{global_step}.pt"
+                checkpoint = {
+                    "step": global_step,
+                    "model_state_dict": accelerator.unwrap_model(model).state_dict(),
+                    "optimizer_state_dict": optimizer.state_dict(),
+                    "scheduler_state_dict": scheduler.state_dict(),
+                    "model_args": args,
+                }
+                torch.save(checkpoint, ckpt_path)
+                print(f"[Checkpoint] Saved complete checkpoint at step {global_step}")
+
+                prompts = [
+                    "Once upon a time",
+                    "In a distant future",
+                    "First step to build a rocket",
+                    "Capital city of France",
+                    "Artificial intelligence will",
+                ]
+
+                samples = generate_text(
+                    accelerator.unwrap_model(model),
+                    tokenizer,
+                    prompts,
+                    max_new_tokens=100,
+                    temperature=0.8,
+                    top_p=0.95,
+                )
+
+                print(f"[Sample generation @ step {global_step}]")
+                for prompt, text in samples.items():
+                    print(f"Prompt: {prompt}")
+                    print(f"Generated: {text}")
+                    print("-" * 50)
+
+            running_loss += loss.item()
+            pbar.set_postfix(
+                loss=f"{running_loss/(step+1):.4f}",
+                lr=f"{scheduler.get_last_lr()[0]:.2e}",
+            )
+
+        # =========================
+        # VALIDATION
+        # =========================
+        model.eval()
+        val_loss = 0.0
+        with torch.no_grad():
+            for x, y in valid_loader:
+                logits = model(x)
+                loss = F.cross_entropy(
+                    logits.view(-1, logits.size(-1)),
+                    y.view(-1),
+                    ignore_index=tokenizer.pad_id(),
+                )
+                val_loss += loss.item()
+
+        val_loss /= len(valid_loader)
+
+        accelerator.print(
+            f"[Epoch {epoch+1}] Train Loss: {running_loss/len(train_loader):.6f} | "
+            f"Val Loss: {val_loss:.6f}"
+        )
+
+        # =========================
+        # END-OF-EPOCH GENERATION
+        # =========================
+        if accelerator.is_main_process:
+            prompts = [
+                "Once upon a time",
+                "In a distant future",
+                "First step to build a rocket",
+                "Capital city of France",
+                "Artificial intelligence will",
+            ]
+
+            samples = generate_text(
+                accelerator.unwrap_model(model),
+                tokenizer,
+                prompts,
+                max_new_tokens=100,
+                temperature=0.8,
+                top_p=0.95,
+            )
+
+            print("[Sample generation]")
+            for prompt, text in samples.items():
+                print(f"Prompt: {prompt}")
+                print(f"Generated: {text}")
+                print("-" * 50)
+
+    # =========================
+    # FINAL SAVE
+    # =========================
+    if accelerator.is_main_process:
+        checkpoint = {
+            "step": global_step,
+            "model_state_dict": accelerator.unwrap_model(model).state_dict(),
+            "optimizer_state_dict": optimizer.state_dict(),
+            "scheduler_state_dict": scheduler.state_dict(),
+            "model_args": args,
+        }
+        torch.save(checkpoint, f"{CHECKPOINT_DIR}/final_model.pt")
+        print("Training complete.")
+
+
+
+# =========================
+# MAIN
+# =========================
+if __name__ == "__main__":
+    args = ModelArgs()
+
+    tokenizer = spm.SentencePieceProcessor(model_file=TOKENIZER_MODEL_PATH)
+    args.vocab_size = tokenizer.vocab_size()
+
+    train_ds = MemmapDataset(TRAIN_BIN, args.max_seq_len)
+    valid_ds = MemmapDataset(VALID_BIN, args.max_seq_len)
+
+    model = Transformer(args)
+
+    RESUME_FROM = "checkpoints/step_200000.pt"
+
+    if os.path.exists(RESUME_FROM):
+        print(f"[Resume] Loading checkpoint from {RESUME_FROM}")
+        checkpoint = torch.load(RESUME_FROM, map_location="cpu")
+        
+        # Support both old format (direct state_dict) and new format (checkpoint dict)
+        if "model_state_dict" in checkpoint:
+            model.load_state_dict(checkpoint["model_state_dict"])
+            print(f"[Resume] Loaded model from step {checkpoint.get('step', 'unknown')}")
+        else:
+            # Old format: checkpoint is directly the model state_dict
+            model.load_state_dict(checkpoint)
+            print(f"[Resume] Loaded model (old format)")
+
+    train(
+        model,
+        train_ds,
+        valid_ds,
+        tokenizer,
+        args,
+        batch_size=1,
+        grad_accum=8,
+        epochs=1,
+        lr=1e-5,
+    )

valid.bin

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