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). ```python 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 ```python 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 ```python 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: ```python dynamics = DemoDynamics( model=model, n_steps=10000, dt=0.5, convergence_hook=ConvergenceHook.from_fmax(threshold=0.05), ) ``` ### With hooks ```python 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 ```python # 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) ```python 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 reprime_on_entry=None, # status codes needing a fresh force compute on entry ) ``` If a sub-stage needs freshly computed forces before it starts integrating (e.g. its `AFTER_COMPUTE` hooks differ from the stage the sample came from), pass its status code in `reprime_on_entry`. Newly entering samples skip one `pre_update`/`post_update` cycle while the shared compute and that stage's `AFTER_COMPUTE` hooks refresh their forces, then integrate normally the following iteration. This is separate from the initial batch-wide force prime and does not affect samples already in that status. ### With torch.compile ```python fused = (opt + md) fused.compile(fullgraph=True, mode="reduce-overhead") with fused: # lazy compilation on context entry result = fused.run(batch) ``` ### CUDA streams: `with` vs bare `run()` Prefer to use context manager: ```python with dynamics: dynamics.run(batch) ``` As best pattern, even without explicit CUDA stream usage, as it in turn supports `compile(mode="reduce-overhead")` Calling `dynamics.run(batch)` without context manager will reuse the current stream. ### 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 ```python 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 ```python 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) ```python pipeline = DistributedPipeline( stages={0: opt_stage, 1: md_stage}, synchronized=False, # True for debugging (adds barriers) ) ``` ### Communication modes Control how inter-rank buffers synchronize: ```python 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: ```python 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 ```python # 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. ```python 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 ```python 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) ```python 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. ```python from nvalchemi.dynamics import GPUBuffer, HostMemory, ZarrData ``` ### DataSink interface ```python 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