# DTYPE DIVERGENCE REPORT

> **HISTORICAL DOCUMENT (December 2025)**
> This report documents the dtype flow fixes that resolved OOM issues in NeuroBrix v0.1.
> File paths have been updated to reflect the February 2026 package restructure.
> - Public package: `neurobrix` (pip install neurobrix)
> - Build tooling: `forge/` (private, requires pip install -e .)
> - Core paths: `core/` → `src/neurobrix/core/`
> - Importer paths: `core/importer/` → `forge/importer/`
> - Cache paths: `.cache/models/` → `~/.neurobrix/cache/`

**Date:** 2025-12-30
**Model Audited:** PixArt-Sigma-XL-2-1024-MS
**Purpose:** Identify dtype mismatch causing OOM

---

## EXECUTIVE SUMMARY

**ROOT CAUSE IDENTIFIED:**
Weights are saved as **float32** during import, but should preserve the **float16** that diffusers uses when loading with `torch_dtype=torch.float16`.

| Stage | Expected (Diffusers) | Actual (NeuroBrix) | Delta |
|-------|---------------------|-------------------|-------|
| text_encoder weights | float16, ~9GB | float32, ~17GB | **+8GB** |
| transformer weights | float16, ~2.3GB | float32, ~2.3GB | - |
| vae weights | float16, ~0.16GB | float32, ~0.31GB | **+0.15GB** |
| **TOTAL** | **~12GB** | **~20GB** | **+8GB** |

---

## 1. DIFFUSERS BEHAVIOR (SOURCE OF TRUTH)

When loading PixArt-Sigma with `torch_dtype=torch.float16`:

```python
pipe = PixArtSigmaPipeline.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
    torch_dtype=torch.float16
)
```

### Component DTYPEs After Load:

| Component | Weight dtype | Memory |
|-----------|-------------|--------|
| text_encoder | **float16** | ~9GB |
| transformer | **float16** | ~2.3GB |
| vae | **float16** | ~0.16GB |

**Key Observation:**
- HuggingFace SOURCE files are stored as **float32** on disk
- Diffusers converts to **float16** during `from_pretrained()` when `torch_dtype=torch.float16`
- Total runtime memory: **~12GB**

---

## 2. NEUROBRIX TRACE BEHAVIOR (CORRECT)

When tracing, we use diffusers with the same config:

```python
# trace is done with diffusers loaded in fp16
# Graph captures ACTUAL runtime dtypes
```

### Graph Recorded DTYPEs:

```json
// components/text_encoder/graph.json
{
  "tensors": {
    "param::text_model.encoder.layers.0.self_attn.q_proj.weight": {
      "dtype": "torch.float16"  // CORRECT - captured from diffusers
    }
  }
}
```

**Trace is CORRECT** - graph records float16 because diffusers loaded as float16.

---

## 3. NEUROBRIX IMPORT BEHAVIOR (BUG)

### Code Path: `forge/importer/builder.py` (formerly `core/importer/nbx_builder.py`)

```python
for sf_path in sf_files:
    with safe_open(str(sf_path), framework="pt") as f:
        for original_name in f.keys():
            tensor = f.get_tensor(original_name)  # <-- Reads from HF SOURCE (float32!)
            # ...
            normalized_tensors[normalized] = tensor  # <-- Saved as-is (float32!)
```

### Problem:
1. Reads weights from **HuggingFace source** safetensors files
2. These files store weights as **float32** on disk
3. **NO conversion** to traced dtype (float16)
4. Weights saved to NBX shards as **float32**

### Evidence - Cached Weight File Sizes:

```
# ~/.neurobrix/cache/<model_name>/components/text_encoder/weights/
text_encoder/weights/shard_000.safetensors: 2.00 GB (float32)
text_encoder/weights/shard_001.safetensors: 2.00 GB (float32)
...
Total text_encoder: 17.74 GB  <-- Should be ~9GB if float16
```

---

## 4. NEUROBRIX WEIGHT LOADING (CURRENT)

### Code Path: `src/neurobrix/core/runtime/weight_loader.py` (formerly `core/runtime/weight_loader.py`)

```python
# Current code - NO dtype conversion
weights = safetensors_load_file(str(cache_file), device=device)
return weights  # Returns float32 as-is
```

### Previous Working Code (Saves/core2):

```python
# dtype conversion was present
weights = safetensors_load_file(str(cache_file), device=device)
if dtype is not None:
    weights = {k: v.to(dtype) if v.is_floating_point() else v for k, v in weights.items()}
return weights
```

---

## 5. PRISM ESTIMATES (AFFECTED)

Prism reads shard sizes from `weights_index.json` to estimate memory:

```json
// weights_index.json
{
  "total_size_bytes": 17740000000,  // 17.74GB (float32)
  "dominant_dtype": "float32"
}
```

Prism correctly reports:
- Component needs ~18GB (based on float32 shard sizes)
- This is **2x** the actual diffusers runtime requirement

---

## 6. FULL FLOW COMPARISON

```
┌─────────────────────────────────────────────────────────────────────────────┐
│                           DIFFUSERS FLOW (CORRECT)                          │
├─────────────────────────────────────────────────────────────────────────────┤
│  HF Disk (fp32) ──► from_pretrained(torch_dtype=fp16) ──► GPU Memory (fp16) │
│                                                              ~12GB          │
└─────────────────────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────────────────────┐
│                          NEUROBRIX FLOW (BROKEN)                            │
├─────────────────────────────────────────────────────────────────────────────┤
│ TRACE:    diffusers(fp16) ──► ATen hooks ──► graph.json records fp16   ✓    │
│                                                                              │
│ IMPORT:   HF Disk (fp32) ──► nbx_builder ──► NBX shards (fp32)         ✗    │
│                              NO conversion!                                  │
│                                                                              │
│ LOAD:     NBX shards (fp32) ──► weight_loader ──► GPU Memory (fp32)    ✗    │
│                                                              ~20GB          │
└─────────────────────────────────────────────────────────────────────────────┘
```

---

## 7. RECOMMENDED FIX

### Option A: Fix at IMPORT time (Preferred)

In `forge/importer/builder.py`, convert weights to traced dtype before saving:

```python
# Get dominant dtype from traced graph
traced_dtype = graph_data.get("dominant_dtype", "float16")

# Convert tensor before adding to sharder
if tensor.is_floating_point():
    tensor = tensor.to(getattr(torch, traced_dtype))
normalized_tensors[normalized] = tensor
```

**Pros:**
- NBX files are correct size (~12GB instead of ~20GB)
- Prism estimates are correct
- No runtime conversion needed

### Option B: Fix at LOAD time

In `src/neurobrix/core/runtime/weight_loader.py`, restore dtype conversion:

```python
# Get dtype from Prism allocation
weights = safetensors_load_file(str(cache_file), device=device)
if dtype is not None:
    weights = {k: v.to(dtype) if v.is_floating_point() else v for k, v in weights.items()}
```

**Pros:**
- Quick fix, restores previous behavior

**Cons:**
- NBX files still 2x larger than needed
- Runtime conversion overhead
- Prism estimates still wrong

---

## 8. VERIFICATION COMMANDS

```bash
# Check HuggingFace cached weight file sizes (source)
ls -lh ~/.cache/huggingface/hub/models--PixArt-alpha--PixArt-Sigma-XL-2-1024-MS/

# Check NBX extracted weight file sizes (runtime cache)
ls -lh ~/.neurobrix/cache/PixArt-Sigma-XL-2-1024-MS/components/text_encoder/weights/

# Check NBX shard sizes in .nbx archive
python -c "
import zipfile
with zipfile.ZipFile('models/image/PixArt-Sigma-XL-2-1024-MS/model.nbx') as z:
    for name in z.namelist():
        if 'shard' in name:
            info = z.getinfo(name)
            print(f'{name}: {info.file_size / 1e9:.2f} GB')
"

# Check graph recorded dtypes
python -c "
import json
import zipfile
with zipfile.ZipFile('models/image/PixArt-Sigma-XL-2-1024-MS/model.nbx') as z:
    with z.open('components/text_encoder/graph.json') as f:
        g = json.load(f)
        for tid, t in list(g.get('tensors', {}).items())[:5]:
            print(f'{tid}: {t.get(\"dtype\")}')"
```

---

## CONCLUSION

The divergence is clear:
1. **DIFFUSERS** converts weights to fp16 at load time (correct, ~12GB)
2. **TRACE** correctly captures fp16 in graph (correct)
3. **IMPORT** saves weights as fp32 from HF source (INTENTIONAL - preserves precision)
4. **LOAD** was returning fp32 weights without conversion (BUG - FIXED)

---

## FIX IMPLEMENTED (2025-12-30)

### Architecture Decision: Runtime Conversion

We chose **runtime conversion** instead of import-time conversion:

```
IMPORT TIME:
  - Save weights as SOURCE dtype (fp32) - preserves precision

RUNTIME:
  - Prism determines optimal dtype from hardware profile (supports_dtypes)
  - WeightLoader converts to Prism dtype during load
```

### Files Changed:

| File | Change | Current Location (Feb 2026) |
|------|--------|-----------------------------|
| `core/graph/dtype_converter.py` | NEW - Graph dtype conversion utility | `src/neurobrix/core/graph/dtype_converter.py` |
| `core/runtime/weight_loader.py` | Accept `dtype` param, convert weights after load | `src/neurobrix/core/runtime/weight_loader.py` |
| `core/runtime/graph_executor.py` | Pass `self.dtype` (from Prism) to weight loader | `src/neurobrix/core/runtime/graph_executor.py` |

### Flow After Fix:

```
Hardware YAML (supports_dtypes: ["float32", "float16"])
    → Prism._resolve_dtype()
    → allocation.dtype = torch.float16
    → Factory extracts dtype
    → GraphExecutor(dtype=torch.float16)
    → load_weights() passes self.dtype
    → WeightLoader converts: weights.to(self.dtype)
    → GPU Memory: ~12GB (not 20GB)
```

### Benefits:

1. **ONE NBX file works on ALL hardware** - V100 gets fp16, A100 gets bf16
2. **Maximum precision preserved** - Source weights stay fp32
3. **ZERO HARDCODE** - dtype embedded from hardware profile, not hardcoded

---

## FIX COMPLETED: Memory Leak (2025-12-30)

### Additional Issue Discovered

After fixing the dtype flow, OOM persisted. Root cause identified:

**`unload_weights()` was NOT clearing `_tensor_store`**

```python
# BEFORE (BROKEN):
def unload_weights(self) -> None:
    self._weights.clear()      # Clears weight dict
    self._op_outputs.clear()   # Clears op outputs
    # _tensor_store NEVER cleared!
    # It still holds references to weight tensors
    # Memory NOT freed!
```

### Why This Caused OOM

```
Memory flow:
1. load_weights() → _weights = {"attn.q": tensor}
2. _tensor_store = {"param::attn.q": tensor}  # SAME tensor, 2nd reference
3. unload_weights():
   - _weights.clear()  → Dict cleared, but tensor still referenced by _tensor_store
   - Memory NOT freed!
4. Next component loads → OOM (previous weights still in GPU memory)
```

### Fix Applied

```python
# AFTER (FIXED):
def unload_weights(self) -> None:
    # Clear tensor store FIRST (holds refs to weights)
    self._tensor_store.clear()  # NEW - Critical fix
    self._weights.clear()
    self._op_outputs.clear()
    self._inputs.clear()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
```

### Additional Fix: inference_mode

Also wrapped `run()` method in `torch.inference_mode()`:

```python
def run(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
    with torch.inference_mode():  # Disables autograd completely
        # All execution code inside this block
        return self._gather_outputs()
```

**Benefits:**
- Disables gradient tracking completely
- Prevents memory doubling from autograd graph
- Faster execution

---

## VERIFICATION SUCCESSFUL (2025-12-30)

```bash
# Original verification command (December 2025)
$ python neurobrix.py run \
    --family image \
    --model PixArt-Sigma-XL-2-1024-MS \
    --hardware v100-32g \
    --prompt "A sunset" \
    --steps 5 \
    --native

# Current command format (February 2026 - package mode)
$ neurobrix run \
    --model PixArt-Sigma-XL-2-1024-MS \
    --hardware v100-32g \
    --prompt "A sunset" \
    --steps 5

# Or dev mode (from project root)
$ PYTHONPATH=src python -m neurobrix run \
    --model PixArt-Sigma-XL-2-1024-MS \
    --hardware v100-32g \
    --prompt "A sunset" \
    --steps 5

# Output:
Step 1/5 completed
Step 2/5 completed
Step 3/5 completed
Step 4/5 completed
Step 5/5 completed
VAE decode complete
Image saved to: outputs/...
```

**Result: Generation completes successfully on V100 32GB with float16.**

**Note:** The `--family` flag was removed in the February 2026 package restructure. The family is auto-detected from the model's manifest.json.

---

## COMPLETE FIX SUMMARY

| Issue | Cause | Fix | File (Current Location) |
|-------|-------|-----|------------------------|
| Weights loaded as fp32 | WeightLoader no dtype param | Pass dtype from Prism | `src/neurobrix/core/runtime/weight_loader.py` |
| Memory doubling | Missing inference_mode | Wrap run() in inference_mode | `src/neurobrix/core/runtime/graph_executor.py` |
| Memory leak on unload | _tensor_store not cleared | Clear _tensor_store first | `src/neurobrix/core/runtime/graph_executor.py` |

**All fixes combined enable PixArt-Sigma to run on V100 32GB, matching diffusers behavior.**
