Agent skill · NVIDIA
physicsnemo-shard-tensor
Official NVIDIA-authored guidance for PhysicsNeMo ShardTensor domain parallelism — integrate domain parallelism into training/inference scripts (new or existing) with DDP or FSDP2, write and register shard patches to enable new layers/ops, and bootstrap multi-GPU correctness tests. Use when working with ShardTensor, scatter_tensor, domain parallelism, sequence/spatial sharding, ring attention, DeviceMesh + DDP/FSDP2 hybrid parallelism, or physicsnemo.domain_parallel. Do NOT use for generic PyTorch DDP/FSDP setup without domain parallelism, picking a PhysicsNeMo model or example (use physicsnemo-discover), or non-distributed training questions.
What it needs
About 7k tokens when loaded.
What this skill does
PhysicsNeMo ShardTensor Development ShardTensor (physicsnemo.domainparallel) 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.shardingshapes). 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 scattertensor, 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 distributemodule, 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. …
How to use it
Reference it in AdaL, Claude Code, Cursor or any coding agent — nothing to install:
@skills NVIDIA/physicsnemo-shard-tensor