File size: 7,409 Bytes
b9301f1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | # How to train ATOM on GPU β 10B tokens on a 5090
This doc is the WTF moment: ATOM is not a transformer, so 10B tokens on a 5090 is hours, not months.
## Why ATOM is 50-100x faster than a transformer on same GPU
| | ATOM d128 1.5M | Llama 7B |
|---|---|---|
| Params | 1.5M (4666x smaller) | 7B |
| Attention | None, O(n) | Softmax O(nΒ²) + KV-cache |
| Tokenizer | None, byte-tick `max_span_bytes=1` | BPE 32k vocab |
| FLOPs / byte | ~350k (Ξ± 32x32 RK4 + Ξ±-MLP 1024β128β256 + last-atom 2-gram + dentate) | ~14B (7B * 2) |
| Theoretical max on 5090 100 TFLOPS | 285M bytes/sec | 7k tokens/sec |
| Real CPU now | 37-50 tps d32 | β |
| Real GPU naive (`.to('cuda')` same loop) | 400-800 tps (10-20x) | β |
| Real GPU batched + torch.compile | **50k-1M tps** (1000-20000x) | **~2k tps** |
**Measured ground truth CPU (this repo):**
- d32 100k: 40 bytes/sec β 1M bytes = 8h CPU
- 10B bytes CPU = 9.1 years β impossible
**10B tokens (10GB) on 5090:**
| Model | Realistic GPU tps | 10B / tps | Time |
|-------|-------------------|-----------|------|
| d32 100k | 500k tps | 20k sec | **5.5 hours** |
| d128 1.5M | 100k tps | 100k sec | **27 hours** |
| d256 6M | 40k tps | 250k sec | **2.9 days** |
| Llama 7B | 2k tps | 5M sec | **58 days** |
ATOM d128 is **50x faster** than Llama 7B for same 10B tokens, with 4666x fewer params, plus growable + CPU fine-tunable + live-surveyable.
**200GB (200B bytes) on a single normal 5090 β outshining server farms:**
| Model | 10B time | 200GB = 20x 10B | Time |
|-------|----------|-----------------|------|
| d32 100k | 5.5h | 110h | **4.5 days** |
| d128 1.5M | 27h | 540h | **22.5 days** |
| d256 6M | 2.9 days | 58 days | **1.9 months** |
| Llama 7B | 58 days | 1160 days | **3.1 years** |
**WTF:** 200GB in 1 month on a normal GPU is 100% feasible with d128 (22.5 days). Llama 7B needs 3.1 years on same single 5090, or 8x H100 cluster costing $500k.
ATOM outshines the need for server farms:
- Train 10s-100s GB ultra-fast on 1x 5090 ($2000) in your living room
- No need for 8x H100 farm for same data
- After GPU pre-train, download 6MB `.pt` (d256) to laptop, continue training on CPU for private domain in 3 min, no breaking, live survey `field_rms 0.26β1.06β3.0`
- Inference CPU 50 tps, no KV-cache explosion
- Growable d16βd32βd64βd128βd256 same lineage, old outputs exact
This is decentralization: train 200GB in 1 month on normal GPU, fine-tune on CPU.
## Current code is CPU-only on purpose
`tools/run_atom_native.py` has no `--device` flag β we banned it to measure true CPU scaling (see `TRAIN.md`: always `PYTHONPATH=. .venv/bin/python`, no `--device`). Model lives on CPU, `torch.Generator(device="cpu")`, `atom-flush-every 64` mandatory or O(nΒ²) atoms β 20Γ slowdown.
## How to enable GPU (10 lines patch)
Add to `tools/run_atom_native.py`:
```python
parser.add_argument("--device", default="cpu", choices=("cpu","cuda"))
# ...
device = torch.device(args.device)
model.to(device)
# in transition_metrics, ensure current.features.to(device)
```
And in `src/atom_native.py` `generate_packets` / `transition_loss`, move `features` to `self.device`.
We will ship this as `--device cuda` + batched mode in next cut.
## Batched mode for true speed (what we need for 10B)
Current loop: 1 Python call = 1 byte forward. Overhead kills GPU.
Batched:
```python
# pseudo
B=256
batch = [next(source) for _ in range(B)] # 256 transitions
loss = model.transition_loss_batched(batch) # one compiled forward
loss.backward()
```
With `torch.compile(model.surface)` + prefetch (`--stream-prefetch` already exists), we go from 800 tps naive to 100k-500k tps.
Steps:
1. Implement `transition_loss_batched` in `AtomNativeModel` β stack 256 `features` [B, F] β one MLP forward
2. `torch.compile` surface head
3. Keep `atom-flush-every 64` and `episode-reset --episode-length 8000`
4. Keep original recipe: `byte-tick, last-atom-readout, dentate, MERGE off, payload-copy off, hard obligatory, efference-every 10, field-loss 0.05, contrast 0, LR 5e-5/1.2e-4, seed 20260913`
## How to train 10GB on 5090 (future command)
```bash
# 1. Generate 10GB shards (seed 1-20000 = ~10GB)
PYTHONPATH=. .venv/bin/python data/gen_corpus_fr_medium.py --seed-start 1 --seed-end 20000 --out-dir data/shards_10GB
# 2. Grow d32 1080k β d128
PYTHONPATH=. .venv/bin/python tools/grow_checkpoint.py \
--checkpoint checkpoints/byte_tick/atom_native_step_1080000_d32_ms1M.pt \
--output checkpoints/byte_tick/atom_native_step_1080000_d128_ms1M.pt \
--d-model 128 --n-modes 128 --n-atoms-max 256
# 3. Train d128 on 10GB on 5090
PYTHONPATH=. .venv/bin/python tools/run_atom_native.py \
--stream --data-glob 'data/shards_10GB/*.txt' --loop-shards --chunk-bytes 1048576 \
--resume checkpoints/byte_tick/atom_native_step_1080000_d128_ms1M.pt \
--output-dir checkpoints/byte_tick --checkpoint-name atom_native_step_11000000_d128_10GB.pt \
--steps 10000000 --d-model 128 --n-modes 128 --n-atoms-max 256 \
--max-span-bytes 1 --field-obligatory-hard --no-enable-merge --no-payload-copy \
--last-atom-readout --efference-every 10 --learning-rate 5e-5 --surface-learning-rate 1.2e-4 \
--seed 20260913 --atom-flush-every 64 --episode-length 8000 --episode-reset \
--device cuda --batch-size 256 --compile
# 4. Live survey without stopping train (read-only)
PYTHONPATH=. .venv/bin/python tools/probe_field_regimes.py --checkpoint checkpoints/byte_tick/atom_native_step_11000000_d128_10GB.pt --primer-packets 3500 --device cuda
PYTHONPATH=. .venv/bin/python tools/diagnose_teacher.py --checkpoint checkpoints/byte_tick/atom_native_step_11000000_d128_10GB.pt --max-span-bytes 1 --data data/corpus_fr_medium_s21.txt --device cuda
# 5. On-demand CPU fine-tune on laptop after GPU pre-train
PYTHONPATH=. .venv/bin/python tools/run_atom_native.py \
--data data/your_private_domain.txt --steps 8000 --resume checkpoints/byte_tick/atom_native_step_11000000_d128_10GB.pt \
--output-dir checkpoints/byte_tick --checkpoint-name atom_native_step_11008000_d128_private.pt \
--d-model 128 --n-modes 128 --n-atoms-max 256 --max-span-bytes 1 --field-obligatory-hard --no-enable-merge --no-payload-copy --last-atom-readout --efference-every 10 --atom-flush-every 64 --episode-length 8000 --episode-reset --device cpu
```
## Comparison vs transformer (for README)
- **Training:** ATOM d128 10B = 27h on 5090 vs Llama 7B 10B = 58 days on 5090 β **50x faster**, 4666x fewer params
- **Fine-tune:** ATOM 8k steps CPU 3 min on laptop, no forgetting, field_rms 0.26β1.06 measurable vs Llama LoRA 1h GPU + forgetting
- **Growth:** ATOM Net2Net d16βd32βd64βd128 same lineage, old outputs exact vs Llama must retrain from scratch
- **Live:** ATOM probe field in real time `probe_field_regimes.py` cold 0.04 / primed 1.06, generation `Peut-Γͺtre`, `Parler,` vs Llama KV-cache only
- **No fluency claim yet:** ATOM cold probe still `tu tu` / spaces attractor since 336k, but weakening (clean stop `irra..\n\n` @1M) β scaling fixes it, not patches
## What to document next
- Add `--device` + `--batch-size` + `--compile` to trainer
- Add `transition_loss_batched`
- Bench d32 vs d128 vs d256 tps on 5090
- Release tip `atom_native_step_1080000_d32_ms1M.pt` β `d128_10GB.pt`
This is the new race: infinite, growable, CPU-scalable, live-modifiable. Train 10s GB ultra-fast on GPU, then teach new stuff on CPU from same .pt.
|