[feat] add zero-overhead-when-off activation trace mode (env-gated)
Add an env-gated activation trace mode that registers PyTorch forward hooks on selected modules, computes per-tensor stats (abs_mean, sum, min, max, mean, std, shape, dtype), and writes JSONL records to a configurable sink path. Designed for parity-debug across model ports: enable trace on both FastVideo's path and the upstream reference, then `diff` the two JSONL files to find the first divergent layer. Inspired by SGLang's --debug-tensor-dump-output-folder pattern and TransformerEngine's DumpTensors selective-dump infra. Zero-overhead-when-off guarantee: the master toggle FASTVIDEO_TRACE_ACTIVATIONS is checked ONCE at pipeline startup. When unset/false, attach_activation_trace() returns None and no hooks are ever registered. The production forward path is untouched. Configuration via 5 env vars (FASTVIDEO_TRACE_LAYERS regex filter, FASTVIDEO_TRACE_STATS, FASTVIDEO_TRACE_OUTPUT path with <pid> templating, FASTVIDEO_TRACE_STEPS step-index filter). Step indexing via the trace_step(step_idx) context manager wires step into thread-local for hook callbacks. Six unit tests cover off/on/filter/stats/step-filter/tuple-flattening /cleanup paths. End-to-end smoke against magi-human's basic example generated 28,864 JSONL records on 451 hooked modules; trace OFF does not create the output file.
This commit is contained in:
SolitaryThinker
@@ -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_<pid>.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 <pid> is replaced at runtime.
|
||||
"FASTVIDEO_TRACE_OUTPUT":
|
||||
lambda: os.getenv("FASTVIDEO_TRACE_OUTPUT", "/tmp/fv_trace_<pid>.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`
|
||||
|
||||
@@ -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("<pid>", 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"<error: {exc!r}>"
|
||||
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()
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user