Storage Systems for
AI Workloads

💿 NVMe SSD Architecture

NVMe (Non-Volatile Memory Express) is the protocol standard for connecting SSDs over PCIe. It replaced older SATA and SAS protocols by taking advantage of PCIe lanes for massive parallelism. Understanding NVMe is foundational to AI storage design — it's the fastest per-node storage available.

Sequential Bandwidth Comparison

DGX H100 Local NVMe Configuration

⚙️ NVMe Key Concepts

NVMe Queues

NVMe supports up to 65,535 I/O queues with up to 65,535 commands per queue. Compare this to AHCI (SATA protocol): only 1 queue with 32 commands. This parallelism is why NVMe can hit 1M+ IOPS while SATA tops out at ~100K IOPS.

NVMe Namespaces

An NVMe namespace is a logical partition of an NVMe drive — analogous to a disk partition but more flexible. Supports multiple namespaces per device. In multi-tenant environments (MIG or virtualized), different namespaces can be assigned to different workloads for isolation.

NVMe-oF (NVMe over Fabrics)

NVMe-oF extends the NVMe protocol over a network fabric (RDMA/RoCEv2 or Fibre Channel), allowing remote NVMe drives to appear as local devices with near-local latency. Used in high-performance shared storage arrays for AI clusters. The WEKA filesystem uses NVMe-oF internally between nodes.

RAID Considerations for AI

• RAID 0 (striping): doubles bandwidth with two drives — used for maximizing checkpoint write speed from local NVMe. No redundancy.
• RAID 5/6: adds parity for redundancy but write penalty reduces performance — rarely used for AI scratch NVMe.
• JBOD / no RAID: use all drives independently in parallel via application-level striping (filesystem or application handles distribution). Common in DGX deployments for maximum bandwidth.

🌐 Parallel Filesystems — Why They Exist

A single DGX H100 has ~50 GB/s of local NVMe bandwidth. A DGX SuperPOD with 32 nodes has 32 × 50 GB/s = 1.6 TB/s of local storage — but those are 32 separate filesystems. Training a large model across all 256 GPUs requires shared access to training data and checkpoint directories. Parallel filesystems solve this by distributing storage across many nodes while presenting a unified namespace to all clients.

🗂️ Lustre — Deep Dive

Lustre is the most widely deployed HPC and AI filesystem. It separates metadata (file names, permissions, directory structure) from data (file contents), scaling each independently. Clients access files by first talking to the MDS for metadata, then reading/writing data directly from OSS nodes in parallel.

Key Lustre Concepts

Lustre Striping — Critical for AI Performance

Lustre striping splits a file across multiple OSTs so reads/writes can happen in parallel. A large dataset file striped across 8 OSTs can be read at 8× the speed of a single-OST file. For AI training, always stripe large files.

# Set stripe count (8 OSTs) and size (4MB chunks) for a directory
lfs setstripe -c 8 -S 4m /lustre/training_data

# Check stripe settings on a file
lfs getstripe /lustre/training_data/dataset.tar

# Check filesystem usage
lfs df -h /lustre

# Find files with suboptimal stripe count
lfs find /lustre/checkpoints -stripe-count 1

Lustre Performance Scale

Lustre performance scales linearly with OSS/OST nodes. A well-configured Lustre system for a DGX SuperPOD targets 1 TB/s+ aggregate read bandwidth for dataset loading, achieved by deploying enough OSS nodes with NVMe-backed OSTs. The MDS must also be sized to handle the metadata load from 256 clients opening millions of small files simultaneously.

⚡ WEKA — Flash-Native AI Filesystem

WEKA is a commercial all-flash parallel filesystem built from scratch for NVMe and RDMA. Unlike Lustre (ported from HDD-era architecture), WEKA was designed entirely around flash storage characteristics. It delivers sub-millisecond latency, native S3 API, and is an official NVIDIA-validated storage platform for DGX.

📋 Other Shared Filesystems

⚡ GPUDirect Technology Suite

GPUDirect is NVIDIA's family of technologies that enable data to move directly to and from GPU HBM memory without passing through the CPU and system DRAM. Each variant eliminates a different copy on the data path, dramatically reducing latency and freeing CPU resources for other work.

💾 GPUDirect Storage — Deep Dive

GPUDirect Storage (GDS) is the most exam-critical GPUDirect technology for NCP-AII storage topics. It enables a DMA engine to transfer data directly between NVMe storage (local or NVMe-oF) and GPU HBM — the CPU is completely uninvolved in the data transfer.

GDS Requirements & Setup

// Minimal GPUDirect Storage read example
#include <cufile.h>

CUfileDescr_t cf_desc = {0};
CUfileHandle_t cf_handle;
void *gpu_ptr;  // already cudaMalloc'd

// Open file with O_DIRECT | O_RDONLY
int fd = open("/lustre/model_weights.bin", O_RDONLY | O_DIRECT);
cf_desc.handle.fd = fd;
cf_desc.type = CU_FILE_HANDLE_TYPE_OPAQUE_FD;

cuFileDriverOpen();
cuFileHandleRegister(&cf_handle, &cf_desc);

// DMA: NVMe → GPU HBM (no CPU copy)
cuFileRead(cf_handle, gpu_ptr, file_size, 0, 0);

cuFileHandleDeregister(cf_handle);
cuFileDriverClose();

🔗 GPUDirect RDMA — Network Path

GPUDirect RDMA enables a ConnectX-7 NIC to DMA data directly from a remote server's GPU HBM into the local GPU HBM, bypassing both CPUs and system DRAMs in the transfer path. This is critical for multi-node AllReduce performance — without RDMA, every inter-node tensor transfer passes through two CPUs and two system memory copies.

GPUDirect P2P — Intra-Node

GPUDirect P2P enables GPU A to read from or write to GPU B's HBM directly over NVLink or PCIe — without staging through system DRAM. In DGX H100 with NVLink 4, P2P transfers at 900 GB/s make tensor parallelism within a node nearly as fast as accessing local HBM. Enable with cudaDeviceEnablePeerAccess(peerDevice, 0) and verify with cudaDeviceCanAccessPeer(&canAccess, devA, devB).

🔄 The Training Data Pipeline

Even with fast NVMe and parallel filesystems, data loading can bottleneck training if the software pipeline is inefficient. The goal is to keep GPU compute 100% utilized by ensuring the next mini-batch of data is always ready before the GPU needs it — this requires pipelining storage reads, preprocessing, and GPU transfer in parallel.

Standard PyTorch Pipeline

NVIDIA DALI — GPU-Accelerated Pipeline

NVIDIA DALI (Data Loading Library) replaces the CPU-side data preprocessing with GPU execution, dramatically reducing data loading overhead for vision and NLP workloads. DALI performs decoding, augmentation, normalization, and format conversion directly on the GPU.

import nvidia.dali.fn as fn
import nvidia.dali.types as types
from nvidia.dali.pipeline import pipeline_def

@pipeline_def(batch_size=256, num_threads=4, device_id=0)
def training_pipeline():
    jpegs, labels = fn.readers.file(file_root="/lustre/imagenet/train",
                                     random_shuffle=True)
    # Decode directly to GPU
    images = fn.decoders.image(jpegs, device="mixed",
                                output_type=types.RGB)
    # Augment on GPU
    images = fn.random_resized_crop(images, size=224, device="gpu")
    images = fn.crop_mirror_normalize(images, device="gpu",
                                       mean=[0.485*255, 0.456*255, 0.406*255],
                                       std=[0.229*255, 0.224*255, 0.225*255])
    return images, labels

Pinned Memory & Prefetching

Pinned (page-locked) memory on the CPU side enables direct DMA from CPU DRAM to GPU HBM over PCIe — no intermediate copy is needed. Paged memory requires the OS to first copy to a pinned staging buffer before PCIe DMA can occur.

# PyTorch: enable pinned memory in DataLoader
loader = torch.utils.data.DataLoader(
    dataset, batch_size=256,
    pin_memory=True,    # ← enables direct DMA to GPU
    num_workers=8,      # ← parallel CPU data loading
    prefetch_factor=2   # ← prefetch 2 batches ahead
)

💼 Checkpointing Strategy

Checkpointing saves model state (weights + optimizer state + training metadata) periodically during training. For large LLMs, checkpoints are hundreds of GB and must be saved quickly to minimize GPU idle time.

Checkpoint Size Estimation

Checkpoint Strategies Comparison

Checkpoint Staging Pattern

• Step 1 (fast, ~seconds): Save checkpoint to local NVMe via GDS or async CPU copy. GPU resumes training immediately.
• Step 2 (background): A background process (or async task) copies the checkpoint from local NVMe to the shared parallel filesystem (Lustre/WEKA) for durability across nodes.
• Step 3 (optional): From the shared filesystem, a separate offload job copies to object storage (S3) for long-term/cloud backup.

📦 Dataset Formats — Avoiding the Small-Files Problem

How you store training data is as important as what storage system you use. Millions of small files destroy filesystem metadata performance (Lustre MDS IOPS) and cause random I/O patterns on HDDs/SSDs. Packaging samples into large sequential archives enables high-throughput streaming reads.

🧪 Practice Quiz — Storage Systems

🃏 Flashcards — Tap to Flip

Click any card to reveal the answer

🤖 Storage Systems Advisor

Select a topic for exam-focused guidance:

🧩 Storage Mnemonics