[trainer] fix: make the multi-trajectory padding row survive THD context parallel

tests/trainer/test_padding_utils_on_cpu.py

@@ -0,0 +1,133 @@
+# Copyright 2025 Meituan Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""The synthetic padding row must survive the THD context-parallel split.
+
+``upsample_batch_to_divisible_size`` appends no-op samples when the batch size is not
+divisible by ``dp_size``, which happens as soon as the number of trajectories per prompt
+varies. Those rows used to be 2 tokens long, which is below what the CP split can handle:
+``preprocess_packed_seqs`` pads each row to ``align_size = tp * cp * 2`` and then hands CP
+rank *r* the slice ``d[half * r : half * (r + 1)]``, where ``half`` comes from the *padded*
+length while ``d`` holds only the *valid* tokens. For a 2-token row at tp=cp=2 the row pads
+to 8, ``half`` is 2, and rank 1 asks for ``d[2:4]`` of a 2-element tensor -- empty, which
+raises on assignment.
+
+verl's own ``preprocess_packed_seqs`` clamps that slice (#6001); the vendored copies in
+mbridge and Megatron-Bridge do not. These tests pin the property that makes the row safe
+regardless of which copy runs: its length is a multiple of the alignment for any realistic
+topology, so it never needs alignment padding and the slice is always in range.
+"""
+
+import pytest
+import torch
+
+from verl.trainer.ppo.padding_utils import (
+    _PADDING_TOKENS_PER_SIDE,
+    construct_minimal_padding_template,
+)
+
+
+def _source_sample() -> tuple[dict, dict]:
+    """Stand-in for one real sample fetched from TransferQueue."""
+    sample = {
+        "prompts": torch.zeros(4, dtype=torch.int64),
+        "responses": torch.zeros(4, dtype=torch.int64),
+        "input_ids": torch.zeros(8, dtype=torch.int64),
+        "attention_mask": torch.ones(8, dtype=torch.int64),
+        "position_ids": torch.arange(8, dtype=torch.int64),
+        "response_mask": torch.ones(4, dtype=torch.int64),
+        "uid": "real",
+    }
+    return sample, {"prompt_len": 4, "response_len": 4, "seq_len": 8}
+
+
+def _first_cp_chunk(valid_len: int, *, tp: int, cp: int, cp_rank: int, clamped: bool):
+    """The first-chunk assignment from ``preprocess_packed_seqs``, both variants.
+
+    ``clamped=True`` is what verl's own copy does; ``clamped=False`` is what the vendored
+    bridge copies still do.
+    """
+    align_size = tp * cp * 2
+    padded = valid_len + (align_size - valid_len % align_size) % align_size
+    half = (padded // cp) // 2
+    out = torch.zeros(padded // cp, dtype=torch.int64)
+    d = torch.arange(valid_len, dtype=torch.int64)
+    first_start, first_end = half * cp_rank, half * (cp_rank + 1)
+    if clamped:
+        first_end = min(first_end, d.shape[0])
+        length = max(first_end - first_start, 0)
+        if length > 0:
+            out[0:length] = d[first_start:first_end]
+    else:
+        out[0:half] = d[first_start:first_end]
+    return out
+
+
+@pytest.mark.parametrize("tp_times_cp", [1, 2, 4, 8, 16, 32, 64])
+def test_padding_row_needs_no_alignment_padding_for_any_topology(tp_times_cp):
+    """A row already a multiple of ``align_size`` is never re-padded, so every CP chunk
+    index stays inside the valid tokens -- with or without the clamp."""
+    sample, _ = construct_minimal_padding_template(*_source_sample(), eos_token_id=7)
+    valid_len = int(sample["attention_mask"].sum())
+    align_size = tp_times_cp * 2
+    assert valid_len % align_size == 0, f"{valid_len} tokens is not a multiple of {align_size}"
+
+
+def test_padding_row_still_contributes_no_gradient_and_no_reward():
+    """Making the row longer must not make it count."""
+    sample, tag = construct_minimal_padding_template(*_source_sample(), eos_token_id=7)
+    assert sample["response_mask"].sum() == 0
+    assert sample["loss_mask"].sum() == 0
+    assert sample["rm_scores"].sum() == 0
+    assert sample["num_turns"] == 0
+    assert tag["is_padding"] is True
+    # The tag must describe the real shape: metrics and the seqlen balancer read it.
+    assert tag["prompt_len"] == _PADDING_TOKENS_PER_SIDE
+    assert tag["response_len"] == _PADDING_TOKENS_PER_SIDE
+    assert tag["seq_len"] == 2 * _PADDING_TOKENS_PER_SIDE
+    assert int(sample["attention_mask"].sum()) == 2 * _PADDING_TOKENS_PER_SIDE
+
+
+def test_padding_row_shapes_stay_internally_consistent():
+    sample, _ = construct_minimal_padding_template(*_source_sample(), eos_token_id=7)
+    n = _PADDING_TOKENS_PER_SIDE
+    assert sample["prompts"].shape == (n,)
+    assert sample["responses"].shape == (n,)
+    assert sample["input_ids"].shape == (2 * n,)
+    assert sample["attention_mask"].shape == (2 * n,)
+    assert sample["position_ids"].shape[-1] == 2 * n
+
+
+def test_a_two_token_row_breaks_the_unclamped_cp_split():
+    """Why the constant exists: the exact failure, at tp=cp=2 on CP rank 1."""
+    with pytest.raises(RuntimeError, match=r"existing size \(0\)"):
+        _first_cp_chunk(2, tp=2, cp=2, cp_rank=1, clamped=False)
+    # verl's own copy survives it; the vendored bridge copies are what still raise.
+    _first_cp_chunk(2, tp=2, cp=2, cp_rank=1, clamped=True)
+
+
+def test_the_padding_row_length_survives_the_unclamped_split():
+    """What the constant buys: no exception even on the unclamped path."""
+    valid_len = 2 * _PADDING_TOKENS_PER_SIDE
+    for cp_rank in range(2):
+        _first_cp_chunk(valid_len, tp=2, cp=2, cp_rank=cp_rank, clamped=False)
+
+
+@pytest.mark.parametrize("valid_len", [128, 4096, 14000])
+@pytest.mark.parametrize("cp_rank", [0, 1])
+def test_clamping_is_a_pure_guard_for_long_rows(valid_len, cp_rank):
+    """Rows long enough to be unaffected must be bit-identical either way, so the guard
+    cannot change training behaviour for real sequences."""
+    unclamped = _first_cp_chunk(valid_len, tp=2, cp=2, cp_rank=cp_rank, clamped=False)
+    clamped = _first_cp_chunk(valid_len, tp=2, cp=2, cp_rank=cp_rank, clamped=True)
+    assert torch.equal(unclamped, clamped)

verl/trainer/ppo/padding_utils.py

@@ -39,6 +39,34 @@
 
 logger = logging.getLogger(__name__)
 
+# Tokens per side (prompt / response) of the synthetic padding sample.
+#
+# This was 1, i.e. a 2-token sequence, which is below what the THD context-parallel split
+# can handle. ``preprocess_packed_seqs`` pads each row to ``align_size = tp * cp * 2`` and
+# then hands CP rank *r* the slice ``d[half * r : half * (r + 1)]``, where ``half`` is
+# derived from the *padded* length while ``d`` holds only the *valid* tokens. For a 2-token
+# row at tp=cp=2 the row pads to 8, ``half`` is 2, and rank 1 asks for ``d[2:4]`` of a
+# 2-element tensor -- an empty slice, which raises on assignment:
+#
+#     RuntimeError: The expanded size of the tensor (2) must match the existing
+#     size (0) at non-singleton dimension 0.  Target sizes: [2].  Tensor sizes: [0]
+#
+# verl's own copy of that function clamps the slice (#6001), but the vendored copies in
+# mbridge and Megatron-Bridge -- which the megatron engine dispatches to -- do not, so the
+# row still crashes there. Fixes are filed for both (ISEEKYAN/mbridge#157,
+# NVIDIA-NeMo/Megatron-Bridge#5614); making the row alignment-safe here removes the
+# dependency on which bridge version is installed.
+#
+# Measured smallest safe valid length, i.e. the length at which the unclamped split stops
+# raising: 2 at TP1/CP2, 3 at TP2/CP2, 8 at TP2/CP4, 16 at TP4/CP4. 64 tokens per side
+# (128 total) is a multiple of ``tp * cp * 2`` for every ``tp * cp`` up to 64, so the row
+# never needs alignment padding and the slice is always in range regardless of topology.
+#
+# The cost is 128 no-op tokens per padding row against a per-GPU token budget in the tens of
+# thousands. These rows still contribute no gradient (``response_mask`` is all zero) and are
+# still excluded from metrics (``is_padding``), so nothing else about them changes.
+_PADDING_TOKENS_PER_SIDE = 64
+
 
 def build_padding_position_ids(source_position_ids: Any, attention_mask: torch.Tensor) -> torch.Tensor:
     """Build padding position ids with the same rank/prefix shape as the source sample."""
@@ -72,7 +100,10 @@ def construct_minimal_padding_template(
     source_tag: dict,
     eos_token_id: int,
 ) -> tuple[dict, dict]:
-    """Construct a minimal text-only padding template of one prompt token and one response token.
+    """Construct a minimal text-only padding template.
+
+    The sequence is ``_PADDING_TOKENS_PER_SIDE`` prompt tokens plus the same number of
+    response tokens; see that constant for why it is not simply one of each.
 
     Args:
         source_td: A single sample dict retrieved from TransferQueue.
@@ -92,7 +123,7 @@ def construct_minimal_padding_template(
     template_tag = copy.deepcopy(source_tag)
 
     # Build minimal sequence
-    prompts = torch.full((1,), eos_token_id, dtype=torch.int64)
+    prompts = torch.full((_PADDING_TOKENS_PER_SIDE,), eos_token_id, dtype=torch.int64)
     input_ids = prompts.repeat(2)
     attention_mask = torch.ones_like(input_ids, dtype=torch.int64)
     response_mask = torch.zeros_like(prompts)
@@ -120,7 +151,12 @@ def construct_minimal_padding_template(
         template_sample.pop("routed_experts", None)
 
     # Padding flag is deployed to protect metrics calculation (e.g. response length, score, reward).
-    template_tag.update(is_padding=True, prompt_len=1, response_len=1, seq_len=2)
+    template_tag.update(
+        is_padding=True,
+        prompt_len=_PADDING_TOKENS_PER_SIDE,
+        response_len=_PADDING_TOKENS_PER_SIDE,
+        seq_len=2 * _PADDING_TOKENS_PER_SIDE,
+    )
     return template_sample, template_tag
 
 
@@ -132,7 +168,7 @@ def upsample_batch_to_divisible_size(
     """Append synthetic no-op samples so the batch size becomes divisible by *batch_multiple*.
 
     The synthetic samples reuse the first real sample as a metadata template,
-    but manually construct a minimal ``prompt_len=1 / response_len=1`` sequence
+    but manually construct a minimal no-op sequence (see ``_PADDING_TOKENS_PER_SIDE``)
     and zero out reward-related fields so they do not contribute to PPO,
     entropy, or KL losses.  An ``is_padding`` flag is added in the tag for
     downstream metrics filtering.