diff --git a/docs/source/contributing/profiling.md b/docs/source/contributing/profiling.md new file mode 100644 index 000000000..d2dacb65f --- /dev/null +++ b/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::` 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 . + +### 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. diff --git a/docs/source/index.md b/docs/source/index.md index 72ef520c6..9027a8429 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -115,6 +115,7 @@ design/overview contributing/overview contributing/developer_env/index +contributing/profiling ::: :::{toctree} diff --git a/fastvideo/envs.py b/fastvideo/envs.py index a42a5c8cd..4842d7f85 100644 --- a/fastvideo/envs.py +++ b/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` diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 5d31abb3f..9b91348e9 100644 --- a/fastvideo/fastvideo_args.py +++ b/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) diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index b15401120..c31aaef9c 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/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( diff --git a/fastvideo/profiler.py b/fastvideo/profiler.py new file mode 100644 index 000000000..914f63a06 --- /dev/null +++ b/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 diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index e00d226a2..dcaa1cc52 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/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() diff --git a/fastvideo/training/self_forcing_distillation_pipeline.py b/fastvideo/training/self_forcing_distillation_pipeline.py index 770683df4..6e97082cf 100644 --- a/fastvideo/training/self_forcing_distillation_pipeline.py +++ b/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() diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 75e8d1858..731febe1b 100644 --- a/fastvideo/training/training_pipeline.py +++ b/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 * diff --git a/fastvideo/utils.py b/fastvideo/utils.py index 50937aa24..d27451ff0 100644 --- a/fastvideo/utils.py +++ b/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