--- title: 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

fastvideo.hooks.activation_trace.ForwardHook

ForwardHook()

Base class for forward hooks. Hooks are used in the way: modified_args, modified_kwargs = hook.pre_forward(module, args, kwargs) output = module.forward(modified_args, **modified_kwargs) modified_output = hook.post_forward(module, output)

fastvideo.hooks.activation_trace.ForwardHook.name

name() -> str

fastvideo.hooks.activation_trace.ForwardHook.on_attach

on_attach(module: nn.Module)

Called once when the hook is attached to the module.

fastvideo.hooks.activation_trace.ForwardHook.on_detach

on_detach(module: nn.Module)

Called once when the hook is detached from the module. Note: this function is not guaranteed to be called if the module is deleted before the hook is detached.

fastvideo.hooks.activation_trace.ForwardHook.pre_forward

pre_forward(module: nn.Module, *args, **kwargs) -> tuple[tuple[Any, ...], dict[str, Any]]

Called before the module's forward method is executed.

fastvideo.hooks.activation_trace.ForwardHook.post_forward

post_forward(module: nn.Module, output: Any) -> Any

Called after the module's forward method is executed.

fastvideo.hooks.activation_trace.ModuleHookManager

ModuleHookManager(module: nn.Module)

fastvideo.hooks.activation_trace.ModuleHookManager.get_from

get_from(module: nn.Module) -> ModuleHookManager | None

fastvideo.hooks.activation_trace.ModuleHookManager.get_from_or_default

get_from_or_default(module: nn.Module) -> ModuleHookManager

fastvideo.hooks.activation_trace.ModuleHookManager.remove_from_manager

remove_from_manager(module: nn.Module) -> None

fastvideo.hooks.activation_trace.ModuleHookManager.append_forward_hook

append_forward_hook(hook: ForwardHook)

fastvideo.hooks.activation_trace.ModuleHookManager.replace_forward_hook

replace_forward_hook(hook_name: str, new_hook: ForwardHook, run_on_attach: bool = True)

fastvideo.hooks.activation_trace.ModuleHookManager.remove_forward_hook

remove_forward_hook(hook_name: str, run_detach: bool = True)

fastvideo.hooks.activation_trace.ModuleHookManager.get_forward_hook

get_forward_hook(hook_name: str) -> ForwardHook | None

fastvideo.hooks.activation_trace.init_logger

init_logger(name: str) -> _FastvideoLogger

The main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root fastvideo logger has already been configured.

fastvideo.hooks.activation_trace.current_step_idx

current_step_idx() -> int | None
Source
def current_step_idx() -> int | None:
    return getattr(_TRACE_STATE, "step_idx", None)

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
@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

fastvideo.hooks.activation_trace.JsonlSink

JsonlSink(path: Path)

Buffered JSONL writer with thread-safe append.

Source
class JsonlSink:
    """Buffered JSONL writer with thread-safe append."""

    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)

    def write(self, record: dict[str, Any]) -> None:
        line = json.dumps(record, default=str) + "\n"
        with self._lock:
            self._fh.write(line)

    def close(self) -> None:
        with self._lock:
            if not self._fh.closed:
                self._fh.close()

fastvideo.hooks.activation_trace.JsonlSink.write

write(record: dict[str, Any]) -> None
Source
def write(self, record: dict[str, Any]) -> None:
    line = json.dumps(record, default=str) + "\n"
    with self._lock:
        self._fh.write(line)

fastvideo.hooks.activation_trace.JsonlSink.close

close() -> None
Source
def close(self) -> None:
    with self._lock:
        if not self._fh.closed:
            self._fh.close()

fastvideo.hooks.activation_trace.ActivationStatHook

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

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

Source
class ActivationStatHook(ForwardHook):
    """Forward hook that emits per-tensor stats to a JSONL sink."""

    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

    def name(self) -> str:
        return "ActivationStatHook"

    def post_forward(self, module: nn.Module, output: Any) -> Any:
        step_idx = current_step_idx()
        if self.step_filter is not None and step_idx not in self.step_filter:
            return output
        for tensor_label, tensor in _flatten_tensors(output):
            record: dict[str, Any] = {
                "module": self.module_name,
                "tensor": tensor_label,
                "step": step_idx,
            }
            for stat_name, stat_fn in self.stats:
                try:
                    record[stat_name] = stat_fn(tensor)
                except Exception as exc:  # pragma: no cover - defensive logging
                    record[stat_name] = f"<error: {exc!r}>"
            self.sink.write(record)
        return output

fastvideo.hooks.activation_trace.ActivationStatHook.name

name() -> str
Source
def name(self) -> str:
    return "ActivationStatHook"

fastvideo.hooks.activation_trace.ActivationStatHook.post_forward

post_forward(module: nn.Module, output: Any) -> Any
Source
def post_forward(self, module: nn.Module, output: Any) -> Any:
    step_idx = current_step_idx()
    if self.step_filter is not None and step_idx not in self.step_filter:
        return output
    for tensor_label, tensor in _flatten_tensors(output):
        record: dict[str, Any] = {
            "module": self.module_name,
            "tensor": tensor_label,
            "step": step_idx,
        }
        for stat_name, stat_fn in self.stats:
            try:
                record[stat_name] = stat_fn(tensor)
            except Exception as exc:  # pragma: no cover - defensive logging
                record[stat_name] = f"<error: {exc!r}>"
        self.sink.write(record)
    return output

fastvideo.hooks.activation_trace.ActivationTraceManager

ActivationTraceManager(managers: list[ModuleHookManager], sink: JsonlSink)
Source
class ActivationTraceManager:

    def __init__(self, managers: list[ModuleHookManager], sink: JsonlSink) -> None:
        self.managers = managers
        self.sink = sink

    def remove_from_manager(self) -> None:
        for manager in self.managers:
            if manager.get_forward_hook("ActivationStatHook") is not None:
                manager.remove_forward_hook("ActivationStatHook")
            if not manager.forward_hooks:
                ModuleHookManager.remove_from_manager(manager.module)
        self.sink.close()

fastvideo.hooks.activation_trace.ActivationTraceManager.remove_from_manager

remove_from_manager() -> None
Source
def remove_from_manager(self) -> None:
    for manager in self.managers:
        if manager.get_forward_hook("ActivationStatHook") is not None:
            manager.remove_forward_hook("ActivationStatHook")
        if not manager.forward_hooks:
            ModuleHookManager.remove_from_manager(manager.module)
    self.sink.close()

fastvideo.hooks.activation_trace.attach_activation_trace

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

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

Source
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.get() or model is None:
        return None

    pattern_spec = envs.FASTVIDEO_TRACE_LAYERS.get()
    pattern = re.compile(pattern_spec) if pattern_spec else re.compile(".*")
    stats = _resolve_stats(envs.FASTVIDEO_TRACE_STATS.get())
    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.get()))
    step_filter = _parse_step_filter(envs.FASTVIDEO_TRACE_STEPS.get())
    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.detach_activation_trace

detach_activation_trace(mgr: ActivationTraceManager | None) -> None
Source
def detach_activation_trace(mgr: ActivationTraceManager | None) -> None:
    if mgr is not None:
        mgr.remove_from_manager()