Streaming 20 TB to 160 GPUs: A Production Architecture for Vision-Language Model Training

How a decoupled, streaming-first data architecture built on MosaicML Streaming delivered reliable, high-throughput multimodal training across hundreds of GPUs for weeks — without ever downloading the full dataset.

Who is this for: ML infrastructure engineers, platform engineers building training pipelines, and applied scientists running multi-week distributed training on cloud infrastructure.

Fine-tuning a 7B-parameter vision-language model on multimodal data sounds straightforward — until you realize your dataset is 20 terabytes, references 600 million images spread across cloud object storage, and your training cluster has nowhere near enough local disk to hold it all.

This was the challenge we faced. The solution was a streaming-first architecture built on MosaicML Streaming (MDS) that transforms cloud storage latency into a deterministic, high-throughput data plane — enabling multi-week training runs with sample-level fault tolerance, zero data duplication, and GPUs that never go idle waiting for data.

In this post, I'll walk through the complete system: why naive approaches fail at this scale, how MDS solves it, the three decoupled pipelines that make it operationally resilient, and the engineering details that kept 160 H100s fed for over three weeks straight.

The Problem: Why Naive Data Loading Fails at Scale

Training a large multimodal model means iterating over billions of tokens and millions of images. The data has to come from somewhere. Let's look at the two obvious approaches and why both break down.

Approach 1: Download everything first

Pre-downloading 20 TB to each training node's local storage simply isn't feasible. Even if the disk were large enough, the download time alone would add hours of idle GPU time before training begins. And if the job crashes mid-training? Start the download over.

Approach 2: Stream individual files on the fly

The alternative — fetching individual images directly from S3 during training — fails catastrophically due to network latency. In our measurements, the Time-To-First-Byte (TTFB) for an S3 GET request ranged from 20–100ms (the exact number varies by region, SDK, and object path, but the order of magnitude is consistent). If a single training batch requires 1,000 images, even parallel fetching hits request rate limits and imposes massive CPU overhead managing thousands of HTTP connections.

What we needed was something in between: a system that appears to be a fast local dataset but actually streams from cloud storage on demand, with intelligent caching and prefetching. That's exactly what MosaicML Streaming provides.

MosaicML Streaming (MDS): The Data Plane

MDS acts as high-performance middleware between cloud object storage and GPU training. It transforms the non-deterministic, latency-bound nature of S3 into a deterministic stream that mimics a local random-access array. Here's how it works.

The Shard Format

The fundamental unit in MDS is the shard — a self-contained binary file holding a contiguous sequence of samples. Each shard contains all the data needed to decode its samples, supports optional compression (we used zstd), and can be addressed independently. In our setup, we generated approximately 39,000 shards from ~5 million training samples across 15 tokenization nodes.

The Index: Fast Sample Lookup

At the heart of MDS sits index.json — a master manifest that enables fast random access to any of the ~5 million samples without scanning. It contains the sample count for each shard, from which MDS builds a cumulative sum table at initialization. Looking up any global sample index is then a binary search over shards — O(log S) where S is the number of shards — followed by O(1) access within the shard itself:

# Cumulative sum table built from index.json
# 5 shards with [128, 128, 100, 128, 50] samples:
cumsum = [0, 128, 256, 356, 484, 534]

def global_to_shard(global_idx, cumsum):
    shard_idx = bisect.bisect_right(cumsum, global_idx) - 1
    local_idx = global_idx - cumsum[shard_idx]
    return shard_idx, local_idx

# global_idx=300 → Shard 2, local sample 44
# O(log N) where N = number of shards

This is what enables MDS's killer feature: instant resumption. When training restarts after a crash, the system can jump directly to the exact sample where it left off — no re-streaming, no re-computation. At our scale, this saved thousands of GPU-hours across multiple node failures during the three-week run.

System Architecture: Three Decoupled Pipelines

The full system consists of three independent pipelines — tokenization, training, and continuous validation — unified by a single data layer (MDS shards on S3) and shared observability (Weights & Biases). Each pipeline can fail, restart, and scale independently.

The key insight behind this architecture is the separation of CPU/IO-bound work (tokenization: downloading 600M images, encoding, sharding) from GPU-bound work (training: forward/backward passes on 160 H100s). Running them together would waste expensive GPU time on image downloads. Decoupling means you tokenize once and run dozens of training experiments with different hyperparameters — iteration time drops from days to minutes.

Pipeline 1: Distributed Tokenization

The tokenization pipeline transforms raw multimodal samples (JSON metadata + S3 image URIs) into self-contained MDS shards. This is the most I/O-intensive phase: each of the 15 nodes downloads images using 80 concurrent threads with 100 S3 connections, achieving ~170 images/second per node — a 40–60× speedup over sequential downloads.

A critical design choice: images are embedded as Base64 directly into each sample. Raw samples contain cloud storage URIs pointing to images. During tokenization, each image is downloaded, JPEG-compressed, and Base64-encoded into the sample payload. The resulting MDS samples are fully self-contained — at training time, zero network calls are needed to image storage. This is what makes the checkpoint resume truly restart-safe: if the job resumes, it doesn't need to re-download any images.

After all nodes complete, a leader node merges the 15 local index.json files into a single master manifest covering all 39K shards, and uploads it to S3. A quality-gate report validates sample counts, image download success rates, and per-node statistics before training is cleared to begin.

Pipeline 2: Distributed Training at Scale

With the tokenized data sitting in S3 as self-contained shards, the training pipeline streams it on-demand to 160 H100 GPUs across 20 nodes. Here's how the data flows.

Deterministic shard assignment

MDS uses a simple, deterministic scheme based on global_rank: GPU 0 gets shards {0, 160, 320, ...}, GPU 1 gets {1, 161, 321, ...}, and so on. Assuming consistent seeds and per-rank partitioning (which MDS manages automatically), this guarantees each rank sees a disjoint subset of samples per epoch — perfect data parallelism with zero coordination overhead.

The data deserialization path

Each MDS sample arrives as a JSON string containing tokenized text (input_ids, attention_mask, labels) and Base64-encoded images. The deserialization path is: json.loads() → base64.b64decode() → PIL.Image.open(). The training framework's data collator then converts PIL images to pixel_values tensors via the vision processor, pads sequences, and moves everything to GPU. This on-the-fly deserialization adds negligible overhead while keeping samples self-contained in storage.

Prefetching hides latency

While the GPU processes current samples, the DataLoader prefetches the next shards in background threads. Combined with multiple S3 connections per node and NVMe caching, S3 latency becomes invisible to training — GPUs stay at full utilization throughout the run.

Fault Tolerance: Exact Resume After Node Failure

Over a three-week training run on 160 GPUs, hardware failures are not a question of "if" but "when." Our run encountered multiple node crashes. Each time, training resumed from the exact same loss value and global step — visible as seamless transitions in the W&B loss curve (each color representing a distinct run after a restart).

The math is straightforward. On restart, MDS reads global_step from the checkpoint's trainer_state.json and calculates:

# Checkpoint resume calculation
total_samples = global_step × batch_size × grad_accum × world_size
            = 500 × 2 × 8 × 160
            = 1,280,000

epoch    = total_samples // dataset_size   # which epoch
position = total_samples %  dataset_size   # position within epoch

# Each GPU skips exactly (position / world_size) samples
# → No wasted compute, no duplicate samples

This is one of MDS's most valuable features at scale. Without sample-level resume, each crash during a multi-week run would mean re-processing millions of samples — potentially adding days to the total wall-clock time.

Pipeline 3: Continuous Validation

Validation runs on a separate 4-node, 32-GPU cluster — completely independent from training. It continuously polls S3 for new checkpoints (every 60 minutes, configurable), downloads them, and runs a forward pass on held-out data to compute validation loss. No text is generated — just loss metrics for checkpoint selection and overfitting detection.

Both training and validation log to the same Weights & Biases project using grouped runs (same WANDB_RUN_GROUP), so scientists see training loss and validation loss side-by-side on aligned step axes. If training loss decreases but validation loss increases — overfitting. Both decreasing — healthy learning.

Operational Model: One Image, YAML-Driven Experiments

A single Docker image supports all three pipelines. An entrypoint.sh script reads the PIPELINE_TYPE environment variable (tokenization, training, or evaluation) and routes to the correct script. All tunable hyperparameters — learning rate, batch size, epochs, LoRA rank — live in YAML files stored in cloud storage. Scientists iterate by changing only two files: training.yaml and eval.yaml. Upload to storage, submit a job with the YAML URI as an environment variable, and the same container image handles everything.

The result: rapid iteration without container rebuilds. The Docker image is built once and never modified for experiments. This is crucial when each training experiment takes weeks — the overhead of rebuilding and redeploying containers would be a significant bottleneck.

MDS vs. Alternatives: Why We Chose MosaicML Streaming

The decisive factors for us were sample-level restart safety and automatic multi-node shard distribution. At 20+ TB and 3+ weeks of training, the ability to resume from any sample without re-streaming was non-negotiable. MDS's index-based seeking and deterministic shard assignment made it the clear choice.

Key Takeaways

Building ML infrastructure at this scale is an exercise in managing complexity through simplicity. Each component — the shard format, the index, the deterministic assignment, the checkpoint arithmetic — is individually simple. The challenge is composing them into a system that runs reliably for weeks at a time, across hundreds of GPUs, on petabytes of data. MosaicML Streaming gave us a solid foundation; the decoupled architecture gave us operational resilience. Together, they made it possible to focus on what actually matters: the science.

Further Reading & References