Add multimodule for multimodal models

[wenqingqian]

Colocated MIMO Training for Multimodal Models

1. What's MIMO

MIMO is an existing Megatron concept for multimodal training: the model is treated as multiple modules (vision encoder + LLM) whose parallel configurations are decoupled — the ViT and the LLM no longer have to share one TP/DP/* grid.

Upstream Megatron implements two layouts: UNIFIED (both modules share one grid, i.e. traditional training) and NON_COLOCATED (modules on disjoint GPU ranks). The COLOCATED layout — same GPUs, different grids per module — is still NotImplementedError, and no upstream layout batches ViT work across microbatches.

2. This PR: colocated MIMO with microbatch-level scheduling

We implement the colocated mode: ViT and LLM run on the same GPUs with heterogeneous parallel configurations (e.g., on 8 GPUs: ViT TP=1/DP=8, LLM TP=2/DP=4), each module with its own DDP and distributed optimizer (ChainedOptimizer).

On top of the heterogeneous layout, a microbatch scheduler (MIMOMicrobatchScheduler) batches vision work across microbatches: one ViT macro forward serves vit_batch_factor LLM microbatches, and ViT backward is delayed until the whole macro batch has produced gradients (detach + grad cache). The net effect: fewer, larger ViT forwards per iteration.

3. Difference from Megatron's MIMO

Capability	Megatron UNIFIED	Megatron NON_COLOCATED	Ours (colocated MIMO)
ViT/LLM on the same ranks	✅ (shared grid)	❌ (disjoint ranks)	✅
Heterogeneous TP/DP per module	❌	✅ (but rank-disjoint)	✅ (e.g., ViT TP=1/DP=8, LLM TP=2/DP=4)
Microbatch-level ViT batching (one ViT forward feeds N LLM forwards)	❌	❌	✅ vit_batch_factor + scheduler
Delayed ViT backward	❌	❌	✅ (detach + grad cache, backward after macro batch completes)

4. How to Use It (Qwen3.5 example)

Enable MIMO in the training YAML:

experiment:
  task:
    entrypoint: ./flagscale/train/megatron/train_qwen35.py

model:
  use_mimo: true
  vision_tensor_model_parallel_size: 1   # ViT TP (currently must be 1)
  vision_micro_batch_size: 4             # ViT micro batch per forward
  micro_batch_size: 2                    # LLM micro batch
  global_batch_size: 32
  tensor_model_parallel_size: 2          # LLM TP

With the layout above (world = 8, LLM TP=2 → LLM DP=4), vit_batch_factor = (8 × 4) / (4 × 2) = 4: one ViT forward with 8 samples per rank feeds 4 LLM microbatches.

5. Performance

Setup: Qwen3.5 4B (hybrid GDN + attention, MTP), 8 × A800, seq_len 2048, global batch 64, LLM TP=2 (DP=4), micro batch 2, 200 iterations.

• baseline: ViT and LLM both TP=2/DP=4
• MIMO vbf=N: ViT TP=1/DP=8, vision_micro_batch_size=N → vit_batch_factor=N for N = 2 / 4 / 8 (vbf=8 needs gbs=64 so that num_microbatches=8; gbs=32 caps at vbf=4)

Steady-state (iter 11–200) per-GPU TFLOPs, estimated from per-iteration time with identical iteration FLOPs across configs (same model, same global batch — the ratios are exact):

Config	Iteration time	per-GPU TFLOPs (est.)	vs baseline
baseline	5001.26 ms	77.5	—
MIMO vbf=2	4022.61 ms	96.3	+24.3%
MIMO vbf=4	3817.43 ms	101.5	+31.0%
MIMO vbf=8	3712.44 ms	104.4	+34.7%

Equivalently, iteration time drops by 19.6% / 23.7% / 25.8%. Gains saturate as vbf grows (vbf=4→8 adds only ~2.1pp), since the ViT share of iteration time is already small.

Loss curves stay aligned with the baseline over the full 200 iterations (left two panels). Point-wise relative difference vs baseline (iter 3–200) across vbf=2/4/8: LM loss mean 0.76–0.96%, max ≤ 4.54%; MTP-1 loss mean 0.82–1.08%, max ≤ 4.30%. The drift does not grow with vit_batch_factor (vbf=8 ≈ vbf=4). Residual differences come from per-module gradient clipping (vs one global clip in the baseline) and macro-batch scheduling order.

6. Profile