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

Class	Description
BaseDynamics	Abstract base — subclass to create integrators
DemoDynamics	Velocity Verlet reference implementation (testing only)

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:

Layer	Location	Purpose
Active batch	_CommunicationMixin.active_batch	Working set being integrated
Communication buffers	send_buffer / recv_buffer	Pre-allocated Batch.empty() for zero-copy inter-rank transfer
Overflow sinks	DataSink list (priority-ordered)	Staging when active batch is full

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

Method	Purpose
_recv_to_batch(incoming)	Route received data through recv buffer into active batch
_buffer_to_batch(incoming)	Append to active batch, overflow to sinks if full
_batch_to_buffer(mask)	Copy graduated samples into send buffer, defrag active batch
_overflow_to_sinks(batch)	Write to first non-full sink in priority order
_drain_sinks_to_batch()	Pull from sinks back into active batch when room available

---

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