Lesson 30 — AI/ML platform: distributed training storage & data loading at scale — dataset streaming, caching tiers & I/O stalls

2026-10-08

What interviewers are really probing

Lesson 29 kept a 16k-GPU synchronous job alive through hardware failure. This lesson covers the other way the same job wastes money while every pod reports Running: the GPUs are waiting for data. In a synchronous job the slowest rank's dataloader sets the pace for every rank, so one node with a cold cache, a CPU-throttled decode worker, or a 503-ing object-store prefix idles the whole fleet. L19 gave you the RDMA backend fabric (which is not the network your data rides on). L23 gave you SM-active and the roofline. L26 priced idle GPU-hours. L27 tiered weights for cold start; this is the same tiering problem for datasets, except you read them every step, for weeks.

What they check for (fuel: wafer-ai/gpu-perf-engineering-resources (opens in a new tab), platform framing, not kernels):

• Bytes-per-second math before architecture. Can you show that pre-tokenized LLM pre-training needs megabytes per second for the whole fleet, while vision/video needs terabytes per second, and design differently for each?
• Data format. Why millions of small files kill a parallel filesystem's metadata server and an object store's request budget, and why everyone converges on 100 MB–1 GB sequential shards (WebDataset tar, Mosaic MDS, Arrow/Parquet, tokenized .bin/.idx).
• Caching tiers. Object store → parallel FS or distributed cache → node NVMe → page cache / pinned RAM. What each tier is for, how each one fails, and why LRU gets a 0% hit rate when the dataset is bigger than the cache.
• Shuffle, sharding and determinism. Shard-level shuffle plus a sample buffer, IterableDataset duplicates, DistributedSampler.set_epoch, and a sample order that survives a world-size change.
• Resume correctness. The dataloader position is checkpoint state (L29 section 3). Resuming without it silently repeats or skips data.
• Finding I/O stalls. Data-wait per rank, worker profiles, CFS throttling, /dev/shm, PSI, FUSE latency, 503 SlowDown, and the startup storm when 2,000 nodes LIST a bucket at once.
• Contention. Checkpoint write bursts and dataset reads on the same filesystem.
• Security. The dataset is part of the model's supply chain: poisoning, provenance, read-only IAM, deserialization in data formats, pipe: URLs, and caches that quietly copy regulated data onto thousands of disks.

Interview sentence: "I start with bytes per second per GPU. Pre-tokenized LLM data is a few kilobytes per second per GPU, so there the problems are startup storms, metadata, tail latency and checkpoint contention. Vision and video are genuinely bandwidth-bound, so I shard into large sequential files, stream them from object storage through an NVMe cache with deep prefetch, push decode onto enough pinned CPU or the GPU, and alert on per-rank data-wait instead of average utilization. The dataset is an immutable, hashed, read-only snapshot, and its ID lands in every checkpoint manifest."

Core diagram: Snapshot → Tiers → Workers → GPU (+ detect, contention, security floor)

How to read it: top to bottom is the data path. Every tier exists to turn high-latency, request-limited reads into low-latency sequential reads close to the GPU. The dotted arrows are where you measure and where things collide. The PNG adds the numbers: bandwidth math, tier properties, what each pipeline stage hides, and the order to hunt a stall.

Idiosyncrasies / operator depth

1. Do the bytes/s math first: most LLM jobs aren't bandwidth-bound

need (bytes/s) = samples/s/GPU × bytes/sample × N_GPUs

Pre-tokenized LLM pre-training (illustrative). A 70B dense model at about 400 TFLOP/s achieved per H100 costs roughly 6 × 7e10 = 4.2e11 FLOPs per token, so about 950 tokens/s per GPU. With a vocabulary over 65,535 (Llama 3 is about 128k) each token ID is 4 bytes, so a GPU needs about 4 KB/s. 16,384 GPUs need about 60 MB/s total. A single object-store prefix can do that.

Vision (illustrative). A mid-size vision model at 2,000 images/s per GPU × ~110 KB per JPEG (roughly ImageNet's average) is about 220 MB/s per GPU, 1.8 GB/s per node, and about 0.9 TB/s across 512 nodes. Now storage, front-end network, and decode CPU are all real limits. Video and interleaved multimodal are 10–100× more bytes and usually need GPU decode.

So the LLM pain is elsewhere:

• Startup storm. 2,000 nodes each LIST a prefix and open their first shards in the same second. That's request rate, not bandwidth.
• Tail latency. A synchronous step waits for the slowest rank. A p99.9 object-store GET of 2 s with shallow prefetch shows up as a periodic step-time spike.
• On-the-fly tokenization and packing. CPU-bound, not I/O-bound. Ten workers per rank tokenizing JSONL can fall behind a small, fast model.
• Checkpoint bursts (L29): a 70B run writes on the order of a terabyte of model plus optimizer state every few minutes to tens of minutes, often to the same filesystem the data is read from.
• Mixture changes and restarts. Every restart (L29: every couple of hours at 16k GPUs) re-warms caches and re-seeks the data position.

The front-end network matters too. Data usually arrives over the node's front-end NIC (often 100–200 Gbps), not the 3.2 Tbps RDMA fabric from L19. Unless you deliberately run storage over RDMA (Lustre o2ib, Weka, GPUDirect Storage), the front-end NIC caps per-node ingest.

2. Format: small files are a metadata problem, not a bandwidth problem

• Object stores charge per request and rate-limit per prefix. S3 documents at least 3,500 PUT/COPY/POST/DELETE and 5,500 GET/HEAD requests per second per partitioned prefix, and it scales partitions gradually. A cold burst gets 503 SlowDown before the bucket adapts. LIST returns at most 1,000 keys per page, so listing a 50M-object dataset at job start costs thousands of sequential round trips per lister.
• Parallel filesystems (Lustre, GPFS) handle huge sequential reads across many OSTs/NSDs, but every open/stat hits the metadata server. Millions of 100 KB JPEGs in one directory saturate one MDT long before any OST is busy.
• The fix everyone converges on: pack samples into 100 MB–1 GB shards read sequentially, with a precomputed index.
  • WebDataset: POSIX tar files; each sample is files sharing a basename (000123.jpg, 000123.json). Streams from https://, s3:// via pipe:, or local disk.
  • Mosaic Streaming (MDS): shards plus an index.json, deterministic shuffle across world sizes, mid-epoch resume, per-shard hashes, and a local cache with cache_limit eviction.
  • Arrow / Parquet (Hugging Face datasets, Ray Data): columnar, memory-mappable; good for tabular and text, and for filtering before training.
  • Megatron-style tokenized .bin + .idx: pre-tokenized uint32 arrays, memory-mapped, random access by document offset. Tiny bandwidth, and the index is everything.
• Shard size tradeoff. Too small and you're back to request overhead. Too big and the shard count drops below ranks × workers, so some workers have nothing to read (or everyone reads the same shard). Rule of thumb: at least 10× more shards than world_size × num_workers, and even so, a resharding pass is cheaper than debugging an uneven epoch.

3. Caching tiers and how each one fails

The LRU trap. If the dataset is bigger than the cache and every epoch reads it in a fresh uniform-random order, LRU evicts each shard just before you need it again. Your hit rate approaches the fraction that fits, and with a cyclic pattern it can approach zero. Fixes: (a) shard affinity: each node keeps a stable subset of shards on its NVMe and shuffles within it plus a cross-node buffer (Mosaic's canonical-node partitioning works roughly this way), (b) size the cache for the whole working set, or (c) accept streaming and make prefetch deep enough that the object store's throughput, not its latency, is the limit.

Lazy-load filesystems. An FSx for Lustre filesystem linked to S3 (data repository association) imports metadata up front, but file contents load from S3 on first read. Epoch one runs at S3 speed and every later epoch runs at Lustre speed. Prewarm with lfs hsm_restore (or a Fluid DataLoad) before the gang is admitted, not while 16k GPUs wait.

4. FUSE mounts: convenient, with sharp edges

Mountpoint for Amazon S3, gcsfuse, and blobfuse2 let a training script open() an object as a file.

• Semantics. Mountpoint is optimized for large sequential reads and new-file writes; it doesn't support modifying existing files, and renames are limited (not supported on general-purpose buckets). Code that writes a temp file and renames it into place fails. For data that's fine; for checkpoints, use the SDK or DCP's storage writers.
• Metadata caching. Default TTLs keep a stat from turning into a HEAD per call. For an immutable snapshot set the metadata TTL to indefinite (--metadata-ttl indefinite on Mountpoint); for a prefix that changes under you, that same setting returns stale listings.
• The FUSE daemon is a process with memory. In the CSI-driver model it runs outside your container, so it isn't counted against your pod's limits the way you'd expect, and when it gets OOM-killed every open file returns Transport endpoint is not connected. That looks like a data bug to the trainer, and it's really a node problem.
• Double caching. FUSE local cache + kernel page cache + your loader's own buffer can hold three copies of a shard. Give the cache an explicit size and put it on NVMe, not the root disk.
• Privilege. FUSE needs /dev/fuse and usually SYS_ADMIN. Use the CSI driver so the trainer pod stays unprivileged (security section).

5. The DataLoader on a GPU node: CPU, memory and /dev/shm

• Prefetch depth. Each rank has up to num_workers × prefetch_factor batches in flight (prefetch_factor defaults to 2 per worker). Depth hides latency spikes. It doesn't help if average decode throughput is below the GPU's consumption rate: then the queue just drains.
• persistent_workers=True. Without it, workers are torn down and re-forked every epoch, and the first batches of each epoch stall while they rebuild state and reopen shards.
• Copy-on-access "leak". Workers are forked. A Dataset that holds a Python list of dicts touches each object's refcount on read, which dirties the page, which copies it into the worker. RSS grows to num_workers × dataset index size and looks exactly like a leak. Keep indexes in numpy arrays, Arrow tables, or memory-mapped files.
• /dev/shm. Workers hand tensors to the main process through shared memory. Docker defaults to 64 MB; Kubernetes gives you whatever the runtime gives you unless you mount an emptyDir with medium: Memory. Too small shows up as DataLoader worker (pid N) is killed by signal: Bus error or No space left on device in torch.multiprocessing. And medium: Memory counts against the pod's memory limit, so size it deliberately.
• CPU limits throttle in bursts. A CFS quota lets workers burn their whole period's budget in the first few milliseconds and then sit throttled, so you get periodic data stalls with average CPU under the limit. Check cpu.stat nr_throttled. For training pods: Guaranteed QoS, integer CPUs, static CPU manager policy, and topology manager aligning CPUs with the GPU's NUMA node.
• NUMA. A loader on socket 0 feeding a GPU on socket 1 crosses the inter-socket link for every H2D copy. nvidia-smi topo -m shows the GPU's CPU affinity; pin workers there.
• Pinned memory and H2D. pin_memory=True plus .to(device, non_blocking=True) lets the copy overlap compute. Pinned memory is page-locked RAM, and L29's async checkpoint staging uses the same pool, so budget both.
• GPU decode. When CPU decode can't keep up (video, high-res images), NVIDIA DALI or nvJPEG/NVDEC moves decode and augmentation onto the GPU. You trade some GPU time for a lot of CPU. GPUDirect Storage (cuFile) can DMA from NVMe or a supported parallel FS straight to GPU memory, skipping the bounce buffer; it pays off for large sequential reads, not small random ones.
• Rough CPU budget. Text tokenization is light; vision decode plus augmentation can need roughly 8–16 cores per GPU. GPU instances ship a fixed CPU:GPU ratio, so measure samples/s per core before you promise a throughput.

6. Sharding, shuffling and determinism

• IterableDataset duplicates. Each rank and each worker gets its own copy of the iterator. If you don't split by rank and by worker, every sample is seen world_size × num_workers times per "epoch". WebDataset's default node splitter is single_node_only, which raises in a multi-node job unless you pass nodesplitter=wds.split_by_node. That's a deliberate tripwire.
• DistributedSampler.set_epoch(epoch). Forget it and every epoch replays the same order. Also, drop_last=False pads by repeating samples so every rank has the same count; at 16k ranks and a small eval set that's a measurable bias.
• Global shuffle at petabyte scale isn't possible, so you approximate it: shuffle the shard list, then shuffle samples inside an in-memory buffer. A buffer that's too small yields batches correlated by source shard (all one language, one crawl date, one camera). Mosaic's py1e/py1br algorithms shuffle within blocks of shuffle_block_size samples spread across shards; bigger blocks are better randomness for more cache.
• Determinism across world sizes. If the sample order is a function of world_size, then L29's elastic restart at 1,920 instead of 2,048 GPUs changes which samples each rank sees, and resuming "at step 50,000" means something different. Mosaic fixes the order to num_canonical_nodes, independent of the physical node count; you want the same property from any loader you build.
• Data mixtures. Pre-training mixes many sources with target weights (web, code, math, multilingual). Mixture state (per-source cursors, tokens seen) must be checkpointed, and changing weights mid-run is a change you version and log, like a hyperparameter.

7. Resume: the data position is checkpoint state

L29 listed the dataloader position as state that silently diverges. In practice:

• Map-style + sampler: checkpoint the epoch, the sampler seed and the number of samples consumed globally; on resume, fast-forward. Naive fast-forward by re-reading data can take minutes at scale, so skip by index.
• torchdata.stateful_dataloader.StatefulDataLoader: drop-in DataLoader with state_dict()/load_state_dict(), aggregated across its workers but not across ranks. You save one state per rank, so you need DCP or a per-rank file, and a world-size change needs a global offset you can recompute.
• Mosaic Streaming: StreamingDataset.state_dict(num_samples, from_beginning) records the global sample offset, so resume works across world sizes as long as num_canonical_nodes is unchanged.
• Prefetched batches are not consumed. If you checkpoint "batches yielded by the loader" instead of "batches the optimizer actually stepped on", you skip up to num_workers × prefetch_factor batches per rank on every resume. Small per restart, large over hundreds of restarts.
• Record the dataset snapshot ID (and mixture version) in the checkpoint manifest. Without it, nobody can say which data a model saw, which matters for evals, data-deletion requests, and incident response.

8. Contention, startup and cost

• Checkpoint versus data. A DCP save (L29) is a write burst from every rank at once. On a shared Lustre/FSx filesystem it saturates the OSTs, reads queue behind it, and every save causes a data stall exactly when you thought async checkpointing made saves free. Fixes: separate filesystems (or buckets) for checkpoints and data, stagger rank uploads, put the first checkpoint tier on local NVMe, and keep enough data prefetched to ride out a save.
• Startup storm. At admission, every node opens the index, lists the prefix, and requests first shards. Precompute the index (index.json, a shard manifest) so nobody lists; spread shards across many prefixes (hash-prefixed keys) so S3 can partition; stagger start by rank; and prewarm the shared tier before the gang starts (Kueue admission check or an init Job). L20 gives you the admission hook.
• Data gravity. A dataset in us-east-1 feeding GPUs elsewhere pays egress on every epoch that streams it. Replicate the snapshot next to the capacity cell once (L24, L26) and charge the copy to the job; count it in the decision to move a job to cheaper GPUs in another region.

Concrete commands & config

Shard the dataset once (Mosaic MDS and WebDataset)

# Write MDS shards with per-shard hashes; upload straight to the versioned snapshot prefix.
from streaming import MDSWriter
 
columns = {"tokens": "ndarray:uint32", "source": "str"}   # explicit types, never 'pkl'
with MDSWriter(out="s3://corp-train-data/ds/v3/train",
               columns=columns,
               compression="zstd",
               hashes=["sha1", "xxh64"],
               size_limit="256mb") as w:
    for sample in packed_sequences():          # pre-tokenized, packed to seq_len
        w.write(sample)

# WebDataset: pack small files into ~512 MB tar shards with a stable naming scheme.
import webdataset as wds
with wds.ShardWriter("shards/train-%06d.tar", maxsize=512 * 1024**2) as sink:
    for key, jpg_bytes, meta in samples():               # meta is a dict -> written as JSON
        sink.write({"__key__": key, "jpg": jpg_bytes, "json": meta})

# Hash every shard and write a manifest the trainer verifies (security section).
find shards -name '*.tar' -print0 | xargs -0 -P 32 sha256sum | sort -k2 > manifest.sha256
aws s3 cp --recursive shards/ s3://corp-train-data/ds/v3/wds/
aws s3 cp manifest.sha256 s3://corp-train-data/ds/v3/wds/manifest.sha256

Stream it: Mosaic Streaming with a bounded NVMe cache

from streaming import StreamingDataset
from torch.utils.data import DataLoader
 
ds = StreamingDataset(
    remote="s3://corp-train-data/ds/v3/train",
    local="/nvme/cache/ds-v3",            # node-local NVMe, not the root disk
    split=None,
    shuffle=True,
    shuffle_algo="py1e",
    shuffle_seed=17,
    shuffle_block_size=1 << 18,           # bigger = better mixing, more cache
    num_canonical_nodes=128,              # FIXED for the life of the run → order survives rescale
    batch_size=8,                         # per-rank batch; must match the DataLoader
    predownload=64,                       # samples ahead per worker
    cache_limit="1.5tb",                  # evict instead of filling the disk
    validate_hash="sha1",                 # check shards against the index hashes
    download_retry=4,
    allow_unsafe_types=False,             # refuse pickle columns
)
loader = DataLoader(ds, batch_size=8, num_workers=8, pin_memory=True,
                    prefetch_factor=4, persistent_workers=True)
 
# Resume (state goes into the checkpoint alongside model/optimizer, L29):
state = ds.state_dict(num_samples=global_samples_stepped, from_beginning=False)
ds.load_state_dict(state)

Stream it: WebDataset done right in a multi-node job

import webdataset as wds
 
urls = "s3-shards.txt"   # a VERIFIED list of https:// or local paths, see security
ds = (
    wds.WebDataset(urls,
                   shardshuffle=1000,                  # shuffle the shard list
                   nodesplitter=wds.split_by_node,     # default single_node_only raises
                   workersplitter=wds.split_by_worker,
                   seed=17)
    .shuffle(10_000)                                   # sample buffer, costs host RAM
    .decode(wds.imagehandler("torchrgb"))              # explicit handler; no pickle members
    .to_tuple("jpg", "json")
    .batched(64, partial=False)
)
loader = wds.WebLoader(ds, batch_size=None, num_workers=12, pin_memory=True,
                       prefetch_factor=4, persistent_workers=True)

Stateful DataLoader for map-style data

from torchdata.stateful_dataloader import StatefulDataLoader
from torch.utils.data.distributed import DistributedSampler
 
sampler = DistributedSampler(ds, num_replicas=world, rank=rank, shuffle=True, seed=17, drop_last=True)
loader = StatefulDataLoader(ds, batch_size=16, sampler=sampler, num_workers=10,
                            pin_memory=True, prefetch_factor=4, persistent_workers=True,
                            in_order=False)            # one slow sample doesn't block the queue
for epoch in range(start_epoch, epochs):
    sampler.set_epoch(epoch)                         # or every epoch replays the same order
    ...
    ckpt["dataloader"] = loader.state_dict()         # per rank; saved via DCP with the model

Kubernetes: shm, NVMe cache, Guaranteed QoS, CSI mounts

# Trainer pod template (inside the JobSet from L29)
spec:
  containers:
  - name: trainer
    resources:
      requests: {cpu: "96", memory: 1600Gi, nvidia.com/gpu: 8}
      limits:   {cpu: "96", memory: 1600Gi, nvidia.com/gpu: 8}   # Guaranteed QoS, integer CPUs
    securityContext:
      allowPrivilegeEscalation: false
      capabilities: {drop: ["ALL"]}                               # no SYS_ADMIN: FUSE lives in the CSI driver
    volumeMounts:
    - {name: dshm, mountPath: /dev/shm}
    - {name: nvme-cache, mountPath: /nvme/cache}
    - {name: ds-v3, mountPath: /data/ds-v3, readOnly: true}
  volumes:
  - name: dshm
    emptyDir: {medium: Memory, sizeLimit: 64Gi}                   # counts against the memory limit
  - name: nvme-cache
    ephemeral:                                                     # local NVMe via a local-volume CSI / LVM driver
      volumeClaimTemplate:
        spec:
          accessModes: ["ReadWriteOnce"]
          storageClassName: local-nvme-encrypted
          resources: {requests: {storage: 3Ti}}
  - name: ds-v3
    persistentVolumeClaim: {claimName: ds-v3-s3, readOnly: true}
---
# Mountpoint for S3 CSI: read-only, one snapshot prefix, indefinite metadata TTL (immutable data only)
apiVersion: v1
kind: PersistentVolume
metadata: {name: ds-v3-s3}
spec:
  capacity: {storage: 1200Gi}       # ignored by the driver, required by the API
  accessModes: [ReadOnlyMany]
  storageClassName: ""
  claimRef: {namespace: train-team-a, name: ds-v3-s3}
  mountOptions:
  - read-only
  - region us-east-1
  - prefix ds/v3/
  - metadata-ttl indefinite
  csi:
    driver: s3.csi.aws.com
    volumeHandle: ds-v3-s3
    volumeAttributes: {bucketName: corp-train-data}
---
# FSx for Lustre (shared tier) — static PV
apiVersion: v1
kind: PersistentVolume
metadata: {name: fsx-train-data}
spec:
  capacity: {storage: 120Ti}
  accessModes: [ReadOnlyMany]
  csi:
    driver: fsx.csi.aws.com
    volumeHandle: fs-0123456789abcdef0
    volumeAttributes:
      dnsname: fs-0123456789abcdef0.fsx.us-east-1.amazonaws.com
      mountname: abcdefgh

Kubelet side, once per GPU node pool (L18): cpuManagerPolicy: static, topologyManagerPolicy: single-numa-node (or restricted), and reserved CPUs for system daemons and the FUSE/CSI daemons so they don't steal decode cores.

Prewarm the shared tier before admission

# FSx for Lustre linked to S3: metadata is imported; contents load lazily on first read.
# Hydrate file contents for this snapshot before the gang starts:
nohup find /fsx/ds/v3 -type f -print0 | xargs -0 -n 64 -P 32 sudo lfs hsm_restore &
lfs hsm_state /fsx/ds/v3/train-000123.tar         # 'released' = not yet hydrated
 
# Stripe large shards across all OSTs (set on the directory before files land):
lfs setstripe -c -1 -S 4M /fsx/ds/v3
lfs getstripe -d /fsx/ds/v3
lfs df -h /fsx                                    # OST balance; one full OST throttles everyone
 
# Straight to node NVMe from object storage (many parallel GETs, spread over prefixes):
s5cmd --numworkers 256 cp 's3://corp-train-data/ds/v3/train/*' /nvme/cache/ds-v3/

# Fluid + Alluxio: NVMe-tiered distributed cache in front of the bucket, prewarmed by a DataLoad
apiVersion: data.fluid.io/v1alpha1
kind: Dataset
metadata: {name: ds-v3, namespace: train-team-a}
spec:
  mounts:
  - mountPoint: s3://corp-train-data/ds/v3/
    name: ds-v3
  accessModes: [ReadOnlyMany]
  nodeAffinity:
    required:
      nodeSelectorTerms:
      - matchExpressions:
        - {key: node.kubernetes.io/instance-type, operator: In, values: ["p5.48xlarge"]}
---
apiVersion: data.fluid.io/v1alpha1
kind: AlluxioRuntime
metadata: {name: ds-v3, namespace: train-team-a}
spec:
  replicas: 64
  tieredstore:
    levels:
    - mediumtype: SSD
      path: /mnt/nvme/alluxio
      quota: 2Ti
      high: "0.95"
      low: "0.7"
---
apiVersion: data.fluid.io/v1alpha1
kind: DataLoad
metadata: {name: ds-v3-warm, namespace: train-team-a}
spec:
  dataset: {name: ds-v3, namespace: train-team-a}
  loadMetadata: true
  target:
  - path: /train
    replicas: 1

Finding the stall on a live node

# 1. Is it the loader at all? Per-rank data_wait from the trainer's own metrics (below).
# 2. What are the workers doing right now?
py-spy dump --pid "$TRAINER_PID" --subprocesses     # TRAINER_PID = the rank's main process
#    fetch (socket read / s3 client) vs decode (PIL / tokenizer) vs collate (torch.stack)
# 3. CPU throttling inside the container's cgroup:
cat /sys/fs/cgroup/cpu.stat | grep -E 'nr_periods|nr_throttled|throttled_usec'
# 4. Local disk and pressure:
iostat -xm 1 5
cat /proc/pressure/io /proc/pressure/memory
# 5. Shared FS client stats:
lctl get_param llite.*.stats | grep -E 'read_bytes|open|getattr'
lfs df -h /fsx
# 6. shm and kills:
df -h /dev/shm
dmesg -T | grep -iE 'bus error|oom|fuse|Transport endpoint'
# 7. GPU side: is SM activity sawtoothing at a data cadence?
nvidia-smi dmon -s u -d 1
nvidia-smi topo -m                                 # loader CPUs on the GPU's NUMA node?

PromQL: the SLI and the detectors

# Data-wait fraction per rank (trainer exports a histogram of time blocked in next(loader))
sum by (job, rank) (rate(trainer_data_wait_seconds_sum[5m]))
  / sum by (job, rank) (rate(trainer_step_seconds_sum[5m]))

# Worst rank vs median — the straggler that sets the pace
max by (job) (rate(trainer_data_wait_seconds_sum[5m]) / rate(trainer_step_seconds_sum[5m]))
  - quantile by (job) (0.5, rate(trainer_data_wait_seconds_sum[5m]) / rate(trainer_step_seconds_sum[5m]))

# CFS throttling on trainer pods (periodic data stalls with average CPU under limit)
sum by (pod) (rate(container_cpu_cfs_throttled_periods_total{namespace=~"train-.*"}[5m]))
  / sum by (pod) (rate(container_cpu_cfs_periods_total{namespace=~"train-.*"}[5m]))

# Node I/O pressure (node_exporter pressure collector)
rate(node_pressure_io_stalled_seconds_total[5m])

# GPU idle dips that line up with data stalls (DCGM, L23)
avg by (Hostname) (DCGM_FI_PROF_SM_ACTIVE)

Alert on max-rank data-wait over 2% for 15 minutes, not on average GPU utilization; a 99% average hides the one rank costing the whole job 5%.

Failures & fixes

Security & hardening

The dataset is part of the model. Whoever can change what a run reads can change what the weights do, and they can do it without touching the code, the image, or the checkpoint bucket from L29. Data loaders also deserialize and fetch untrusted bytes on every GPU node with the training job's identity, and caches copy regulated data onto thousands of disks. L11, L12, L13, L14, L27 and L29 built the controls for images, weights and checkpoints. Here they get applied to the data path, and tied back to weights so you can say which data produced which model.

Dataset, model-weight & checkpoint lineage (required)

Data-path and deployment controls

# Kyverno: training pods may only mount the dataset read-only, no privileged FUSE, shm sized
apiVersion: kyverno.io/v1
kind: ClusterPolicy
metadata: {name: training-data-path}
spec:
  validationFailureAction: Enforce
  rules:
  - name: no-privileged-or-sysadmin
    match: {any: [{resources: {kinds: [Pod], namespaces: ["train-*"]}}]}
    validate:
      message: "Training pods must not be privileged or add SYS_ADMIN; mount data via CSI."
      pattern:
        spec:
          containers:
          - =(securityContext):
              =(privileged): "false"
              =(capabilities):
                X(add): "null"
  - name: hostpath-only-nvme
    match: {any: [{resources: {kinds: [Pod], namespaces: ["train-*"]}}]}
    validate:
      message: "hostPath is limited to /mnt/nvme."
      deny:
        conditions:
          any:
          - key: "{{ request.object.spec.volumes[?hostPath].hostPath.path || `[]` }}"
            operator: AnyNotIn
            value: ["/mnt/nvme/*"]

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "TrainerReadsDeclaredSnapshotOnly",
      "Effect": "Allow",
      "Action": "s3:GetObject",
      "Resource": "arn:aws:s3:::corp-train-data/ds/${aws:PrincipalTag/dataset}/${aws:PrincipalTag/dataset-version}/*"
    },
    {
      "Sid": "DenyOutsideVpce",
      "Effect": "Deny",
      "Principal": "*",
      "Action": "s3:*",
      "Resource": ["arn:aws:s3:::corp-train-data", "arn:aws:s3:::corp-train-data/*"],
      "Condition": {"StringNotEquals": {"aws:SourceVpce": "vpce-0abc123def4567890"}}
    }
  ]
}

The first statement belongs in the trainer role's identity policy (the principal tags come from the job's workload identity); the deny belongs in the bucket policy. Deletion on snapshot prefixes is handled by Object Lock and a break-glass role, never by the trainer or the pipeline.

Interview soundbites

• "Do the bytes-per-second math first. A 70B model on sixteen thousand GPUs reads about sixty megabytes a second of pre-tokenized data. That job isn't bandwidth-bound; it's startup-, metadata- and tail-bound."
• "Small files are a metadata problem. Pack into hundred-megabyte-to-gigabyte shards, precompute the index, and never LIST a bucket at job start."
• "In synchronous training the slowest rank's dataloader sets the pace for everyone, so I alert on max-rank data-wait, not average GPU utilization."
• "Prefetch hides latency, not throughput. If decode can't keep up on average, a deeper queue just drains slower."
• "LRU over a dataset bigger than the cache, read in random order, can get a near-zero hit rate. Either size the cache for the working set, use shard affinity, or stream and accept object-store throughput."
• "The dataloader position is checkpoint state. It has to be a global sample offset, or an elastic restart at a different world size repeats or skips data."
• "Checkpoint bursts and dataset reads on one filesystem means every async save is a data stall. I separate them, or put the first checkpoint tier on local NVMe."
• "The dataset is part of the model's supply chain: immutable snapshots, signed hash manifests, a read-only trainer identity, no pickle in data formats, and the snapshot ID recorded in every checkpoint and released model."
• "WebDataset pipe: URLs run through a shell, so a shard list is code. I build it from a signed manifest."
• "Next: storage and data are done; Lesson 31 is evaluation and release gating at fleet scale."

How this maps to the JD

JDs that say "build the data and storage layer for large-scale training," "maximize GPU utilization," or "own the ML platform from data to deployment securely" are asking whether you can keep expensive accelerators fed without overbuilding storage, and whether you can prove what data went into a model. This lesson connects L19 (front-end versus RDMA networks), L20 (admission gates for prewarm), L23 (SM-active and per-rank signals), L24/L26 (data gravity and egress in placement), L27 (the same tiering for weights) and L29 (resume state, checkpoint contention), with the security floor from L11–L14. Fuel: wafer-ai/gpu-perf-engineering-resources (opens in a new tab).

Watch (talks & demos)

Real talks. Watch them after Failures & fixes.

Optimizing Data Locality and GPU Utilization for Training Workloads in Kubernetes

Bin Fan (Alluxio), KubeCon + CloudNativeCon Japan 2025. A Kubernetes-native distributed cache built on the NVMe already sitting in GPU nodes, with consistent hashing to place data and production examples of decoupling GPU clusters from the data lake. Operator takeaway: the shared-tier row of section 3, and why "copy the dataset to every cluster" doesn't scale once GPU capacity moves between regions. Watch it skeptically: it's a vendor talk, so check its claims against your own bytes/s math.

Open on YouTube (opens in a new tab)

A Distributed Stateful Dataloader for Large-Scale Pretraining

Davis Wertheimer & Linsong Chu (IBM Research), PyTorch Conference 2024. A PyTorch-native loader that is distributed, checkpointable and rescalable: streaming, stateless shuffling with a custom LCG, dynamic data mixing, and mid-epoch resume, proven on a 7B Llama run over 4T tokens. Operator takeaway: sections 6–7 done properly, especially "same data order at a different GPU count".

Open on YouTube (opens in a new tab)

Loading Training Data with WebDataset

Thomas Breuel (WebDataset author). Why sequential tar shards beat random file access, how sharding, shard shuffle and sample buffers fit together, and how the same pipeline reads from local disk, web servers or object stores. Operator takeaway: the format half of section 2 and the reason "small files" is a design smell at scale.

Open on YouTube (opens in a new tab)

References

• Formats & loaders: Mosaic Streaming (opens in a new tab) · Streaming docs (shuffling, canonical nodes) (opens in a new tab) · WebDataset (opens in a new tab) · Efficient PyTorch I/O for large datasets (WebDataset + AIStore) (opens in a new tab) · StatefulDataLoader (opens in a new tab) · torch.utils.data (opens in a new tab) · NVIDIA DALI (opens in a new tab) · GPUDirect Storage (opens in a new tab)
• Storage: S3 performance guidelines (per-prefix request rates) (opens in a new tab) · S3 Express One Zone (opens in a new tab) · Mountpoint for Amazon S3 (opens in a new tab) · Mountpoint S3 CSI driver (opens in a new tab) · FSx for Lustre data repositories (opens in a new tab) · FSx for Lustre CSI driver (opens in a new tab) · Cloud Storage FUSE CSI (GKE) (opens in a new tab) · Fluid (opens in a new tab) · s5cmd (opens in a new tab)
• Kubernetes node tuning: CPU management policies (opens in a new tab) · Topology manager (opens in a new tab) · emptyDir (medium: Memory) (opens in a new tab) · PSI (pressure stall information) (opens in a new tab)
• Security: Poisoning Web-Scale Training Datasets is Practical (Carlini et al.) (opens in a new tab) · OpenSSF model signing (sigstore/model-transparency) (opens in a new tab) · Sigstore cosign sign-blob (opens in a new tab) · SLSA provenance (opens in a new tab) · Kyverno policies (opens in a new tab) · S3 Object Lock (opens in a new tab)
• Talk slides: PyTorch Conf 2024 — distributed stateful dataloader (PDF) (opens in a new tab)
• wafer-ai/gpu-perf-engineering-resources (opens in a new tab) — AI/ML platform fuel
• Prior: 29 — Training reliability / checkpoint / elastic · 28 — Inference autoscaling · 27 — GPU supply chain / hydrate · 26 — GPU cost / reclaim / spot · 25 — Multi-tenant GPU isolation · 24 — Multi-cluster routing / capacity · 23 — GPU observability · 20 — Kueue / gang scheduling · 19 — GPU topology / NCCL · 18 — GPU node pools · 14 — ML deploy hardening · 13 — ML weights · 12 — Supply chain · 11 — PSS / admission / RBAC · 9 — NetworkPolicy / multi-NIC
• Next: Lesson 31 — AI/ML platform: evaluation & release gating at fleet scale — eval clusters, regression gates, canary rollouts of new weights

Quick self-check

Out loud, in under two minutes each:

1. Estimate the dataset read bandwidth for (a) an 8B LLM on 2,048 H100s with pre-tokenized uint32 data and (b) a ViT on 512 GPUs at 2,500 images/s per GPU. Which one needs a parallel filesystem, and what actually breaks in the other one?
2. Your 1,024-node job's first epoch is 3× slower than the rest and the first ten minutes are full of 503 SlowDown. Name the two causes and the fix for each, including where in the admission flow you'd put the fix.
3. Every rank reports 1.3 s per step, SM-active sawtooths every 40 steps, and average CPU is 55% of the limit. Walk through how you'd find the stall and the three most likely culprits.
4. An elastic restart moves you from 256 to 240 nodes. What has to be true of your data order and your checkpointed loader state so you neither repeat nor skip data? Why isn't StatefulDataLoader alone enough?
5. Design the IAM, bucket layout and manifest flow so a compromised trainer can't change any dataset, a poisoned shard is caught before training, and six months later you can list every released model that trained on snapshot ds/web/v7. Then name two ways a data loader itself can execute attacker-controlled code, and the control for each.