docs/OPTIMIZATION_2026-08-22.md

# Optimisation de l'entraînement Fractus-1B — 2026-08-22

**Principe cardinal : chirurgie à cœur ouvert.**
Toute modification du code doit être **mathématiquement équivalente** à
l'implémentation de référence, prouvée par tests d'équivalence numérique,
sans changer la moindre forme de paramètre. Le cerveau (`.pt`) ne bouge pas ;
seul le corps (le code) est opéré. La reprise au token exact via les manifests
reste valide — un pod peut passer à ce code **en cours de run**, sans jeter la
digestion.

---

## 1. Diagnostic : pourquoi ~1000 tok/s/GPU alors que la 5090 peut faire bien plus

Config 1B par bloc : 20 têtes × d_head=64, 2 niveaux, chunk C=128.
Le moteur aplatit `(B, niv, H)` en `G = B·niv·H = B·40` groupes.

### Goulot n°1 — le cumsum déguisé en matmul (`attention.py`, ancien code)

```python
S = torch.einsum("tj,bjpq->btpq", mask_tril, outer)   # Σ_{j≤t} outer[j]
```

C'est un cumsum causal écrit comme une multiplication par un masque
triangulaire inférieur :

| Forme | Opérations | À C=128, D=64 |
|---|---|---|
| masque (ancienne) | O(C²·D²) MACs + rétro O(C²·D²) | ≈ 67M MACs/groupe |
| cumsum (nouvelle) | O(C·D²) additions | ≈ 0.5M adds/groupe → **~128× moins** |

La rétropropagation profite pareil : dérivée d'un cumsum = cumsum inverse,
pas un nouveau matmul masqué géant.

### Goulot n°2 — matérialisation des tenseurs `(G, C, D, D)`

`outer` et `S` font `G·128·64·64` éléments. À B=4 : G=160 → **84M éléments
(~168 Mo en bf16)** chacun, retenus pour la backward, **par bloc ×16 blocs**.
C'est LA pression VRAM qui force `BATCH=2-4` et interdit `torch.compile`
(commentaire "compile disabled for VRAM" dans fast4gpu_boost.py).

Forme *chunked* (bloc=64) : on ne matérialise que `(G, 64, 64)` (scores)
+ l'état courant `(G, D, D)` → **~32× moins d'activations dominantes**.

### Goulot n°3 — la tête liée 50257

Les logits `(B, C, 50257)` fp32 (~103 Mo à B=4) + softmax backward sont
matérialisés deux fois avec SS. La tête ≈ **la moitié des FLOPs actifs**
(64.3M des ~192M MACs/token). CE par morceaux checkpointée → mémoire
transitoire plate (`ce_chunk·vocab`), sans toucher aux FLOPs.

### Goulot n°4 — pipeline données

`np.load(mmap).to(torch.int64)` copiait **tout le shard en RAM** :
430M tokens × 8 octets = **3,4 Go/process × 8 process**. Le slicing
par chunk int32→long supprime ça (et la page cache fait le reste).

### Estimation d'utilisation GPU

FLOPs actifs/token ≈ 1.15 GFLOP (fwd+bwd) → à 1000 tok/s ≈ 1.15 TFLOPS
sur une 5090 (~200 TFLOPS bf16 crête) : **<1% d'utilisation**. Le modèle est
borné par la mémoire/les petits batches/l'overhead Python, pas par le calcul.
Tout l'espace d'optimisation est là.

---

## 2. Les optimisations (toutes prouvées équivalentes)

### Opt 1 — noyau cumsum (`fractus/nn/attention.py`)
- `_linear_attention_causal_einsum` : **référence conservée** (vérité terrain).
- `_linear_attention_causal_cumsum` : production par défaut.
- `_linear_attention_causal_vectorized` : dispatcher inchangé pour les appelants.
- Preuves : `tests/test_attention_equivalence.py::test_forward_matches_einsum_*`,
  `test_gradients_match_with_carry`, `test_both_match_looped_reference`.

### Opt 2 — forme chunked mémoire-plate (`_linear_attention_causal_chunked`)
Intra-bloc via matrice de scores `(bloc×bloc)` masquée + inter-bloc via l'état
courant `(S_run, z_run)` ; mise à jour inclusive APRÈS lecture (= sémantique
exacte S_t incluant le token t). Sélection : `set_attention_impl('chunked')`
ou env `FRACTUS_ATTN_IMPL=chunked`. Ragged → fallback cumsum exact.
- Recommandation pod : **chunked sur GPU** (VRAM → compile + gros batch),
  cumsum reste parfait sur CPU.

### Opt 3 — CE par morceaux checkpointée (`fractus/nn/ce.py`)
`chunked_cross_entropy(h, W, targets, ce_chunk)` : perte moyenne identique,
grads identiques (tolérance ordre de sommation fp32), pics mémoire plats grâce
à `torch.utils.checkpoint` par morceau (recompute en backward).
Nouvelle entrée moteur : `tick_chunk_train_ce(obs, targets[, return_hidden])`.
`sample_tokens_chunked(...)` remplace le multinomial dense pour le SS
(distribution identique ; flux RNG consommé par morceaux → tirages non
bit-identiques, statistiquement équivalents).

### Opt 4 — sémantique SS préservée + accumulation optionnelle
`scripts/fast4gpu_boost_v2.py` réplique EXACTEMENT v1 quand `ACCUM=1`
(défaut) : step TF puis step SS séparés. `ACCUM>1` = capability nouvelle,
déviation documentée (grads TF+SS accumulés, un clip+step tous les N lots).

### Opt 5 — pipeline données zéro-copie
Slice memmap int32 par chunk → transfert → `.long()` sur GPU.
Un seul fetch `(B·SEQ+1)` fournit chunk ET target. RAM économisée :
~27 Go sur un pod 8×5090.

---

## 3. Comment déployer sur le pod (cœur ouvert)

```bash
# 1. Sauvegarder l'état (rien à faire de plus : HF = source de vérité)
#    Les checkpoints gpu*.pt + RESUME_MANIFEST restent valides tels quels.

# 2. Remplacer le corps :
#    fractus/nn/attention.py, fractus/nn/ce.py (nouveau),
#    fractus/continuous_engine.py (méthode ajoutée, rien retiré),
#    scripts/fast4gpu_boost_v2.py (nouveau).

# 3. Relancer chaque GPU avec les MÊMES offsets qu'avant l'arrêt :
CUDA_VISIBLE_DEVICES=$i GPU_ID=$i \
  START_TOKEN=$(python -c "import json;print(json.load(open('checkpoints/RESUME_MANIFEST_8GPU.json'))['gpu_$i']['start_token'])") \
  BATCH=8 CE_CHUNK=2048 FRACTUS_ATTN_IMPL=chunked COMPILE=1 \
  python -u scripts/fast4gpu_boost_v2.py
```

Ordre de montée en puissance recommandé (valider à chaque étape) :
1. Code swap, mêmes réglages que v1 (`BATCH=4 CE_CHUNK=0`) → vérifier ema_tf
   continue exactement sa courbe (équivalence en conditions réelles).
2. `CE_CHUNK=2048` → perte identique, VRAM ↓.
3. `FRACTUS_ATTN_IMPL=chunked` → VRAM ↓↓ (mesuré §4 : ×15–25 vs référence,
   mémoire plate là où la référence explose).
4. `COMPILE=1` puis monter `BATCH` (8, 16…) — surveiller tok/s et VRAM.
5. Optionnel `ACCUM` si on veut un lot effectif plus grand sans OOM.

Critères de non-régression (cf. docs/TRUSTED_LOSS.md) : ema_tf continue de
descendre sans NaN, lb stable ~14, sondes gén inchangées de comportement.

---

## 4. Résultats mesurés (CPU local, torch 2.9.1+cpu, 2026-08-23)

Micro-bench attention formes 1B réelles (G=B·40 groupes, C=128, dH=64),
carry actif ; médiane sur 8 itérations. Chaque cellule tourne dans son
propre sous-processus : un crash natif d'une cellule n'emporte pas la table.

| Impl | B | G | fwd ms | fwd+bwd ms | RSS Δ MB |
|---|---|---|---|---|---|
| einsum (réf) | 2 | 80  | 267.4 | 572.1 | +3 |
| einsum (réf) | 4 | 160 | 584.7 | 1126.6 | 0 |
| einsum (réf) | 8 | 320 | — | — | crash natif (commit) |
| cumsum       | 2 | 80  | 322.4 | 799.1 | −8 |
| cumsum       | 4 | 160 | — | — | crash natif (commit) |
| chunked      | 2 | 80  | **12.3** | **31.9** | 0 |
| chunked      | 4 | 160 | **23.0** | **72.4** | 0 |
| chunked      | 8 | 320 | **40.5** | **150.0** | −10 |

Lecture honnête de ces chiffres :

- **Chunked vs einsum à B égal** : ×21.7 (fwd) et ×17.9 (fwd+bwd) à B=2 ;
  ×25.4 / ×15.6 à B=4.
- **Mémoire** : einsum et cumsum matérialisent le tenseur (G, C, dH, dH)
  (~0.34 GB à G=160, ~1.3 GB à G=320, multiplié par les buffers de
  backward). Sur cette machine au commit mémoire limité ils segfaultent
  (rc=3221225477) dès G≥160–320 ; chunked est memory-flat ((G, block²+dH²))
  et traverse toutes les tailles testées. C'est exactement la propriété qui
  déverrouille BATCH≥8 + torch.compile sur pod.
- **Cumsum sur CPU est plus lent qu'einsum** : l'einsum masqué descend en
  bmm BLAS très optimisé, tandis que le scan élément-par-élément est
  borné par la bande passante. La réduction de FLOPs O(C²·dH²)→O(C·dH²)
  se paie réellement sur GPU/compile, pas sur ce CPU — c'est pour ça que
  le défaut repo reste `cumsum` (mathématiquement prouvé, sémantique
  simple) et que la recommandation pod passe directement à `chunked`.
- Les timings absolus CPU ne se transfèrent pas à CUDA ; ce qui se
  transfère : l'ordre relatif des kernels, le profil mémoire, et le fait
  mesuré que chunked scale là où la référence explose. Bench GPU à faire
  sur pod (même script).

End-to-end moteur CPU (bench_engine.py, d=128, 2 blocs, E8, B=8, SEQ=128,
20 steps) :

| Kernel | tok/s |
|---|---|
| cumsum (défaut) | 703 |
| chunked | 764 (+8.7 %) |

## 5. Tests

```bash
py -m pytest tests/ -q        # 44 passed = suite repo (28) + équivalences (14) + smoke v2 (2)
py benchmarks/bench_attention.py --iters 8
py benchmarks/bench_engine.py --steps 20
```

Statut : **44/44 passent** (2026-08-23). Les deux preuves moteur
(`test_engine_tick_chunk_train_ce_matches_train`,
`test_engine_end_to_end_chunk_equivalence`) exigent le clonage explicite des
poids (`load_state_dict`) entre instances comparées — deux constructions sous
un même seed ont des poids DIFFÉRENTS, piège documenté dans les docstrings.

---

*Document généré pendant la session d'optimisation 2026-08-22/23.
Règle : ne jamais merger dans fractus-cte sans que la section 4 soit remplie.*