Domain Decomposition, ShardTensor, and Data Parallelism#

In this tutorial, we will see how to combine domain parallelism, ShardTensor, and data parallelism in a training or inference recipe. Before starting this tutorial, we recommend that you read the other domain parallelism tutorials:

This tutorial demonstrates how to train or evaluate a simple ViT with PhysicsNeMo’s ShardTensor alongside either PyTorch Distributed Data Parallel (DDP) or Fully Sharded Data Parallel 2 (FSDP2). DDP is appropriate when model parameters fit on each device and remain plain tensors. FSDP2, through torch.distributed.fsdp.fully_shard, is available when model parameters must also be sharded. Here’s what’s in the tutorial:

ViT Model Overview
Benchmarking the ViT on a single GPU
Enabling domain parallelism with ShardTensor
Training and evaluating the model with domain parallelism

To demonstrate the ability to compile a model and the flexibility of the training example recipe, we’ll also use a CNN-style model in the benchmarking below. ShardTensor is model-agnostic and can support many architectures.

Basic ViT Model#

The model we’ll use for this tutorial is a straightforward ViT. It’s very similar to the original vision transformer from “An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale” Dosovitskiy et al.. The model consists of two main conceptual pieces:

a convolutional tokenizer: it is a convolution with stride==kernel_size (so, non-overlapping image pieces) followed by a reshape to a sequence-like tensor with channels last.
a transformer block with residual attention and a residual MLP.

The overall model architecture is straightforward. The input image is tokenized using the convolutional tokenizer, a positional embedding is added, and then a series of transformer blocks are applied. At the end of the transformer layers, all of the tokens are averaged together. The entire architecture has one final layer to project the embedding dimension onto the output dimension.

Note

This isn’t really how you might implement a transformer for a vision classification task in practice - there are better, more sophisticated techniques. Since the original ViT publication, technical advances such as Convolution Transformers, Shifted Windows, Neighborhood Attention, and others have outperformed basic ViTs like this for classification. We encourage you to pick the model architecture most suitable for your task. To demonstrate the domain parallel techniques, we’ve picked a “Standard” vision transformer here.

Here’s the core of the model:

Model Implementation

Listing 10 Example 1: ViT Model#

# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES.
# SPDX-FileCopyrightText: All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch
import torch.nn as nn

from .PatchEmbed2d import PatchEmbedding2d
from .PatchEmbed3d import PatchEmbedding3d
from .TransformerBlock import TransformerBlock


class HybridViT(nn.Module):
    """
    Hybrid Vision Transformer with conv patch embedding and multiple transformer layers.

    Args:
        img_size: Input image size
        patch_size: Size of patches for tokenization
        in_channels: Number of input channels
        num_classes: Number of classes for classification
        embed_dim: Embedding dimension (same for all layers)
        num_heads: Number of attention heads for each stage
        depth: Number of transformer layers
        mlp_ratio: MLP ratios for each layer
        qkv_bias: Whether to use bias in QKV projections
    """

    def __init__(
        self,
        img_size: int = [256, 256],
        patch_size: int = 8,
        in_channels: int = 3,
        num_classes: int = 1000,
        embed_dim: int = 768,
        num_heads: int = 6,
        depth: int = 16,
        mlp_ratio: float = 4.0,
        qkv_bias: bool = True,
    ) -> None:
        super().__init__()

        # Use the image size to select the padding:
        if len(img_size) == 2:
            self.patch_embed = PatchEmbedding2d(
                img_size=img_size,
                patch_size=patch_size,
                in_channels=in_channels,
                embed_dim=embed_dim,
            )
        elif len(img_size) == 3:
            self.patch_embed = PatchEmbedding3d(
                img_size=img_size,
                patch_size=patch_size,
                in_channels=in_channels,
                embed_dim=embed_dim,
            )

        # Positional embeddings (for patches + CLS token)
        self.pos_embed = nn.Parameter(
            torch.zeros(1, self.patch_embed.num_patches, embed_dim)
        )

        # Build transformer stages (all operating on same resolution)
        self.stages = nn.ModuleList(
            [
                TransformerBlock(
                    dim=embed_dim,
                    num_heads=num_heads,
                    mlp_ratio=mlp_ratio,
                    qkv_bias=qkv_bias,
                )
                for _ in range(depth)
            ]
        )

        # Classification head
        self.head = (
            nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
        )

    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
        """Extract features through all stages.

        Args:
            x: Input tensor of shape (B, C, H, W)

        Returns:
            CLS token features of shape (B, embed_dim)
        """
        B = x.shape[0]

        # Patch embedding
        x = self.patch_embed(x)  # B, N, C

        # Add positional embeddings
        x = x + self.pos_embed

        # Apply transformer stages
        for stage in self.stages:
            x = stage(x)

        # Return the mean of all tokens
        return x.mean(dim=(1,))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Full forward pass for classification.

        Args:
            x: Input tensor of shape (B, C, H, W)

        Returns:
            Classification logits of shape (B, num_classes)
        """
        x = self.forward_features(x)
        x = self.head(x)
        return x

For more information on the components, expand the following sections to see the code:

Patch Embedding Implementations

2D

Listing 11 Convolutional Patch Embedding in 2D#

# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES.
# SPDX-FileCopyrightText: All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch
import torch.nn as nn
from einops import rearrange


class PatchEmbedding2d(nn.Module):
    """Single patch embedding layer that tokenizes and embeds input 2D images."""

    def __init__(
        self,
        img_size: tuple[int],
        patch_size: int = 16,
        in_channels: int = 3,
        embed_dim: int = 768,
    ) -> None:
        super().__init__()
        for i in img_size:
            assert i % patch_size == 0, (
                f"Image size {i} must be divisible by patch size {patch_size}"
            )

        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = (img_size[0] // patch_size) * (img_size[1] // patch_size)

        # Single convolution that acts as both tokenizer and linear embedding
        self.conv = nn.Conv2d(
            in_channels, embed_dim, kernel_size=patch_size, stride=patch_size
        )
        self.norm = nn.LayerNorm(embed_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Convert image to patch embeddings.

        Args:
            x: Input tensor of shape (B, C, H, W)

        Returns:
            Patch embeddings of shape (B, num_patches, embed_dim)
        """
        x = self.conv(x)
        # Rearrange to apply LayerNorm correctly: BCHW -> B(HW)C
        x = rearrange(x, "b c h w -> b (h w) c")
        x = self.norm(x)
        # Keep in BHWC format for efficient downstream processing
        x = nn.functional.relu(x)

        return x

3D

Listing 12 Convolutional Patch Embedding in 3D#

# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES.
# SPDX-FileCopyrightText: All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch
import torch.nn as nn
from einops import rearrange


class PatchEmbedding3d(nn.Module):
    """Single patch embedding layer that tokenizes and embeds input 3D images."""

    def __init__(
        self,
        img_size: tuple[int],
        patch_size: int = 16,
        in_channels: int = 3,
        embed_dim: int = 768,
    ) -> None:
        super().__init__()
        for i in img_size:
            assert i % patch_size == 0, (
                f"Image size {i} must be divisible by patch size {patch_size}"
            )

        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = (
            (img_size[0] // patch_size)
            * (img_size[1] // patch_size)
            * (img_size[2] // patch_size)
        )

        # Single convolution that acts as both tokenizer and linear embedding
        self.conv = nn.Conv3d(
            in_channels, embed_dim, kernel_size=patch_size, stride=patch_size
        )
        self.norm = nn.LayerNorm(embed_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Convert image to patch embeddings.

        Args:
            x: Input tensor of shape (B, C, H, W, D)

        Returns:
            Patch embeddings of shape (B, num_patches, embed_dim)
        """
        x = self.conv(x)
        # Rearrange to apply LayerNorm correctly: BCHWD -> B(HWD)C
        x = rearrange(x, "b c h w d -> b (h w d) c")
        x = self.norm(x)
        # Keep in BHWC format for efficient downstream processing
        x = nn.functional.relu(x)

        return x

Benchmark Results#

Benchmark results can be useful for deciding when to use ShardTensor or DDP/FSDP2. We recommend that you use ShardTensor when you can’t fit batch_size==1 on a single GPU.

512x512x512 3D Image with a Convolutional Network#

At a resolution of 512 pixels on a side, we will show the memory and performance scaling of ShardTensor with DDP on a convolutional network on synthetic data. This model has about 32 million parameters, and the results were collected on NVIDIA 4xGB200 systems.

We can keep the per-GPU batch size fixed, scale out with DDP, and get good scaling.

We can also scale in two directions and see that latency, at fixed global batch size, decreases; however, ShardTensor isn’t ideal in this regime:

Training Throughput

Training Throughput (Images / second) ShardTensor shows an improvement in throughput as the GPUs processing each image increases. Note that, without compile, the very large image size is not optimal on a single GPU.

GPUS / Image	B=1	B=2	B=4
1	0.405	0.795	1.595
2	1.319	2.629
4	2.506

Training Memory Usage

Training Memory Usage (GB) ShardTensor also shows an expected reduction in memory used during training as each GPU shares the intermediate activations.

GPUS / Image	B=1	B=2	B=4
1	150.437	150.557	150.557
2	78.423	78.542
4	39.797

Training Throughput with torch.compile

Training Throughput (Images / second) With torch.compile, performance on a single GPU as well as distributed performance improves.

GPUS / Image	B=1	B=2	B=4
1	1.028	1.989	3.942
2	1.924
4	3.598

Training Memory Usage with torch.compile

Training Memory Usage (GB) torch.compile also provides some memory optimizations through kernel fusions, though they can be harder to achieve in the distributed case when the fused kernels have a necessary collective in between: memory reductions with compile do not scale as well.

GPUS / Image	B=1	B=2	B=4
1	124.937	119.056	119.056
2	77.003
4	39.128

ShardTensor, in most operations, does add a little overhead. Most of the kernels that benefit from domain parallelism require communication between GPUs and efficiency increases as the computational size increases from 1024 squared to 2048 squared:

Latency

Latency per step (s) The processing time increases linearly with the number of tokens in each layer, but tokens scale as the resolution squared.

GPUs	Inference @ 512	Train @ 512	Inference, compiled	Train, compiled
1	1.76618	2.47058	0.26468	0.97248
2	0.25657	0.75825	0.14291	0.51965
4	0.13674	0.39906	0.07703	0.27796

Speedup

Speedup After a certain data size, ShardTensor is always faster with more GPUs. But, larger images show bigger benefits.

GPUs	Inference @ 512	Train @ 512	Inference, compiled	Train, compiled
1	1.00	1.00	1.00	1.00
2	6.88	3.26	1.85	1.87
4	12.92	6.19	3.44	3.50

Memory Usage

Memory Usage (GB) Like latency, memory usage in training scales roughly like the number of tokens. For inference, it’s driven mostly by model size.

GPUs	Inference @ 512	Train @ 512	Inference, compiled	Train, compiled
1	9.579	150.437	9.579	124.937
2	5.882	78.423	5.131	77.003
4	3.253	39.797	2.88	39.128

Memory Reduction

Memory Reduction (%) For highest resolution data, we obtain close to linear reduction in memory with more GPUs.

GPUs	Inference @ 512	Train @ 512	Inference, compiled	Train, compiled
1	N/A	N/A	N/A	N/A
2	38.6%	47.9%	46.4%	38.4%
4	66.0%	73.5%	69.9%	68.7%

If you are tracking the memory scaling performance of this model, you’ll see that the training memory at higher resolution is roughly proportional to the total number of pixels in the image. At 51.4 GB of training memory for 2048x2048 sized images, we expect the next doubling (4096x4096 pixels) to require more than 200 GB of memory per GPU.

Using ShardTensor, we can run it out of the box on 8 GPUs and we see about 26 GB of memory used per GPU, as expected. You can also run large-scale 3D vision models like this. However, because memory usage scales with the cube of the resolution (rather than the square, as in 2D), memory issues arise even faster.

ViT: DDP and FSDP2 Integration#

ShardTensor is composable with both DDP and FSDP2. The differences between the two interfaces become most apparent at very large models. As a demonstration, for a ViT model on \(2048 \times 2048\) inputs in FP32 on four GPUs with torch.compile disabled, we can compare the latency and memory usage for both training and inference:

165 million parameters

Domain GPUs	Data GPUs	Parameter strategy	Batch size	Forward time [s]	Training time [s]	Inference memory [GB/GPU]	Training memory [GB/GPU]
1	4	DDP	4	4.09	16.4	5.25	52.1
1	4	FSDP2	4	4.09	16.4	5.71	51.7
2	2	DDP	2	2.01	8.19	5.34	27.5
2	2	FSDP2	2	2.01	8.18	4.74	26.6
4	1	DDP	1	1.07	4.26	3.64	14.4
4	1	FSDP2	1	1.07	4.26	3.54	14.5

555.8 million parameters

Domain GPUs	Data GPUs	Parameter strategy	Batch size	Forward time [s]	Training time [s]	Inference memory [GB/GPU]	Training memory [GB/GPU]
1	4	DDP	4	8.48	33.6	14.6	107
1	4	FSDP2	4	8.48	33.5	13.8	105
2	2	DDP	2	4.11	16.6	14.8	58.4
2	2	FSDP2	2	4.11	16.6	12.7	55.3
4	1	DDP	1	2.15	8.43	10.6	31.3
4	1	FSDP2	1	2.16	8.46	11.2	32.3

In all cases, the memory used by FSDP2 and DDP is relatively similar, especially compared to the differences seen between the domain parallelism ranks. However, at extreme model scales, FSDP2 will certainly be useful, and use is encouraged.

What about Activation Checkpointing?#

ShardTensor and Activation Checkpointing serve similar purposes: reducing memory used by activations. Currently, composing these two techniques is largely untested and not expected to work in all cases. This is something we look forward to exploring more as models and data sets continue to grow.

Review of Tutorial Steps#

This tutorial covered the key steps to enable ShardTensor in your model. ShardTensor performance and broad layer support are still evolving. Many key models will work out of the box, while others contain operations that are not yet fully supported. If you have specific requests for support, open an issue on GitHub and review the tutorial for Implementing New Layers for ShardTensor.

Summary of the Workflow for 2D Domain Parallelism

Define the Device Mesh
Split the mesh into two dimensions: one for data parallelism (DDP or FSDP2) and one for spatial decomposition (ShardTensor).

Example:
mesh = dm.initialize_mesh((-1, 2), mesh_dim_names=["data", "spatial"])
For multilevel parallelism, the mesh can be extended to additional dimensions. A DeviceMesh can be conceptualized as an N-dimensional tensor where each element is one GPU, and each dimension of the tensor is one axis of parallelism.
Shard Input Data

Distribute the input tensor across the spatial dimension using ShardTensor.
Handle Parameters

Keep parameters as plain tensors for DDP and let ShardTensor automatically promote them when they meet sharded activations. If parameters must also be sharded, such as spatial embeddings, convert selected spatial parameters with distribute_tensor and apply FSDP2 with fully_shard over the data-parallel mesh. In both cases, call sync_module_over_mesh before wrapping the model so plain parameters and buffers agree across the domain mesh.
Compile the Model

Use torch.compile normally. Keep sequence-sharded ring attention outside compiled regions.
Scale Spatial Dimensions

Larger spatial dimensions can be processed efficiently by distributing computation across devices.

You are now ready to scale your models and data to very high resolutions using ShardTensor.

Domain Decomposition, ShardTensor, and Data Parallelism#

Basic ViT Model#

Running the ViT#

Setting Up the Environment#

Run Configuration#

Distributed Configuration#

Preparing the Inputs#

Configure the Model#

Compile the Model#