# Activation Trace Mode

!!! note
    This page covers Extension 0 (module forward hooks), which is the implemented
    tracing mechanism. Extensions 1-3 are design sketches for future work and are
    **not yet implemented**.

## Overview

Activation trace mode is a zero-overhead-when-off, env-gated mechanism for
dumping per-layer activation statistics during FastVideo inference. Its primary
use case is **parity debugging across model ports**: enable tracing on both
FastVideo and the upstream reference implementation, then `diff` the resulting
JSONL files to find the first divergent layer.

The mechanism is intentionally narrow. It doesn't replace general logging,
profiling, or function tracing. It answers one question: "at which layer do
FastVideo and the reference model first produce different numbers?"

## When to use

- Investigating numerical drift between FastVideo and an upstream reference.
- Debugging mid-pipeline divergence (e.g., one block produces wrong output while earlier blocks match).
- Validating that a refactor preserves bf16 noise-floor behavior across many layers.

## When NOT to use

| Goal | Use instead |
|---|---|
| General logging | `init_logger(__name__)` |
| Per-stage timing | `FASTVIDEO_STAGE_LOGGING` |
| Profiling kernel timings | `FASTVIDEO_TORCH_PROFILER_DIR` (see [Profiling](profiling.md)) |
| Function-call tracing | `fastvideo.logger.enable_trace_function_call()` (heavy) |

## Quickstart

```bash
FASTVIDEO_TRACE_ACTIVATIONS=1 \
FASTVIDEO_TRACE_LAYERS="^block\.layers\.[0-9]+$" \
FASTVIDEO_TRACE_STATS="abs_mean,sum,max,shape" \
FASTVIDEO_TRACE_OUTPUT="/tmp/fv_trace.jsonl" \
python examples/inference/basic/basic_magi_human.py
```

Each line in `/tmp/fv_trace.jsonl` is a JSON record:

```json
{"module": "block.layers.0", "tensor": "out", "step": 0, "abs_mean": 1.234, "sum": -5.678, "max": 9.012, "shape": [1, 4096, 5120]}
```

## Configuration

| Env var | Default | Description |
|---|---|---|
| `FASTVIDEO_TRACE_ACTIVATIONS` | `False` | Master toggle. When unset or false, **zero overhead** in the production hot path. |
| `FASTVIDEO_TRACE_LAYERS` | `""` (all) | Python regex filter applied to `model.named_modules()` names. Empty string matches all modules. |
| `FASTVIDEO_TRACE_STATS` | `"abs_mean,sum"` | Comma-separated stats to compute. Available: `abs_mean`, `sum`, `min`, `max`, `mean`, `std`, `shape`, `dtype`. |
| `FASTVIDEO_TRACE_OUTPUT` | `"/tmp/fv_trace_<pid>.jsonl"` | Output file path. `<pid>` is replaced with the process ID at runtime. |
| `FASTVIDEO_TRACE_STEPS` | `""` (all) | Comma-separated denoising step indices to capture. Empty string captures all steps. |

## Workflow: parity-debug a model port

1. Set up a tightly-controlled comparison: a parity test or a small standalone
   script that loads both the FastVideo model and the upstream reference with
   identical inputs and seeds.

2. Run the FastVideo side with tracing on:

   ```bash
   FASTVIDEO_TRACE_ACTIVATIONS=1 \
   FASTVIDEO_TRACE_LAYERS="<your regex>" \
   FASTVIDEO_TRACE_OUTPUT="/tmp/fv_trace_fv.jsonl" \
   python <fv_runner.py>
   ```

3. Run the upstream side. The upstream repo needs separate instrumentation. See
   "Hooking the upstream side" below.

4. Sort both files by `(module, step)` if needed, then diff:

   ```bash
   diff /tmp/fv_trace_fv.jsonl /tmp/fv_trace_upstream.jsonl
   ```

5. The first divergent line identifies the first layer where FastVideo and the
   upstream produce different outputs. Start debugging there.

## Performance impact

When the master toggle is off, overhead is nil. When on, the cost is
proportional to how broadly the layer regex matches.

### What runs when tracing is on

For every match against `model.named_modules()`, FastVideo registers an
`ActivationStatHook` that runs after the module's forward returns:

1. Walks the (possibly nested) output `tuple` / `list` / `dict` and extracts
   every `torch.Tensor` leaf.
2. Computes each enabled stat on the tensor's `.detach().float()` view — the
   cast is required for bf16 inputs because some reductions are not stable
   in bf16.
3. Writes one JSON record per tensor to the line-buffered JSONL sink.

Cost is roughly `O(num_matched_modules × num_output_tensors × num_stats × tensor_numel)`
per forward, dominated by `tensor_numel` for value-stats (`abs_mean`, `sum`,
`min`, `max`, `mean`, `std`). The `shape` and `dtype` stats are O(1).

### Cost-shaping knobs

The default config (empty `FASTVIDEO_TRACE_LAYERS` is treated as `.*`, default
two stats, all steps) is intentionally blunt — useful only for a one-shot
smoke run. For real debugging, scope down:

| Knob | Effect |
|---|---|
| Tighten `FASTVIDEO_TRACE_LAYERS` to a regex matching <50 modules | Linear reduction in hook count |
| Drop unused stats from `FASTVIDEO_TRACE_STATS` (`mean`, `std`, `min`, `max`) if you only need divergence detection | One full reduction per stat saved per matched module per forward |
| Set `FASTVIDEO_TRACE_STEPS="0,15,31"` for a 32-step run | ~10x reduction vs all-steps (the hook still fires but exits early when the step doesn't match) |
| Use `shape` + `dtype` only on layers where you only care about layout | Skips tensor reductions entirely on those layers |

### Disk

Output is line-buffered (`open(..., buffering=1)`), so every record flushes
on write. A typical 32-step run that traces 40 DiT blocks with 4 stats
writes roughly 5K records — about 1 MB of JSONL. Point
`FASTVIDEO_TRACE_OUTPUT` at a fast local disk for parity runs; slow network
mounts will dominate runtime once tracing is on.

## Troubleshooting

### "I set the env var but no JSONL file appears"

Three things to check, in order:

1. **Toggle semantics.** `FASTVIDEO_TRACE_ACTIVATIONS` uses a strict
   not-equal-to-`"0"` test. Setting it to `1`, `true`, or even `""` all
   enable tracing. Only an unset variable or `FASTVIDEO_TRACE_ACTIVATIONS=0`
   disables it.
2. **Module exposure.** The hook attaches to `pipeline.modules.get("transformer")`
   at the end of `post_init`. If your pipeline does not expose a module
   under that key (e.g. a non-standard custom pipeline whose DiT is
   reachable only via `pipeline.modules["sr_transformer"]`), the trace is
   silently a no-op. Check the
   `Activation trace attached to N modules` log line at startup; if
   `N=0`, either the regex didn't match anything or the expected module
   isn't exposed.
3. **Output path failure.** The parent of `FASTVIDEO_TRACE_OUTPUT` is
   auto-created. If creation fails (permissions, read-only mount),
   `JsonlSink.__init__` raises at startup — look for an `OSError` early
   in the log.

### "My regex isn't filtering the way I expect"

`FASTVIDEO_TRACE_LAYERS` is compiled with `re.compile(spec)` and matched
with `pattern.search(name)`. Two consequences:

- `search`, not `fullmatch`. `block.layers` matches
  `transformer.block.layers.0.attn`. Anchor with `^...$` if you want
  exact matches.
- Module names use Python dot notation (`transformer.block.layers.0`),
  not slashes. The `.` in your regex is a metacharacter — escape it as
  `\.` if you want a literal dot.

Run once with `FASTVIDEO_LOGGING_LEVEL=DEBUG` to see the
`Activation trace attached to N modules (pattern=...)` line. `N` is the
ground truth for how many modules survived your regex.

### "Stats are NaN or `<error: ...>` for some layers"

Some outputs hit numerical edge cases:

- `mean` / `std` on a 0-dim tensor returns NaN.
- `abs_mean` on an empty tensor or one full of `inf` returns NaN.
- Non-tensor outputs (e.g. a Python `bool` from a verification gate) are
  silently skipped — the hook only walks `torch.Tensor` leaves.

Dropping `mean`/`std` and using `abs_mean`/`max` is more robust. For
layout-only debugging, `shape`+`dtype` never fail.

### "Tracing slows the run by 5x"

You're probably matching too broadly. An empty `FASTVIDEO_TRACE_LAYERS` is
treated as `.*` and matches every named module — for a 15B-param DiT
that's hundreds of submodules, each running stat reductions on every
forward. Tighten to a single block depth: `^transformer\.blocks\.\d+$`
typically matches a few dozen modules, which is a manageable trace.

### "Trace records are missing tensors I expect"

The hook walks `tuple` / `list` / `dict` outputs recursively but does
**not** unpack custom dataclasses or named tuples — those are silently
skipped. If your module returns
`BlockOutput(hidden_states=..., attn_logits=...)`, no records are
emitted. Workaround: either return a plain `dict` (`{"hidden_states": ..., "attn_logits": ...}`)
or attach the hook to a deeper module that already returns a raw tensor.

### "I want tracing inside a `torch.compile`'d region"

Module forward hooks run on the eager wrapper. If the entire module is
compiled, the hook sees only the wrapped op's output, not internal FX
nodes. This is by design — Extension 1 (FX backend rewrite) in
[Future extensions](#future-extensions-design-only-not-yet-implemented)
is the planned path when inside-graph granularity is required.

## Architecture (Extension 0: module forward hooks)

At pipeline initialization, `attach_activation_trace()` reads the env vars once.
If `FASTVIDEO_TRACE_ACTIVATIONS` is unset or false, the function returns
immediately and no hooks are registered. If tracing is on, it walks
`model.named_modules()`, filters by the layer regex, and registers an
`ActivationStatHook` on each matching module.

During the forward pass, each hook fires after its module completes, computes
the requested stats on the output tensor, and appends a JSON record to the
output file.

```
ComposedPipelineBase
  └─ attach_activation_trace()
       ├─ reads env vars (once at startup)
       ├─ if off: returns None immediately
       └─ if on: walks named_modules()
            └─ registers ActivationStatHook on matching modules
                 └─ on each forward: compute stats → append JSONL
```

### Zero-overhead-when-off guarantee

- The env var check happens **once at startup** inside `attach_activation_trace()`.
- If the env var is unset or false, the function returns `None` immediately.
- No hooks are registered. No branches are added to the production forward path.
- The only cost when tracing is off is one env var lookup at pipeline
  initialization, which takes under a microsecond.

### Hooking the upstream side

The upstream reference repo isn't part of FastVideo, so it can't read FastVideo
env vars directly. Two options:

**Option 1: Inline patch** in your local clone of the upstream repo. Add
`register_forward_hook` calls in the same shape as `ActivationStatHook`. Clean
up afterward with `git stash` or `git checkout HEAD -- <file>`.

**Option 2: Wrapper script**. Write a small Python harness that imports the
upstream model, walks its `named_modules()`, and attaches hooks externally.
This is the same pattern used in
`tests/local_tests/transformers/_debug_magi_human_block_parity.py`.

The `add-model-08-trace` skill at `.agents/skills/add-model-08-trace/`
provides a script template for this purpose.

## Future extensions (design only, not yet implemented)

### Extension 1: FX/Dynamo backend graph rewrite

**Granularity**: per-FX-node (every matmul, every add).

**Mechanism**: a `torch.compile` backend that takes the captured `GraphModule`
and inserts logger nodes after each op. Compiles into a separate artifact from
the production graph.

**Off semantics**: zero overhead. The production compile path is untouched.

**When to add**: if you need to trace inside a `torch.compile`'d graph and
Extension 0 is too coarse.

**Build cost**: roughly 1-2 days. Reference:
`torchao.quantization.pt2e._numeric_debugger`.

### Extension 2: AST source injection at import time

**Granularity**: per-line (between any two Python statements).

**Mechanism**: an importlib loader hook rewrites Python source AST at module
import time, inserting `if TRACE: dump(...)` statements. The decision is made
once at import.

**Off semantics**: zero overhead. If the env var is off at import time, source
is loaded as-is.

**When to add**: if you need per-line granularity that even FX-node-level can't
provide. This is almost never the right choice.

**Build cost**: roughly 1 week. Brittle and hard to debug.

### Extension 3: `__torch_dispatch__` / `TorchDispatchMode`

**Granularity**: per-op (every dispatcher call: matmul, add, view, etc.).

**Mechanism**: a `TorchDispatchMode` context manager that intercepts all ops at
the dispatcher level.

**Off semantics**: zero overhead. PyTorch's dispatcher only invokes mode hooks
when a mode is active.

**When on**: significant overhead. Every op pays a Python callback cost. Triton
kernels bypass it.

**When to add**: useful for quantization or dtype debugging where module-level
granularity isn't enough.

**Build cost**: roughly 1 day. Reference:
`torch.utils._python_dispatch.TorchDispatchMode`.

## Comparison with similar tools

| Tool | Pattern | FastVideo equivalent |
|---|---|---|
| SGLang `--debug-tensor-dump-output-folder` | env-gated forward hooks at startup | Extension 0 (this) |
| TransformerEngine `DumpTensors` | config-driven selective dumps | Extension 0 (env-driven) |
| HuggingFace `output_hidden_states=True` | source-level boolean gating | Not used; Extension 0 avoids model code edits |
| torchao numeric debugger | FX pass + node-level loggers | Extension 1 (future) |
| W&B `wandb.watch()` | runtime forward hooks (always on once registered) | Extension 0 has a similar mechanism, but gated off by default |

## Implementation references

- Module: `fastvideo/hooks/activation_trace.py`
- Env vars: `fastvideo/envs.py` (`FASTVIDEO_TRACE_ACTIVATIONS` and friends)
- Pipeline integration: `fastvideo/pipelines/composed_pipeline_base.py`
- Tests: `fastvideo/tests/hooks/test_activation_trace.py`
- Companion skill (for ad-hoc port investigations): `.agents/skills/add-model-08-trace/`

## Changelog

| Date | Change |
|---|---|
| 2026-05-01 | Initial Extension 0 (module forward hooks) implementation. Extensions 1-3 designed but not implemented. |