nvalchemi-dynamics-api

nvalchemi Dynamics API

Overview

The dynamics API provides tools to discover available dynamics classes, configure them, and scale simulations up (single GPU) and out (multi-rank pipelines).

from nvalchemi.dynamics import (
    BaseDynamics,
    DemoDynamics,
    FusedStage,
    DistributedPipeline,
    ConvergenceHook,
    Hook,
    DynamicsStage,
    SizeAwareSampler,
    DataSink, GPUBuffer, HostMemory, ZarrData,
    hooks,
)

Available dynamics classes

To find all dynamics classes in a codebase, search for subclasses of BaseDynamics.

Configuring a dynamics run

Basic setup

from nvalchemi.dynamics import DemoDynamics, ConvergenceHook
from nvalchemi.models.demo import DemoModelWrapper
from nvalchemi.data import AtomicData, Batch
import torch

model = DemoModelWrapper()
dynamics = DemoDynamics(
    model=model,
    n_steps=1000,       # total steps for run()
    dt=0.5,             # timestep (fs)
)

# Create batch with required state
data = AtomicData(
    atomic_numbers=torch.tensor([6, 8, 1], dtype=torch.long),
    positions=torch.randn(3, 3),
)
batch = Batch.from_data_list([data])
batch.forces = torch.zeros(3, 3)
batch.energy = torch.zeros(1, 1)

result = dynamics.run(batch)

With convergence detection

dynamics = DemoDynamics(
    model=model,
    n_steps=10000,
    dt=0.5,
    convergence_hook=ConvergenceHook(
        criteria=[
            {"key": "fmax", "threshold": 0.05, "reduce_op": "max"},
            {"key": "energy_change", "threshold": 1e-6},
        ],
        frequency=1,
    ),
)

Shorthand for force-based convergence:

dynamics = DemoDynamics(
    model=model,
    n_steps=10000,
    dt=0.5,
    convergence_hook=ConvergenceHook.from_fmax(threshold=0.05),
)

With hooks

from nvalchemi.dynamics.hooks import MaxForceClampHook, LoggingHook

dynamics = DemoDynamics(
    model=model,
    n_steps=1000,
    dt=0.5,
    hooks=[
        MaxForceClampHook(max_force=10.0),
        LoggingHook(frequency=100),
    ],
)

Scaling up: FusedStage (single GPU, multiple stages)

FusedStage composes multiple dynamics stages on one GPU with a single shared model forward pass per step. Samples migrate between stages via convergence.

Composition with + operator

# Stage 0: geometry optimization
opt = DemoDynamics(
    model=model, n_steps=100, dt=1.0,
    convergence_hook=ConvergenceHook.from_fmax(0.05),
)

# Stage 1: production MD
md = DemoDynamics(model=model, n_steps=500, dt=0.5)

# Compose — auto-registers convergence hook to migrate status 0 → 1
fused = opt + md

result = fused.run(batch)  # runs until all samples reach exit_status

Constructor (explicit)

fused = FusedStage(
    sub_stages=[(0, opt), (1, md)],
    entry_status=0,          # initial sample status
    exit_status=2,           # auto-set to len(sub_stages) if -1
    compile_step=False,      # enable torch.compile
    compile_kwargs=None,     # kwargs for torch.compile
)

With torch.compile

fused = (opt + md)
fused.compile(fullgraph=True, mode="reduce-overhead")

with fused:  # lazy compilation on context entry
    result = fused.run(batch)

How it works

1. Each sample has a status field (integer)
2. Each sub-stage processes only samples matching its status code
3. When a sub-stage's ConvergenceHook fires, converged samples' status increments
4. Samples at exit_status are graduated (no longer updated)
5. run() loops until all samples reach exit_status or sampler is exhausted

Chaining more stages

fused_3 = opt + md + analysis   # 3 sub-stages: status 0, 1, 2
fused_3 = fused + extra_stage   # append to existing FusedStage

Scaling out: DistributedPipeline (multi-rank)

DistributedPipeline chains dynamics stages across multiple ranks using torch.distributed. Each rank runs one stage; converged samples are sent to the next rank.

Composition with | operator

opt_stage = DemoDynamics(model=model, n_steps=100, dt=1.0)   # rank 0
md_stage = DemoDynamics(model=model, n_steps=500, dt=0.5)    # rank 1

pipeline = opt_stage | md_stage

with pipeline:   # initializes torch.distributed + setup
    pipeline.run()

Constructor (explicit)

pipeline = DistributedPipeline(
    stages={0: opt_stage, 1: md_stage},
    synchronized=False,   # True for debugging (adds barriers)
)

Communication modes

Control how inter-rank buffers synchronize:

stage = DemoDynamics(
    model=model, n_steps=100, dt=0.5,
    comm_mode="async_recv",   # default: deferred blocking
    # comm_mode="sync",       # immediate blocking (debugging)
    # comm_mode="fully_async", # maximum overlap
)

The default comm_mode is "async_recv". The three modes differ in when blocking occurs:

• "sync": irecv completes inline in _prestep_sync_buffers; simplest and good for debugging.
• "async_recv": irecv is posted in _prestep_sync_buffers, but wait() is deferred to _complete_pending_recv for communication overlap.
• "fully_async": send and receive are both deferred for maximum overlap. Pending sends from the prior step are drained at the start of the next _prestep_sync_buffers.

Pre-allocated buffers

For high-throughput pipelines, pre-allocate send/recv buffers:

from nvalchemi.dynamics.base import BufferConfig

buffer_cfg = BufferConfig(
    num_systems=100,   # max graphs
    num_nodes=5000,    # total node capacity
    num_edges=20000,   # total edge capacity
)

stage = DemoDynamics(
    model=model, n_steps=100, dt=0.5,
    buffer_config=buffer_cfg,
)

Buffers are lazily initialized on the first step using the first concrete batch as a template for attribute keys, dtypes, and shapes. This means the first step has slightly more overhead.

Adjacent stages must use identical BufferConfig values. This is validated in DistributedPipeline.setup().

Buffer semantics and communication

Three buffer layers

The dynamics framework manages data flow through three layers:

Communication protocol (DistributedPipeline)

Each pipeline step follows a four-phase protocol:

1. _prestep_sync_buffers() zeros the send buffer and posts irecv from the prior rank.
2. _complete_pending_recv() waits on deferred receive, routes into the active batch, and drains overflow sinks.
3. step() runs dynamics integration.
4. _poststep_sync_buffers(converged_indices) extracts converged samples into the send buffer and sends them to the next rank.

Deadlock prevention: when no samples converge, an empty send buffer is still sent so the downstream irecv completes.

Back-pressure

When send_buffer has limited capacity (via BufferConfig):

• Only min(converged_count, remaining_capacity) samples are extracted
• Excess converged samples remain in the active batch as no-ops. Their positions and velocities are saved before the integrator and restored after it runs.
• Without BufferConfig, all converged samples are sent without constraints (backward compatible).

Buffer lifecycle: put/defrag/zero

# Pre-allocated buffer created via Batch.empty()
buffer = Batch.empty(num_systems=100, num_nodes=5000, num_edges=20000, template=batch)

# Copy selected graphs into buffer (Warp GPU kernels, float32 only)
mask = converged_mask  # bool tensor, True = copy this graph
buffer.put(src_batch, mask)

# Remove copied graphs from source in-place
src_batch.defrag()

# Reset buffer for reuse (preserves allocated memory)
buffer.zero()

Important: Batch.put() uses Warp GPU kernels that only handle float32 attributes. Adjacent pipeline stages must have identical BufferConfig values.

Data routing methods

Inflight batching with SizeAwareSampler

For streaming workflows, SizeAwareSampler manages dataset access with bin-packing for size-matched batching. As samples converge and leave the batch, new samples are pulled from the dataset.

from nvalchemi.dynamics import SizeAwareSampler

sampler = SizeAwareSampler(
    dataset=dataset,                # must have __len__, __getitem__, get_metadata
    max_atoms=1000,                 # max total atoms per batch (None = auto from GPU)
    max_edges=5000,                 # max total edges per batch
    max_batch_size=32,              # max graphs per batch
    bin_width=10,                   # group samples by atom count bins
    shuffle=False,
    max_gpu_memory_fraction=0.8,    # for auto max_atoms estimation
)

# Build initial batch via greedy bin-packing
batch = sampler.build_initial_batch()

# Request a replacement sample that fits constraints
replacement = sampler.request_replacement(num_atoms=50, num_edges=200)

# Check if all samples consumed
sampler.exhausted  # bool

How inflight replacement works (_refill_check)

When refill_frequency triggers (every N steps), _refill_check():

1. Identifies graduated graphs (status >= exit_status)
2. Writes graduated graphs to sinks
3. Extracts remaining graphs via Batch.index_select
4. Requests replacements from sampler (one per graduated slot, matching atom/edge budget)
5. Appends replacements via Batch.append
6. Rebuilds status (replacements get 0) and fmax (replacements get inf) tensors

This produces a new Batch object, not an in-place mutation. It returns None when the sampler is exhausted and no active samples remain.

With FusedStage

opt = DemoDynamics(
    model=model, n_steps=100, dt=1.0,
    sampler=sampler,
    refill_frequency=1,    # check for replacements every N steps
    convergence_hook=ConvergenceHook.from_fmax(0.05),
)
md = DemoDynamics(model=model, n_steps=500, dt=0.5)

fused = opt + md
with fused:
    result = fused.run()  # no batch arg — built from sampler

With DistributedPipeline (first stage only)

first_stage = DemoDynamics(
    model=model, n_steps=100, dt=1.0,
    sampler=sampler,
    refill_frequency=1,
)
pipeline = first_stage | second_stage
with pipeline:
    pipeline.run()  # first stage auto-refills from sampler

Data sinks

Sinks store graduated/snapshot data. Used by SnapshotHook and the communication layer.

from nvalchemi.dynamics import GPUBuffer, HostMemory, ZarrData

DataSink interface

class DataSink(ABC):
    def write(self, batch: Batch, mask: Tensor | None = None) -> None: ...
    def read(self) -> Batch: ...
    def zero(self) -> None: ...           # clear contents
    def drain(self) -> Batch: ...         # read() + zero()
    def is_full(self) -> bool: ...        # len(self) >= capacity
    def __len__(self) -> int: ...
    @property
    def capacity(self) -> int: ...

Implementations

GPUBuffer — GPU-resident, pre-allocated.

gpu_sink = GPUBuffer(
    capacity=100,       # max graphs
    max_atoms=5000,     # total node capacity
    max_edges=20000,    # total edge capacity
    device="cuda",
)

HostMemory — CPU-resident list.

cpu_sink = HostMemory(capacity=1000)

ZarrData — Disk-backed persistent storage.

zarr_sink = ZarrData(
    store="trajectory.zarr",   # path, S3 URI, or dict
    capacity=1_000_000,
)

Using sinks with dynamics

from nvalchemi.dynamics.hooks import SnapshotHook

dynamics = DemoDynamics(
    model=model, n_steps=1000, dt=0.5,
    sinks=[gpu_sink, zarr_sink],   # for pipeline communication
    hooks=[
        SnapshotHook(sink=zarr_sink, frequency=10),
    ],
)

ConvergenceHook reference

ConvergenceHook(
    criteria=[                              # AND semantics (all must pass)
        {
            "key": "fmax",                  # batch attribute to check
            "threshold": 0.05,              # convergence threshold
            "reduce_op": "max",             # min, max, norm, mean, sum
            "reduce_dims": -1,              # dimensions to reduce
            "custom_op": None,              # custom callable
        },
    ],
    source_status=0,                        # check samples with this status
    target_status=1,                        # migrate converged samples
    frequency=1,                            # check every N steps
)

# Shorthand
ConvergenceHook.from_fmax(
    threshold=0.05,
    source_status=None,
    target_status=None,
    frequency=1,
)

Full workflow example

import torch
from nvalchemi.data import AtomicData, Batch
from nvalchemi.data.datapipes import AtomicDataZarrReader, Dataset
from nvalchemi.models.demo import DemoModelWrapper
from nvalchemi.dynamics import (
    DemoDynamics, ConvergenceHook, SizeAwareSampler, ZarrData,
)
from nvalchemi.dynamics.hooks import MaxForceClampHook, LoggingHook, SnapshotHook

# Model
model = DemoModelWrapper()

# Dataset + sampler for inflight batching
reader = AtomicDataZarrReader("structures.zarr")
dataset = Dataset(reader, device="cuda")
sampler = SizeAwareSampler(
    dataset=dataset, max_atoms=1000, max_edges=5000, max_batch_size=32,
)

# Output sink
output = ZarrData("results.zarr", capacity=100000)

# Stage 0: optimize
opt = DemoDynamics(
    model=model, n_steps=500, dt=1.0,
    sampler=sampler, refill_frequency=1,
    convergence_hook=ConvergenceHook.from_fmax(0.05),
    hooks=[MaxForceClampHook(max_force=10.0)],
)

# Stage 1: MD
md = DemoDynamics(
    model=model, n_steps=1000, dt=0.5,
    convergence_hook=ConvergenceHook(
        criteria=[{"key": "fmax", "threshold": 0.01}],
    ),
    hooks=[
        LoggingHook(frequency=100),
        SnapshotHook(sink=output, frequency=50),
    ],
)

# Compose and run
fused = opt + md
with fused:
    fused.run()  # batch built from sampler, runs until exhausted