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): JLa_hat, frozen + α-MLP,--last-atom-readout2-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)
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)
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)
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:
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
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)
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