From 3baeff9f5cdc4dbfb3db908751921640764b5131 Mon Sep 17 00:00:00 2001 From: khazic Date: Mon, 24 Aug 2026 19:54:27 +0800 Subject: [PATCH] fix(peagle): validate the packed document lengths instead of reusing them per row seq_lengths is a flat list of document lengths for ONE packed sequence, which is what base_trainer builds for P-EAGLE. The forward loops over the batch but hands the whole tensor to every row, so a multi-row batch would silently give every row row zero's document layout. base_trainer always produces a single flat row today, so this is latent, but the loop is written for batch_size > 1 and the fallback branch does index per row. Nothing checked that the lengths cover the sequence either. The COD mask gives any tail past sum(seq_lengths) document id -1, and its is_not_padding term then makes those queries attend to nothing at all, so a short seq_lengths drafts on garbage rather than failing. Reject both, and rename the local to document_lengths since that is what the mask consumes. Signed-off-by: khazic --- .../test_peagle_backend_contract.py | 51 +++++++++++++++++++ verl_speco/backends/peagle_trainer_backend.py | 27 ++++++++-- 2 files changed, 75 insertions(+), 3 deletions(-) diff --git a/tests/integration/test_peagle_backend_contract.py b/tests/integration/test_peagle_backend_contract.py index 094800fe..e0db4c12 100644 --- a/tests/integration/test_peagle_backend_contract.py +++ b/tests/integration/test_peagle_backend_contract.py @@ -294,6 +294,57 @@ def _model(**kwargs): assert float(out["accuracy"]) == pytest.approx(0.75) +def _peagle_training_model_and_batch(batch_size: int, seq_len: int, seq_lengths): + import torch + + from verl_speco.backends.peagle_trainer_backend import PEagleTrainingModel + from verl_speco.models.peagle import LlamaForCausalLMPeagle + + config = _tiny_peagle_config() + model = PEagleTrainingModel( + LlamaForCausalLMPeagle(config), num_depths=2, down_sample_ratio=0.5 + ) + kwargs = dict( + input_ids=torch.zeros(batch_size, seq_len, dtype=torch.long), + aux_hidden=torch.zeros(batch_size, seq_len, config.target_hidden_size * 3), + loss_mask=torch.ones(batch_size, seq_len), + attention_mask=torch.ones(batch_size, seq_len, dtype=torch.long), + target_logits=torch.zeros(batch_size, seq_len, config.vocab_size), + seq_lengths=seq_lengths, + ) + return model, kwargs + + +def test_peagle_rejects_packed_lengths_for_a_multi_row_batch() -> None: + """seq_lengths describes one packed sequence, so it cannot span rows. + + The forward loops over the batch but applies the whole seq_lengths tensor to + every row, which would silently give row 1 row 0's document layout. + """ + torch = pytest.importorskip("torch") + pytest.importorskip("transformers") + + model, kwargs = _peagle_training_model_and_batch( + batch_size=2, seq_len=8, seq_lengths=torch.tensor([4, 4]) + ) + + with pytest.raises(ValueError, match="single packed sequence"): + model(**kwargs) + + +def test_peagle_rejects_packed_lengths_that_do_not_cover_the_sequence() -> None: + """A short seq_lengths silently turns the tail into unattendable padding.""" + torch = pytest.importorskip("torch") + pytest.importorskip("transformers") + + model, kwargs = _peagle_training_model_and_batch( + batch_size=1, seq_len=8, seq_lengths=torch.tensor([3, 2]) + ) + + with pytest.raises(ValueError, match="sum to 5"): + model(**kwargs) + + def test_peagle_checkpoint_export_unwraps_the_training_model() -> None: pytest.importorskip("torch") from types import SimpleNamespace diff --git a/verl_speco/backends/peagle_trainer_backend.py b/verl_speco/backends/peagle_trainer_backend.py index 8fbb09b9..78fd1f60 100644 --- a/verl_speco/backends/peagle_trainer_backend.py +++ b/verl_speco/backends/peagle_trainer_backend.py @@ -122,6 +122,25 @@ def forward( batch_size, seq_len = input_ids.shape device = input_ids.device + if seq_lengths is not None: + # seq_lengths is a flat list of document lengths for ONE packed + # sequence, which is what base_trainer builds for P-EAGLE. There is no + # per-row structure to index, so a multi-row batch would silently + # apply one row's document layout to all of them. + if batch_size > 1: + raise ValueError( + "P-EAGLE seq_lengths describe a single packed sequence, but the batch " + f"has {batch_size} rows; pack the documents into one row or drop seq_lengths" + ) + # Any tail past sum(seq_lengths) gets document id -1 in the COD mask, + # which makes those queries attend to nothing at all rather than + # failing, so check the invariant instead of drafting on garbage. + total_length = int(seq_lengths.sum()) + if total_length != seq_len: + raise ValueError( + f"P-EAGLE seq_lengths sum to {total_length} but the packed sequence is " + f"{seq_len} tokens long" + ) loss_num = torch.zeros((), device=device, dtype=torch.float32) loss_den = torch.zeros((), device=device, dtype=torch.float32) correct = torch.zeros((), device=device, dtype=torch.float32) @@ -137,9 +156,11 @@ def forward( ) orig_positions = anchor_pos + depth if seq_lengths is not None: - row_length = seq_lengths.to(device) + document_lengths = seq_lengths.to(device) else: - row_length = attention_mask[b].sum().clamp_min(1).reshape(1).to(device) + document_lengths = ( + attention_mask[b].sum().clamp_min(1).reshape(1).to(device) + ) loss_positions = row_loss_mask[0, orig_positions].bool() is_depth0 = depth == 0 @@ -161,7 +182,7 @@ def forward( block_mask = draft.build_peagle_block_mask( anchor_pos=anchor_pos, depth=depth, - lengths=row_length, + lengths=document_lengths, total_seq_len=seq_len, ) hidden = draft.forward_peagle(