Heterogeneous Parallelism for Multimodal Large Language Model Training

Unknown

Conceptio › Archive › arXiv CS

arXiv CSopen access

Heterogeneous Parallelism for Multimodal Large Language Model Training

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

cloud distributedcomputing parallelcomputing

distributed computing, parallel computing, cloud

2026-5-28

Heterogeneous Parallelism for Multimodal Large Language Model Training

arXiv:2605.27678v1 [cs.LG] 26 May 2026

Yashaswi Karnati, Kamran Jafari, Akash Mehra, Li Ding, Pranav Prashant Thombre, Ali Roshan Ghias, Shifang Xu, Parth Mannan, Yu Yao, Hao Wu, Eric Harper, Ashwath Aithal, Nima Tajbakhsh

Abstract: Foundation model training is becoming multimodal, from post-training pipelines to large-scale pretraining. As modality coverage broadens, context windows grow, and encoder–LLM scales diverge, a single LLM-centric TP/CP/PP/DP/EP layout increasingly limits throughput. This coupling forces encoders to inherit LLM-driven sharding and placement choices that can add communication, limit encoder parallelism, or constrain the LLM schedule; the mismatch is most pronounced at long contexts, where LLM context parallelism is needed for the fused multimodal sequence but encoder inputs remain bounded. We present heterogeneous parallelism for multimodal large language model training, an abstraction that lets modules in one end-to-end graph use independent layouts and rank placements, supporting colocated execution on shared GPUs and non-colocated execution on disjoint rank sets. The key challenge is preserving boundary tensor semantics across independent layouts: forward activations must be materialized for the destination layout, while backward gradients must be routed back to the source layout. We address this with boundary communicators that implement forward and backward layout transforms, plus scheduling extensions for both placement modes. We evaluate optimized homogeneous, colocated heterogeneous, and non-colocated heterogeneous configurations across multimodal workloads and GPU scales to characterize when added layout and placement freedom exposes a better operating point. Across this sweep, colocated heterogeneity improves TFLOPS/GPU by up to 49.3%, while non-colocated heterogeneity improves aggregate token throughput by up to 13.0% and TFLOPS/GPU by up to 9.6%. We validate loss convergence parity against homogeneous baselines and release the system as an open-source Megatron-LM extension.

1. Introduction As foundation-model training expands from text-only transformers to multimodal systems spanning images, audio, video, and documents, the training graph has grown heterogeneous while the parallelism stack has not kept pace. The resulting graph consists of modules with fundamentally different depth, width, and context size — and therefore different requirements for TP, CP, PP, DP, and EP1 — yet in practice all modules are constrained to a single shared layout. While choosing this layout for the LLM2 is natural, encoder and LLM parallelism requirements increasingly diverge at scale. These mismatches motivate heterogeneous multimodal parallelism: layout should be chosen per module rather than as a model-wide constant. First, the high TP degrees required by large LLMs introduce all-reduce collectives around otherwise local encoder computation, and this overhead can outweigh the compute and memory savings for small encoders. Sec1 TP, CP, PP, DP, and EP denote tensor, context, pipeline, data, and expert parallelism, respectively. 2 LLM denotes large language model; MLLM denotes multimodal large language model.

ond, and most consequentially, CP shards the fused language sequence to manage attention memory, but encoder sequence length is bounded by modality resolution rather than the LLM context window. As a result, inheriting LLM CP shards a fixed-size encoder sequence without memory relief, while those ranks could instead act as independent encoder DP replicas and process more modality samples per step. Third, PP confines encoder execution to the first PP stage of LLM, forcing the encoder to inherit that stage’s TP degree and encoder-heavy workloads can overload those ranks and prevent the LLM from using a more efficient pipeline layout. Which mismatch matters most is workload-dependent. Changing encoder and LLM scale, context length, modality-token density, world size, or optimizer-state sharding shifts the balance among compute, activation memory, optimizer-state footprint, communication volume, and pipeline balance. Heterogeneous parallelism therefore does not prescribe one universally better layout; it exposes additional degrees of freedom so the system can tune encoder and LLM layouts to the workload’s operating point. Decoupling module layouts, however, is not just a con-

Heterogeneous Parallelism for Multimodal Large Language Model Training

figuration change: adjacent modules must still agree on boundary tensor semantics. Forward activations must be materialized in the destination layout, and backward gradients must be returned to the source layout. This requirement raises three system challenges. First, every inter-module edge needs a paired forward and backward layout transform. The central correctness challenge is that an incorrect layout mapping can produce tensors with valid shapes, allowing training to proceed while silently routing gradients to the wrong ranks and corrupting optimization without a runtime error. Second, colocated placement requires each rank to interpret its position in two logical grids and apply the correct forward layout transform, with the matching reverse transform in backward. Third, pipeline scheduling must generalize beyond a single linear stage chain: multimodal graphs can contain modules with different PP depths and boundary edges that require communicators distinct from ordinary pipeline P2P. Solving all three together, without changing the model or training semantics, is the core technical problem this paper addresses. We present heterogeneous parallelism for multimodal large language model training as an open-source Megatron-LM extension that lets each module choose its own TP/CP/PP/DP/EP layout and physical rank set while preserving a single end-to-end training graph. Its core primitive is the boundary communicator, an autograd-aware operator attached to an inter-module edge: in forward, it materializes the source activation in the destination module’s layout; in backward, it applies the inverse transform so the gradient returns to the source module’s layout. The same abstraction supports two complementary placement modes. In colocated execution, modules share physical GPUs but interpret those ranks under different logical grids, which is useful when encoder and LLM can share hardware but should not share a layout. In non-colocated execution, modules occupy disjoint rank sets, which is useful when encoder state on the LLM ranks restricts the LLM memory or pipeline layout. Scheduling extensions compose these boundary transforms with pipeline-parallel training. A graph-aware 1F1B3 dispatcher coordinates non-colocated modules, while a three-phase schedule for colocated modules prevents encoder collectives from conflicting with LLM pipeline P2P. Together, boundary communicators, placement modes, and scheduling extensions provide a runtime substrate for module-local parallelism in one end-to-end MLLM training graph. We make the following contributions: 1. A heterogeneous-parallel runtime abstrac31F1B denotes one-forward-one-backward pipeline scheduling; P2P denotes point-to-point communication.

tion for MLLM training. We let modules in one end-to-end multimodal graph choose independent TP/CP/PP/DP/EP layouts and physical rank placements, covering both colocated execution on shared GPUs and non-colocated execution on disjoint rank sets. 2. Boundary communicators and heterogeneous pipeline schedules. We introduce autograd-aware boundary communicators that materialize forward activations in the destination layout and return backward gradients to the source layout. We compose these transforms with pipeline-parallel training using graph-aware 1F1B dispatch for non-colocated modules and a three-phase schedule for colocated modules. 3. An operating-regime evaluation against optimized homogeneous baselines. We evaluate the two modes in the regimes they target: shared-GPU colocated workloads across encoder size, context length, vision-token density, and GPU scale, and workloads that pair larger encoders with a 120B LLM, where noncolocation separates encoder state from LLM ranks. Across these regimes, colocated execution improves TFLOPS/GPU4 by up to 49.3%, while non-colocated execution improves aggregate token throughput by up to 13.0% and TFLOPS/GPU by up to 9.6%. 4. Correctness, convergence, and open-source release. Step-level FP32 parity checks and FP32/BF165 LLaVA-style training sweeps validate correctness and convergence, and we release the implementation in Megatron-LM for community use in large-scale multimodal training.

2. Related work Large-scale training frameworks such as MegatronLM, DeepSpeed/ZeRO, Megatron-DeepSpeed, and PyTorch FSDP provide the core mechanisms for scaling modern transformer training, including tensor, data, pipeline, sequence/context, expert, and optimizer-state parallelism [9, 12–14, 21]. These systems are highly effective for homogeneous transformer stacks, but typically expose a single parallel layout for the composed graph. In multimodal training, this shared-layout assumption can force encoders and LLMs into the same TP/DP/PP/CP choices despite their different sequence lengths, memory pressure, and communication patterns. Several recent systems improve the efficiency of het4 TFLOPS denotes tera floating-point operations per second; TFLOPS/GPU denotes per-GPU throughput. 5 FP32 denotes single-precision 32-bit floating point; BF16 denotes bfloat16.

Heterogeneous Parallelism for Multimodal Large Language Model Training

erogeneous multimodal training graphs without relaxing this shared-layout constraint. Prior work reduces pipeline bubbles or wavefront overheads [3, 6, 16, 18], addresses sequence or workload imbalance [7, 8, 10, 15], and mitigates multimodal dataplane bottlenecks [17, 20]. These techniques improve scheduling, balancing, or data movement around a selected layout, but do not let adjacent modules choose and compose independent parallel layouts inside one end-to-end training graph. More closely related systems explicitly target multimodal heterogeneity, but their abstractions remain tied to particular topologies or placement models. DISTMM [4] partitions heterogeneous multimodal submodules, but focuses on contrastive or fusion-style models where modality towers meet at aggregation, interaction, or loss modules. MLLM training instead has feed-through boundaries: encoder activations become projector, embedding, and LLM inputs, and gradients must return through those same edges. DistTrain [19] disaggregates the modality encoder, LLM backbone, and generator into broker-connected parallelism units, but centers on physically separated units rather than colocated modules that share GPUs while using different logical layouts. Cornstarch exposes modular MLLM construction and modality-aware parallelism [2, 5], but its public interface still restricts independent module layouts, including different encoder/LLM TP sizes, LLM-only CP, and DP inferred after composition. In contrast, our work makes heterogeneous parallelism a first-class runtime abstraction for MLLM training. Boundary communicators pair forward activation materialization with backward gradient transforms, preserving tensor semantics across independently parallelized modules. The same abstraction supports both shared-rank colocated execution and disjointrank non-colocated execution, while heterogeneous pipeline scheduling composes the modules into one end-to-end training graph. By implementing this in Megatron-LM, we provide a reusable runtime substrate for modular MLLM frameworks, disaggregated training systems, and future placement planners.

3. Method We view multimodal training as a graph of modules and graph operators connected by activation edges. As shown in Figure 1(a), the graph contains modality encoders, projection layers, an embedding assembly step, and an autoregressive language model. Rather than assigning one parallel layout to the entire graph, our method assigns each module 𝑚 its own logical layout 𝐺𝑚 = (TP𝑚 , DP𝑚 , PP𝑚 , CP𝑚 , . . .) and phys-

ical rank set 𝑅𝑚 . The logical layout describes how the module shards or replicates its computation; the rank set describes where that module executes. A boundary edge 𝑢 → 𝑣 carries one tensor with two distributed representations. The source module produces it according to 𝐺𝑢 on ranks 𝑅𝑢 , while the destination module consumes it according to 𝐺𝑣 on ranks 𝑅𝑣 . A boundary communicator is the runtime primitive that converts between these representations: forward materializes the activation in the destination layout, and backward applies the matching reverse layout transform so that gradients return to the source layout. This preserves boundary tensor semantics while allowing adjacent modules to use independent TP/DP/PP/CP grids inside one end-to-end training graph. The boundary communicator implements a logical batch-shard transform for every inter-module edge. Equal-DP boundaries pair source and destination shards one-to-one. Fan-in, where DP𝑢 > DP𝑣 , maps multiple source batch intervals to one destination shard; backward splits the destination gradient along the same recorded intervals. Fan-out, where DP𝑢 < DP𝑣 , maps one source shard to multiple destination intervals; backward concatenates sibling destinationgradient intervals to reconstruct the source-layout gradient. How this logical transform is physically realized depends on whether source and destination modules share ranks or occupy disjoint rank sets; we therefore introduce colocated and non-colocated communicators next. 3.1. Non-colocated communicator A non-colocated boundary realizes the layout transform through explicit transfer between the source module’s last pipeline stage and the destination module’s first pipeline stage. The performance-sensitive step is the cross-boundary exchange: module-local TP/CP groups are often mapped within a highbandwidth locality domain, while disjoint module islands may communicate over lower-bandwidth fabric. To avoid a cross-product exchange between source and destination TP/CP ranks, the communicator routes each boundary DP shard through a deterministic boundary leader, usually the rank with canonical TP/CP coordinates on the boundary stage. Leaders perform all inter-module P2P; local TP/CP broadcast then materializes the activation or gradient on the non-leader ranks required by the receiving layout. This design makes cross-boundary communication follow the logical batch-shard relation rather than the product of source and destination TP/CP widths. The source and destination leaders are derived from

Heterogeneous Parallelism for Multimodal Large Language Model Training

(a) multimodal graph

image

enc DP0

proj

audio

proj

video

proj

(b) colocated ranks

embed

LLM

enc DP1

TP0

TP1

TP0

TP1

TP0

TP1

TP2

TP3

encoder Enc DP0

LLM PP=2 NC

PP0 TP0

P2P

PP1 TP0

bcast Enc DP1

PP0 TP1

PP1 TP1

LLM DP0

text

Enc: TP=2, DP=2; LLM: TP=4, DP=1

Enc: TP=1, DP=2; LLM: TP=2, PP=2

Figure 1 | Multimodal graph and heterogeneous placement modes. Modality encoders produce projected tokens that are inserted into the LLM input sequence. Colocated execution reinterprets the same physical ranks under different encoder and LLM logical grids, while non-colocated execution places encoder and LLM on disjoint rank sets and uses non-colocated communicators before ordinary LLM pipeline P2P. 𝐺𝑢 , 𝐺𝑣 , 𝑅𝑢 , and 𝑅𝑣 , so both sides construct the same static routing pattern for the edge. Equal-DP boundaries use one-to-one P2P between paired source and destination leaders. Fan-in uses a many-to-one route: source leaders whose batch intervals form one destination shard send to the same destination leader, which concatenates in batch order. Fan-out uses the dual one-to-many route: a source leader partitions its shard into destination intervals and sends each interval to the corresponding destination leader. Backward uses the inverse leader routes. For equalDP, the destination leader returns the gradient shard to the paired source leader. For fan-in, the destination leader splits the gradient along the recorded batch intervals and returns each slice to the source leader that supplied the corresponding forward shard. For fan-out, the source leader receives sibling destinationgradient intervals and concatenates them to reconstruct the source-layout gradient. In all cases, the reconstructed leader result is broadcast within the local TP/CP group required by the receiving side. Figure 2 visualizes these leader routes for unequal TP widths.

ery rank’s source coordinate, destination coordinate, batch interval, and local collective group from 𝐺𝑢 , 𝐺𝑣 , and the shared rank set. Equal-DP boundaries are rank-local: each destination rank consumes the source shard already present on that physical rank, reinterpreted under the destination layout. Fan-in uses local all-gather among ranks whose source intervals form one destination shard, followed in backward by interval selection from the destination-gradient shard. Fan-out is the dual: forward selects the destination batch interval from the larger source shard locally, while backward all-gathers sibling destinationgradient intervals to reconstruct the original sourcelayout gradient. Figure 3 visualizes these shared-rank transforms. Unlike the non-colocated protocol, colocated communication does not move tensors across rank sets; it changes the logical view of the shared rank set and invokes local collectives only when the layout transform spans multiple batch intervals. We profile these communicators and find that their overhead is less than 1%; Appendix A.4 provides details. Section 3.3 describes how this shared-rank layout operator is ordered with the LLM pipeline schedule.

3.2. Colocated communicator A colocated boundary realizes the same layout transform on a shared rank set. The source and destination modules use the same physical ranks but assign each rank different logical coordinates under 𝐺𝑢 and 𝐺𝑣 . The communicator is an autograd-aware sharedrank layout operator: its forward pass constructs the destination-layout representation on the shared ranks, and its backward pass reconstructs the corresponding source-layout gradient representation. For each boundary edge, the runtime derives ev-

3.3. Pipeline orchestration over heterogeneous modules The boundary communicators define the layout transform for one edge; the scheduler decides when those edge transforms are invoked inside the full training loop. This is not ordinary linear pipeline parallelism. A conventional 1F1B schedule assumes one linear PP chain, where each stage communicates with a fixed predecessor and successor. In a heterogeneous multimodal graph, modules can have different PP depths and layouts, a node can have multiple incoming or 4

Heterogeneous Parallelism for Multimodal Large Language Model Training

(a) Fan-in: Enc DP=4, TP=2 → LLM DP=2, TP=4

Encoder DP leaders forward: P2P + concat + broadcast TP0 leader

DP0

B/4

TP1 member

B/2

B/4

TP0 leader

DP1

concat

(b) Fan-out: Enc DP=2, TP=4 → LLM DP=4, TP=2

LLM DP leaders

Encoder DP leaders

LLM DP leaders

forward: split + P2P + broadcast

TP0 leader TP1

TP2

DP0

TP3

B/2

TP0 leader TP1

TP2

split

TP0 leader

B/4

TP1 member B/4

TP3

TP0 leader

TP1 member

backward: split + P2P + broadcast TP0 leader

DP2

B/4

TP1 member

TP0 leader

DP3

split

B/2

B/4

DP0

DP1

backward: P2P + concat + broadcast

TP0 leader TP1

TP2

DP1

TP3

TP0 leader TP1

TP2

B/2

concat

TP0 leader

B/4

TP1 member B/4

TP3

TP0 leader

TP1 member

backward P2P

forward P2P

DP2

DP3

local TP/CP broadcast

Figure 2 | Non-colocated boundary communicator. Source and destination modules occupy disjoint rank sets, so boundary tensors cross the module boundary only through deterministic boundary leaders. Solid arrows show inter-module P2P among leaders; dashed arrows show local TP/CP broadcast to non-leader ranks. 𝐵 denotes the global boundary batch. In fan-in, two 𝐵/4 source shards are concatenated into one 𝐵/2 destination shard in forward, and the 𝐵/2 gradient is split back into two 𝐵/4 source gradients in backward. Fan-out is the dual case: one 𝐵/2 source shard is split into two 𝐵/4 destination shards in forward, and the returned gradients are concatenated in backward. Unequal TP widths change the local broadcast domains around each leader, not the number of cross-boundary rank pairs. (a) Fan-in: Enc DP=4, TP=2 → LLM DP=2, TP=4 Encoder view TP=2, DP=4

D0 {r0,r2}

r0: D0/T0

(b) Fan-out: Enc DP=2, TP=4 → LLM DP=4, TP=2

LLM view

Encoder view

TP=4, DP=2 forward all-gather by TP lane B/4

B/2

r1: D0/T1

r0: D0/T0

r1: D0/T1

r1: D0/T1 D0

{r1,r3}

r2: D1/T0

B/4

B/2

r3: D1/T1

B/4

B/2

r5: D2/T1

r2: D0/T2

r3: D0/T3

r4: D1/T0

r7: D3/T1

B/4

B/2

forward interval select B/2

B/4

r0: D0/T0 r1: D0/T1

r5: D1/T1

B/2

B/4

r2: D1/T0 r3: D1/T1

backward all-gather by dest TP lane

r6: D3/T0

TP=2, DP=4

r2: D0/T2

backward interval select r4: D2/T0

LLM view

TP=4, DP=2

r6: D1/T2

{r4,r6} D1 {r5,r7}

r7: D1/T3

r6: D1/T2

B/2

backward

r4: D2/T0 r5: D2/T1

B/2

r7: D1/T3

forward

B/4

r6: D3/T0 r7: D3/T1

all-gather pair

Figure 3 | Colocated boundary communicator. Source and destination modules share the same physical ranks but interpret them under different logical layouts. 𝐵 denotes the global boundary batch. In fan-in, ranks whose source slices form one destination shard all-gather locally: two 𝐵/4 source slices become one 𝐵/2 destination shard in forward, and backward selects from the 𝐵/2 gradient the batch interval contributed by each source rank. In fan-out, forward is rank-local batch-interval selection from a 𝐵/2 source shard to a 𝐵/4 destination shard, while backward all-gathers sibling 𝐵/4 gradients to reconstruct the original 𝐵/2 source-gradient shard. Thus, colocated communication changes the logical view of a shared rank set rather than moving tensors across rank sets.

Heterogeneous Parallelism for Multimodal Large Language Model Training

(a) non-colocated graph-aware PP

(b) colocated three-phase schedule phase 1

E1 PP=2

E1 PP0 R0

P2P

E1 PP1

enc. fwd

NC join

LLM PP0

LLM PP1

E2 PP0

P2P

phase 2

phase 3

coloc. fwd

coloc. bwd grad

detach

LLM PP2

LLM PP

warm

E2 PP=1

steady 1F1B

cool

detached microbatch views

LLM PP=3

enc. bwd

CPU

D2H

H2D

sync time

Figure 4 | Pipeline orchestration for heterogeneous modules. (a) Non-colocated graph-aware pipeline dispatch: Encoder 1 has two PP stages, Encoder 2 has one, and the LLM has three; intra-module edges use module-local P2P, while encoder–LLM edges invoke edge-local non-colocated communicators. (b) Colocated LLM-PP training uses three phases: encoder forward and colocated forward transform, detached LLM 1F1B, then boundary-gradient handoff followed by colocated backward transform and encoder backward. outgoing activation edges, and a rank’s next communication may be either module-local P2P or a boundary communicator that changes layouts across modules. A pipeline action is therefore an edge-dispatch decision over a PP-stage graph, not simply “send to the next PP rank.” For non-colocated execution, the runtime expands the module graph into PP-stage nodes, as in Figure 4(a). At each schedule action and microbatch, it resolves the active module edge. Internal edges use the module’s ordinary PP communicator; edges crossing disjoint module rank sets invoke the non-colocated communicator from Section 3.1 between the source module’s final PP stage and the destination module’s first PP stage. Preserving edge identity is essential at joins and branches: an LLM join becomes runnable only after all required encoder activations have been materialized in the LLM layout, and backward routes each boundary gradient through the communicator attached to the forward edge that produced that activation. The scheduler also generalizes warmup, steady state, and cooldown. In a linear pipeline, these phases are derived from a stage’s position in one chain. In a PP-stage graph, different nodes can have different downstream distances to the sink, and these distances determine when each node warms up, enters paired forward/backward actions, and drains. At steady state, the two halves of a paired schedule call may resolve to different edge types: a rank can send a new activation through one edge while receiving an older gradient through the reverse of another. Appendix D expands Figure 4(a) into the corresponding dispatch table. Colocated execution introduces a different ordering constraint. The encoder and LLM share physical

ranks, but their communication order is not the same. Encoder execution may require collectives over the encoder layout, where all ranks in a collective group must enter together. LLM 1F1B deliberately staggers ranks across PP stages and microbatches. If encoder collectives are invoked inside the LLM pipeline loop, ranks in the same encoder collective group can arrive at different times; adding global synchronization would avoid the conflict but serialize the LLM pipeline. We therefore use the three-phase colocated schedule in Figure 4(b). First, the runtime batches the samples assigned to the colocated rank group for the current LLM pipeline window, runs encoder forward once over those microbatches, and applies the colocated forward transform from Section 3.2 to produce a packed boundary tensor in the LLM PP0 layout. Second, this tensor is detached as a leaf boundary variable and exposed to LLM PP0 as per-microbatch views; ordinary LLM 1F1B then indexes the active view, inserts it into that microbatch’s text embeddings, and proceeds with standard LLM pipeline P2P without re-entering encoder collectives. Third, after the LLM pipeline drains, accumulated gradients on the packed boundary tensor are handed through the colocated backward transform, and encoder backward runs over the saved batched encoder activations. The explicit gradient handoff reconnects the detached LLM phase to the encoder-side autograd graph while preserving the paired forward/backward layout transform. The phase boundary also creates a memory window: encoder parameters can be offloaded after encoder forward, reloaded during LLM cooldown, and synchronized before encoder backward consumes them.

Heterogeneous Parallelism for Multimodal Large Language Model Training

4. Experimental results We evaluate how independent encoder and LLM layouts improve end-to-end MLLM training throughput. A homogeneous shared TP/CP/PP/DP grid imposes one set of layout choices on both encoder and LLM, which can over-constrain the system. Heterogeneous parallelism relaxes this constraint by allowing each module to use its own layout. This flexibility matters because the optimization space is large: encoder scale, context length, vision-token density, world size, and optimizer-state sharding each shift pressure among compute, memory, pipeline balance, and communication. As these pressures shift, the best operating point can also change, including whether heterogeneity should be colocated or non-colocated. Our evaluation samples representative points across this landscape to characterize where each execution mode is most effective. For each workload, we tune homogeneous (H), colocated heterogeneous (C), and non-colocated heterogeneous (NC) execution to their best valid throughput configurations.6 The search includes module layouts, recomputation, batch-size constraints, distributed optimizer instances, and LLM pipeline-layer balance. Colocated runs use the same physical world size and global batch size as homogeneous training, whereas non-colocated runs are iso-GBS7 ; when exact worldsize matching is infeasible, we keep the LLM world size comparable to homogeneous training, add a compact encoder island, and report both aggregate tokens/s and TFLOPS/GPU. Appendix B shows how these module layouts and placements are expressed in the user-facing configuration. We first evaluate colocated heterogeneous execution, which is the relevant setting when the encoder and LLM can fit on the same GPUs but should not share the same logical grid. In this regime, heterogeneous parallelism improves throughput by remapping the encoder and LLM to different layouts. The encoder can reduce unnecessary TP, avoid unnecessary CP, increase useful DP, and choose its own number of distributed optimizer instances. This last choice is important because increasing encoder DP can improve encoder efficiency, but it also changes how optimizer state is sharded and how much optimizer communication is introduced. Colocated heterogeneity lets the encoder make this tradeoff independently of the LLM. Figure 5(a) summarizes this shared-GPU regime at 8K context, with full measurements in Table A2. Across 6 Throughput is measured as the median over 15 steady-state iterations after warmup; repeated measurements on representative configurations showed low run-to-run variation, typically below 1%. 7 GBS denotes global batch size.

the short-context model-size sweep, colocated execution consistently improves throughput when the shared homogeneous grid over-constrains the encoder or the LLM pipeline. The tuned 256-GPU rows show the same mechanism at larger scale, where encoder layout and encoder optimizer instances are selected per workload. As an example, for the 1B encoder + 14B LLM at 8K, 50% vision tokens, and 72 GPUs, the best homogeneous configuration uses a shared 4/1/1/18 TP/CP/PP/DP layout. The best colocated heterogeneous configuration instead uses encoder 2/1/1/36 and LLM 2/1/2/18. This lets the LLM use PP=2 with lower TP, while the encoder remains a singlestage, higher-DP module rather than being tied to the LLM pipeline structure. Step time improves at the same GPU count and essentially the same peak memory usage. As context length grows, the LLM needs CP to fit and parallelize the fused multimodal sequence, but the encoder sequence length is fixed by image resolution and patching; colocated heterogeneity keeps encoder CP=1 and converts those ranks into encoder DP replicas. This is the strongest colocated setting in our sweep (Figure 5(b) and Table A3). We next evaluate non-colocated heterogeneous execution, which targets a different bottleneck. For larger encoders paired with the 120B LLM, the active constraint shifts from logical layout coupling to memory residency. Encoder parameters, activations, and optimizer state now compete with the LLM for the same rank memory and can block a better LLM pipeline layout. In this regime, non-colocated heterogeneity becomes the better operating point: the encoder runs on a compact island, while the LLM recovers an independently optimized layout over the remaining ranks. Figure 5(c) and (d) show how the best operating point for a 120B LLM changes between colocated and non-colocated modes, with full measurements in Tables A4 and A5. For 1B and 3B encoders, colocated execution remains the better operating point because the encoder can share the LLM GPUs and benefit from layout remapping. Starting at the 6B encoder scale, the best mode switches to non-colocated execution. The NC rows remain positive despite averaging over the additional 8-GPU encoder island. This indicates that, once encoder residency constrains the LLM memory or pipeline layout, separating the encoder can improve both end-to-end throughput and normalized per-GPU efficiency. Overall, the experiments show that heterogeneous parallelism improves MLLM throughput by making

Heterogeneous Parallelism for Multimodal Large Language Model Training

(a) Model size: Colocated vs Homogeneous 50% vision

75% vision

Gain (%)

14B LLM

70B LLM

1B+7B

TFLOPS/GPU gain (%)

8K context, TFLOPS/GPU gain over homogeneous

(b) Long context: Colocated vs Homogeneous

40 30 20 10 0 16K

Tokens/s gain (%)

+13.0

+5.9

+12.7 +7.3

+7.1 +4.4

7.2B 50% vision

64K

12B 75% vision

128K

16K

32K

64K

128K

LLM context length 50% vision 75% vision

(d) 120B LLM: Best Heterogeneous Mode Tokens/s gain over H* (%)

120B LLM, 8K, 256/264 GPUs

32K

120B LLM

1B+14B

8K, 50% vision 30 +21.5

+12.9

+7.1

+5.9

+7.3

6B *

7.2B

12B

Figure 5 | Operating regimes for heterogeneous training. When memory headroom is sufficient, colocated execution is the right first step: it gives the encoder and LLM separate logical layouts while keeping them on the same GPUs. Long context strengthens this effect because the LLM needs context parallelism while the encoder does not; keeping encoder CP = 1 turns LLM context-parallel ranks into encoder data-parallel replicas. As encoder pressure grows with a 120B LLM, non-colocated execution becomes the better operating point: it separates encoder work from the LLM rank set, and the best heterogeneous mode switches from colocated for small encoders to non-colocated from the 6B encoder scale onward. Appendices A.1, A.2, and A.3 report the full measurements. module-level layout and placement tunable. Colocated execution extracts more throughput from a fixed GPU budget when encoder and LLM can share hardware but need different logical grids. Non-colocated execution becomes effective when encoder residency should be separated from the LLM rank set. This flexibility lets the system select the right operating point for each workload.

5. Convergence validation Heterogeneous parallelism should change only how tensors are partitioned, placed, and communicated across ranks, not the underlying training semantics. Because training loss curves can hide subtle communicator-level errors, we first validated this invariant with step-level numerical parity tests in FP32. Starting from the same initialization and identical inputs, the homogeneous baseline, a non-colocated variant, and a colocated heterogeneous variant execute identical training steps; after each optimizer update,

we compare all parameters, gradients, and optimizer states within tight numerical tolerance. This directly checks boundary-communicator correctness, including forward activation materialization and backward gradient routing. For convergence parity, we compare loss trajectories across colocated and non-colocated layouts using a LLaVA-1.5-style workload: a CLIP ViT-L/14 vision encoder [11], a two-layer projection MLP, and Vicuna7B [1] trained on LLaVA-Pretrain. LLaVA serves as a representative validation case, but our boundary communication mechanism also supports crossattention-based MLLM such as Qwen3VL. We use FP32 (vision encoder frozen, projection MLP and LLM trainable), disable TF32, and enforce deterministic PyTorch/cuDNN execution and fixed NCCL collective selection to minimize numerical noise. Because TP reductions remain layout-order-dependent, we compare within numerical tolerance rather than requiring bitwise agreement. Appendix C provides the full validation protocol and extends the study

Heterogeneous Parallelism for Multimodal Large Language Model Training

LLM training. FP32/BF16 validation confirms training behavior. The open-source Megatron-LM extension already supports multiple encoders, audio/video, and MoE; future work will evaluate their performance regimes, including how extra encoders and module boundaries shift layout and placement choices.

Colocated homogeneous (baseline) Non-colocated homogeneous Colocated heterogeneous

LM loss

6 5 4 3

abs. diff.

2 0.10

max |non-colo. homo|=0.034 max |colo. hetero|=0.057

0.05 0 0

100

Training step

Figure 6 | Controlled FP32 convergence comparison. Baseline: colocated homogeneous on four GPUs with the vision encoder and LLM both using TP = 4, DP = 1; non-colocated same-layout keeps TP = 4, DP = 1 for both modules but places the LLM on ranks 0– 3 and the vision encoder on ranks 4–7. Colocated heterogeneous keeps the same four ranks but uses vision TP = 1, DP = 4 and LLM TP = 4, DP = 1; both variants follow the baseline loss curve. to longer BF16 runs and broader TP/PP/DP layout sweeps. Figure 6 reports the main controlled convergence comparison. Starting from the homogeneous baseline, the same-layout non-colocated variant changes physical placement while keeping module layouts fixed, and the colocated heterogeneous variant changes modulelocal layout while keeping placement fixed. All other training settings are held constant. The loss curves closely overlap, supporting convergence parity under both placement and layout changes.

6. Conclusion We presented heterogeneous parallelism for multimodal training, a runtime abstraction that lets modules in one end-to-end graph use independent layouts and rank placements while boundary communicators and scheduling extensions preserve activation and gradient semantics across colocated and non-colocated pipeline execution. This flexibility lets the training system tune its operating point as model scale, context length, modality-token density, memory pressure, communication volume, and pipeline balance change. In our evaluated workloads, colocated execution is effective when the same GPU allocation can support different module grids, whereas non-colocated execution becomes useful when encoder state on the LLM ranks restricts the LLM memory or pipeline layout. Across these workloads, this yields up to 49.3% TFLOPS/GPU for colocated execution and 13.0% aggregate token throughput for large-encoder 120B 9

Heterogeneous Parallelism for Multimodal Large Language Model Training

References [1] W.-L. Chiang, Z. Li, Z. Lin, Y. Sheng, Z. Wu, H. Zhang, L. Zheng, S. Zhuang, Y. Zhuang, J. E. Gonzalez, I. Stoica, and E. P. Xing. Vicuna: An open-source chatbot impressing GPT-4 with 90%* ChatGPT quality. https://lmsys.org/blog/2023-03-30-vicuna/, March 2023. [2] Cornstarch Project. Using Cornstarch 5d parallelism. https://cornstarch-org.github.io/ parallelization/cornstarch_parallel/, 2025. Accessed: 2026-05-06. [3] W. Feng, Y. Chen, S. Wang, Y. Peng, H. Lin, and M. Yu. Optimus: Accelerating large-scale multi-modal LLM training by bubble exploitation. In 2025 USENIX Annual Technical Conference (USENIX ATC 25). USENIX Association, 2025. URL https://www.usenix.org/conference/atc25/ presentation/feng. [4] J. Huang, Z. Zhang, S. Zheng, F. Qin, and Y. Wang. DISTMM: Accelerating distributed multimodal model training. In 21st USENIX Symposium on Networked Systems Design and Implementation (NSDI 24), pages 1157–1171. USENIX Association, 2024. URL https://www.usenix.org/conference/ nsdi24/presentation/huang. [5] I. Jang, R. Lu, N. Bansal, A. Chen, and M. Chowdhury. Efficient distributed MLLM training with Cornstarch. arXiv preprint arXiv:2503.11367, 2025. [6] B. Jeon, M. Wu, S. Cao, S. Kim, S. Park, N. Aggarwal, C. Unger, D. Arfeen, P. Liao, X. Miao, M. Alizadeh, G. R. Ganger, T. Chen, and Z. Jia. GraphPipe: Improving performance and scalability of DNN training with graph pipeline parallelism. arXiv preprint arXiv:2406.17145, 2024. [7] C. Jiang, Z. Cai, Y. Tian, Z. Jia, Y. Wang, and C. Wu. DCP: Addressing input dynamism in long-context training via dynamic context parallelism. In ACM SIGOPS 31st Symposium on Operating Systems Principles (SOSP 25), 2025. doi: 10.1145/3731569.3764849. [8] H. Li, F. Fu, S. Lin, H. Ge, X. Wang, J. Niu, J. Xue, Y. Tao, D. Wang, J. Jiang, and B. Cui. Hydraulis: Balancing large transformer model training via co-designing parallel strategies and data assignment. arXiv preprint arXiv:2412.07894, 2024.

[11] A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, G. Krueger, and I. Sutskever. Learning transferable visual models from natural language supervision. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 8748–8763. PMLR, 2021. URL https: //proceedings.mlr.press/v139/radford21a.html. [12] S. Rajbhandari, J. Rasley, O. Ruwase, and Y. He. ZeRO: Memory optimizations toward training trillion parameter models. arXiv preprint arXiv:1910.02054, 2020. [13] M. Shoeybi et al. Megatron-LM training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019. [14] S. Smith et al. Using DeepSpeed and Megatron to train Megatron-Turing NLG 530B, a large-scale generative language model. arXiv preprint arXiv:2201.11990, 2022. [15] Y. Wang, S. Wang, S. Zhu, F. Fu, X. Liu, X. Xiao, H. Li, J. Li, F. Wu, and B. Cui. FlexSP: Accelerating large language model training via flexible sequence parallelism. arXiv preprint arXiv:2412.01523, 2024. [16] Y. Wang, S. Zhu, F. Fu, X. Miao, J. Zhang, J. Zhu, F. Hong, Y. Li, and B. Cui. Efficient multi-task large model training via data heterogeneity-aware model management. Proceedings of the VLDB Endowment, 18(1), 2025. [17] B. Xiao, Y. Zheng, L. Shi, X. Li, F. Wu, T. Li, X. Xiao, Y. Zhang, Y. Wang, and S. Liu. OrchMLLM: Orchestrate multimodal data with batch post-balancing to accelerate multimodal large language model training. arXiv preprint arXiv:2503.23830, 2025. [18] Z. Xue, H. Hu, X. Chen, Y. Jiang, Y. Song, Z. Mi, Y. Zhu, D. Jiang, Y. Xia, and H. Chen. PipeWeaver: Addressing data dynamicity in large multimodal model training with dynamic interleaved pipeline. In Proceedings of the 31st ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS), 2026. doi: 10.1145/3779212.3790154.

[9] D. Narayanan et al. Efficient large-scale language model training on GPU clusters using Megatron-LM. arXiv preprint arXiv:2104.04473, 2021.

[19] Z. Zhang, Y. Zhong, Y. Jiang, H. Hu, J. Sun, Z. Ge, Y. Zhu, D. Jiang, and X. Jin. DistTrain: Addressing model and data heterogeneity with disaggregated training for multimodal large language models. arXiv preprint arXiv:2408.04275, 2024.

[10] Y. Niu, H. Xiao, D. Liu, W. Zhou, and J. Li. DHP: Efficient scaling of MLLM training with dynamic hybrid parallelism. arXiv preprint arXiv:2602.21788, 2026.

[20] J. Zhao, Q. Lu, W. Jia, B. Wan, L. Zuo, J. Feng, J. Jiang, Y. Chen, S. Cao, J. He, K. Jiang, Y. Hu, S. Nong, Y. Peng, H. Lin, and C. Wu. MegaScale-Data: Scaling dataloader for multisource 10

Heterogeneous Parallelism for Multimodal Large Language Model Training

large foundation model training. arXiv preprint arXiv:2504.09844, 2025. [21] Y. Zhao et al. PyTorch FSDP: Experiences on scaling fully sharded data parallel. arXiv preprint arXiv:2304.11277, 2023.

Heterogeneous Parallelism for Multimodal Large Language Model Training

A. Appendix This appendix reports the layouts, peak memory, step times, tokens/s gains, and TFLOPS/GPU gains supporting Section 4. All parallelism columns use TP/CP/PP/DP. Encoder rows use CP = 1 and PP = 1 unless otherwise stated. Metric columns report percentage gains over homogeneous, and step-time columns report absolute milliseconds. Memory columns report peak max-rank memory in GB. Unless stated otherwise, LLM recompute is off. For the encoder-size sweeps, 1B encoder rows run without encoder recompute; larger encoder rows use full encoder recompute. Notation and acronyms Table A1 | Notation and acronyms used in the throughput tables. Symbol

Meaning

H C NC H-E / H-L H Layout C-E / C-L NC-E / NC-L TP / CP / PP / DP TP/CP/PP/DP SP Seq. Len. Vis% GBS NMB Tokens/s TFLOPS/GPU Peak Mem (GB) Step (ms) F+B (ms) Enc. Rec. / LLM Rec. Enc. opt. inst. W5/S15

Homogeneous colocated execution Heterogeneous colocated execution Heterogeneous non-colocated execution Encoder / LLM layout in an H run Shared encoder / LLM layout in an H run Encoder / LLM layout in a C run Encoder / LLM layout in an NC run Tensor / context / pipeline / data parallelism Layout tuple order used in all tables Sequence parallelism LLM sequence length Vision-token percentage Global batch size Number of microbatches Aggregate training throughput in tokens/s Per-GPU tera floating-point operations per second Peak max-rank GPU memory in GB Iteration step time in milliseconds Forward plus backward time in milliseconds Encoder / LLM recomputation setting Number of encoder distributed optimizer instances Five warmup steps followed by fifteen measured steps

Compute infrastructure. The empirical results reported in this paper were measured on a GPU cluster with NVIDIA H100 80GB GPUs. Cluster nodes contain eight GPUs; the throughput tables report world size directly in GPUs. Within a node, GPUs are connected with fourth-generation NVLink, providing 900 GB/s bidirectional GPU-to-GPU NVLink bandwidth with 18 NVLink connections per GPU. Across nodes, each node uses eight 400 Gb/s ConnectX-7 InfiniBand cards. This hardware setup is used for all empirical measurements reported in the main text and for the tokens/s gain, TFLOPS/GPU gain, and step-time measurements reported below.

Heterogeneous Parallelism for Multimodal Large Language Model Training

A.1. Colocated results Table A2 | Colocated 8K gain, memory, and step-time results. GBS is 288 for the 72-GPU rows and 256 for the tuned 256-GPU 70B/120B rows. Model

Vis% GPUs

H-E H-L C-E C-L Enc. opt. inst.* TFLOPS/GPU Gain Peak Mem (GB) TP/CP/PP/DP TP/CP/PP/DP TP/CP/PP/DP TP/CP/PP/DP Δ% H/C

Step (ms) H/C

1B+14B 1B+14B 1B+14B 3B+14B 3B+14B 3B+14B

50 75 87.5 50 75 87.5

72 72 72 72 72 72

4/1/1/18 4/1/1/18 4/1/1/18 4/1/1/18 4/1/1/18 4/1/1/18

2/1/1/36 1/1/1/72 1/1/1/72 2/1/1/36 1/1/1/72 1/1/1/72

2/1/2/18 2/1/2/18 2/1/2/18 2/1/2/18 2/1/2/18 2/1/2/18

1 1 1 1 1 1

+12.1 +13.5 +12.7 +12.9 +12.7 +14.0

57.2 / 57.1 57.2 / 57.6 57.2 / 57.8 59.7 / 57.3 59.7 / 57.7 59.7 / 57.9

7,570 / 6,750 7,780 / 6,840 7,660 / 6,800 7,850 / 6,970 8,020 / 7,140 8,130 / 7,150

1B+32B 3B+32B 3B+32B 3B+32B

75 50 75 87.5

72 72 72 72

4/1/1/18 4/1/1/18 4/1/1/18 4/1/1/18

1/1/1/72 1/1/1/72 1/1/1/72 1/1/1/72

4/1/2/9 4/1/2/9 4/1/2/9 4/1/2/9

1 1 1 1

on par on par +11.1 +13.1

66.2 / 55.2 68.7 / 54.2 68.7 / 55.3 68.7 / 55.8

14,210 / 14,380 14,330 / 14,750 14,370 / 12,950 14,630 / 12,910

1B+70B 1B+70B 3B+70B 3B+70B 1B+120B 1B+120B 3B+120B 3B+120B 6B+120B 7.2B+120B 12B+120B

50 75 50 75 50 75 50 75 50 50 50

256 256 256 256 256 256 256 256 256 256 256

8/1/1/32 8/1/1/32 8/1/1/32 8/1/1/32 8/1/1/8 8/1/1/8 8/1/1/8 8/1/1/8 8/1/1/8 8/1/1/8 8/1/1/8

8/1/1/32 8/1/1/32 8/1/1/32 8/1/1/32 8/1/4/8 8/1/4/8 8/1/4/8 8/1/4/8 8/1/4/8 8/1/4/8 8/1/4/8

1/1/1/256 1/1/1/256 1/1/1/256 1/1/1/256 1/1/1/256 1/1/1/256 2/1/1/128 2/1/1/128 2/1/1/128 4/1/1/64 4/1/1/64

8/1/2/16 8/1/2/16 8/1/2/16 8/1/2/16 8/1/4/8 8/1/4/8 8/1/4/8 8/1/4/8 8/1/4/8 8/1/4/8 8/1/4/8

4 4 4 4 8 8 2 2 1 1 1

+16.5 +15.1 +13.9 +12.7 +16.6 +13.9 +21.5 +21.6 −0.7 −0.3 −8.0

61.5 / 47.2 61.5 / 47.8 62.7 / 48.8 62.7 / 49.4 57.2 / 54.0 63.2 / 55.4 60.5 / 66.2 60.5 / 67.6 53.6 / 68.0 53.6 / 64.4 66.9 / 69.6

8,441 / 7,244 8,415 / 7,312 8,471 / 7,434 8,452 / 7,501 15,258 / 13,081 15,169 / 13,319 15,093 / 12,420 15,303 / 12,579 17,890 / 18,020 17,980 / 18,040 15,860 / 18,190

The 8K colocated sweep shows where module-local layout freedom is useful on the same GPUs. The 14B LLM rows gain 12.1%–14.0% TFLOPS/GPU; for 32B, gains concentrate at higher vision-token share. The tuned 256-GPU rows show a second colocated operating point: reducing encoder TP, adding LLM pipeline parallelism where useful, and tuning encoder optimizer instances gives +12.7%–21.6% for the audited 1B/3B encoder rows with 70B and 120B LLMs. For larger 120B encoder rows, C converges back toward H or slows down, motivating the NC crossover summarized in Table A5. Table A3 | Colocated long-context gain, memory, and step-time results. Encoder CP remains 1 under heterogeneous parallelism, while homogeneous training forces the encoder to inherit LLM CP. F+B is forward plus backward time in ms. Because H and C use the same GPU count, tokens/s and TFLOPS/GPU gains are identical; Gain reports both. Model 1B+7B 1B+7B 1B+7B 1B+7B 1B+7B 1B+7B 1B+7B 1B+7B 1B+14B 1B+14B 1B+14B 1B+14B 1B+14B 1B+14B 1B+14B 1B+14B

Seq. Len. Vis% GPUs GBS NMB 16K 16K 32K 32K 64K 64K 128K 128K 16K 16K 32K 32K 64K 64K 128K 128K

50 75 50 75 50 75 50 75 50 75 50 75 50 75 50 75

16 16 16 16 32 32 64 64 16 16 32 32 32 32 64 64

8 8 4 4 4 4 4 4 4 4 4 4 4 4 4 4

2 2 2 2 2 2 2 2 1 1 1 1 2 2 2 2

H Layout C-E C-L Gain Peak Mem (GB) F+B (ms) TP/CP/PP/DP TP/CP/PP/DP TP/CP/PP/DP Δ% H/C H/C 2/2/1/4 2/2/1/4 2/4/1/2 2/4/1/2 4/8/1/1 4/8/1/1 4/16/1/1 4/16/1/1 4/1/1/4 8/1/1/2 4/2/1/4 8/2/1/2 8/4/1/1 8/4/1/1 4/8/2/1 4/8/2/1

1/1/1/16 1/1/1/16 1/1/1/16 1/1/1/16 1/1/1/32 1/1/1/32 1/1/1/64 1/1/1/64 2/1/1/8 2/1/1/8 2/1/1/16 2/1/1/16 2/1/1/16 2/1/1/16 2/1/1/32 2/1/1/32

2/2/1/4 2/2/1/4 2/4/1/2 2/4/1/2 2/8/1/2 2/8/1/2 2/16/1/2 2/16/1/2 4/1/1/4 4/1/1/4 4/2/1/4 4/2/1/4 4/4/1/2 4/4/1/2 4/8/1/2 4/8/1/2

+6.6 +8.5 +11.7 +20.3 +34.3 +49.3 +17.8 +24.8 +0.1 +16.3 +2.9 +21.3 +23.9 +30.3 +29.8 +34.0

54.5 / 50.2 59.6 / 52.2 65.3 / 50.7 75.6 / 52.7 54.0 / 48.4 68.9 / 50.3 25.1 / 48.6 26.6 / 50.4 75.1 / 73.8 50.4 / 76.3 76.6 / 67.8 53.5 / 70.3 61.2 / 69.1 73.2 / 71.6 55.4 / 63.0 55.8 / 62.9

1,300 / 1,221 1,352 / 1,248 1,836 / 1,644 1,980 / 1,640 3,966 / 2,937 4,457 / 2,973 6,321 / 5,381 6,614 / 5,291 1,202 / 1,205 1,438 / 1,211 1,628 / 1,583 1,962 / 1,585 5,687 / 4,579 5,848 / 4,479 9,397 / 7,184 9,692 / 7,175

Long context gives the clearest case for C: the LLM needs context parallelism as sequence length grows, but the encoder sequence length is fixed by image resolution rather than LLM context. H makes the encoder inherit LLM CP, while C keeps the encoder at CP=1 and converts those ranks into encoder DP replicas. Gains therefore grow with context, reaching 34.3%–49.3% for 1B+7B at 64K and 29.8%–34.0% for 1B+14B at 128K.

∗ Enc. opt. inst. reports the number of encoder distributed optimizer instances when swept; default is 1.

Heterogeneous Parallelism for Multimodal Large Language Model Training

A.2. Non-colocated results Table A4 | Non-colocated 120B LLM gain, memory, and step-time results. Rows compare homogeneous shared-grid and non-colocated (NC) execution with a dedicated encoder island. H Layout is the shared encoder/LLM layout for the H run. Rows are grouped by encoder size. All rows use sequence parallelism (SP) on and LLM recompute off; encoder full recompute is enabled. Model 6B+120B 6B+120B 6B+120B 6B+120B 6B+120B 6B+120B 7.2B+120B 7.2B+120B 7.2B+120B 7.2B+120B 7.2B+120B 12B+120B 12B+120B 12B+120B 12B+120B 12B+120B

Seq. Len. Vis% GPUs GBS 8K 8K 8K 16K 16K 16K 8K 8K 16K 16K 16K 8K 8K 8K 16K 16K

50 75 50 50 75 50 50 75 50 75 50 50 75 87.5 50 50

256/264 256/264 512/520 256/264 256/264 512/520 256/264 256/264 256/264 256/264 512/520 256/264 256/264 256/264 256/264 512/520

256 256 512 256 256 512 256 256 256 256 512 256 256 256 256 512

H Layout NC-E NC-L Tokens/s Gain TFLOPS/GPU Gain Peak Mem (GB) TP/CP/PP/DP TP/CP/PP/DP TP/CP/PP/DP Δ% Δ% H / NC 8/1/4/8 8/1/4/8 8/1/4/16 8/2/4/4 8/2/4/4 8/2/4/8 8/1/4/8 8/1/2/16 8/2/4/4 8/2/4/4 8/2/4/8 8/1/4/8 8/1/4/8 8/1/4/8 8/2/4/4 8/2/4/8

1/1/1/8 1/1/1/8 1/1/1/8 1/1/1/8 2/1/1/4 1/1/1/8 1/1/1/8 1/1/1/8 1/1/1/8 2/1/1/4 1/1/1/8 2/1/1/4 2/1/1/4 2/1/1/4 2/1/1/4 2/1/1/4

8/1/2/16 8/1/2/16 8/1/2/32 8/2/4/4 8/2/4/4 8/2/2/16 8/1/2/16 8/1/2/16 8/2/4/4 8/2/4/4 8/2/2/16 8/1/2/16 8/1/2/16 8/1/2/16 8/2/4/4 8/2/2/16

+5.9 +4.4 +4.0 +0.7 +3.3 +5.1 +7.1 +13.0 +5.5 +3.6 +4.9 +7.3 +12.7 +9.5 +5.6 +2.5

+2.7 +1.3 +2.4 −2.3 +0.2 +3.5 +3.8 +9.6 +2.3 +0.4 +3.3 +4.1 +9.3 +6.2 +2.4 +0.9

Step (ms) H / NC 17,700 / 16,710 17,600 / 16,860 17,870 / 17,180 38,380 / 38,110 38,120 / 36,900 38,400 / 36,550 17,980 / 16,790 16,210 / 14,340 38,260 / 36,260 38,380 / 37,070 38,450 / 36,650 15,860 / 14,780 16,690 / 14,810 17,020 / 15,540 38,430 / 36,400 38,370 / 37,430

64.8 / 76.6 53.6 / 76.8 50.4 / 73.5 53.6 / 53.6 53.6 / 63.8 50.4 / 74.6 53.6 / 76.6 77.2 / 73.3 53.6 / 63.2 53.6 / 63.8 50.4 / 74.6 66.9 / 73.1 67.0 / 73.3 67.0 / 73.4 56.3 / 63.2 52.6 / 74.6

The 120B NC rows show physical separation becoming useful once encoder state constrains the LLM layout. At 8K and 50% vision tokens, NC is positive from 6B onward and sharpens for the 7.2B and 12B encoders; at 16K, the 512/520-GPU rows show the scale-out benefit of preserving the LLM island. Because NC averages over the extra 8-GPU encoder island, TFLOPS/GPU gains are conservative relative to the aggregate tokens/s gains.

+4.0%, +2.4%

250 200 150 100

+5.9%, +2.7%

+4.4%, +1.3%

50%

75%

50%†

tokens/s gain, TFLOPS/GPU gain

(a) 6B+120B 8K scale-out

512/520-GPU scale-out at 16K Throughput (K tokens/s)

Throughput (K tokens/s)

6B scale-out at 8K

6B encoder + 120B LLM, H/NC GPUs

120B LLM, 50% vision

230

+5.1%, +3.5%

+4.9%, +3.3% +2.5%, +0.9%

220

210

6B H

7.2B NC

12B

tokens/s gain, TFLOPS/GPU gain

(b) 120B LLM 16K scale-out

Figure 7 | Non-colocated 120B scale-out results. (a) Non-colocation remains positive for the 6B+120B case when expanding from 256/264 to 512/520 GPUs. (b) At 512/520 GPUs, NC preserves the PP2 LLM layout and gives +5.1%, +4.9%, and +2.5% tokens/s gains for 6B, 7.2B, and 12B encoders. Labels report tokens/s and TFLOPS/GPU gains over H. A.3. Operating ranges for different modes Table A5 | H/C/NC operating-range comparison at 8K. C and NC gains are relative to H for the same workload; C vs NC is the gap of C relative to NC. Each gain cell reports tokens/s / TFLOPS/GPU. When C converges to H, we report C=H because H is a valid special case of colocated heterogeneous execution. Small-encoder rows use 72 total GPUs for all modes, with NC splitting 64 LLM GPUs plus an 8-GPU encoder island. The 70B/120B rows use 256 GPUs for H/C and 264 GPUs for NC. The 70B/120B C rows report the best tuned colocated configuration, including encoder offload where available.

Heterogeneous Parallelism for Multimodal Large Language Model Training

Model

Vis%

GPUs H / C / NC

1B+14B 1B+14B 3B+14B 3B+14B 1B+32B 3B+32B 3B+32B 1B+70B 3B+70B 1B+120B 3B+120B 6B+120B 7.2B+120B 12B+120B

50 87.5 50 87.5 75 50 87.5 50 50 50 50 50 50 50

72 / 72 / 72 72 / 72 / 72 72 / 72 / 72 72 / 72 / 72 72 / 72 / 72 72 / 72 / 72 72 / 72 / 72 256 / 256 / 264 256 / 256 / 264 256 / 256 / 264 256 / 256 / 264 256 / 256 / 264 256 / 256 / 264 256 / 256 / 264

C-E NC-E Enc. opt. inst.* C Gain NC Gain C vs NC Best TP/CP/PP/DP TP/CP/PP/DP tokens/s / TFLOPS/GPU tokens/s / TFLOPS/GPU tokens/s / TFLOPS/GPU 2/1/1/36 1/1/1/72 2/1/1/36 1/1/1/72 1/1/1/72 1/1/1/72 1/1/1/72 1/1/1/256 1/1/1/256 1/1/1/256 2/1/1/128 H layout H layout H layout

1/1/1/8 1/1/1/8 1/1/1/8 1/1/1/8 1/1/1/8 1/1/1/8 1/1/1/8 1/1/1/8 1/1/1/8 1/1/1/8 1/1/1/8 1/1/1/8 1/1/1/8 2/1/1/4

+12.1 / +12.1 +12.7 / +12.7 +12.9 / +12.9 +14.0 / +14.0 on par / -0.9 on par / -0.1 +13.1 / +13.1 +16.5 / +16.5 +13.9 / +13.9 +16.6 / +16.6 +21.5 / +21.5 C=H C=H C=H

1 1 1 1 1 1 1 4 4 8 2 1 1 1

-11.3 / -11.3 -10.5 / -10.5 -8.6 / -8.6 -5.4 / -5.4 -10.2 / -10.2 -10.1 / -10.1 -7.8 / -7.8 +7.3 / +4.0 +9.3 / +6.0 -3.0 / -0.5 -9.9 / -7.7 +5.9 / +2.7 +7.1 / +3.8 +7.3 / +4.1

+26.4 / +26.3 +25.9 / +25.8 +23.5 / +23.5 +20.5 / +20.6 +11.4 / +10.4 +11.2 / +11.1 +22.7 / +22.7 +8.6 / +12.0 +4.2 / +7.5 +20.2 / +17.2 +34.9 / +31.6 -5.6 / -2.6 -6.6 / -3.7 -6.8 / -3.9

C C C C H/C H/C C C C C C NC NC NC

Table A5 shows the operating ranges of H, C, and NC. When the combined model has enough memory headroom, C is the best use of the fixed GPU budget: it remaps the encoder to a module-local layout with more data parallelism while keeping the LLM colocated. That larger encoder-DP group increases cross-node optimizer communication, so the tuned C rows use multiple encoder optimizer instances where memory allows. Because encoder parameters can be offloaded, the added optimizer-state replication is affordable and improves throughput. As LLM and encoder memory pressure increase, this colocated flexibility diminishes: C converges back to H because H is a valid special case of C, and both modules still share the same physical ranks. NC becomes preferable once the encoder and LLM should stop carrying each other’s memory and scheduling pressure. At the 120B LLM scale, the verified 1B/3B rows still favor C, while the separation point appears at the 6B encoder scale, or equivalent encoder-side compute and memory from multiple encoders: from 6B+120B onward, the dedicated encoder island unlocks the PP2 LLM layout and becomes the preferred operating point.

(a) Bridge communication vs. vision tokens 0.62%

0.6

0.56%

0.46%

0.4 0.23%

0.2

0.08%

0.15%

0.0

25%

50% Vision tokens

Bridge payload / step (GB)

Bridge communication / step time (%)

A.4. Bridge communication

3.22/3.22

2.15/2.15

4 1.07/1.07

0.54/0.27

Non-colocated

1.07/0.54

1.61/0.81

25%

75% Colocated

(b) Fwd/Bwd payload vs. vision tokens

Forward

50% Vision tokens Backward

75%

Figure 8 | Bridge communication profile for the 3B encoder + 70B LLM payload sweep at 8K sequence length. Panel (a) reports median bridge communication time as a percentage of measured step time. Panel (b) reports local bridge payload per step as F/B, where F is forward payload and B is backward payload, both in GB. In the 3B+70B colocated setup, the encoder uses DP = 32 and the LLM uses DP = 16, so the fan-in scale is 2. Each colocated bridge group is therefore a 2-rank all-gather, not a global DP32 or DP16 collective. This small-group fan-in explains why colocated forward bridge time is low. NC forward is larger because it includes both cross-island activation transfer and LLM-side activation fanout: the encoder island sends activations to the LLM receiver ranks, and those activations are then broadcast within the relevant LLM-side group. NC backward is simpler in this setup because the encoder island uses TP = 1; there is no encoder-side TP broadcast, so the backward bridge is the LLM-to-encoder gradient transport.

∗ Enc. opt. inst. reports the number of encoder distributed optimizer instances when swept; default is 1.

Heterogeneous Parallelism for Multimodal Large Language Model Training

B. Open-source user configuration example Non-colocated: disjoint GPU sets

Colocated: shared GPU set

... parallelism = MegatronMIMOParallelismConfig( module_parallelisms={ "language": ModuleParallelismConfig( tensor_model_parallel_size=2, pipeline_model_parallel_size=2, data_parallel_size=1, rank_offset=0, # ranks [0, 4) ), "images": ModuleParallelismConfig( tensor_model_parallel_size=1, pipeline_model_parallel_size=1, data_parallel_size=4, rank_offset=4, # ranks [4, 8) ), }, ) ...

... parallelism = MegatronMIMOParallelismConfig( module_parallelisms={ "language": ModuleParallelismConfig( tensor_model_parallel_size=4, pipeline_model_parallel_size=1, data_parallel_size=2, rank_offset=0, # ranks [0, 8) ), "images": ModuleParallelismConfig( tensor_model_parallel_size=1, pipeline_model_parallel_size=1, data_parallel_size=8, rank_offset=0, # ranks [0, 8) ), }, ) ...

Figure 9 | Minimal user-facing configuration for heterogeneous parallelism in multimodal training. The left example places the language model and vision encoder on disjoint rank ranges, so the runtime uses the non-colocated communicator. The right example places both modules on the same physical rank range, so the runtime uses colocated execution while still allowing different module-local TP/PP/DP layouts. The open-source implementation exposes this layout choice through a compact configuration interface. Users provide one entry per module, specifying its TP/PP/DP layout and rank offset; together, these fields determine the module’s logical process groups and physical rank range. At model construction time, the runtime uses the resulting ranges to select colocated or non-colocated boundary communicators, while keeping the model definition, loss, and training loop shared across both placement modes.

C. LLaVA validation details and additional convergence results C.1. Validation setup This appendix gives the convergence-validation details that complement Section 5. Appendix C.1 describes the LLaVA setup and deterministic FP32 validation protocol, Appendix C.2 reports the additional frozen-LLM and BF16 convergence results, and Appendices C.3 and C.4 give the trainable-LLM non-colocated and colocated layout sweeps. The validation model follows LLaVA-1.5. We use a CLIP ViT-L/14 vision encoder, a two-layer projection MLP, and a Vicuna-7B v1.5 language model based on Llama-2-7B. The vision encoder and language model are initialized from pretrained CLIP ViT-L/14 and Vicuna-7B checkpoints, respectively. Across all validation runs, the vision encoder is frozen and the projection MLP is trainable; the LLM is either frozen for projector-only runs or trainable for trainable-LLM runs. The projection MLP initializes the first linear√layer weights from 𝒩 (0, 0.02) with zero bias; the second linear layer weights are initialized from 𝒩 (0, 0.02/ 2), also with zero bias. The vision tower has 24 layers and 16 attention heads, processes 336 × 336 images with 14 × 14 patches, and contributes 576 patch features of width 1024 after dropping the CLS token. The projection MLP maps these features to the 4096-dimensional language-model hidden space. Vicuna-7B v1.5 uses 32 decoder layers, 32 attention heads, hidden size 4096, RoPE base 10,000, and a 4096-token context length. For TP = 4, Megatron-LM pads the Vicuna vocabulary plus the reserved LLaVA image token to 32,256 entries. For the FP32 parity runs, we further reduce numerical nondeterminism so loss curves can be compared more directly across layouts. We switch these runs to FP32 precision, disable TF32, and enable PyTorch deterministic execution with torch.use_deterministic_algorithms(True). We also disable cuDNN autotuning with cudnn.benchmark=False and set cudnn.deterministic=True for deterministic cuDNN kernels. For NCCL reductions, we pin the algorithm to Ring so collective reductions use a deterministic order across runs. These settings reduce loss-curve variation caused by floating-point non-associativity, making discrepancies more 16

Heterogeneous Parallelism for Multimodal Large Language Model Training

likely to reflect communicator errors rather than runtime nondeterminism. Table A6 | Optimization hyperparameters for the LLaVA convergence validation. Parameter

Value

Global batch size Optimizer Learning rate Min learning rate Schedule / warmup Weight decay 𝛽1 /𝛽2 Gradient clipping

96 Distributed Adam 1.0e-3 (1.0e-4 for trainable LLM) 2.0e-5 (1.0e-5 for trainable LLM) Cosine decay / 60 steps 0.0 0.9 / 0.95 0.0 (1.0 for BF16 runs)

C.2. Additional convergence validation results Section 5 and Appendices C.3–C.4 focus on FP32 trainable-LLM runs. Here we report the remaining convergence validation results: the FP32 projector-only frozen-LLM placement/layout comparison, the BF16 placement/layout comparisons, and the frozen-LLM layout sweeps. The shorter validation runs use 100 training steps in FP32 precision, whereas the larger-scale comparisons run for 1K training steps in mixed-precision BF16. Runs within each experiment set use the same initialization, data order, and optimization setup, with hyperparameters summarized in Table A6. The frozen-LLM layout sweep uses the configurations in Table A7. Figure 10 shows the FP32 projector-only placement/layout comparison, Figure 11 reports the corresponding BF16 placement/layout comparisons, and Figure 12 summarizes the frozen-LLM layout sweeps. Table A7 | Non-colocated parallelism sweep with frozen Table A8 | Non-colocated parallelism sweep with trainlanguage model. All configurations use 8 GPUs. able language model. All configurations use 8 GPUs. Config name tp4_both tp2_dp2_both tp2_pp2_llm_tp4_vision tp1_dp4_both pp4_llm_tp4_vision pp4_llm_dp4_vision pp2_dp2_llm_tp2_dp2 tp4_llm_dp4_vision tp2_pp2_llm_tp2_dp2_vis tp2_dp2_llm_dp4_vision

LLM Vision (TP/PP/DP) (TP/PP/DP) 4/1/1 2/1/2 2/2/1 1/1/4 1/4/1 1/4/1 1/2/2 4/1/1 2/2/1 2/1/2

4/1/1 2/1/2 4/1/1 1/1/4 4/1/1 1/1/4 2/1/2 1/1/4 2/1/2 1/1/4

Config name tp4_both tp2_dp2_both tp2_pp2_llm_tp4_vision tp2_pp2_llm_tp2_dp2_vis pp4_llm_tp4_vision pp4_llm_tp2dp2_vision pp2_dp2_llm_tp2dp2_vis tp4_llm_tp2dp2_vision

LLM Vision (TP/PP/DP) (TP/PP/DP) 4/1/1 2/1/2 2/2/1 2/2/1 1/4/1 1/4/1 1/2/2 4/1/1

4/1/1 2/1/2 4/1/1 2/1/2 4/1/1 2/1/2 2/1/2 2/1/2

C.3. Non-colocated layout sweep The second experiment set sweeps TP/PP/DP choices for non-colocated layouts across the language model and vision encoder. These configurations preserve the same end-to-end LLaVA training task while changing the ownership layout presented to the boundary communicator. Each non-colocated sweep configuration assigns four GPUs to the LLM and four GPUs to the vision encoder; the listed TP/PP/DP triples show how each module partitions its own four-rank group. Table A8 lists the trainable-LLM sweep, which uses micro batch size (MBS) 2. Figure 13(a) shows that the loss curves remain aligned across these non-colocated layouts when the LLM and projection MLP are trainable. C.4. Colocated layout sweep The third experiment set validates the colocated communicator under the same trainable-LLM layout patterns tested in the non-colocated sweep. We reuse the configurations from Table A8, but run the vision encoder and language model on the same eight physical ranks. Because each module now spans the shared 8-GPU rank set rather than its own 4-GPU rank set, we keep each module’s TP and PP degrees unchanged and double its DP degree. Thus, for example, a non-colocated configuration with TP = 4, PP = 1, DP = 1 for a module becomes TP = 4, PP = 1, DP = 2 in the colocated sweep. 17

Heterogeneous Parallelism for Multimodal Large Language Model Training

Colocated homogeneous (baseline) Non-colocated homogeneous Colocated heterogeneous

7 6

LM loss

5 4 3

abs. diff.

2 0.10

max |non-colo. homo|=0.078 max |colo. hetero|=0.087

0.05 0 0

100

Training step Figure 10 | Frozen-LLM FP32 placement and layout convergence comparison. The vision encoder and LLM are frozen, so only the projection MLP is trained. The modified runs follow the corresponding colocated homogeneous baseline. 8

Colocated homogeneous (baseline) Non-colocated homogeneous Colocated heterogeneous

LM loss

abs. diff.

2 0.20

max |non-colo. homo|=0.134 max |colo. hetero|=0.100

0.10 0 0

250

500

750

1000

2 0.40

abs. diff.

LM loss

Colocated homogeneous (baseline) Non-colocated homogeneous Colocated heterogeneous

max |non-colo. homo|=0.226 max |colo. hetero|=0.246

0.20 0 0

250

Training step

(a) Frozen LLM, BF16.

500

750

1000

Training step

(b) Trainable LLM, BF16.

Figure 11 | BF16 placement and layout convergence comparisons. Panel (a) uses projector-only training with the vision encoder and LLM frozen; panel (b) keeps the vision encoder frozen while training the LLM and projection MLP. In both cases, the modified runs follow the corresponding colocated homogeneous baseline. Figure 13(b) shows that the colocated loss curves remain tightly aligned across the tested trainable-LLM configurations, indicating that colocated rank reinterpretation and the corresponding forward/backward ownership transforms preserve the LLaVA training signal across the tested TP/PP/DP layouts.

Heterogeneous Parallelism for Multimodal Large Language Model Training

pp2_dp2_llm_tp2_dp2

tp4_llm_dp4

pp4_llm_dp4_vision

tp2_dp2_both

pp4_llm_tp4_vision

tp2_dp2_llm_dp4_vision

tp4_both

tp2_pp2_llm_tp2_dp2_vision

7 6

tp2_pp2_llm_tp4_vision

LM loss

tp1_dp4_both

4 3

tp1_dp8_both

tp2_dp4_both

tp2_dp4_llm_dp8_vision

pp2_dp4_llm_tp2_dp4_vision

tp4_dp2_llm_dp8_vision

tp2_pp2_dp2_llm_tp2_dp4_vision

tp4_dp2_both

pp4_dp2_llm_dp8_vision

pp4_dp2_llm_tp4_dp2_vision

tp2_pp2_dp2_llm_tp4_dp2_vision

5 4 3

max range=0.094

max range=0.077

2 0

100

Training step

100

Training step

(a) Non-colocated frozen-LLM layout sweep.

(b) Colocated frozen-LLM layout sweep.

Figure 12 | Frozen-LLM configuration sweeps. Panel (a) sweeps the non-colocated configurations from Table A7; panel (b) applies the corresponding colocated layouts by sharing the physical ranks and doubling each module’s DP degree. The overlapping curves show that convergence remains insensitive to the tested parallelism layout when only the projection MLP is trained. 8

8 tp2_pp2_llm_tp2_dp2_vision

pp2_dp2_llm_tp2dp2_vision

pp4_llm_tp2dp2_vision

tp4_llm_tp2dp2_vision

pp4_llm_tp4_vision

tp4_both

tp2_pp2_llm_tp4_vision

LM loss

tp2_dp2_both

4 3

tp4_dp2_llm_tp2_dp4_vision

pp2_dp4_llm_tp2_dp4_vision

pp4_dp2_llm_tp4_dp2_vision

tp4_dp2_both

pp4_dp2_llm_tp2_dp4_vision

tp2_pp2_dp2_llm_tp2_dp4_vision

tp2_dp4_both

tp2_pp2_dp2_llm_tp4_dp2_vision

4 3

max range=0.199

max range=0.172

2 0

Training step

(a) Non-colocated layout sweep.

100

Training step

(b) Colocated layout sweep.

Figure 13 | Trainable-LLM FP32 configuration sweeps. The vision encoder is frozen, while the LLM and projection MLP are trained. Panel (a) sweeps non-colocated TP/PP/DP configurations from Table A8; panel (b) sweeps the corresponding colocated configurations by sharing the physical ranks and doubling each module’s DP degree. The overlapping trends indicate that convergence is insensitive to the tested parallelism layout. The annotations report the maximum step-wise loss range across all plotted configurations.

D. Pipeline schedule details Section 3.3 describes pipeline orchestration at the abstraction level. This appendix gives the execution details for both placement modes. In both cases, the boundary communicator defines the layout transform for an edge, while the schedule determines when that edge is active and which ranks must participate. For non-colocated execution, the scheduler extends ordinary 1F1B by resolving each pipeline action to an edge in the expanded module graph. Figure 14 expands the topology in Figure 4(a), where Encoder 1 has two PP stages, Encoder 2 has one PP stage, and the LLM has three PP stages. Rows correspond to the communication calls issued during warmup, steady state, and cooldown; columns correspond to rank roles in the expanded graph. A p2p cell means that the active edge is internal to one module and uses that module’s ordinary pipeline communicator. An NC cell means that the active edge crosses from an encoder rank set to the LLM rank set and therefore invokes the edge-local non-colocated communicator. Blank cells indicate that the rank has no active neighbor for that call. Warmup illustrates why the dispatch must be graph-aware. Encoder 1 first sends activations from E1P0 to E1P1 through module-local P2P; its next forward send crosses the encoder–LLM boundary and switches

Heterogeneous Parallelism for Multimodal Large Language Model Training

phase

schedule call

R0 E1P0

R1 E1P1

R2 E2P0

p2p

recv_fwd

R3 L0

R4 L1

R5 L2

p2p

warm p2p

send_fwd

warm end

steady

p2p

recv_fwd send_fwd recv_bwd

p2p

send_bwd recv_fwd

p2p p2p

recv_bwd

p2p

cool p2p

send_bwd

p2p

NC = non-colocated comm.

Figure 14 | Graph-aware communication dispatch for the non-colocated topology in Figure 4(a). Rows are pipeline schedule calls during warmup, steady state, and cooldown; columns are rank roles in the expanded module graph. Each nonempty cell shows the communicator selected for the active edge: p2p for module-internal pipeline communication and NC for an encoder–LLM boundary edge using the non-colocated communicator. Paired steady-state calls resolve the forward and backward halves independently, so different halves of the same schedule call may use different edge types. to the non-colocated communicator. Encoder 2 has only one PP stage, so its first forward send is already a boundary transfer. LLM PP0 receives the two encoder boundary tensors through their corresponding non-colocated communicators and becomes runnable for a microbatch only after the required activations have been materialized in the LLM layout. The remaining LLM stages then proceed with ordinary LLM pipeline P2P. The steady-state rows apply the same dispatch rule to paired 1F1B calls. A call such as send_fwd/recv_bwd contains a forward action for a newer microbatch and a backward action for an older microbatch. The runtime resolves the active edge for each half independently. Thus, different ranks in the same schedule phase may use different communication paths, and a boundary rank can send a new activation through an NC edge while receiving an older gradient through the reverse of an NC edge. The scheduler selects the active edge and microbatch; the non-colocated communicator then performs the leader exchange, DP fan-in or fan-out, and local TP/CP broadcast described in Section 3.1. Cooldown uses the same edge identities in reverse. LLM gradients move by P2P until they reach LLM PP0. Boundary gradients then return through the non-colocated communicator attached to the original encoder–LLM edge, while encoder-internal gradients continue through module-local P2P, such as from E1P1 back to E1P0. This is the key invariant at graph joins: even though LLM PP0 consumes multiple encoder activations for the same microbatch, backward routes each gradient through the boundary communicator

Heterogeneous Parallelism for Multimodal Large Language Model Training

associated with the forward edge that produced that activation. For colocated execution, the schedule detail is different because there is no inter-rank boundary P2P between modules. The issue is ordering. Encoder execution may require collectives over the encoder layout, while the LLM 1F1B schedule staggers ranks across pipeline stages and microbatches. The implementation therefore uses the three-phase schedule shown in Figure 4(b). Phase 1 runs encoder forward once over the local microbatch group for the current LLM pipeline window, then applies the colocated forward transform to produce a packed boundary tensor in the LLM PP0 layout. The runtime records the metadata needed for the reverse colocated transform, including the source and destination layout coordinates and the batch intervals used by any fan-in or fan-out transform. Phase 2 runs the ordinary LLM pipeline schedule on detached boundary views. The packed boundary tensor is detached as a leaf variable for the LLM phase and exposed to LLM PP0 as per-microbatch views. As 1F1B advances, PP0 selects the view for the active microbatch, inserts it into that microbatch’s text embeddings, and then uses standard LLM pipeline P2P for the remaining LLM stages. During this phase, the LLM does not enter encoder collectives, which avoids mixing encoder-layout collectives with the staggered LLM pipeline order. Phase 3 reconnects the LLM phase to the encoder side. After the LLM pipeline drains, gradients accumulated on the packed boundary tensor are handed to the colocated backward transform. If the boundary used fan-in, the transform selects from the destination-gradient shard the batch interval contributed by each source rank. If the boundary used fan-out, sibling destination-gradient slices are all-gathered and concatenated to reconstruct the original source-layout gradient shard. The runtime then invokes encoder backward over the saved batched encoder activations for the local microbatch group. The detach in Phase 2 is therefore only a scheduling device, not a change in training semantics. The LLM treats the packed boundary tensor as the leaf input to its pipeline phase, while the encoder-side graph, saved activations, and communicator metadata remain live. The explicit gradient handoff at the start of Phase 3 supplies the boundary gradients needed by the colocated backward transform and reconnects them to the encoder backward computation. The phase boundary also enables encoder-state offload. Encoder parameters are needed during Phase 1 and Phase 3, but not while Phase 2 executes the detached LLM pipeline. When offload is enabled, the runtime starts the device-to-host transfer after encoder forward, overlaps it with the LLM warmup and steady-state pipeline, starts the host-to-device reload during LLM cooldown, and synchronizes before encoder backward consumes the parameters. Optimizer state can use the same split lifecycle when enabled, with transfers scheduled around the encoder forward/backward phases rather than inside the LLM pipeline loop.

Record · ID 238583 · SHA-256 471604cf574b44e2
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.