nvalchemi.dynamics.DistributedPipeline

class nvalchemi.dynamics.DistributedPipeline(stages, synchronized=False, debug_mode=False, mesh=None, **dist_kwargs)

Orchestrates multi-rank pipeline execution.

Maps GPU ranks to pipeline stages and coordinates the distributed step loop. Each rank executes only its assigned stage.

Parameters:

• stages (dict[int, BaseDynamics]) – Mapping from rank to its assigned pipeline stage.

• synchronized (bool) –

  If True, insert a global dist.barrier() across all pipeline ranks after every step() call, forcing every rank to complete its current step before any rank proceeds to the next one. This is primarily useful for debugging ordering or deadlock issues because it eliminates all inter-rank timing skew.

  Note

  This is distinct from the per-stage comm_mode parameter on _CommunicationMixin, which controls the blocking behavior of pairwise isend/irecv between adjacent stages. synchronized enforces a global synchronization point across the entire pipeline and will significantly reduce throughput; it should be disabled (False) in production.

• debug_mode (bool)

• mesh (Any)

• dist_kwargs (Any)

stages

Rank-to-stage mapping.

Type:

dict[int, BaseDynamics]

synchronized

Whether a global dist.barrier() is inserted after every step.

Type:

bool

_dist_initialized

Whether this DistributedPipeline instance initialized the distributed process group (used to determine cleanup responsibility).

Type:

bool

Examples

>>> # Context manager usage (recommended):
>>> pipeline = DistributedPipeline(stages={0: opt_stage, 1: md_stage})
>>> with pipeline:
...     pipeline.run()
...
>>> # Manual usage:
>>> pipeline = DistributedPipeline(stages={0: opt_stage, 1: md_stage})
>>> pipeline.init_distributed()
>>> pipeline.setup()
>>> pipeline.run()
>>> pipeline.cleanup()
>>> # Composing multiple pipelines together
>>> full_pipeline = pipe1 | pipe2 | pipe3
>>> with full_pipeline:
...     pipeline.run()
...

cleanup()

Destroy the torch.distributed process group.

Only destroys the process group if it was initialized by this DistributedPipeline instance (via init_distributed()). If the process group was externally initialized (e.g., by torchrun), this method is a no-op.

Return type:

None

property global_rank: int

Get the global rank for this process.

init_distributed()

Initialize the torch.distributed process group.

If torch.distributed is already initialized, this method is a no-op. Otherwise, it calls torch.distributed.init_process_group(**self._dist_kwargs).

The backend and other distributed options are configured via the constructor’s **dist_kwargs parameter.

Notes

When launching with torchrun, the process group is typically already initialized. This method provides a convenient fallback for scripts that do not use torchrun.

Return type:

None

property local_rank: int

Get the local rank for this process.

property local_stage: BaseDynamics

Get the stage associated with the rank this is executed on.

run()

Run the pipeline loop until all stages report done.

After each step(), an all_reduce synchronizes the done flags across all ranks so that every process can observe the global termination state.

Return type:

None

setup()

Wire up prior_rank / next_rank between adjacent stages.

Sorts stages by rank and connects each stage to its predecessor and successor.

Raises:

• ValueError – If fewer than 2 stages are provided, or if adjacent stages have mismatched buffer configurations.

• RuntimeError – If the world size does not match the number of configured pipeline stages.

Return type:

None

step()

Execute one timestep for the local rank’s stage.

The stage (a BaseDynamics subclass) handles both the dynamics step and buffer synchronization.

Supports two modes for the first stage:

Mode 1 (external batch loop): Standard flow where the first stage receives from _prestep_sync_buffers like other stages.

Mode 2 (inflight batching): When the first stage has inflight_mode=True (i.e., a sampler is configured), it builds the initial batch from the sampler and refills graduated samples instead of receiving from a prior stage.

When self.synchronized is True, a global dist.barrier() is issued at the end of each step so that no rank advances until every rank in the pipeline has finished the current step.

Raises:

• RuntimeError – If torch.distributed is not initialized, or if the world size does not match the number of configured pipeline stages.

• KeyError – If the current rank is not in the global rank stage mapping.

Return type:

None