[feature] Add torch profiler (#827)

docs/source/contributing/profiling.md

@@ -0,0 +1,53 @@
+# Profiling FastVideo
+
+!!! warning
+    Profiling is only intended for FastVideo developers and maintainers to understand the proportion of time spent in different parts of the codebase. **FastVideo end-users should never turn on profiling** as it will significantly slow down the inference.
+
+## Profiling with PyTorch
+
+FastVideo exposes a process-wide torch profiler that you can enable via environment variables. Set `FASTVIDEO_TORCH_PROFILER_DIR` to an absolute directory path to start collecting traces, and specify the regions you want recorded with `FASTVIDEO_TORCH_PROFILE_REGIONS`:
+
+```bash
+FASTVIDEO_TORCH_PROFILER_DIR=/mnt/traces/fastvideo \
+FASTVIDEO_TORCH_PROFILE_REGIONS="profiler_region_model_loading,profiler_region_training_step"
+```
+
+All profiled regions must be registered in `fastvideo.profiler`; the current list includes:
+
+- `profiler_region_model_loading` — pipeline/module loading
+- `profiler_region_inference_pre_denoising`
+- `profiler_region_inference_denoising`
+- `profiler_region_inference_post_denoising`
+- `profiler_region_training_checkpoint_saving`
+- `profiler_region_training_dit`
+- `profiler_region_training_validation`
+- `profiler_region_training_epoch`
+- `profiler_region_training_step`
+- `profiler_region_training_backward`
+- `profiler_region_training_optimizer`
+- `profiler_region_distillation_teacher_forward`
+- `profiler_region_distillation_student_forward`
+- `profiler_region_distillation_loss`
+- `profiler_region_distillation_update`
+
+While profiling is enabled, FastVideo records additional annotations:
+
+- `fastvideo.region::<name>` spans are emitted when entering a region.
+- `fastvideo.profiler.enable_collection` / `fastvideo.profiler.disable_collection` events mark when torch profiler collection is toggled on or off.
+
+Only one profiler instance is created per process; subsequent pipelines reuse the same controller. If you set `FASTVIDEO_TORCH_PROFILE_REGIONS` incorrectly (e.g. misspelled name), FastVideo logs a warning and ignores that entry.
+
+Additional knobs:
+
+- `FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES`
+- `FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY`
+- `FASTVIDEO_TORCH_PROFILER_WITH_STACK`
+- `FASTVIDEO_TORCH_PROFILER_WITH_FLOPS`
+
+Traces can be visualized using <https://ui.perfetto.dev/>.
+
+### Best Practices
+
+- Keep the profiled step count small; traces can be large and slow down job shutdown while the profiler flushes data.
+- After profiling, clean up trace directories to avoid filling disks.
+- When adding new regions, register them in `fastvideo.profiler` and wrap the corresponding code block with `with self.profiler_controller.region("your_region"):` or the `@profile_region` decorator.

docs/source/index.md

@@ -115,6 +115,7 @@ design/overview
 
 contributing/overview
 contributing/developer_env/index
+contributing/profiling
 :::
 
 :::{toctree}

fastvideo/envs.py

@@ -26,6 +26,12 @@ if TYPE_CHECKING:
     NVCC_THREADS: str | None = None
     CMAKE_BUILD_TYPE: str | None = None
     VERBOSE: bool = False
+    FASTVIDEO_TORCH_PROFILER_DIR: str | None = None
+    FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES: bool = False
+    FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY: bool = False
+    FASTVIDEO_TORCH_PROFILER_WITH_STACK: bool = True
+    FASTVIDEO_TORCH_PROFILER_WITH_FLOPS: bool = False
+    FASTVIDEO_TORCH_PROFILE_REGIONS: str = ""
     FASTVIDEO_SERVER_DEV_MODE: bool = False
     FASTVIDEO_STAGE_LOGGING: bool = False
 
@@ -197,6 +203,34 @@ environment_variables: dict[str, Callable[[], Any]] = {
              if os.getenv("FASTVIDEO_TORCH_PROFILER_DIR", None) is None else os.
              path.expanduser(os.getenv("FASTVIDEO_TORCH_PROFILER_DIR", "."))),
 
+    # Enable torch profiler to record shapes if set
+    # FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES=1. If not set, torch profiler will
+    # not record shapes.
+    "FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES":
+    lambda: bool(
+        os.getenv("FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES", "0") != "0"),
+
+    # Enable torch profiler to profile memory if set
+    # FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY=1. If not set, torch profiler
+    # will not profile memory.
+    "FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY":
+    lambda: bool(
+        os.getenv("FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY", "0") != "0"),
+
+    # Enable torch profiler to profile stack if set
+    # FASTVIDEO_TORCH_PROFILER_WITH_STACK=1. If not set, torch profiler WILL
+    # profile stack by default.
+    "FASTVIDEO_TORCH_PROFILER_WITH_STACK":
+    lambda: bool(os.getenv("FASTVIDEO_TORCH_PROFILER_WITH_STACK", "1") != "0"),
+
+    # Enable torch profiler to profile flops if set
+    # FASTVIDEO_TORCH_PROFILER_WITH_FLOPS=1. If not set, torch profiler will
+    # not profile flops.
+    "FASTVIDEO_TORCH_PROFILER_WITH_FLOPS":
+    lambda: bool(os.getenv("FASTVIDEO_TORCH_PROFILER_WITH_FLOPS", "0") != "0"),
+    "FASTVIDEO_TORCH_PROFILE_REGIONS":
+    lambda: os.getenv("FASTVIDEO_TORCH_PROFILE_REGIONS", ""),
+
     # If set, fastvideo will run in development mode, which will enable
     # some additional endpoints for developing and debugging,
     # e.g. `/reset_prefix_cache`

fastvideo/fastvideo_args.py

@@ -421,6 +421,7 @@ class FastVideoArgs:
             "--init-weights-from-safetensors-2",
             type=str,
             help="Path to safetensors file for initial weight loading")
+
         # Add pipeline configuration arguments
         PipelineConfig.add_cli_args(parser)
 

fastvideo/pipelines/composed_pipeline_base.py

@@ -14,12 +14,14 @@ import torch
 
 from fastvideo.configs.pipelines import PipelineConfig
 from fastvideo.distributed import (
-    maybe_init_distributed_environment_and_model_parallel)
+    maybe_init_distributed_environment_and_model_parallel, get_world_group)
 from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.logger import init_logger
+from fastvideo.profiler import get_or_create_profiler
 from fastvideo.models.loader.component_loader import PipelineComponentLoader
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.pipelines.stages import PipelineStage
+import fastvideo.envs as envs
 from fastvideo.utils import (maybe_download_model,
                              verify_model_config_and_directory)
 
@@ -69,9 +71,18 @@ class ComposedPipelineBase(ABC):
         maybe_init_distributed_environment_and_model_parallel(
             fastvideo_args.tp_size, fastvideo_args.sp_size)
 
+        # Torch profiler. Enabled and configured through env vars:
+        # FASTVIDEO_TORCH_PROFILER_DIR=/path/to/save/trace
+        trace_dir = envs.FASTVIDEO_TORCH_PROFILER_DIR
+        self.profiler_controller = get_or_create_profiler(trace_dir)
+        self.profiler = self.profiler_controller.profiler
+
+        self.local_rank = get_world_group().local_rank
+
         # Load modules directly in initialization
         logger.info("Loading pipeline modules...")
-        self.modules = self.load_modules(fastvideo_args, loaded_modules)
+        with self.profiler_controller.region("profiler_region_model_loading"):
+            self.modules = self.load_modules(fastvideo_args, loaded_modules)
 
     def set_trainable(self) -> None:
         # Only train DiT
@@ -352,6 +363,18 @@ class ComposedPipelineBase(ABC):
         self._stage_name_mapping[stage_name] = stage
         setattr(self, stage_name, stage)
 
+    def profile(self, is_start: bool = True):
+        if self.profiler is None:
+            raise RuntimeError("Profiler is not enabled.")
+        if is_start:
+            self.profiler.start()
+        else:
+            self.profiler.stop()
+            # only print profiler results on rank 0
+            if self.local_rank == 0:
+                print(self.profiler.key_averages().table(
+                    sort_by="self_cuda_time_total"))
+
     # TODO(will): don't hardcode no_grad
     @torch.no_grad()
     def forward(

fastvideo/profiler.py

@@ -0,0 +1,451 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Utilities for managing the PyTorch profiler within FastVideo.
+
+The profiler is shared across the process; this module adds a light-weight
+controller that gates collection based on named *regions*. Regions may be
+enabled through dedicated environment variables (e.g.
+``FASTVIDEO_TORCH_PROFILE_MODEL_LOADING=1``) or via the consolidated
+``FASTVIDEO_TORCH_PROFILE_REGIONS`` comma-separated list (e.g.
+``FASTVIDEO_TORCH_PROFILE_REGIONS=model_loading,training_dit``).
+
+Typical usage from client code::
+
+    controller = TorchProfilerController(profiler, activities)
+    with controller.region("training_dit"):
+        run_training_step()
+
+To introduce a new region, register it via :func:`register_profiler_region`
+and wrap the corresponding code in :meth:`TorchProfilerController.region`.
+"""
+
+from __future__ import annotations
+
+import contextlib
+from dataclasses import dataclass
+from typing import Any
+from collections.abc import Callable
+import functools
+from collections.abc import Iterable
+
+import torch
+
+import fastvideo.envs as envs
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+_GLOBAL_PROFILER: torch.profiler.profile | None = None
+_GLOBAL_CONTROLLER: TorchProfilerController | None = None
+
+
+@dataclass(frozen=True)
+class ProfilerRegion:
+    """Metadata describing a profiler region."""
+
+    name: str
+    description: str
+    default_enabled: bool = False
+
+    def __post_init__(self) -> None:
+        if not self.name or self.name.strip() != self.name:
+            raise ValueError(
+                f"Profiler region name must be non-empty without surrounding whitespace: {self.name!r}"
+            )
+        if not self.name.islower():
+            raise ValueError(
+                f"Profiler region name must be lower-case: {self.name!r}")
+
+
+_REGISTERED_REGIONS: dict[str, ProfilerRegion] = {}
+
+
+def _normalize_token(token: str) -> str:
+    return token.strip().lower()
+
+
+def register_profiler_region(
+    name: str,
+    description: str,
+    *,
+    default_enabled: bool = False,
+) -> None:
+    """Register a profiler region so configuration can validate inputs."""
+
+    canonical = _normalize_token(name)
+    if canonical in _REGISTERED_REGIONS:
+        raise ValueError(f"Profiler region {name!r} is already registered")
+
+    region = ProfilerRegion(
+        name=canonical,
+        description=description,
+        default_enabled=bool(default_enabled),
+    )
+    _REGISTERED_REGIONS[canonical] = region
+
+
+def resolve_profiler_region(name: str) -> ProfilerRegion | None:
+    """Return the registered region matching ``name`` or ``None`` if absent."""
+
+    canonical = _normalize_token(name)
+    return _REGISTERED_REGIONS.get(canonical)
+
+
+def list_profiler_regions() -> list[ProfilerRegion]:
+    """Return all registered profiler regions sorted by canonical name."""
+
+    return [_REGISTERED_REGIONS[name] for name in sorted(_REGISTERED_REGIONS)]
+
+
+_DEFAULT_ACTIVITIES: tuple[torch.profiler.ProfilerActivity, ...] = (
+    torch.profiler.ProfilerActivity.CPU,
+    torch.profiler.ProfilerActivity.CUDA,
+)
+
+
+def get_global_profiler() -> torch.profiler.profile | None:
+    """Return the global profiler instance if one was created."""
+
+    return _GLOBAL_PROFILER
+
+
+def set_global_profiler(profiler: torch.profiler.profile | None) -> None:
+    global _GLOBAL_PROFILER
+    _GLOBAL_PROFILER = profiler
+
+
+def get_global_controller() -> TorchProfilerController | None:
+    return _GLOBAL_CONTROLLER
+
+
+def set_global_controller(controller: TorchProfilerController | None) -> None:
+    global _GLOBAL_CONTROLLER
+    _GLOBAL_CONTROLLER = controller
+
+
+register_profiler_region(
+    name="profiler_region_model_loading",
+    description="Module/model loading during pipeline initialization.",
+    default_enabled=False,
+)
+# register_profiler_region(
+#     name="profiler_region_inference_pre_denoising",
+#     description="Pre-denoising inference steps (conditioning, preprocessing).",
+# )
+# register_profiler_region(
+#     name="profiler_region_inference_denoising",
+#     description="The main inference denoising loop.",
+# )
+# register_profiler_region(
+#     name="profiler_region_inference_post_denoising",
+#     description=
+#     "Post-processing after denoising (decoder, conditioning restores).",
+# )
+register_profiler_region(
+    name="profiler_region_training_save_checkpoint",
+    description="Training save checkpoint operations.",
+)
+
+# general training related regions
+register_profiler_region(
+    name="profiler_region_training_validation",
+    description="Validation loop during training.",
+)
+register_profiler_region(
+    name="profiler_region_training_train_one_step",
+    description="High-level step orchestration in the training loop.",
+)
+register_profiler_region(
+    name="profiler_region_training_train",
+    description="Single optimizer step including forward/backward passes.",
+)
+
+# distillation specific regions
+register_profiler_region(
+    name="profiler_region_distillation_teacher_forward",
+    description="Teacher model forward pass in distillation pipelines.",
+)
+register_profiler_region(
+    name="profiler_region_distillation_student_forward",
+    description="Student model forward pass in distillation pipelines.",
+)
+register_profiler_region(
+    name="profiler_region_distillation_loss",
+    description="Distillation loss computation and aggregation.",
+)
+register_profiler_region(
+    name="profiler_region_distillation_update",
+    description="Parameter updates specific to distillation workflows.",
+)
+
+
+def get_or_create_profiler(trace_dir: str | None) -> TorchProfilerController:
+    """Create or reuse the process-wide torch profiler controller."""
+
+    existing = get_global_controller()
+    if existing is not None:
+        if trace_dir:
+            logger.info("Reusing existing global torch profiler controller")
+        return existing
+
+    if not trace_dir:
+        logger.info("Torch profiler disabled; returning no-op controller")
+        return TorchProfilerController(None, _DEFAULT_ACTIVITIES, disabled=True)
+
+    logger.info("Profiling enabled. Traces will be saved to: %s", trace_dir)
+    logger.info(
+        "Profiler config: record_shapes=%s, profile_memory=%s, with_stack=%s, with_flops=%s",
+        envs.FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES,
+        envs.FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY,
+        envs.FASTVIDEO_TORCH_PROFILER_WITH_STACK,
+        envs.FASTVIDEO_TORCH_PROFILER_WITH_FLOPS,
+    )
+    logger.info("FASTVIDEO_TORCH_PROFILE_REGIONS=%s",
+                envs.FASTVIDEO_TORCH_PROFILE_REGIONS)
+
+    profiler = torch.profiler.profile(
+        activities=_DEFAULT_ACTIVITIES,
+        record_shapes=envs.FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES,
+        profile_memory=envs.FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY,
+        with_stack=envs.FASTVIDEO_TORCH_PROFILER_WITH_STACK,
+        with_flops=envs.FASTVIDEO_TORCH_PROFILER_WITH_FLOPS,
+        on_trace_ready=torch.profiler.tensorboard_trace_handler(trace_dir,
+                                                                use_gzip=True),
+    )
+    controller = TorchProfilerController(profiler, _DEFAULT_ACTIVITIES)
+    controller.start()
+    logger.info("Torch profiler started")
+    return controller
+
+
+@dataclass
+class TorchProfilerConfig:
+    """Configuration for torch profiler region control.
+
+    Use :meth:`from_env` to construct an instance with defaults inherited from
+    registered regions and optional overrides from the
+    ``FASTVIDEO_TORCH_PROFILE_REGIONS`` environment variable. The resulting
+    ``regions`` map is consumed by :class:`TorchProfilerController` to decide
+    when collection should be enabled.
+    """
+
+    regions: dict[str, bool]
+
+    @classmethod
+    def from_env(cls) -> TorchProfilerConfig:
+        """Build a configuration from process environment variables."""
+
+        requested_regions = {
+            token.strip()
+            for token in (getattr(envs, "FASTVIDEO_TORCH_PROFILE_REGIONS", "")
+                          or "").split(",") if token.strip()
+        }
+
+        if not requested_regions:
+            available = ", ".join(region.name
+                                  for region in list_profiler_regions())
+            raise ValueError(
+                "FASTVIDEO_TORCH_PROFILE_REGIONS must list at least one region; "
+                f"available regions: {available}")
+
+        regions: dict[str, bool] = {}
+        available_regions = list_profiler_regions()
+        available_names = ", ".join(region.name for region in available_regions)
+
+        for token in requested_regions:
+            resolved = resolve_profiler_region(token)
+            if resolved is None:
+                logger.warning(
+                    "Unknown profiler region '%s'; available regions: %s",
+                    token, available_names)
+                continue
+            regions[resolved.name] = True
+
+        if not regions:
+            raise ValueError(
+                "FASTVIDEO_TORCH_PROFILE_REGIONS did not match any known regions; "
+                f"requested={sorted(requested_regions)}, available={available_names}"
+            )
+
+        return cls(regions=regions)
+
+    def __str__(self) -> str:
+        return f"TorchProfilerConfig(regions={self.regions})"
+
+
+class TorchProfilerController:
+    """Helper that toggles torch profiler collection for named regions.
+
+    Parameters
+    ----------
+    profiler:
+        The shared :class:`torch.profiler.profile` instance, or ``None`` if
+        profiling is disabled.
+    activities:
+        Iterable of :class:`torch.profiler.ProfilerActivity` recorded by the
+        profiler.
+    config:
+        Optional :class:`TorchProfilerConfig`. If omitted, :meth:`from_env`
+        constructs one during initialization.
+
+    Examples
+    --------
+    Enabling an existing region from the command line::
+
+        FASTVIDEO_TORCH_PROFILE_REGIONS=model_loading,training_dit \
+        python fastvideo/training/wan_training_pipeline.py ...
+
+    Wrapping a code block in a custom region::
+
+        controller = TorchProfilerController(profiler, activities)
+        with controller.region("training_validation"):
+            run_validation_epoch()
+
+    Adding a new region requires three steps:
+      1. Define an env var in ``envs.py``.
+      2. Add a default entry to ``register_profiler_region`` in this module.
+      3. Wrap the target code in :meth:`region` using the new name.
+    """
+
+    def __init__(
+        self,
+        profiler: Any,
+        activities: Iterable[torch.profiler.ProfilerActivity],
+        config: TorchProfilerConfig | None = None,
+        disabled: bool = False,
+    ) -> None:
+        activities_tuple = tuple(activities)
+        existing = get_global_controller()
+        if existing is not None and not disabled:
+            raise RuntimeError(
+                "TorchProfilerController already initialized globally. Use get_global_controller()."
+            )
+        if disabled:
+            self._profiler = None
+            return
+
+        self._profiler = profiler
+        self._activities = activities_tuple
+        self._config = config or TorchProfilerConfig.from_env()
+        self._collection_enabled = False
+        self._active_region_depth = 0
+        logger.info(
+            "PROFILER: TorchProfilerController initialized with config: %s",
+            self._config)
+        set_global_profiler(self._profiler)
+        set_global_controller(self)
+
+    @property
+    def is_enabled(self) -> bool:
+        """Return ``True`` when the underlying profiler is collecting."""
+
+        if self._profiler is None:
+            return False
+        return self._collection_enabled
+
+    def is_region_enabled(self, region: str) -> bool:
+        """Return ``True`` if ``region`` should be collected."""
+
+        if self._profiler is None:
+            return False
+        return self._config.regions.get(region, False)
+
+    def _set_collection(self, enabled: bool) -> None:
+        if self._profiler is None:
+            return
+        if self._collection_enabled == enabled:
+            return
+        event = ("fastvideo.profiler.enable_collection"
+                 if enabled else "fastvideo.profiler.disable_collection")
+        with torch.profiler.record_function(event):
+            self._profiler.toggle_collection_dynamic(enabled, self._activities)
+        self._collection_enabled = enabled
+
+    @contextlib.contextmanager
+    def region(self, region: str):
+        """Context manager that enables profiling for ``region`` if configured."""
+
+        if self._profiler is None:
+            yield
+            return
+
+        if not self.is_region_enabled(region):
+            yield
+            return
+
+        with torch.profiler.record_function(f"fastvideo.region::{region}"):
+            self._active_region_depth += 1
+            if self._active_region_depth == 1:
+                logger.info(
+                    "PROFILER: Setting collection to True (depth=%s) for region %s",
+                    self._active_region_depth, region)
+                self._set_collection(True)
+            try:
+                yield
+            finally:
+                self._active_region_depth -= 1
+                logger.info("PROFILER: Decreasing active region depth to %s",
+                            self._active_region_depth)
+                if self._active_region_depth == 0:
+                    logger.info(
+                        "PROFILER: Setting collection to False upon exiting region %s",
+                        region)
+                    self._set_collection(False)
+
+    def start(self) -> None:
+        """Start the profiler and pause collection until a region is entered."""
+
+        logger.info("PROFILER: Starting profiler...")
+        if self._profiler is None:
+            return
+        self._profiler.start()
+        logger.info("PROFILER: Profiler started")
+        # Profiler starts with collection disabled by default.
+        logger.info("PROFILER: Setting collection to False")
+        self._set_collection(False)
+        logger.info("PROFILER: Profiler started with collection disabled")
+
+    def stop(self) -> None:
+        """Stop the profiler after disabling collection and clearing state."""
+
+        if self._profiler is None:
+            return
+
+        logger.info("PROFILER: Stopping profiler...")
+        self._profiler.stop()
+        logger.info("PROFILER: Profiler stopped")
+        self._active_region_depth = 0
+        set_global_profiler(None)
+        set_global_controller(None)
+
+    @property
+    def has_profiler(self) -> bool:
+        """Return ``True`` when a profiler instance is available."""
+
+        return self._profiler is not None
+
+    @property
+    def activities(self) -> tuple[torch.profiler.ProfilerActivity, ...]:
+        return tuple(self._activities)
+
+    @property
+    def profiler(self) -> torch.profiler.profile | None:
+        return self._profiler
+
+
+def profile_region(
+        region: str) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    """Wrap a bound method so it runs inside a profiler region if available."""
+
+    def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
+
+        @functools.wraps(fn)
+        def wrapped(self, *args, **kwargs):
+            controller = getattr(self, "profiler_controller", None)
+            if controller is None or not controller.has_profiler:
+                return fn(self, *args, **kwargs)
+            with controller.region(region):
+                return fn(self, *args, **kwargs)
+
+        return wrapped
+
+    return decorator

fastvideo/training/distillation_pipeline.py

@@ -1725,5 +1725,10 @@ class DistillationPipeline(TrainingPipeline):
             self.save_ema_weights(self.training_args.output_dir,
                                   self.training_args.max_train_steps)
 
+        if envs.FASTVIDEO_TORCH_PROFILER_DIR:
+            logger.info("Stopping profiler...")
+            self.profiler_controller.stop()
+            logger.info("Profiler stopped.")
+
         if get_sp_group():
             cleanup_dist_env_and_memory()

fastvideo/training/self_forcing_distillation_pipeline.py

@@ -26,6 +26,7 @@ from fastvideo.training.distillation_pipeline import DistillationPipeline
 from fastvideo.training.training_utils import (EMA_FSDP,
                                                save_distillation_checkpoint)
 from fastvideo.utils import is_vsa_available, set_random_seed
+from fastvideo.profiler import profile_region
 
 logger = init_logger(__name__)
 
@@ -940,6 +941,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
         if self.global_rank == 0:
             wandb.log(wandb_loss_dict, step=step)
 
+    @profile_region("profiler_region_training_train")
     def train(self) -> None:
         """Main training loop with self-forcing specific logging."""
         assert self.training_args.seed is not None, "seed must be set"
@@ -1221,5 +1223,10 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
             self.save_ema_weights(self.training_args.output_dir,
                                   self.training_args.max_train_steps)
 
+        if envs.FASTVIDEO_TORCH_PROFILER_DIR:
+            logger.info("Stopping profiler...")
+            self.profiler_controller.stop()
+            logger.info("Profiler stopped.")
+
         if get_sp_group():
             cleanup_dist_env_and_memory()

fastvideo/training/training_pipeline.py

@@ -7,7 +7,7 @@ from abc import ABC, abstractmethod
 from collections import deque
 from collections.abc import Iterator
 from typing import Any
-
+from fastvideo.profiler import profile_region
 import imageio
 import numpy as np
 import torch
@@ -281,6 +281,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
 
     def _prepare_dit_inputs(self,
                             training_batch: TrainingBatch) -> TrainingBatch:
+        assert self.training_args is not None, "training_args must be set"
         latents = training_batch.latents
         batch_size = latents.shape[0]
         noise = torch.randn(latents.shape,
@@ -478,6 +479,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
         training_batch.grad_norm = grad_norm
         return training_batch
 
+    @profile_region("profiler_region_training_train_one_step")
     def train_one_step(self, training_batch: TrainingBatch) -> TrainingBatch:
         training_batch = self._prepare_training(training_batch)
 
@@ -536,8 +538,10 @@ class TrainingPipeline(LoRAPipeline, ABC):
             logger.warning("Failed to load checkpoint, starting from step 0")
             self.init_steps = 0
 
+    @profile_region("profiler_region_training_train")
     def train(self) -> None:
         assert self.seed is not None, "seed must be set"
+        assert self.training_args is not None, "training_args must be set"
         set_random_seed(self.seed + self.global_rank)
         logger.info('rank: %s: start training',
                     self.global_rank,
@@ -632,24 +636,29 @@ class TrainingPipeline(LoRAPipeline, ABC):
                     step=step,
                 )
             if step % self.training_args.training_state_checkpointing_steps == 0:
-                save_checkpoint(self.transformer, self.global_rank,
-                                self.training_args.output_dir, step,
-                                self.optimizer, self.train_dataloader,
-                                self.lr_scheduler, self.noise_random_generator)
+                with self.profiler_controller.region(
+                        "profiler_region_training_save_checkpoint"):
+                    save_checkpoint(self.transformer, self.global_rank,
+                                    self.training_args.output_dir, step,
+                                    self.optimizer, self.train_dataloader,
+                                    self.lr_scheduler,
+                                    self.noise_random_generator)
                 self.transformer.train()
                 self.sp_group.barrier()
             if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
-                if self.training_args.log_visualization:
-                    self.visualize_intermediate_latents(training_batch,
-                                                        self.training_args,
-                                                        step)
-                self._log_validation(self.transformer, self.training_args, step)
-                gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
-                trainable_params = round(
-                    count_trainable(self.transformer) / 1e9, 3)
-                logger.info(
-                    "GPU memory usage after validation: %s MB, trainable params: %sB",
-                    gpu_memory_usage, trainable_params)
+                with self.profiler_controller.region(
+                        "profiler_region_training_validation"):
+                    if self.training_args.log_visualization:
+                        self.visualize_intermediate_latents(
+                            training_batch, self.training_args, step)
+                    self._log_validation(self.transformer, self.training_args,
+                                         step)
+                    gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
+                    trainable_params = round(
+                        count_trainable(self.transformer) / 1e9, 3)
+                    logger.info(
+                        "GPU memory usage after validation: %s MB, trainable params: %sB",
+                        gpu_memory_usage, trainable_params)
 
         wandb.finish()
         save_checkpoint(self.transformer, self.global_rank,
@@ -658,10 +667,16 @@ class TrainingPipeline(LoRAPipeline, ABC):
                         self.train_dataloader, self.lr_scheduler,
                         self.noise_random_generator)
 
+        if envs.FASTVIDEO_TORCH_PROFILER_DIR:
+            logger.info("Stopping profiler...")
+            self.profiler_controller.stop()
+            logger.info("Profiler stopped.")
+
         if get_sp_group():
             cleanup_dist_env_and_memory()
 
     def _log_training_info(self) -> None:
+        assert self.training_args is not None, "training_args must be set"
         total_batch_size = (self.world_size *
                             self.training_args.gradient_accumulation_steps /
                             self.training_args.sp_size *

fastvideo/utils.py

@@ -19,6 +19,7 @@ import traceback
 from collections.abc import Callable
 from dataclasses import dataclass, fields, is_dataclass
 from functools import lru_cache, partial, wraps
+from pathlib import Path
 from typing import Any, TypeVar, cast
 
 import cloudpickle
@@ -904,3 +905,73 @@ def save_decoded_latents_as_video(decoded_latents: list[torch.Tensor],
 
     os.makedirs(os.path.dirname(output_path), exist_ok=True)
     imageio.mimsave(output_path, frames, fps=fps, format="mp4")
+
+
+def _format_bytes(num_bytes: int | float | None) -> str:
+    if num_bytes is None:
+        return "N/A"
+    return f"{num_bytes / (1024 ** 3):.2f} GB"
+
+
+def log_torch_cuda_memory(
+        tag: str | None = None,
+        *,
+        log_fn: Callable[[str], None] | None = None,
+        log_file_path: str | os.PathLike[str] | None = "memory_trace.txt"
+) -> None:
+    """Log CUDA memory statistics via logger and append to a trace file."""
+
+    log_fn = log_fn or logger.info
+    prefix = f"[{tag}] " if tag else ""
+
+    if not torch.cuda.is_available():
+        message = f"{prefix}CUDA not available on this host."
+        log_fn(message)
+        _append_to_memory_trace(message, log_file_path)
+        return
+
+    try:
+        device_index = torch.cuda.current_device()
+        device_name = torch.cuda.get_device_name(device_index)
+        allocated = torch.cuda.memory_allocated(device_index)
+        reserved = torch.cuda.memory_reserved(device_index)
+        max_allocated = torch.cuda.max_memory_allocated(device_index)
+        max_reserved = torch.cuda.max_memory_reserved(device_index)
+        free_mem, total_mem = torch.cuda.mem_get_info(device_index)
+    except Exception as exc:  # noqa: BLE001
+        message = f"{prefix}Unable to query CUDA memory stats: {exc}"
+        log_fn(message)
+        _append_to_memory_trace(message, log_file_path)
+        return
+
+    used_mem = total_mem - free_mem
+
+    stats = [
+        f"device={device_name} (index={device_index})",
+        f"allocated={_format_bytes(allocated)}",
+        f"reserved={_format_bytes(reserved)}",
+        f"max_allocated={_format_bytes(max_allocated)}",
+        f"max_reserved={_format_bytes(max_reserved)}",
+        f"used={_format_bytes(used_mem)}",
+        f"free={_format_bytes(free_mem)}",
+        f"total={_format_bytes(total_mem)}",
+    ]
+
+    message = f"{prefix}CUDA memory stats: {' | '.join(stats)}"
+    log_fn(message)
+    _append_to_memory_trace(message, log_file_path)
+
+
+def _append_to_memory_trace(
+        message: str, log_file_path: str | os.PathLike[str] | None) -> None:
+    if not log_file_path:
+        return
+
+    try:
+        path = Path(log_file_path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with path.open("a", encoding="utf-8") as outfile:
+            outfile.write(message + "\n")
+    except Exception:  # noqa: BLE001
+        # Avoid raising/logging recursively if writing to the file fails.
+        pass