fix(peagle): refuse a reduced draft vocabulary without a t2d/d2t mapping

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

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}"
             )
 

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)