hao-ai-lab-FastVideo/docs/contributing/activation_trace.md

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)
Function-call tracing	FASTVIDEO_TRACE_FUNCTION (heavy)

Quickstart

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:

{"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:

   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:

   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.

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.