Optimizing Data Loading for ML Workloads

Introduction

Welcome back. In our previous lesson, we established the crucial role of containerization in creating reproducible, isolated software environments for our GPU-accelerated workloads. We've essentially packaged our compute engine. Now, we must address an equally critical question: how do we efficiently feed this engine with data?

This lesson directly addresses the learning outcome: Analyze data loading bottlenecks and compare the suitability of parallel file systems vs. object stores for ML workloads.

We'll begin by identifying why the data pipeline is often a hidden and costly bottleneck in large-scale AI clusters. Then, we will define and contrast the two primary storage architectures in this space: high-performance parallel file systems and massively scalable object stores. The core of our lesson will be a deep analysis of how each architecture suits the different stages of a modern ML workload, from data preparation to model training. You will see that the "obvious" choice isn't always the optimal one, and a more nuanced, disaggregated architecture is emerging as the standard for cost-effective, large-scale training.

1. The Data Loading Bottleneck

In a multi-million dollar GPU cluster, the single most expensive resource is idle time. If a GPU, capable of trillions of operations per second, is waiting for data, that is a direct and significant financial loss. The entire data path, from durable storage to the GPU's VRAM, is fraught with potential bottlenecks.

As the image shows, feeding the GPU is not a simple matter of reading a file. It involves moving data across networks, from disks to system RAM, processing it with the CPU, and finally transferring it to the GPU. Your goal as a system architect is to design a storage and data infrastructure that minimizes the time GPUs spend waiting.

2. The Contenders: A Tale of Two Storage Paradigms

At the heart of any large-scale data strategy are two fundamentally different approaches to storing and accessing data: file storage and object storage.

Let's briefly define our two main contenders in the context of large-scale systems.

2.1 Parallel File Systems (PFS)

A traditional file server, often using protocols like NFS (Network File System), presents a simple, familiar interface but becomes a performance bottleneck because all requests go through a single machine.

Parallel file systems (e.g., Lustre, BeeGFS, IBM Spectrum Scale) solve this by distributing both the data and the I/O load across many servers.

• Architecture: They consist of dedicated metadata servers (which manage the directory structure) and multiple object storage servers (which store the actual data blocks). Clients can talk to many storage servers in parallel, achieving extremely high aggregate bandwidth.
• Interface: They typically present a standard POSIX-compliant filesystem. This means applications can interact with them just like a local disk (open, read, write, seek), which is a major advantage for legacy HPC applications.

The following video gives a concise overview of this evolution from a simple NFS server to a more scalable parallel cluster file system.

2.2 Object Stores

Object stores (e.g., Amazon S3, Google Cloud Storage, Ceph) take a different approach, abandoning the hierarchical filesystem structure.

• Architecture: They manage data in a flat namespace. Each "object" contains data, a rich set of user-defined metadata, and a globally unique ID.
• Interface: Access is via a RESTful API (e.g., GET, PUT, DELETE operations over HTTP). There are no in-place modifications; an object is typically read or written in its entirety.
• Key Strengths: Their design allows for virtually limitless scalability, exceptional data durability, and generally lower cost per gigabyte compared to high-performance file systems.

The trade-off is clear: PFS offers extreme performance and a familiar interface at a high cost and complexity, while object storage offers immense scale and cost-efficiency with a different, API-driven access model and typically higher latency for individual operations.

3. Analyzing the ML Workload: A Stage-by-Stage Breakdown

So, which is better for ML? The answer depends entirely on which part of the ML lifecycle you're looking at. A monolithic "one-size-fits-all" storage system is rarely the most efficient or cost-effective solution.

To structure our analysis, we'll adopt the framework proposed by Glenn K. Lockwood in his excellent article, "LLM training without a parallel file system," which breaks the process into four stages.

3.1 Stages 1 & 2: Data Ingestion and Preparation

These initial stages involve collecting raw data (web scrapes, images, etc.) and processing it into a clean, tokenized format suitable for training. The I/O patterns here are characterized by:

• Massively parallel reads of raw data.
• Large, sequential writes of the processed data.
• Data being written once and read many times (immutability).

These characteristics are an almost perfect match for object stores. The API-driven access is well-suited to distributed processing frameworks like Apache Spark, and the write-once nature aligns with object immutability. Furthermore, a crucial architectural and economic argument comes into play: this work is I/O-bound and does not require expensive GPUs.

The key takeaway here is system disaggregation: it is more cost-effective to perform data preparation on a separate, CPU-based cluster optimized for I/O against a scalable object store, rather than using a single, expensive, GPU-centric cluster with a parallel file system for all tasks.

3.2 Stage 3: Model Training

This is the most performance-critical stage, where GPUs are fully engaged. It has two distinct I/O components: reading the training dataset and writing model checkpoints.

Reading Training Data

Conventional wisdom from HPC suggests that the repeated reading of data batches requires a high-bandwidth, low-latency parallel file system to keep the GPUs fed. However, the reality of modern LLM training is different.

1. Data size is manageable: After tokenization, the entire training dataset for even a massive model can be surprisingly small (tens of terabytes). For a large cluster, this is a tiny amount of data per node.
2. Variability is the enemy, not bandwidth: The biggest threat to GPU utilization isn't the raw speed of the storage system, but unpredictable I/O latency (the "noisy neighbor" problem) from sharing a central resource.

The emerging best practice is therefore to cache the data locally. At the beginning of a training job, the entire dataset is copied from a backend object store onto the fast, local NVMe SSDs present in each GPU node. The training loop then reads from this local, uncontended storage, completely isolating it from any network or shared storage bottlenecks.

This architecture effectively makes the performance of the shared storage system irrelevant for reads during the training loop, undermining a primary justification for an expensive, all-flash parallel file system.

Writing Model Checkpoints

Checkpointing is essential for fault tolerance in long-running jobs, but it presents a massive write burden. A single checkpoint for a large model can be terabytes in size. Synchronously writing this data to a shared file system would stall thousands of GPUs, costing a fortune in idle time.

The solution is asynchronous, multi-level checkpointing.

This strategy creates a hierarchy of storage tiers, balancing performance, cost, and durability.

1. Tier 1 (Fast/Local): After a training step, the model state is copied from GPU memory to a faster local medium like CPU RAM or NVMe SSDs. This is an extremely fast operation, allowing the GPU to resume computation almost immediately.
2. Tier 2 (Durable/Shared): In the background, a separate process asynchronously copies the checkpoint from the local tier to a durable, cost-effective object store.

This decouples the training loop from the slow I/O to durable storage. The high-performance file system is again bypassed in favor of a more sophisticated, tiered workflow that leverages local storage for speed and object storage for durable, cost-effective persistence.

The AWS blog post on checkpointing provides a complementary view, quantifying the impact of different strategies and discussing implementation details within cloud environments.

4. An Ideal AI Cluster Storage Architecture

By analyzing the I/O requirements of each stage, a clear architectural pattern emerges—one that deviates significantly from the traditional HPC model of a single, monolithic parallel file system. The optimal design is disaggregated and hierarchical.

In summary, the comparison is not simply "PFS vs. Object Store." It's about a fundamental shift in system design:

• Traditional PFS-centric Architecture: A single, high-performance parallel file system serves all I/O needs. It's simple to reason about but is expensive, can be a point of contention, and is often over-provisioned for many parts of the workload.
• Modern Disaggregated Architecture: A scalable object store acts as the central, durable data lake.
  • A separate CPU cluster handles data preparation against the object store.
  • Node-local NVMe SSDs in the GPU cluster act as a high-performance cache for training data.
  • Sophisticated framework logic orchestrates data movement, including initial data staging and asynchronous, multi-level checkpointing.

This modern architecture intelligently maps the I/O requirements of each workload component to the most cost-effective storage tier, maximizing performance where it matters most—the training loop—while leveraging cheaper storage for bulk data.

Conclusion

In this lesson, we dissected the complex data loading pipeline for large-scale ML and contrasted the two major storage paradigms used to feed it. We've moved beyond a simple "which is faster?" comparison to a nuanced architectural analysis.

Key Takeaways:

• Data loading bottlenecks are a primary cause of inefficient GPU utilization. Minimizing them is a key goal of AI system design.
• Parallel file systems (PFS) offer high performance and POSIX compliance but at a high cost and management complexity. Object stores provide massive scalability and cost-efficiency via an API-based model.
• A "one-size-fits-all" storage solution is suboptimal. A disaggregated architecture that matches storage tiers to workload requirements is more effective.
• The emerging best practice for large-scale training involves using an object store as the durable backend, a separate CPU cluster for data prep, and node-local NVMe SSDs as a high-performance cache during the training loop.
• Asynchronous, multi-level checkpointing is a critical technique that leverages fast local storage and durable object stores to provide fault tolerance without stalling the expensive GPU cluster.

Preview of the Next Lesson:

We have established an architecture that brings data to the compute node efficiently. However, the final step—moving data from the node's memory or local SSD into the GPU's VRAM—is its own critical bottleneck. Our next lesson, "Describe techniques for mitigating data pipeline bottlenecks, such as prefetching, caching, and GPUDirect Storage," will focus on the technologies that optimize this final, crucial data path, ensuring the GPU is fed with maximum efficiency.