nvalchemi.training.hooks.TorchProfilerHook

class nvalchemi.training.hooks.TorchProfilerHook(*, output_dir, activities=None, schedule=None, record_shapes=True, profile_memory=True, with_flops=True, with_stack=False, on_trace_ready_path=None, frequency=1, name='torch', rank_subdirs=True)

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.

Parameters:

• output_dir (Path)

• activities (tuple[ProfilerActivity, ...] | None)

• schedule (Callable[[...], Any] | None)

• record_shapes (bool)

• profile_memory (bool)

• with_flops (bool)

• with_stack (bool)

• on_trace_ready_path (Path | None)

• frequency (Annotated[int, Ge(ge=1)])

• name (str)

• rank_subdirs (bool)

__init__(**data)

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Parameters:

data (Any)

Return type:

None

Methods

__init__(**data)	Create a new model by parsing and validating input data from keyword arguments.
close()	Finalize profiler outputs once.
construct([_fields_set])
copy(*[, include, exclude, update, deep])	Returns a copy of the model.
dict(*[, include, exclude, by_alias, ...])
from_orm(obj)
json(*[, include, exclude, by_alias, ...])
model_construct([_fields_set])	Creates a new instance of the Model class with validated data.
model_copy(*[, update, deep])	!!! abstract "Usage Documentation"
model_dump(*[, mode, include, exclude, ...])	!!! abstract "Usage Documentation"
model_dump_json(*[, indent, ensure_ascii, ...])	!!! abstract "Usage Documentation"
model_json_schema([by_alias, ref_template, ...])	Generates a JSON schema for a model class.
model_parametrized_name(params)	Compute the class name for parametrizations of generic classes.
model_post_init(context, /)	This function is meant to behave like a BaseModel method to initialise private attributes.
model_rebuild(*[, force, raise_errors, ...])	Try to rebuild the pydantic-core schema for the model.
model_validate(obj, *[, strict, extra, ...])	Validate a pydantic model instance.
model_validate_json(json_data, *[, strict, ...])	!!! abstract "Usage Documentation"
model_validate_strings(obj, *[, strict, ...])	Validate the given object with string data against the Pydantic model.
parse_file(path, *[, content_type, ...])
parse_obj(obj)
parse_raw(b, *[, content_type, encoding, ...])
schema([by_alias, ref_template])
schema_json(*[, by_alias, ref_template])
update_forward_refs(**localns)
validate(value)

Attributes

model_computed_fields
model_config	Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].
model_extra	Get extra fields set during validation.
model_fields
model_fields_set	Returns the set of fields that have been explicitly set on this model instance.
stage
output_dir
activities
schedule
record_shapes
profile_memory
with_flops
with_stack
on_trace_ready_path
frequency
name
rank_subdirs