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() -> strfastvideo.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) -> AnyCalled 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 | Nonefastvideo.hooks.activation_trace.ModuleHookManager.get_from_or_default
get_from_or_default(module: nn.Module) -> ModuleHookManagerfastvideo.hooks.activation_trace.ModuleHookManager.remove_from_manager
remove_from_manager(module: nn.Module) -> Nonefastvideo.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 | Nonefastvideo.hooks.activation_trace.init_logger
init_logger(name: str) -> _FastvideoLoggerThe 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 | NoneSource
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 = prevfastvideo.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()</code></pre></details>
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</code></pre></details>
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()</code></pre></details>
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)</code></pre></details>
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()