# Lesson 002: DtypeEngine — Prism Is The Master

**Date:** 2026-02-02
**Severity:** Critical (NaN outputs, black images, dtype crashes)
**Files:** `src/neurobrix/core/dtype/engine.py`, `src/neurobrix/core/runtime/graph/compiled_ops.py`, `src/neurobrix/core/flow/iterative_process.py`, `forge/tracer/worker.py`

## Symptom

Sana 1600M generation produces NaN outputs or completely black images after implementing a centralized DtypeEngine with AMP-style wrapping (FP32_OPS, COMPUTE_OPS classification).

## Root Cause

Three interconnected issues:

### 1. DtypeEngine trop complexe — AMP wrapping inutile

Le DtypeEngine tentait de reproduire PyTorch autocast (FP32_OPS upcast fp32, COMPUTE_OPS downcast fp16, passthrough ops sans cast). Cette approche est incompatible avec notre runtime car :

- Le graph capture DEJA les upcasts/downcasts internes du modele (via `_to_copy` ops)
- Les poids sont deja au bon dtype via WeightLoader (Prism dicte)
- Les constantes sont converties au chargement
- Ajouter des casts AMP par-dessus cree des conflits (double cast, upcasts annules)

**Regle : le graph est la source de verite pour le flux de dtype. Le DtypeEngine ne doit intervenir que sur `_to_copy` et les constantes.**

### 2. `_to_copy` — ne JAMAIS remapper les targets fp32

Les graphs contiennent des `_to_copy` avec target fp32 pour la stabilite numerique (RoPE freq computation, RMSNorm internals, MoE routing scores). Ces targets fp32 doivent etre preservees.

```python
# CORRECT — ancien code et DtypeEngine simplifie :
# Remap UNIQUEMENT fp16/bf16 targets, JAMAIS fp32
if target_dtype in (torch.float16, torch.bfloat16):
    if compute_dtype is not None and target_dtype != compute_dtype:
        target_dtype = compute_dtype

# FAUX — remapper les targets fp32 :
if target_dtype.is_floating_point and target_dtype != compute_dtype:
    target_dtype = compute_dtype  # Casse les upcasts de stabilite
```

### 3. `convert_constant` — respecter les constantes hors graph_dtype

Les graphs fp16 contiennent des constantes fp32 intentionnelles (ex: `rotary_emb.inv_freq` dans Gemma2). Ces constantes alimentent des chaines de calcul fp32 (RoPE). Les convertir en fp16 casse la chaine.

```python
# CORRECT — ne convertir que les constantes qui matchent graph_dtype :
if self.graph_dtype is not None and tensor.dtype != self.graph_dtype:
    return tensor  # Constante fp32 dans graph fp16 = intentionnelle

# FAUX — convertir TOUTES les constantes float :
if tensor.is_floating_point() and tensor.dtype != self.compute_dtype:
    return tensor.to(self.compute_dtype)  # Casse inv_freq fp32
```

## Lecons connexes

### Tracer : bf16 → fp16 sur V100, pas bf16 → fp32

Le worker convertissait bf16 → fp32 sur V100 par precaution. Mais cela produit un graph fp32 que le runtime ne sait pas gerer (toutes les `_to_copy` ont target fp32, impossible de distinguer "dtype par defaut" de "upcast stabilite").

**Fix :** bf16 → fp16 dans le worker. Prism verifie independamment que les poids bf16 sont dans les limites fp16 (±65504) avant de les charger en fp16.

### `encoder_dtype` doit venir de Prism, pas de DtypeAdapter

`iterative_process.py` utilisait `DtypeAdapter.resolve_runtime_dtype()` pour deduire le dtype de l'encoder depuis la topologie. DtypeAdapter resolvait bf16 → fp32 sur V100 (ancienne regle obsolete), ce qui contredit Prism qui a verifie via `_scan_bf16_fp16_safety()` et decide bf16 → fp16.

**Fix :** `encoder_dtype` vient directement de `plan.components[encoder].dtype` (allocation Prism).

### Phase A du traceur doit etre fp32 (topologie seulement)

Phase A charge le pipeline complet pour extraire la topologie (inputs/outputs entre composants). Pas besoin de precision. Charger en fp32 sur CPU :
- Evite le warning "fp16 cannot run with cpu"
- Evite de forcer un dtype global sur tous les composants
- Les workers tracent chaque composant dans son dtype natif

## Architecture finale du DtypeEngine

```
DtypeEngine (simplifie) :
  compile_op():
    - _to_copy → remap fp16/bf16 targets vers compute_dtype
    - Toute autre op → passthrough (zero wrapping)

  convert_constant():
    - Constante dtype == graph_dtype → convertir en compute_dtype
    - Constante dtype != graph_dtype → preserver (intentionnelle)

  PAS de :
    - FP32_OPS / COMPUTE_OPS classification
    - AMP wrapping (upcast/downcast)
    - enabled flag
    - get_op_dtype()
    - cast_inputs()
```

## Fichiers modifies

| Fichier | Changement |
|---------|-----------|
| `src/neurobrix/core/dtype/engine.py` | Simplifie : _to_copy + convert_constant seulement |
| `src/neurobrix/core/runtime/graph/compiled_ops.py` | Utilise DtypeEngine.compile_op() au lieu de logique locale |
| `src/neurobrix/core/flow/iterative_process.py` | encoder_dtype depuis Prism plan, plus DtypeAdapter |
| `forge/tracer/worker.py` | bf16 → fp16 sur V100 (pas fp32) |
| `forge/tracer/orchestrator.py` | Phase A en fp32/CPU |

## Regle d'or

**Prism dicte le compute_dtype. Le graph dicte le flux de dtype (upcasts/downcasts internes). Le DtypeEngine ne fait que traduire les `_to_copy` du graph vers le compute_dtype de Prism. Rien de plus.**
