nvalchemi.dynamics.FusedStage

class nvalchemi.dynamics.FusedStage(sub_stages, *, entry_status=0, exit_status=-1, compile_step=False, compile_kwargs=None, init_fn=None, **kwargs)

Composite dynamics engine fusing multiple sub-stages on a single GPU.

FusedStage composes multiple BaseDynamics sub-stages to share one Batch and one model forward pass per step, avoiding redundant forward passes when multiple simulation phases (e.g., relaxation then MD) operate on the same batch.

Unlike BaseDynamics, ``step(batch)`` is overridden. Instead of the standard pre_update → compute → post_update loop, FusedStage performs: (1) a single compute() call on the full batch, then (2) iterates over sub-stages, applying masked_update(batch, mask) on each sub-stage’s dynamics for samples whose batch.status matches that sub-stage’s status code. Only ONE forward pass happens per step regardless of the number of sub-stages. ``run(batch)`` is also overridden — the n_steps attribute (inherited from BaseDynamics) and any n_steps argument passed to run() are both the maximum number of steps; the loop runs until all samples have migrated to the exit_status, the sampler is exhausted, or n_steps is reached. Convergence-driven migration is handled by ConvergenceHook instances auto-registered between adjacent sub-stages: when samples converge in sub-stage i, their batch.status is updated to sub-stage i+1’s code, causing them to be processed by the next dynamics on the following step. The + operator composes sub-stages: dyn_a + dyn_b creates a FusedStage, and fused + dyn_c appends a third sub-stage. The | operator (inherited from BaseDynamics via _CommunicationMixin) creates a DistributedPipeline for multi-rank execution instead.

Developers generally do NOT subclass FusedStage. Instead, create BaseDynamics subclasses (integrators) and compose them using +. FusedStage handles orchestration automatically. The key requirement is that sub-stage dynamics must implement masked_update correctly (inherited from BaseDynamics) and that the batch must have a status tensor.

Hook Firing Semantics

Because FusedStage shares a single forward pass across all sub-stages, hook firing differs from standalone BaseDynamics execution. The following hooks fire on each sub-stage during _step_impl:

Fired on sub-stages (in order):

• BEFORE_STEP — at the start of each fused step, before any work.

• AFTER_COMPUTE — after the shared model forward pass completes.

• BEFORE_PRE_UPDATE — before each sub-stage’s masked_update (fires even when no samples match the sub-stage’s status code).

• AFTER_POST_UPDATE — after each sub-stage’s masked_update (fires even when no samples match the sub-stage’s status code).

• AFTER_STEP — after all masked updates are complete.

• ON_CONVERGE — when a sub-stage’s _check_convergence detects converged samples.

NOT fired on sub-stages:

• BEFORE_COMPUTE — the forward pass is shared across all sub-stages, not executed per-sub-stage; there is no meaningful “before compute” point for individual sub-stages.

• AFTER_PRE_UPDATE — masked_update combines pre_update and post_update atomically; there is no intermediate hook point.

• BEFORE_POST_UPDATE — same reason as AFTER_PRE_UPDATE.

Step count semantics: Each sub-stage’s step_count is incremented alongside the FusedStage’s own step_count after every fused step, ensuring that hook frequency (e.g., every_n_steps) is respected correctly across all sub-stages.

param sub_stages:

Ordered (status_code, dynamics) pairs. Status codes are auto-assigned starting from 0 when using the + operator.

type sub_stages:

list[tuple[int, BaseDynamics]]

param entry_status:

Status code assigned to incoming samples (default: 0).

type entry_status:

int

param exit_status:

Status code that triggers graduation to the next pipeline stage. Auto-set to len(sub_stages) (one past the last sub-stage code).

type exit_status:

int

param compile_step:

If True, replace self.step with torch.compile(self.step, **compile_kwargs).

type compile_step:

bool

param compile_kwargs:

Keyword arguments forwarded to torch.compile.

type compile_kwargs:

dict

param **kwargs:

Additional keyword arguments forwarded to BaseDynamics.

sub_stages

Ordered (status_code, dynamics) pairs.

Type:

list[tuple[int, BaseDynamics]]

entry_status

Status code for incoming samples.

Type:

int

exit_status

Status code that triggers graduation.

Type:

int

compile_step

Whether the step method is compiled.

Type:

bool

compile_kwargs

Arguments passed to torch.compile.

Type:

dict

__needs_keys__

Union of all sub-stage __needs_keys__ sets. Populated automatically during __init__.

Type:

set[str]

__provides_keys__

Union of all sub-stage __provides_keys__ sets. Populated automatically during __init__.

Type:

set[str]

Examples

>>> from nvalchemi.dynamics import FusedStage, BaseDynamics
>>> dynamics0 = BaseDynamics(model=model)
>>> dynamics1 = BaseDynamics(model=model)
>>> fused = FusedStage(sub_stages=[(0, dynamics0), (1, dynamics1)])
>>> fused.exit_status
2

static all_complete(batch, exit_status)

Check if all samples have reached the exit status.

Parameters:

• batch (Batch) – The current batch.

• exit_status (int) – The status code that indicates completion.

Returns:

True if every sample has status == exit_status.

Return type:

bool

compile(**kwargs)

Compile the fused step with torch.compile.

Merges kwargs with any compile_kwargs stored at construction time (values passed here take precedence), then wraps _step_impl with torch.compile. Calling this method also sets compile_step = True so that the step dispatch path uses the compiled callable.

This method is idempotent in intent but will re-compile if called again (e.g. with different kwargs).

Parameters:

**kwargs (Any) – Keyword arguments forwarded to torch.compile. Merged with compile_kwargs from __init__; values here win.

Returns:

This instance, enabling fluent chaining such as fused.compile(fullgraph=True).run(batch).

Return type:

FusedStage

refill_check(batch, exit_status)

Replace graduated samples and clear stale convergence indices.

Delegates to the parent BaseDynamics.refill_check() to remove graduated graphs and append replacements from the sampler. When the batch composition changes, _last_converged is cleared on this FusedStage and all its sub-stages so that subsequent hooks do not receive an invalid converged_mask.

Parameters:

• batch (Batch) – The current batch with a status field.

• exit_status (int) – Status code indicating graduation.

Returns:

A new batch with graduated graphs replaced by fresh samples, or None if no active samples remain.

Return type:

Batch | None

register_fused_hook(hook)

Register a hook that fires at the FusedStage level on the full batch.

Unlike hooks registered on individual sub-stages (which only receive the sub-batched view), fused hooks observe the complete batch at the BEFORE_STEP, AFTER_STEP, BEFORE_COMPUTE, and AFTER_COMPUTE stages of every fused step.

Parameters:

hook (Hook) – The hook to register. Only BEFORE_STEP and AFTER_STEP stages are meaningful at the fused level; other stages are silently accepted but will not fire during normal execution.

Raises:

ValueError – If hook.frequency is not a positive integer.

Return type:

None

run(batch=None, n_steps=None)

Run the fused stage until all samples converge or the sampler is exhausted.

Supports two modes of execution:

Mode 1 (external batch loop): When batch is provided, runs the dynamics until all_complete or until n_steps have been executed (whichever comes first).

Mode 2 (inflight batching): When batch is None and a sampler is configured, builds the initial batch from the sampler and replaces graduated samples every refill_frequency steps.

Note

In Mode 2, refill_check replaces graduated samples by extracting remaining graphs via Batch.index_select(), requesting replacements from the sampler, and appending them via Batch.append(). This produces a new Batch object; the batch = result reassignment in the loop body updates the local reference. None is returned when the sampler is exhausted and no active samples remain, which triggers termination.

Parameters:

• batch (Batch | None, optional) – The initial batch. If None, uses the sampler to build one.

• n_steps (int | None, optional) – Maximum number of steps to run. When None, falls back to self.n_steps. When both are None, the loop runs until all_complete (Mode 1) or sampler exhaustion (Mode 2). Sub-stages that have no exit criterion (e.g. a plain MD stage) will loop forever without a step limit, so always pass n_steps when such a stage is the final sub-stage. Note: sub-stages with n_steps set use that value as a per-system step budget for automatic migration to the next stage.

Returns:

The batch after all steps, or None if the sampler was exhausted and all samples graduated.

Return type:

Batch | None

Raises:

ValueError – If batch is None and no sampler is configured.

step(batch)

Execute one fused step: single forward pass + masked updates.

If compile_step=True was set, this delegates to the compiled step implementation.

Parameters:

batch (Batch) – The batch with a status field.

Returns:

The updated batch, and a 1-D integer tensor of sample indices that newly graduated (reached exit_status) during this step, or None if no samples graduated.

Return type:

tuple[Batch, torch.Tensor | None]

Parameters:

• sub_stages (list[tuple[int, BaseDynamics]])

• entry_status (int)

• exit_status (int)

• compile_step (bool)

• compile_kwargs (dict[str, Any] | None)

• init_fn (Callable[[Batch], None] | None)

• kwargs (Any)