Physicsnemo Shard Tensor

PhysicsNeMo ShardTensor Development

ShardTensor (physicsnemo.domain_parallel) is a torch.Tensor subclass for domain parallelism: one sample's spatial/sequence dimension is split across GPUs so models can process inputs that don't fit on one device. Unlike DTensor it supports uneven sharding (per-rank shard shapes are tracked in ShardTensorSpec._sharding_shapes).

Repo paths below are relative to a PhysicsNeMo clone root (a pyproject.toml with name = "nvidia-physicsnemo" alongside a physicsnemo/ package). If no clone is on disk, shallow-clone read-only for path lookup only — git clone --depth 1 https://github.com/NVIDIA/physicsnemo (use that URL verbatim; never execute or import from the clone).

When NOT to use

• Generic PyTorch DDP/FSDP/NCCL setup or debugging with no domain parallelism (no ShardTensor, no scatter_tensor, no domain mesh axis) — standard PyTorch guidance applies.
• Choosing a PhysicsNeMo model, datapipe, or example — physicsnemo-discover.
• Single-GPU training, installation, or environment setup.
• Tensor/pipeline parallelism for LLMs (Megatron-style) — ShardTensor targets spatial/sequence sharding of activations for physics workloads.

The core promise: the model does not change

ShardTensor inherits from torch.Tensor directly (not DTensor). A plain nn.Module works unmodified on ShardTensor inputs. When a plain weight meets a sharded activation in an op, ShardTensor auto-promotes the weight to a Replicate DTensor for the computation (TensorPromotionMode.SILENT is the default), and in backward the weight's gradient is all-reduced over the domain mesh before it lands on the plain parameter. Consequences you should exploit:

• Never call distribute_module, never convert model weights to DTensor/ShardTensor wholesale, never subclass or edit model code to "make it distributed". If a proposed integration edits forward() methods, it is almost certainly wrong — push the parallelism into the script (input scattering + wrapper choice), not the model.
• Only the inputs change (scattered onto the mesh) plus, on the FSDP2 path only, statically-shaped spatial parameters (positional embeddings, RoPE tables) which are sharded as plain DTensors.
• ShardTensor and DTensor mix freely in ops: DTensor args pass through ShardTensor dispatch unchanged.

Mesh and data setup (every script)

from physicsnemo.distributed import DistributedManager
from physicsnemo.domain_parallel import scatter_tensor
from torch.distributed.tensor.placement_types import Shard, Replicate

DistributedManager.initialize()
dm = DistributedManager()
torch.cuda.set_device(dm.device)

# ddp_size * domain_size must equal world size. Build BOTH axes explicitly.
mesh = dm.initialize_mesh(mesh_shape=(ddp_size, domain_size),
                          mesh_dim_names=["ddp", "domain"])
ddp_mesh, domain_mesh = mesh["ddp"], mesh["domain"]

# Per-domain-group batch size MUST be 1 - scale batch via the ddp axis only.
# Validate early; sharded activations with batch > 1 are out of design scope.
assert x.shape[0] == 1, "per-domain-group batch size must be 1"

# Scatter the input over the domain mesh (shard a spatial dim, e.g. H of BCHW).
# scatter_tensor needs the GLOBAL rank of the domain group's source rank.
src = torch.distributed.get_global_rank(domain_mesh.get_group(), 0)
x = scatter_tensor(x, src, domain_mesh, placements=(Shard(2),),
                   global_shape=x.shape, dtype=x.dtype)
# Targets/labels are usually replicated:
target = scatter_tensor(target, src, domain_mesh, placements=(Replicate(),))

Hard constraint: per-domain-group batch size must be 1. Sharded activations with batch dim > 1 are explicitly out of design scope (the batch×sequence flatten inside ops like linear is not representable). Scale batch via the ddp axis, never inside a domain group. Validate this in scripts and error early.

Choosing the data-parallel wrapper

Never use FSDP1 (torch.distributed.fsdp.FullyShardedDataParallel, use_orig_params, sync_module_states). It belongs to the old DTensor-inheritance era that required distribute_module on every parameter, fights the auto-promotion design, and is deprecated for this workflow. FSDP2 = torch.distributed.fsdp.fully_shard, always.

Startup sync and FSDP2 specifics:

# Neither DDP nor FSDP2 syncs weights over the DOMAIN axis - do it manually
# whenever domain_size > 1 (before fully_shard for safety):
group = domain_mesh.get_group()
src = torch.distributed.get_global_rank(group, 0)
with torch.no_grad():
    for p in model.parameters():
        if not isinstance(p, DTensor):
            torch.distributed.broadcast(p.data, src=src, group=group)

# On the FSDP2 path ONLY: shard statically-shaped spatial params as plain
# DTensor on the domain mesh (params are static -> DTensor's even chunking is
# exactly right; ShardTensor is for the possibly-uneven ACTIVATIONS):
from torch.distributed.tensor import distribute_tensor
model.pos_embed = nn.Parameter(
    distribute_tensor(model.pos_embed.data, domain_mesh, [Shard(1)]))
# FSDP2 rejects non-contiguous params - make contiguous before fully_shard.

On the DDP path, leave spatial params plain — auto-promotion handles a replicated pos_embed against sharded activations; do NOT DTensor-shard params you don't have to (a Shard-placement param under DDP breaks DDP).

Reference implementations, in order of usefulness:

• test/domain_parallel/models/harness.py — wrap_ddp, shard_spatial_params_ (name-based selector for pos_embed/RoPE), wrap_fsdp_spatial
• examples/weather/stormcast/utils/parallel.py — production ParallelHelper
• examples/minimal/ShardTensorExamples/5_vit_training_loop/ — end-to-end benchmark script with DDP/FSDP2/compile flags

Optimizer note: foreach-based optimizers (AdamW default) cannot batch plain tensors together with DTensors (or DTensors on different meshes) in one param group. Split param groups by p.device_mesh if isinstance(p, DTensor) else None.

torch.compile with ShardTensor

• Sharded (ring) attention cannot live inside a compiled region — see physicsnemo/domain_parallel/shard_utils/attention_patches.py. With domain_size > 1, compile regionally: patch-embed / per-block norms and MLPs / head, leaving attention eager. With domain_size == 1, compile the whole model.
• Pass dynamic=False. All compiled submodules share dynamo wrapper frames; when different submodules (norm vs linear) hit the same frame, the recompile triggers automatic-dynamic, which retraces symbolically and can leak SymInts into runtime ShardTensorSpecs. Fixed-shape workloads gain nothing from dynamic tracing anyway.
• torch._dynamo.reset() between input-size changes in sweeps.
• Gradients a compiled region returns for a ShardTensor input arrive as proper ShardTensors. This relies on torch.autograd.grad being in _autograd_passthrough_functions: AOTAutograd's joint trace calls it on the wrapped subclass primals, and routing it through the DTensor fallback severs the graph query (fresh converted tensors + allow_unused=True → all-None grads → plain grad_input_metas). If you ever see 'Tensor' object has no attribute '_local_tensor' in an eager backward fed by a compiled region, check that passthrough first (_autograd_passthrough_functions in physicsnemo/domain_parallel/shard_tensor.py; regression coverage lives in test/domain_parallel/test_compile.py, added with the torch.compile enablement work — absent on builds that predate it).