PhysicsNeMo domain_parallel#

In scientific AI applications, the parallelization techniques to enable state of the art models are different from those used in training large language models. PhysicsNeMo introduces a parallelization primitive called a ShardTensor that is designed for large-input AI applications to enable domain parallelization.

ShardTensor provides a distributed tensor implementation that supports uneven sharding across devices. It is a subclass of torch.Tensor that interoperates with PyTorch’s DTensor and plain tensors, while adding flexibility for cases where different ranks may have different local tensor sizes.

A key feature of ShardTensor is automatic promotion. When a plain nn.Module weight meets a sharded activation, ShardTensor promotes the weight to a replicated distributed tensor and reduces its gradient over the domain mesh in the backward pass. Standard nn.Module models nn.Module models work unmodified on sharded inputs - distribute_module is no longer needed therefore work unmodified on sharded inputs. distribute_module is no longer needed or recommended. For an additional data parallel axis, use Distributed Data Parallel (DDP) when parameters are plain tensors, or Fully Sharded Data Parallel 2 (FSDP2) (torch.distributed.fsdp.fully_shard) when parameters are sharded. torch.compile is supported, with the caveat that sequence-sharded ring attention must stay outside compiled regions.

Note

PhysicsNeMo 2.3 changes the ``ShardTensor`` input contract. Prior to version 2.3, ShardTensor expected all inputs to be either ShardTensor or DTensor. Those inputs are still supported, but plain torch.Tensor inputs enables DDP and FSDP2 support now as well. We recommend this newer way of using ShardTensor.

ShardTensor#

class physicsnemo.domain_parallel.ShardTensor(
local_tensor: Tensor,
spec: ShardTensorSpec,
*,
requires_grad: bool,
)[source]#

Bases: Tensor

A distributed tensor class with support for uneven data sharding.

Similar to PyTorch’s native DTensor but with more flexibility for uneven data sharding. Leverages a very similar API to DTensor (identical where possible) but deliberately tweaks routines to avoid implicit assumptions about tensor sharding.

The key differences from DTensor are:

  • Supports uneven sharding where different ranks can have different local tensor sizes

  • Tracks and propagates shard size information across operations

  • Handles redistribution of unevenly sharded tensors

  • Provides custom collective operations optimized for uneven sharding

A Partial placement always describes a pending reduction. Its local tensor is this rank’s contribution to the logical value, not a replicated copy of that value. Resolve it with redistribute() or full_tensor() before treating the local data as complete. This rule also applies to gradients: backward paths that retain a fully replicated gradient label it Replicate rather than using Partial as layout-only metadata.

Like DTensor, operations are dispatched through PyTorch’s dispatcher system. Most operations work by:

  1. Converting inputs to local tensors

  2. Performing the operation locally

  3. Constructing a new ShardTensor with appropriate sharding spec

  4. Handling any needed communication between ranks

The class provides methods for:

  • Converting to/from local tensors

  • Redistributing between different sharding schemes

  • Performing collective operations like all_gather and reduce_scatter

  • Basic tensor operations that maintain sharding information

_local_tensor#

The local tensor data on this rank.

Type:

torch.Tensor

_spec#

The specification defining sharding scheme and metadata.

Type:

ShardTensorSpec

backward(*args, **kwargs)[source]#

Perform backward pass for ShardTensor.

Handles the redistribution of the tensor to resolve any partial placements before calling backward on the local tensor.

Parameters:

  • *args – Positional arguments passed to torch.Tensor.backward.

  • **kwargs – Keyword arguments passed to torch.Tensor.backward.

property device_mesh: DeviceMesh#

Return the DeviceMesh that this tensor is distributed over.

classmethod from_dtensor(
dtensor: DTensor,
) → ShardTensor[source]#

Convert a DTensor to a ShardTensor.

Differentiable when dtensor is non-leaf (has a grad_fn). Spec is inferred from the DTensor (chunk-based, no communication).

Parameters:

dtensor (DTensor) – DTensor to convert.

Returns:

Equivalent ShardTensor with the same local tensor and inferred spec.

Return type:

ShardTensor

static from_local(
local_tensor: Tensor,
device_mesh: DeviceMesh | None = None,
placements: Sequence[Placement] | None = None,
sharding_shapes: str | dict[int, list[tuple[int, ...]]] = 'infer',
global_shape: tuple[int, ...] | None = None,
) → ShardTensor[source]#

Generate a new ShardTensor from local torch tensors.

Uses device mesh and placements to infer global tensor properties. No restriction is made on forcing tensors to have equal shapes locally. Instead, the requirement is that tensor shapes could be concatenated into a single tensor according to the placements.

Parameters:

  • local_tensor (torch.Tensor) – Local chunk of tensor. All participating tensors must be of the same rank and concatenatable across the mesh dimensions.

  • device_mesh (Optional[DeviceMesh], optional) – Target device mesh. If not specified, will use the current mesh.

  • placements (Optional[Sequence[Placement]], optional) – Target placements. Must have same number of elements as device_mesh.ndim.

  • sharding_shapes (Union[str, Dict[int, List[Tuple[int, ...]]]], default="infer") –

    Controls how shard tensor spec is generated:

    • "chunk": Use torch.chunk shapes to infer shapes from global shape (no communication). Requires global_shape.

    • "infer": Use collective communication to infer shapes from mesh neighbors.

    • Manual dict mapping mesh dim to list of shard shapes: Use provided shapes. Must pass on each rank.

  • global_shape (Optional[Tuple[int, ...]], optional) – Global shape of the full tensor across all ranks. Required when sharding_shapes="chunk" (it is what makes that mode communication-free); ignored for "infer" and dict modes.

Returns:

A new ShardTensor instance.

Return type:

ShardTensor

full_tensor(
*,
grad_placements: Sequence[Placement] | None = None,
) → Tensor[source]#

Gather the full tensor from all ranks.

Redistributes to Replicate placement on all mesh dimensions and returns the local tensor.

Parameters:

grad_placements (Optional[Sequence[Placement]], optional) – Future layout of gradients. If provided, gradients will be constructed with this placement scheme during backward pass. Specifying Partial declares that each returned local gradient is an additive contribution with a pending reduction.

Returns:

The full gathered tensor, identical on all ranks.

Return type:

torch.Tensor

classmethod get_promotion_mode() → TensorPromotionMode[source]#

Return the active plain-tensor promotion mode (defaults to SILENT).

property grad: ShardTensor | None#

Return the accumulated gradient, wrapped as a ShardTensor.

If the autograd engine stored a distributed gradient, its placements are preserved rather than inherited from the primal tensor. Plain local gradients cannot carry a pending reduction and therefore normalize any primal Partial placement to Replicate. If no gradient has been accumulated yet, returns None.

property grad_dtype#

the local tensor’s dtype.

Overriding shields the read from __torch_function__, which would otherwise fall back to a non-leaf DTensor whose grad_dtype getter raises during Dynamo fake conversion. Mirrors .grad_fn / .is_leaf / .grad.

Type:

dtype of this tensor’s gradient (newer PyTorch)

property grad_fn#

Return the stored grad_fn without re-entering __torch_function__. Without this override, .grad_fn (a C-level getset_descriptor on torch.Tensor) re-enters ShardTensor.__torch_function__ whenever someone reads it, falls back via _torch_function_fallback_via_dtensor(), and the fallback constructs a new temporary DTensor via _ShardTensorToDTensor.apply(self) – whose .grad_fn (a _ShardTensorToDTensorBackward BackwardCFunction instance) is what the caller actually receives. On newer PyTorch that node’s .next_functions accessor raises a “legacy access pattern” error, which is exactly what makes AOTAutograd.setup_stacktrace_preservation_hooks (and our own diagnostic dump_grad_fn_chain) fail when they try to walk the autograd graph of a ShardTensor output. Mirrors the same shielding pattern already used by .is_leaf and .grad.

property is_leaf: bool#

Whether this tensor is a leaf in the autograd graph.

offsets(mesh_dim: int | None = None) → list[int] | int[source]#

Get offsets of shards along a mesh dimension.

Parameters:

mesh_dim (Optional[int], optional) – Mesh dimension to get offsets for. If None, returns all offsets.

Returns:

List of offsets for shards along all dimensions, or single offset if mesh_dim is specified.

Return type:

Union[List[int], int]

classmethod patches_enabled() → bool[source]#

Check whether patches are enabled for this class.

Returns:

True if shard patches are enabled, False otherwise. Default is False until a ShardTensor is constructed.

Return type:

bool

property placements: tuple[Placement, ...]#

Return the placement strategy for each mesh dimension.

classmethod promotion_mode(
mode: TensorPromotionMode,
)[source]#

Temporarily set the promotion mode, restoring the previous one on exit.

redistribute(
device_mesh: DeviceMesh | None = None,
placements: Sequence[Placement] | None = None,
*,
async_op: bool = False,
) → ShardTensor[source]#

Redistribute tensor across device mesh with new placement scheme.

Like DTensor.redistribute but uses custom layer for shard redistribution that supports uneven sharding.

Parameters:

  • device_mesh (Optional[DeviceMesh], optional) – Target device mesh. Uses current mesh if None.

  • placements (Optional[Sequence[Placement]], optional) – Target placement scheme. Required.

  • async_op (bool, default=False) – Whether to run asynchronously.

Returns:

Redistributed ShardTensor with new placement scheme.

Return type:

ShardTensor

Raises:

RuntimeError – If placements is not specified or contains invalid placements (e.g., Partial placements or negative shard dimensions).

classmethod register_dispatch_handler(
op: OpOverload,
handler: Callable,
) → None[source]#

Register a handler for a specific PyTorch operator in the dispatch system.

Parameters:

  • op (torch._ops.OpOverload) – The PyTorch operator to register a handler for.

  • handler (Callable) – The handler function to call when the operator is invoked.

classmethod register_function_handler(
func: Callable,
handler: Callable,
) → None[source]#

Register a handler for a Python-level function or method.

Before the handler is called, plain non-scalar torch.Tensor arguments are promoted to Replicate DTensors on the mesh of the accompanying ShardTensor arguments (honoring the promotion mode) – the same contract the DTensor fallback provides – so handlers uniformly receive distributed tensors with a _spec.

Parameters:

  • func (Callable) – The Python function to register a handler for.

  • handler (Callable) – The handler function to call when the function is invoked.

classmethod register_named_function_handler(
func_name: str,
handler: Callable,
) → None[source]#

Register a named function registered via torch.library.custom_op.

Parameters:

  • func_name (str) – The string name of the custom op (e.g., "module.function_name.default").

  • handler (Callable) – The handler function to call when the function is invoked.

property requires_grad: bool#

Whether this tensor requires gradient computation.

Returns True if either the wrapper tensor or the underlying local tensor has requires_grad set.

requires_grad_(
requires_grad: bool = True,
) → ShardTensor[source]#

Set requires_grad in-place on both the wrapper and local tensor.

Parameters:

requires_grad (bool, optional) – Whether to enable gradient tracking. Default is True.

Returns:

self, for method chaining.

Return type:

ShardTensor

classmethod set_promotion_mode(
mode: TensorPromotionMode,
) → None[source]#

Set the plain-tensor promotion mode.

mode may be a TensorPromotionMode or an equivalent string ("disabled", "warn", "silent"), which is coerced.

to_local(
*,
grad_placements: Sequence[Placement] | None = None,
) → Tensor[source]#

Get local tensor from this ShardTensor.

Parameters:

grad_placements (Optional[Sequence[Placement]], optional) – Future layout of gradients. If provided, gradients will be constructed with this placement scheme during backward pass. Specifying Partial declares that each returned local gradient is an additive contribution with a pending reduction.

Returns:

Local tensor. Shape may vary between ranks for sharded tensors.

Return type:

torch.Tensor

Notes

A Partial placement is not resolved: this returns the unreduced local contribution. Use full_tensor() if you need a reduced value.

class physicsnemo.domain_parallel.ShardTensorSpec(mesh: ~torch.distributed.device_mesh.DeviceMesh, placements: tuple[~torch.distributed.tensor.placement_types.Placement, ...], tensor_meta: ~torch.distributed.tensor._dtensor_spec.TensorMeta | None = None, shard_order: tuple[~torch.distributed.tensor._dtensor_spec.ShardOrderEntry, ...] = None, use_strided_shard_as_shard_order: bool | None = None, *, _local_shape: ~torch.Size | None = <factory>, _sharding_shapes: dict[int, tuple[tuple[int, ...], ...]] | None = <factory>)[source]#

Bases: DTensorSpec

A distributed tensor specification that tracks sharding information.

This class extends DTensorSpec to include information about global placements of shards. This is useful when the tensor is distributed in an uneven or unexpected way.

Placement metadata describes the local data’s current semantic state. In particular, Partial means the local tensor is an additive contribution with a pending collective reduction. Fully reduced data must use Replicate; Partial is never a layout-only label, including for gradients.

_local_shape#

The shape of the local shard of the tensor.

Type:

Optional[torch.Size]

_sharding_shapes#

Mapping from mesh dimension to shard shapes. Keys are mesh dimensions, values are tuples of plain int tuples representing shard shapes along that dimension. Shard shapes are only tracked along the sharded dimensions, not replicated dimensions.

Storage type note: we deliberately use plain tuple[int, ...] rather than torch.Size here. torch.Size is special-cased by PyTorch’s symbolic shape machinery: when a ShardTensor is fakeified by dynamo, any torch.Size stored in this dict has its contained ints converted into unbacked SymInt``s. Those SymInts then orphan whenever an op's output drops or filters ``_sharding_shapes (e.g. Partial-only outputs from reductions), producing PendingUnbackedSymbolNotFound errors during AOT tracing. Plain Python int tuples don’t trigger this path.

Type:

Optional[dict[int, Tuple[Tuple[int, …], …]]]

property local_shape: Size#

Get the shape of the local shard.

Returns:

Shape of local tensor shard.

Return type:

torch.Size

Raises:

RuntimeError – If local shape has not been set.

offsets(
mesh_dim: int | None = None,
) → tuple[int, ...] | int[source]#

Calculate offsets for the local shard within the global tensor.

Returns the effective offset of this tensor along sharded dimensions, as if it was all collected into one device and you wanted to slice it to recover the local slice.

Parameters:

mesh_dim (Optional[int], optional) – If provided, return offset only for this mesh dimension.

Returns:

Tuple of offsets for each mesh dimension, or single offset if mesh_dim is specified.

Return type:

Union[Tuple[int, …], int]

sharding_shapes(
mesh_dim: int | None = None,
) → dict[int, tuple[tuple[int, ...], ...]] | tuple[tuple[int, ...], ...][source]#

Get the shapes of shards along specified mesh dimensions.

Parameters:

mesh_dim (Optional[int], optional) – If provided, return shapes only for this mesh dimension.

Returns:

Dictionary of shard shapes by mesh dim if mesh_dim is None, or tuple of shapes for the specific mesh dimension.

Return type:

Union[Dict[int, Tuple[Tuple[int, …], …]], Tuple[Tuple[int, …], …]]

For debugging purposes, you can modify the promotion behavior of ShardTensor.

class physicsnemo.domain_parallel.TensorPromotionMode(*values)[source]#

Bases: Enum

How a plain torch.Tensor is handled when it meets a ShardTensor in an intercepted op.

Such a plain tensor is typically an unsharded model weight (all-gathered by FSDP2 in its pre-forward hook, or replicated under DDP).

DISABLED#

No promotion; plain tensors pass through to DTensor routing unchanged (mixing a non-scalar plain tensor with sharded data raises – the historical behavior).

Type:

TensorPromotionMode

WARN#

Promote each plain tensor to a Replicate distributed tensor on the accompanying distributed argument’s mesh, warning on every promotion.

Type:

TensorPromotionMode

SILENT#

Same as WARN but without emitting warnings. The default.

Type:

TensorPromotionMode

Utility Functions#

physicsnemo.domain_parallel.scatter_tensor(
tensor: Tensor,
global_src: int,
mesh: DeviceMesh,
placements: tuple[Placement, ...],
global_shape: Size | None = None,
dtype: dtype | None = None,
requires_grad: bool = False,
) → ShardTensor[source]#

Distribute a tensor from source rank across devices on the mesh.

Takes a tensor that exists on a single source rank and distributes it across a device mesh according to the specified placement scheme. For multi-dimensional meshes, it performs a flattened scatter operation before constructing the sharded tensor.

Parameters:

  • tensor (torch.Tensor) – The tensor to distribute. Must exist on source rank; can be None on other ranks.

  • global_src (int) – Global rank ID of the source process.

  • mesh (DeviceMesh) – Device mesh defining the process topology.

  • placements (Tuple[Placement, ...]) – Tuple of placement specifications defining how to distribute the tensor.

  • global_shape (Optional[torch.Size], optional) – Global shape of the tensor. If None, will be broadcast from source.

  • dtype (Optional[torch.dtype], optional) – Data type of the tensor. If None, will be broadcast from source.

  • requires_grad (bool, default=False) – Whether the resulting ShardTensor requires gradients.

Returns:

The distributed tensor with specified placements.

Return type:

ShardTensor

Raises:

ValueError – If global_src is not an integer or not in the mesh.

physicsnemo.domain_parallel.sync_module_over_mesh(
module: Module,
mesh: DeviceMesh,
src_mesh_rank: int = 0,
verify: bool = False,
) → None[source]#

Broadcast a module’s plain parameters and buffers over a mesh axis.

Synchronizes every plain (non-distributed) parameter and buffer of module across mesh’s process group, from the rank at position src_mesh_rank of the mesh. Distributed tensors (DTensor, ShardTensor) are skipped because arbitrary distributed layouts cannot safely be broadcast by this utility. Their initialization remains the caller’s responsibility. Source-based constructors such as distribute_tensor and scatter_tensor() synchronize values, while DTensor.from_local and ShardTensor.from_local do not.

Call this whenever domain_size > 1, regardless of the data-parallel wrapper: DDP broadcasts weights only over the data-parallel group at construction, and FSDP2 (fully_shard) does not synchronize initial weights on any axis. The ordering relative to DDP construction does not affect correctness (the two axes compose), but on the FSDP2 path this must run before fully_shard, while the parameters are still plain.

Parameters:

  • module (torch.nn.Module) – The module whose plain parameters/buffers are synchronized in place.

  • mesh (DeviceMesh) – A 1-D device mesh (typically the "domain" submesh of a larger mesh) whose process group the broadcast runs over.

  • src_mesh_rank (int, default=0) – The source position within the mesh to broadcast from.

  • verify (bool, default=False) – If True, collectively compare tensor metadata before the broadcast and per-tensor diagnostic checksums afterward. The checksum detects rank disagreement but is not a cryptographic integrity check.

Raises:

  • ValueError – If mesh is not 1-D. For a multi-dim mesh, pass the axis you want explicitly, e.g. mesh["domain"].

  • RuntimeError – If called after FSDP2 wrapping, if plain state is incompatible with the mesh, if all registered state is already distributed, or if verify=True finds differing metadata or values.

Synchronization Responsibilities#

In a training script, DDP synchronizes weights across its entire process group. With a 2D mesh, that process group does not include the domain axis of the mesh.

Call sync_module_over_mesh after creating the model and before converting its weights to distributed tensors. It copies the plain parameters and buffers from one process to the others in the mesh, so every model copy starts with the same values.

This function does not synchronize gradients. ShardTensor automatically sums the gradient contributions produced by domain-parallel operations.

sync_module_over_mesh skips DTensor and ShardTensor values. scatter_tensor and distribute_tensor create consistent distributed tensors from one source tensor. In contrast, from_local uses the value already present on each process. When using from_local, you are responsible for providing the correct local piece on every process.

The startup synchronization happens only once. It does not keep buffers in sync if they change during training. Synchronize changing buffers separately when your model requires it. Checkpoint loading is also separate. Use the distributed checkpoint utilities for distributed tensors, and synchronize plain model state if only one process loaded it.

Example: FSDP2 with Domain Parallelism#

Create a two-dimensional mesh. The "ddp" dimension holds different training samples and shards the model weights. The "domain" dimension splits one sample across multiple GPUs. The product of the two dimensions must equal the number of distributed processes.

For example, with eight GPUs, data_parallel_size=4 and domain_size=2 creates a \((4, 2)\) mesh:

import torch
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.fsdp import fully_shard

from physicsnemo.domain_parallel import sync_module_over_mesh

data_parallel_size = 4
domain_size = 2
mesh = init_device_mesh(
    "cuda",
    (data_parallel_size, domain_size),
    mesh_dim_names=("ddp", "domain"),
)
data_mesh = mesh["ddp"]
domain_mesh = mesh["domain"]

model = MyModel().cuda()

# All domain ranks must start with the same plain weights and buffers.
sync_module_over_mesh(model, domain_mesh)

# FSDP2 shards weights only across the data-parallel dimension.
fully_shard(model, mesh=data_mesh)

Call sync_module_over_mesh before fully_shard. If the model has parameters that must themselves be sharded over the domain, synchronize the plain model first, then convert those selected parameters with distribute_tensor before calling fully_shard.

Example: DDP with Domain Parallelism#

When the weights fit on each GPU, use DDP on the "ddp" dimension. The domain mesh is still used for ShardTensor inputs:

from torch.nn.parallel import DistributedDataParallel as DDP

model = MyModel().cuda()
sync_module_over_mesh(model, domain_mesh)
model = DDP(
    model,
    device_ids=[torch.cuda.current_device()],
    process_group=data_mesh.get_group(),
)

For detailed information on ShardTensor and domain parallelism, please refer to the Domain Parallelism tutorial.