Source-linked AI summary
FoldPipe: Bounded Remote Streaming of Native Molecular Shards with Asynchronous Prefetch
Dhiren Mukesh Khatri
TL;DR
Remote molecular training can require repeatedly loading preprocessed shards under constrained accelerator environments. FoldPipe streams existing native PyTorch/PyG shards with one-stage asynchronous prefetch and bounded buffering, then evaluates the mechanism and runtime on paired MD17/SchNet passes. It measured overlap, but the runtime interval included no effect, leaving reliable wall-clock advantage unresolved.
Problem
Remote preprocessed molecular graphs must be repeatedly retrieved when ephemeral or memory-constrained environments lack sufficient local storage and memory.
Method
FoldPipe provides a small layer that retrieves the next native PyTorch/PyG shard in the background while the current shard is consumed, with live payloads bounded independently of total dataset size.
Results
16.33 s mean retrieval–computation overlap was measured, while the geometric mean paired runtime ratio was 1.059× with a 95% interval of 0.878×–1.288×.
Takeaways & Limitations
FoldPipe verifies asynchronous overlap for native remote molecular shards, but observed overlap alone does not guarantee lower end-to-end runtime.
Takeaways & Limitations
The principal benchmark uses a public remote network, so alternating order does not ensure identical network conditions; it also covers one molecular system, model family, and accelerator configuration.
Abstract
from arXiv · showhide
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.
1 Introduction
FoldPipe targets remote streaming of existing native PyTorch/PyG molecular shards when memory and persistent storage are limited. Its contribution combines bounded asynchronous retrieval with an instrumented experiment that separates overlap from end-to-end speed.
- Remote molecular training data may be stored in many serialized shards while ephemeral environments provide limited host memory and persistent disk.
- FoldPipe consumes existing bounded-size .pt shards through a small source interface while retrieving shard i + 1 during computation on shard i.
- The system deliberately applies established prefetch and bounded-buffering techniques rather than claiming a new scheduling algorithm.
- Its contributions include a backend-independent native-shard interface, constant-with-shard-count live payloads, and a paired MD17/SchNet evaluation measuring overlap, memory, and runtime.
2 Related Work
FoldPipe occupies a narrow compatibility boundary among existing data-loading and streaming systems. It preserves already-generated native .pt shards instead of requiring conversion or adopting a broader input-processing framework.
- PyTorch and PyG provide general loading, OnDiskDataset, and batch prefetch facilities, whereas FoldPipe retrieves and deserializes remote shard files before ordinary PyG batching.
- WebDataset and the Hugging Face Hub support shard packaging or remote repository retrieval, but FoldPipe consumes existing .pt shards without conversion to tar records or a database.
- Unlike broader systems such as tf.data, FoldPipe intentionally remains a small orchestration layer rather than a general input-processing framework.
3 FoldPipe Design
FoldPipe abstracts remote shard enumeration and retrieval, then overlaps one-shard-ahead prefetch with current-shard consumption. Its working set is bounded by shard size and prefetch depth, subject to per-shard and deserialization constraints.
- Source abstraction: A FoldPipe source lazily enumerates shard identifiers and retrieves and deserializes one shard, with Hugging Face and Google Drive backends supported.
- Source abstraction: Remote bytes are streamed into CPU memory and deserialized, avoiding a required local-disk cache but allowing serialized bytes and objects to coexist briefly.
- Source abstraction: FoldPipe requires trusted native PyTorch/PyG shard files because deserialization uses torch.load(weights only=False).
- One-stage asynchronous prefetch: The single background worker submits retrieval of shard i + 1 before yielding batches from shard i to the consumer.
- One-stage asynchronous prefetch: Sequential streaming adds retrieval and consumption times, whereas ideal one-stage prefetch approaches a per-stage maximum and can hide part of retrieval latency.
- One-stage asynchronous prefetch: These overlap equations describe standard pipeline properties, not a FoldPipe-specific algorithmic result.
- Bounded working set: At steady state, the consumer holds the current shard while the background future holds the next shard or its in-progress buffer.
- Bounded working set: With bounded shards, live data-path memory remains independent of total shard count, although an individually oversized shard can still exceed available RAM.
4 Experimental Method
The evaluation uses a pinned MD17 aspirin SchNet workload and paired, order-alternating passes on a Tesla T4. It measures pipeline timing and resource behavior while accounting for paired uncertainty.
- Workload: The real workload transforms the MD17 aspirin trajectory into PyG graph objects and uses five benchmark shards containing 25,000 structures per timed pass.
- Workload: The repository is pinned to revision f779686deb9217877dd7ddde99b2522bd441492a.
- Workload: The model is PyG SchNet with 128 hidden channels, 128 filters, six interaction blocks, 50 Gaussian basis functions, and a 10 ˚A cutoff.
- Workload: Each batch contains 32 molecular graphs, and optimization combines energy mean-squared error with ten times force mean-squared error.
- Paired protocol: Twenty paired passes ran on a Tesla T4, using identical shards, initialization, and seed within pairs while alternating pipeline order.
- Measurement: The benchmark timestamps download, deserialization, training, and pass events, while excluding remote dataset discovery from timed comparisons.
- Measurement: Paired runtime ratios are summarized by their geometric mean, with a deterministic percentile bootstrap of paired log ratios estimating the 95% interval.
- Measurement: The benchmark summary is reported for twenty paired MD17/SchNet passes.
5 Results
FoldPipe consistently performs the intended retrieval–computation overlap, but the experiment does not establish a reliable wall-clock speed advantage. Prefetching slightly increases peak memory while keeping the working set independent of total shard count.
- 11 of 20 pairs were faster with FoldPipe, while the mean paired time saved was 6.59 s with a 95% interval from −7.69 to 21.11 s.
- 1.059× was the geometric mean paired runtime ratio, with a 95% paired-bootstrap interval of 0.878×–1.288×.
- Because the runtime interval includes 1.0, the experiment does not establish a reliable wall-clock speed advantage.
- 16.33 s mean retrieval–computation overlap per pass was measured for FoldPipe, while the sequential implementation had none by construction.
- FoldPipe has slightly higher peak RSS because a prefetched payload may coexist with the consumed shard, although both working sets remain independent of total shard count.
6 Discussion
FoldPipe’s observed overlap demonstrates that its pipeline operates as intended, while end-to-end throughput depends on remote I/O latency relative to per-shard computation. Its narrower role is to consume existing native molecular shards without the migration required by broader input systems.
- Remote retrieval overlaps with current-shard computation, but variable public-network latency can prevent lower end-to-end runtime.
- When next-shard computation time is comparable to retrieval time, much of retrieval can be hidden; when retrieval dominates, the consumer still waits.
- The experiment supports the intended mechanism rather than a universal speedup, because throughput value depends on remote I/O latency relative to per-shard computation.
- FoldPipe is most useful when existing native .pt shards and minimal format migration are primary constraints.
7 Limitations
The evaluation is constrained by public-network variability, limited replication, and a single molecular workload and accelerator configuration. FoldPipe also operates only at shard granularity and requires trusted serialized files.
- Public-network conditions were not identical across pipelines despite alternating their order, limiting direct end-to-end comparison.
- Twenty paired runs leave wide uncertainty intervals, so larger experiments are needed for tighter throughput-effect estimates.
- One molecular system, model family, and accelerator configuration characterize this workload rather than all molecular training pipelines.
- Shard granularity excludes sample-level streaming, automatic distributed sharding, generalized transformation graphs, and mature storage-system features.
- Native PyTorch object deserialization is unsafe for untrusted files, so only trusted serialized shards should be used.
8 Software and Data Availability
FoldPipe is available as open-source MIT-licensed software with a preserved research benchmark artifact and public reproducibility materials. The benchmark data derive from MD17 aspirin rather than constituting a new molecular dataset.
- FoldPipe is publicly available as open-source software under the MIT license.
- The current release is version 0.3.2, while the reported MD17 measurements are preserved as the frozen v0.3.0 research benchmark artifact.
- The public repository includes the benchmark report, raw statistics, per-shard traces, source manifest, execution log, benchmark image, and shard-construction scripts.
- The benchmark data are derived from the MD17 aspirin dataset rather than introduced as a new molecular dataset.
9 Conclusion
FoldPipe streams existing native molecular shards with bounded memory while prefetching the next shard during computation. The experiment demonstrates overlap but remains inconclusive about reliable wall-clock gains.
- FoldPipe provides bounded-working-set streaming for existing native PyTorch and PyG molecular-data shards stored remotely.
- Its single-stage prefetcher retrieves the next shard during computation, with memory bounded by shard size and prefetch depth rather than total dataset size.
- 16.33 s mean retrieval–computation overlap was measured across the 20-pair MD17/SchNet experiment.
- The experiment demonstrates the overlap mechanism while remaining inconclusive about a reliable wall-clock speed advantage under variable public-network conditions.
- FoldPipe is a low-conversion integration whose performance benefit depends on the I/O–compute regime, not a new prefetch algorithm or universally faster replacement.