ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints
9f1076a verified | # 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 | |
| ``` | |