Rubick: Exploiting Job Reconfigurability for Deep Learning Cluster Scheduling

Rubick: Exploiting Job Reconfigurability for Deep Learning Cluster Scheduling

Speakers: Xinyi Zhang, Hanyu Zhao, Wencong Xiao, Xianyan Jia, Fei Xu, Fangming Liu

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

In the rapidly evolving landscape of deep learning, the training of large models, particularly Large Language Models (LLMs) built on transformer architectures, demands immense computational resources. This talk introduces Rubick, a novel deep learning cluster scheduler designed to significantly boost training performance and resource efficiency by intelligently exploiting the inherent reconfigurability of deep learning jobs. Traditionally, deep learning jobs are treated as "black boxes" by schedulers, with their execution plans and resource requirements fixed at launch time. This static approach leads to a considerable mismatch in dynamic, shared GPU clusters where available resources fluctuate unpredictably. Rubick challenges this paradigm by adopting a "white-box" approach, allowing the scheduler transparent access to a job's execution plan and multi-dimensional resource needs, enabling dynamic adjustments to both.

The core innovation of Rubick lies in its ability to jointly optimize a job's execution plan and resource allocation in response to changing cluster conditions. This is achieved through a sophisticated performance model that accurately predicts training throughput under various configurations, coupled with a scheduling policy guided by resource sensitivity curves. By understanding how different execution strategies (e.g., 3D parallelism, gradient checkpointing, CPU offload) interact with varying resource allocations, Rubick can dynamically reconfigure running jobs and adapt waiting jobs to the current cluster state. This dynamic adaptability allows for more concurrent job execution, leading to substantial improvements in average job completion time and overall cluster throughput, a critical advancement for shared, multi-tenant deep learning infrastructure.

Background

▶ Watch: Introduction: Problem of static Deep Learning job scheduling (0:00)

The proliferation of large language models has fundamentally shifted the demands on deep learning infrastructure. These models, often comprising billions or even trillions of parameters, necessitate advanced training strategies like data parallelism (DP), model parallelism (MP), and pipeline parallelism (PP), collectively known as 3D parallelism, alongside memory optimization techniques such as gradient checkpointing and CPU offload (e.g., in systems like DeepSpeed's Zero-Offload). While these strategies are crucial for scaling training, their optimal application is highly dependent on the available hardware resources.

In shared GPU clusters, which are now standard for deep learning development, resources are inherently dynamic and often unpredictable from a user's perspective. The prevailing scheduling paradigm, however, treats deep learning jobs as static entities. Users typically define an execution plan and resource requirements (e.g., number of GPUs) at job submission, and these parameters remain fixed throughout the job's lifecycle. This "black box" approach means the scheduler simply allocates the requested fixed resources, queuing jobs if insufficient resources are available.

This static setup creates a significant problem: a constant mismatch between a job's fixed resource requirements and the fluctuating availability of resources in a real-world cluster. For instance, a job requesting 8 GPUs might be stuck in a queue if only 4 are available, even if it could efficiently run on 4 GPUs with a different execution plan. The talk illustrates this with GPT-2 training, showing how different execution plans vary significantly in both resource types and quantities demanded. This oversight severely impacts overall job completion times and cluster utilization.

The fundamental insight behind Rubick is that deep learning jobs are inherently reconfigurable. Their execution plans (e.g., switching from pure data parallelism to a strategy involving gradient checkpointing or CPU offload) and resource allocations can be adjusted dynamically at runtime. For example, a job initially needing 8 GPUs could be reconfigured to use 2 GPUs by employing gradient checkpointing and CPU offload, freeing up resources for other waiting jobs. Exploiting this reconfigurability, traditionally managed by users, presents a valuable opportunity for intelligent cluster schedulers to adapt to dynamic multi-resource availability and significantly improve efficiency.

Key Findings

▶ Watch: Rubick's core functions: inspecting plans and resource allocation (2:40)

Rubick's primary contribution is demonstrating that exploiting job reconfigurability at the cluster scheduling level can lead to substantial improvements in deep learning training efficiency and resource utilization. The key findings include:

• Dynamic Adaptation via White-Box Scheduling: Rubick successfully transitions from traditional "black box" job scheduling to a "white box" approach. By transparently inspecting job launch commands and model structures, it can dynamically determine and adjust both the execution plan and multi-dimensional resource allocations (GPUs, CPU, memory) for deep learning jobs. This allows jobs to adapt to the constantly changing availability of resources in shared clusters.
• Performance Model for Enormous Search Space: The system addresses the challenge of an "enormous search space" for optimal execution plan and resource combinations by introducing an accurate and practical performance model. This model efficiently predicts training performance (throughput) for any given model architecture, training strategy, and resource allocation, significantly reducing the need for costly real training runs for profiling.
• Resource Sensitivity Curve-Guided Scheduling: Rubick introduces the concept of a resource sensitivity curve, which plots the best achievable performance (with the optimal execution plan) against varying resource amounts. This curve simplifies multi-dimensional resource allocation decisions for the scheduler and allows for the identification of jobs most sensitive to resource changes based on the slope of their curve.
• Improved Cluster Throughput and Job Completion Time: Through end-to-end cluster experiments on a 64-GPU cluster, Rubick consistently achieved the shortest average job completion time and makespan compared to state-of-the-art baselines like Sia and Synergy. This improvement is attributed to its ability to schedule more jobs concurrently and utilize resources more flexibly and efficiently.
• Minimal Impact on Training Accuracy: The dynamic reconfiguration process implemented by Rubick, specifically keeping the global batch size unchanged, was verified to have "little impact on training accuracy," ensuring that efficiency gains do not come at the cost of model quality.
• Holistic Optimization: Breakdown studies confirmed that the combination of both execution plan and resource reconfiguration yields the most significant improvements, underscoring the importance of Rubick's joint optimization approach over isolated adjustments.

Technical Deep Dive

▶ Watch: Challenge: Optimal plan and resource combination is complex (3:20)

Rubick's architecture and scheduling policy are built upon two fundamental components: an accurate performance model and a dynamic, slope-comparison-based resource allocation strategy.

At a high level, Rubick takes over the user's traditional role of selecting an execution plan and specifying resources. When a job is submitted, Rubick first inspects the user's launch command to extract the model structure. Based on this, it builds or retrieves a performance model for that specific model type. The system then continuously monitors cluster status and, using the performance models, makes smart decisions about how to allocate multi-dimensional resources (GPUs, CPU, memory) and select appropriate execution plans for both newly submitted jobs and currently running ones. Rubick differentiates between guaranteed jobs, which require an SLO guarantee and consume a specific resource quota, and best-effort jobs, which opportunistically utilize free cluster resources.

Performance Model

The core challenge in dynamic reconfiguration is predicting how a job's performance will change under different execution plans and resource allocations, given the "enormous search space." Rubick addresses this with a lightweight yet accurate performance model. This model estimates the training performance (throughput) by considering the model architecture and the characteristics of various training strategies across each stage of an iteration (forward, backward, communication, optimization).

The model incorporates specific "tricks" to balance accuracy and practicality. For instance, when modeling data parallelism, the total time (Tcc) involves forward, backward computation, and communication. However, backward computation and communication often overlap. To account for this, Rubick uses an overlap function (F_overlap) and a fittable parameter (K_sync) to represent the degree of this overlapping, thus estimating the actual effective time. A similar approach is applied to strategies like Zero-Offload, where CPU offloading overlaps with gradient synchronization and optimizer steps. Here, another overlap function and fittable parameters like K_off and K_swap (implied from the Q&A) are used.

Crucially, the performance model is designed for practicality, requiring only seven fittable parameters in its modeling process. This limits the initial profiling runs to a minimal seven, making the model efficient to build. The model is also capable of online updates, continuously refining its predictions using metrics gathered from real training runs. This adaptability ensures the model remains accurate even as system characteristics or workloads evolve.

Scheduling Policy and Resource Sensitivity Curves

Rubick's scheduling policy aims to provide strong performance guarantees for guaranteed jobs while maximizing overall cluster throughput for best-effort jobs by leveraging job reconfigurability. A central concept is the resource sensitivity curve. For a given model, this curve plots the best achievable training performance (throughput) against different amounts of allocated resources. Each point on the curve represents the optimal execution plan for that specific resource level.

These curves serve several purposes:

1. Reduced Complexity: By pre-determining the best execution plan for each resource level, the scheduler can focus primarily on multi-dimensional resource allocation, significantly reducing the complexity of the decision space.
2. Performance Guarantees: For guaranteed jobs, the curve ensures that Rubick can identify resource allocations that meet user-specified performance baselines, even when using fewer resources by adopting a more efficient execution plan.
3. Prioritization: The slope of the resource sensitivity curve is used to identify which jobs are most sensitive to changes in resource allocation. Jobs with steeper slopes will experience a greater performance gain from additional resources (or a greater loss from resource reduction).

Rubick's dynamic resource reallocation strategy is based on comparing these slopes:

• New Job Allocation: When a new job (e.g., job A) arrives, Rubick first tries to allocate any free resources on a node.
• Shrinking and Reallocating: If more resources are needed or if a more optimal distribution can be achieved, Rubick compares the slopes of all running jobs on the node. It then identifies the job with the lowest slope – meaning the job whose performance is least sensitive to a reduction in resources. This "least sensitive" job (e.g., job B) is then shrunk, freeing up GPUs.
• Guaranteed Minimum: A critical constraint is that even after shrinking, the job must still receive at least the minimum number of resources required to meet its guaranteed performance baseline.
• Reallocation to High-Slope Jobs: The freed resources are then reallocated to jobs with higher slopes, which will benefit more significantly from the additional GPUs, thus maximizing overall cluster throughput.

This continuous adjustment, guided by the performance model and resource sensitivity curves, allows Rubick to dynamically reconfigure both the resources and execution plans of jobs, adapting to the fluid nature of shared deep learning clusters.

Experimental Setup & Results

▶ Watch: Rubick's cluster scheduling workflow overview (5:20)

The Rubick system was rigorously evaluated in a controlled cluster environment and against various baselines to validate its effectiveness.

Experimental Environment

• Cluster Size: Experiments were conducted in a cluster comprising 64 GPUs. The specific GPU models were not explicitly mentioned but are typical for deep learning training.
• Models: The evaluation primarily focused on large models like Llama (for performance model validation and reconfiguration studies) and GPT-2 (used in background for diverse resource demands).
• Reconfiguration: Rubick ensures that the global batch size remains unchanged during reconfiguration, which is crucial for maintaining training accuracy. This was verified to have "little impact on training accuracy."

Baselines for Comparison

Rubick's performance was compared against several existing scheduling approaches:

1. Sia: A baseline scheduler that focuses on accelerating training by adjusting GPU allocation and tuning hyperparameters to quickly reach a target accuracy.
2. Synergy: Another baseline that optimizes CPU and memory allocations across jobs, highlighting the importance of multi-dimensional resource management.
3. Antman: For multi-tenant trace experiments, Rubick was compared against Antman, a scheduler that guarantees the requested resources for jobs, representing a more traditional static allocation approach.

Metrics

The primary metrics used for evaluation included:

• Average Job Completion Time (JCT): A crucial metric for user experience and overall cluster efficiency.
• Makespan: The total time taken to complete all jobs in a given workload, reflecting cluster throughput.
• Performance Model Prediction Quality: Accuracy of predicted throughput versus measured throughput.
• Training Accuracy: To ensure reconfigurations do not degrade model quality.

Headline Results

• Performance Model Accuracy: The performance model demonstrated "very good prediction quality," accurately predicting model throughput across various training settings. Even when continuously decreasing available resources for Llama training, Rubick's model consistently chose the best execution plan, proving its robustness.
• End-to-End Cluster Performance: In comprehensive end-to-end cluster experiments, Rubick consistently achieved the shortest average job completion time and makespan when compared against both Sia and Synergy. This highlights Rubick's superior ability to manage resources and schedules in dynamic environments.
• Impact of Reconfiguration Capabilities: Breakdown studies revealed that the combination of both execution plan and resource reconfiguration brought the "most significant improvement." This validates Rubick's joint optimization approach, confirming that adapting both aspects of a job is more effective than adjusting only one.
• Multi-Tenant Efficiency: In experiments using a multi-tenant trace, Rubick achieved a lower average job completion time compared to Antman. This advantage stems from Rubick's more flexible and efficient use of resources through dynamic reconfiguration, allowing it to adapt to varying demands and resource availability better than a static resource guarantee model.

Practical Implications

▶ Watch: Core scheduling policy: Resource Sensitivity Curves (6:20)

Rubick's innovative approach to deep learning cluster scheduling carries significant practical implications for various stakeholders in the ML ecosystem.

For practitioners and model builders, Rubick offers a path towards faster experimentation and training cycles. By dynamically adapting execution plans and resource allocations, jobs can run more efficiently and complete sooner, even in resource-constrained or highly utilized shared clusters. This means less time spent manually tuning parallelism strategies or waiting in queues for specific resource configurations to become available. The system effectively abstracts away much of the complexity of deep learning system optimization, allowing users to focus more on model development rather than infrastructure intricacies. The guarantee of "little impact on training accuracy" during reconfiguration is also critical, ensuring that performance gains do not compromise model quality.

For infrastructure teams and cluster operators, Rubick promises higher GPU utilization and significantly improved cluster throughput. By intelligently reallocating resources and reconfiguring jobs, the system can schedule more jobs concurrently, reducing overall makespan and maximizing the return on expensive GPU investments. The ability to shrink less sensitive jobs to free up resources for more sensitive ones, while still meeting performance baselines for guaranteed jobs, introduces a new level of flexibility and responsiveness to cluster management. This dynamic balancing act can lead to a more stable and predictable environment for critical workloads, even under fluctuating demand.

However, implementing such a system comes with its own set of tradeoffs and considerations:

• Complexity of Integration: Adopting a "white-box" scheduler like Rubick requires deeper integration with deep learning frameworks and an understanding of their internal execution characteristics. This is a departure from traditional black-box schedulers and necessitates careful design and implementation.
• Performance Model Accuracy and Maintenance: The reliance on an accurate performance model is paramount. While Rubick's model is designed to be practical with few fittable parameters and online update capabilities, maintaining its accuracy across a diverse and evolving landscape of models, frameworks, and hardware is an ongoing challenge. Inaccurate predictions could lead to suboptimal reconfigurations. The Q&A session highlighted this, with a question probing the reliance on online prediction versus simpler heuristics. While heuristics might be easier to implement, Rubick's approach aims for more fine-grained, adaptive optimization that can exploit nuanced interactions between execution plans and resources that heuristics might miss.
• Overhead of Reconfiguration: While the talk states "little impact on training accuracy," dynamic reconfiguration itself introduces some overhead, which needs to be carefully managed to ensure net performance gains.
• Fairness Considerations: The question of fairness in resource allocation for hundreds of competing jobs was raised. Rubick addresses this by explicitly defining guaranteed jobs with SLO guarantees and resource quotas, ensuring their performance. Best-effort jobs are then sorted by GPU utilization, with the implicit understanding that they opportunistically use remaining resources. This tiered approach provides a structured way to manage fairness in a multi-tenant environment.

In summary, Rubick represents a significant step towards more intelligent and adaptive deep learning infrastructure. By embracing job reconfigurability, it offers a powerful mechanism to overcome the limitations of static scheduling, leading to more efficient, higher-performing, and ultimately more productive deep learning clusters.

Key Takeaways

• Deep learning jobs are reconfigurable: Their execution plans and resource allocations can be dynamically adjusted at runtime, a property largely overlooked by traditional static schedulers.
• White-box scheduling is crucial for efficiency: Rubick's ability to transparently inspect job characteristics and adapt them allows for significantly better resource utilization and throughput compared to "black box" approaches.
• Performance modeling is essential for optimization: An accurate yet practical performance model (with few fittable parameters and online updates) is key to navigating the enormous search space of execution plans and resource combinations without extensive profiling.
• Resource sensitivity curves guide intelligent allocation: By characterizing how job performance responds to varying resources, Rubick can prioritize jobs that will benefit most from additional GPUs and efficiently reallocate resources based on slope comparisons.
• Rubick significantly improves cluster metrics: The system consistently achieves shorter average job completion times and makespan, demonstrating its effectiveness in boosting overall cluster throughput and resource efficiency on a 64-GPU cluster.
• Joint optimization of plan and resources is most impactful: Breakdown studies confirm that combining both execution plan and resource reconfiguration yields the most substantial performance improvements, validating Rubick's holistic approach.

About the Speaker(s)

The talk was presented by Xinyi Zhang, representing a joint research effort supported by their supervisor from their university and engineers from Alibaba. While specific individual bios were not detailed, the collaboration highlights a strong foundation in both academic research and industry application. The team's expertise lies at the intersection of deep learning training strategies and cluster scheduling, aiming to innovate how large-scale AI models are trained efficiently in shared environments. The work presented in Rubick reflects a deep understanding of the challenges in managing modern GPU clusters and the potential for intelligent systems to overcome them.