From 335dae21c7976d5923d65824f5080bccbae44e70 Mon Sep 17 00:00:00 2001 From: khazic Date: Mon, 24 Aug 2026 19:36:13 +0800 Subject: [PATCH] fix(peagle): refuse a reduced draft vocabulary without a t2d/d2t mapping The EAGLE-3 backend refuses draft_vocab_size != vocab_size unless the draft checkpoint supplies valid t2d/d2t buffers, then validates them. P-EAGLE overrides build_model entirely and does neither. Setting peagle_draft_vocab_size below the target vocabulary therefore falls through to the model constructor's default, t2d[:draft_vocab_size] = True, which means "the draft vocabulary is target token ids 0..N-1". On any real tokenizer that is an arbitrary slice with no relation to token frequency, and the draft can never emit anything outside it. Training reports nothing unusual because the loss restricts the target logits to the same slice. Give P-EAGLE the same guard, loading the checkpoint with output_loading_info so a supplied mapping is detected the way EAGLE-3 detects it. The shared _validate_vocab_mapping is inherited by every EAGLE-family backend, so its messages now name the algorithm that failed instead of always saying EAGLE3. Signed-off-by: khazic --- .../test_peagle_backend_contract.py | 95 +++++++++++++++++++ verl_speco/backends/eagle3_trainer_backend.py | 11 ++- verl_speco/backends/peagle_trainer_backend.py | 35 ++++++- 3 files changed, 135 insertions(+), 6 deletions(-) diff --git a/tests/integration/test_peagle_backend_contract.py b/tests/integration/test_peagle_backend_contract.py index 094800fe..23ec08d3 100644 --- a/tests/integration/test_peagle_backend_contract.py +++ b/tests/integration/test_peagle_backend_contract.py @@ -170,6 +170,101 @@ def test_peagle_vllm_guardrail() -> None: _speculative_method_from_drafter({"speculative_algorithm": "PEAGLE"}) +def _peagle_backend_and_target(draft_vocab_size=None): + from omegaconf import OmegaConf + from transformers import LlamaConfig + + from verl_speco.backends.peagle_trainer_backend import PEagleTrainerBackend + + training: dict = { + "peagle_num_draft_layers": 1, + "peagle_num_aux_hidden_states": 3, + "peagle_num_depths": 2, + } + if draft_vocab_size is not None: + training["peagle_draft_vocab_size"] = draft_vocab_size + + target_hf_config = LlamaConfig( + hidden_size=8, + intermediate_size=16, + num_attention_heads=2, + num_key_value_heads=2, + num_hidden_layers=2, + vocab_size=32, + max_position_embeddings=64, + ) + backend = PEagleTrainerBackend( + OmegaConf.create( + { + "rollout": {"drafter": {"model_path": None, "training": training}}, + "model": {"path": "/tmp/none"}, + } + ), + target_hf_config, + ) + return backend, target_hf_config + + +def test_peagle_rejects_reduced_draft_vocab_without_mapping() -> None: + """A reduced draft vocabulary needs a real frequency-derived t2d/d2t pair. + + Without one the model constructor falls back to "the first draft_vocab_size + target ids", which is an arbitrary slice of any real tokenizer and leaves the + draft unable to ever emit the rest. EAGLE-3 already refuses this + configuration; P-EAGLE used to accept it silently. + """ + pytest.importorskip("torch") + pytest.importorskip("transformers") + + backend, _ = _peagle_backend_and_target(draft_vocab_size=16) + + with pytest.raises(ValueError) as excinfo: + backend.build_model() + + message = str(excinfo.value) + assert "draft_vocab_size differs from target vocab_size" in message + # Name the algorithm the user configured, not the EAGLE-3 backend it reuses. + assert message.startswith("PEAGLE") + + +def test_peagle_identity_vocab_mapping_passes_validation() -> None: + pytest.importorskip("torch") + pytest.importorskip("transformers") + from verl_speco.models.peagle import LlamaForCausalLMPeagle + + backend, target_hf_config = _peagle_backend_and_target() + draft_config = backend._build_draft_config(None, target_hf_config) + assert draft_config.draft_vocab_size == draft_config.vocab_size + + # Must not raise: the full-vocabulary default needs no frequency mapping. + backend._validate_vocab_mapping(LlamaForCausalLMPeagle(draft_config)) + + +def test_vocab_mapping_validation_names_the_failing_algorithm() -> None: + pytest.importorskip("torch") + pytest.importorskip("transformers") + from types import SimpleNamespace + + from omegaconf import OmegaConf + + from verl_speco.backends.eagle3_trainer_backend import Eagle3TrainerBackend + from verl_speco.backends.peagle_trainer_backend import PEagleTrainerBackend + + empty_config = OmegaConf.create( + {"rollout": {"drafter": {"training": {}}}, "model": {"path": "/tmp/none"}} + ) + drafter_without_mapping = SimpleNamespace(vocab_size=32, draft_vocab_size=32) + + for backend_cls, expected_label in ( + (Eagle3TrainerBackend, "EAGLE3"), + (PEagleTrainerBackend, "PEAGLE"), + ): + backend = backend_cls(empty_config, OmegaConf.create({})) + with pytest.raises(AttributeError) as excinfo: + backend._validate_vocab_mapping(drafter_without_mapping) + assert str(excinfo.value).startswith(expected_label) + + def test_peagle_batch_assembly_matches_reference_shift() -> None: """base_trainer must apply the reference target-wrapper shift for P-EAGLE: row p pairs unshifted aux[p] with token x[p+1], supervised by the diff --git a/verl_speco/backends/eagle3_trainer_backend.py b/verl_speco/backends/eagle3_trainer_backend.py index 039d625f..ef1a6ac4 100644 --- a/verl_speco/backends/eagle3_trainer_backend.py +++ b/verl_speco/backends/eagle3_trainer_backend.py @@ -786,26 +786,29 @@ def _has_valid_vocab_mapping(self, drafter_module) -> bool: return False def _validate_vocab_mapping(self, drafter_module) -> None: + # Subclasses share this validator, so name the algorithm that actually + # failed instead of always blaming EAGLE3. + label = str(self.model_type).upper() if not hasattr(drafter_module, "t2d") or not hasattr(drafter_module, "d2t"): raise AttributeError( - "EAGLE3 draft model does not have t2d/d2t vocab mapping buffers" + f"{label} draft model does not have t2d/d2t vocab mapping buffers" ) if drafter_module.t2d.numel() != drafter_module.vocab_size: raise ValueError( - f"EAGLE3 t2d shape mismatch: expected {drafter_module.vocab_size}, " + f"{label} t2d shape mismatch: expected {drafter_module.vocab_size}, " f"got {drafter_module.t2d.numel()}" ) if drafter_module.d2t.numel() != drafter_module.draft_vocab_size: raise ValueError( - f"EAGLE3 d2t shape mismatch: expected {drafter_module.draft_vocab_size}, " + f"{label} d2t shape mismatch: expected {drafter_module.draft_vocab_size}, " f"got {drafter_module.d2t.numel()}" ) selected_vocab_size = int(drafter_module.t2d.sum().item()) if selected_vocab_size != drafter_module.draft_vocab_size: raise ValueError( - f"EAGLE3 vocab mapping selects {selected_vocab_size} tokens, " + f"{label} vocab mapping selects {selected_vocab_size} tokens, " f"but draft_vocab_size is {drafter_module.draft_vocab_size}" ) diff --git a/verl_speco/backends/peagle_trainer_backend.py b/verl_speco/backends/peagle_trainer_backend.py index 8fbb09b9..43530469 100644 --- a/verl_speco/backends/peagle_trainer_backend.py +++ b/verl_speco/backends/peagle_trainer_backend.py @@ -256,18 +256,49 @@ def build_model(self): draft_config = self._build_draft_config(spec_model_path, target_hf_config) self.vocab_size = draft_config.vocab_size + checkpoint_has_vocab_mapping = False if spec_model_path and os.path.exists( os.path.join(spec_model_path, "config.json") ): log_drafter_checkpoint_step( logger, spec_model_path, action="Loading P-EAGLE drafter weights" ) - drafter_module = LlamaForCausalLMPeagle.from_pretrained( - spec_model_path, config=draft_config + loaded = LlamaForCausalLMPeagle.from_pretrained( + spec_model_path, + config=draft_config, + output_loading_info=True, ) + if isinstance(loaded, tuple): + drafter_module, loading_info = loaded + missing_keys = set(loading_info.get("missing_keys", [])) + checkpoint_has_vocab_mapping = not {"t2d", "d2t"}.intersection( + missing_keys + ) + else: + drafter_module = loaded + checkpoint_has_vocab_mapping = self._has_valid_vocab_mapping( + drafter_module + ) else: drafter_module = LlamaForCausalLMPeagle(draft_config) + # A reduced draft vocabulary is only meaningful with a real t2d/d2t pair + # derived from token frequency. The model's constructor falls back to + # "the first draft_vocab_size target ids", which is an arbitrary slice of + # any real tokenizer, so refuse it exactly like the EAGLE-3 backend does + # instead of silently training a draft that can never emit the rest. + if drafter_module.draft_vocab_size != drafter_module.vocab_size: + if checkpoint_has_vocab_mapping and self._has_valid_vocab_mapping( + drafter_module + ): + logger.debug("Using P-EAGLE vocab mapping loaded from draft checkpoint") + else: + raise ValueError( + "PEAGLE draft_vocab_size differs from target vocab_size, but the draft " + "checkpoint does not provide valid t2d/d2t vocab mapping buffers" + ) + self._validate_vocab_mapping(drafter_module) + # P-EAGLE trains the draft embeddings (speculators sets embed_requires_grad=True), # so seed them from the target but do NOT freeze. drafter_module.load_embedding(self.config.model.path)