diff --git a/fastvideo/envs.py b/fastvideo/envs.py index b6de4d588..cdab38026 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -35,6 +35,11 @@ if TYPE_CHECKING: FASTVIDEO_TORCH_PROFILER_WARMUP_STEPS: int = 1 FASTVIDEO_TORCH_PROFILER_ACTIVE_STEPS: int = 2 FASTVIDEO_TORCH_PROFILE_REGIONS: str = "" + FASTVIDEO_TRACE_ACTIVATIONS: bool = False + FASTVIDEO_TRACE_LAYERS: str = "" + FASTVIDEO_TRACE_STATS: str = "abs_mean,sum" + FASTVIDEO_TRACE_OUTPUT: str = "/tmp/fv_trace_.jsonl" + FASTVIDEO_TRACE_STEPS: str = "" FASTVIDEO_SERVER_DEV_MODE: bool = False FASTVIDEO_STAGE_LOGGING: bool = False FASTVIDEO_HOST_IP: str = "" @@ -252,6 +257,22 @@ environment_variables: dict[str, Callable[[], Any]] = { "FASTVIDEO_TORCH_PROFILE_REGIONS": lambda: os.getenv("FASTVIDEO_TORCH_PROFILE_REGIONS", ""), + # Enable activation trace hooks if set. + "FASTVIDEO_TRACE_ACTIVATIONS": + lambda: bool(os.getenv("FASTVIDEO_TRACE_ACTIVATIONS", "0") != "0"), + # Regex filter for traced module names. Empty means all modules. + "FASTVIDEO_TRACE_LAYERS": + lambda: os.getenv("FASTVIDEO_TRACE_LAYERS", ""), + # Comma-separated activation stats to dump for each output tensor. + "FASTVIDEO_TRACE_STATS": + lambda: os.getenv("FASTVIDEO_TRACE_STATS", "abs_mean,sum"), + # JSONL sink path. The literal is replaced at runtime. + "FASTVIDEO_TRACE_OUTPUT": + lambda: os.getenv("FASTVIDEO_TRACE_OUTPUT", "/tmp/fv_trace_.jsonl"), + # Comma-separated denoise step indices. Empty means all steps. + "FASTVIDEO_TRACE_STEPS": + lambda: os.getenv("FASTVIDEO_TRACE_STEPS", ""), + # If set, fastvideo will run in development mode, which will enable # some additional endpoints for developing and debugging, # e.g. `/reset_prefix_cache` diff --git a/fastvideo/hooks/activation_trace.py b/fastvideo/hooks/activation_trace.py new file mode 100644 index 000000000..c6a47a23b --- /dev/null +++ b/fastvideo/hooks/activation_trace.py @@ -0,0 +1,221 @@ +# SPDX-License-Identifier: Apache-2.0 +"""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 +""" +from __future__ import annotations + +import json +import os +import re +import threading +from contextlib import contextmanager +from pathlib import Path +from typing import Any +from collections.abc import Callable, Iterator + +import torch +from torch import nn + +from fastvideo import envs +from fastvideo.hooks.hooks import ForwardHook, ModuleHookManager +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + +_TRACE_STATE = threading.local() + + +def current_step_idx() -> int | None: + return getattr(_TRACE_STATE, "step_idx", None) + + +@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 + + +_STAT_FNS: dict[str, Callable[[torch.Tensor], Any]] = { + "abs_mean": lambda t: float(t.detach().float().abs().mean().item()), + "sum": lambda t: float(t.detach().float().sum().item()), + "min": lambda t: float(t.detach().float().min().item()), + "max": lambda t: float(t.detach().float().max().item()), + "mean": lambda t: float(t.detach().float().mean().item()), + "std": lambda t: float(t.detach().float().std().item()), + "shape": lambda t: list(t.shape), + "dtype": lambda t: str(t.dtype), +} + + +def _resolve_stats(spec: str) -> list[tuple[str, Callable[[torch.Tensor], Any]]]: + stats = [] + for name in [s.strip() for s in spec.split(",") if s.strip()]: + stat_fn = _STAT_FNS.get(name) + if stat_fn is None: + logger.warning( + "FASTVIDEO_TRACE_STATS contains unknown stat %r; valid: %s", + name, + sorted(_STAT_FNS), + ) + continue + stats.append((name, stat_fn)) + return stats + + +def _resolve_output_path(template: str) -> Path: + return Path(template.replace("", str(os.getpid()))) + + +def _parse_step_filter(spec: str) -> set[int] | None: + if not spec.strip(): + return None + return {int(s.strip()) for s in spec.split(",") if s.strip()} + + +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() + + +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"" + self.sink.write(record) + return output + + +def _flatten_tensors(obj: Any, prefix: str = "out") -> list[tuple[str, torch.Tensor]]: + """Yield (label, tensor) pairs from arbitrarily-nested forward outputs.""" + if isinstance(obj, torch.Tensor): + return [(prefix, obj)] + if isinstance(obj, tuple | list): + out = [] + for idx, item in enumerate(obj): + out.extend(_flatten_tensors(item, f"{prefix}[{idx}]")) + return out + if isinstance(obj, dict): + out = [] + for key, value in obj.items(): + out.extend(_flatten_tensors(value, f"{prefix}.{key}")) + return out + return [] + + +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() + + +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) + + +def detach_activation_trace(mgr: ActivationTraceManager | None) -> None: + if mgr is not None: + mgr.remove_from_manager() diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index 079ea169a..f0b000252 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -16,6 +16,7 @@ from fastvideo.configs.pipelines import PipelineConfig from fastvideo.distributed import (maybe_init_distributed_environment_and_model_parallel, get_world_group) from fastvideo.distributed.communication_op import (warmup_sequence_parallel_communication) from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs +from fastvideo.hooks.activation_trace import attach_activation_trace, detach_activation_trace from fastvideo.logger import init_logger from fastvideo.profiler import get_or_create_profiler from fastvideo.models.loader.component_loader import PipelineComponentLoader @@ -62,6 +63,7 @@ class ComposedPipelineBase(ABC): self.model_path: str = model_path self._stages: list[PipelineStage] = [] self._stage_name_mapping: dict[str, PipelineStage] = {} + self._trace_mgr = None if required_config_modules is not None: self._required_config_modules = required_config_modules @@ -183,6 +185,8 @@ class ComposedPipelineBase(ABC): ) logger.info("Torch Compile enabled for DiT") + self._trace_mgr = attach_activation_trace(self.modules.get("transformer")) + if not self.fastvideo_args.training_mode: logger.info("Creating pipeline stages...") self.create_pipeline_stages(self.fastvideo_args) @@ -455,3 +459,10 @@ class ComposedPipelineBase(ABC): def train(self) -> None: raise NotImplementedError("if training_mode is True, the pipeline must implement this method") + + def close(self) -> None: + detach_activation_trace(getattr(self, "_trace_mgr", None)) + self._trace_mgr = None + + def __del__(self): + self.close() diff --git a/fastvideo/tests/hooks/test_activation_trace.py b/fastvideo/tests/hooks/test_activation_trace.py new file mode 100644 index 000000000..e5da540ef --- /dev/null +++ b/fastvideo/tests/hooks/test_activation_trace.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 + +import json +from pathlib import Path + +import torch +from torch import nn + +from fastvideo.hooks.activation_trace import ( + attach_activation_trace, + detach_activation_trace, + trace_step, +) +from fastvideo.hooks.hooks import ModuleHookManager + + +class ToyModel(nn.Module): + + def __init__(self) -> None: + super().__init__() + self.block = nn.Sequential(nn.Linear(2, 2), nn.ReLU()) + self.other = nn.Linear(2, 2) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.other(self.block(x)) + + +class TupleLayer(nn.Module): + + def __init__(self) -> None: + super().__init__() + self.proj = nn.Linear(2, 2) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + out = self.proj(x) + return out, out + 1 + + +class TupleOutputModel(nn.Module): + + def __init__(self) -> None: + super().__init__() + self.tuple = TupleLayer() + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return self.tuple(x) + + +def _read_jsonl(path: Path) -> list[dict]: + return [json.loads(line) for line in path.read_text().splitlines()] + + +def test_attach_activation_trace_off_returns_none(monkeypatch) -> None: + monkeypatch.delenv("FASTVIDEO_TRACE_ACTIVATIONS", raising=False) + model = ToyModel() + + manager = attach_activation_trace(model) + + assert manager is None + assert len(model._forward_hooks) == 0 + assert ModuleHookManager.get_from(model.block[0]) is None + + +def test_attach_activation_trace_on_respects_layer_filter( + monkeypatch, + tmp_path, +) -> None: + monkeypatch.setenv("FASTVIDEO_TRACE_ACTIVATIONS", "1") + monkeypatch.setenv("FASTVIDEO_TRACE_LAYERS", r"block\.0.*") + monkeypatch.setenv("FASTVIDEO_TRACE_OUTPUT", str(tmp_path / "trace.jsonl")) + model = ToyModel() + + manager = attach_activation_trace(model) + + try: + assert manager is not None + assert ModuleHookManager.get_from(model.block[0]) is not None + assert ModuleHookManager.get_from(model.block[1]) is None + assert ModuleHookManager.get_from(model.other) is None + finally: + detach_activation_trace(manager) + + +def test_activation_trace_writes_configured_stats(monkeypatch, tmp_path) -> None: + path = tmp_path / "trace.jsonl" + monkeypatch.setenv("FASTVIDEO_TRACE_ACTIVATIONS", "1") + monkeypatch.setenv("FASTVIDEO_TRACE_LAYERS", r"block\.0") + monkeypatch.setenv("FASTVIDEO_TRACE_STATS", "abs_mean,sum,shape,dtype") + monkeypatch.setenv("FASTVIDEO_TRACE_OUTPUT", str(path)) + model = ToyModel() + manager = attach_activation_trace(model) + + try: + with trace_step(3): + model(torch.ones(1, 2)) + finally: + detach_activation_trace(manager) + + records = _read_jsonl(path) + assert len(records) == 1 + record = records[0] + assert record["module"] == "block.0" + assert record["tensor"] == "out" + assert record["step"] == 3 + assert {"abs_mean", "sum", "shape", "dtype"}.issubset(record) + assert record["shape"] == [1, 2] + assert record["dtype"] == "torch.float32" + + +def test_activation_trace_step_filter(monkeypatch, tmp_path) -> None: + path = tmp_path / "trace.jsonl" + monkeypatch.setenv("FASTVIDEO_TRACE_ACTIVATIONS", "1") + monkeypatch.setenv("FASTVIDEO_TRACE_LAYERS", r"block\.0") + monkeypatch.setenv("FASTVIDEO_TRACE_OUTPUT", str(path)) + monkeypatch.setenv("FASTVIDEO_TRACE_STEPS", "0,2") + model = ToyModel() + manager = attach_activation_trace(model) + + try: + for step_idx in range(4): + with trace_step(step_idx): + model(torch.ones(1, 2)) + finally: + detach_activation_trace(manager) + + assert [record["step"] for record in _read_jsonl(path)] == [0, 2] + + +def test_activation_trace_flattens_tuple_outputs(monkeypatch, tmp_path) -> None: + path = tmp_path / "trace.jsonl" + monkeypatch.setenv("FASTVIDEO_TRACE_ACTIVATIONS", "1") + monkeypatch.setenv("FASTVIDEO_TRACE_LAYERS", "tuple$") + monkeypatch.setenv("FASTVIDEO_TRACE_OUTPUT", str(path)) + model = TupleOutputModel() + manager = attach_activation_trace(model) + + try: + model(torch.ones(1, 2)) + finally: + detach_activation_trace(manager) + + records = _read_jsonl(path) + assert [record["tensor"] for record in records] == ["out[0]", "out[1]"] + + +def test_detach_activation_trace_removes_hooks(monkeypatch, tmp_path) -> None: + monkeypatch.setenv("FASTVIDEO_TRACE_ACTIVATIONS", "1") + monkeypatch.setenv("FASTVIDEO_TRACE_LAYERS", r"block\.0") + monkeypatch.setenv("FASTVIDEO_TRACE_OUTPUT", str(tmp_path / "trace.jsonl")) + model = ToyModel() + manager = attach_activation_trace(model) + + assert ModuleHookManager.get_from(model.block[0]) is not None + + detach_activation_trace(manager) + + assert ModuleHookManager.get_from(model.block[0]) is None