arXiv:2605.11111cs.DCcs.LG2026-05被引 1

提出ShardTensor,让科学机器学习可任意扩展数据规模。

ShardTensor: Domain Parallelism for Scientific Machine Learning

论文配图:ShardTensor: Domain Parallelism for Scientific Machine Learning
图 1 · 摘自论文原文
  • 通过解耦空间维度与硬件限制实现灵活数据并行
  • 训练和推理中均实现强弱扩展,降低延迟、处理更大数据
  • 适合超大规模科学数据的模型训练与推断场景

科学机器学习(SciML)在极端分辨率数据上面临独特挑战,现有缓解方法往往难以扩展或降低模型精度。尽管某些专用方法已在大规模空间数据上取得显著成果,但尚无通用框架支持单设备输入小于批量大小的并行化。本文提出ShardTensor:一种新型域并行范式,可将输入数据灵活扩展至任意规模。通过解耦输入数据的空间维度与硬件约束,ShardTensor使科学机器学习工作负载达到更高保真度的训练与推理水平。我们展示了训练与推理中的强弱扩展性能,强扩展下降低延迟,弱扩展下支持更大数据规模。此外,还实现了多维度并行,突破了科学机器学习在极端规模输入下的障碍。

原文摘要 · Abstract (English)

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.

科学计算数据并行模型训练

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。