Multimodal LLM Training, in depth: the encoder-connector-decoder pipeline, staged freezing, packing and the GPU load-balancing problem

The architecture in one picture

• Vision encoder. Usually a pre-trained ViT, such as a CLIP or SigLIP model. It cuts the image into patches, typically 14 by 14 pixels, and outputs one vector per patch.
• Connector. Maps encoder vectors to the LLM's hidden size. LLaVA-1.5 uses a two-layer MLP. Other designs use a cross-attention resampler that outputs a fixed number of query tokens, or merge neighbouring patches before projecting so the LLM sees fewer tokens.
• LLM decoder. An ordinary pre-trained language model. The tokenizer inserts placeholder tokens where images go, and after embedding lookup those positions are overwritten with connector outputs.

The important property for GPU planning is that the encoder's cost depends on the number of patches, and the decoder's cost depends on the total sequence length. A batch can be balanced on one and badly unbalanced on the other.

How many tokens is an image?

At fixed resolution the answer is simple. A ViT-L/14 at 336 by 336 pixels produces 24 by 24, or 576, patch vectors, and LLaVA-1.5 feeds all 576 to the LLM. Every image costs the same, which makes batching easy and throws away detail on documents and charts.

Dynamic-resolution models keep closer to the native image size. Qwen2-VL, for example, uses 14-pixel patches at roughly native resolution, then merges each 2 by 2 group of adjacent patches into one token before the LLM, so an image of H by W pixels becomes about (H/28) times (W/28) LLM tokens. It also replaces 1D positions with multimodal rotary positions (M-RoPE) that carry separate temporal, height and width components. A 1344 by 896 image therefore produces 6,144 patches in the encoder and 1,536 tokens in the LLM. Tiling schemes such as LLaVA-NeXT's AnyRes reach a similar result differently: a low-resolution overview plus several fixed-size tiles, each encoded separately.

Dynamic resolution is better for OCR and fine detail, but it turns image cost into a variable that ranges over more than an order of magnitude, and that variable is what the rest of this article has to manage.

Training stages and what to freeze

The widely copied recipe has two stages.

1. Alignment. Freeze the encoder and the LLM and train only the connector on image-caption pairs. LLaVA-1.5 used about 558K pairs. The connector learns to produce vectors the LLM can read. Only the connector, about 20 million parameters for a 7B-class model, receives gradients, so optimizer memory is tiny and a high learning rate (around 1e-3 in LLaVA-style recipes) is normal.
2. Instruction tuning. Unfreeze the LLM and keep training the connector on multimodal conversations. LLaVA-1.5's mix had about 665K examples and included text-only conversations. The learning rate drops to full fine-tuning levels, around 2e-5.

Larger efforts add a broad multimodal pre-training stage between the two and often unfreeze the encoder, usually with a smaller learning rate than the LLM, so it adapts to higher resolutions and to text-heavy images. Unfreezing changes the GPU picture: the encoder now needs backward passes, gradients and optimizer state, and its activations must be kept or recomputed.

Keep text-only data in every stage after alignment. Without it the model's language ability degrades, which shows up as worse instruction following on plain text prompts and is easy to miss if you evaluate only on image benchmarks.

The splice and the loss mask

The core of a training step is short. The parts people get wrong are the placeholder count check and the label mask.

import torch
import torch.nn.functional as F

IGNORE = -100

def mllm_step(batch, encoder, connector, llm, image_token_id, train_encoder=False):
    # batch["pixel_values"]: patches for every image in the micro-batch
    # batch["input_ids"]:    text with N placeholder ids per image
    # batch["labels"]:       copy of input_ids with prompt, system and image positions = IGNORE
    with torch.set_grad_enabled(train_encoder):
        feats = encoder(batch["pixel_values"])          # [num_image_tokens_total, d_vis]
    img_embeds = connector(feats)                       # [num_image_tokens_total, d_llm]

    embeds = llm.get_input_embeddings()(batch["input_ids"])   # [B, T, d_llm]
    mask = batch["input_ids"] == image_token_id
    n_slots = int(mask.sum())
    if n_slots != img_embeds.shape[0]:                  # fail loudly, never truncate
        raise ValueError(f"{n_slots} placeholders but {img_embeds.shape[0]} image tokens")
    embeds = embeds.masked_scatter(mask.unsqueeze(-1), img_embeds.to(embeds.dtype))

    logits = llm(inputs_embeds=embeds,
                 attention_mask=batch["attention_mask"],
                 position_ids=batch["position_ids"]).logits
    return F.cross_entropy(logits[:, :-1].flatten(0, 1).float(),
                           batch["labels"][:, 1:].flatten(), ignore_index=IGNORE)

The label mask sets every position that is not assistant output to the ignore index: system prompt, user turn, and every image placeholder. Training on image placeholders asks the model to predict a token that is overwritten on input anyway; training on user turns teaches it to write questions. The position ids are passed explicitly because packed sequences restart positions at each sample boundary.

Worked example: why high-resolution images unbalance ranks

Assumptions, stated so you can redo the sums: a ViT-L/14 encoder with about 304M parameters, 24 layers and width 1024; a 7B-parameter LLM; 300 text tokens per sample. Encoder forward FLOPs are 2 times parameters times patches for the dense layers plus 4 times patches squared times width times layers for attention. LLM training FLOPs use the usual 6 times parameters times tokens. The numbers below are the output of a short script with exactly those formulas.

At fixed resolution the encoder is a rounding error. At high resolution it becomes a real fraction of the step, and because attention grows with the square of patch count, it grows faster than the image's token count.

Now pack two data-parallel ranks to the same token budget. Rank A gets four fixed-resolution samples, 3,504 LLM tokens. Rank B gets two high-resolution samples, 3,672 LLM tokens. A token-based balancer calls that even. The encoder forward on rank A costs 1.5e12 FLOPs and on rank B 1.49e13, almost ten times more. Every step, rank A finishes early and waits at the gradient all-reduce. With hundreds of ranks the step time is set by the unluckiest one.

The fixes, roughly in order of effort: balance micro-batches on an estimated cost that adds encoder and decoder terms instead of tokens alone; cap the maximum pixels per image during training; bucket by resolution so ranks get similar mixes; and, for the largest jobs, run the encoder on a separate pool of GPUs and ship features to the LLM ranks.

Parallelism and memory

The encoder is small relative to the LLM, so the usual layout replicates it on every data-parallel rank and shards only the LLM with FSDP, tensor parallelism or both. When frozen, run it under no_grad in bf16 with no optimizer state; its memory cost is then just weights and one micro-batch of activations.

Pipeline parallelism needs more care. The encoder and connector naturally sit on the first stage, which makes that stage heavier than the others and creates bubbles that move with image resolution. Options are to give the first stage fewer LLM layers, put the encoder on its own stage, or move it out of the pipeline altogether. The general bubble mechanics are in Pipeline parallelism.

Precision follows LLM practice: bf16 compute with fp32 master weights and fp32 loss, as described in Mixed precision training. Some vision encoders were trained in fp16 or fp32 and produce occasional large activations; if loss spikes coincide with particular images, check encoder outputs for infinities before blaming the LLM.

Feeding the GPUs

Image decode and resize are CPU work, and a node with eight GPUs can easily need more JPEG decoding throughput than its data-loader workers provide. Symptoms are low, saw-toothed GPU utilisation and step times that improve when you add workers.

• Store data as large sequential shards (tar or similar) with image bytes, text and pre-computed metadata, including image size, so the batcher can estimate cost without decoding.
• Pre-tokenize text offline. Tokenization is cheap per sample but adds up across millions.
• Decode at the target resolution where the library supports it, rather than decoding full size and downscaling.
• Pack multiple samples per sequence to remove padding, with position ids reset per sample and an attention layout that stops samples attending to each other. Variable-length attention kernels support this directly.
• Prefetch at least two batches ahead and measure the loader alone: iterate it without the model and confirm it exceeds the training step rate.

Trade-offs to decide explicitly

Evaluate on both sides of the model throughout training. A small fixed set of image questions, a document or chart set if you care about OCR, and one text-only benchmark, run at every checkpoint, catch most regressions earlier than a final evaluation would, including the common case where a resolution change helps documents and hurts natural images.

Failure modes

• Placeholder mismatch. The processor's token count and the encoder's output disagree after a resolution or merge change. Silent truncation produces a model that sees half an image; raise instead.
• Language regression. No text-only data in instruction tuning; image scores rise while text benchmarks fall.
• Connector learning rate carried into stage two. A rate suited to training only the connector applied to the full LLM destroys it in a few hundred steps.
• Rank stragglers. High step-time variance with token-balanced batches; per-rank timing shows the same ranks waiting at all-reduce.
• Loader starvation. GPU utilisation drops whenever a high-resolution shard comes up.
• Cross-sample attention in packing. A causal mask over the whole packed row lets sample two attend to sample one's image.

What to do next

1. Write down, for your model, the encoder patch count and LLM token count as a function of image size, and plot the distribution over your dataset.
2. Run the FLOP arithmetic from the worked example with your own encoder and LLM sizes to see whether the encoder is 1% or 30% of the step.
3. Add the placeholder-count check and a unit test that builds labels for one conversation and asserts which positions are ignored.
4. Benchmark the data loader alone and size CPU workers until it outruns the GPUs with margin.
5. Log per-rank step time; if variance is high, switch to cost-based batching or cap pixels per image.
6. Keep text-only data in the mix and track at least one text-only benchmark through every stage.