💻 computer science

ShardTensor: Domain Parallelism for Scientific Machine Learning

This paper introduces ShardTensor, a novel domain parallelism framework that decouples input data spatial dimensionality from hardware constraints to enable scalable, high-fidelity training and inference of Scientific Machine Learning models on extreme-resolution datasets.

Original authors: Corey Adams, Peter Harrington, Akshay Subramaniam, Mohammad Shoaib Abbas, Jaideep Pathak, Mike Pritchard, Sanjay Choudhry

Published 2026-05-13

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Corey Adams, Peter Harrington, Akshay Subramaniam, Mohammad Shoaib Abbas, Jaideep Pathak, Mike Pritchard, Sanjay Choudhry

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Too Big to Fit" Box

Imagine you are a scientist trying to simulate a hurricane, a black hole, or a human brain. To get accurate results, you need to look at these things in extreme detail—like zooming in on a photo until you can see every single pixel.

In the world of Artificial Intelligence (AI), this is called high-resolution data.

The problem is that AI models run on special computers called GPUs (Graphics Processing Units). These GPUs have a limited amount of memory, like a small backpack.

The Issue: When scientists try to feed a massive, high-resolution image (like a 3D scan of a storm) into the AI, the data is so huge it simply doesn't fit in the backpack.
The Old Solution: Scientists used to have to "downsample" the data. This is like taking a 4K movie and shrinking it to a tiny, blurry 144p thumbnail just so it fits in the backpack. The AI can run, but the results are blurry and inaccurate.
The Other Old Solution: They could try to split the data across many computers (Data Parallelism), but this only works if you have many separate "chunks" of data (like 100 different storms). If you only have one giant storm to analyze, the old methods hit a wall. You can't split one storm into 100 pieces easily without breaking the math.

The Solution: ShardTensor (The "Team Backpack" Strategy)

The authors from NVIDIA introduced a new tool called ShardTensor.

Think of ShardTensor as a magical way to cut a single, giant pizza (your high-resolution data) into slices and hand those slices to a team of chefs (GPUs) standing in a circle.

How it works: Instead of trying to fit the whole pizza on one plate, the system slices the pizza spatially (by area, not by separate pizzas).
- Chef A holds the top-left slice.
- Chef B holds the top-right slice.
- Chef C holds the bottom-left slice.
The Magic: When the chefs need to mix ingredients (do math), they can pass the edges of their slices to their neighbors. If Chef A needs to know what's happening at the very edge of their slice, they ask Chef B. They work together to solve the whole puzzle without ever needing to put the whole pizza on one plate.

This is called Domain Parallelism. It allows scientists to process data that is larger than the memory of a single computer, even if they only have one single dataset to work with.

Why This Matters (The "Why Bother?" Section)

The paper explains that in scientific AI, the biggest memory hog isn't the "brain" of the AI (the model weights); it's the intermediate steps (activations).

Analogy: Imagine you are solving a giant math problem on a whiteboard. You don't just need space for the final answer; you need space for every single scratch calculation you make along the way.
The Result: With high-resolution data, these "scratch calculations" fill up the memory instantly. ShardTensor spreads these scratch calculations across many GPUs.
The Benefit:
1. Strong Scaling: If you have a massive dataset, adding more GPUs makes the job finish much faster (like adding more workers to a construction site).
2. Weak Scaling: If you have a dataset so huge it was previously impossible to run, you can now run it by adding more GPUs to share the load.

Real-World Examples from the Paper

The authors tested this on two specific scientific problems to prove it works:

StormScope (Weather Forecasting):
- The Challenge: Predicting individual thunderstorms requires looking at a map of the entire United States with very high detail (3km resolution).
- The Problem: A single computer's memory (80GB) couldn't hold the data for the whole US map at that detail. It was like trying to carry the entire US map in a single folder.
- The Fix: Using ShardTensor, they split the US map across 32 GPUs. The AI could now "see" the whole country in high definition and predict storms accurately without crashing.
Transolver (Aerodynamics):
- The Challenge: Simulating air flowing over a car to improve fuel efficiency.
- The Result: They were able to train the AI on a mesh (a 3D grid) containing over 1.2 million points representing the car's shape. This level of detail was previously impossible to train on a single machine.

What the Paper Doesn't Say (Limitations)

The authors are honest about the trade-offs:

Communication Overhead: Because the GPUs have to constantly talk to each other to share the edges of their data slices, there is a tiny bit of "talking time" lost.
Small Data: If the data is small, this talking time makes the system slower than just using one computer. ShardTensor is only for huge data.
Not Magic for Everything: It works best with specific types of math operations. Some older or very specific operations might not be supported yet.

Summary

ShardTensor is a new software tool that lets scientists slice up massive, high-resolution scientific data and spread it across multiple computers. This allows them to train AI models on data so detailed that it previously wouldn't fit in a single computer's memory, leading to more accurate weather forecasts, better car designs, and deeper scientific discoveries.

Technical Summary: ShardTensor – Domain Parallelism for Scientific Machine Learning

1. Problem Statement

Scientific Machine Learning (SciML) faces a critical bottleneck: the inability to train and perform inference on high-resolution scientific data due to GPU memory constraints. Unlike Large Language Models (LLMs), where memory is often dominated by parameter counts, SciML workloads are frequently constrained by intermediate activations.

The paper identifies a "computational curse of dimensionality" in scientific data. While doubling text length in an LLM roughly doubles input tokens, doubling the resolution of $N$ -dimensional scientific data increases the data size by $2^N$ . For 3D or higher-dimensional data (e.g., fluid dynamics, medical imaging, climate modeling), the memory required to store intermediate activations during the backward pass of training grows exponentially.

Existing solutions often fail to scale or degrade accuracy:

Reduced Precision: Offers only modest (2x) memory savings and can suffer from dynamic range issues in scientific surrogates.
Spatial Downsampling: Reduces memory but sacrifices the high-fidelity resolution required for scientific breakthroughs.
Model Reduction: Decreases accuracy.
Activation Checkpointing/CPU Offloading: Effective but introduces significant computational overhead (recomputation or host-device transfers).
Standard Parallelism: Data parallelism (splitting the batch) hits a hard limit at batch size 1. Model parallelism (splitting weights) is designed for parameters, not input data, and assumes static tensor shapes that do not hold for dynamic input data processing (e.g., convolutions changing global shapes).

Consequently, there is no generalized framework for parallelizing input data below a batch size of one per device, preventing SciML from reaching the native resolutions scientists require.

2. Methodology: ShardTensor

The authors introduce ShardTensor, a novel paradigm of domain parallelism that decouples the spatial dimensionality of input data from hardware constraints. It is an extension of PyTorch's DTensor (Distributed Tensor) abstraction, specifically designed to handle input data rather than model weights.

Core Design Principles

Domain Parallelism: Defined as parallelism across devices for input data, extending beyond the batch dimension. When the batch size reaches 1, ShardTensor further subdivides the spatial dimensions (e.g., height, width, depth) across multiple GPUs.
Dynamic Shape Handling: Unlike DTensor, which assumes static distribution shapes suitable for model weights, ShardTensor supports dynamic shapes. It introduces a fourth component of metadata: sharding shapes, which describe the local chunk shape of each tensor along the sharded mesh axis. This allows for arbitrary chunking of unstructured or non-uniform data.
Imperative Dispatch: The framework is imperative rather than static. It utilizes PyTorch's __torch_dispatch__ and __torch_function__ mechanisms to route operations dynamically. This allows the system to handle layer-by-layer domain parallelism without requiring pre-compiled computational graphs or bespoke patches to every layer.
Extensibility: Users can extend the system via high-level functional interfaces (@custom_op) or low-level dispatch logic (aten operations). This ensures compatibility with evolving model architectures and custom kernels.

Implementation Details

ShardTensor is built as an extension to PyTorch's DTensor.

Dispatch Architecture: It intercepts operations at the Python level. If a specific operation has a domain-parallel handler, it is executed; otherwise, it falls back to DTensor semantics or standard PyTorch behavior.
Communication: The framework manages collective operations (e.g., all-reduce, all-gather) and "halo" exchanges (fetching adjacent pixels from neighboring devices) required for operations like convolution or normalization to maintain numerical consistency.
Integration: It is designed to be composable with existing parallelism strategies (Data Parallelism and Model Parallelism), allowing for 2D or higher-dimensional parallelization strategies (e.g., splitting the batch on one axis and spatial dimensions on another).

3. Key Contributions

ShardTensor Framework: A generalized, open-source framework (available via NVIDIA PhysicsNeMo) enabling domain parallelism for input data below batch size one.
Novel Abstraction: The introduction of "sharding shapes" to handle dynamic input tensor distributions, distinguishing it from weight-focused parallelism tools like DTensor.
Non-Invasive Programming Model: A design that allows users to wrap model inputs with minimal code changes, avoiding the need to rewrite layers or models for parallelism.
Composability: Demonstrated ability to combine domain parallelism with data and model parallelism on orthogonal mesh axes.

4. Results and Benchmarks

The paper validates ShardTensor through synthetic benchmarks and real-world industrial applications on NVIDIA Blackwell (GB200) and H100 GPUs.

Ring Attention: On very small sequence sizes, parallelization overhead dominates. However, at very large sequence sizes, the system achieves near-linear strong scaling with GPU count in both training and inference.
Vision Transformer (ViT):
- Strong Scaling: At high resolutions (e.g., $4096^2$ ), inference is 15x faster on 16 GPUs compared to a single GPU. Training on $2048^2$ is 5x faster with 8 GPUs.
- Memory Efficiency: Memory usage scales quadratically for 2D data and cubically for 3D data, confirming that intermediate activations dominate. Strong scaling with ShardTensor reduces per-device memory proportionally, enabling the training of 3D models with over 1 billion input points on standard hardware.
Transolver (Automotive Aerodynamics):
- Trained on the DrivaerML dataset with up to 1.2 million points per domain.
- The system maintained numerical stability and competitive accuracy (L2 error for pressure and velocity) compared to single-GPU baselines, even as the domain size increased by factors of 2, 4, and 8.
- The entire preprocessing pipeline was also parallelized, ensuring end-to-end scalability.
StormScope (Weather Forecasting):
- A continental-scale model (CONUS) at 3 km resolution ( $1024 \times 1792$ grid).
- The peak memory usage (114 GB) exceeded the capacity of a single H100 GPU (80 GB).
- ShardTensor enabled stable training by splitting activations across 2 GPUs within 16 data-parallel groups. The model converged stably, tracking the loss trajectory of a lower-resolution (6 km) single-GPU baseline, enabling the resolution of individual convective storms previously inaccessible without sacrificing domain extent.

5. Significance and Claims

The paper claims that ShardTensor addresses a fundamental bottleneck in SciML: the inability to train and infer on data at the native resolution required by scientists.

Unblocking Resolution-Limited Workloads: It converts previously impossible computations (due to memory ceilings) into tractable ones, enabling high-fidelity training in fields like volumetric medical imaging, computational fluid dynamics, and climate modeling.
Composability: It removes the trade-off between data volume, model size, and resolution. Scientists can now scale along multiple parallelism axes simultaneously.
Accessibility: By providing a non-invasive programming model, it lowers the barrier to entry for practitioners who are not distributed systems experts, contrasting with previous bespoke, architecture-specific parallelization efforts.

Limitations Acknowledged

The authors modestly note that the imperative, layer-by-layer dispatch model introduces Python-level overhead, making it less efficient for small operations or low-resolution data. Performance is also dependent on interconnect bandwidth. Furthermore, not all PyTorch operations currently have domain-parallel dispatch paths, and full integration with torch.compile for static graph optimizations is ongoing.

In conclusion, ShardTensor represents a step toward making high domain parallelism as commonplace in scientific machine learning as data parallelism is today, facilitating the era of scientific foundation models.