# ShardTensor: Domain Parallelism for Scientific Machine Learning

arXiv: https://arxiv.org/abs/2605.11111 · HTML: https://mlsanity.com/paper/2605.11111/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 32 of its 72 references that are in the map.
- First posted on arXiv 11 May 2026.

## Builds on (most cited first, 30 of 32)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [4D Spatio-Temporal ConvNets: Minkowski Convolutional Neural Networks](https://mlsanity.com/paper/1904.08755/) (2019)
- [Large Batch Optimization for Deep Learning: Training BERT in 76 minutes](https://mlsanity.com/paper/1904.00962/) (2019)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [Large Batch Training of Convolutional Networks](https://mlsanity.com/paper/1708.03888/) (2017)
- [Machine learning accelerated computational fluid dynamics](https://mlsanity.com/paper/2102.01010/) (2021)
- [Ring Attention with Blockwise Transformers for Near-Infinite Context](https://mlsanity.com/paper/2310.01889/) (2023)
- [DeepSpeed Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale](https://mlsanity.com/paper/2207.00032/) (2022)
- [Machine Learning for Fluid Mechanics](https://mlsanity.com/paper/1905.11075/) (2019)
- [Transolver: A Fast Transformer Solver for PDEs on General Geometries](https://mlsanity.com/paper/2402.02366/) (2024)
- [Horovod: fast and easy distributed deep learning in TensorFlow](https://mlsanity.com/paper/1802.05799/) (2018)
- [Neighborhood Attention Transformer](https://mlsanity.com/paper/2204.07143/) (2022)
- [A Study of BFLOAT16 for Deep Learning Training](https://mlsanity.com/paper/1905.12322/) (2019)
- [Submanifold Sparse Convolutional Networks](https://mlsanity.com/paper/1706.01307/) (2017)
- [FP8 Formats for Deep Learning](https://mlsanity.com/paper/2209.05433/) (2022)
- [Transolver++: An Accurate Neural Solver for PDEs on Million-Scale Geometries](https://mlsanity.com/paper/2502.02414/) (2025)
- [FourCastNet 3: A geometric approach to probabilistic machine-learning weather forecasting at scale](https://mlsanity.com/paper/2507.12144/) (2025)
- [A Probabilistic Graphical Model Foundation for Enabling Predictive Digital Twins at Scale](https://mlsanity.com/paper/2012.05841/) (2020)
- [AB-UPT: Scaling Neural CFD Surrogates for High-Fidelity Automotive Aerodynamics Simulations via Anchored-Branched Universal Physics Transformers](https://mlsanity.com/paper/2502.09692/) (2025)
- [DrivAerML: High-Fidelity Computational Fluid Dynamics Dataset for Road-Car External Aerodynamics](https://mlsanity.com/paper/2408.11969/) (2024)
- [DoMINO: A Decomposable Multi-scale Iterative Neural Operator for Modeling Large Scale Engineering Simulations](https://mlsanity.com/paper/2501.13350/) (2025)
- [GeoTransolver: Learning Physics on Irregular Domains Using Multi-scale Geometry Aware Physics Attention Transformer](https://mlsanity.com/paper/2512.20399/) (2025)

## Built on by (most cited first, 1 of 1)

- [TERRA: A Hierarchical Parallel Training and Memory Orchestration Framework for High-Resolution AI-based Earth Modeling](https://mlsanity.com/paper/2608.15211/) (2026)

## Abstract (arXiv)

Scientific Machine Learning (SciML) faces unique challenges for extreme-resolution data, with mitigations that often fail to scale or degrade the accuracy of trained models. While some specialized methods have achieved remarkable results in training models or performing inference on massive spatial datasets with bespoke techniques, there is no generalized framework for parallelization over input data below batch size one per device. In this work we introduce ShardTensor: a novel paradigm of domain parallelism that enables flexible scaling of input data to arbitrary sizes. By decoupling the spatial dimensionality of input data from hardware constraints, ShardTensor enables scientific machine learning workloads to reach new levels of high fidelity training and inference. We demonstrate both strong and weak scaling of workloads during training and inference, showing improved latency with strong scaling and demonstrating the capacity to process higher data sizes with weak scaling. Additionally, we demonstrate multiple dimensions of parallelization, removing barriers to SciML on extreme-scale inputs.