Drawing from various experiences around setting up AI storage infrastructure, this guide aims to help engineers choose the most fitting solution for implementing storage for deep learning pipelines.
Use Case: Storage for Deep Learning
Problem Statement
Storage in deep learning is a tricky problem with multiple challenges across the stages of the deep learning pipeline, and one size does not fit all.
Realization Approach
The requirements for AI storage in deep learning pipelines need to be designed so that they can serve a specific purpose from initial data preparation to final inference stages. Each stage calls for different read and write capabilities, storage formats, timing, and bandwidth considerations.
Solution Space
A combination of high-performance filesystem, object storage, and local SSD cache serves as a flexible artificial intelligence storage solution that assists the various use cases in deep learning pipeline processing, such as large dataset storage, checkpointing, sharing data between GPUs, and high-speed read operations for model weights.
Featured AI Storage Solution
Nebius offers managed GPU cloud, storage and inference services for training and deploying AI models. Their AI storage capabilities include object store and shared filesystem on SSD, which compliments the GPU ML workloads with a secure, scalable and cost-effective solution to share and manage large amounts of unstructured data.
Good-Better-Best (GBB) Framework For Evaluating AI Storage
The rapid evolution of deep learning models has brought about unprecedented growth in both their size and complexity. This trend, while pushing the boundaries of what is technologically possible, has also placed immense demands on the underlying AI storage infrastructure, particularly in terms of data management and access speeds.
As organizations seek to benchmark their current setups and plan future infrastructure investments, here is a simplified Good-Better-Best (GBB) approach that can be help in laying out the benchmarks for evaluating different AI storage options for deep learning pipelines.
- Good: Meets baseline requirements for effective operations. Such solutions might be sufficient for smaller, less complex models or for organizations just beginning to scale their deep learning initiatives.
- Better: Provides significant improvements in performance, scalability, or efficiency. These are often the sweet spot for organizations that have outgrown their initial setups and need more robust infrastructure to handle increasing demands.
- Best: Represents state-of-the-art solutions offering top-tier performance and scalability. These solutions are typically adopted by organizations at the forefront of deep learning, where cutting-edge performance is vital for maintaining a competitive edge.
Please keep in mind that optimal solutions are not one-size-fits-all; they constantly change based on specific use cases and evolve with technological advancements. Therefore, the specific values provided here should only be considered within the current context, as they may shift with future developments.
AI Storage Solution Alternatives
A tiered AI storage approach often provides the best balance for managing the diverse demands of deep learning scenarios.
A high-performance shared filesystem is ideal for active data and recent checkpoints, offering the speed and reliability needed for demanding training and inference tasks. Object storage serves as a cost-effective solution for long-term storage and archival. Local NVMe SSDs or high-performance network block storage provide the necessary speed for caching and quick data retrieval.
| Storage type | Use cases |
|---|---|
| Shared filesystem | Data streaming, checkpointing, and sharing any data between GPU hosts. |
| Object storage (S3-compatible) | Large dataset storage, sharing inference results, potential for data streaming. |
| Network block storage and local NVMe SSDs | Boot disks, SSD cache, additional storage for self-managed solutions. |
Let’s look at some of the critical stages in the deep learning pipeline and evaluate the possible AI storage options with benchmarks as per the GBB framework.
Data preparation and tokenization
The first stage in any deep learning pipeline is data preparation. It is a critical phase that transforms raw data into a format suitable for model training. This process involves several key tasks, including data cleaning, concatenation, formatting, and updating metadata. For tasks involving computer vision, data augmentation is also a vital step that helps improve model generalization.
- Pre-computed augmentation: Applies transformations in advance, increasing storage requirements but potentially speeding up training.
- On-the-fly augmentation: Applies transformations dynamically during training, saving storage costs but increasing computational load.
A more advanced technique gaining traction is the use of heterogeneous clusters that combine CPU and GPU nodes. In this setup, CPU nodes handle preprocessing and augmentation tasks, freeing up GPUs to focus exclusively on model training.
From a storage perspective, data preparation presents several unique challenges. The file sizes involved can vary widely, from as small as 4KB to several gigabytes. Additionally, the read/write patterns are often unpredictable, with frequent read-modify-write operations that require robust storage solutions. In many cases, distributed computing resources are involved, further complicating the storage requirements. Given these challenges, the choice of storage solutions becomes critical.
Recommended AI storage solutions
- S3-compatible object storage like the one Nebius provides enables excellent scalability and compatibility with various data processing frameworks. It is particularly advantageous when prepared data needs to be streamed to different GPU providers for training. Object storage’s durability and cost-effectiveness is good for large datasets.
- Shared filesystem: Best used for heterogeneous clusters, which implies that data, CPU, and GPU compute must be within one provider’s network.
Data streaming for training
Efficient data streaming to GPU accelerators is key for maintaining high utilization rates and minimizing training time. This process typically involves transferring datasets from storage to the host machine’s RAM and then moving the data into GPU memory in batches.
The choice of storage solutions at this stage can significantly impact overall performance, with key factors including:
- Dataset size (ranging from 1GB to 100TB+)
- File sizes (from 39KB image files to 110MB TFRecord files)
- Model size (smaller models require more frequent data streaming)
Recent developments in storage technology have made significant strides in improving the interfacing with object storage systems. For example, AWS and Mosaic have introduced connectors that allow to optimize performance when streaming data from S3-compatible storage, thereby reducing transfer overhead and simplifying data pipelines. This is particularly beneficial when dealing with large-scale datasets, where efficient data shuffling between epochs is essential for preventing biases and ensuring that the model generalizes well.
Performance targets for data streaming:
| Compute components | Operation | Good, GB/s | Better, GB/s (H100/H200) | Best, GB/s (B100/B200) |
|---|---|---|---|---|
| Single GPU | Read | 0.5 | 1 | 2 |
| Single node | Read | 4 | 8 | 16 |
Recommended AI storage solutions
- High-performance shared filesystem: The default solution for most training workloads. Nebius’s file storage, with its distributed nature and high performance, excels in such scenarios, offering efficient shuffling capabilities that are crucial for large-scale setups. For example, Nebius offers a shared filesystem that provides lower latency compared to object storage and POSIX compatibility, which can be crucial for certain applications.
- S3-compatible object storage: For those dealing with large-scale data but with less stringent performance needs — especially when dataset size exceeds 1TB.
- Local SSD cache: Local storage could be placed between S3 and GPU instances and used as a cache to accelerate subsequent epochs, further optimizing the training process.
Bandwidth
As deep learning models and datasets grow in size, the bandwidth requirements for storage solutions become increasingly critical. Both aggregate bandwidth, which is the maximum data transfer capacity between the storage cluster and all its clients, and client bandwidth, the maximum data transfer capacity for a single virtual machine, must be carefully considered.
During studies, it has been found that in larger clusters, aggregate bandwidth requirements do not scale linearly, which can present challenges.
Performance targets for bandwidth:
| Number of nodes (8 GPUs per node) in training cluster | 8 | 16 | 32 | 64 | 256 |
|---|---|---|---|---|---|
| Good — text-only datasets | |||||
| Read, GB/s | 0.134 | 0.534 | 4.667 | 12 | 20 |
| Better — multimodal LLM training | |||||
| Read, GB/s | 0.267 | 1.067 | 9.334 | 24 | 40 |
| Best — image/video generation models | |||||
| Read, GB/s | 0.667 | 2.667 | 23.336 | 60 | 100 |
A tiered storage approach can be particularly effective in managing these varying bandwidth requirements.
Recommended storage solutions
- Compressed images, compressed audio, and text data can be stored in both shared filesystem and S3-compatible object storage solutions, as the bandwidth requirement is not very high.
- Image and video generation models might necessitate a high-performance shared filesystem or S3-compatible object storage with local SSD cache solution capable of handling massive bandwidth and I/O requirements.
For more AI storage guidelines related to other stages in deep learning pipeline, refer to the original post on Nebius.
Key Takeaways for Deploying AI Storage
One of the main lessons here is the importance of flexibility. What works well for one stage of the machine learning lifecycle might not be suitable for another. For instance, the storage requirements for data preparation are vastly different from those needed for inference.
Another important aspect to consider is the integration of storage solutions with other components of the deep learning infrastructure. Whether it’s ensuring that storage systems are aligned with the latest GPU accelerators, or that they can seamlessly integrate with data processing frameworks, the ability to create a cohesive and well-integrated system is crucial for maximizing performance.
Ultimately, the goal is to create a storage foundation that not only meets today’s needs but is also adaptable enough to handle tomorrow’s challenges. The decisions you make about storage will have a profound impact on your ability to train and deploy cutting-edge models, manage large-scale datasets, and ultimately deliver value from your AI initiatives.
