Accelerating Growth Through Optimizing GPU Usage // Sahil Khanna // AI in Production 2025

Training large-scale generative AI models, like those powering Adobe Firefly, requires a massive and efficiently managed GPU infrastructure. At Adobe, this involves managing hundreds of distributed jobs across thousands of GPU nodes. This scale introduces significant challenges in resource allocation, cost optimization, and reliability.

Core Challenges in GPU Optimization

The journey to optimize GPU usage at Adobe was driven by four primary challenges:

1. High Demand: Research and development for generative models that create videos, audio, and images require constant access to powerful and scarce GPUs (e.g., A100s, H100s).
2. Limited Supply: GPU capacity is finite and must be reserved years in advance, creating a fixed resource pool that must be managed carefully.
3. Resource Management: With reserved capacity, the platform must handle cost optimization, fair allocation of resources among different projects, and ensure the efficient use of every allocated GPU.
4. High Failure Rate: At scale, hardware failures are common. Nodes can overheat, lose connectivity, or suffer other issues, leading to performance degradation, training disruptions, and increased costs for long-running jobs.

Solution 1: An In-House, Multi-Cluster Scheduler

To address the high demand and complex resource management needs, the team developed a custom in-house scheduler. Off-the-shelf solutions couldn’t meet their specific requirements, which included:

• Cross-Cluster Scheduling: The ability to schedule jobs across multiple Kubernetes clusters to avoid bottlenecks and manage a large state of the world regarding node availability.
• Quota Management: A system to provide guaranteed, dedicated capacity for critical production workloads.
• Preemption: The capability to reclaim resources from lower-priority, non-critical jobs when a high-priority job needs them, ensuring that idle capacity is always put to use.

Solution 2: Supporting Development Workflows with Snapshotting

While the scheduler worked well for production training, it created problems for development workflows, which often involve interactive Jupyter sessions. These sessions are frequently idle and contain significant custom state (e.g., environment configurations). Simple preemption was too disruptive, as it would cause developers to lose hours of setup work.

The team engineered two key features to solve this:

• Snapshotting: Before preempting an interactive job, the system automatically takes a snapshot of the environment’s state. This allows a developer to resume their session within seconds, preserving their entire workspace and making preemption a seamless experience.
• Reclaim Policies: Automated policies were implemented to identify and reclaim resources from idle jobs, returning them to the general pool for other applications to use. This significantly improved the utilization of otherwise idle GPUs.

Solution 3: An Auto-Recovery System for Reliability

Generative model training can take days or even weeks, making it highly vulnerable to hardware failures. To combat this, an in-house auto-recovery system was built. This system acts as a centralized brain with three key functions:

1. It constantly monitors the progress of active training jobs.
2. It gathers health data from nodes and failure data from other systems.
3. It uses this information to make intelligent decisions on how to recover from various failures (e.g., bad hardware, connectivity issues) and automatically resumes the training process.

This system is critical for ensuring that long-running jobs complete successfully and efficiently, minimizing time lost to infrastructure instability.

AI Compute Platform Architecture

The components work together in a layered architecture:

• User Interfaces: Users interact with the platform via a UI, Python SDK, or CLI.
• API & Events Layer: User requests are translated into events that trigger actions within the system.
• Global Scheduler: This is the core decision-making component. It maintains a global state of all clusters, nodes, and quotas to decide where and when to schedule a job.
• Cluster Manager: Each Kubernetes cluster has a dedicated manager that receives instructions from the Global Scheduler and interacts with the Kubernetes master to schedule the actual pods.
• Node Agents: Agents run on each node and pod to collect health data and execute specific actions, such as taking a snapshot and saving it to S3.

• Tags:
• Generative AI
• Infrastructure
• MLOps
• Kubernetes
• GPU
• Distributed Training
• Adobe Firefly
• Scheduling