# Lesson 005: Block-Level Offload Tracing — Pause/Resume + Companion Fixes

**Date:** 2026-02-11
**Severity:** Critical (blocks tracing of models > single GPU VRAM)
**Files:** `forge/tracer/capture.py`, `forge/tracer/worker.py`, `forge/tracer/stimulus/structure.py`

## Context

Tracing models whose largest component exceeds a single GPU's VRAM (e.g., FLUX.2-dev transformer at 32GB fp16 on V100-32GB) requires **block-level CPU offload** via accelerate's `dispatch_model()`. This places one transformer block on GPU at a time (~570MB) instead of the full component (~32GB).

## Problem 1: TorchDispatchMode Captures Weight Transfer Ops

### Symptom

FLUX transformer trace produced **395 orphan constant tensors** (60GB+ graph.json) instead of the expected ~17MB. The graph was polluted with hundreds of copy operations from accelerate's weight transfers.

### Root Cause

`ATenHookManager` uses `TorchDispatchMode` which intercepts **ALL** PyTorch ATen operations, including:
- `aten::_to_copy` (weight dtype conversion)
- `aten::copy_` (CPU to GPU transfer)
- `aten::empty` (GPU buffer allocation)

These happen inside accelerate's `_hf_hook.pre_forward()` when moving weights from CPU to GPU before each block's computation. Since the weight tensors haven't been registered in `_weight_data_ptr_map` yet (they're on CPU with different data_ptrs), every transferred weight is captured as an orphan constant.

```
Accelerate dispatch flow per block:
  pre_forward(module):       ← Copies weights CPU→GPU (ATen ops!)
    module.weight = weight.to(device)
    module.bias = bias.to(device)
  module.forward(inputs):    ← Actual compute (ATen ops we WANT)
  post_forward(module):      ← Moves weights back to CPU (ATen ops!)
    module.weight = weight.to("cpu")
```

### Fix: Mutable Pause Flag

Added `_pause_flag = [False]` (mutable list) to `ATenHookManager`. The flag is checked early in `__torch_dispatch__`:

```python
# capture.py — ATenHookManager.__init__
self._pause_flag = [False]

# capture.py — CaptureMode.__torch_dispatch__
if pause_flag[0]:
    return func(*args, **kwargs)  # Execute but don't record
```

In the worker, accelerate's hooks are wrapped to toggle the flag:

```python
# worker.py — _run_dispatched_trace()
pause_flag = aten_hooks._pause_flag

def make_pre_wrapper(graph_ref, prefix, orig_fn, module_ref, pf):
    def wrapped_pre_forward(module, *args, **kwargs):
        pf[0] = True          # PAUSE capture
        try:
            result = orig_fn(module, *args, **kwargs)
        finally:
            pf[0] = False     # RESUME capture
        # Weights are now on GPU — refresh data_ptr map
        graph_ref.refresh_weight_data_ptrs(module_ref, prefix)
        return result
    return wrapped_pre_forward

def make_post_wrapper(orig_fn, pf):
    def wrapped_post_forward(module, output):
        pf[0] = True          # PAUSE capture
        try:
            result = orig_fn(module, output)
        finally:
            pf[0] = False     # RESUME capture
        return result
    return wrapped_post_forward
```

### Why Mutable List, Not Bool

Python closures capture variables by reference, but rebinding (`flag = True`) creates a local variable in the inner scope. A mutable container (`flag[0] = True`) modifies the shared object, visible to all closures.

### Why refresh_weight_data_ptrs After pre_forward

When accelerate copies a weight from CPU to GPU, the tensor's `data_ptr()` changes (new GPU memory address). The `GraphBuilder` identifies weights by `data_ptr`, so the mapping must be refreshed after each transfer to correctly attribute the tensor to its parameter name instead of creating an orphan.

**Result:** 6936 ops, 0 orphan constants, 17MB graph.json (was 60GB+).

---

## Problem 2: Nested vocab_size in Multimodal Text Encoders

### Symptom

FLUX text_encoder (Pixtral/Mistral3) trace fails with `vocab_size not found in model config`.

### Root Cause

Standard models store `vocab_size` at the root of `config.json`. Multimodal models (Pixtral, Mistral3) nest it inside `text_config`:

```json
{
  "architectures": ["PixtralModel"],
  "text_config": {
    "vocab_size": 131072,
    "hidden_size": 4096
  }
}
```

### Fix

Fallback lookup in `worker.py`:

```python
vocab_size = comp_config.get("vocab_size")
if vocab_size is None:
    text_cfg = comp_config.get("text_config", {})
    vocab_size = text_cfg.get("vocab_size") if isinstance(text_cfg, dict) else None
```

### Rule

**Never assume config key locations are flat.** Always check nested config sections (`text_config`, `vision_config`) when the root key is missing.

---

## Problem 3: Opaque Signatures from Accelerate Wrappers

### Symptom

FLUX VAE trace fails: "All layers failed for AutoencoderKLFlux2". The archetype system cannot discover the `z` parameter for the `decode()` method.

### Root Cause

Accelerate wraps module methods with hooks that have opaque `(*args, **kwargs)` signatures:

```python
# After dispatch_model(), model.decode is actually:
def hooked_method(*args, **kwargs):
    pre_forward(module, *args, **kwargs)  # Load weights
    result = original_decode(*args, **kwargs)
    post_forward(module, result)          # Offload weights
    return result
```

`inspect.signature(model.decode)` returns `(*args, **kwargs)` — no parameter names. The archetype system sees an empty param set and can't generate inputs.

### Fix

Added closure unwrapping in `structure.py:build_archetype_inputs()`:

```python
if not forward_params and hasattr(target_method, '__closure__') and target_method.__closure__:
    for cell in target_method.__closure__:
        try:
            val = cell.cell_contents
            if callable(val):
                inner_sig = inspect.signature(val)
                inner_params = set(inner_sig.parameters.keys()) - {'self', 'return_dict', 'kwargs'}
                for p, pi in inner_sig.parameters.items():
                    if pi.kind in (inspect.Parameter.VAR_KEYWORD, inspect.Parameter.VAR_POSITIONAL):
                        inner_params.discard(p)
                if inner_params:
                    forward_params = inner_params
                    break
        except ValueError:
            pass
```

This recovers the real parameters (`z`, `return_dict`) from the original method stored in the closure.

### Rule

**When `inspect.signature()` returns opaque `(*args, **kwargs)`, always check `__closure__` cells for the original unwrapped callable.** This pattern occurs with any decorator or hook wrapper (accelerate, DeepSpeed, custom wrappers).

---

## Verification Matrix

| Model | Component | Issue Fixed | Result |
|-------|-----------|-------------|--------|
| FLUX.2-dev | transformer (32GB fp16) | Pause/resume | 6936 ops, 0 orphans, 17MB |
| FLUX.2-dev | text_encoder (Pixtral) | Nested vocab_size | 6473 ops, 8 tiny constants |
| FLUX.2-dev | vae (AutoencoderKLFlux2) | Opaque signature | 167 ops, 3 tiny constants |
| PixArt-XL-2-1024-MS | all 3 | Regression test | No regression |
| Sana 4K BF16 | all 3 | Regression test | No regression |

## Key Takeaways

1. **TorchDispatchMode is global** — it sees ALL tensor ops, not just compute. Any framework that moves tensors (accelerate, DeepSpeed, FSDP) will inject unwanted ops unless capture is paused.

2. **data_ptr is the ground truth** for weight identity. When weights move between devices, data_ptr changes. The mapping MUST be refreshed after every device transfer.

3. **Config structures vary across model families.** Never assume flat configs — always check nested sections.

4. **Hook wrappers destroy introspection.** `inspect.signature()` on a wrapped method is useless. The real signature is recoverable from `__closure__` cells.
