nvalchemi.hooks.TorchProfilerHook

pydantic model nvalchemi.hooks.TorchProfilerHook

Capture PyTorch profiler traces through PhysicsNeMo’s profiler wrapper.

TorchProfilerHook drives PhysicsNeMo’s Profiler (backed by TorchProfileWrapper) so that torch.profiler traces are collected for an nvalchemi workflow without hand-rolling profiler setup, stepping, and finalization. The same hook attaches to both training and dynamics workflows: it recognizes TrainingStage.BEFORE_TRAINING, BEFORE_BATCH, AFTER_BATCH, and AFTER_TRAINING, plus DynamicsStage.BEFORE_STEP and AFTER_STEP.

The profiler starts when the hook enters its context (__enter__) or, if it is dispatched by a workflow without being used as a context manager, lazily on the first supported start stage. It advances the torch.profiler schedule once per batch or dynamics step (at AFTER_BATCH / AFTER_STEP) and finalizes traces at AFTER_TRAINING or when the hook context closes. Register it like any other hook by adding it to a strategy’s or dynamics object’s hooks=[...] list; for dynamics runs it is also valid to wrap the run in a with block so start/finalize bracket exactly the profiled region.

Outputs are written under output_dir (named by name). In distributed runs, or whenever rank_subdirs is set, per-process outputs land in output_dir / rank_<global_rank>, and the optional on_trace_ready_path TensorBoard handler directory is rank-suffixed the same way. Activity selection accepts either ProfilerActivity values or the string aliases "cpu" / "cuda"; None lets PhysicsNeMo pick CPU and CUDA when available.

Examples

Profile a training run by registering the hook alongside the strategy’s other hooks:

>>> import torch
>>> from nvalchemi.hooks.physicsnemo_profiling import TorchProfilerHook
>>> from nvalchemi.training import (
...     EnergyMSELoss, OptimizerConfig, TrainingStrategy, default_training_fn,
... )
>>> profiler = TorchProfilerHook(
...     output_dir="prof/train",
...     activities=("cpu", "cuda"),
...     record_shapes=True,
...     profile_memory=True,
...     with_flops=True,
... )
>>> strategy = TrainingStrategy(
...     models=model,
...     optimizer_configs=OptimizerConfig(
...         optimizer_cls=torch.optim.Adam, optimizer_kwargs={"lr": 1e-3},
...     ),
...     training_fn=default_training_fn,
...     loss_fn=EnergyMSELoss(),
...     num_epochs=1,
...     devices=[torch.device("cuda")],
...     hooks=[profiler],
... )
>>> strategy.run(train_loader)

For dynamics, use the hook as a context manager so the profiler brackets the exact steps you care about:

>>> hook = TorchProfilerHook(output_dir="prof/md", activities=("cuda",))
>>> with hook:
...     dynamics.run(batch, num_steps=100)

Notes

Only one PhysicsNeMo profiler may be active at a time: _start raises a RuntimeError if the global Profiler is already initialized or enabled, so construct and register this hook before any other PhysicsNeMo profiler configuration. The hook is single-use — once finalized it cannot be restarted, and calling it (or re-entering it) after close raises. Finalization happens at AFTER_TRAINING or on context exit; dynamics workflows that never emit an AFTER_TRAINING stage should be run under the with block (or have close called) to flush traces. frequency is a ClassVar-style workflow field, and stage is None because the hook handles multiple stages itself rather than binding to a single one.

field output_dir: Path [Required]

Root directory for PhysicsNeMo profiler outputs.

field activities: tuple[ProfilerActivity, ...] | None = None

PyTorch profiler activities, or None to let PhysicsNeMo choose CPU and CUDA when available.

field schedule: Callable[[...], Any] | None = None

Optional torch.profiler schedule.

field record_shapes: bool = True

Record input tensor shapes in the trace.

field profile_memory: bool = True

Profile memory allocations.

field with_flops: bool = True

Estimate FLOPs for supported operations.

field with_stack: bool = False

Record Python stack traces.

field on_trace_ready_path: Path | None = None

Optional path for PyTorch tensorboard trace handler output.

field frequency: int = 1

Run every N workflow steps.

Constraints:

• ge = 1

field name: str = 'torch'

PhysicsNeMo profiler output name.

field rank_subdirs: bool = True

Write nvalchemi-managed outputs under rank_<global_rank>.