Blame for docs/BATCHED_TRAINING.md

# Batched byte-tick training (B streams, CPU / CUDA)

Goal: make multi-GB French training feasible on a GPU box by training **B independent byte streams in
lockstep** (B = 1..1024+), instead of the single stream / one tick per byte of `tools/run_atom_native.py`.
No Transformers, no attention, no HF, no DDP. Inference is unchanged and stays CPU-only
(`tools/chat_atom_native.py`).

Files:

| file | role |
|---|---|
| `src/atom_batched.py` | `BatchedByteTick(model, B, atom_flush_every)`: batched runtime state + batched per-tick forward/loss; `check_batched_supported()` / `BatchedUnsupported` |
| `src/io/batched_stream.py` | `ByteCorpus` (memmapped shard set = one circular byte string), `BatchedByteStream` (B offsets, vectorised Atomizer byte-tick features, resumable) |
| `tools/train_batched.py` | trainer CLI (BPTT windows, episode resets, checkpoints, val, chat samples, progress TSV, exact resume) |
| `tools/bench_batched.py` | bytes/s vs B benchmark (one subprocess per B; `--legacy` = single-stream baseline) |
| `test/test_batched_train.py` | equivalence / roundtrip / resume tests |
| `scripts/setup_gpu_box.sh`, `scripts/train_batched.sh` | GPU runbook |

`src/atom_native.py` and `tools/run_atom_native.py` are **not modified**: the old path is bit-identical by
construction and all defaults are unchanged.

## 1. Design

**No new parameters; batch lives only in runtime state.** `BatchedByteTick` wraps an ordinary
`AtomNativeModel` and uses *its* parameters (same modules, same names/shapes). It re-implements, with a
leading batch dimension, exactly the subset of `AtomNativeModel._advance` + `transition_loss` that the
byte-tick recipe executes:

* field state `z[B, n_modes]` (complex), atom buffers (`e_buf/h_buf[B, n_atoms_max, ·]`, `n_atoms[B]`),
  persistence EMA, amplitude controller;
* compile of the byte packet → drive, leak (scalar or learnable `field_leak`, graph kept across the BPTT
  window), projection, RK4 dynamics with per-stream drive, phase sync, RMS cap (`field_max_rms`);
* consolidation (energy history → persistence), per stream;
* hard readout (`--field-obligatory-hard`): JL `a_hat`, frozen + α-MLP, `--last-atom-readout` 2-gram with
  per-stream last-atom dropout;
* losses per stream (length CE, byte CE, `0.05·(1−cos)` field persistence), averaged over streams.

Atomizer features for `max_span_bytes == 1`: each tick is one byte, so packet features depend only on
(byte, small context) — `BatchedByteStream` computes them with a lookup table + vectorised context
update, **bit-exact** vs `Atomizer.encode_bytes` in both pack modes (tested).

**Lockstep schedule** (identical for all streams, identical to the legacy trainer for B=1):
episode reset when `tick % episode_length == 0`, atom flush when `tick % atom_flush_every == 0`,
BPTT window closes (backward + optimizer step + detach) at K ticks, at episode end, at checkpoint ticks and
at the last tick. `reset_state(mask)` / `flush_atoms(mask)` accept per-stream masks (tested) for future
desynchronised schedules. Gradient = mean over streams (so LR 2e-3 remains a sane starting point; with
larger B the gradient is less noisy — LR can be raised cautiously, e.g. 3e-3..5e-3 for B ≥ 256, watch val).

**Paths skipped in batched mode** (computed by the legacy path but have no effect on the hard-path loss or
gradients of the recipe): aggregation/abstraction bookkeeping, interaction energy diagnostics, α-bank
statistics. Everything that *does* affect the loss raises `BatchedUnsupported` if enabled:
`max_span_bytes > 1`, non-hard readout, merge, payload copy, efference, free-run aux, leak spectrum,
slow-mode readout / slow_block / `slow_every ≠ 1`, contrast / ignorance / printable / next-packet /
slow-memory loss weights ≠ 0, adaptive last-atom scale, hard couple > 0, `atom_flush_every < 1`.

**Data loader.** `ByteCorpus(glob)` memmaps sorted shards (np.memmap, no loading into RAM) as one
circular byte string; `BatchedByteStream` holds B offsets (evenly spaced by default, or random with
`--offset-seed`), refills `--chunk` bytes per stream, warms the Atomizer context up on 64 bytes, and yields
`(cur[B], feats[B,320], tgt[B])`. Its `state_dict` (corpus signature = shard names + sizes, positions,
context, prev_len) is saved in every checkpoint → exact resume.

**Checkpoints** (`atom_native_step_{tick}_{suffix}.pt`): `model._checkpoint()` moved to CPU (identical
format to the single-stream trainer), written atomically, never overwritten. Extra `training` block:
`step`, `optimizer`, `bytes_seen`, `batched = {B, stream offsets/ctx, runtime state (field, atoms, trimmed to
live prefix), args, torch RNG}`. Loadable unchanged by `AtomNativeModel.load`,
`tools/chat_atom_native.py` and `tools/run_atom_native.py --resume` (tested). `--init` accepts a
single-stream checkpoint to continue it batched.

**GPU friendliness.** constants cached per device, pinned `non_blocking` host→device copies, lockstep fast
paths (no per-stream Python loops in the hot path), optional `--tf32`. RNG for last-atom dropout is drawn
on CPU (`torch.rand(B)`) so CPU and CUDA runs see the same masks.

## 2. Equivalence tests (`test/test_batched_train.py`)

| test | what | result |
|---|---|---|
| features_equal_atomizer | vectorised features == `Atomizer.encode_bytes` (both pack modes) | exact |
| multishard_circular_equals_stream_source | memmapped multi-shard circular reads == `StreamingPacketSource` | exact |
| offsets_and_state_roundtrip | offsets, loader state_dict resume | exact |
| b1_matches_single_stream_bptt | B=1 vs legacy, 300 ticks, K=16 (window cuts), learnable leak, last-atom dropout 0.5, episode reset + flushes, consolidation threshold lowered to exercise persistence | max \|Δlogit\| = 0, \|Δloss\| = 0, rel Δgrad = 0 (bit-identical) |
| b1_matches_single_stream_no_bptt | same, K=0, scalar leak 0.3 | bit-identical |
| b4_equals_four_b1 | B=4 == 4 independent B=1 runs (logits, per-stream loss, grad = mean of the 4) | logits/loss ≤ 1e-5, grads rel ≤ 1e-4 |
| b4_masked_reset | per-stream masked reset/flush | matches |
| unsupported_raise | every unsupported feature raises `BatchedUnsupported` | ok |
| checkpoint_roundtrip_into_old_code | batched ckpt → `AtomNativeModel.load`, `chat_atom_native.load_model` + `generate_reply`, `run_atom_native.py --resume` (optimizer restored, same keys/shapes) | ok |
| exact_resume | run 2N ticks vs N + resume + N (mid-flush, n_atoms>0) | identical weights/loss |
| cuda_matches_cpu | CUDA vs CPU (skipped when no CUDA; run by `setup_gpu_box.sh`) | skipped here |

Full suite: **162 passed, 2 skipped, 3 failed** — see `docs/artifacts/batched_training/pytest_full.txt` (only the 3 known pre-existing failures:
`test_chat_merge_off` payload_copy, 2 in `test_speech_lock`).

## 3. CPU benchmark (this box)

`tools/bench_batched.py`, d_model 32, K=64, 1 torch thread, 64 ticks × best-of-3.
The box was heavily loaded by a concurrent training (load avg ≈ 9–10 on 8 cores), so wall time is noisy;
**bytes per CPU-second** (process_time) is the reliable metric. JSON:
`docs/artifacts/batched_training/bench_cpu_1thread.json`.

| B | bytes / CPU-s | vs legacy | peak RSS |
|---|---|---|---|
| legacy single stream | 160 | 1× | 347 MB |
| 1 | 304 | 1.9× | 345 MB |
| 8 | 1 691 | 10.6× | 424 MB |
| 32 | 3 108 | 19× | 716 MB |
| 128 | 6 540 | 41× | 1 826 MB |

Cost model (1 thread): ≈ 3.3 ms fixed per tick (Python + ~400 small ops) + ≈ 0.13 ms per stream.
On an unloaded 8-thread CPU, B=128 should reach ~15–30 K bytes/s (~55–110 MB/h) — useful for testing,
not for multi-GB.

**Memory.** BPTT activations ≈ 180 KB per stream·tick (d32) ⇒ at K=128: ≈ 23 MB per stream →
B=256 ≈ 6 GB, B=512 ≈ 12 GB, B=1024 ≈ 24 GB (+ ~1 GB overhead). Halve K to halve memory.

**GPU estimate.** At d32 each tick is ~400 tiny kernels; the GPU is launch/Python-bound at ≈ 3–5 ms per tick
nearly independent of B until B ≈ 4–8 K. So throughput ≈ B / 4 ms:
B=256 → ~50–85 K bytes/s, B=1024 → ~200–340 K bytes/s (≈ 0.7–1.2 GB/h). A 5 GB corpus epoch ≈ 4–7 h at
B=1024 on a 40–80 GB GPU (24 GB GPU: B≈768, K=128, or B=1024 with K=96). `setup_gpu_box.sh` measures the
real numbers (B=64/256/1024) into `docs/artifacts/batched_training/bench_gpu_<host>.json`.
Future speedups if needed: CUDA graphs / `torch.compile` of the tick, which remove launch overhead.

Note: the data-efficiency trade-off — B streams × fewer optimizer steps per byte (one step per K ticks
covers B·K bytes). Watch val CE per byte seen, not per step; if it lags, raise LR or lower K.

## 4. GPU runbook

### Needs on the GPU box
* Linux x86_64, NVIDIA GPU with driver (nvidia-smi), ideally ≥ 24 GB VRAM (40–80 GB for B=1024 at K=128);
* Python ≥ 3.10 (tested 3.13), torch ≥ 2.1 CUDA build matching the driver (tested code on torch 2.14 CPU);
  numpy, pytest. No HF/transformers/tokenizers;
* disk: corpus size (bytes, plain UTF-8) + checkpoints (≈ 1 MB weights+optimizer at d32, + runtime
  state ≈ 20 KB × B; B=1024 ≈ 25 MB/ckpt; 500 ckpts ≈ 12 GB) + logs. RAM: a few GB (corpus is memmapped).

### 0. Push code + data (from this box)
```bash
REMOTE=user@gpu-box; RDIR=atom-ai
# code (no venv, no checkpoints, no logs); or: git bundle create atom.bundle batched-train && scp
rsync -avz --exclude .venv --exclude checkpoints --exclude logs --exclude '__pycache__' --exclude data \
  /workspace/exports/atom-ai-batched/ $REMOTE:$RDIR/
# corpus shards (resumable), keep a held-out val shard separate
rsync -avP --partial /workspace/exports/atom-ai/data/corpus_fr_mix_00*.txt $REMOTE:$RDIR/data/fr/
rsync -avP --partial /workspace/exports/atom-ai/data/heldout/ $REMOTE:$RDIR/data/fr/val/
# optional: a single-stream checkpoint to continue from
rsync -avP <ckpt.pt> $REMOTE:$RDIR/checkpoints/init/
```
Data format: any set of plain UTF-8 text files; shards are concatenated (sorted glob) into one circular
byte stream. Don't include the val shard in `DATA_GLOB`.

### 1. Setup (once, on the GPU box)
```bash
cd ~/atom-ai && bash scripts/setup_gpu_box.sh
# if torch can't see the GPU: TORCH_INDEX=https://download.pytorch.org/whl/cu128 bash scripts/setup_gpu_box.sh
```
Checks nvidia-smi/disk/python, creates `.venv`, installs torch+numpy+pytest, asserts CUDA, runs the batched
tests (incl. CUDA==CPU) and a GPU bench (B=64/256/1024, K=128).

### 2. Launch (resumable, detached)
```bash
cd ~/atom-ai && mkdir -p logs/batched
nohup env DATA_GLOB='data/fr/corpus_fr_mix_*.txt' VAL_GLOB='data/fr/val/*.txt' B=512 DEVICE=cuda \
  TOTAL=4000000 CKPT_EVERY=8000 \
  bash scripts/train_batched.sh >> logs/batched/master.log 2>&1 < /dev/null &
# continue an existing single-stream model instead of fresh init: add INIT=checkpoints/init/<ckpt>.pt
```
Bytes consumed = TOTAL × B (B=512, 4 M ticks ≈ 2 GB). Re-running the same command resumes from the newest
checkpoint in `OUT` (default `checkpoints/batched_b$B`): weights, optimizer, RNG, stream offsets and
per-stream field state (exact resume). A flock prevents two drivers on the same `OUT`; the driver restarts
the trainer up to `RESTARTS=3` times on crash. Env reference is in the script header.

Monitor / stop:
```bash
tail -f logs/batched/master.log
column -t -s$'\t' checkpoints/batched_b512/progress.tsv | tail   # loss, byte CE, val CE (normal/field-alone), leak, bytes/s, chat samples
nvidia-smi
kill $(cat checkpoints/batched_b512/driver.pid)
```

### 3. Pull checkpoints back and chat on CPU
```bash
rsync -avP --ignore-existing $REMOTE:$RDIR/checkpoints/batched_b512/ /workspace/exports/atom-ai/checkpoints/batched_b512/
rsync -avP $REMOTE:$RDIR/checkpoints/batched_b512/progress.tsv ./
cd /workspace/exports/atom-ai && PYTHONPATH=. .venv/bin/python tools/chat_atom_native.py \
  --checkpoint checkpoints/batched_b512/atom_native_step_<tick>_batched_b512.pt --prompt "Bonjour"
```
`--ignore-existing` guarantees local checkpoints are never overwritten.

### Direct trainer use (CPU smoke test)
```bash
PYTHONPATH=. .venv/bin/python tools/train_batched.py --data-glob 'data/corpus_fr_medium.txt' -B 8 \
  --device cpu --threads 2 --output-dir /tmp/bt_smoke --total-ticks 400 --ckpt-every 200 --log-every 50
```