arXiv:2608.27029cs.PFcs.LG2026-08

用异步预取让分子数据训练不卡顿,同时控制内存占用。

FoldPipe: Bounded Remote Streaming of Native Molecular Shards with Asynchronous Prefetch

论文配图:FoldPipe: Bounded Remote Streaming of Native Molecular Shards with Asynchronous Prefetch
图 1 · 摘自论文原文
  • 后台线程提前加载下一数据块,与计算并行进行。
  • 实测平均重叠时间16.33秒,比串行方案提升明显。
  • 适合在显存受限的云端训练分子模型的研究者。

在临时或内存受限的加速器实例上训练分子机器学习模型时,需反复从远程存储获取预处理的分子图数据。FoldPipe 是一个轻量级的 Python 协调层,用于已分片的 PyTorch 与 PyTorch Geometric 数据。它在消费者训练当前数据块的同时,于后台线程提前获取下一个数据块,确保活载荷数量相对于总数据集大小保持有界。异步预取与有界缓冲是成熟系统技术,并非新型调度算法。FoldPipe 的贡献在于针对原生 .pt 分片格式的简易集成,以及对其运行状态的源码定位实证分析。我们在 Tesla T4 上对 MD17 阿司匹林数据集上的 SchNet 能量与力预测任务进行了 20 次配对、顺序交替的基准测试,每次处理包含 25,000 个结构的五个固定分片。FoldPipe 平均 I/O-计算重叠时间为 16.33 秒,而串行有界基线为零。平均单次运行时间为 76.78 秒(FoldPipe)对比 83.37 秒(基线)。但配对几何均值加速比为 1.059×,95% 置信区间为 0.878× 至 1.288×。实验验证了重叠机制有效性,但在公开网络波动下未明确证明可靠的时钟速度优势。

原文摘要 · Abstract (English)

Training molecular machine-learning models on ephemeral or memory-constrained accelerator instances can require repeatedly retrieving preprocessed molecular graphs from remote storage. FoldPipe is a lightweight Python orchestration layer for already-sharded PyTorch and PyTorch Geometric data. It retrieves one shard ahead in a background thread while the consumer trains on the current shard, keeping the number of live shard payloads bounded with respect to total dataset size. Asynchronous prefetch and bounded buffering are established systems techniques rather than novel scheduling algorithms. FoldPipe's contribution is a small integration targeted at native .pt molecular shards together with a source-pinned empirical characterization of its operating regime. We evaluate a SchNet energy-and-force workload on MD17 aspirin using 20 paired, order-alternating benchmark passes on a Tesla T4. Each pass processes five pinned shards containing 25,000 structures. FoldPipe records 16.33 s mean I/O-compute overlap, compared with zero by construction for the sequential bounded baseline. Mean pass time is 76.78 s for FoldPipe and 83.37 s for the baseline. However, the geometric mean paired speedup is $1.059\times$ with a 95% bootstrap interval from $0.878\times$ to $1.288\times$. The experiment therefore verifies the overlap mechanism but is inconclusive about a reliable wall-clock speed advantage under the observed public-network variability.

分子建模数据流异步加载

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