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).

Debugging pitfalls (each of these cost real time — check them first)

1. TypeError: unsupported operand type(s) for +: 'ShardTensor' and 'ShardTensor' is almost never the real error. Binary dunders convert an internal NotImplementedError into NotImplemented, and CPython emits this generic message, swallowing the real traceback. Temporarily replace x + y with torch.add(x, y) to surface the true exception.
2. In-place x.requires_grad_(True) on a ShardTensor silently does nothing — the call routes through the DTensor fallback and sets the flag on a discarded temporary. Use scatter_tensor(..., requires_grad=True) or thread gradients through parameters.
3. torch.autograd.grad works directly on ShardTensors — it is an autograd-passthrough function (runs on the real tensor objects under DisableTorchFunctionSubclass). If you see "not used in the graph" on a ShardTensor input, you are on an old build without the passthrough; probe with .backward() + tensor.register_hook(...) there instead. Beware that monkeypatching torch.autograd.grad (e.g. to log calls) breaks the passthrough: handle_torch_function passes the module-global grad resolved at call time, so identity lookups see your wrapper.
4. Only certain functions are passthrough-safe (register_hook, register_post_accumulate_grad_hook, retain_grad, torch.autograd.grad — see _autograd_passthrough_functions in shard_tensor.py). Any other identity-sensitive method may act on a converted temporary.
5. Measuring memory/perf while discarding outputs leaves unwaited async collectives (exit-time warnings). Resolve with to_local()/AsyncCollectiveTensor.wait() on discarded results.
6. CommDebugMode (torch.distributed.tensor.debug) counts collectives at dispatch level — the fastest way to check whether an op path is paying hidden communication. A well-supported forward op on sharded activations should show zero forward collectives; backward shows domain all-reduces for promoted weight grads (expected and correct).

Enabling new layers / ops

Read references/new-op-patterns.md before writing any patch. Summary of the decision process:

1. Try the model unmodified first. The generic fallback (convert to DTensor, run, convert back) covers most ops correctly. Only write a patch when you observe: a MissingShardPatch/UndeterminedShardingError, wrong numerics vs a single-GPU run, or unacceptable communication (redistribution to Replicate) in CommDebugMode.
2. Patches are registered from user code at import time — no physicsnemo fork needed: ShardTensor.register_function_handler(torch.nn.functional.foo, wrapper) (Python/__torch_function__ level), ShardTensor.register_dispatch_handler(aten.foo.default, fn) (__torch_dispatch__ level), and ShardTensor.register_named_function_handler("lib.op.default", wrapper) for torch.library.custom_ops.
3. Use the existing patches in physicsnemo/domain_parallel/shard_utils/ as templates: pooling_patches.py (config gating + MissingShardPatch), conv_patches.py + halo.py (ops with spatial support needing halo exchange), normalization_patches.py (explicit autograd.Function with custom backward), view_ops.py (dual-level registration; shape-only ops).

Testing new layers

Read references/testing.md. The one-line summary: scatter a full input, run the module distributed and single-GPU, and compare outputs and gradients with numerical_shard_tensor_check(mesh, module, [sharded_x], {}, check_grads=True) under the multigpu_static marker, launched as

torchrun --nproc-per-node 4 -m pytest test/... --multigpu-static -m multigpu_static

A forward-only test proves almost nothing — the weight gradient is where sharding bugs live (it is Partial over the domain mesh and must be reduced). Always check_grads=True, always disable TF32 for the comparison.

Related resources

• references/integration-checklist.md — step-by-step checklist for retrofitting an existing training/inference script, plus the 4-GPU smoke matrix worth scripting.
• references/new-op-patterns.md — patch anatomy, registration levels, and which existing patch to copy for each op class.
• references/testing.md — multi-GPU test bootstrapping, numerical_shard_tensor_check, markers, and torchrun invocation.
• physicsnemo-discover — for choosing models, datapipes, and examples.