Optimus: Accelerating Large-Scale Multi-Modal LLM Training by Bubble Exploitation

Multimodal LLMs are increasingly important. These models inherit the foundational principles of LLMs, integrating advanced natural language processing techniques while expanding their scope to encompass diverse data modalities. GPT-4[22] represents a prominent example of a multimodal model that extends the capabilities and success of its predecessors to encompass multimodal understanding and generation, demonstrating human-level performance in various benchmark tests with inputs of both images and text.

Multimodal large language model (MLLM) comprises three key parts: one or multiple modality encoders, input projectors, and a large language model backbone[36]. The Modality Encoders are designed to encode inputs from non-textual modalities into respective features, while the input projector aligns features from these modalities with the text feature space. Ultimately, the LLM backbone utilizes aligned features from various modalities and textual features as its input. Figure 1 illustrates the architecture of the MLLM. We exclude the input projector from our discussion due to its relatively minor computational demand compared to the encoder and the LLM (refer to Llava [19]). Additionally, we treat the input projector as the final layer of the modality encoder in our analysis.

Different from homogeneous LLM architecture, multimodal LLM has the following unique characteristics.

Dominant Model Size of LLM Backbone: In multimodal LLMs, the LLM backbone has a significantly larger number of parameters compared to other components such as encoders and projectors. For instance, Flamingo [4] boasts a total of 80 billion parameters, with its LLM backbone alone comprising 70 billion parameters.

Dependency between Encoders and LLM Backbone: In MLLM training, there are two types of data dependencies between encoders and LLM. During the forward pass, encoders must complete the generation of encoded features before the LLM backbone can proceed with forwarding. Conversely, in the backward pass, the LLM backbone calculates gradients before the encoders initiate the backward pass.

Existing LLM pipeline optimizations are not model-agnostic, and fall short in MLLM training tasks. In our internal large-scale MLLM training tasks with ViT encoder and GPT backbone (over 100B parameters), we train Megatron-LM with more than 3,000 NVIDIA GPUs and observe more than 48% GPU cycle idleness when applying multiple SOTA techniques, including MegaScale [14], Zero Bubble Pipeline [24], fine-grained communication-computation overlapping [32]. We analyze the profiled timeline to identify and investigate the occurrences of GPU idleness (i.e., bubbles). Table 1 shows the total time and percentage of average training step time (5.12s) occupied by different types of bubbles.

These bubbles can be classified into three categories based on their underlying causes.

(1) Communication in Data Parallelism (DP). Data parallelism requires communication to aggregate gradients, leading to GPU idle time during the communication. Specifically, MegaScale [14] and Megatron-LM [26] use the distributed optimizer (similar to $P_{os+g}$ in ZeRO [25]) to save memory for large model training, which performs two collective communications (all-gather and reduce-scatter). At the start of each training step, an all-gather operation gathers updated parameters from all data parallel (DP) ranks, resulting in a DP all-gather bubble (occupying 3.3% of the training time). At the end of the training step, reduce-scatter is performed to aggregate gradients, leading to a DP reduce-scatter bubble (occupying 8.9% of the training time). It should be noted that overlapping optimization in data parallelism proposed in Megascale [14] have already been applied and above DP communications are required for the first model chunk which can not be hidden because of the nature of synchronous training  [14]. Figure 2 illustrates DP bubbles that occur due to all-gather and reduce-scatter operations at the start and conclusion of each training step.

(2) Dependency in Pipeline Parallelism (PP). Despite applying pipeline send-receive overlap optimization from Megascale [14], pipeline bubbles still occur due to the inherent data dependencies between stages during the forward and backward passes. It should be noted that Zero Bubble Pipeline cannot eliminate pipeline bubbles in MLLM training, owing to the required changes in the optimizer [24] (refer to discussions in §7). Figure 2 illustrates the MLLM training pipeline schedule, which consists of three phases: warm-up (forward only), steady (one forward and one backward), and cool-down (backward only). Throughout pipeline training, three types of bubbles arise:

1. •

   PP warm-up bubbles occur at all stages except the initial one due to the forward dependency of the first forward pass, averaging 5.0% of the training time.

2. •

   PP cool-down bubbles occur at all stages except the initial one due to the backward dependency of the final backward pass, averaging 9.2% of the training time.

3. •

   Other PP bubbles manifest at all stages except the last one due to dependencies of other forward and backward passes, occupying 8.7% of training time. For instance, PP bubbles emerge immediately after the PP warm-up phase due to the backward dependency of the initial backward pass. Additionally, in cases of imbalanced pipeline stages caused by MLLM’s heterogeneous model, there are additional pipeline bubbles not depicted in Figure 2.

(3) Communications in Tensor Parallelism (TP). Tensor parallelism entails partitioning individual model layers across multiple GPUs, necessitating communication during forward and backward passes to synchronize between GPUs. In Megatron-LM, each forward or backward pass of a transformer layer involves two all-gather and two reduce-scatter kernels [15]. Figure 3 provides a detailed view of CUDA computation and communication kernels during two GPT-175B[6] layer forward passes. In the CUDA communication stream, green kernels represent all-gather communications, while blue kernels denote reduce-scatter communications. The compute stream idles during these communications. Typically, these TP bubbles last for sub-millisecond durations, averaging around 300 $\mu$s. However, during MLLM training, there are thousands of TP bubbles, totaling 11.2% of the training time.

To minimize bubbles in MLLM training, we aim to leverage the distinct dual-component structure of MLLM, which includes encoders and the LLM backbone. We have noted two key observations. Firstly, the majority of bubbles during MLLM training tend to occur during the forward and backward passes of the LLM backbone, with around 90% of these bubbles arising from LLM communication, as indicated in Table 1. Secondly, the encoders require fewer computational operations (FLOPs) than the LLM backbone due to their smaller number of parameters[5, 8, 19, 18, 11].

In response, we propose to schedule encoder computation in LLM bubbles (occurring during communication in LLM) to reduce bubbles throughout the MLLM training process.

We identify three main challenges of scheduling encoder computation to LLM bubbles.

Challenge 1: Only a few GPUs have both encoder and LLM model states. Current training systems [38, 21] use pipeline parallelism to parallelize the MLLM as a single pipeline. Due to the dependency between the encoder and LLM, encoder layers are assigned to earlier pipeline stages, while LLM layers are assigned to later pipeline stages. Consequently, only one pipeline stage typically contains both encoder and LLM layers. To illustrate, Figure 4 demonstrates the application of 3D parallelism (DP=1, PP=4, TP=2) to parallelize MLLM across 8 GPUs, where only 2 GPUs in pipeline stage 1 possess both encoder and LLM model states. The remaining 6 GPUs are incapable of executing encoder computations during LLM bubbles because they lack encoder model states.

Challenge 2: Complex Dependencies in MLLM Training. The intricate dependencies inherent in MLLM training pose significant challenges when scheduling encoder computation within LLM bubbles. Firstly, in synchronous training, the utilization of LLM bubbles is restricted to executing the required encoder computation solely within the current training iteration (iteration dependency). Secondly, the dependency within the encoder pipeline requires scheduling the forward computation of the current encoder pipeline stage $i$ after the completion of the previous encoder stage, and scheduling the backward computation after the subsequent encoder stage concludes. Lastly, the encoder-LLM dependency entails a microbatch-level dependency, where the encoder must complete the forward pass of microbatch $i$ before the LLM pipeline initiates the forward pass of microbatch $i$, and similarly, the encoder can commence the backward pass of microbatch $i$ after the LLM pipeline completes the backward pass of microbatch $i$.

Challenge 3: Sub-millisecond LLM bubbles. Existing frameworks like MegaScale[14] and Megatron-LM[21] typically schedule in the unit of layers. However, bubbles in the LLM exhibit a wide range of durations, spanning from sub-milliseconds (TP bubbles) to hundreds of milliseconds (DP bubbles). For instance, TP bubbles in Figure 3 average around 300$\mu$s across different LLM layers during forward and backward passes and they are too short to complete even a single encoder layer forward or backward. For example, a single ViT-22B layer typically requires around 1.4 milliseconds to complete forward propagation and 2.0 milliseconds to complete backward propagation.

Design decision 1: Colocate encoders and LLM with separate parallelism. To ensure that each GPU possesses both encoder and LLM model states, we propose assigning separate parallel plans to encoders and LLMs across all GPUs. This strategy is illustrated in Figure 5, where using parallel plan (DP=2, PP=2, TP=2) for encoders and (DP=1, PP=4, TP=2) for LLM. Each GPU retains both encoder and LLM model states, and then it becomes feasible for all GPUs to execute encoder computations during LLM bubbles. Note that colocating both the encoder and LLM states may require more GPU memory and we analyze the memory overhead in Section 4.5.

Design decision 2: Dual-Stage Dependency Management. We use two stages to handle complex dependencies in MLLM training: local scheduling and global ordering. Each encoder pipeline undergoes local scheduling, which schedules encoder computations with available LLM bubbles, adhering to the iteration-dependency and encoder-internal dependencies. Global ordering ensures microbatch-level dependency between encoders and LLM by sequencing the encoder’s ending times forward and the encoder’s starting times backward across microbatches. This involves comparing timestamps to verify encoder-LLM dependency compliance. As shown in Figure 6, local scheduling is applied independently to two encoder pipelines, maintaining iteration dependency and encoder-internal dependency. In global ordering, timestamps across all microbatches (totaling 8) are checked to confirm that encoder-LLM dependencies are met.

Design Decision 3: Schedule encoder computation at Kernel Level. Decomposing the encoder layer into kernels enables efficient utilization of sub-millisecond bubbles. However, TP communication kernels in the encoder layer compete for link bandwidth during LLM TP bubbles, causing longer time per iteration. To resolve this, we must additionally schedule encoder communication kernels during LLM compute (see Figure 7).

Optimus is a distributed training system designed for MLLM, enabling the scheduling of encoder computation within LLM bubbles to improve end-to-end training latency. To tackle challenges in Section 3.1, Optimus has two components, which are the model planner and bubble scheduler.

Model Planner. The model planner partitions encoders and the LLM backbone separately to all given GPUs (addressing Challenge 1 in §3.1). Initially, the planner determines the 3D parallelism plan ($\textit{DP}_{llm},\textit{PP}_{llm},\textit{TP}_{llm}$) for the LLM backbone based on insights in Megatron-LM[21]. Subsequently, the planner enumerates potential 3D parallelism plans ($\textit{DP}_{enc},\textit{PP}_{enc},\textit{TP}_{enc}$) for the encoders, considering the available GPU memory after the deployment of the LLM. With the model planner, each GPU holds both LLM and encoder model states, enabling encoder computation during LLM bubbles. The encoder and LLM model parallel plans are provided as input to the bubble scheduler, where Optimus selects parallel plans based on the output schedule with the shortest execution time.

Bubble Scheduler. Bubble scheduler is responsible for scheduling encoder computation into LLM bubbles. Given that the LLM training pipeline divides data into multiple microbatches, the scheduler schedules encoder computations on a per-microbatch basis and satisfies encoder-LLM data dependency at microbatch level (addressing Challenge 2 in §3.1). In addition, the scheduler breaks down encoder computation into kernel granularity, to enable the utilization of sub-millisecond bubbles (TP bubbles) during LLM training (addressing Challenge 3 in §3.1).

Optimus uses the model planner to devise parallel plans for both encoders and LLMs. Subsequently, for each encoder parallel plan, Optimus utilizes the bubble scheduler to generate a schedule and estimate the latency. Ultimately, Optimus selects the schedule with the shortest training time to schedule encoder computation into LLM bubbles. The workflow of Optimus is outlined in Algorithm 1.

Searching separate parallel plans. Initially, the planner determines the 3D parallelism plan ($\textit{DP}_{llm},\textit{PP}_{llm},\textit{TP}_{llm}$) for the LLM backbone based on insights in Megatron-LM[21]. Subsequently, the planner enumerates potential 3D parallelism plans ($\textit{DP}_{enc},\textit{PP}_{enc},\textit{TP}_{enc}$), ensuring that $\textit{PP}_{enc}$ is a factor of $\textit{PP}_{llm}$ and $\textit{TP}_{enc}$ is a factor of $\textit{TP}_{llm}$. In practice, $\textit{PP}_{llm}$ can reach up to 64 and $\textit{TP}_{llm}$ up to 8 for training large language models (LLMs) [21]. Consequently, there are generally no more than 28 encoder parallel plans available, with up to 7 options for $PP_{enc}$ and 4 for $TP_{enc}$.

Colocating encoders and LLM. To guarantee that each GPU can perform encoder computations during LLM downtime, the model planner assigns both encoder and LLM model states to every GPU. As illustrated in Figure 5, all GPUs contain model states for both the encoder (depicted in green) and the LLM (shown in red). Without such colocation, many GPUs would lack the necessary encoder model states to execute encoder computations.

Prune parallel plans based on memory constraint. As we colocate the encoder and LLM stages on GPUs, we calculate the memory requirements for both encoder and LLM states based on the chosen parallelism plan, referencing memory analysis in [15]. Plans that violate GPU memory capacity are immediately pruned.

Constructing separate microbatches. Due to the different parallel plans for encoders and LLMs, there are $m=\frac{\textit{DP}{enc}}{\textit{DP}{llm}}$ times more encoder pipelines than LLM pipelines for a given set of GPUs (e.g. $m=2$ in Figure 5). For GPUs belonging to the same LLM pipeline, there are $m$ encoder pipelines colocated. Depending on the number of microbatches $N_{mb}$ utilized in LLM pipeline training, the data from these $N_{mb}$ microbatches needs to be distributed among these $m$ encoder pipelines, where each encoder pipeline $i$ handles forward and backward computations for $N_{enc_{i}}$ microbatch data. The model planner enumerates possible ways to partition these $N_{mb}$ microbatches among the $m$ encoder pipelines. For instance, if there are 8 microbatches in the LLM training and $m=2$ encoder pipelines, there are a total of 7 possible partitioning options, such as $[1,7]$, $[2,6]$, …, $[7,1]$.

Although LLM bubbles in different GPUs have different start times and duration, there is one common pattern of LLM bubbles as shown in Figure 8. There is one single big bubble (the sum of DP all-gather bubble and PP-warm bubble) before any LLM computation starts and one single big bubble (the sum of PP-cooldown bubble and reduce-scatter bubble) after all LLM computation finishes. And there are many small bubbles (PP bubbles and TP bubbles)[21, 26, 15] interleaved with LLM computation.

Design decision 2: The bubble scheduler, as described in Algorithm 2, initially engages in coarse-grained bubble exploitation by creating initial schedules that incorporate encoder computations within the bubbles positioned before and after LLM computations (line 2). However, it’s possible that these two bubbles may not allow sufficient time to complete all encoder computations, leading to some encoder computations being unscheduled within bubbles. To reduce the total training time, the bubble scheduler then executes fine-grained bubble exploitation. This involves refining the schedule by allocating encoder forward computations to the bubbles that alternate with LLM computations (line 7), followed by assigning encoder backward computations to these same bubbles (line 8). The final output of the bubble scheduler is the schedule that achieves the shortest possible runtime.

Coarse-grained bubble exploitation. For each potential data partitioning approach, the bubble scheduler initializes the schedule by scheduling encoder forward operations to occur before LLM computations and encoder backward operations to occur after LLM computations. Figure 9 illustrates the initialized schedule when there are $m=2$ encoder pipelines and the data partitioning approach is $[3,5]$, i.e., 3 microbatches is allocated to the first encoder pipeline and 5 for the second encoder pipeline.

Fine-grained bubble exploitation. The OptimizeSchedule function (line 15 at Algorithm 2) refines the initial schedule through an iterative approach. Initially, the bubble scheduler employs findCritical to identify the encoder pipeline whose computation is on the critical path of the end-to-end MLLM training (line 17). Subsequently, it utilizes AssignKernels to allocate one microbatch of this encoder computation to bubbles interleaved with LLM computations (line 18). If there are sufficient bubbles available for scheduling encoder computation and encoder-LLM data dependencies are met, the bubble scheduler repeats this process. Otherwise, it returns the current optimized schedule.

When optimizing the schedule for encoder forward computation (line 7 in Algorithm 2), findCritical identifies the encoder pipeline whose forward computation is critical. As shown in the left portion of Figure 10, encoder pipeline 2’s forward computation (microbatch 8 forward) is initially on the critical path in the initial schedule. After successfully scheduling that microbatch forward to later bubbles, encoder pipeline 1 assumes the critical path position. This iterative process leads to a reduction in the end-to-end MLLM training time after each step. Similarly, encoder pipelines whose backward computation is critical are illustrated in the right portion of Figure 10. After each step, the bubble scheduler must verify if it still satisfies the encoder-LLM data dependency before proceeding with the next steps.

When scheduling encoder computation to bubbles interleaved with LLM compute (AssignKernels at line 18), the bubble scheduler decomposes the encoder computation into kernel granularity and schedules these kernels based on the duration of the bubble. For each bubble, the bubble scheduler schedules multiple kernels while ensuring that the total execution time of these kernels is within the bubble duration. Additionally, the bubble scheduler must satisfy the encoder’s internal data dependencies. As illustrated in Figure 11, device 1 holds the first two layers of the encoder, while device 2 holds the next two layers. When scheduling kernels during the forward pass, device 2 can only utilize bubbles that occur after device 1 completes its forward pass to execute encoder computation. For the forward computation, the bubble scheduler schedules encoder computation from upstream encoder pipeline stages to downstream encoder pipeline stages. Conversely, for backward computation, the bubble scheduler schedules encoder computation in the reverse order. While each encoder layer also includes communication kernels, the scheduler ensures that these kernels are not assigned to TP bubbles that occur during LLM communication. Instead, the scheduler identifies long-duration computation kernels within the LLM layers and overlaps them with encoder communication kernels. As the LLM and encoder layers alternately perform computation and communication tasks, they make efficient use of GPU bandwidth and Streaming Multiprocessors (SMs). This design strategy helps to minimize resource contention and improves overall GPU utilization[16].

Complexity. Our bubble scheduling algorithm has low complexity. Given $n$ GPUs and the number of prime factors of $n$ is $n_{p}$, the search space of parallel plans is $C_{n_{p}+1}^{2}$. The number of microbatch partitioning is $O(N_{mb}^{m-1})$. Hence, the complexity for scheduling bubbles is $O(C_{n_{p}+1}^{2}*N_{mb}^{m}*(F+B))$. For our experimented settings, it usually takes around several minutes to calculate the optimal schedule (see §5.3.2), which is also a one-time cost.

The model planner provides different parallel strategies for encoders and LLM backbone, including the number of microbatches, resulting in complex data dependencies both between and within the encoder and LLM. Also, the communication and computation of the encoder and LLM are executed by interleaving, and this may introduce additional pipeline bubbles, if not orchestrated effectively, intensifying the complexity of dependencies in the system.

The bubble scheduler addresses encoder-LLM dependencies at the microbatch level by examining the encoder-LLM forward and backward dependency points for each microbatch $i$. These dependency points, denoted as $F_{i}$ and $B_{i}$ respectively, represent the time when the LLM requires the corresponding activations $A_{i}$ (output by the encoder) for forward propagation, and when the LLM generates the corresponding gradients $G_{i}$ (input for the encoder) during backward propagation. To ensure the satisfaction of encoder-LLM dependencies, the bubble scheduler employs two functions: GetEncLLMDep (line 3 at Algorithm 2) and CheckEncLLMDep (line 19 at Algorithm 2), as described below.

GetEncLLMDep gets encoder-LLM forward and backward dependency points. Given that the interleaved 1F1B schedule [21] stands out as one of the most efficient pipeline schedules for LLM training, we delve into the specifics of the data dependency points $F_{i}$ and $B_{i}$ within this schedule. The top illustration in Figure 12 depicts an instance of the interleaved 1F1B schedule featuring two model chunks. Here, the forward dependency points denote the instances when the first pipeline stage (device 1) commences forward execution for the first model chunk (depicted in dark blue), while the backward dependency points signify the moments when the first pipeline stages (device 1) complete backward execution for the first model chunk (depicted in dark green).
We observe that deferring forward data dependency points for the last four microbatches ($F5$ through $F8$) is feasible without exerting any adverse effects on the overall pipeline latency. To accomplish this, we can adjust the number of warmup microbatches at each pipeline stage, as illustrated in the bottom portion of Figure 12. This adjustment enables the bubble scheduler to leverage bubbles during the phase transition from the warmup phase to the 1F1B-steady phase for scheduling encoder forward computation when optimizing initial schedules. GetEncLLMDep yields the adjusted forward and backward data dependency points for 1F1B interleave schedules.

CheckEncLLMDep verifies the satisfaction of microbatch-level encoder-LLM dependencies. By considering the scheduled encoder computation into bubbles, the bubble scheduler estimates when the encoder finishes the forward pass for microbatches distributed over different encoder pipelines. The bubble scheduler sorts these finishing times in ascending order as $EF_{i}$ (global ordering), representing when the encoder forward operation ends for microbatch $i$ involved in LLM pipeline training. The forward dependency for encoder-LLM is considered met if the encoder completes its forward operation before the specified $F_{i}$ timepoint ($EF_{i}\leq F_{i}$) for all microbatches ($i=1...N_{mb}$). Similarly, the backward dependency is satisfied if the encoder’s backward operation begins no earlier than the $B_{i}$ timepoint ($EB_{i}\geq B_{i}$) for each microbatch ($i=1...N_{mb}$). CheckEncLLMDep returns true when it confirms that both the forward and backward dependencies are successfully met. To illustrate this, Figure 13 provides an example of evaluating encoder-LLM dependency with two encoder pipelines, each handling four microbatches. The order in which the encoder completes its forward pass dictates how the activations are used in the LLM pipeline: activations from encoder pipeline 1 are designated as the 1st, 3rd, 7th, and 8th microbatches, while activations from encoder pipeline 2 are used as the 2nd, 4th, 5th, and 6th microbatches. The bubble scheduler then verifies microbatch-level dependency by ensuring that each encoder’s forward operation concludes before the start of the corresponding LLM forward pass and that each encoder’s backward operation does not commence until after the LLM has ended, for each microbatch.

When dependencies are satisfied, the bubble scheduler integrates necessary peer-to-peer (P2P) communications into the training schedule between the last stage of the encoder pipeline and the first stage of the LLM pipeline. For instance, if encoder pipeline $j$ completes the forward pass for microbatch $i$, the scheduler will insert a P2P send (sending activations) at the last stage of encoder pipeline $j$ and a P2P receive (receiving activations) at the first stage of the LLM pipeline. Similarly, when the LLM pipeline completes the backward pass for microbatch $i$, the scheduler adds a P2P send (sending gradients) at the first stage of the LLM pipeline and a P2P receive (receiving gradients) at the last stage of encoder pipeline $j$. In the scenario depicted in Figure 13, the scheduler inserts 8 pairs of P2P send-receive at devices 1 and 2 to manage the dependencies between encoder pipeline 1 and the LLM pipeline, with 4 pairs allocated for forward dependencies and 4 pairs for backward dependencies. Likewise, an additional 8 pairs of P2P send-receive are inserted at devices 3 and 4 to address the dependencies between encoder pipeline 2 and the LLM pipeline.

To support MLLM with multiple encoders[7, 35], the model planner applies an encoder parallelism plan $(\textit{DP}_{enc},\textit{PP}_{enc},\textit{TP}_{enc})$. independently for each encoder. For pipeline parallelism, layers within each encoder are divided into $\textit{PP}_{enc}$ stages (as illustrated in Figure 14). Each layer of every encoder is then parallelized according to $\textit{TP}_{enc}$. The bubble scheduler breaks down the layers of distinct encoders into kernel-level granularity and arranges their scheduling as if these kernels were part of a single encoder. This is because the encoders within MLLM operate independently, without any data dependencies between them.

When utilizing $n_{gpu}$ GPUs for MLLM training, the model planner requires $\textit{DP}_{enc}$ replicated encoder model states and $\textit{DP}_{llm}$ replicated LLM model states based on parallel plans. Suppose the number of parameters in the encoder is $\phi_{enc}$ and the number of parameters in the LLM is $\phi_{llm}$, with each parameter requiring $k$ bytes of memory. The average GPU memory usage $MEM_{model}$ for storing model states is calculated as follows:

In comparison to existing 3D parallel training solutions, where $DP_{enc}=DP_{llm}$, the estimated memory overhead $MEM_{overhead}$ can be expressed as:

With a larger value of $\textit{DP}_{enc}$, there is a higher memory overhead due to more replicated encoder model states. However, this results in less complex encoder internal dependencies during scheduling (indicated by a smaller $PP_{enc}$). Model planner filters the encoder parallel plans based on the estimated memory usage $MEM_{model}$, ensuring adherence to GPU memory constraints. In practice, the memory overhead typically amounts to less than 12% in our evaluation (§5.3.1) because $\phi_{enc}$ is small (e.g., the largest vision encoder has 22 billion parameters [10]) and $k$ is small (e.g., $k=6$ when using $bf16$ parameters and $fp32$ gradients with distributed optimizer[1]).

Testbed. We conduct our experiments in a production training cluster with thousands of NVIDIA Hopper GPUs. Each GPU has 80GB memory and 989TFLOPS computing performance. The intra-server connection is NVLink and the inter-server connection is a high-bandwidth RDMA network.

MLLM models. We examine the performance of Optimus using various sizes of image encoders and LLM backbones. The image encoders include three sizes: ViT-22B [10], ViT-11B, and ViT-5B, which are scaled-down versions of ViT-22B with smaller hidden sizes. For the language models, we employ two sizes: LLAMA-70B [31] and GPT-175B [6]. Appendix A includes detailed model configurations.

Baselines. We use three open-sourced MLLM training systems with one strawman method as our baselines for comparison.

$\bullet$ PyTorch FSDP [37]: FSDP is a distributed data-parallel training module designed to scale PyTorch models across multiple GPUs with minimal code changes. It shards the model across GPUs, runs $All\_Gather$ to collect all shards from all ranks to recover the full parameter for forward and backward computation, and runs $Reduce\_Scatter$ to synchronize gradients.

$\bullet$ Alpa [38]: Alpa is a compiler system for distributed DL training that automatically generates parallel execution plans covering 3D parallelisms.

$\bullet$ Megatron-LM [21]: Megatron-LM is a state-of-the-art LLM training framework that integrates 3D parallelism techniques. Megatron-LM is designed for symmetric transformer models, and we place multimodal encoders to the preprocess in the first pipeline stage to adapt to MLLM training.

$\bullet$ Megatron-LM balanced: In this strawman method, we balance the layer partitioning among different pipeline stages with an interleaved 1F1B pipeline schedule. Considering the heterogeneity in MLLM submodules, we use a dynamic programming algorithm to assign different layers of submodules to pipeline stages and achieve approximately the same computation amount. The DP algorithm is a simplified version of Alpa’s inter-operator DP algorithm and is included in Appendix B.

We use iteration time and Model Flops Utilization (MFU) [9] as the performance metrics. The reported performance numbers are averaged over 300 training iterations after a warm-up of 10 iterations. The detailed Megatron-LM configurations across experiments are included in Appendix D.