A multimodal model has two main parts: encoders that turn images, video, or audio into embeddings, and the LLM backbone, which is the language model that takes those embeddings together with the text tokens and does most of the computation.
Memory profile
a. The encoder for images, video, or audio is usually small, but long samples create a lot of activation memory.
b. The LLM backbone is large and needs tensor, pipeline, and expert parallelism. When both are forced into the same parallelism strategy, one of them ends up running inefficiently.
c. The data mix also changes during training. Teams adjust the ratio of text, image, video, and audio samples between training phases, and sometimes every few steps.
As such the training system has to treat the data mix as a dynamic entity while the job is running.
MegaScale-Omni from Bytedance runs the encoder and the LLM on the same GPUs but give each its own parallelism strategy.
In my previous post I covered architecture where MoE gets its own parallelism strategy as well.
Encoder work is spread evenly across all GPUs and scheduled right before the LLM needs it. Samples are reordered by length across nearby GPUs before any compute starts, so no single GPU gets all the long samples.
Two very interesting ideas from the paper
a. The LLM backbone keeps the standard combination of data, tensor, pipeline, sequence, and expert parallelism.
b. The encoder runs data parallel across all GPUs, and each GPU switches to sequence parallelism only for long samples.
c. Every pipeline stage gets an equal share of the encoder work, and that work runs right before the LLM backbone needs it.
This avoids the two common problems with current systems. The first pipeline stage no longer becomes the bottleneck when the share of images goes up, and there is no dedicated set of encoder GPUs sitting idle when the batch is mostly text.
The second idea - Balancing work before it reaches the GPUs.
a. Data loaders usually give each GPU the same number of samples, but multimodal samples vary a lot in length, so one GPU can end up with several long images while the rest wait.
b. MegaScale-Omni has NVLinked GPUs exchange sample sizes, reorder samples so each GPU gets a similar amount of encoder work, and restore the original order before the LLM backbone sees the data.
c. It then redistributes encoder outputs to match the layout the LLM backbone expects.
Observations
1. I have not seen a open weight multi-modal model trained with heterogeneous parallelism strategies for encoders, attention and experts.
2. As far as I know megatron-LM is the only framework that supports these heterogeneous parallelism strategies as standard concepts while pytorch and JAX gives you building blocks and you have to put them together.