diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md index 144b9975c..354c936d4 100644 --- a/docs/contributing/profiling.md +++ b/docs/contributing/profiling.md @@ -9,33 +9,82 @@ FastVideo exposes a process-wide torch profiler that you can enable via environm ```bash FASTVIDEO_TORCH_PROFILER_DIR=/mnt/traces/fastvideo \ -FASTVIDEO_TORCH_PROFILE_REGIONS="profiler_region_model_loading,profiler_region_training_step" +FASTVIDEO_TORCH_PROFILE_REGIONS="model_loading,training_train_one_step" \ +bash examples/train/run.sh /path/to/config.yaml ``` -All profiled regions must be registered in `fastvideo.profiler`; the current list includes: +The `profiler_region_` prefix is optional. 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_train` — the complete training run +- `profiler_region_training_train_one_step` — one complete optimizer step +- `profiler_region_training_dataloader` — fetch the next batch in the trainer process +- `profiler_region_training_forward` — method forward and loss computation +- `profiler_region_training_backward` — backward pass +- `profiler_region_training_optimizer` — gradient clipping, optimizer/scheduler step, and zeroing gradients +- `profiler_region_training_callbacks` — end-of-step callbacks such as EMA updates - `profiler_region_training_validation` -- `profiler_region_training_epoch` -- `profiler_region_training_step` -- `profiler_region_training_backward` -- `profiler_region_training_optimizer` +- `profiler_region_training_save_checkpoint` - `profiler_region_distillation_teacher_forward` - `profiler_region_distillation_student_forward` - `profiler_region_distillation_loss` - `profiler_region_distillation_update` +- `profiler_region_dmd2_student_rollout` — one DMD2 student rollout, including simulated prefix steps +- `profiler_region_dmd2_generator_loss` — teacher/critic scoring for the generator loss +- `profiler_region_dmd2_critic_loss` — critic flow-matching loss, including its student rollout + +### Profiling modular training + +The YAML-driven trainer under `fastvideo/train/` initializes and flushes the +profiler automatically. For example, this captures six DMD2 steps: one warmup +step, four steady-state critic updates, and the configured 1-in-5 generator +update, without changing the method's update cadence: + +```bash +TRACE_DIR="$(pwd)/profiler_traces/wan_dmd2" +mkdir -p "$TRACE_DIR" + +NUM_GPUS=4 \ +WANDB_MODE=offline \ +FASTVIDEO_TORCH_PROFILER_DIR="$TRACE_DIR" \ +FASTVIDEO_TORCH_PROFILE_REGIONS="training_train,training_train_one_step,training_dataloader,training_forward,training_backward,training_optimizer,training_callbacks,dmd2_student_rollout,dmd2_generator_loss,dmd2_critic_loss" \ +bash examples/train/run.sh \ + examples/train/configs/distribution_matching/wan/dmd2_t2v.yaml \ + --training.loop.max_train_steps 6 \ + --training.checkpoint.training_state_checkpointing_steps 0 \ + --callbacks.validation.every_steps 0 +``` + +The two overrides disable validation and checkpoint writes so their I/O does +not contaminate training-step measurements. Remove them when profiling those +regions. Methods that manage optimization internally emit the enclosing +`training_train_one_step` region but do not emit the generic +dataloader/forward/backward/optimizer child regions because those boundaries +belong to the method. With `dataloader_num_workers > 0`, the +`training_dataloader` region measures the training process waiting for and +receiving the next batch; work performed inside dataloader worker subprocesses +does not appear in that process's torch-profiler trace. + +The modular trainer also logs `dataloader_time_sec` on every ordinary training +step (including DMD2), so routine runs can monitor rank 0's +`next(dataloader)` wait without enabling the torch profiler. It is +intentionally not reduced across ranks; cross-rank synchronization would +perturb the hot path being measured. 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. +Only one profiler controller is created per process; subsequent pipelines +reuse it. Each outermost enabled region invocation produces one complete +CPU/CUDA trace segment. Enabled regions nested inside it are annotations in +that segment. Include an enclosing region such as `training_train_one_step` +when selecting its child regions so they share one trace and profiler startup +does not perturb each phase independently. If you set +`FASTVIDEO_TORCH_PROFILE_REGIONS` incorrectly (e.g. misspelled name), FastVideo +logs a warning and ignores that entry. Additional knobs: @@ -44,13 +93,18 @@ Additional knobs: - `FASTVIDEO_TORCH_PROFILER_WITH_STACK` - `FASTVIDEO_TORCH_PROFILER_WITH_FLOPS` -Traces can be visualized using . +Traces can be visualized using . Each rank also +writes `summary_rank_segment_.txt` and the corresponding JSON +beside every trace segment for a quick operator-level view without loading the +full timeline. The JSON contains CPU and device totals; input-shape grouping is +only enabled with `FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES=1` because it makes +summary generation materially more expensive on large traces. ### 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 disk storage. -- 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. +- When adding new regions, register them in `fastvideo.profiler` and wrap the corresponding code block with `profiler_region("your_region")` or the `@profile_region` decorator. ## Related: Activation Trace Mode diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 037762fa4..5871ca3b3 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -250,12 +250,7 @@ environment_variables: dict[str, Callable[[], Any]] = { # not profile flops. "FASTVIDEO_TORCH_PROFILER_WITH_FLOPS": lambda: bool(os.getenv("FASTVIDEO_TORCH_PROFILER_WITH_FLOPS", "0") != "0"), - # Wait steps per profiling cycle (torch.profiler.schedule wait parameter) - # Defaults to 2 if not set. - # Warmup steps per profiling cycle (torch.profiler.schedule warmup parameter) - # Defaults to 1 if not set. - # Active steps per profiling cycle (torch.profiler.schedule active parameter) - # Defaults to 2 if not set. + # Comma-separated names of registered profiler regions to capture. "FASTVIDEO_TORCH_PROFILE_REGIONS": lambda: os.getenv("FASTVIDEO_TORCH_PROFILE_REGIONS", ""), diff --git a/fastvideo/profiler.py b/fastvideo/profiler.py index 3303d4600..74ce8379b 100644 --- a/fastvideo/profiler.py +++ b/fastvideo/profiler.py @@ -2,18 +2,17 @@ """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. Short names work +controller that gates collection based on named *regions*. Regions are enabled +through the ``FASTVIDEO_TORCH_PROFILE_REGIONS`` comma-separated list. Short names work (``FASTVIDEO_TORCH_PROFILE_REGIONS=model_loading,training_train`` resolves the ``profiler_region_`` prefix automatically). Typical usage from client code:: - controller = TorchProfilerController(profiler, activities) - with controller.region("training_dit"): + controller = get_or_create_profiler("/tmp/fastvideo-traces") + with controller.region("training_train"): run_training_step() + controller.stop() To introduce a new region, register it via :func:`register_profiler_region` and wrap the corresponding code in :meth:`TorchProfilerController.region`. @@ -22,11 +21,11 @@ and wrap the corresponding code in :meth:`TorchProfilerController.region`. from __future__ import annotations import contextlib +import functools +import os +from collections.abc import Callable, Iterable from dataclasses import dataclass from typing import Any -from collections.abc import Callable -import functools -from collections.abc import Iterable import torch @@ -142,6 +141,26 @@ register_profiler_region( name="profiler_region_training_train_one_step", description="Single optimizer step including forward/backward passes.", ) +register_profiler_region( + name="profiler_region_training_forward", + description="Training method forward pass and loss computation.", +) +register_profiler_region( + name="profiler_region_training_dataloader", + description="Fetch the next training batch in the trainer process.", +) +register_profiler_region( + name="profiler_region_training_backward", + description="Training backward pass.", +) +register_profiler_region( + name="profiler_region_training_optimizer", + description="Gradient clipping, optimizer/scheduler steps, and zero_grad.", +) +register_profiler_region( + name="profiler_region_training_callbacks", + description="End-of-step training callbacks such as EMA updates.", +) register_profiler_region( name="profiler_region_training_train", description="High-level step orchestration in the training loop.", @@ -165,6 +184,21 @@ register_profiler_region( description="Parameter updates specific to distillation workflows.", ) +# DMD2 method regions. These sit inside ``training_forward`` and make the +# method's multi-model forward path distinguishable in a single trace. +register_profiler_region( + name="profiler_region_dmd2_student_rollout", + description="DMD2 student rollout, including its simulated prefix steps.", +) +register_profiler_region( + name="profiler_region_dmd2_generator_loss", + description="DMD2 generator loss, including teacher and critic scoring.", +) +register_profiler_region( + name="profiler_region_dmd2_critic_loss", + description="DMD2 critic flow-matching loss, including its student rollout.", +) + def get_or_create_profiler(trace_dir: str | None) -> TorchProfilerController: """Create or reuse the process-wide torch profiler controller.""" @@ -189,25 +223,28 @@ def get_or_create_profiler(trace_dir: str | None) -> TorchProfilerController: ) 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, - # No schedule: nothing in the codebase calls profiler.step(), so a - # wait/warmup schedule never advances and the profiler records nothing. - # Region toggling gates collection; the single trace exports at stop(). - on_trace_ready=torch.profiler.tensorboard_trace_handler(trace_dir, use_gzip=True), + def profiler_factory() -> Any: + return 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( + None, + _DEFAULT_ACTIVITIES, + profiler_factory=profiler_factory, + trace_dir=trace_dir, ) - controller = TorchProfilerController(profiler, _DEFAULT_ACTIVITIES) - controller._trace_dir = trace_dir controller.start() - # The trace only exports at stop(); inference paths have no shutdown hook - # that calls it, so register one. stop() is idempotent. + # Region exit normally exports each trace segment. Keep an atexit hook for + # exceptions or process shutdown while a region is still active. import atexit atexit.register(controller.stop) - logger.info("Torch profiler started") + logger.info("Torch profiler armed; collection starts at the first enabled region") return controller @@ -260,7 +297,14 @@ class TorchProfilerConfig: class TorchProfilerController: - """Helper that toggles torch profiler collection for named regions. + """Create complete torch-profiler trace segments for named regions. + + PyTorch's dynamic CUDA collection toggle can fail to re-enable CUPTI on + some supported stacks. In that failure mode it emits CPU operators while + silently dropping every CUDA kernel. This controller therefore starts a + fresh profiler at each outermost enabled region and stops it when that + region exits. Nested enabled regions become annotations in the same + CPU/CUDA trace segment. Parameters ---------- @@ -273,12 +317,17 @@ class TorchProfilerController: config: Optional :class:`TorchProfilerConfig`. If omitted, :meth:`from_env` constructs one during initialization. + profiler_factory: + Factory for fresh profiler instances. Required to profile more than + one outermost region invocation. + trace_dir: + Directory for per-segment summaries. Examples -------- Enabling an existing region from the command line:: - FASTVIDEO_TORCH_PROFILE_REGIONS=model_loading,training_dit \ + FASTVIDEO_TORCH_PROFILE_REGIONS=model_loading,training_train \ python fastvideo/training/wan_training_pipeline.py ... Wrapping a code block in a registered region:: @@ -299,24 +348,29 @@ class TorchProfilerController: activities: Iterable[torch.profiler.ProfilerActivity], config: TorchProfilerConfig | None = None, disabled: bool = False, + profiler_factory: Callable[[], Any] | None = None, + trace_dir: str | None = None, ) -> 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().") + self._profiler = profiler + self._profiler_factory = profiler_factory + self._activities = activities_tuple + self._trace_dir = trace_dir + self._segment_index = 0 + self._segment_region: str | None = None + self._active_region_depth = 0 + self._collection_enabled = False if disabled: - self._profiler = None + self._configured = False + self._armed = False return - self._profiler = profiler - self._activities = activities_tuple self._config = config or TorchProfilerConfig.from_env() - # torch.profiler collects from start(); reflect that so the initial - # _set_collection(False) in start() actually toggles it off instead of - # short-circuiting (which captured everything before the first region). - self._collection_enabled = True - self._active_region_depth = 0 - self._trace_dir: str | None = None + self._configured = True + self._armed = False logger.info("PROFILER: TorchProfilerController initialized with config: %s", self._config) set_global_controller(self) @@ -324,29 +378,61 @@ class TorchProfilerController: 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: + if not self.has_profiler: return False resolved = resolve_profiler_region(region) if resolved is None: return False return self._config.regions.get(resolved.name, False) - def _set_collection(self, enabled: bool) -> None: - if self._profiler is None: + def _new_profiler(self) -> Any: + if self._profiler is not None: + profiler = self._profiler + self._profiler = None + return profiler + if self._profiler_factory is None: + raise RuntimeError("Torch profiler cannot start another trace segment without a profiler_factory") + return self._profiler_factory() + + def _start_segment(self, region: str) -> None: + if self._collection_enabled: return - if self._collection_enabled == enabled: + self._profiler = self._new_profiler() + logger.info( + "PROFILER: Starting segment %d for region %s", + self._segment_index, + region, + ) + self._profiler.start() + self._segment_region = region + self._collection_enabled = True + + def _finish_segment(self) -> None: + if self._profiler is None or not self._collection_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 + profiler = self._profiler + segment_index = self._segment_index + segment_region = self._segment_region or "unknown" + logger.info( + "PROFILER: Stopping segment %d for region %s", + segment_index, + segment_region, + ) + profiler.stop() + self._write_summary( + profiler, + segment_index=segment_index, + segment_region=segment_region, + ) + self._profiler = None + self._collection_enabled = False + self._segment_region = None + self._segment_index += 1 _warned_unregistered: set[str] = set() @@ -354,7 +440,7 @@ class TorchProfilerController: def region(self, region: str): """Context manager that enables profiling for ``region`` if configured.""" - if self._profiler is None: + if not self.has_profiler: yield return @@ -371,64 +457,70 @@ class TorchProfilerController: yield return - # NVTX range so the same region names are visible in nsys timelines + if self._active_region_depth == 0: + self._start_segment(region) + + # NVTX range so the same region names are visible in nsys timelines. + # Push after profiler startup so Kineto also records the annotation. nvtx = torch.cuda.is_available() if nvtx: torch.cuda.nvtx.range_push(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: - # record_function opens after collection is enabled so the region - # marker itself lands in the trace. with torch.profiler.record_function(f"fastvideo.region::{region}"): 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) - if nvtx: - torch.cuda.nvtx.range_pop() + try: + if nvtx: + torch.cuda.nvtx.range_pop() + finally: + # Close NVTX before stopping Kineto so both profilers see a + # balanced outermost range in the exported segment. + if self._active_region_depth == 0: + self._finish_segment() def start(self) -> None: - """Start the profiler and pause collection until a region is entered.""" + """Arm the controller; collection begins at an enabled region.""" - logger.info("PROFILER: Starting profiler...") - if self._profiler is None: + if not self._configured: 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") + self._armed = True + logger.info("PROFILER: Controller armed") - def _write_summary(self) -> None: + def _write_summary( + self, + profiler: Any, + *, + segment_index: int, + segment_region: str, + ) -> None: """Compact per-rank op summary next to the trace: a key_averages table and a JSON with input shapes, so operator-split analysis does not require parsing multi-GB chrome traces.""" - if self._profiler is None or not self._trace_dir: + if not self._trace_dir: return try: import json as _json - import os as _os - rank = _os.environ.get("RANK", "0") - averages = self._profiler.key_averages(group_by_input_shape=True) - stem = _os.path.join(self._trace_dir, f"summary_rank{rank}") - with open(f"{stem}.txt", "w") as fh: + rank = os.environ.get("RANK", "0") + short_region = segment_region.removeprefix("profiler_region_") + stem = os.path.join( + self._trace_dir, + f"summary_rank{rank}_segment{segment_index:04d}_{short_region}", + ) + averages = profiler.key_averages(group_by_input_shape=envs.FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES, ) + with open(f"{stem}.txt", "w", encoding="utf-8") as fh: fh.write(averages.table(sort_by="self_device_time_total", row_limit=60)) rows = [{ "name": e.key, "shapes": str(e.input_shapes), + "self_cpu_us": e.self_cpu_time_total, + "cpu_us": e.cpu_time_total, "self_device_us": e.self_device_time_total, "device_us": e.device_time_total, "count": e.count, } for e in averages] - with open(f"{stem}.json", "w") as fh: + with open(f"{stem}.json", "w", encoding="utf-8") as fh: _json.dump(rows, fh) if rank == "0": logger.info("PROFILER: summary written to %s.txt", stem) @@ -436,24 +528,25 @@ class TorchProfilerController: logger.exception("PROFILER: summary generation failed") def stop(self) -> None: - """Stop the profiler after disabling collection and clearing state.""" + """Flush any active segment and disable this controller.""" - if self._profiler is None: + if not self._configured: return logger.info("PROFILER: Stopping profiler...") - self._profiler.stop() - self._write_summary() - self._profiler = None # makes stop() idempotent (atexit may re-enter) + self._finish_segment() + self._profiler = None + self._configured = False + self._armed = False logger.info("PROFILER: Profiler stopped") self._active_region_depth = 0 set_global_controller(None) @property def has_profiler(self) -> bool: - """Return ``True`` when a profiler instance is available.""" + """Return ``True`` when this controller is configured and armed.""" - return self._profiler is not None + return self._configured and self._armed @property def activities(self) -> tuple[torch.profiler.ProfilerActivity, ...]: @@ -476,13 +569,21 @@ def profiler_region(region: str): def profile_region(region: str) -> Callable[[Callable[..., Any]], Callable[..., Any]]: - """Wrap a bound method so it runs inside a profiler region if available.""" + """Wrap a bound method so it runs inside a profiler region if available. + + Prefer a controller attached to the instance, then fall back to the + process-wide controller. The fallback lets lightweight owners such as the + modular trainer and its callbacks add regions without threading profiler + plumbing through their public constructors. + """ 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: + controller = get_global_controller() if controller is None or not controller.has_profiler: return fn(self, *args, **kwargs) with controller.region(region): diff --git a/fastvideo/tests/contract/test_profiler_regions.py b/fastvideo/tests/contract/test_profiler_regions.py index de18fa7a5..3ac2be1e3 100644 --- a/fastvideo/tests/contract/test_profiler_regions.py +++ b/fastvideo/tests/contract/test_profiler_regions.py @@ -16,6 +16,9 @@ import os import subprocess import sys +import pytest +import torch + # Five-window child: ops before any region, inside a region, between regions, # inside a second (short-named) region, after the last region. Exits without # calling stop() — export must happen via the atexit hook. Each window is @@ -61,19 +64,21 @@ def _run_child(tmp_path): return trace_dir, proc.stdout + proc.stderr -def _trace_event_names(trace_dir): +def _trace_events(trace_dir): traces = glob.glob(os.path.join(trace_dir, "**", "*.json*"), recursive=True) traces = [t for t in traces if "summary" not in os.path.basename(t)] assert traces, f"no trace exported to {trace_dir} (atexit hook missing?)" - opener = gzip.open if traces[0].endswith(".gz") else open - with opener(traces[0], "rt") as fh: - events = json.load(fh).get("traceEvents", []) - return {e.get("name", "") for e in events} + events = [] + for trace in traces: + opener = gzip.open if trace.endswith(".gz") else open + with opener(trace, "rt") as fh: + events.extend(json.load(fh).get("traceEvents", [])) + return events def test_regions_gate_collection_and_atexit_exports(tmp_path): trace_dir, output = _run_child(tmp_path) - names = _trace_event_names(trace_dir) + names = {event.get("name", "") for event in _trace_events(trace_dir)} assert "win_region1" in names, "op inside an enabled region was not captured" assert "win_region2" in names, "short region name did not resolve/capture" @@ -87,8 +92,61 @@ def test_regions_gate_collection_and_atexit_exports(tmp_path): assert output.count("is not registered") == 1 # per-rank op summary written next to the trace - summaries = glob.glob(os.path.join(trace_dir, "summary_rank0.*")) - assert sorted(os.path.splitext(s)[1] for s in summaries) == [".json", ".txt"] + summaries = glob.glob(os.path.join(trace_dir, "summary_rank0_segment*.*")) + assert sorted(os.path.splitext(s)[1] for s in summaries) == [ + ".json", + ".json", + ".txt", + ".txt", + ] + with open(next(s for s in summaries if s.endswith(".json")), encoding="utf-8") as fh: + summary_rows = json.load(fh) + assert summary_rows + assert { + "name", + "shapes", + "self_cpu_us", + "cpu_us", + "self_device_us", + "device_us", + "count", + } <= summary_rows[0].keys() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_cuda_region_exports_kernel_events(tmp_path): + trace_dir = str(tmp_path / "cuda_traces") + child = r""" +import torch +from fastvideo.profiler import get_or_create_profiler, profiler_region + +controller = get_or_create_profiler({trace_dir!r}) +x = torch.randn(1024, 1024, device="cuda") +y = torch.randn(1024, 1024, device="cuda") +torch.cuda.synchronize() +with profiler_region("training_forward"): + torch.mm(x, y) + torch.cuda.synchronize() +controller.stop() +""".format(trace_dir=trace_dir) + env = os.environ.copy() + env["FASTVIDEO_TORCH_PROFILER_DIR"] = trace_dir + env["FASTVIDEO_TORCH_PROFILE_REGIONS"] = "training_forward" + proc = subprocess.run( + [sys.executable, "-c", child], + env=env, + capture_output=True, + text=True, + timeout=300, + ) + assert proc.returncode == 0, proc.stderr + + events = _trace_events(trace_dir) + categories = {event.get("cat", "") for event in events} + names = {event.get("name", "") for event in events} + assert "kernel" in categories + assert "cuda_runtime" in categories + assert "fastvideo.region::training_forward" in names def test_noop_without_profiler_dir(tmp_path): diff --git a/fastvideo/tests/train/test_profiler_entrypoint.py b/fastvideo/tests/train/test_profiler_entrypoint.py new file mode 100644 index 000000000..9239adae1 --- /dev/null +++ b/fastvideo/tests/train/test_profiler_entrypoint.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only tests for modular training profiler lifecycle ownership.""" + +from __future__ import annotations + +from contextlib import contextmanager +from types import SimpleNamespace + +from fastvideo.train.entrypoint.train import run_training_from_config + + +class _RecordingProfiler: + + def __init__(self) -> None: + self.events: list[tuple[str, str] | tuple[str]] = [] + + @contextmanager + def region(self, name: str): + self.events.append(("enter", name)) + try: + yield + finally: + self.events.append(("exit", name)) + + def stop(self) -> None: + self.events.append(("stop", )) + + +def test_dry_run_profiles_model_build_and_flushes(monkeypatch) -> None: + profiler = _RecordingProfiler() + training = SimpleNamespace( + distributed=SimpleNamespace(tp_size=1, sp_size=1), + vsa_sparsity=0.0, + model_path="model", + ) + cfg = SimpleNamespace(training=training) + + monkeypatch.setattr( + "fastvideo.train.utils.config.load_run_config", + lambda *args, **kwargs: cfg, + ) + monkeypatch.setattr( + "fastvideo.distributed.maybe_init_distributed_environment_and_model_parallel", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "fastvideo.train.utils.builder.build_from_config", + lambda loaded_cfg: (loaded_cfg.training, object(), object(), 0), + ) + monkeypatch.setattr( + "fastvideo.train.entrypoint.train.get_or_create_profiler", + lambda trace_dir: profiler, + ) + + run_training_from_config("unused.yaml", dry_run=True) + + assert profiler.events == [ + ("enter", "profiler_region_model_loading"), + ("exit", "profiler_region_model_loading"), + ("stop", ), + ] diff --git a/fastvideo/tests/train/trainer/test_profiler.py b/fastvideo/tests/train/trainer/test_profiler.py new file mode 100644 index 000000000..b70f00cca --- /dev/null +++ b/fastvideo/tests/train/trainer/test_profiler.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only tests for modular Trainer profiler boundaries.""" + +from __future__ import annotations + +from contextlib import contextmanager +from types import SimpleNamespace +from typing import Any + +import torch + +from fastvideo.profiler import list_profiler_regions +from fastvideo.train.trainer import Trainer +from fastvideo.train.utils.training_config import TrainingConfig + + +class _RecordingProfiler: + + def __init__(self) -> None: + self.events: list[tuple[str, str]] = [] + + @property + def has_profiler(self) -> bool: + return True + + @contextmanager + def region(self, name: str): + self.events.append(("enter", name)) + try: + yield + finally: + self.events.append(("exit", name)) + + +class _DummyTracker: + + def log(self, metrics: dict[str, float], step: int) -> None: + del metrics, step + + def finish(self) -> None: + pass + + +class _DummyMethod: + + def __init__(self) -> None: + self.weight = torch.nn.Parameter(torch.tensor(1.0)) + + def set_tracker(self, tracker: Any) -> None: + del tracker + + def on_train_start(self) -> None: + pass + + def manages_optimization(self) -> bool: + return False + + def single_train_step( + self, + batch: dict[str, Any], + iteration: int, + ) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, float]]: + del batch, iteration + return {"total_loss": self.weight.square()}, {}, {} + + def backward( + self, + loss_map: dict[str, torch.Tensor], + outputs: dict[str, Any], + *, + grad_accum_rounds: int, + ) -> None: + del outputs + (loss_map["total_loss"] / grad_accum_rounds).backward() + + def optimizers_schedulers_step(self, iteration: int) -> None: + del iteration + + def optimizers_zero_grad(self, iteration: int) -> None: + del iteration + self.weight.grad = None + + +def test_modular_trainer_emits_nested_step_regions(monkeypatch) -> None: + group = SimpleNamespace(rank=0, local_rank=0, rank_in_group=0, world_size=1) + profiler = _RecordingProfiler() + + monkeypatch.setattr("fastvideo.train.trainer.get_world_group", lambda: group) + monkeypatch.setattr("fastvideo.train.trainer.get_sp_group", lambda: group) + monkeypatch.setattr( + "fastvideo.train.trainer.build_tracker", + lambda *args, **kwargs: _DummyTracker(), + ) + monkeypatch.setattr("fastvideo.profiler._GLOBAL_CONTROLLER", profiler) + + trainer = Trainer(TrainingConfig()) + trainer.run( + _DummyMethod(), + dataloader=[{}], + max_steps=1, + ) + + assert profiler.events == [ + ("enter", "profiler_region_training_train"), + ("enter", "profiler_region_training_train_one_step"), + ("enter", "profiler_region_training_dataloader"), + ("exit", "profiler_region_training_dataloader"), + ("enter", "profiler_region_training_forward"), + ("exit", "profiler_region_training_forward"), + ("enter", "profiler_region_training_backward"), + ("exit", "profiler_region_training_backward"), + ("enter", "profiler_region_training_optimizer"), + ("exit", "profiler_region_training_optimizer"), + ("enter", "profiler_region_training_callbacks"), + ("exit", "profiler_region_training_callbacks"), + ("exit", "profiler_region_training_train_one_step"), + ("exit", "profiler_region_training_train"), + ] + + +def test_modular_training_regions_are_registered() -> None: + names = {region.name for region in list_profiler_regions()} + assert { + "profiler_region_model_loading", + "profiler_region_training_train", + "profiler_region_training_train_one_step", + "profiler_region_training_dataloader", + "profiler_region_training_forward", + "profiler_region_training_backward", + "profiler_region_training_optimizer", + "profiler_region_training_callbacks", + "profiler_region_training_save_checkpoint", + "profiler_region_training_validation", + "profiler_region_dmd2_student_rollout", + "profiler_region_dmd2_generator_loss", + "profiler_region_dmd2_critic_loss", + } <= names diff --git a/fastvideo/tests/train/trainer/test_validation.py b/fastvideo/tests/train/trainer/test_validation.py index 5a08a36f8..f996c6d3f 100644 --- a/fastvideo/tests/train/trainer/test_validation.py +++ b/fastvideo/tests/train/trainer/test_validation.py @@ -136,4 +136,5 @@ def test_trainer_runs_validation_callback_during_training(monkeypatch, ) -> None assert method.zero_grad_steps == [0, 1, 2, 3] assert method.optimizer_steps == [1, 2, 3] assert [step for _, step in tracker.logs] == [1, 2, 3] + assert all("dataloader_time_sec" in metrics for metrics, _ in tracker.logs) assert tracker.finished is True diff --git a/fastvideo/train/callbacks/validation.py b/fastvideo/train/callbacks/validation.py index 4b635e827..fd8a6e66d 100644 --- a/fastvideo/train/callbacks/validation.py +++ b/fastvideo/train/callbacks/validation.py @@ -31,6 +31,7 @@ from fastvideo.distributed import ( get_world_group, ) from fastvideo.logger import init_logger +from fastvideo.profiler import profile_region from fastvideo.pipelines import ForwardBatch from fastvideo.train.callbacks.callback import Callback from fastvideo.train.utils.instantiate import resolve_target @@ -285,6 +286,7 @@ class ValidationCallback(Callback): # Core validation logic # ---------------------------------------------------------- + @profile_region("profiler_region_training_validation") def _run_validation( self, method: TrainingMethod, diff --git a/fastvideo/train/entrypoint/train.py b/fastvideo/train/entrypoint/train.py index 3f71e2a72..b77471d10 100644 --- a/fastvideo/train/entrypoint/train.py +++ b/fastvideo/train/entrypoint/train.py @@ -24,7 +24,9 @@ from typing import Any import torch +import fastvideo.envs as envs from fastvideo.logger import init_logger +from fastvideo.profiler import get_or_create_profiler logger = init_logger(__name__) @@ -74,47 +76,55 @@ def run_training_from_config( tc.distributed.sp_size, ) - _, method, dataloader, start_step = build_from_config(cfg) + profiler_controller = get_or_create_profiler(envs.FASTVIDEO_TORCH_PROFILER_DIR, ) + try: + with profiler_controller.region("profiler_region_model_loading"): + _, method, dataloader, start_step = build_from_config(cfg) - if dry_run: - logger.info("Dry-run: config parsed and " - "build_from_config succeeded.") - return + if dry_run: + logger.info("Dry-run: config parsed and " + "build_from_config succeeded.") + return - trainer = Trainer( - tc, - config=cfg.resolved_config(), - callback_configs=cfg.callbacks, - ) + trainer = Trainer( + tc, + config=cfg.resolved_config(), + callback_configs=cfg.callbacks, + ) - # Attach the exact YAML used for this run to the - # tracker (e.g., W&B Files). - trainer.tracker.log_file( - os.path.abspath(os.path.expanduser(config_path)), - name="run.yaml", - ) + # Attach the exact YAML used for this run to the + # tracker (e.g., W&B Files). + trainer.tracker.log_file( + os.path.abspath(os.path.expanduser(config_path)), + name="run.yaml", + ) - ckpt_config = CheckpointConfig( - save_steps=int(tc.checkpoint.training_state_checkpointing_steps or 0), - keep_last=int(tc.checkpoint.checkpoints_total_limit or 0), - ) + ckpt_config = CheckpointConfig( + save_steps=int(tc.checkpoint.training_state_checkpointing_steps or 0), + keep_last=int(tc.checkpoint.checkpoints_total_limit or 0), + ) - checkpoint_manager = CheckpointManager( - method=method, - dataloader=dataloader, - output_dir=tc.checkpoint.output_dir, - config=ckpt_config, - callbacks=trainer.callbacks, - raw_config=cfg.raw, - ) + checkpoint_manager = CheckpointManager( + method=method, + dataloader=dataloader, + output_dir=tc.checkpoint.output_dir, + config=ckpt_config, + callbacks=trainer.callbacks, + raw_config=cfg.raw, + ) - trainer.run( - method, - dataloader=dataloader, - max_steps=tc.loop.max_train_steps, - start_step=start_step, - checkpoint_manager=checkpoint_manager, - ) + trainer.run( + method, + dataloader=dataloader, + max_steps=tc.loop.max_train_steps, + start_step=start_step, + checkpoint_manager=checkpoint_manager, + ) + finally: + # torch.profiler exports its trace at stop(). Keep atexit as a fallback + # for abrupt exits, but flush here so normal training returns only after + # all per-rank traces and summaries are complete. + profiler_controller.stop() def main( diff --git a/fastvideo/train/methods/distribution_matching/dmd2.py b/fastvideo/train/methods/distribution_matching/dmd2.py index ee1890bcb..6d2ddd80a 100644 --- a/fastvideo/train/methods/distribution_matching/dmd2.py +++ b/fastvideo/train/methods/distribution_matching/dmd2.py @@ -8,6 +8,7 @@ from typing import Any, Literal import torch import torch.nn.functional as F +from fastvideo.profiler import profile_region from fastvideo.train.methods.base import TrainingMethod, LogScalar from fastvideo.train.models.base import ModelBase from fastvideo.train.utils.optimizer import ( @@ -476,6 +477,7 @@ class DMD2Method(TrainingMethod): self._score_max_timestep, ) + @profile_region("profiler_region_dmd2_student_rollout") def _student_rollout( self, batch: Any, @@ -598,6 +600,7 @@ class DMD2Method(TrainingMethod): batch.dmd_latent_vis_dict["generator_timestep"] = target_timestep.float().detach() return pred_x0 + @profile_region("profiler_region_dmd2_critic_loss") def _critic_flow_matching_loss( self, batch: Any, @@ -638,6 +641,7 @@ class DMD2Method(TrainingMethod): outputs, ) + @profile_region("profiler_region_dmd2_generator_loss") def _dmd_loss( self, generator_pred_x0: torch.Tensor, diff --git a/fastvideo/train/trainer.py b/fastvideo/train/trainer.py index be6054f77..2b379ac59 100644 --- a/fastvideo/train/trainer.py +++ b/fastvideo/train/trainer.py @@ -11,6 +11,7 @@ import torch from tqdm.auto import tqdm from fastvideo.distributed import get_sp_group, get_world_group +from fastvideo.profiler import profile_region, profiler_region from fastvideo.train.callbacks.callback import CallbackDict from fastvideo.train.methods.base import LogScalar, TrainingMethod from fastvideo.train.utils.tracking import build_tracker @@ -98,6 +99,113 @@ class Trainer: if self.global_rank == 0 and validation_metrics: self.tracker.log(validation_metrics, iteration) + @profile_region("profiler_region_training_train_one_step") + def _run_train_step( + self, + method: TrainingMethod, + *, + data_stream: Iterator[dict[str, Any]], + step: int, + grad_accum: int, + method_manages_optimization: bool, + ) -> None: + t0 = time.perf_counter() + + # Accumulate on GPU during grad-accum; materialise to CPU once per + # step right before logging. + loss_sums: dict[str, float | torch.Tensor] = {} + metric_sums: dict[str, float | torch.Tensor] = {} + dataloader_time_sec = 0.0 + if method_manages_optimization: + # Managed methods own their forward/backward/optimizer boundaries, + # so the enclosing training_train_one_step region is the truthful + # granularity available to the trainer. + loss_map, outputs, step_metrics = method.managed_train_step( + data_stream, + step, + ) + for k, v in loss_map.items(): + if isinstance(v, torch.Tensor): + loss_sums[k] = v.detach() + for k, v in step_metrics.items(): + if k in loss_sums: + raise ValueError(f"Metric key {k!r} collides " + "with loss key. Use a " + "different name (e.g. prefix " + "with 'train/').") + metric_sums[k] = _coerce_log_scalar( + v, + where=("method.managed_train_step()" + f".metrics[{k!r}]"), + ) + else: + for _ in range(grad_accum): + dataloader_t0 = time.perf_counter() + with profiler_region("profiler_region_training_dataloader"): + batch = next(data_stream) + dataloader_time_sec += (time.perf_counter() - dataloader_t0) + with profiler_region("profiler_region_training_forward"): + loss_map, outputs, step_metrics = (method.single_train_step( + batch, + step, + )) + + with profiler_region("profiler_region_training_backward"): + method.backward( + loss_map, + outputs, + grad_accum_rounds=grad_accum, + ) + + for k, v in loss_map.items(): + if isinstance(v, torch.Tensor): + prev = loss_sums.get(k, 0.0) + loss_sums[k] = prev + v.detach() + for k, v in step_metrics.items(): + if k in loss_sums: + raise ValueError(f"Metric key {k!r} collides " + "with loss key. Use a " + "different name (e.g. prefix " + "with 'train/').") + prev = metric_sums.get(k, 0.0) + metric_sums[k] = (prev + _coerce_log_scalar( + v, + where=("method.single_train_step()" + f".metrics[{k!r}]"), + )) + + if not method_manages_optimization: + with profiler_region("profiler_region_training_optimizer"): + self.callbacks.on_before_optimizer_step( + method, + iteration=step, + ) + method.optimizers_schedulers_step(step) + method.optimizers_zero_grad(step) + + # Single CPU sync point: materialise GPU tensors to float right before + # logging. + divisor = 1 if method_manages_optimization else grad_accum + metrics = {k: float(v) / divisor for k, v in loss_sums.items()} + metrics.update({k: float(v) / divisor for k, v in metric_sums.items()}) + metrics["step_time_sec"] = (time.perf_counter() - t0) + if not method_manages_optimization: + # This is the local training process's wait for next(data_stream). + # Track it without a cross-rank reduction to avoid adding a + # synchronization to every training step. + metrics["dataloader_time_sec"] = dataloader_time_sec + metrics["vsa_sparsity"] = float(self.training_config.vsa_sparsity) + if self.global_rank == 0 and metrics: + self.tracker.log(metrics, step) + + with profiler_region("profiler_region_training_callbacks"): + self.callbacks.on_training_step_end( + method, + metrics, + iteration=step, + ) + + @profile_region("profiler_region_training_train") def run( self, method: TrainingMethod, @@ -150,84 +258,12 @@ class Trainer: # Allow method-specific optimization flow (e.g. DiffusionNFT). method_manages_optimization = bool(method.manages_optimization()) for step in progress: - t0 = time.perf_counter() - - # Accumulate on GPU during grad-accum; materialise - # to CPU once per step right before logging. - loss_sums: dict[str, float | torch.Tensor] = {} - metric_sums: dict[str, float | torch.Tensor] = {} - if method_manages_optimization: - loss_map, outputs, step_metrics = method.managed_train_step( - data_stream, - step, - ) - for k, v in loss_map.items(): - if isinstance(v, torch.Tensor): - loss_sums[k] = v.detach() - for k, v in step_metrics.items(): - if k in loss_sums: - raise ValueError(f"Metric key {k!r} collides " - "with loss key. Use a " - "different name (e.g. prefix " - "with 'train/').") - metric_sums[k] = _coerce_log_scalar( - v, - where=("method.managed_train_step()" - f".metrics[{k!r}]"), - ) - else: - for accum_iter in range(grad_accum): - batch = next(data_stream) - loss_map, outputs, step_metrics = (method.single_train_step( - batch, - step, - )) - - method.backward( - loss_map, - outputs, - grad_accum_rounds=grad_accum, - ) - - for k, v in loss_map.items(): - if isinstance(v, torch.Tensor): - prev = loss_sums.get(k, 0.0) - loss_sums[k] = prev + v.detach() - for k, v in step_metrics.items(): - if k in loss_sums: - raise ValueError(f"Metric key {k!r} collides " - "with loss key. Use a " - "different name (e.g. prefix " - "with 'train/').") - prev = metric_sums.get(k, 0.0) - metric_sums[k] = (prev + _coerce_log_scalar( - v, - where=("method.single_train_step()" - f".metrics[{k!r}]"), - )) - - if not method_manages_optimization: - self.callbacks.on_before_optimizer_step( - method, - iteration=step, - ) - method.optimizers_schedulers_step(step) - method.optimizers_zero_grad(step) - - # Single CPU sync point: materialise GPU tensors - # to float right before logging. - divisor = 1 if method_manages_optimization else grad_accum - metrics = {k: float(v) / divisor for k, v in loss_sums.items()} - metrics.update({k: float(v) / divisor for k, v in metric_sums.items()}) - metrics["step_time_sec"] = (time.perf_counter() - t0) - metrics["vsa_sparsity"] = float(tc.vsa_sparsity) - if self.global_rank == 0 and metrics: - self.tracker.log(metrics, step) - - self.callbacks.on_training_step_end( + self._run_train_step( method, - metrics, - iteration=step, + data_stream=data_stream, + step=step, + grad_accum=grad_accum, + method_manages_optimization=method_manages_optimization, ) if checkpoint_manager is not None: diff --git a/fastvideo/train/utils/checkpoint.py b/fastvideo/train/utils/checkpoint.py index 7a3d6c105..264b7e5d6 100644 --- a/fastvideo/train/utils/checkpoint.py +++ b/fastvideo/train/utils/checkpoint.py @@ -23,6 +23,7 @@ from torch.distributed.checkpoint.state_dict import ( from torch.distributed.checkpoint.stateful import Stateful from fastvideo.logger import init_logger +from fastvideo.profiler import profile_region logger = init_logger(__name__) @@ -246,6 +247,7 @@ class CheckpointManager: return self.save(step) + @profile_region("profiler_region_training_save_checkpoint") def save(self, step: int) -> None: checkpoint_dir = self._checkpoint_dir(step) dcp_dir = self._dcp_dir(step)