# Lesson 004: FP16 Matmul Overflow — Inf Fix (Not Clamp!)

**Date:** 2026-02-03
**Severity:** Critical (NaN propagation, marble effect images)
**Files:** `src/neurobrix/core/dtype/engine.py`, `src/neurobrix/core/runtime/graph/compiled_sequence.py`

## Symptom

Sana 4K (4096x4096) generation on V100 with fp16 produces:
1. NaN values in transformer attention outputs
2. "Marble effect" images when using CLAMP approach
3. 2.5x performance overhead when using per-op NaN-guard

## Root Cause

### The Overflow Chain

```
addmm (Q/K/V projection) produces values > 65504
    ↓
fp16 cannot represent > 65504 → becomes inf
    ↓
SDPA receives inf in Query tensor
    ↓
Attention computation produces NaN
    ↓
NaN propagates through all subsequent ops
```

### Why CLAMP Was Wrong

Initial fix attempted:
```python
# BROKEN — clips VALID values too!
result = result.clamp(-65504, 65504)
```

Problem: `clamp()` modifies ALL values > 65504, including valid intermediate computations that would naturally reduce below 65504 in subsequent ops. This corrupts the model's internal representations → marble/incoherent images.

### Why FP32_OPS Wrapping Was Wrong

Second attempt — wrap matmul in fp32:
```python
# BROKEN — changes model semantics
def fp32_wrapper(*args):
    args_fp32 = [a.float() for a in args]
    result = func(*args_fp32)
    return result.half()  # ← clamp needed here too!
```

Problem: Even with fp32 compute, the downcast back to fp16 requires clamping → same marble effect.

## Correct Solution: `nan_to_num` on Inf Only

```python
# CORRECT — replaces ONLY inf values, preserves valid values
MATMUL_OPS = frozenset({"addmm", "mm", "bmm", "baddbmm", "matmul", "linear"})

def _make_inf_fix_wrapper(self, func):
    def inf_fix_func(*args, **kwargs):
        result = func(*args, **kwargs)
        if isinstance(result, torch.Tensor) and result.dtype == torch.float16:
            if torch.isinf(result).any():
                result = torch.nan_to_num(result, nan=0.0, posinf=65504.0, neginf=-65504.0)
        return result
    return inf_fix_func
```

**Critical difference:**
- `clamp(-65504, 65504)` → modifies ALL values outside range
- `nan_to_num(posinf=65504)` → replaces ONLY inf values, valid values unchanged

## Diagnostic Journey

### Step 1: Per-Op NaN-Guard (Expensive)

Added NaN-guard in `src/neurobrix/core/runtime/graph/compiled_sequence.py` that checks EVERY op output:
```python
if torch.isnan(result).any() or torch.isinf(result).any():
    result = torch.nan_to_num(result, ...)
```

**Result:** Works but 2.5x slower (5434ms vs 2199ms for transformer)

### Step 2: Identify Culprit Op

With `NBX_NAN_GUARD_VERBOSE=1`, traced the NaN origin:
```
[NaN-GUARD TUPLE] Op 1816/2454: aten::_scaled_dot_product_efficient_attention
  Output[0]: NaN=3655120
  Input[0] (Query): range=[-inf, inf], inf=True  ← INF HERE!

★ CREATOR: Op 1797: aten::addmm
  Input[1]: range=[-1884, 2009]  ← Large values overflow in fp16
```

### Step 3: Targeted Fix

Wrap ONLY matmul ops (addmm, mm, bmm, etc.) with the inf-fix wrapper in `src/neurobrix/core/dtype/engine.py`.

## Performance Impact

| Configuration | Transformer Time | Notes |
|--------------|------------------|-------|
| Per-op NaN-guard (old) | 5434ms | Check every op = slow |
| Targeted inf-fix (new) | 2199ms | Check only matmul ops |
| **Improvement** | **-59%** | |

## Files Modified

| File | Change |
|------|--------|
| `src/neurobrix/core/dtype/engine.py` | Added `MATMUL_OPS`, `_make_inf_fix_wrapper()`, applied in `compile_op()` BEFORE legacy mode check |
| `src/neurobrix/core/runtime/graph/compiled_sequence.py` | Changed NaN-guard default: OFF (was ON). Enable with `NBX_NAN_GUARD=1` for debugging |

## Environment Variables

| Variable | Effect |
|----------|--------|
| `NBX_NAN_GUARD=1` | Enable per-op NaN-guard (slow, for debugging) |
| `NBX_NAN_GUARD_VERBOSE=1` | Show which ops create/propagate NaN |
| `NBX_LEGACY_AMP=1` | Enable old FP32_OPS wrapping (broken, for testing) |

## Key Insight

**fp16 overflow in matmul is the ROOT CAUSE. Fix it at the source (matmul wrapper), not at every consumer (per-op guard).**

The `nan_to_num` approach is correct because:
1. Inf values are INVALID — they represent overflow, not valid computation
2. Replacing inf with max_fp16 (65504) is mathematically sound — it's the closest representable value
3. Valid values (even large ones like 60000) are preserved unchanged

## Rule

**Never use `clamp()` for fp16 overflow protection. Use `nan_to_num(posinf=max, neginf=-max)` which only replaces invalid inf values.**

## Tentative AMP échouée (même session)

Après le fix inf, une tentative d'implémenter AMP (FP32_OPS wrapping) a produit des résultats catastrophiques (rayures diagonales).

**Cause:** L'implémentation faisait un downcast fp32→fp16 immédiat après chaque op, alors que PyTorch autocast laisse la sortie en fp32 et laisse l'op suivante gérer le cast.

**Leçon:** AMP nécessite une implémentation au niveau du tracer, pas du runtime. Voir `docs/roadmap/amp-mixed-precision.md` pour le plan futur.

## État final

| Composant | État | Impact |
|-----------|------|--------|
| `_make_inf_fix_wrapper` | ✅ Actif | Fixe overflow matmul |
| `_make_fp32_wrapper` (AMP) | ❌ Désactivé | Cassait les images |
| NaN-guard | ❌ Off par défaut | Disponible pour debug |

La qualité reste légèrement dégradée (proportions visage) due à bf16→fp16. AMP correct pourrait améliorer cela — voir roadmap.
