mlp-storage Documentation
June 22, 2026 · View on GitHub
This directory contains reference documentation for mlp-storage and its dlio_benchmark submodule.
Benchmark Catalog
mlp-storage hosts four benchmark workloads:
| Benchmark | What It Measures | Where to Start |
|---|---|---|
| Training I/O | Storage throughput under AI training data loading patterns | QUICK_START.md |
| Checkpointing | Checkpoint save/restore performance (file and object store) | Streaming-Chkpt-Guide.md |
| KV-Cache | Storage performance for LLM KV-cache offloading (GPU → CPU → NVMe) | kv_cache_benchmark/README.md |
| Vector DB | Vector similarity search storage performance (Milvus) | vdb_benchmark/README.md |
Where to Start
| Your goal | Start here |
|---|---|
| First time — install and run any benchmark | QUICK_START.md |
| Run or understand any test (unit, integration, object-store) | ../tests/README.md |
| Benchmark LLM KV-cache offload storage | kv_cache_benchmark/README.md |
| Benchmark vector database storage (Milvus) | vdb_benchmark/README.md |
| Set up object storage (S3 / MinIO / Azure / GCS) | OBJECT_STORAGE_GUIDE.md |
| Test object storage locally or in CI | OBJECT_STORAGE_TESTING.md |
| Compare object storage libraries (s3dlio, minio, s3torchconnector) | OBJECT_STORAGE_GUIDE.md |
| Understand map-style vs. iterable DataLoader tradeoffs for S3 | DATALOADER_ARCHITECTURE.md |
| Benchmark NVMe with O_DIRECT (bypass page cache) | DATALOADER_ARCHITECTURE.md — O_DIRECT section |
| Understand AIStore gaps, reader/checkpoint issues, rationalization options | dlio_benchmark/docs/AIStore_Analysis.md |
| Test streaming checkpointing | Streaming-Chkpt-Guide.md |
| Configure multi-endpoint / load-balanced object storage | MULTI_ENDPOINT_GUIDE.md |
| Understand the system architecture | ARCHITECTURE.md |
| Add a new workload or benchmark | ADDING_BENCHMARKS.md |
Document Reference
Getting Started
QUICK_START.md
First steps for all four benchmark types: training I/O (local + S3, all three object storage libraries), checkpointing (file and object-store), KV-Cache, and Vector DB. Quick-start commands with links to full documentation for each.
ARCHITECTURE.md
System architecture overview: how mlpstorage, dlio_benchmark, and the object storage library layer fit together. Explains the reader plugin model, MPI execution, and data-flow from storage to the training loop.
KV-Cache Benchmark
kv_cache_benchmark/README.md ← Full KV-Cache documentation
The KV-Cache benchmark simulates LLM inference KV-cache offloading — the process by which production inference systems move intermediate attention state (Key-Value tensors) from expensive GPU VRAM to CPU RAM or NVMe storage when memory is exhausted. It answers:
- What is the real latency impact of each storage tier (GPU vs. CPU vs. NVMe)?
- Is your NVMe fast enough to sustain cache spillover at your target user count?
- How many concurrent users can your storage tier support at a given throughput?
Workload types: synthetic multi-user conversation traffic, ShareGPT trace replay, BurstGPT trace replay.
Quick start:
cd kv_cache_benchmark
pip install ".[full]"
python3 kv-cache.py --model llama3.1-8b --num-users 50 --duration 120 \
--gpu-mem-gb 0 --cpu-mem-gb 4 --cache-dir /mnt/nvme --output results.json
- Location:
mlp-storage/kv_cache_benchmark/ - Unit tests:
pytest kv_cache_benchmark/tests/ -v - See kv_cache_benchmark/README.md for full configuration, ShareGPT/BurstGPT replay, result interpretation, and MLPerf submission guidelines.
Vector Database Benchmark
vdb_benchmark/README.md ← Full Vector DB documentation
The Vector DB benchmark measures storage subsystem performance for vector similarity search workloads. It currently supports Milvus with three index types: DiskANN (disk-based ANN), HNSW (in-memory graph), and AISAQ (quantization). Use it to compare NVMe, NFS, or object-backed storage for vector search.
Benchmark steps: load vectors → build index → run similarity queries → measure throughput, latency, and recall.
Quick start:
cd vdb_benchmark
docker compose up -d # starts Milvus + MinIO + etcd
# then follow vdb_benchmark/README.md for load/index/query steps
- Location:
mlp-storage/vdb_benchmark/ - Tests:
vdb_benchmark/tests/ - See vdb_benchmark/README.md for Docker setup, Milvus configuration, benchmark execution, and result interpretation.
Training I/O Benchmark (DLIO)
Uses the DLIO benchmark to simulate deep learning training data loading patterns across multiple storage backends.
OBJECT_STORAGE_GUIDE.md ← Main object storage reference
Primary guide for running mlpstorage training workloads against S3-compatible
object storage. Covers MinIO quick start, .env configuration, the current
file|object CLI grammar, --data-dir prefix semantics, supported libraries
(s3dlio, minio, s3torchconnector), library capability comparison,
multi-protocol notes, and common troubleshooting.
dlio_benchmark/docs/AIStore_Analysis.md
Detailed gap analysis of the native AIStore support (storage_type: aistore)
versus the S3 multi-library path. Covers four specific gaps — checkpointing
(silently falls back to local-disk PT_SAVE), per-format reader routing
(JPEG/PNG broken; NPY/NPZ loses streaming reader; Parquet untested), config
validation gaps, and zero checkpoint test coverage. Includes a full feature-parity
table and three concrete rationalization options (A: S3 gateway, B: fill gaps,
C: consolidate as 4th library) with a pros/cons comparison and a per-option file
change list.
OBJECT_STORAGE_TESTING.md
Small repeatable object-storage tests: Python library smoke tests, parser checks, tiny MinIO datagen/run cycles, multi-endpoint smoke validation, and the most useful unit tests to run when changing object-storage behavior.
MULTI_ENDPOINT_GUIDE.md
Multi-endpoint load balancing for object storage: comma-separated URI lists, template expansion, file-based endpoint lists, and MPI rank-based distribution. Compares native multi-endpoint (s3dlio) vs. MPI rank selection across all three object storage libraries.
Streaming-Chkpt-Guide.md
The two checkpoint optimizations: dgen-py integration (155× faster data generation) and StreamingCheckpointing (producer-consumer pipeline, 192× memory reduction). Architecture diagrams, tuning parameters, and expected output.
Performance and Data Formats
PARQUET_FORMATS.md
Parquet format support via two new DLIO reader classes: ParquetReader
(local/NFS filesystem, pyarrow native, row-group LRU cache) and
ParquetReaderS3Iterable (S3 object storage, byte-range GETs, all three
object storage libraries). Includes YAML config examples and unit test commands.
DATALOADER_ARCHITECTURE.md
Architecture and tradeoff analysis for map-style vs. iterable-style data loaders on both object storage and local NVMe. Two major topics:
Part 1 — Map-style vs. iterable on S3 (implemented via TorchIterableDatasetSimple):
Explains why the conventional "iterable is better for large datasets" advice originates from
spinning-disk seek patterns and does not transfer directly to S3. Covers the real argument
for iterable on object storage (pipeline depth: 64 in-flight GETs per worker, up to 256 total),
the tradeoffs (shuffling, worker partitioning, prefetch memory), and current implementation
status for NPZ/NPY/JPEG/PNG workloads.
Part 2 — O_DIRECT on local NVMe (two independent paths): Why O_DIRECT is required for accurate NVMe benchmarking (page cache bypass). Detailed comparison of both available O_DIRECT mechanisms:
odirect: true— legacy Pythonos.open + os.readv, map-style, 1 read/worker (baseline)storage_library: direct— Rust/Tokiolibc::O_DIRECT, iterable-style, 64 reads/worker
Includes a full 12-property comparison table and guidance on when to use each path (and why keeping both enables a direct comparison isolating I/O concurrency depth and GIL contention). Essential reading before any DataLoader refactor or NVMe benchmarking run.
Extending the Benchmark Suite
ADDING_BENCHMARKS.md
How to add new benchmark workloads: DLIO config structure, workload parameters, dataset format registration, and integrating custom storage readers.
Test Scripts
For a complete guide to running tests — including environment setup, unit tests, integration tests, and object-store performance scripts — see tests/README.md.
testing/TEST_README.md lists legacy quick-run
commands for the major benchmark workloads. Run those scripts from the project
root (not from inside docs/).
The quick-link tables below list the most commonly used scripts.
Quick Links — Test Scripts
Training I/O and Object Storage Tests
| What | Script |
|---|---|
| Object-storage smoke guide | docs/OBJECT_STORAGE_TESTING.md |
| Full object-store script guide | tests/object-store/README.md |
| DLRM Parquet benchmark | tests/object-store/run_dlrm_bench.sh |
| Flux Parquet benchmark | tests/object-store/run_flux_bench.sh |
| RetinaNet JPEG datagen + run | tests/object-store/gen_retinanet_jpeg.sh, tests/object-store/test_retinanet.sh |
| UNet3D NPZ datagen + run | tests/object-store/gen_unet3d_npz.sh, tests/object-store/test_unet3d.sh |
| Object checkpoint write/read | tests/object-store/run_checkpointing.sh |
| Unit tests (no infrastructure needed) | pytest tests/unit/ |
| Integration tests (requires S3 endpoint) | pytest tests/integration/ |
Checkpointing Tests
| What | Script |
|---|---|
| File checkpoint demo | tests/checkpointing/demo_checkpoint_methods.sh |
| Object-store checkpoint demo (all 3 libraries) | tests/object-store/demo_streaming_checkpoint.sh |
| s3dlio checkpoint test | tests/object-store/test_s3dlio_checkpoint.py |
| minio checkpoint test | tests/object-store/test_minio_checkpoint.py |
| s3torchconnector checkpoint test | tests/object-store/test_s3torch_checkpoint.py |
| Streaming backend comparison | tests/checkpointing/test_streaming_backends.py |
KV-Cache Tests
| What | Script |
|---|---|
| KV-Cache unit tests | pytest kv_cache_benchmark/tests/test_kv_cache.py -v |
Vector DB Tests
| What | Script |
|---|---|
| Vector DB tests | vdb_benchmark/tests/ |