Upload Slayer GPT tokenizer model archive

.gitattributes

@@ -1,35 +1,5 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+tokenizers/*.json filter=lfs diff=lfs merge=lfs -text
+metadata/*.csv filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+.DS_Store
+.venv/
+venv/
+*.pyc
+

README.md

@@ -0,0 +1,101 @@
+---
+license: other
+library_name: pytorch
+tags:
+  - pytorch
+  - gpt
+  - gpt2-style
+  - polish
+  - tokenizer
+  - byte-level-bpe
+  - causal-lm
+language:
+  - pl
+pipeline_tag: text-generation
+---
+
+# Slayer GPT Tokenizer Model
+
+Teaching archive for the Slayer GPT-style Polish language-model experiment.
+
+This repo is meant to help people replicate the workflow, not just store artifacts. It includes a runnable local GPT checkpoint, the custom tokenizer it was trained with, a later tokenizer variant, the Slayer H100 training scripts, run logs, and replication docs.
+
+This is a raw PyTorch/custom-code checkpoint, not a Transformers-native `AutoModelForCausalLM` repository.
+
+## Start Here
+
+Read:
+
+- `docs/REPLICATION_GUIDE.md` - full step-by-step replication lesson.
+- `docs/TOKENIZER_NOTES.md` - tokenizer-specific teaching notes.
+- `metadata/artifact_manifest.json` - exact artifact provenance and model/tokenizer metadata.
+
+Run the saved model:
+
+```bash
+python3 -m venv .venv
+source .venv/bin/activate
+pip install -r requirements.txt
+python scripts/sample_mac.py "Polska jest" 80
+```
+
+## What Is Included
+
+- `model/ckpt.pt` - runnable nanoGPT-style checkpoint from `/Users/kacper/Local/Ventures/Slayer/gpt2-pl-mac/ckpt.pt`.
+- `tokenizers/polish_bpe_32k.json` - custom byte-level BPE tokenizer paired with `model/ckpt.pt`.
+- `tokenizers/rxlm_polish_bpe_65k.json` - separate later 65k custom tokenizer from RXLM/Slayer work.
+- `scripts/model.py` - GPT model definition for the checkpoint.
+- `scripts/sample_mac.py` - local sampler.
+- `scripts/knowbench_mac.py`, `scripts/syntaxbench_mac.py` - simple probes.
+- `examples/prepare_corpus.py` - reference corpus/tokenizer/shard preparation script.
+- `training/` - Slayer remote nanoGPT training code and launch script.
+- `logs/` and `metadata/` - run evidence.
+
+## Key Compatibility Rule
+
+Use this pairing:
+
+```text
+model/ckpt.pt -> tokenizers/polish_bpe_32k.json
+```
+
+Do not sample `model/ckpt.pt` with `tokenizers/rxlm_polish_bpe_65k.json`. That tokenizer is a separate later artifact.
+
+## Checkpoint Summary
+
+`model/ckpt.pt`:
+
+- 12 layers
+- 12 attention heads
+- 768 embedding dimension
+- 1024 token context
+- 32768 vocabulary size
+- about 136M state-dict parameters
+- checkpoint step stored as `iter=500`
+
+## Remote Slayer Training States
+
+The full remote training states are not committed because each is about 3.6 GB and includes optimizer/runtime state.
+
+Known local copies:
+
+```text
+/Users/kacper/Local/Ventures/Slayer/slayer-nanogpt/ckpt/state_step000500.pt
+/Users/kacper/Local/Ventures/Slayer/slayer-nanogpt/ckpt/state_step001000.pt
+/Users/kacper/Local/Ventures/Slayer/slayer-nanogpt/ckpt/state_step001500.pt
+```
+
+Known remote copies:
+
+```text
+ssh slayer:/home/ubuntu/modded-nanogpt/logs/f143baf2-dbfc-406e-b13b-a9fcc166b31b/state_step000500.pt
+ssh slayer:/home/ubuntu/modded-nanogpt/logs/f143baf2-dbfc-406e-b13b-a9fcc166b31b/state_step001000.pt
+ssh slayer:/home/ubuntu/modded-nanogpt/logs/f143baf2-dbfc-406e-b13b-a9fcc166b31b/state_step001500.pt
+```
+
+Fetch explicitly when needed:
+
+```bash
+mkdir -p training-states
+rsync -avP slayer:/home/ubuntu/modded-nanogpt/logs/f143baf2-dbfc-406e-b13b-a9fcc166b31b/state_step001500.pt training-states/
+```

docs/REPLICATION_GUIDE.md

@@ -0,0 +1,246 @@
+# Replication Guide
+
+This guide explains how to replicate the Slayer GPT-style Polish language-model experiment from raw text to a runnable checkpoint.
+
+The repo contains two related tracks:
+
+1. `model/ckpt.pt` is a small, runnable GPT checkpoint paired with `tokenizers/polish_bpe_32k.json`.
+2. `training/` contains the larger Slayer H100 training code derived from modded-nanoGPT. Its full optimizer checkpoints are documented but not committed.
+
+Use the small track for teaching and local demos. Use the H100 track to explain how the larger remote run was structured.
+
+## 1. What Was Built
+
+The local model is a GPT-2-style decoder-only Transformer:
+
+- 12 layers
+- 12 attention heads
+- 768 embedding dimension
+- 1024 token context
+- 32768-token custom Polish byte-level BPE vocabulary
+- bias-free linear/layernorm setup
+- about 136M parameters in the checkpoint state dict
+
+It is not a fine-tune of OpenAI GPT-2 weights. The important idea is the recipe:
+
+1. collect Polish text,
+2. train a custom byte-level BPE tokenizer,
+3. tokenize the corpus into token-id shards,
+4. train a causal language model on next-token prediction,
+5. sample and evaluate with the exact same tokenizer.
+
+## 2. Repository Map
+
+- `model/ckpt.pt` - runnable model checkpoint.
+- `tokenizers/polish_bpe_32k.json` - tokenizer paired with `model/ckpt.pt`.
+- `tokenizers/rxlm_polish_bpe_65k.json` - separate later 65k custom tokenizer. Do not use it with `model/ckpt.pt`.
+- `scripts/model.py` - GPT implementation for the local checkpoint.
+- `scripts/sample_mac.py` - generation script.
+- `scripts/knowbench_mac.py` and `scripts/syntaxbench_mac.py` - simple evaluation probes.
+- `examples/prepare_corpus.py` - reference script for tokenizer training and `.bin` shard creation.
+- `training/train_gpt.py` - H100 Slayer training script.
+- `training/run_polish.sh` - remote launch script used on `ssh slayer`.
+- `logs/` and `metadata/` - evidence from the saved local and remote runs.
+
+## 3. Local Demo
+
+```bash
+python3 -m venv .venv
+source .venv/bin/activate
+pip install -r requirements.txt
+python scripts/sample_mac.py "Polska jest" 80
+```
+
+Expected behavior:
+
+- On Apple Silicon, the sampler uses MPS.
+- On other machines, it falls back to CPU.
+- The model should load `model/ckpt.pt` and `tokenizers/polish_bpe_32k.json` without path changes.
+
+Run the lightweight probes:
+
+```bash
+python scripts/syntaxbench_mac.py
+python scripts/knowbench_mac.py
+```
+
+## 4. Corpus Preparation
+
+For teaching, start with a plain UTF-8 text corpus:
+
+```text
+data/raw/doc_0001.txt
+data/raw/doc_0002.txt
+...
+```
+
+One document per file is easiest to reason about. Clean enough to remove boilerplate and encoding damage, but do not over-normalize language. The tokenizer here is byte-level BPE, so it can represent arbitrary text.
+
+Recommended minimum for a useful class demo:
+
+- tokenizer demo: 10 MB to 100 MB of text,
+- tiny GPT training demo: 100 MB to 1 GB,
+- meaningful GPT run: many GB.
+
+Keep validation data separate before tokenization.
+
+## 5. Tokenizer Training
+
+The model checkpoint in this repo uses:
+
+- byte-level BPE,
+- vocab size 32768,
+- no Unicode normalizer,
+- `add_prefix_space=False`,
+- `<|endoftext|>` as the document separator / BOS-style token.
+
+Train a compatible tokenizer and create token-id shards with:
+
+```bash
+python examples/prepare_corpus.py \
+  --raw-dir data/raw \
+  --out-dir data/processed \
+  --vocab-size 32768 \
+  --train-tokenizer
+```
+
+This writes:
+
+- `data/processed/tokenizer.json`
+- `data/processed/shards/polish_train_000000.bin`
+- `data/processed/shards/polish_val_000000.bin`
+
+The `.bin` shards are raw `uint16` token ids. This matters because `training/train_gpt.py` expects token ids to fit under 65536 and loads shards as `torch.uint16`.
+
+For a fresh 65k experiment you may use a 65536-token tokenizer, but then the model config and training code must match that vocabulary size. Do not use the 65k tokenizer with the checked-in 32768-vocab checkpoint.
+
+## 6. Training A Small Local GPT
+
+This repo includes inference code for the saved checkpoint, not a full clean-room tiny trainer. For teaching, the simplest path is:
+
+1. Use `examples/prepare_corpus.py` to create tokenizer and shards.
+2. Use an existing nanoGPT trainer or your own minimal causal LM trainer.
+3. Match these model settings for compatibility with `scripts/model.py`:
+
+```python
+GPTConfig(
+    n_layer=12,
+    n_head=12,
+    n_embd=768,
+    block_size=1024,
+    vocab_size=32768,
+    bias=False,
+    dropout=0.0,
+)
+```
+
+4. Save checkpoints in this format:
+
+```python
+torch.save(
+    {
+        "model": model.state_dict(),
+        "model_args": config_dict,
+        "iter": step,
+    },
+    "ckpt.pt",
+)
+```
+
+Then put the checkpoint at `model/ckpt.pt`, the tokenizer at `tokenizers/polish_bpe_32k.json`, and run:
+
+```bash
+python scripts/sample_mac.py "Dawno temu w Polsce" 100
+```
+
+## 7. Replicating The Slayer H100 Run
+
+The remote run used a modified fast GPT trainer under:
+
+```text
+ssh slayer:/home/ubuntu/modded-nanogpt
+```
+
+The local copy of the relevant training files is in `training/`.
+
+Important paths from the run:
+
+```text
+~/dynaword/shards/polish_train_*.bin
+~/dynaword/shards/polish_val_*.bin
+~/modded-nanogpt/logs/f143baf2-dbfc-406e-b13b-a9fcc166b31b/state_step001500.pt
+```
+
+The H100 trainer expects:
+
+- CUDA-capable Linux host,
+- PyTorch 2.10-compatible environment,
+- Triton,
+- `kernels`,
+- token-id shards as raw `uint16`,
+- train files matching `~/dynaword/shards/polish_train_*.bin`,
+- validation files matching `~/dynaword/shards/polish_val_*.bin`.
+
+The launch script:
+
+```bash
+cd training
+sed -n '1,120p' run_polish.sh
+```
+
+The core launch pattern is:
+
+```bash
+export TORCHINDUCTOR_CACHE_DIR="$HOME/.cache/torchinductor_polish"
+export TORCHINDUCTOR_FX_GRAPH_CACHE=1
+export TORCHINDUCTOR_AUTOGRAD_CACHE=1
+cd "$HOME/modded-nanogpt"
+.venv/bin/torchrun --standalone --nproc_per_node=1 train_gpt.py
+```
+
+The saved full training states include model and optimizer state:
+
+```text
+state_step000500.pt
+state_step001000.pt
+state_step001500.pt
+```
+
+They are about 3.6 GB each and are not committed. Fetch them only when needed:
+
+```bash
+mkdir -p training-states
+rsync -avP slayer:/home/ubuntu/modded-nanogpt/logs/f143baf2-dbfc-406e-b13b-a9fcc166b31b/state_step001500.pt training-states/
+```
+
+## 8. Evaluation And Sanity Checks
+
+Always validate these before trusting a run:
+
+- Tokenizer round trip: encode and decode sample Polish text.
+- Vocabulary compatibility: checkpoint `vocab_size` equals tokenizer vocab size.
+- Loss curve: training loss should decrease smoothly, not only memorize a tiny sample.
+- Sample quality: inspect repeated n-grams and broken Unicode.
+- Validation loss: keep validation shards separate from training shards.
+
+The included metadata shows the local training loss dropping from about `10.54` at step 0 to around `4.63` at step 500, with later probe rows in `metadata/traj.csv`.
+
+## 9. Common Mistakes
+
+- Mixing tokenizer files. A model trained with `polish_bpe_32k.json` must be sampled with that tokenizer.
+- Saving only weights but losing `model_args`. The loader needs architecture parameters.
+- Tokenizing train and validation together. Split first, tokenize second.
+- Using `int32` shards with the Slayer trainer. Its loader is built around raw `uint16` token ids.
+- Treating full optimizer checkpoints as deployable model artifacts. For inference, export a model-only checkpoint when possible.
+- Teaching from the H100 script first. Start with the local checkpoint and tokenizer, then show the larger training script as the scaled version.
+
+## 10. Suggested Lesson Flow
+
+1. Show `scripts/sample_mac.py` generating from the saved model.
+2. Open `metadata/artifact_manifest.json` and explain the model/tokenizer pairing.
+3. Train a toy tokenizer on a small Polish corpus with `examples/prepare_corpus.py`.
+4. Inspect tokenization of Polish words, punctuation, and diacritics.
+5. Explain next-token prediction and the checkpoint format.
+6. Show how `training/train_gpt.py` scales the same idea to H100 training.
+7. End with failure modes: wrong tokenizer, data leakage, repeated text, and no validation split.
+

docs/TOKENIZER_NOTES.md

@@ -0,0 +1,75 @@
+# Tokenizer Notes
+
+The key teaching point is that the tokenizer is part of the model. It is not an interchangeable preprocessing detail.
+
+## Tokenizers In This Repo
+
+| File | Vocab | Use |
+| --- | ---: | --- |
+| `tokenizers/polish_bpe_32k.json` | 32768 | Paired with `model/ckpt.pt` |
+| `tokenizers/rxlm_polish_bpe_65k.json` | 65536 | Separate later custom tokenizer artifact |
+
+## 32k Polish BPE
+
+Properties:
+
+- byte-level BPE,
+- 32768 vocabulary entries,
+- 32511 merges,
+- no normalizer,
+- `add_prefix_space=False`,
+- `<|endoftext|>` is the special document separator token.
+
+This tokenizer is small enough that token ids fit safely in `uint16`, which is why the training shards can be compact raw binary files.
+
+## 65k RXLM BPE
+
+Properties:
+
+- byte-level BPE,
+- 65536 vocabulary entries,
+- 65283 merges,
+- NFKC normalization,
+- `add_prefix_space=True`,
+- 12 added tokens.
+
+This is a different tokenizer. It should be taught as a later design variant, not as the tokenizer for `model/ckpt.pt`.
+
+## Why Custom BPE
+
+For Polish, a custom tokenizer can reduce awkward fragmentation of common morphemes, diacritics, inflected forms, and domain-specific text. The lesson is not that 32k is always best; the lesson is that tokenizer choice changes:
+
+- effective context length,
+- training cost,
+- model embedding size,
+- output fluency,
+- evaluation comparability.
+
+## Quick Inspection Snippet
+
+```python
+from tokenizers import Tokenizer
+
+tok = Tokenizer.from_file("tokenizers/polish_bpe_32k.json")
+text = "Zażółć gęślą jaźń. Polska jest częścią Europy."
+enc = tok.encode(text)
+
+print(enc.ids)
+print(enc.tokens)
+print(tok.decode(enc.ids))
+```
+
+## Compatibility Rule
+
+For inference:
+
+```text
+checkpoint vocab_size == tokenizer vocab size
+```
+
+For this repo:
+
+```text
+model/ckpt.pt -> tokenizers/polish_bpe_32k.json
+```
+

examples/prepare_corpus.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+"""Train a byte-level BPE tokenizer and write uint16 token-id shards.
+
+This is a teaching/reference script for reproducing the data shape expected by
+the Slayer GPT-style training code. Input is a directory of UTF-8 .txt files.
+"""
+
+from __future__ import annotations
+
+import argparse
+import random
+from pathlib import Path
+
+import numpy as np
+from tokenizers import Tokenizer
+from tokenizers.decoders import ByteLevel as ByteLevelDecoder
+from tokenizers.models import BPE
+from tokenizers.pre_tokenizers import ByteLevel
+from tokenizers.trainers import BpeTrainer
+
+
+SPECIAL_TOKEN = "<|endoftext|>"
+
+
+def iter_text_files(raw_dir: Path) -> list[Path]:
+    files = sorted(raw_dir.rglob("*.txt"))
+    if not files:
+        raise SystemExit(f"No .txt files found under {raw_dir}")
+    return files
+
+
+def train_tokenizer(files: list[Path], out_path: Path, vocab_size: int) -> Tokenizer:
+    tokenizer = Tokenizer(BPE(unk_token=None))
+    tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False)
+    tokenizer.decoder = ByteLevelDecoder()
+
+    trainer = BpeTrainer(
+        vocab_size=vocab_size,
+        min_frequency=2,
+        special_tokens=[SPECIAL_TOKEN],
+        show_progress=True,
+    )
+    tokenizer.train([str(path) for path in files], trainer)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    tokenizer.save(str(out_path))
+    return tokenizer
+
+
+def load_tokenizer(path: Path) -> Tokenizer:
+    return Tokenizer.from_file(str(path))
+
+
+def write_shard(tokenizer: Tokenizer, files: list[Path], out_path: Path) -> int:
+    ids: list[int] = []
+    eot = tokenizer.token_to_id(SPECIAL_TOKEN)
+    if eot is None:
+        raise SystemExit(f"Tokenizer is missing {SPECIAL_TOKEN!r}")
+
+    for path in files:
+        text = path.read_text(encoding="utf-8", errors="replace").strip()
+        if not text:
+            continue
+        ids.append(eot)
+        ids.extend(tokenizer.encode(text).ids)
+
+    arr = np.asarray(ids, dtype=np.uint16)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    arr.tofile(out_path)
+    return int(arr.size)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--raw-dir", type=Path, required=True)
+    parser.add_argument("--out-dir", type=Path, required=True)
+    parser.add_argument("--vocab-size", type=int, default=32768)
+    parser.add_argument("--val-frac", type=float, default=0.01)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--train-tokenizer", action="store_true")
+    args = parser.parse_args()
+
+    if args.vocab_size > 65536:
+        raise SystemExit("This shard writer uses uint16; vocab size must be <= 65536")
+
+    files = iter_text_files(args.raw_dir)
+    random.Random(args.seed).shuffle(files)
+    val_count = max(1, int(len(files) * args.val_frac))
+    val_files = sorted(files[:val_count])
+    train_files = sorted(files[val_count:])
+
+    tokenizer_path = args.out_dir / "tokenizer.json"
+    if args.train_tokenizer or not tokenizer_path.exists():
+        tokenizer = train_tokenizer(train_files, tokenizer_path, args.vocab_size)
+    else:
+        tokenizer = load_tokenizer(tokenizer_path)
+
+    train_tokens = write_shard(
+        tokenizer,
+        train_files,
+        args.out_dir / "shards" / "polish_train_000000.bin",
+    )
+    val_tokens = write_shard(
+        tokenizer,
+        val_files,
+        args.out_dir / "shards" / "polish_val_000000.bin",
+    )
+
+    print(f"tokenizer={tokenizer_path}")
+    print(f"train_files={len(train_files)} train_tokens={train_tokens}")
+    print(f"val_files={len(val_files)} val_tokens={val_tokens}")
+
+
+if __name__ == "__main__":
+    main()
+

logs/train.out

@@ -0,0 +1,65 @@
+[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0
+
+logs/b9ed9818-4ffe-4f5a-aff8-0df5b4a150d1.txt
+Compiling model and warming up kernels (~7 minutes on first execution)
+Sampling steps [0, 1, 4398, 4399, 4400, 4401, 8798, 8799, 8800, 8801, 13198, 13199, 13200, 13201] for warmup
+W0615 01:45:46.376000 198186 torch/distributed/elastic/agent/server/api.py:739] Received 15 death signal, shutting down workers
+W0615 01:45:46.376000 198186 torch/distributed/elastic/multiprocessing/api.py:1010] Sending process 198206 closing signal SIGTERM
+W0615 01:45:46.377000 198186 torch/distributed/elastic/multiprocessing/api.py:1010] Sending process 198206 closing signal SIGTERM
+Traceback (most recent call last):
+  File "/home/ubuntu/modded-nanogpt/.venv/lib/python3.14/site-packages/torch/distributed/elastic/agent/server/api.py", line 731, in run
+    result = self._invoke_run(role)
+  File "/home/ubuntu/modded-nanogpt/.venv/lib/python3.14/site-packages/torch/distributed/elastic/agent/server/api.py", line 908, in _invoke_run
+    time.sleep(monitor_interval)
+    ~~~~~~~~~~^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/modded-nanogpt/.venv/lib/python3.14/site-packages/torch/distributed/elastic/multiprocessing/api.py", line 86, in _terminate_process_handler
+    raise SignalException(f"Process {os.getpid()} got signal: {sigval}", sigval=sigval)
+torch.distributed.elastic.multiprocessing.api.SignalException: Process 198186 got signal: 15
+
+During handling of the above exception, another exception occurred:
+
+Traceback (most recent call last):
+  File "/home/ubuntu/modded-nanogpt/.venv/bin/torchrun", line 6, in <module>
+    sys.exit(main())
+             ~~~~^^
+  File "/home/ubuntu/modded-nanogpt/.venv/lib/python3.14/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 362, in wrapper
+    return f(*args, **kwargs)
+  File "/home/ubuntu/modded-nanogpt/.venv/lib/python3.14/site-packages/torch/distributed/run.py", line 991, in main
+    run(args)
+    ~~~^^^^^^
+  File "/home/ubuntu/modded-nanogpt/.venv/lib/python3.14/site-packages/torch/distributed/run.py", line 982, in run
+    elastic_launch(
+    ~~~~~~~~~~~~~~~
+        config=config,
+        ~~~~~~~~~~~~~~
+        entrypoint=cmd,
+        ~~~~~~~~~~~~~~~
+    )(*cmd_args)
+    ~^^^^^^^^^^^
+  File "/home/ubuntu/modded-nanogpt/.venv/lib/python3.14/site-packages/torch/distributed/launcher/api.py", line 170, in __call__
+    return launch_agent(self._config, self._entrypoint, list(args))
+  File "/home/ubuntu/modded-nanogpt/.venv/lib/python3.14/site-packages/torch/distributed/launcher/api.py", line 308, in launch_agent
+    result = agent.run()
+  File "/home/ubuntu/modded-nanogpt/.venv/lib/python3.14/site-packages/torch/distributed/elastic/metrics/api.py", line 134, in wrapper
+    result = f(*args, **kwargs)
+  File "/home/ubuntu/modded-nanogpt/.venv/lib/python3.14/site-packages/torch/distributed/elastic/agent/server/api.py", line 740, in run
+    self._shutdown(e.sigval)
+    ~~~~~~~~~~~~~~^^^^^^^^^^
+  File "/home/ubuntu/modded-nanogpt/.venv/lib/python3.14/site-packages/torch/distributed/elastic/agent/server/local_elastic_agent.py", line 413, in _shutdown
+    self._pcontext.close(death_sig)
+    ~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^
+  File "/home/ubuntu/modded-nanogpt/.venv/lib/python3.14/site-packages/torch/distributed/elastic/multiprocessing/api.py", line 659, in close
+    self._close(death_sig=death_sig, timeout=timeout)
+    ~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/modded-nanogpt/.venv/lib/python3.14/site-packages/torch/distributed/elastic/multiprocessing/api.py", line 1022, in _close
+    handler.proc.wait(time_to_wait)
+    ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^
+  File "/usr/lib/python3.14/subprocess.py", line 1279, in wait
+    return self._wait(timeout=timeout)
+           ~~~~~~~~~~^^^^^^^^^^^^^^^^^
+  File "/usr/lib/python3.14/subprocess.py", line 2078, in _wait
+    time.sleep(delay)
+    ~~~~~~~~~~^^^^^^^
+  File "/home/ubuntu/modded-nanogpt/.venv/lib/python3.14/site-packages/torch/distributed/elastic/multiprocessing/api.py", line 86, in _terminate_process_handler
+    raise SignalException(f"Process {os.getpid()} got signal: {sigval}", sigval=sigval)
+torch.distributed.elastic.multiprocessing.api.SignalException: Process 198186 got signal: 1

metadata/artifact_manifest.json

@@ -0,0 +1,44 @@
+{
+  "created_from": {
+    "local_gpt_checkpoint": "/Users/kacper/Local/Ventures/Slayer/gpt2-pl-mac/ckpt.pt",
+    "local_gpt_tokenizer": "/Users/kacper/Local/Ventures/Slayer/gpt2-pl-mac/polish_bpe_32k.json",
+    "local_rxlm_tokenizer": "/Users/kacper/Local/DueDil/rxlm/beta-training/tokenizer.json",
+    "remote_rxlm_tokenizer": "slayer:/home/ubuntu/rxlm_dd/beta-training/tokenizer.json",
+    "remote_nanogpt_tree": "slayer:/home/ubuntu/modded-nanogpt"
+  },
+  "model_ckpt": {
+    "path": "model/ckpt.pt",
+    "format": "torch checkpoint",
+    "checkpoint_keys": ["model", "model_args", "iter"],
+    "model_args": {
+      "n_layer": 12,
+      "n_head": 12,
+      "n_embd": 768,
+      "block_size": 1024,
+      "bias": false,
+      "vocab_size": 32768,
+      "dropout": 0.0
+    },
+    "iter": 500,
+    "state_dict_parameters": 136071936,
+    "tokenizer": "tokenizers/polish_bpe_32k.json"
+  },
+  "tokenizers": {
+    "polish_bpe_32k": {
+      "path": "tokenizers/polish_bpe_32k.json",
+      "type": "ByteLevel BPE",
+      "vocab_size": 32768,
+      "merges": 32511,
+      "add_prefix_space": false
+    },
+    "rxlm_polish_bpe_65k": {
+      "path": "tokenizers/rxlm_polish_bpe_65k.json",
+      "type": "ByteLevel BPE",
+      "vocab_size": 65536,
+      "merges": 65283,
+      "normalizer": "NFKC",
+      "add_prefix_space": true
+    }
+  }
+}
+

metadata/loss_train.csv

@@ -0,0 +1,161 @@
+0,10.5432
+10,8.9096
+20,8.3818
+30,7.7284
+40,7.3585
+50,7.2123
+60,6.8500
+70,6.7956
+80,6.4644
+90,6.4187
+100,6.4978
+110,6.2566
+120,6.3528
+130,6.0993
+140,6.0455
+150,6.0754
+160,5.8164
+170,5.8069
+180,5.7759
+190,5.6976
+200,5.6238
+210,5.6289
+220,5.5407
+230,5.4001
+240,5.4523
+250,5.4548
+260,5.2646
+270,5.2490
+280,5.2464
+290,5.2203
+300,5.2502
+310,5.1590
+320,5.0595
+330,5.1221
+340,5.0980
+350,4.9837
+360,4.9870
+370,4.8676
+380,5.0423
+390,4.8983
+400,4.8116
+410,4.7852
+420,4.7880
+430,4.7554
+440,4.7762
+450,4.7746
+460,4.8073
+470,4.5162
+480,4.5992
+490,4.6830
+500,4.6345
+510,4.3883
+520,4.6188
+530,4.4315
+540,4.4713
+550,4.4083
+560,4.3543
+570,4.3069
+580,4.2223
+590,4.3264
+600,4.3473
+610,4.1376
+620,4.2780
+630,4.2489
+640,4.1217
+650,4.1767
+660,4.0496
+670,4.0011
+680,4.0010
+690,4.0702
+700,4.0163
+710,4.0544
+720,4.1402
+730,4.0240
+740,4.1338
+750,4.0968
+760,3.9717
+770,3.8710
+780,3.9123
+790,3.9936
+800,3.9854
+810,3.9391
+820,3.8748
+830,3.9396
+840,4.0900
+850,3.9185
+860,3.9237
+870,3.9972
+880,3.8443
+890,3.8706
+900,3.9335
+910,3.8034
+920,3.8431
+930,3.8501
+940,3.9286
+950,3.8670
+960,3.8986
+970,3.6916
+980,3.7584
+990,3.7107
+1000,3.5749
+1010,3.7844
+1020,3.8467
+1030,3.6829
+1040,3.7354
+1050,3.9265
+1060,3.7477
+1070,3.6859
+1080,3.7451
+1090,3.8840
+1100,3.7716
+1110,3.6441
+1120,3.7806
+1130,3.6817
+1140,3.7985
+1150,3.7247
+1160,3.7286
+1170,3.7495
+1180,3.7451
+1190,3.7496
+1200,3.7041
+1210,3.7436
+1220,3.5851
+1230,3.6694
+1240,3.5732
+1250,3.7169
+1260,3.7615
+1270,3.7332
+1280,3.6454
+1290,3.7745
+1300,3.5835
+1310,3.6660
+1320,3.7584
+1330,3.6219
+1340,3.6977
+1350,3.5445
+1360,3.6224
+1370,3.6865
+1380,3.6163
+1390,3.8143
+1400,3.6447
+1410,3.6732
+1420,3.5276
+1430,3.6848
+1440,3.7317
+1450,3.7915
+1460,3.6741
+1470,3.6490
+1480,3.6448
+1490,3.5571
+1500,3.6427
+1510,3.7507
+1520,3.6749
+1530,3.7123
+1540,3.7059
+1550,3.5544
+1560,3.6306
+1570,3.7105
+1580,3.7773
+1590,3.7557
+1600,3.6184

metadata/traj.csv

@@ -0,0 +1,8 @@
+iter,gram,know
+1000,90.9,28.4
+1100,90.9,32.6
+1200,90.9,25.3
+1300,90.9,26.3
+1400,90.9,30.5
+1500,90.9,28.4
+1600,90.9,27.4

model/ckpt.pt

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

requirements.txt

@@ -0,0 +1,4 @@
+torch
+tokenizers
+numpy
+

scripts/knowbench_mac.py

@@ -0,0 +1,86 @@
+"""Mini knowledge bench (likelihood MCQ) for GPT-2-PL. Run: python knowbench_mac.py [-v]
+Scores each item by length-normalized log-prob of each answer; picks argmax. No generation."""
+import os, sys, math, torch, torch.nn.functional as F
+from model import GPTConfig, GPT
+from tokenizers import Tokenizer
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+ROOT = os.path.dirname(HERE)
+VERBOSE = "-v" in sys.argv
+dev = "mps" if torch.backends.mps.is_available() else "cpu"
+
+# (prompt, correct, [distractors])  — correct is index 0; choices shuffled deterministically by pos
+ITEMS = [
+ ("Stolicą Polski jest", "Warszawa", ["Kraków", "Gdańsk", "Poznań"]),
+ ("Najdłuższą rzeką w Polsce jest", "Wisła", ["Odra", "Warta", "Bug"]),
+ ("Autorem „Pana Tadeusza” jest", "Adam Mickiewicz", ["Juliusz Słowacki", "Henryk Sienkiewicz", "Bolesław Prus"]),
+ ("Pierwszym koronowanym królem Polski był", "Bolesław Chrobry", ["Kazimierz Wielki", "Władysław Łokietek", "Jan Sobieski"]),
+ ("Polska leży w", "Europie", ["Azji", "Afryce", "Ameryce"]),
+ ("Walutą Polski jest", "złoty", ["euro", "dolar", "frank"]),
+ ("Tatry to", "góry", ["rzeka", "jezioro", "miasto"]),
+ ("Mikołaj Kopernik był", "astronomem", ["malarzem", "kompozytorem", "pisarzem"]),
+ ("Fryderyk Chopin komponował", "muzykę", ["obrazy", "powieści", "rzeźby"]),
+ ("Druga wojna światowa wybuchła w roku", "1939", ["1914", "1945", "1918"]),
+ ("Polska wstąpiła do Unii Europejskiej w roku", "2004", ["1999", "2010", "1989"]),
+ ("Słońce jest", "gwiazdą", ["planetą", "księżycem", "kometą"]),
+ ("Woda składa się z wodoru i", "tlenu", ["azotu", "węgla", "żelaza"]),
+ ("Największym oceanem jest Ocean", "Spokojny", ["Atlantycki", "Indyjski", "Arktyczny"]),
+ ("Pszczoły produkują", "miód", ["mleko", "jedwab", "wełnę"]),
+ ("Stolicą Francji jest", "Paryż", ["Londyn", "Berlin", "Madryt"]),
+ ("Wisła wpada do Morza", "Bałtyckiego", ["Czarnego", "Śródziemnego", "Czerwonego"]),
+ ("Lech Wałęsa był przywódcą", "Solidarności", ["PZPR", "Sejmu", "wojska"]),
+ ("Kraków leży nad", "Wisłą", ["Odrą", "Wartą", "Bugiem"]),
+ ("Zakopane leży w", "Tatrach", ["Bieszczadach", "Sudetach", "Karkonoszach"]),
+ ("Bursztyn powstaje z", "żywicy", ["kamienia", "metalu", "piasku"]),
+ ("Człowiek do oddychania potrzebuje", "tlenu", ["azotu", "wodoru", "helu"]),
+ ("Księżyc krąży wokół", "Ziemi", ["Słońca", "Marsa", "Wenus"]),
+ ("Orzeł biały znajduje się w godle", "Polski", ["Niemiec", "Francji", "Czech"]),
+]
+
+ck = torch.load(os.path.join(ROOT, "model", "ckpt.pt"), map_location="cpu")
+m = GPT(GPTConfig(**ck["model_args"]))
+sd = ck["model"]
+for k in list(sd):
+    if k.startswith("_orig_mod."): sd[k[len("_orig_mod."):]] = sd.pop(k)
+m.load_state_dict(sd); m.eval().to(dev)
+block = ck["model_args"]["block_size"]
+tok = Tokenizer.from_file(os.path.join(ROOT, "tokenizers", "polish_bpe_32k.json"))
+
+@torch.no_grad()
+def score(prompt, answer):
+    """length-normalized log-prob of answer tokens given prompt."""
+    pids = tok.encode(prompt).ids
+    fids = tok.encode(prompt + " " + answer).ids
+    # common prefix length (BPE boundary safety)
+    plen = 0
+    for a, b in zip(pids, fids):
+        if a != b: break
+        plen += 1
+    ans = fids[plen:]
+    if not ans: return -1e9
+    x = torch.tensor(fids, dtype=torch.long, device=dev)[None]
+    logits, _ = m(x, x)                      # targets triggers full-position logits
+    logp = F.log_softmax(logits[0].float(), dim=-1)
+    total = sum(logp[plen + i - 1, ans[i]].item() for i in range(len(ans)))
+    return total / len(ans)
+
+correct = 0
+rows = []
+for prompt, ans, dist in ITEMS:
+    cands = [ans] + dist
+    scores = [(c, score(prompt, c)) for c in cands]
+    pred = max(scores, key=lambda s: s[1])[0]
+    ok = pred == ans
+    correct += ok
+    rows.append((prompt, ans, pred, ok))
+
+n = len(ITEMS)
+base = sum(1/(1+len(d)) for _,_,d in ITEMS) / n
+print(f"\nMini-bench wiedzy PL  (likelihood MCQ, {n} pytań, 4 opcje)")
+print(f"checkpoint iter {ck.get('iter')}  ·  device {dev}")
+print("=" * 50)
+print(f"ACCURACY: {correct}/{n} = {100*correct/n:.1f}%   (losowy baseline {100*base:.1f}%)")
+if VERBOSE:
+    print("-" * 50)
+    for p, a, pred, ok in rows:
+        print(f"{'✓' if ok else '✗'}  {p} → {pred}" + ("" if ok else f"  [poprawnie: {a}]"))

scripts/model.py

@@ -0,0 +1,330 @@
+"""
+Full definition of a GPT Language Model, all of it in this single file.
+References:
+1) the official GPT-2 TensorFlow implementation released by OpenAI:
+https://github.com/openai/gpt-2/blob/master/src/model.py
+2) huggingface/transformers PyTorch implementation:
+https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py
+"""
+
+import math
+import inspect
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+class LayerNorm(nn.Module):
+    """ LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False """
+
+    def __init__(self, ndim, bias):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(ndim))
+        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None
+
+    def forward(self, input):
+        return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)
+
+class CausalSelfAttention(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        # key, query, value projections for all heads, but in a batch
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
+        # output projection
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
+        # regularization
+        self.attn_dropout = nn.Dropout(config.dropout)
+        self.resid_dropout = nn.Dropout(config.dropout)
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+        self.dropout = config.dropout
+        # flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0
+        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
+        if not self.flash:
+            print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
+            # causal mask to ensure that attention is only applied to the left in the input sequence
+            self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
+                                        .view(1, 1, config.block_size, config.block_size))
+
+    def forward(self, x):
+        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)
+
+        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
+        q, k, v  = self.c_attn(x).split(self.n_embd, dim=2)
+        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+
+        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
+        if self.flash:
+            # efficient attention using Flash Attention CUDA kernels
+            y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)
+        else:
+            # manual implementation of attention
+            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+            att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
+            att = F.softmax(att, dim=-1)
+            att = self.attn_dropout(att)
+            y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
+        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
+
+        # output projection
+        y = self.resid_dropout(self.c_proj(y))
+        return y
+
+class MLP(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
+        self.gelu    = nn.GELU()
+        self.c_proj  = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = self.gelu(x)
+        x = self.c_proj(x)
+        x = self.dropout(x)
+        return x
+
+class Block(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
+        self.attn = CausalSelfAttention(config)
+        self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)
+        self.mlp = MLP(config)
+
+    def forward(self, x):
+        x = x + self.attn(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+@dataclass
+class GPTConfig:
+    block_size: int = 1024
+    vocab_size: int = 50304 # GPT-2 vocab_size of 50257, padded up to nearest multiple of 64 for efficiency
+    n_layer: int = 12
+    n_head: int = 12
+    n_embd: int = 768
+    dropout: float = 0.0
+    bias: bool = True # True: bias in Linears and LayerNorms, like GPT-2. False: a bit better and faster
+
+class GPT(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        assert config.vocab_size is not None
+        assert config.block_size is not None
+        self.config = config
+
+        self.transformer = nn.ModuleDict(dict(
+            wte = nn.Embedding(config.vocab_size, config.n_embd),
+            wpe = nn.Embedding(config.block_size, config.n_embd),
+            drop = nn.Dropout(config.dropout),
+            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
+            ln_f = LayerNorm(config.n_embd, bias=config.bias),
+        ))
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        # with weight tying when using torch.compile() some warnings get generated:
+        # "UserWarning: functional_call was passed multiple values for tied weights.
+        # This behavior is deprecated and will be an error in future versions"
+        # not 100% sure what this is, so far seems to be harmless. TODO investigate
+        self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying
+
+        # init all weights
+        self.apply(self._init_weights)
+        # apply special scaled init to the residual projections, per GPT-2 paper
+        for pn, p in self.named_parameters():
+            if pn.endswith('c_proj.weight'):
+                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))
+
+        # report number of parameters
+        print("number of parameters: %.2fM" % (self.get_num_params()/1e6,))
+
+    def get_num_params(self, non_embedding=True):
+        """
+        Return the number of parameters in the model.
+        For non-embedding count (default), the position embeddings get subtracted.
+        The token embeddings would too, except due to the parameter sharing these
+        params are actually used as weights in the final layer, so we include them.
+        """
+        n_params = sum(p.numel() for p in self.parameters())
+        if non_embedding:
+            n_params -= self.transformer.wpe.weight.numel()
+        return n_params
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, idx, targets=None):
+        device = idx.device
+        b, t = idx.size()
+        assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
+        pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t)
+
+        # forward the GPT model itself
+        tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
+        pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
+        x = self.transformer.drop(tok_emb + pos_emb)
+        for block in self.transformer.h:
+            x = block(x)
+        x = self.transformer.ln_f(x)
+
+        if targets is not None:
+            # if we are given some desired targets also calculate the loss
+            logits = self.lm_head(x)
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+        else:
+            # inference-time mini-optimization: only forward the lm_head on the very last position
+            logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
+            loss = None
+
+        return logits, loss
+
+    def crop_block_size(self, block_size):
+        # model surgery to decrease the block size if necessary
+        # e.g. we may load the GPT2 pretrained model checkpoint (block size 1024)
+        # but want to use a smaller block size for some smaller, simpler model
+        assert block_size <= self.config.block_size
+        self.config.block_size = block_size
+        self.transformer.wpe.weight = nn.Parameter(self.transformer.wpe.weight[:block_size])
+        for block in self.transformer.h:
+            if hasattr(block.attn, 'bias'):
+                block.attn.bias = block.attn.bias[:,:,:block_size,:block_size]
+
+    @classmethod
+    def from_pretrained(cls, model_type, override_args=None):
+        assert model_type in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}
+        override_args = override_args or {} # default to empty dict
+        # only dropout can be overridden see more notes below
+        assert all(k == 'dropout' for k in override_args)
+        from transformers import GPT2LMHeadModel
+        print("loading weights from pretrained gpt: %s" % model_type)
+
+        # n_layer, n_head and n_embd are determined from model_type
+        config_args = {
+            'gpt2':         dict(n_layer=12, n_head=12, n_embd=768),  # 124M params
+            'gpt2-medium':  dict(n_layer=24, n_head=16, n_embd=1024), # 350M params
+            'gpt2-large':   dict(n_layer=36, n_head=20, n_embd=1280), # 774M params
+            'gpt2-xl':      dict(n_layer=48, n_head=25, n_embd=1600), # 1558M params
+        }[model_type]
+        print("forcing vocab_size=50257, block_size=1024, bias=True")
+        config_args['vocab_size'] = 50257 # always 50257 for GPT model checkpoints
+        config_args['block_size'] = 1024 # always 1024 for GPT model checkpoints
+        config_args['bias'] = True # always True for GPT model checkpoints
+        # we can override the dropout rate, if desired
+        if 'dropout' in override_args:
+            print(f"overriding dropout rate to {override_args['dropout']}")
+            config_args['dropout'] = override_args['dropout']
+        # create a from-scratch initialized minGPT model
+        config = GPTConfig(**config_args)
+        model = GPT(config)
+        sd = model.state_dict()
+        sd_keys = sd.keys()
+        sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')] # discard this mask / buffer, not a param
+
+        # init a huggingface/transformers model
+        model_hf = GPT2LMHeadModel.from_pretrained(model_type)
+        sd_hf = model_hf.state_dict()
+
+        # copy while ensuring all of the parameters are aligned and match in names and shapes
+        sd_keys_hf = sd_hf.keys()
+        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')] # ignore these, just a buffer
+        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')] # same, just the mask (buffer)
+        transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
+        # basically the openai checkpoints use a "Conv1D" module, but we only want to use a vanilla Linear
+        # this means that we have to transpose these weights when we import them
+        assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}"
+        for k in sd_keys_hf:
+            if any(k.endswith(w) for w in transposed):
+                # special treatment for the Conv1D weights we need to transpose
+                assert sd_hf[k].shape[::-1] == sd[k].shape
+                with torch.no_grad():
+                    sd[k].copy_(sd_hf[k].t())
+            else:
+                # vanilla copy over the other parameters
+                assert sd_hf[k].shape == sd[k].shape
+                with torch.no_grad():
+                    sd[k].copy_(sd_hf[k])
+
+        return model
+
+    def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):
+        # start with all of the candidate parameters
+        param_dict = {pn: p for pn, p in self.named_parameters()}
+        # filter out those that do not require grad
+        param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
+        # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
+        # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.
+        decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
+        nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
+        optim_groups = [
+            {'params': decay_params, 'weight_decay': weight_decay},
+            {'params': nodecay_params, 'weight_decay': 0.0}
+        ]
+        num_decay_params = sum(p.numel() for p in decay_params)
+        num_nodecay_params = sum(p.numel() for p in nodecay_params)
+        print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")
+        print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")
+        # Create AdamW optimizer and use the fused version if it is available
+        fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
+        use_fused = fused_available and device_type == 'cuda'
+        extra_args = dict(fused=True) if use_fused else dict()
+        optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args)
+        print(f"using fused AdamW: {use_fused}")
+
+        return optimizer
+
+    def estimate_mfu(self, fwdbwd_per_iter, dt):
+        """ estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS """
+        # first estimate the number of flops we do per iteration.
+        # see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311
+        N = self.get_num_params()
+        cfg = self.config
+        L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd//cfg.n_head, cfg.block_size
+        flops_per_token = 6*N + 12*L*H*Q*T
+        flops_per_fwdbwd = flops_per_token * T
+        flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter
+        # express our flops throughput as ratio of A100 bfloat16 peak flops
+        flops_achieved = flops_per_iter * (1.0/dt) # per second
+        flops_promised = 312e12 # A100 GPU bfloat16 peak flops is 312 TFLOPS
+        mfu = flops_achieved / flops_promised
+        return mfu
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
+        """
+        Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
+        the sequence max_new_tokens times, feeding the predictions back into the model each time.
+        Most likely you'll want to make sure to be in model.eval() mode of operation for this.
+        """
+        for _ in range(max_new_tokens):
+            # if the sequence context is growing too long we must crop it at block_size
+            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
+            # forward the model to get the logits for the index in the sequence
+            logits, _ = self(idx_cond)
+            # pluck the logits at the final step and scale by desired temperature
+            logits = logits[:, -1, :] / temperature
+            # optionally crop the logits to only the top k options
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float('Inf')
+            # apply softmax to convert logits to (normalized) probabilities
+            probs = F.softmax(logits, dim=-1)
+            # sample from the distribution
+            idx_next = torch.multinomial(probs, num_samples=1)
+            # append sampled index to the running sequence and continue
+            idx = torch.cat((idx, idx_next), dim=1)
+
+        return idx

scripts/sample_mac.py

@@ -0,0 +1,56 @@
+"""GPT-2-PL probe on Apple Silicon (MPS). Usage: python sample_mac.py "Prompt" [n_tok]"""
+import os, sys, time, torch, torch.nn.functional as F
+from model import GPTConfig, GPT
+from tokenizers import Tokenizer
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+ROOT = os.path.dirname(HERE)
+TEMP, TOP_K, TOP_P, REP_PEN, NGRAM = 0.7, 40, 0.92, 1.15, 3
+EOT = 0
+dev = "mps" if torch.backends.mps.is_available() else "cpu"
+
+prompt = sys.argv[1] if len(sys.argv) > 1 else "Polska jest"
+maxtok = int(sys.argv[2]) if len(sys.argv) > 2 else 90
+
+ck = torch.load(os.path.join(ROOT, "model", "ckpt.pt"), map_location="cpu")
+m = GPT(GPTConfig(**ck["model_args"]))
+sd = ck["model"]
+for k in list(sd):
+    if k.startswith("_orig_mod."):
+        sd[k[len("_orig_mod."):]] = sd.pop(k)
+m.load_state_dict(sd); m.eval().to(dev)
+block = ck["model_args"]["block_size"]
+tok = Tokenizer.from_file(os.path.join(ROOT, "tokenizers", "polish_bpe_32k.json"))
+
+def banned(seq, n):
+    if len(seq) < n - 1: return set()
+    pre = tuple(seq[-(n-1):]); bad = set()
+    for i in range(len(seq)-n+1):
+        if tuple(seq[i:i+n-1]) == pre: bad.add(seq[i+n-1])
+    return bad
+
+@torch.no_grad()
+def gen(prompt, maxtok):
+    idx = torch.tensor(tok.encode(prompt).ids, dtype=torch.long, device=dev)[None]
+    t0 = time.time(); n = 0
+    for _ in range(maxtok):
+        logits, _ = m(idx[:, -block:]); logits = logits[:, -1, :].float()
+        for t in set(idx[0].tolist()):
+            logits[0, t] /= REP_PEN if logits[0, t] > 0 else 1/REP_PEN
+        for t in banned(idx[0].tolist(), NGRAM):
+            logits[0, t] = -float("inf")
+        logits /= TEMP
+        kth = torch.topk(logits, TOP_K)[0][..., -1, None]; logits[logits < kth] = -float("inf")
+        sl, si = torch.sort(logits, descending=True)
+        cum = torch.cumsum(F.softmax(sl, dim=-1), dim=-1)
+        rm = cum > TOP_P; rm[..., 1:] = rm[..., :-1].clone(); rm[..., 0] = False
+        logits[0, si[0][rm[0]]] = -float("inf")
+        nxt = torch.multinomial(F.softmax(logits, dim=-1), 1); n += 1
+        if nxt.item() == EOT: break
+        idx = torch.cat([idx, nxt], dim=1)
+    dt = time.time() - t0
+    return tok.decode(idx[0].tolist()), n/dt
+
+txt, tps = gen(prompt, maxtok)
+print(f"[device={dev}  {tps:.1f} tok/s]\n")
+print(txt)

scripts/syntaxbench_mac.py

@@ -0,0 +1,70 @@
+"""Mini SYNTAX bench (grammatical minimal pairs, BLiMP-style) for GPT-2-PL.
+For each pair, model should give higher total log-prob to the grammatical sentence.
+Run: python syntaxbench_mac.py [-v]"""
+import os, sys, torch, torch.nn.functional as F
+from model import GPTConfig, GPT
+from tokenizers import Tokenizer
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+ROOT = os.path.dirname(HERE)
+VERBOSE = "-v" in sys.argv
+dev = "mps" if torch.backends.mps.is_available() else "cpu"
+EOT = 0
+
+# (grammatical, ungrammatical) — differ by one agreement/case/conjugation error
+PAIRS = [
+ ("Czerwony samochód stoi przed domem.", "Czerwona samochód stoi przed domem."),
+ ("Mały pies głośno szczeka.", "Mała pies głośno szczeka."),
+ ("Idę do szkoły.", "Idę do szkoła."),
+ ("Mieszkam w Warszawie.", "Mieszkam w Warszawa."),
+ ("Dzieci bawią się w parku.", "Dzieci bawi się w parku."),
+ ("Kot śpi na kanapie.", "Kot śpią na kanapie."),
+ ("Ona poszła do domu.", "Ona poszedł do domu."),
+ ("On czytał książkę.", "On czytała książkę."),
+ ("Widzę dużego psa.", "Widzę dużego pies."),
+ ("Lubię mocną kawę.", "Lubię mocną kawa."),
+ ("Duże domy stoją w mieście.", "Duży domy stoją w mieście."),
+ ("Piszę długopisem.", "Piszę długopis."),
+ ("Nie mam czasu.", "Nie mam czas."),
+ ("Nie lubię gorzkiej herbaty.", "Nie lubię gorzkiej herbata."),
+ ("Trzy koty śpią na dworze.", "Trzy kot śpią na dworze."),
+ ("Rozmawiam z bratem.", "Rozmawiam z brat."),
+ ("To jest moja siostra.", "To jest mój siostra."),
+ ("To jest mój brat.", "To jest moja brat."),
+ ("Ja idę do domu.", "Ja idzie do domu."),
+ ("Ty masz rację.", "Ty ma rację."),
+ ("Wczoraj byłem w kinie.", "Wczoraj byłem w kino."),
+ ("Ten wysoki mężczyzna śpiewa.", "Ten wysoka mężczyzna śpiewa."),
+]
+
+ck = torch.load(os.path.join(ROOT, "model", "ckpt.pt"), map_location="cpu")
+m = GPT(GPTConfig(**ck["model_args"]))
+sd = ck["model"]
+for k in list(sd):
+    if k.startswith("_orig_mod."): sd[k[len("_orig_mod."):]] = sd.pop(k)
+m.load_state_dict(sd); m.eval().to(dev)
+tok = Tokenizer.from_file(os.path.join(ROOT, "tokenizers", "polish_bpe_32k.json"))
+
+@torch.no_grad()
+def logprob(sentence):
+    ids = [EOT] + tok.encode(sentence).ids
+    x = torch.tensor(ids, dtype=torch.long, device=dev)[None]
+    logits, _ = m(x, x)
+    lp = F.log_softmax(logits[0].float(), dim=-1)
+    return sum(lp[i, ids[i+1]].item() for i in range(len(ids)-1))
+
+correct = 0; rows = []
+for good, bad in PAIRS:
+    lg, lb = logprob(good), logprob(bad)
+    ok = lg > lb; correct += ok
+    rows.append((good, bad, ok))
+
+n = len(PAIRS)
+print(f"\nMini-bench SKŁADNI PL  (pary minimalne, {n} par)")
+print(f"checkpoint iter {ck.get('iter')}  ·  device {dev}")
+print("=" * 50)
+print(f"ACCURACY: {correct}/{n} = {100*correct/n:.1f}%   (losowy baseline 50.0%)")
+if VERBOSE:
+    print("-" * 50)
+    for good, bad, ok in rows:
+        print(f"{'✓' if ok else '✗'}  {good}   vs   {bad}")

scripts/upload_to_hf.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python3
+"""Upload this archive to Hugging Face Hub.
+
+Requires an HF token with write access to the target namespace.
+Default target: SlayerLab/slayer-gpt-tokenizer-model
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from huggingface_hub import HfApi, create_repo
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--repo-id", default="SlayerLab/slayer-gpt-tokenizer-model")
+    parser.add_argument("--private", action="store_true")
+    parser.add_argument("--repo-type", default="model", choices=["model", "dataset", "space"])
+    args = parser.parse_args()
+
+    root = Path(__file__).resolve().parents[1]
+    create_repo(
+        args.repo_id,
+        repo_type=args.repo_type,
+        private=args.private,
+        exist_ok=True,
+    )
+
+    commit = HfApi().upload_folder(
+        repo_id=args.repo_id,
+        repo_type=args.repo_type,
+        folder_path=str(root),
+        commit_message="Upload Slayer GPT tokenizer model archive",
+        ignore_patterns=[
+            ".git/*",
+            ".venv/*",
+            "__pycache__/*",
+            "**/__pycache__/*",
+            ".DS_Store",
+        ],
+    )
+    print(commit)
+
+
+if __name__ == "__main__":
+    main()
+

training/polish_nanogpt.patch

@@ -0,0 +1,90 @@
+diff --git a/train_gpt.py b/train_gpt.py
+index a9b3f62..9ae8e7a 100644
+--- a/train_gpt.py
++++ b/train_gpt.py
+@@ -1494,8 +1494,8 @@ def _load_data_shard(file: Path):
+         assert nbytes == 2 * num_tokens, "number of tokens read does not match header"
+     return tokens
+ 
+-BOS_ID = 50256
+-TRAIN_MAX_NUM_DOCS = {16384: 64, 32768: 96, 49152: 128}
++BOS_ID = 0  # Polish BPE <|endoftext|>
++TRAIN_MAX_NUM_DOCS = {16384: 384, 32768: 768, 49152: 1152}  # bumped: dense short Polish docs
+ 
+ class Shard:
+     def __init__(self, tokens: Tensor, world_size: int = 1):
+@@ -1602,7 +1602,7 @@ def distributed_data_generator(filename_pattern: str, num_tokens: int, max_seq_l
+ 
+     while True:
+         num_tokens_local = num_tokens // world_size
+-        max_num_docs = TRAIN_MAX_NUM_DOCS.get(num_tokens_local, next_multiple_of_n(num_tokens_local // 300, n=128))
++        max_num_docs = TRAIN_MAX_NUM_DOCS.get(num_tokens_local, next_multiple_of_n(num_tokens_local // 48, n=128))
+ 
+         if align_to_bos:
+             try:
+@@ -1669,13 +1669,13 @@ def distributed_data_generator(filename_pattern: str, num_tokens: int, max_seq_l
+ class Hyperparameters:
+     # data
+     data_path = os.environ.get("DATA_PATH", ".")
+-    train_files: str = os.path.join(data_path, "data/fineweb10B/fineweb_train_*.bin") # input .bin to train on
+-    val_files: str = os.path.join(data_path, "data/fineweb10B/fineweb_val_*.bin") # input .bin to eval validation loss on
++    train_files: str = os.path.expanduser("~/dynaword/shards/polish_train_*.bin") # input .bin to train on
++    val_files: str = os.path.expanduser("~/dynaword/shards/polish_val_*.bin") # input .bin to eval validation loss on
+     val_tokens: int = 10485760 # how many tokens of validation data? it's important to keep this fixed for consistent comparisons
+     # batch sizes
+     val_batch_size: int = 4 * 64 * 1024 * 8
+     # schedule
+-    num_scheduled_iterations: int = 1375  # number of steps to complete lr and ws schedule
++    num_scheduled_iterations: int = 13200  # ~1 epoch of 3.47B Polish tokens  # number of steps to complete lr and ws schedule
+     num_extension_iterations: int = 10  # number of steps to continue training at final lr and ws
+     # evaluation and logging
+     run_id: str = f"{uuid.uuid4()}"
+@@ -1684,7 +1684,8 @@ class Hyperparameters:
+     #   - (1 + m_r9) * x self-reference fuse on layer 9
+     #   - backout_lambda fully removed (slot dropped from self.scalars; absorbed into MUDD bias init)
+     val_loss_every: int = 250  # every how many steps to evaluate val loss? 0 for only at the end
+-    save_checkpoint: bool = False
++    save_checkpoint: bool = True
++    checkpoint_every: int = 500  # save every N steps for crash-resume
+     run_evals: bool = False  # run additional evaluations after training is completed
+     # bigram hash embedding
+     bigram_vocab_size: int = 50304 * 15
+@@ -2014,7 +2015,7 @@ print0(nvidia_smi())
+ print0("="*100)
+ 
+ model: nn.Module = GPT(
+-    vocab_size=50257,
++    vocab_size=32896,  # mult of 128, not power-of-2 (Karpathy); tokenizer stays 32768
+     num_layers=11,
+     num_heads=6,
+     head_dim=128,
+@@ -2118,12 +2119,11 @@ for step in range(train_steps + 1):
+         torch.cuda.synchronize()
+         t0 = time.perf_counter()
+ 
++    if master_process and args.save_checkpoint and (last_step or (step > 0 and step % args.checkpoint_every == 0)):
++        log = dict(step=step, code=code, model=model.state_dict(), optimizer=training_manager.get_state())
++        os.makedirs(f"logs/{run_id}", exist_ok=True)
++        torch.save(log, f"logs/{run_id}/state_step{step:06d}.pt")
+     if last_step:
+-        if master_process and args.save_checkpoint:
+-            log = dict(step=step, code=code, model=model.state_dict(), optimizer=training_manager.get_state())
+-            os.makedirs(f"logs/{run_id}", exist_ok=True)
+-            torch.save(log, f"logs/{run_id}/state_step{step:06d}.pt")
+-        # the last step only has the validation loop, so break to avoid training
+         break
+ 
+     # --------------- TRAINING SECTION -----------------
+diff --git a/triton_kernels.py b/triton_kernels.py
+index 4f377ce..6b40884 100644
+--- a/triton_kernels.py
++++ b/triton_kernels.py
+@@ -898,7 +898,7 @@ ce_fwd_bwd_kernel = torch.cuda._compile_kernel(
+     CE_KERNEL_DECLS + CE_KERNEL_SOURCE,
+     "ce_fwd_bwd_kernel",
+     compute_capability="90",
+-    cuda_include_dirs=["/usr/local/cuda/include/"],
++    cuda_include_dirs=['/home/ubuntu/modded-nanogpt/.venv/lib/python3.14/site-packages/triton/backends/nvidia/include', '/home/ubuntu/modded-nanogpt/.venv/lib/python3.14/site-packages/nvidia/cuda_runtime/include'],
+     nvcc_options=["-lineinfo", "--use_fast_math"],
+ )
+ ce_fwd_bwd_kernel.set_shared_memory_config(CE_KERNEL_VOCAB_SIZE * 2)

training/requirements.txt

@@ -0,0 +1,9 @@
+numpy
+tqdm
+torch==2.10
+huggingface-hub
+kernels
+setuptools
+typing-extensions==4.15.0
+datasets
+tiktoken

training/run_polish.sh

@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+# Single clean training launch with persistent compile cache.
+export TORCHINDUCTOR_CACHE_DIR="$HOME/.cache/torchinductor_polish"
+export TORCHINDUCTOR_FX_GRAPH_CACHE=1
+export TORCHINDUCTOR_AUTOGRAD_CACHE=1
+mkdir -p "$TORCHINDUCTOR_CACHE_DIR"
+cd "$HOME/modded-nanogpt"
+exec .venv/bin/torchrun --standalone --nproc_per_node=1 train_gpt.py

training/triton_kernels.py

@@ -0,0 +1,1007 @@
+import torch
+import triton
+import triton.language as tl
+from triton.tools.tensor_descriptor import TensorDescriptor
+
+# -----------------------------------------------------------------------------
+# Triton kernel for symmetric matrix multiplication by @byronxu99
+
+@triton.jit
+def _pid_to_block(
+    pid,
+    M,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+):
+    # Split output matrix into blocks of size (BLOCK_SIZE_M, BLOCK_SIZE_N)
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(M, BLOCK_SIZE_N)
+
+    # Map PID to a single matrix in batch
+    batch_idx = pid // (num_pid_m * num_pid_n)
+    pid = pid % (num_pid_m * num_pid_n)
+
+    # Map PID to 2D grid of blocks
+    pid_m = pid // num_pid_n
+    pid_n = pid % num_pid_n
+    pid_m, pid_n = tl.swizzle2d(pid_m, pid_n, num_pid_m, num_pid_n, GROUP_SIZE_M)
+
+    m_idx = pid_m * BLOCK_SIZE_M
+    n_idx = pid_n * BLOCK_SIZE_N
+    return batch_idx, m_idx, n_idx
+
+@triton.jit
+def XXT_kernel(
+    A_ptr, C_ptr,
+    M, K,
+    a_stride_b, a_stride_r, a_stride_c,
+    c_stride_b, c_stride_r, c_stride_c,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+    LOWER_UPPER: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    batch_idx, m_idx, n_idx = _pid_to_block(
+        pid, M, BLOCK_SIZE_M, BLOCK_SIZE_N, GROUP_SIZE_M
+    )
+
+    # Skip blocks that don't need to be computed
+    skip_block_below_diag = (LOWER_UPPER == 0) and (n_idx + BLOCK_SIZE_N <= m_idx)
+    skip_block_above_diag = (LOWER_UPPER != 0) and (m_idx + BLOCK_SIZE_M <= n_idx)
+    if skip_block_below_diag or skip_block_above_diag:
+        return
+
+    # Index into one matrix of batch
+    A_ptr += batch_idx * a_stride_b
+    C_ptr += batch_idx * c_stride_b
+
+    # Create pointer arrays for A and A.T
+    offs_m = (m_idx + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_n = (n_idx + tl.arange(0, BLOCK_SIZE_N)) % M
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    
+    # Load A blocks for C[m,n] = A[m,:] @ A[n,:].T
+    # Load A[m, k] -> shape (BM, BK)
+    a_ptrs = A_ptr + (offs_m[:, None] * a_stride_r + offs_k[None, :] * a_stride_c)
+    # Load A[n, k] -> shape (BN, BK). Transpose to get (BK, BN) for accumulation.
+    # Loading (BN, BK) is coalesced because stride_c is 1 (contiguous dim is k).
+    at_ptrs = A_ptr + (offs_n[:, None] * a_stride_r + offs_k[None, :] * a_stride_c)
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+
+    # Accumulate over blocks of K
+    for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        k_remaining = K - k * BLOCK_SIZE_K
+        a = tl.load(a_ptrs, mask=offs_k[None, :] < k_remaining, other=0.0)
+        at_temp = tl.load(at_ptrs, mask=offs_k[None, :] < k_remaining, other=0.0)
+        at = tl.trans(at_temp)
+        accumulator = tl.dot(a, at, accumulator)
+        a_ptrs += BLOCK_SIZE_K * a_stride_c
+        at_ptrs += BLOCK_SIZE_K * a_stride_c
+
+    out_dtype = C_ptr.dtype.element_ty
+    output = accumulator.to(out_dtype)
+
+    # Store block of C
+    offs_cm = m_idx + tl.arange(0, BLOCK_SIZE_M)
+    offs_cn = n_idx + tl.arange(0, BLOCK_SIZE_N)
+    c_ptrs = C_ptr + (offs_cm[:, None] * c_stride_r + offs_cn[None, :] * c_stride_c)
+    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M)
+    tl.store(c_ptrs, output, mask=c_mask)
+
+    # Store block of C mirrored across the diagonal
+    c_ptrs_t = C_ptr + (offs_cn[:, None] * c_stride_r + offs_cm[None, :] * c_stride_c)
+    c_mask_t = (offs_cn[:, None] < M) & (offs_cm[None, :] < M)
+    tl.store(c_ptrs_t, output.T, mask=c_mask_t)
+
+def XXT(A: torch.Tensor, out: torch.Tensor):
+    """
+    Launch Triton kernel to compute C = A @ A.T
+    """
+    assert A.ndim == 2 or A.ndim == 3
+    M, K = A.shape[-2:]
+    assert out.size(-2) == M, "Output matrix has incorrect shape"
+    assert out.size(-1) == M, "Output matrix has incorrect shape"
+
+    batch_size = A.size(0) if A.ndim == 3 else 1
+    input_batch_stride = A.stride(0) if A.ndim == 3 else 0
+    output_batch_stride = out.stride(0) if out.ndim == 3 else 0
+
+    # Hardcoded configs based on H100 autotuning
+    if K == 768:
+        BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 128, 128, 64
+        num_stages, num_warps = 4, 8
+    else:
+        BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 64, 128, 128
+        num_stages, num_warps = 4, 8
+
+    grid = (batch_size * triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(M, BLOCK_SIZE_N),)
+    XXT_kernel[grid](
+        A_ptr=A,
+        C_ptr=out,
+        M=M,
+        K=K,
+        a_stride_b=input_batch_stride,
+        a_stride_r=A.stride(-2),
+        a_stride_c=A.stride(-1),
+        c_stride_b=output_batch_stride,
+        c_stride_r=out.stride(-2),
+        c_stride_c=out.stride(-1),
+        BLOCK_SIZE_M=BLOCK_SIZE_M,
+        BLOCK_SIZE_N=BLOCK_SIZE_N,
+        BLOCK_SIZE_K=BLOCK_SIZE_K,
+        GROUP_SIZE_M=8,
+        LOWER_UPPER=1,
+        num_stages=num_stages,
+        num_warps=num_warps,
+    )
+    return out
+
+# -----------------------------------------------------------------------------
+# Triton kernel for X.T @ X (tall matrices)
+# Computes C = A.T @ A where A is (M, K) and output C is (K, K)
+
+@triton.jit
+def XTX_kernel(
+    A_ptr, C_ptr,
+    M, K,
+    a_stride_b, a_stride_r, a_stride_c,
+    c_stride_b, c_stride_r, c_stride_c,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+    LOWER_UPPER: tl.constexpr,
+):
+    """
+    Compute C = A.T @ A where A is (M, K) and C is (K, K).
+    This is the transpose variant of XXT for tall matrices.
+    
+    The output matrix C is symmetric, so we compute upper triangle and mirror.
+    We iterate over blocks of M (the reduction dimension after transpose).
+    """
+    pid = tl.program_id(axis=0)
+    # Note: Output is (K, K), so we use K for the output grid
+    batch_idx, k_idx, n_idx = _pid_to_block(
+        pid, K, BLOCK_SIZE_M, BLOCK_SIZE_N, GROUP_SIZE_M
+    )
+
+    # Skip blocks that don't need to be computed (symmetry optimization)
+    skip_block_below_diag = (LOWER_UPPER == 0) and (n_idx + BLOCK_SIZE_N <= k_idx)
+    skip_block_above_diag = (LOWER_UPPER != 0) and (k_idx + BLOCK_SIZE_M <= n_idx)
+    if skip_block_below_diag or skip_block_above_diag:
+        return
+
+    # Index into one matrix of batch
+    A_ptr += batch_idx * a_stride_b
+    C_ptr += batch_idx * c_stride_b
+
+    # For A.T @ A:
+    # - A.T has shape (K, M), so A.T[k, m] = A[m, k]
+    # - We load blocks from columns k_idx and n_idx of A (which are rows of A.T)
+    # - We reduce over M (the shared dimension)
+    offs_k = (k_idx + tl.arange(0, BLOCK_SIZE_M)) % K  # Output row indices (columns of A)
+    offs_n = (n_idx + tl.arange(0, BLOCK_SIZE_N)) % K  # Output col indices (columns of A)
+    offs_m = tl.arange(0, BLOCK_SIZE_K)  # Reduction dimension (rows of A)
+
+    # Pointers for loading A[:, k_idx:k_idx+BLOCK] (transposed view is A.T[k_idx:, :])
+    # at_ptrs loads A.T block: A.T[offs_k, offs_m] = A[offs_m, offs_k]
+    at_ptrs = A_ptr + (offs_m[:, None] * a_stride_r + offs_k[None, :] * a_stride_c)
+    # a_ptrs loads A block for the other factor: A.T[offs_m, offs_n].T = A[offs_m, offs_n]
+    a_ptrs = A_ptr + (offs_m[:, None] * a_stride_r + offs_n[None, :] * a_stride_c)
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+
+    # Accumulate over blocks of M (the reduction dimension)
+    for m in tl.range(0, tl.cdiv(M, BLOCK_SIZE_K)):
+        m_remaining = M - m * BLOCK_SIZE_K
+        # Load A.T[offs_k, offs_m] = A[offs_m, offs_k] -> shape (BLOCK_K, BLOCK_M)
+        at = tl.load(at_ptrs, mask=offs_m[:, None] < m_remaining, other=0.0)
+        # Load A[offs_m, offs_n] -> shape (BLOCK_K, BLOCK_N)
+        a = tl.load(a_ptrs, mask=offs_m[:, None] < m_remaining, other=0.0)
+        # C[k, n] = sum_m A.T[k, m] * A[m, n] = sum_m A[m, k] * A[m, n]
+        # at.T @ a: (BLOCK_M, BLOCK_K) @ (BLOCK_K, BLOCK_N) = (BLOCK_M, BLOCK_N)
+        accumulator = tl.dot(at.T, a, accumulator)
+        at_ptrs += BLOCK_SIZE_K * a_stride_r
+        a_ptrs += BLOCK_SIZE_K * a_stride_r
+
+    out_dtype = C_ptr.dtype.element_ty
+    output = accumulator.to(out_dtype)
+
+    # Store block of C
+    offs_ck = k_idx + tl.arange(0, BLOCK_SIZE_M)
+    offs_cn = n_idx + tl.arange(0, BLOCK_SIZE_N)
+    c_ptrs = C_ptr + (offs_ck[:, None] * c_stride_r + offs_cn[None, :] * c_stride_c)
+    c_mask = (offs_ck[:, None] < K) & (offs_cn[None, :] < K)
+    tl.store(c_ptrs, output, mask=c_mask)
+
+    # Store block of C mirrored across the diagonal (symmetry)
+    c_ptrs_t = C_ptr + (offs_cn[:, None] * c_stride_r + offs_ck[None, :] * c_stride_c)
+    c_mask_t = (offs_cn[:, None] < K) & (offs_ck[None, :] < K)
+    tl.store(c_ptrs_t, output.T, mask=c_mask_t)
+
+
+def XTX(A: torch.Tensor, out: torch.Tensor):
+    """
+    Launch Triton kernel to compute C = A.T @ A
+    
+    For tall matrices (M > K), this is more efficient than transposing
+    and using XXT because the intermediate products are smaller (K x K vs M x M).
+    
+    Args:
+        A: Input tensor of shape (M, K) or (batch, M, K)
+        out: Output tensor of shape (K, K) or (batch, K, K)
+    
+    Returns:
+        out: The same output tensor, filled with A.T @ A
+    """
+    assert A.ndim == 2 or A.ndim == 3
+    M, K = A.shape[-2:]
+    assert out.size(-2) == K, f"Output matrix has incorrect shape: expected ({K}, {K}), got {tuple(out.shape[-2:])}"
+    assert out.size(-1) == K, f"Output matrix has incorrect shape: expected ({K}, {K}), got {tuple(out.shape[-2:])}"
+
+    batch_size = A.size(0) if A.ndim == 3 else 1
+    input_batch_stride = A.stride(0) if A.ndim == 3 else 0
+    output_batch_stride = out.stride(0) if out.ndim == 3 else 0
+
+    # Hardcoded configs based on H100 autotuning
+    if K == 768:
+        BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 128, 128, 64
+        num_stages, num_warps = 4, 8
+    else:
+        BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 64, 128, 128
+        num_stages, num_warps = 4, 8
+
+    grid = (batch_size * triton.cdiv(K, BLOCK_SIZE_M) * triton.cdiv(K, BLOCK_SIZE_N),)
+    XTX_kernel[grid](
+        A_ptr=A,
+        C_ptr=out,
+        M=M,
+        K=K,
+        a_stride_b=input_batch_stride,
+        a_stride_r=A.stride(-2),
+        a_stride_c=A.stride(-1),
+        c_stride_b=output_batch_stride,
+        c_stride_r=out.stride(-2),
+        c_stride_c=out.stride(-1),
+        BLOCK_SIZE_M=BLOCK_SIZE_M,
+        BLOCK_SIZE_N=BLOCK_SIZE_N,
+        BLOCK_SIZE_K=BLOCK_SIZE_K,
+        GROUP_SIZE_M=8,
+        LOWER_UPPER=1,
+        num_stages=num_stages,
+        num_warps=num_warps,
+    )
+    return out
+
+
+@triton.jit
+def ba_plus_cAA_kernel(
+    A_ptr, C_ptr,
+    M,
+    a_stride_b, a_stride_r, a_stride_c,
+    c_stride_b, c_stride_r, c_stride_c,
+    alpha, beta,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+    LOWER_UPPER: tl.constexpr,
+):
+    # This is mostly duplicated from XXT_kernel, but also loads and adds a block of A
+    # Performance is slightly slower than XXT_kernel, so we use two separate kernels
+    pid = tl.program_id(axis=0)
+    batch_idx, m_idx, n_idx = _pid_to_block(
+        pid, M, BLOCK_SIZE_M, BLOCK_SIZE_N, GROUP_SIZE_M
+    )
+
+    # Skip blocks that don't need to be computed
+    skip_block_below_diag = (LOWER_UPPER == 0) and (n_idx + BLOCK_SIZE_N <= m_idx)
+    skip_block_above_diag = (LOWER_UPPER != 0) and (m_idx + BLOCK_SIZE_M <= n_idx)
+    if skip_block_below_diag or skip_block_above_diag:
+        return
+
+    # Index into one matrix of batch
+    A_ptr += batch_idx * a_stride_b
+    C_ptr += batch_idx * c_stride_b
+
+    # Create pointer arrays for A and A.T
+    offs_m = (m_idx + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_n = (n_idx + tl.arange(0, BLOCK_SIZE_N)) % M
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    
+    # Coalesced loads similar to XXT_kernel
+    a_ptrs = A_ptr + (offs_m[:, None] * a_stride_r + offs_k[None, :] * a_stride_c)
+    at_ptrs = A_ptr + (offs_n[:, None] * a_stride_r + offs_k[None, :] * a_stride_c)
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+
+    # Accumulate over blocks of K
+    for k in tl.range(0, tl.cdiv(M, BLOCK_SIZE_K)):
+        k_remaining = M - k * BLOCK_SIZE_K
+        a = tl.load(a_ptrs, mask=offs_k[None, :] < k_remaining, other=0.0)
+        at_temp = tl.load(at_ptrs, mask=offs_k[None, :] < k_remaining, other=0.0)
+        at = tl.trans(at_temp)
+        accumulator = tl.dot(a, at, accumulator)
+        a_ptrs += BLOCK_SIZE_K * a_stride_c
+        at_ptrs += BLOCK_SIZE_K * a_stride_c
+
+    # Load block of A to add (corresponds to the current block of C)
+    offs_am = m_idx + tl.arange(0, BLOCK_SIZE_M)
+    offs_an = n_idx + tl.arange(0, BLOCK_SIZE_N)
+    a_add_ptrs = A_ptr + (offs_am[:, None] * a_stride_r + offs_an[None, :] * a_stride_c)
+    a_add_mask = (offs_am[:, None] < M) & (offs_an[None, :] < M)
+    a_add = tl.load(a_add_ptrs, mask=a_add_mask, other=0.0).to(tl.float32)
+
+    # Apply alpha and beta
+    accumulator *= alpha
+    accumulator += a_add * beta
+
+    out_dtype = C_ptr.dtype.element_ty
+    output = accumulator.to(out_dtype)
+
+    # Store block of C
+    offs_cm = m_idx + tl.arange(0, BLOCK_SIZE_M)
+    offs_cn = n_idx + tl.arange(0, BLOCK_SIZE_N)
+    c_ptrs = C_ptr + (offs_cm[:, None] * c_stride_r + offs_cn[None, :] * c_stride_c)
+    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M)
+    tl.store(c_ptrs, output, mask=c_mask)
+
+    # Store block of C mirrored across the diagonal
+    c_ptrs_t = C_ptr + (offs_cn[:, None] * c_stride_r + offs_cm[None, :] * c_stride_c)
+    c_mask_t = (offs_cn[:, None] < M) & (offs_cm[None, :] < M)
+    tl.store(c_ptrs_t, output.T, mask=c_mask_t)
+
+def ba_plus_cAA(A: torch.Tensor, alpha: float, beta: float, out: torch.Tensor):
+    """
+    Launch Triton kernel to compute C = alpha * A @ A.T + beta * A
+    """
+    assert A.ndim == 2 or A.ndim == 3
+    M, K = A.shape[-2:]
+    assert M == K, "Input matrix must be square"
+    assert out.size(-2) == M
+    assert out.size(-1) == M
+
+    batch_size = A.size(0) if A.ndim == 3 else 1
+    input_batch_stride = A.stride(0) if A.ndim == 3 else 0
+    output_batch_stride = out.stride(0) if out.ndim == 3 else 0
+
+    # Hardcoded config based on H100 autotuning (M=768)
+    BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 128, 128, 64
+    num_stages, num_warps = 4, 8
+
+    grid = (batch_size * triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(M, BLOCK_SIZE_N),)
+    ba_plus_cAA_kernel[grid](
+        A_ptr=A,
+        C_ptr=out,
+        M=M,
+        a_stride_b=input_batch_stride,
+        a_stride_r=A.stride(-2),
+        a_stride_c=A.stride(-1),
+        c_stride_b=output_batch_stride,
+        c_stride_r=out.stride(-2),
+        c_stride_c=out.stride(-1),
+        alpha=alpha,
+        beta=beta,
+        BLOCK_SIZE_M=BLOCK_SIZE_M,
+        BLOCK_SIZE_N=BLOCK_SIZE_N,
+        BLOCK_SIZE_K=BLOCK_SIZE_K,
+        GROUP_SIZE_M=8,
+        LOWER_UPPER=1,
+        num_stages=num_stages,
+        num_warps=num_warps,
+    )
+    return out
+
+# -----------------------------------------------------------------------------
+# Triton kernel for MLP: relu(x @ W1.T)^2, by @andrewbriand, @jrauvola
+
+@triton.jit
+def linear_relu_square_kernel(a_desc, b_desc, c_desc, aux_desc,
+                                 M, N, K,
+                                 BLOCK_SIZE_M: tl.constexpr,
+                                 BLOCK_SIZE_N: tl.constexpr,
+                                 BLOCK_SIZE_K: tl.constexpr,
+                                 GROUP_SIZE_M: tl.constexpr,
+                                 NUM_SMS: tl.constexpr,
+                                 FORWARD: tl.constexpr,
+                                 ):
+    dtype = tl.bfloat16
+    start_pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
+    num_tiles = num_pid_m * num_pid_n
+
+    tile_id_c = start_pid - NUM_SMS
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+
+    for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten=True):
+        pid_m = tile_id // num_pid_n
+        pid_n = tile_id % num_pid_n
+        offs_am = pid_m * BLOCK_SIZE_M
+        offs_bn = pid_n * BLOCK_SIZE_N
+
+        accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+        for ki in range(k_tiles):
+            offs_k = ki * BLOCK_SIZE_K
+            a = a_desc.load([offs_am, offs_k])
+            b = b_desc.load([offs_bn, offs_k])
+            accumulator = tl.dot(a, b.T, accumulator)
+
+        tile_id_c += NUM_SMS
+        pid_m = tile_id // num_pid_n
+        pid_n = tile_id % num_pid_n
+        offs_am_c = pid_m * BLOCK_SIZE_M
+        offs_bn_c = pid_n * BLOCK_SIZE_N
+
+        acc = tl.reshape(accumulator, (BLOCK_SIZE_M, 2, BLOCK_SIZE_N // 2))
+        acc = tl.permute(acc, (0, 2, 1))
+        acc0, acc1 = tl.split(acc)
+
+        c0 = acc0.to(dtype)
+        if not FORWARD:
+            c0_pre = aux_desc.load([offs_am_c, offs_bn_c])
+            c0 = 2 * c0 * tl.where(c0_pre > 0, c0_pre, 0)
+
+        c_desc.store([offs_am_c, offs_bn_c], c0)
+
+        if FORWARD:
+            c0_post = tl.maximum(c0, 0)
+            c0_post = c0_post * c0_post
+            aux_desc.store([offs_am_c, offs_bn_c], c0_post)
+
+        c1 = acc1.to(dtype)
+        if not FORWARD:
+            c1_pre = aux_desc.load([offs_am_c, offs_bn_c + BLOCK_SIZE_N // 2])
+            c1 = 2 * c1 * tl.where(c1_pre > 0, c1_pre, 0)
+
+        c_desc.store([offs_am_c, offs_bn_c + BLOCK_SIZE_N // 2], c1)
+
+        if FORWARD:
+            c1_post = tl.maximum(c1, 0)
+            c1_post = c1_post * c1_post
+            aux_desc.store([offs_am_c, offs_bn_c + BLOCK_SIZE_N // 2], c1_post)
+
+
+def linear_relu_square(a, b, aux=None):
+    M, K = a.shape
+    N, K = b.shape
+    dtype = a.dtype
+
+    c = torch.empty((M, N), device=a.device, dtype=dtype)
+
+    FORWARD = False
+    if aux is None:
+        FORWARD = True
+        aux = torch.empty((M, N), device=a.device, dtype=dtype)
+
+    NUM_SMS = torch.cuda.get_device_properties("cuda").multi_processor_count
+
+    BLOCK_SIZE_M = 128
+    BLOCK_SIZE_N = 256
+    BLOCK_SIZE_K = 64
+    num_stages = 4 if FORWARD else 3
+    num_warps = 8
+
+    a_desc = TensorDescriptor.from_tensor(a, [BLOCK_SIZE_M, BLOCK_SIZE_K])
+    b_desc = TensorDescriptor.from_tensor(b, [BLOCK_SIZE_N, BLOCK_SIZE_K])
+    c_desc = TensorDescriptor.from_tensor(c, [BLOCK_SIZE_M, BLOCK_SIZE_N // 2])
+    aux_desc = TensorDescriptor.from_tensor(aux, [BLOCK_SIZE_M, BLOCK_SIZE_N // 2])
+
+    def grid(META):
+        return (min(
+            NUM_SMS,
+            triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N),
+        ), )
+
+    linear_relu_square_kernel[grid](
+        a_desc, b_desc, c_desc, aux_desc,
+        M, N, K,
+        BLOCK_SIZE_M=BLOCK_SIZE_M,
+        BLOCK_SIZE_N=BLOCK_SIZE_N,
+        BLOCK_SIZE_K=BLOCK_SIZE_K,
+        GROUP_SIZE_M=1,
+        NUM_SMS=NUM_SMS,
+        FORWARD=FORWARD,
+        num_stages=num_stages,
+        num_warps=num_warps
+    )
+
+    if FORWARD:
+        return c, aux
+    else:
+        return c
+
+class FusedLinearReLUSquareFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, W1, W2):
+        pre, post = linear_relu_square(x.view((-1, x.shape[-1])), W1)
+        x3 = post @ W2
+        ctx.save_for_backward(x, W1, W2, pre, post)
+        return x3.view(x.shape)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        x, W1, W2, pre, post = ctx.saved_tensors
+        dW2 = post.T @ grad_output
+        dpre = linear_relu_square(grad_output.view((-1, grad_output.shape[-1])), W2, aux=pre)
+        dW1 = dpre.T @ x
+        dx = dpre @ W1
+        return dx.view(x.shape), dW1, dW2
+
+
+# -----------------------------------------------------------------------------
+# Tiled transpose copy kernel: dst (N, M) = src (M, N).T
+# Uses coalesced reads from src and coalesced writes to dst via tl.trans().
+# Replaces PyTorch's elementwise copy_ which uses a naive 75k-block kernel
+# with non-coalesced writes, saturating all SMs and blocking NCCL.
+
+@triton.jit
+def _transpose_copy_kernel(
+    src_ptr, dst_ptr,
+    M, N,
+    src_stride_m, src_stride_n,
+    dst_stride_0, dst_stride_1,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    offs_m = (pid_m * BLOCK_M + tl.arange(0, BLOCK_M)).to(tl.int64)
+    offs_n = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N)).to(tl.int64)
+
+    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
+
+    # Coalesced read from src (M, N)
+    tile = tl.load(
+        src_ptr + offs_m[:, None] * src_stride_m + offs_n[None, :] * src_stride_n,
+        mask=mask, other=0.0,
+    )
+
+    # Coalesced write to dst (N, M): dst[n, m] = src[m, n]
+    mask_T = (offs_n[:, None] < N) & (offs_m[None, :] < M)
+    tl.store(
+        dst_ptr + offs_n[:, None] * dst_stride_0 + offs_m[None, :] * dst_stride_1,
+        tl.trans(tile), mask=mask_T,
+    )
+
+
+def transpose_copy(src: torch.Tensor, dst: torch.Tensor):
+    """Tiled transpose copy: dst = src.T where src is (M, N) and dst is (N, M).
+
+    Uses a 64x128 tiled Triton kernel with coalesced reads AND writes,
+    achieving near memory-bandwidth-limited performance.
+    """
+    assert src.ndim == 2 and dst.ndim == 2
+    M, N = src.shape
+    assert dst.shape == (N, M), f"Expected dst shape ({N}, {M}), got {dst.shape}"
+
+    BLOCK_M, BLOCK_N = 64, 128
+    grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
+
+    _transpose_copy_kernel[grid](
+        src, dst,
+        M, N,
+        src.stride(0), src.stride(1),
+        dst.stride(0), dst.stride(1),
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        num_warps=8,
+        num_stages=2,
+    )
+
+
+# -----------------------------------------------------------------------------
+# Tiled transpose-add kernel: dst (M, N) += src (N, M).T
+# Same tiling strategy as transpose_copy but with a fused read-add-write.
+# Replaces PyTorch's .add_(src.T) which uses the same 75k-block elementwise
+# kernel with non-coalesced reads from the transposed operand.
+
+@triton.jit
+def _transpose_add_kernel(
+    src_ptr, dst_ptr,
+    M, N,
+    src_stride_m, src_stride_n,
+    dst_stride_0, dst_stride_1,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
+
+    # Coalesced read from src (M, N)
+    src_tile = tl.load(
+        src_ptr + offs_m[:, None] * src_stride_m + offs_n[None, :] * src_stride_n,
+        mask=mask, other=0.0,
+    )
+
+    # Coalesced read-add-write on dst (N, M): dst[n, m] += src[m, n]
+    mask_T = (offs_n[:, None] < N) & (offs_m[None, :] < M)
+    dst_ptrs = dst_ptr + offs_n[:, None] * dst_stride_0 + offs_m[None, :] * dst_stride_1
+    dst_tile = tl.load(dst_ptrs, mask=mask_T, other=0.0)
+    tl.store(dst_ptrs, dst_tile + tl.trans(src_tile), mask=mask_T)
+
+
+def transpose_add(src: torch.Tensor, dst: torch.Tensor):
+    """Tiled transpose-add: dst += src.T where src is (M, N) and dst is (N, M).
+
+    Uses a 32x32 tiled Triton kernel with coalesced access on both src and dst,
+    replacing PyTorch's .add_(src.T) which has non-coalesced reads from the
+    transposed operand.
+    """
+    assert src.ndim == 2 and dst.ndim == 2
+    M, N = src.shape
+    assert dst.shape == (N, M), f"Expected dst shape ({N}, {M}), got {dst.shape}"
+
+    BLOCK_M, BLOCK_N = 32, 32
+    grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
+
+    _transpose_add_kernel[grid](
+        src, dst,
+        M, N,
+        src.stride(0), src.stride(1),
+        dst.stride(0), dst.stride(1),
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        num_warps=4,
+        num_stages=2,
+    )
+
+
+CE_KERNEL_BLOCK_SIZE = 256
+CE_KERNEL_VOCAB_SIZE = 50304
+
+CE_KERNEL_DECLS = f"""
+constexpr int VOCAB_SIZE = {CE_KERNEL_VOCAB_SIZE};
+constexpr int BLOCK_SIZE = {CE_KERNEL_BLOCK_SIZE};
+"""
+
+CE_KERNEL_SOURCE = """
+#include <cuda_bf16.h>
+#include <math_constants.h>
+
+#define __nv_fp8_e5m2 char
+#define uint16_t unsigned short
+#define uint8_t unsigned char
+#define int64_t long long
+
+__device__ __forceinline__ __nv_fp8_e5m2 f32_to_fp8_e5m2(float x) {
+    uint16_t packed;
+    asm volatile(
+        "cvt.rn.satfinite.e5m2x2.f32 %0, %1, %2;"
+        : "=h"(packed)
+        : "f"(x), "f"(0.0f)
+    );
+    __nv_fp8_e5m2 result;
+    *reinterpret_cast<uint8_t*>(&result) = (packed & (0xFF << 8)) >> 8;
+    return result;
+}
+
+struct __align__(16) __nv_bfloat168 {
+    __nv_bfloat16 data[8];
+    __device__ __nv_bfloat16& operator[](int i) { return data[i]; }
+    __device__ const __nv_bfloat16& operator[](int i) const { return data[i]; }
+};
+
+struct __align__(8) __nv_fp8_e5m28 {
+    __nv_fp8_e5m2 data[8];
+    __device__ __nv_fp8_e5m2& operator[](int i) { return data[i]; }
+    __device__ const __nv_fp8_e5m2& operator[](int i) const { return data[i]; }
+};
+
+template<typename T> __device__ constexpr T CEIL_DIV(T a, T b) { return (a + b - 1) / b; }
+
+//__device__ float sigmoid(float x) {
+//  return 1.0f / (1.0f + __expf(-x));
+//}
+__device__ float sigmoid(float x) {
+  return 0.5f + __tanhf(x * 0.5f) * 0.5f;
+}
+
+extern "C"
+__launch_bounds__(BLOCK_SIZE, 2)
+__global__ void ce_fwd_bwd_kernel(
+    const __nv_bfloat16* __restrict__ logits,
+    const int64_t* __restrict__ targets,
+    const float* __restrict__ mtp_weights,
+    float* __restrict__ losses,
+    __nv_fp8_e5m2* grad_input,
+    int batch_size,
+    int n_predict,
+    double A_param,
+    double B_param,
+    double C_param,
+    double grad_s_param,
+    double grad_scale_param)
+{
+  constexpr int VEC_WIDTH = 8;
+  constexpr int NUM_FULL_LOADS = VOCAB_SIZE / (BLOCK_SIZE * VEC_WIDTH);
+  constexpr int NUM_LOADS = CEIL_DIV(VOCAB_SIZE, BLOCK_SIZE * VEC_WIDTH);
+
+  float A = (float)A_param;
+  float B = (float)B_param;
+  float C = (float)C_param;
+  float grad_s = (float)grad_s_param;
+  float grad_scale = (float)grad_scale_param;
+
+  extern __shared__ __nv_bfloat16 smem[];
+
+  static_assert(VEC_WIDTH == 8);
+
+  const __nv_bfloat16 *block_logit_ptr = logits + VOCAB_SIZE * blockIdx.x;
+
+  float inv_C = 1 / C;
+  float B_div_C = B * inv_C;
+  float thread_max = -CUDART_INF_F;
+
+  #pragma unroll 25
+  for (int i = 0; i < NUM_LOADS; i++) {
+    int idx = i * BLOCK_SIZE * VEC_WIDTH + threadIdx.x * VEC_WIDTH;
+    if (i < NUM_FULL_LOADS || idx < VOCAB_SIZE) {
+      __nv_bfloat168 result = *(__nv_bfloat168*)(&block_logit_ptr[idx]);
+      __nv_bfloat168 result_sigmoid;
+      #pragma unroll
+      for (int k = 0; k < VEC_WIDTH; k++) {
+        float tmp = __bfloat162float(result[k]);
+        tmp = sigmoid(tmp * inv_C + B_div_C);
+        result_sigmoid[k] = __float2bfloat16(tmp);
+        tmp = A * tmp;
+        thread_max = max(tmp, thread_max);
+      }
+      *(__nv_bfloat168*)(&smem[idx]) = result_sigmoid;
+    }
+  }
+
+  constexpr int NUM_WARPS = BLOCK_SIZE / 32;
+  int warp_id = threadIdx.x / 32;
+  __shared__ float block_maxs[NUM_WARPS];
+  __shared__ float block_sums[NUM_WARPS];
+
+  for (int offset = 16; offset > 0; offset >>= 1)
+    thread_max = fmaxf(thread_max, __shfl_down_sync(0xFFFFFFFF, thread_max, offset));
+
+  if (threadIdx.x % 32 == 0) {
+    block_maxs[warp_id] = thread_max;
+  }
+
+  __syncthreads();
+
+  float block_max = -CUDART_INF_F;
+  for (int i = 0; i < NUM_WARPS; i++) {
+    block_max = fmaxf(block_max, block_maxs[i]);
+  }
+
+  float thread_sum = 0.0f;
+  #pragma unroll 2
+  for (int i = 0; i < NUM_LOADS; i++) {
+    int idx = i * BLOCK_SIZE * VEC_WIDTH + threadIdx.x * VEC_WIDTH;
+    __nv_bfloat168 l;
+    if (i < NUM_FULL_LOADS || idx < VOCAB_SIZE) {
+      l = *(__nv_bfloat168*)(&smem[idx]);
+    }
+    #pragma unroll
+    for (int k = 0; k < VEC_WIDTH; k++) {
+      float tmp = A * __bfloat162float(l[k]);
+      tmp = __expf(tmp - block_max);
+      if (i < NUM_FULL_LOADS || idx < VOCAB_SIZE) {
+        thread_sum += tmp;
+      }
+    }
+  }
+
+  for (int offset = 16; offset > 0; offset >>= 1)
+    thread_sum += __shfl_down_sync(0xFFFFFFFF, thread_sum, offset);
+
+  if (threadIdx.x % 32 == 0) {
+    block_sums[warp_id] = thread_sum;
+  }
+
+  __syncthreads();
+
+  float block_sum = 0.0f;
+  for (int i = 0; i < NUM_WARPS; i++) {
+    block_sum += block_sums[i];
+  }
+
+  float lse = block_max + __logf(block_sum);
+
+  if (threadIdx.x == 0) {
+    float total_loss = 0.0f;
+    for (int k = 0; k < n_predict; k++) {
+      int64_t target_idx = blockIdx.x + k;
+      if (target_idx < batch_size) {
+        float weight = mtp_weights[k];
+        int64_t target = targets[target_idx];
+        if (target >= 0 && target < VOCAB_SIZE) {
+          float z_target = A * __bfloat162float(smem[target]);
+          total_loss += weight * (lse - z_target);
+        }
+      }
+    }
+    losses[blockIdx.x] = total_loss;
+  }
+
+  float S_w = 0.0f;
+
+  for (int i = 0; i < n_predict; i++) {
+    S_w += mtp_weights[i];
+  }
+
+  #pragma unroll 4
+  for (int i = 0; i < NUM_LOADS; i++) {
+    int idx = i * BLOCK_SIZE * VEC_WIDTH + threadIdx.x * VEC_WIDTH;
+    __nv_fp8_e5m28 result;
+
+    if (i < NUM_FULL_LOADS || idx < VOCAB_SIZE) {
+      __nv_bfloat168 sigmoid_us = *(__nv_bfloat168*)(&smem[idx]);
+      #pragma unroll
+      for (int j = 0; j < VEC_WIDTH; j++) {
+        float sigmoid_u = __bfloat162float(sigmoid_us[j]);
+        float z = A * sigmoid_u;
+        float p = __expf(z - lse);
+
+        float term1 = S_w * p;
+        float term2 = 0.0f;
+
+        float grad_z = term1 - term2;
+        float grad_x = grad_scale * (1.0f / C * A) * (1.0f / grad_s) * grad_z * sigmoid_u * (1.0f - sigmoid_u);
+        auto result_tmp = f32_to_fp8_e5m2(grad_x);
+        result[j] = *reinterpret_cast<__nv_fp8_e5m2*>(&result_tmp);
+      }
+      *(__nv_fp8_e5m28*)(&grad_input[blockIdx.x * VOCAB_SIZE + idx]) = result;
+    }
+  }
+
+  __syncthreads();
+
+  if (threadIdx.x < n_predict && blockIdx.x + threadIdx.x < batch_size) {
+    int i = threadIdx.x;
+    int64_t target = targets[blockIdx.x + i];
+
+    float sigmoid_u = __bfloat162float(smem[target]);
+    float z = A * sigmoid_u;
+    float p = __expf(z - lse);
+
+    float term1 = S_w * p;
+    float term2 = 0.0f;
+
+    #pragma unroll
+    for (int k = 0; k < 3; k++) {
+      int64_t target_idx = blockIdx.x + k;
+      if (target_idx < batch_size && k < n_predict) {
+        if (targets[target_idx] == target) {
+          term2 += mtp_weights[k];
+        }
+      }
+    }
+
+    float grad_z = term1 - term2;
+    float grad_x = grad_scale * (1.0f / C * A) * (1.0f / grad_s) * grad_z * sigmoid_u * (1.0f - sigmoid_u);
+    auto result_tmp = f32_to_fp8_e5m2(grad_x);
+    auto result = *reinterpret_cast<__nv_fp8_e5m2*>(&result_tmp);
+    grad_input[blockIdx.x * VOCAB_SIZE + target] = result;
+  }
+}
+"""
+
+ce_fwd_bwd_kernel = torch.cuda._compile_kernel(
+    CE_KERNEL_DECLS + CE_KERNEL_SOURCE,
+    "ce_fwd_bwd_kernel",
+    compute_capability="90",
+    cuda_include_dirs=['/home/ubuntu/modded-nanogpt/.venv/lib/python3.14/site-packages/triton/backends/nvidia/include', '/home/ubuntu/modded-nanogpt/.venv/lib/python3.14/site-packages/nvidia/cuda_runtime/include'],
+    nvcc_options=["-lineinfo", "--use_fast_math"],
+)
+ce_fwd_bwd_kernel.set_shared_memory_config(CE_KERNEL_VOCAB_SIZE * 2)
+
+@torch.library.custom_op("nanogpt::ce_fwd_bwd", mutates_args={"losses", "grad_input"})
+def ce_fwd_bwd(
+    logits: torch.Tensor,
+    targets: torch.Tensor,
+    mtp_weights: torch.Tensor,
+    losses: torch.Tensor,
+    grad_input: torch.Tensor,
+    n_rows: int,
+    n_predict: int,
+    A: float,
+    B: float,
+    C: float,
+    grad_s: float,
+    grad_scale: float,
+) -> None:
+    grid = (n_rows, 1, 1)
+    ce_fwd_bwd_kernel(
+        grid,
+        (CE_KERNEL_BLOCK_SIZE, 1, 1),
+        (logits, targets, mtp_weights, losses, grad_input,
+         n_rows, n_predict, A, B, C, grad_s, grad_scale),
+        shared_mem=CE_KERNEL_VOCAB_SIZE * 2,
+    )
+
+class FusedSoftcappedCrossEntropy(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, targets, mtp_weights, lm_head_weight, x_s, w_s, grad_s, grad_scale, A=23.0, B=5.0, C=7.5):
+
+        x_f8 = x.div(x_s).to(torch.float8_e4m3fn)
+        w_f8 = lm_head_weight.div(w_s).to(torch.float8_e4m3fn)
+
+        w_f8_col_major = w_f8.T.contiguous().T
+
+        logits = torch._scaled_mm(
+            x_f8,
+            w_f8_col_major,
+            out_dtype=torch.bfloat16,
+            scale_a=x.new_tensor(x_s, dtype=torch.float32),
+            scale_b=x.new_tensor(w_s, dtype=torch.float32),
+            use_fast_accum=True,
+        )
+
+        n_rows, n_cols = logits.shape
+        if mtp_weights is None:
+             mtp_weights = torch.tensor([1.0], device=logits.device, dtype=torch.float32)
+        n_predict = mtp_weights.shape[0]
+
+        losses = torch.empty(n_rows, dtype=torch.float32, device=logits.device)
+        lse = torch.empty(n_rows, dtype=torch.float32, device=logits.device)
+
+        logits = logits.contiguous()
+        targets = targets.contiguous()
+        mtp_weights = mtp_weights.contiguous()
+
+        grad_input = torch.empty((n_rows, n_cols), dtype=torch.float8_e5m2, device=logits.device)
+
+        ce_fwd_bwd(logits, targets, mtp_weights, losses, grad_input,
+             n_rows, n_predict, A, B, C, grad_s, grad_scale)
+
+        ctx.save_for_backward(logits, targets, mtp_weights, lse, x, lm_head_weight, x_f8, w_f8, grad_input)
+        ctx.params = (A, B, C, x_s, w_s, grad_s)
+        return losses
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        logits, targets, mtp_weights, lse, x, lm_head_weight, x_f8, w_f8, grad_input = ctx.saved_tensors
+        A, B, C, x_s, w_s, grad_s = ctx.params
+        n_rows, n_cols = logits.shape
+        n_predict = mtp_weights.shape[0]
+
+        grad_output = grad_output.contiguous()
+
+        x_scale = grad_input.new_tensor(x_s, dtype=torch.float32)
+        w_scale = grad_input.new_tensor(w_s, dtype=torch.float32)
+        grad_scale = grad_input.new_tensor(grad_s, dtype=torch.float32)
+
+        grad_x = torch._scaled_mm(
+            grad_input,
+            w_f8.T,
+            out_dtype=torch.bfloat16,
+            scale_a=grad_scale,
+            scale_b=w_scale,
+            use_fast_accum=False,
+        )
+
+        x_f8_T = torch.empty((x_f8.shape[1], x_f8.shape[0]), dtype=x_f8.dtype, device=x_f8.device)
+        transpose_copy(x_f8, x_f8_T)  # (768, n_rows) row-major
+
+        grad_input_T = torch.empty((n_cols, n_rows), dtype=grad_input.dtype, device=grad_input.device)
+        transpose_copy(grad_input, grad_input_T)  # (50304, n_rows) row-major
+
+        grad_w = torch._scaled_mm(
+            x_f8_T,            # (768, n_rows) row-major
+            grad_input_T.T,    # (n_rows, 50304) column-major view
+            out_dtype=torch.float32,
+            scale_a=x_scale,
+            scale_b=grad_scale,
+            use_fast_accum=False,
+        )
+
+        return grad_x, None, None, grad_w, None, None, None
+