Agent skill · software engineering · 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.
Why this skill is useful
Adds executable scripts and domain-specific knowledge for integrating domain parallelism in PhysicsNeMo, which the AI wouldn't reliably generate on its own.
What it needs
Requires torch installed locally. About 7k tokens when loaded. Last updated 2026-08-07. 2,816 stars on the source repository.
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