activation_trace

Zero-overhead-when-off activation trace mode for FastVideo pipelines.

Enable by setting FASTVIDEO_TRACE_ACTIVATIONS=1. When off, this module adds zero overhead — no hooks are registered, no branches exist in the production forward path. When on, registers forward hooks on modules whose name matches FASTVIDEO_TRACE_LAYERS, computes the requested stats (FASTVIDEO_TRACE_STATS) on each output tensor, and writes JSONL records to FASTVIDEO_TRACE_OUTPUT.

Useful for parity debugging across model ports — log on both the FastVideo path and the upstream reference, diff the two JSONL files to find the first divergent layer.

Example:

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

Classes

fastvideo.hooks.activation_trace.ActivationStatHook

ActivationStatHook(module_name: str, stats: list[tuple[str, Callable[[Tensor], Any]]], sink: JsonlSink, step_filter: set[int] | None)

Bases: ForwardHook

Forward hook that emits per-tensor stats to a JSONL sink.

Source code in fastvideo/hooks/activation_trace.py

def __init__(
    self,
    module_name: str,
    stats: list[tuple[str, Callable[[torch.Tensor], Any]]],
    sink: JsonlSink,
    step_filter: set[int] | None,
) -> None:
    self.module_name = module_name
    self.stats = stats
    self.sink = sink
    self.step_filter = step_filter

fastvideo.hooks.activation_trace.JsonlSink

JsonlSink(path: Path)

Buffered JSONL writer with thread-safe append.

Source code in fastvideo/hooks/activation_trace.py

def __init__(self, path: Path) -> None:
    self.path = path
    self.path.parent.mkdir(parents=True, exist_ok=True)
    self._fh = open(self.path, "a", buffering=1)  # noqa: SIM115
    self._lock = threading.Lock()
    logger.info("Activation trace JSONL sink: %s", self.path)

Functions

fastvideo.hooks.activation_trace.attach_activation_trace

attach_activation_trace(model: Module | None) -> ActivationTraceManager | None

Attach activation-stat hooks to model. Returns None if trace is off.

Source code in fastvideo/hooks/activation_trace.py

def attach_activation_trace(model: nn.Module | None) -> ActivationTraceManager | None:
    """Attach activation-stat hooks to model. Returns None if trace is off."""
    if not envs.FASTVIDEO_TRACE_ACTIVATIONS or model is None:
        return None

    pattern_spec = envs.FASTVIDEO_TRACE_LAYERS
    pattern = re.compile(pattern_spec) if pattern_spec else re.compile(".*")
    stats = _resolve_stats(envs.FASTVIDEO_TRACE_STATS)
    if not stats:
        logger.warning("FASTVIDEO_TRACE_STATS yielded no valid stats; trace disabled.")
        return None

    sink = JsonlSink(_resolve_output_path(envs.FASTVIDEO_TRACE_OUTPUT))
    step_filter = _parse_step_filter(envs.FASTVIDEO_TRACE_STEPS)
    managers = []
    for name, module in model.named_modules():
        if not name or not pattern.search(name):
            continue
        manager = ModuleHookManager.get_from_or_default(module)
        manager.append_forward_hook(ActivationStatHook(name, stats, sink, step_filter))
        managers.append(manager)

    logger.info(
        "Activation trace attached to %d modules (pattern=%r, stats=%s)",
        len(managers),
        pattern_spec,
        [stat_name for stat_name, _ in stats],
    )
    return ActivationTraceManager(managers, sink)

fastvideo.hooks.activation_trace.trace_step

trace_step(step_idx: int) -> Iterator[None]

Context manager that sets the current denoise step for trace records.

Source code in fastvideo/hooks/activation_trace.py

@contextmanager
def trace_step(step_idx: int) -> Iterator[None]:
    """Context manager that sets the current denoise step for trace records."""
    prev = getattr(_TRACE_STATE, "step_idx", None)
    _TRACE_STATE.step_idx = step_idx
    try:
        yield
    finally:
        _TRACE_STATE.step_idx = prev