21.3c Dataset pre-caching strategies and prefetch pipelines¶

🧭 Context Introduction¶

When training large AI models, the GPU often sits idle waiting for data to arrive from storage. This waiting time — called I/O stall — can waste up to 40% of GPU compute cycles. Dataset pre-caching and prefetch pipelines solve this by bringing data to the GPU before it is needed, keeping the GPU fully utilized.

Think of it like a kitchen chef: instead of waiting for each ingredient to be fetched from the pantry, the chef's assistant pre-stages all ingredients for the next dish while the chef is still cooking the current one. In AI infrastructure, pre-caching and prefetching do exactly this for data.

⚙️ What is Dataset Pre-Caching?¶

Pre-caching means loading a copy of frequently used training data into a faster storage tier (like GPU memory or high-speed NVMe) before training begins or during idle cycles.

Key characteristics: - Data is loaded once and reused many times - Reduces repeated reads from slow storage (e.g., HDDs or network storage) - Works best when the dataset fits within the cache capacity

Common pre-caching locations: - GPU memory (HBM): Fastest, but limited capacity (40–80 GB per GPU) - CPU system RAM: Larger capacity (256 GB–2 TB), but slower than GPU memory - Local NVMe SSD: Very large capacity (1–30 TB), but slower than RAM

🛠️ What is a Prefetch Pipeline?¶

A prefetch pipeline is a software mechanism that predicts which data samples the GPU will need next and loads them into a buffer in advance, overlapping I/O with computation.

How it works step-by-step: 1. The GPU processes batch N 2. While batch N runs, the prefetch pipeline loads batch N+1 into a staging buffer 3. When batch N finishes, batch N+1 is immediately available in GPU memory 4. The pipeline then starts loading batch N+2

Benefits: - Eliminates I/O wait time between batches - Keeps GPU utilization at 95–100% - Enables scaling to larger datasets without performance loss

📊 Comparison: Pre-Caching vs. Prefetch Pipeline¶

📊 Visual Representation: Dataset Pre-Caching pipeline¶

This flowchart maps out standard caching: staging dataset chunks in fast local NVMe SSDs before training starts to hide network latency.

🕵️ Common Pre-Caching Strategies¶

1. Full dataset pre-cache - Copy entire dataset to local NVMe before training - Best when dataset < 80% of local storage capacity - Eliminates all network I/O during training

2. Shuffle buffer pre-cache - Pre-load a large random sample into RAM - Shuffle within the buffer for each epoch - Reduces random access patterns to storage

3. Hot data pre-cache - Identify frequently accessed samples (e.g., validation set) - Keep only these in fast storage - Useful when dataset has skewed access patterns

🔄 Prefetch Pipeline Implementation Patterns¶

Pattern 1: Double buffering - Two buffers in GPU memory - GPU reads from buffer A while pipeline fills buffer B - Swap buffers when batch completes

Pattern 2: Multi-stage prefetch - Stage 1: Load from network storage to local SSD - Stage 2: Load from local SSD to CPU RAM - Stage 3: Load from CPU RAM to GPU memory - Each stage runs in parallel on different hardware

Pattern 3: Adaptive prefetch depth - Start with prefetching 1 batch ahead - Monitor GPU idle time - Increase prefetch depth (up to 4–8 batches) if GPU is still waiting - Decrease depth if memory pressure is high

🧪 Practical Considerations for Engineers¶

When to use pre-caching: - Dataset size is less than 500 GB - Training runs repeat the same data multiple times (multiple epochs) - Storage network is slow or congested

When to use prefetch pipelines: - Dataset is larger than available cache (terabytes) - Training uses random sampling each epoch - You need to maximize GPU utilization on expensive hardware

Common pitfalls to avoid: - Pre-caching a dataset that doesn't fit in memory causes swapping, which is worse than no caching - Prefetch pipelines that prefetch too far ahead can waste memory and cause cache thrashing - Mixing both strategies without coordination can lead to duplicate data in memory

✅ Summary Checklist for New Engineers¶

• [ ] Measure your current GPU idle time (use nvidia-smi and look at GPU-Util)
• [ ] If idle time > 10%, implement either pre-caching or prefetching
• [ ] For datasets < 100 GB, try full pre-cache to local NVMe
• [ ] For datasets > 1 TB, implement a multi-stage prefetch pipeline
• [ ] Always monitor memory usage — don't let caching starve the model
• [ ] Test with a small subset first, then scale up

Remember: The goal is to keep the GPU fed with data at all times. Pre-caching and prefetching are your two primary tools to eliminate I/O as the bottleneck in AI training pipelines.