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(backend="csv", log_path="md_log.csv", 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. This is pipeline parallelism for dynamics, distinct from data-parallel multi-GPU training (DDP).

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,

---

*Content truncated.*