💻 computer science

ShardTensor: Domain Parallelism for Scientific Machine Learning

本文介绍了 ShardTensor，这是一种新颖的域并行框架，它将输入数据的空间维度与硬件约束解耦，从而支持在超高分辨率数据集上对科学机器学习模型进行可扩展、高保真的训练与推理。

原作者： Corey Adams, Peter Harrington, Akshay Subramaniam, Mohammad Shoaib Abbas, Jaideep Pathak, Mike Pritchard, Sanjay Choudhry

发布于 2026-05-13

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Corey Adams, Peter Harrington, Akshay Subramaniam, Mohammad Shoaib Abbas, Jaideep Pathak, Mike Pritchard, Sanjay Choudhry

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

以下是论文《ShardTensor：面向科学机器学习的领域并行》的通俗解读，通过日常类比拆解为简单概念。

核心难题：“大得装不下”的盒子

想象你是一位科学家，试图模拟飓风、黑洞或人脑。为了获得准确的结果，你需要以极致的细节观察这些事物——就像放大一张照片，直到能看清每一个像素。

在人工智能（AI）领域，这被称为高分辨率数据。

问题在于，AI 模型运行在一种名为 GPU（图形处理器）的专用计算机上。这些 GPU 的内存有限，就像一个小背包。

问题所在： 当科学家试图将巨大的高分辨率图像（例如风暴的 3D 扫描图）输入 AI 时，数据量过于庞大，根本无法装进这个“背包”。
旧方案一： 科学家过去不得不“下采样”数据。这就像把一部 4K 电影压缩成一个微小的、模糊的 144p 缩略图，仅仅为了它能塞进背包。AI 可以运行，但结果模糊且不准确。
旧方案二： 他们可以尝试将数据拆分到多台计算机上（数据并行），但这仅在你拥有许多独立的数据“块”（例如 100 场不同的风暴）时才有效。如果你只有一个巨大的风暴需要分析，旧方法就会碰壁。你无法轻易将一个风暴拆分成 100 份而不破坏数学逻辑。

解决方案：ShardTensor（“团队背包”策略）

来自 NVIDIA 的作者们介绍了一种名为ShardTensor的新工具。

将ShardTensor想象成一种神奇的方法，它能把一个巨大的披萨（你的高分辨率数据）切成片，并将这些切片分发给围成一圈的一群厨师（GPU）。

工作原理： 系统不是试图把整个披萨放在一个盘子里，而是按空间（按区域，而非按独立的披萨）将披萨切片。
- 厨师 A 拿着左上角的切片。
- 厨师 B 拿着右上角的切片。
- 厨师 C 拿着左下角的切片。
神奇之处： 当厨师们需要混合食材（进行数学运算）时，他们可以将自己切片的边缘传递给邻居。如果厨师 A 需要知道其切片边缘处发生了什么，他会询问厨师 B。他们协同工作以解决整个谜题，而无需将整张披萨放在一个盘子上。

这被称为领域并行。它允许科学家处理比单台计算机内存更大的数据，即使他们只有一个数据集可供使用。

为何重要（“何必费心”部分）

论文指出，在科学 AI 中，最占用内存的并非 AI 的“大脑”（模型权重），而是中间步骤（激活值）。

类比： 想象你在白板上解决一个巨大的数学问题。你不仅需要空间来写最终答案，还需要空间来存放沿途所做的每一个草稿计算。
结果： 对于高分辨率数据，这些“草稿计算”会瞬间占满内存。ShardTensor 将这些草稿计算分散到多个 GPU 上。
优势：
1. 强扩展性： 如果你拥有海量数据集，增加更多 GPU 会使任务完成得更快（就像在建筑工地增加更多工人）。
2. 弱扩展性： 如果你拥有一个大到以前根本无法运行的数据集，现在你可以通过增加更多 GPU 来分担负载从而运行它。

论文中的现实案例

作者通过两个具体的科学问题测试了该方法，以证明其有效性：

StormScope（天气预报）：
- 挑战： 预测单个雷暴需要以极高的细节（3 公里分辨率）查看整个美国的地图。
- 问题： 单台计算机的内存（80GB）无法容纳该细节下整个美国地图的数据。这就像试图把整张美国地图装进一个文件夹里。
- 解决： 使用 ShardTensor，他们将美国地图拆分到 32 个 GPU 上。AI 现在能够以高清“看到”整个国家，并准确预测风暴而不会崩溃。
Transolver（空气动力学）：
- 挑战： 模拟气流流过汽车以改善燃油效率。
- 结果： 他们能够在包含超过120 万个点的网格（3D 网格）上训练 AI，这些点代表了汽车的形状。这种细节水平以前在单台机器上是无法训练的。

论文未提及的内容（局限性）

作者诚实地说明了权衡取舍：

通信开销： 由于 GPU 必须不断相互通信以共享数据切片的边缘，因此会损失一点点“交谈时间”。
小数据： 如果数据量很小，这种交谈时间会使系统比仅使用一台计算机更慢。ShardTensor 仅适用于巨大的数据。
并非万能魔法： 它最适合特定类型的数学运算。某些较旧或非常具体的运算可能尚未得到支持。

总结

ShardTensor是一款新的软件工具，它允许科学家将巨大的高分辨率科学数据切片并分散到多台计算机上。这使得他们能够在以前无法装入单台计算机内存的超详细数据上训练 AI 模型，从而带来更准确的天气预报、更优的汽车设计以及更深入的科学发现。

技术摘要：ShardTensor——面向科学机器学习的领域并行

1. 问题陈述

科学机器学习（SciML）面临一个关键瓶颈：由于 GPU 显存限制，无法对高分辨率科学数据进行训练和推理。与大型语言模型（LLM）不同（LLM 的显存占用通常由参数量主导），SciML 工作负载经常受限于中间激活值。

该论文指出了科学数据中存在的“计算维度灾难”。对于 LLM，文本长度加倍大致使输入令牌数量加倍；而对于 $N$ 维科学数据，分辨率加倍会使数据量增加 $2^N$ 倍。对于 3D 或更高维度的数据（例如流体动力学、医学成像、气候建模），在训练反向传播过程中存储中间激活值所需的显存呈指数级增长。

现有解决方案往往无法扩展或会降低精度：

降低精度：仅提供适度的（2 倍）显存节省，且可能在科学代理模型中遭遇动态范围问题。
空间下采样：减少了显存占用，但牺牲了科学突破所需的高保真分辨率。
模型缩减：降低了精度。
激活检查点/CPU 卸载：虽然有效，但引入了显著的计算开销（重计算或主机 - 设备传输）。
标准并行：数据并行（拆分批次）在批次大小为 1 时达到硬性限制。模型并行（拆分权重）是为参数设计的，而非输入数据，且假设张量形状是静态的，这不适用于动态输入数据处理（例如改变全局形状的卷积操作）。

因此，目前缺乏一个通用框架，能够在每个设备的批次大小小于 1 的情况下对输入数据进行并行化，这阻碍了 SciML 达到科学家所需的原生分辨率。

2. 方法论：ShardTensor

作者提出了ShardTensor，这是一种领域并行的新范式，它将输入数据的空间维度与硬件约束解耦。它是 PyTorch DTensor（分布式张量）抽象的扩展，专门设计用于处理输入数据而非模型权重。

核心设计原则

领域并行：定义为跨设备对输入数据进行的并行化，超越了批次维度。当批次大小达到 1 时，ShardTensor 进一步将空间维度（例如高度、宽度、深度）在多个 GPU 上进行细分。
动态形状处理：与假设静态分布形状（适用于模型权重）的 DTensor 不同，ShardTensor 支持动态形状。它引入了第四个元数据组件：分片形状（sharding shapes），用于描述每个张量沿分片网格轴的局部块形状。这使得对非结构化或不均匀数据进行任意分块成为可能。
命令式分发：该框架是命令式的而非静态的。它利用 PyTorch 的 __torch_dispatch__ 和 __torch_function__ 机制动态路由操作。这使得系统能够逐层处理领域并行，而无需预编译计算图或对每一层进行定制修补。
可扩展性：用户可以通过高级功能接口（@custom_op）或低级分发逻辑（aten 操作）扩展系统。这确保了与不断演进的模型架构和自定义内核的兼容性。

实现细节

ShardTensor 是作为 PyTorch DTensor 的扩展构建的。

分发架构：它在 Python 层面拦截操作。如果特定操作具有领域并行处理程序，则执行该操作；否则，回退到 DTensor 语义或标准 PyTorch 行为。
通信：该框架管理集体操作（例如 all-reduce、all-gather）以及“光环（halo）”交换（从相邻设备获取相邻像素），以维持卷积或归一化等操作所需的数值一致性。
集成：它设计为可与现有并行策略（数据并行和模型并行）组合，允许在正交网格轴上实施 2D 或更高维度的并行策略（例如，在一个轴上拆分批次，在另一个轴上拆分空间维度）。

3. 主要贡献

ShardTensor 框架：一个通用的开源框架（可通过 NVIDIA PhysicsNeMo 获取），支持对小于 1 的批次大小的输入数据进行领域并行。
新颖抽象：引入“分片形状”来处理动态输入张量分布，使其区别于 DTensor 等专注于权重的并行化工具。
非侵入式编程模型：一种允许用户以最小代码更改包装模型输入的设计，无需为了并行性而重写层或模型。
可组合性：展示了在正交网格轴上将领域并行与数据并行和模型并行相结合的能力。

4. 结果与基准测试

该论文通过合成基准测试以及在 NVIDIA Blackwell (GB200) 和 H100 GPU 上的真实工业应用验证了 ShardTensor。

Ring Attention：在非常小的序列长度下，并行化开销占主导地位。然而，在非常大的序列长度下，该系统在训练和推理中均实现了与 GPU 数量的近线性强扩展。
Vision Transformer (ViT)：
- 强扩展：在高分辨率下（例如 $4096^2$ ），与单 GPU 相比，16 个 GPU 上的推理速度快15 倍。在 $2048^2$ 分辨率下的训练，使用 8 个 GPU 时速度快5 倍。
- 显存效率：2D 数据的显存使用量呈二次方增长，3D 数据呈三次方增长，证实了中间激活值占主导地位。使用 ShardTensor 的强扩展按比例减少了每个设备的显存占用，使得在标准硬件上训练具有超过10 亿个输入点的 3D 模型成为可能。
Transolver（汽车空气动力学）：
- 在 DrivaerML 数据集上进行训练，每个域包含多达120 万个点。
- 即使域大小增加 2 倍、4 倍和 8 倍，该系统仍保持了数值稳定性，并与单 GPU 基线相比保持了竞争力精度（压力和速度的 L2 误差）。
- 整个预处理流水线也实现了并行化，确保了端到端的可扩展性。
StormScope（天气预报）：
- 一个大陆尺度模型（CONUS），分辨率为3 公里（ $1024 \times 1792$ 网格）。
- 峰值显存使用量（114 GB）超过了单块 H100 GPU 的容量（80 GB）。
- ShardTensor 通过在 16 个数据并行组内的 2 个 GPU 上拆分激活值，实现了稳定训练。模型收敛稳定，跟踪了较低分辨率（6 公里）单 GPU 基线的损失轨迹，使得解析以前无法在不牺牲域范围的情况下获得的单个对流风暴成为可能。

5. 意义与主张

该论文声称，ShardTensor 解决了 SciML 中的一个根本性瓶颈：无法在科学家所需的原生分辨率下对数据进行训练和推理。

解除分辨率受限工作负载的束缚：它将此前因显存上限而无法进行的计算转化为可处理的计算，从而实现了体素医学成像、计算流体动力学和气候建模等领域的高保真训练。
可组合性：它消除了数据量、模型大小和分辨率之间的权衡。科学家现在可以同时沿多个并行轴进行扩展。
可及性：通过提供非侵入式编程模型，它降低了非分布式系统专家从业者的入门门槛，这与以往定制的、特定架构的并行化工作形成对比。

承认的局限性

作者谦逊地指出，命令式的逐层分发模型引入了 Python 层面的开销，使其在处理小操作或低分辨率数据时效率较低。性能也取决于互连带宽。此外，并非所有 PyTorch 操作目前都具有领域并行分发路径，与 torch.compile 的完全集成以实现静态图优化正在进行中。

总之，ShardTensor 代表了向使高领域并行在科学机器学习中像当今数据并行一样普及迈出的重要一步，促进了科学基础模型时代的到来。