RetainZ improves AI checkpoint storage on zoned SSDs
RetainZ: Reclaim-Time-Aware Placement for AI Checkpoints on Zoned SSDs
Operating SystemsHardware Architecture
Summary
Saving AI model checkpoints often fills up storage too quickly because of how zoned SSDs manage space, making deletion inefficient. The authors show that data retention policies can predict when storage space can be reused without complicated predictions. They build RetainZ, a system that organizes saved data to better group files by when they can be deleted, freeing up space more efficiently. Their experiments demonstrate that RetainZ significantly reduces wasted storage and speeds up saving checkpoints compared to previous systems.
What this means in practice
- •For cloud storage providers: Improve storage efficiency and reduce latency for AI model checkpointing on zoned SSDs in data centers.
- •For ai infrastructure engineers: Manage checkpoint data with retention-aware allocation to reduce wasted SSD space and write overhead during model training.
Authors
Minxing Chu, Ruoxi Yang
Abstract
Solid State Drives (SSDs) are becoming the dominant medium for performance-critical storage. Zoned Namespace (ZNS) SSDs are getting more and more attractive because sequential writes and explicit zone resets reduce address-mapping, over-provisioning, and internal garbage-collection costs. However, reclaiming space requires resetting an entire zone, so deleting one file does not free its space while other files there must be kept. This erase-before-write constraint becomes a bottleneck for write-intensive workloads, particularly AI checkpointing during model training, where frequent saves of model and optimizer state compete for space with long-retained checkpoints. We observe that retention policies inherently encode when data can be retired, avoiding the need for a learned lifetime predictor. We present RetainZ, a storage backend that translates these policies into explicit reclaim epochs. Its allocator isolates long-retained data and groups objects with nearby reclaim epochs, balancing delayed space reclamation against unused space at zone ends. Evaluated on FEMU with multi-tenant streams and checkpoint manifests up to Pythia-1B, RetainZ reduces space occupied by deleted data and unused zone tails by 41.7% and mean per-input checkpoint p99 latency by 47.5% at 84% target load compared with the allocator of ZenFS, an established lifetime-aware ZNS backend, and lowers host write amplification from 1.165 to 1.003.