Enable VLM context parallel training

src/prime_rl/trainer/batch.py

@@ -196,6 +196,8 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch
         assert len(mm_token_type_ids) == len(input_ids), (
             f"mm_token_type_ids: {len(mm_token_type_ids)}, input_ids: {len(input_ids)}"
         )
+    if mm_kwargs is not None and "image_grid_thw" in mm_kwargs and mm_token_type_ids is None:
+        raise ValueError("image_grid_thw multimodal samples require mm_token_type_ids")
     assert len(env_names) == len(input_ids), f"env_names: {len(env_names)}, input_ids: {len(input_ids)}"
 
     return MicroBatch(
@@ -214,6 +216,7 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch
         rl_weights=rl_weights,
         ce_weights=ce_weights,
         ref_kl_weights=ref_kl_weights,
+        seq_lens=[len(input_ids)],
     )
 
 
@@ -237,15 +240,14 @@ def first_sample(self) -> MicroBatch:
 
     def can_add(self, sample: MicroBatch, max_seq_len: int) -> bool:
         # Loss routing is per token (component weight streams), so samples of
-        # different loss types pack together freely — only modality, length and
-        # routed-experts presence constrain packing.
+        # different loss types pack together freely. The packer groups by
+        # run/LoRA before this point.
         first_sample = self.first_sample
-        return (
-            not _is_multimodal_sample(first_sample)
-            and not _is_multimodal_sample(sample)
-            and self.length + len(sample.input_ids) <= max_seq_len
-            and (first_sample.routed_experts is None) == (sample.routed_experts is None)
-        )
+        if self.length + len(sample.input_ids) > max_seq_len:
+            return False
+        if (first_sample.routed_experts is None) != (sample.routed_experts is None):
+            return False
+        return True
 
     def add(self, lora_idx: int, sample: MicroBatch) -> None:
         self.samples.append((lora_idx, sample))
@@ -290,6 +292,8 @@ def _materialize_bin(bin_content: _MicroBatchBin, num_loras: int) -> MicroBatch:
     env_names: list[str] = []
     ref_logprobs: list[float] | None = [] if has_ref_logprobs else None
     mm_token_type_ids: list[int] | None = [] if has_mm_token_type_ids else None
+    mm_kwargs: dict[str, EncodedTensor] | None = None
+    seq_lens: list[int] = []
     streams: dict[str, list[float] | None] = {name: ([] if has_stream[name] else None) for name in STREAM_FILL}
     routed_experts: RoutedExperts | None = None
     lora_num_tokens = [0] * num_loras
@@ -322,11 +326,19 @@ def _materialize_bin(bin_content: _MicroBatchBin, num_loras: int) -> MicroBatch:
                 assert routed_experts.shape[1:] == sample.routed_experts.shape[1:]
                 routed_experts.data += sample.routed_experts.data
                 routed_experts.shape[0] += sample.routed_experts.shape[0]
+        if sample.mm_kwargs is not None:
+            if mm_kwargs is None:
+                mm_kwargs = copy.deepcopy(sample.mm_kwargs)
+            else:
+                for key in mm_kwargs:
+                    mm_kwargs[key].data += sample.mm_kwargs[key].data
+                    mm_kwargs[key].shape[0] += sample.mm_kwargs[key].shape[0]
+        seq_lens.extend(sample.seq_lens)
         lora_num_tokens[lora_idx] += sample_len
 
     sequence_lengths = [len(sample.input_ids) for _, sample in bin_content.samples]
     assert sum(sequence_lengths) == len(input_ids), (sequence_lengths, len(input_ids))
-    first_sample = bin_content.first_sample
+    assert sum(seq_lens) == len(input_ids), (seq_lens, len(input_ids))
 
     return MicroBatch(
         input_ids=input_ids,
@@ -341,10 +353,11 @@ def _materialize_bin(bin_content: _MicroBatchBin, num_loras: int) -> MicroBatch:
         routed_experts=routed_experts,
         mm_token_type_ids=mm_token_type_ids,
         env_names=env_names,
-        mm_kwargs=first_sample.mm_kwargs if _is_multimodal_sample(first_sample) else None,
+        mm_kwargs=mm_kwargs,
         rl_weights=streams["rl_weights"],
         ce_weights=streams["ce_weights"],
         ref_kl_weights=streams["ref_kl_weights"],
+        seq_lens=seq_lens,
     )
 
 
@@ -374,19 +387,17 @@ def packed_samples_into_micro_bs(
 ) -> list[MicroBatch]:
     """
     Pack samples into micro_batch efficiently.
-    We follow the First Fit Decreasing algorithm to pack the samples into bins and minimize potential padding while never truncating.
-    With per-token temperatures, samples can be packed together regardless of their temperature values.
-
-    NOTE: Multimodal samples (with mm_kwargs) are NOT packed together as they have variable-sized
-    vision data that doesn't pack well. Each multimodal sample becomes its own micro batch.
+    We follow the First Fit Decreasing algorithm to pack samples into bins and
+    minimize potential padding while never truncating. Packed batches preserve
+    sample boundaries in ``seq_lens``.
     """
     # Sort by (lora_idx, -length) for packing efficiency
     samples.sort(key=lambda x: (x[0], -len(x[1].input_ids)))
 
     bins: list[_MicroBatchBin] = []
 
     for idx, sample in samples:
-        # Try to find a bin that can fit this sequence (only pack text-only samples)
+        # Try to find a bin that can fit this sequence.
         for bin_content in bins:
             if bin_content.can_add(sample, max_seq_len):
                 bin_content.add(idx, sample)
@@ -464,6 +475,7 @@ def pad_micro_batch(micro_batch: MicroBatch, pad_to_multiple_of: int) -> MicroBa
         )
     if micro_batch.mm_token_type_ids is not None:
         micro_batch.mm_token_type_ids.extend([0] * padding_size)
+    micro_batch.seq_lens.append(padding_size)
     if micro_batch.routed_experts is not None:
         _pad_routed_experts(micro_batch, padding_size)
     micro_batch.env_names.extend([""] * padding_size)
@@ -538,32 +550,24 @@ def prepare_batch(
     Prepare a batch of problems for each GPU. Each batch is a list of micro batches.
     Each micro batch is shape [1, seq_len], the number of samples is not fixed per micro batch.
 
-    FSDP requires all ranks to execute the same operations at each step. If one rank
-    processes a multimodal batch (triggering the vision encoder) while another processes
-    a text-only batch, the all-gather will hang. We separate micro batches by modality
-    and distribute them so that at each step index, all ranks see the same modality.
+    VLM models handle text-only samples with a synthetic dummy vision path, so
+    microbatches can be distributed through one generic packing path.
     """
     all_samples = [(idx, prepare_sample(rollout, seq_len)) for idx, rollout in zip(idxs, rollouts)]
 
-    micro_batches = packed_samples_into_micro_bs(all_samples, seq_len, num_loras, num_train_workers, bin_cost)
+    micro_batches = packed_samples_into_micro_bs(
+        all_samples,
+        seq_len,
+        num_loras,
+        num_train_workers,
+        bin_cost,
+    )
     micro_batches = [pad_micro_batch(micro_batch, pad_to_multiple_of) for micro_batch in micro_batches]
 
-    # Separate by modality so each step index has uniform modality across all ranks
-    mm_batches = [b for b in micro_batches if _is_multimodal_sample(b)]
-    text_batches = [b for b in micro_batches if not _is_multimodal_sample(b)]
+    micro_batches = _pad_group_for_distribution(micro_batches, num_train_workers)
 
-    # Pad each group independently so its count is divisible by num_train_workers
-    mm_batches = _pad_group_for_distribution(mm_batches, num_train_workers)
-    text_batches = _pad_group_for_distribution(text_batches, num_train_workers)
-
-    # Alignment check after distribution padding so the dummy batches are covered too
-    for micro_batch in (*mm_batches, *text_batches):
+    # Alignment check after distribution padding so the dummy batches are covered too.
+    for micro_batch in micro_batches:
         _assert_token_arrays_aligned(micro_batch)
 
-    batches_per_gpu: list[list[MicroBatch]] = [[] for _ in range(num_train_workers)]
-    for group in (mm_batches, text_batches):
-        group_batches_per_gpu = _distribute_group(group, num_train_workers, bin_cost)
-        for worker_idx, worker_batches in enumerate(group_batches_per_gpu):
-            batches_per_gpu[worker_idx].extend(worker_batches)
-
-    return batches_per_gpu
+    return _distribute_group(micro_batches, num_train_workers, bin_cost)

src/prime_rl/trainer/model.py

@@ -1,3 +1,4 @@
+import inspect
 import logging
 import os
 import time
@@ -56,7 +57,11 @@
 from prime_rl.utils.logger import get_logger
 from prime_rl.utils.sequence import get_cu_seqlens_from_position_ids
 from prime_rl.utils.utils import format_time
-from prime_rl.utils.vlm import get_language_model, get_vision_encoder, is_vlm_architecture
+from prime_rl.utils.vlm import (
+    get_language_model,
+    get_vision_encoder,
+    is_vlm_architecture,
+)
 
 
 def pre_download_model(model_name: str) -> None:
@@ -476,6 +481,12 @@ def get_model(
             raise ValueError(
                 "VLM models must use optimization_dtype='bfloat16' and reduce_dtype='bfloat16' to match vLLM inference."
             )
+        # Context parallelism patches the global flash-attention dispatch to do the
+        # Ulysses all-to-all, which is causal- and sharded-sequence-only.
+        vision_config = getattr(model_config, "vision_config", None)
+        if config.cp > 1 and vision_config is not None:
+            vision_config._attn_implementation = "sdpa"
+            logger.info("CP enabled for VLM: pinning vision encoder attention to 'sdpa' (excluded from Ulysses CP).")
 
     # GPT-OSS only supports FlashAttention via kernels-community/vllm-flash-attn3, which requires Hopper (SM 90).
     # On other architectures (e.g. Blackwell), users must fall back to eager attention.
@@ -583,6 +594,12 @@ def get_model(
     else:
         impl_to_use = config.impl
 
+    if config.vlm is not None and config.cp > 1 and not (is_vlm_arch and impl_to_use == "custom" and custom_vlm_cls):
+        raise ValueError(
+            "VLM context parallelism requires a custom PrimeRL VLM implementation; "
+            f"{getattr(model_config, 'model_type', config.name)!r} is only supported with cp=1."
+        )
+
     with device:
         if impl_to_use == "custom" and custom_vlm_cls is not None:
             model_cls = custom_vlm_cls
@@ -663,13 +680,31 @@ def setup_fsdp(model: nn.Module, config: ModelConfig, parallel_dims: ParallelDim
         dp_mod_ep_mesh = parallel_dims.world_mesh[tuple(dp_mod_ep_mesh_dim_names)]
 
     is_vlm_training = config.vlm is not None
+    cp_vlm = is_vlm_training and config.cp > 1
+    # Params kept out of FSDP entirely (full/replicated on every rank).
+    fsdp_ignored_params: set[nn.Parameter] = set()
+
+    def _all_params_frozen(module: nn.Module) -> bool:
+        params = list(module.parameters())
+        return bool(params) and not any(p.requires_grad for p in params)
+
+    def _ignore_params(module: nn.Module, name: str) -> None:
+        params = list(module.parameters())
+        if not params:
+            return
+        fsdp_ignored_params.update(params)
+        get_logger().info(f"CP+VLM: keeping frozen {name} out of FSDP (ignored_params).")
+
     if is_vlm_training:
         vision_encoder = get_vision_encoder(model, override=config.vlm.vision_encoder_attr)
         if vision_encoder is None:
             raise ValueError(f"VLM model {config.name} has no recognized vision encoder")
 
-        fully_shard(vision_encoder, mesh=hsdp_mesh, **fsdp_config)
-        get_logger().info(f"Applied FSDP to vision encoder (frozen={config.vlm.freeze_vision_encoder})")
+        if cp_vlm and _all_params_frozen(vision_encoder):
+            _ignore_params(vision_encoder, "vision encoder")
+        else:
+            fully_shard(vision_encoder, mesh=hsdp_mesh, **fsdp_config)
+            get_logger().info(f"Applied FSDP to vision encoder (frozen={config.vlm.freeze_vision_encoder})")
 
     language_model = get_language_model(model, override=config.vlm.language_model_attr if is_vlm_training else None)
     transformer_layers = language_model.layers
@@ -692,11 +727,12 @@ def setup_fsdp(model: nn.Module, config: ModelConfig, parallel_dims: ParallelDim
     if shard_norm_and_lm_head:
         # This optimization breaks weight tying
         embed_module = getattr(language_model, "embed_tokens", None) or getattr(language_model, "embeddings", None)
-        fully_shard(
-            embed_module,
-            mesh=hsdp_mesh,
-            **fsdp_config,
-        )
+        if embed_module is not None:
+            fully_shard(
+                embed_module,
+                mesh=hsdp_mesh,
+                **fsdp_config,
+            )
         norm_module = getattr(language_model, "norm", None) or language_model.norm_f
         fully_shard(
             [model.lm_head, norm_module],
@@ -714,6 +750,7 @@ def setup_fsdp(model: nn.Module, config: ModelConfig, parallel_dims: ParallelDim
         mp_policy=mp_policy,
         offload_policy=offload_policy,
         reshard_after_forward=config.reshard_after_forward,
+        ignored_params=fsdp_ignored_params or None,
     )
 
     if not parallel_dims.ep_enabled:
@@ -1164,28 +1201,47 @@ def setup_model(
     return model
 
 
+def _model_accepts_seq_lens(model: nn.Module) -> bool:
+    accepts_seq_lens = getattr(model, "_prime_rl_accepts_seq_lens", None)
+    if accepts_seq_lens is None:
+        accepts_seq_lens = "seq_lens" in inspect.signature(model.forward).parameters
+        setattr(model, "_prime_rl_accepts_seq_lens", accepts_seq_lens)
+    return accepts_seq_lens
+
+
 def forward(
     model: nn.Module,
-    input_ids: Int[Tensor, "batch seq"],
-    position_ids: Int[Tensor, "batch seq"],
+    input_ids: Int[Tensor, "batch seq"] | None,
+    position_ids: Int[Tensor, "batch seq"] | None,
     labels: Int[Tensor, "batch seq"] | None = None,
     temperature: Tensor | None = None,
     routed_experts: Int[Tensor, "batch seq layers topk"] | None = None,
+    inputs_embeds: Tensor | None = None,
     # Generic multimodal kwargs (e.g. {"pixel_values": ...,
     # "image_grid_thw": ...} for Qwen3-VL; just {"pixel_values": ...}
     # for Gemma3). Passed straight through to ``model(**kwargs)`` so
     # the model's HF forward signature is the schema. ``mm_token_type_ids``
-    # is split out because it's prime-rl-computed (from token ids),
-    # not a renderer/processor output.
+    # is split out because it's prime-rl-computed (from token ids).
     mm_kwargs: dict[str, Tensor] | None = None,
     mm_token_type_ids: Int[Tensor, "batch seq"] | None = None,
+    seq_lens: Int[Tensor, "segments"] | None = None,
+    seq_lens_are_global: bool = False,
+    cp_group: object | None = None,
+    cp_rank: int | None = None,
+    cp_world_size: int | None = None,
+    cp_style: str | None = None,
 ) -> PrimeLmOutput:
+    forwards_seq_lens = seq_lens is not None and _model_accepts_seq_lens(model)
+
     # Build kwargs for model forward
     kwargs = {
-        "input_ids": input_ids,
         "labels": labels,
         "temperature": temperature,
     }
+    if inputs_embeds is None:
+        kwargs["input_ids"] = input_ids
+    else:
+        kwargs["inputs_embeds"] = inputs_embeds
 
     if mm_kwargs:
         # Forward the per-model multimodal tensors verbatim, plus the
@@ -1194,18 +1250,30 @@ def forward(
         kwargs.update(mm_kwargs)
         if mm_token_type_ids is not None:
             kwargs["mm_token_type_ids"] = mm_token_type_ids
-        # ``position_ids`` for MRoPE families: Qwen3-VL's HF forward
-        # recomputes 3D positions from ``image_grid_thw`` and breaks if
-        # given the trainer's pre-computed 1D ``position_ids``. Detect
-        # via the mm_kwargs shape so we don't enumerate model_types.
         if "image_grid_thw" not in mm_kwargs:
             kwargs["position_ids"] = position_ids
+        elif forwards_seq_lens:
+            kwargs["seq_lens"] = seq_lens
     else:
         kwargs["position_ids"] = position_ids
+        if forwards_seq_lens:
+            kwargs["seq_lens"] = seq_lens
+        if seq_lens_are_global and forwards_seq_lens:
+            kwargs["seq_lens_are_global"] = True
 
     if routed_experts is not None:
         kwargs["routed_experts"] = routed_experts
 
+    if cp_group is not None:
+        kwargs.update(
+            {
+                "cp_group": cp_group,
+                "cp_rank": cp_rank,
+                "cp_world_size": cp_world_size,
+                "cp_style": cp_style,
+            }
+        )
+
     out = model(**kwargs)
 
     # PrimeLmOutput is a TypedDict (dict at runtime), HF outputs are dataclass-like objects