diff --git a/tests/trainer/test_padding_utils_on_cpu.py b/tests/trainer/test_padding_utils_on_cpu.py new file mode 100644 index 00000000000..b8bdfce76aa --- /dev/null +++ b/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) diff --git a/verl/trainer/ppo/padding_utils.py b/verl/trainer/ppo/padding_utils.py index 436e17bc602..1d070c4c9ed 100644 --- a/verl/trainer/ppo/padding_utils.py +++ b/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.