AI Training Dataset & Checkpoint Storage: Parallel Byte-Range Reads, Throughput Limits, and Egress Control

Model Checkpoint Footprint Mathematics (Weights vs. AdamW Optimizer States)

During distributed training (PyTorch FSDP or DeepSpeed ZeRO), a full training checkpoint does not merely save the 16-bit (`BF16`) inference weights; it must persist the 32-bit master weights, `FP32` first-order momentum ($m$), and `FP32` second-order variance ($v$) states required by the AdamW optimizer to resume training without loss spikes:

Retention Rule: Never retain full 14-byte-per-parameter AdamW optimizer states for historical evaluation checkpoints. Persist the full 980 GB optimizer state strictly for the latest 2 rolling recovery steps (in case of hardware node failure), and strip optimizer states from historical milestone checkpoints so only the 140 GB `BF16` `.safetensors` inference weights are archived—cutting checkpoint storage costs by **85.7%**.

The Multi-Cloud GPU Training Egress Trap

Suppose a machine learning team stores a **40 TB multimodal training dataset** inside an AWS S3 bucket (`us-east-1`), and spins up a 16-node H100 cluster on an external GPU cloud (such as CoreWeave, Lambda, or GCP) where each node streams shards or syncs datasets locally across 5 training runs per month ($200\text{ TB}$ total outbound transfer):

Preventing GPU Starvation: WebDataset `.tar` Shards vs. Individual Files

An 8×H100 node training a vision-language model consumes roughly **12,000 to 24,000 image-text pairs per second**. If your dataset is stored in object storage as 500 million individual `50 KB` `.jpg` and `.json` files, your PyTorch `DataLoader` will spend 85% of wall-clock time waiting on HTTP TLS handshakes and `GetObject` TTFB latency—leaving your $25/hour GPU cluster idling at 15% compute utilization while racking up $200 per epoch in Class B `GET` fees.

The WebDataset Sequential Shard Pattern: Package training samples into sequential POSIX `.tar` archives (`shard-000000.tar` to `shard-009999.tar`), each sized between **256 MB and 1.0 GB**:

import webdataset as wds
from torch.utils.data import DataLoader

# Stream 512MB sequential TAR shards directly from zero-egress object storage
dataset_url = "https://datasets.internal-r2.daktan.com/laion-filtered/shard-{000000..004095}.tar"

dataset = (
    wds.WebDataset(dataset_url, shardshuffle=True)
    .shuffle(5000)
    .decode("pil")
    .to_tuple("jpg", "json")
)

loader = DataLoader(dataset, batch_size=256, num_workers=16, pin_memory=True)

Because each worker streams a contiguous 512 MB HTTP response sequentially, TCP congestion windows scale to full line rate (10–25 Gbps per stream) and S3/R2 Class B `GET` request volume drops by **10,000x**.

Saturating High-Bandwidth NICs with `s5cmd` & Asynchronous Checkpointing

When pulling a 140 GB model checkpoint (`29 × 4.8 GB .safetensors` shards) onto local RAID-0 NVMe scratch disks (`/mnt/local-nvme`) at container startup, standard Python `aws s3 cp` or single-threaded `boto3` scripts bottleneck due to Python GIL contention. Below is our measured cold-start download benchmark across three CLI tools on an 8×H100 node with dual 100 Gbps NICs:

# Pull 140GB Safetensors checkpoint to local NVMe scratch at 40+ Gbps
s5cmd --endpoint-url https://${ACCOUNT_ID}.r2.cloudflarestorage.com \
  --numworkers 256 \
  cp --concurrency 32 --part-size 64 \
  "s3://model-checkpoints/llama3-70b/step-48000/*.safetensors" \
  /mnt/local-nvme/checkpoints/step-48000/

Similarly, when saving checkpoints during training, use **`torch.distributed.checkpoint.async_save`**. This copies GPU HBM tensors into pinned host CPU RAM in under 2 seconds (allowing GPUs to immediately resume the next training step) while background CPU threads stream the 980 GB checkpoint to object storage asynchronously.