Saturn: Efficient Multi-Large-Model Deep Learning

Fine-tuning large-scale deep learning (DL) models has become a common practice in a variety of domains. It is now common for a developer to download an open-source model spanning tens of billions of parameters from a model hub (e.g. HuggingFace [21]) and fine-tune it on their application-specific data [20]. Unfortunately, models of these sizes tend to impose considerable time & cost burdens, often making them impractical for users in smaller companies & the domain sciences [8, 16].

We identify a few key systems-oriented headaches for users seeking to fine-tune large models: (1) GPU memory remains a bottleneck. Large-memory GPUs are expensive, and even public cloud vendors still ration them. (2) Multi-GPU parallelism is needed, but dozens of parallelism techniques exist [6, 7, 1, 18, 19, 24, 2, 13, 11, 14, 12, 9, 4], and understanding the performance behaviors of these complex approaches can be difficult for DL users. (3) Model selection, i.e. tuning hyper-parameters, model layers, etc., only amplifies the computational load by producing multiple models to train.

These challenges result in a three-part technical problem for practitioners. First, resource allocation [15]. Given a cluster of GPUs, how should they be distributed across jobs? Second, parallelism selection. For each model in a multi-job, which technique should be applied? Third, scheduling. What job ordering will produce the lowest makespan?

These three questions are deeply intertwined. GPU allocations affect the optimal parallelism selection for each job, and constrain schedule orderings. Per-job parallelism selections affect runtimes, thus impacting scheduling & allocation decisions. In this paper, we describe a new approach to tackling this intertwined problem, and demonstrate how taking a novel “joint optimization” approach can significantly improve performance.

Saturn has 3 main modules — the Parallelism Library, the Trial Runner, & the Solver. Figure 1(A) illustrates the overall design.

The Parallelism Library allows users to register and apply various model parallelization techniques with minimal effort. Techniques can be added to the Library by implementing a simple two-function interface, as shown in Figure 1(B). They can then be registered to the library to be reused in different execution sessions (even across different cluster users).

The Trial Runner is our empirical approach to profiling & evaluating user-provided black-boxes (models, parallelisms). Given a set of models, the Trial Runner profiles each model under each possible parallelism under each possible GPU count it could be assigned. Since profiling only requires processing one or two mini-batches, this profiling time tends to be negligible in the context of a larger job. The information from the Trial Runner is critical for our next component, the Solver.

The Solver uses the empirical estimates from the Trial Runner to formulate our joint optimization problem — parallelism selection, resource distribution, and scheduling — as a mixed-integer linear program (MILP). We apply the popular MILP solver Gurobi [5] to produce solutions. The output is used to determine: (1) which parallelism technique from the Library each model should use, (2) how many GPUs each model should get, and (3) in what order & when the models should be submitted to the cluster. We augment the solver with an “introspection” mechanism, adapted from prior art [22, 23]. As models are trained, remaining runtimes per-model will change and shift the workload. So, we re-run the solver on fixed intervals. When a new plan is produced, executing jobs are checkpointed and re-launched under the new scheme.

We evaluate Saturn’s performance against four baselines from standard practice & prior art. The first, “Current Practice”, involves allocating all GPUs per node to one job at a time, with models running in sequence. Task parallelism is applied across nodes. The second baseline, “Random”, randomizes allocations, parallelisms, and schedule orderings. The third baseline ”Optimus” [17], greedily assigns GPUs one-at-a-time based on the estimated marginal runtime improvement. The fourth baseline, Optimus-Dynamic, augments Optimus with our introspection mechanism. Since all approaches optimize on a higher orchestration level, they do not affect correctness/accuracy. We register 4 parallelism techniques in the Library for Saturn: FSDP & DDP from PyTorch Distributed, GPipe, and model offloading from FairScale. Table 1 provides the details of our workloads.

Table LABEL:tb:experiments demonstrates our results from both workloads. We evaluate first on a single 8-GPU node, then on two such nodes. In both cases, Saturn significantly outperforms the baselines. We see speedups of 1.64-1.96X versus the Current Practice baseline, with training time reductions of 39-48%. We find that the allocations Saturn chooses are often unintuitive to a human (e.g. giving 5 GPUs to one model and applying GPipe, but 3 GPUs to another with FSDP. But the overall result is lower costs & runtimes. These results indicate that using Saturn for multi-large-model training may help encourage adoption of large models by time- & and cost- constrained users.