Pratik Mishra

Generative AI training is rapidly scaling in model size, data volume, and sequence length, requiring multiple instances for larger models. Distributed and parallel training strategies partition the training state across GPUs to support large-scale model training. As models and datasets grow, scaling infrastructure becomes a critical challenge. However, as AI infrastructure scales, the Mean Time Between Failures (MTBF) decreases, leading to more frequent job failures. Efficient recovery from failures is crucial, especially for resuming AI training. Existing systems offer limited support for reconfiguring parallelism mid-training, slowing progress during hardware failures or GPU re-allocation due to tightly coupled distributed checkpoints.

In this talk, we present Universal Checkpointing (UCP), a novel system enabling flexible and efficient generative AI training with reconfigurable parallelism on large-scale AMD GPU clusters. UCP addresses storage and memory performance challenges through careful hardware and software architecture considerations. Through under-the-hood analysis of the PyTorch GPU-Storage data path, UCP achieves optimal performance between AMD GPU clusters and high-performance remote storage systems.  Through our optimizations, UCP enables reconfiguration for broad set of popular parallelism strategies across varying GenAI models (size and type) with minimal reconfiguration costs, enhancing flexibility and resilience. These findings have wide applicability across the entire AI data pipelines, for instance, reducing the cold start overhead during inference or loading checkpoints to downstream post-training tasks, such as Supervised Fine-Tuning or Reinforcement Learning, etc.

This is a joint effort between UIUC and AMD.

Pratik Mishra

Senior Staff - Research

AMD

The performance gap between compute and storage has always been fairly considerable. This mismatch between applications need from storage and what they deliver in large-scale deployments is ever increasing, especially those catering AI workloads. The storage requirements encompasses nearly everything and anything, i.e. capacity/$, availability, reliability, IOPS, throughput, security, etc. Moreover, these requirements are highly dynamic for different phases of AI pipelines. For instance, emerging storage-centric AI applications such as Vector DBs and RAGs have unique requirements of performance, capacity and throughput compared to the most talked upon, i.e. AI Training and Inference. However, the main challenge common to all is to reduce data-entropy across the network, i.e., from storage to compute and vice-versa. All this while balancing the foreground and background IO processing.

In this talk, we would discuss on the core-issues and requirements from storage deployed in data-centers doing AI, with the focus on emerging application use-cases. It also aims to provide a holistic full-stack view of the infrastructure and the under-the-hood of how AI impacts each layer of computing, i.e. compute, network, and storage. This is followed with discussions on myraid of already existing solutions and optimizations which are currently deployed as piece-meal solutions but needs a refresh for enabling storage to meet the challenges posed by AI.

Pratik Mishra

Senior Staff - Research

AMD