fix(peagle): publish the drafter head and embedding P-EAGLE actually trains

tests/integration/test_rollout_publish_contract.py

@@ -247,7 +247,11 @@ def test_publish_state_filter_keeps_eagle3_trainable_lm_head() -> None:
     DrafterBaseTrainer = base_trainer.DrafterBaseTrainer
 
     trainer = DrafterBaseTrainer.__new__(DrafterBaseTrainer)
-    trainer.backend = SimpleNamespace(model_type="eagle3")
+    trainer.backend = SimpleNamespace(
+        model_type="eagle3",
+        trains_draft_lm_head=True,
+        trains_draft_embeddings=False,
+    )
     trainer.training_device_mesh = None
     trainer._frozen_param_names = ["target_model."]
     trainer.model = SimpleNamespace(
@@ -275,7 +279,11 @@ def test_publish_state_filter_skips_non_eagle_lm_head() -> None:
     DrafterBaseTrainer = base_trainer.DrafterBaseTrainer
 
     trainer = DrafterBaseTrainer.__new__(DrafterBaseTrainer)
-    trainer.backend = SimpleNamespace(model_type="dflash")
+    trainer.backend = SimpleNamespace(
+        model_type="dflash",
+        trains_draft_lm_head=False,
+        trains_draft_embeddings=False,
+    )
     trainer.training_device_mesh = None
     trainer._frozen_param_names = []
     trainer.model = SimpleNamespace(
@@ -297,7 +305,11 @@ def test_publish_state_filter_excludes_block_drafter_embedding() -> None:
     DrafterBaseTrainer = base_trainer.DrafterBaseTrainer
 
     trainer = DrafterBaseTrainer.__new__(DrafterBaseTrainer)
-    trainer.backend = SimpleNamespace(model_type="dspark")
+    trainer.backend = SimpleNamespace(
+        model_type="dspark",
+        trains_draft_lm_head=False,
+        trains_draft_embeddings=False,
+    )
     trainer.training_device_mesh = None
     trainer._frozen_param_names = []
     trainer.model = SimpleNamespace(
@@ -310,6 +322,77 @@ def test_publish_state_filter_excludes_block_drafter_embedding() -> None:
     assert set(trainer._get_trainable_state_dict()) == {"draft_model.fc.weight"}
 
 
+def test_publish_state_filter_keeps_peagle_trained_head_and_embedding() -> None:
+    """P-EAGLE owns its lm_head and fine-tunes the draft embedding.
+
+    Both are trained every drafter step, so hot publish has to ship them; the
+    generic filter used to drop them and leave the rollout engine on the initial
+    weights forever.
+    """
+    torch = pytest.importorskip("torch")
+    base_trainer = pytest.importorskip(
+        "verl_speco.trainer.base_trainer",
+        reason="publish state filtering needs the trainer dependency stack",
+    )
+    DrafterBaseTrainer = base_trainer.DrafterBaseTrainer
+
+    trainer = DrafterBaseTrainer.__new__(DrafterBaseTrainer)
+    trainer.backend = SimpleNamespace(
+        model_type="peagle",
+        trains_draft_lm_head=True,
+        trains_draft_embeddings=True,
+    )
+    trainer.training_device_mesh = None
+    trainer._frozen_param_names = ["target_model."]
+    trainer.model = SimpleNamespace(
+        state_dict=lambda: {
+            "embed_tokens.weight": torch.ones(2, 2),
+            "lm_head.weight": torch.ones(2, 2),
+            "fc.weight": torch.ones(2, 2),
+            "mask_hidden": torch.ones(1, 1, 2),
+            "target_model.fc.weight": torch.ones(2, 2),
+            "t2d": torch.ones(2, dtype=torch.bool),
+        }
+    )
+
+    assert set(trainer._get_trainable_state_dict()) == {
+        "embed_tokens.weight",
+        "lm_head.weight",
+        "fc.weight",
+        "mask_hidden",
+    }
+
+
+def test_backend_publish_contract_matches_what_each_backend_trains() -> None:
+    """The declared flags must track the real freeze/build calls in each backend.
+
+    ``model_type`` is not a usable proxy here: EAGLE-1/2 deliberately report
+    ``"eagle3"`` to reuse the data plumbing, and P-EAGLE is the only backend that
+    skips ``freeze_embedding()``.
+    """
+    pytest.importorskip("torch")
+    pytest.importorskip("transformers")
+
+    from verl_speco.backends.dflash_trainer_backend import DFlashTrainerBackend
+    from verl_speco.backends.domino_trainer_backend import DominoTrainerBackend
+    from verl_speco.backends.dspark_trainer_backend import DSparkTrainerBackend
+    from verl_speco.backends.eagle1_trainer_backend import Eagle1TrainerBackend
+    from verl_speco.backends.eagle3_trainer_backend import Eagle3TrainerBackend
+    from verl_speco.backends.peagle_trainer_backend import PEagleTrainerBackend
+
+    expected = {
+        Eagle3TrainerBackend: (True, False),
+        Eagle1TrainerBackend: (False, False),
+        PEagleTrainerBackend: (True, True),
+        DFlashTrainerBackend: (False, False),
+        DSparkTrainerBackend: (False, False),
+        DominoTrainerBackend: (False, False),
+    }
+    for backend_cls, (lm_head, embeddings) in expected.items():
+        assert backend_cls.trains_draft_lm_head is lm_head, backend_cls.__name__
+        assert backend_cls.trains_draft_embeddings is embeddings, backend_cls.__name__
+
+
 def test_target_lm_head_device_helper_handles_dflash_style_backend() -> None:
     base_trainer = pytest.importorskip(
         "verl_speco.trainer.base_trainer",

verl_speco/backends/dflash_trainer_backend.py

@@ -617,6 +617,12 @@ def forward(
 
 
 class DFlashTrainerBackend:
+    # Hot-publish contract: the block drafters read logits off the frozen target
+    # head and keep the target-seeded embedding frozen, so neither belongs in the
+    # published delta. DSpark and Domino inherit this.
+    trains_draft_lm_head = False
+    trains_draft_embeddings = False
+
     def __init__(self, config, target_model_config):
         self.config = config
         self.target_model_config = target_model_config

verl_speco/backends/eagle1_trainer_backend.py

@@ -72,6 +72,10 @@ def model_type(self):
     # forces ulysses_sequence_parallel_size = 1 instead of aborting mid-run.
     supports_ulysses_sp = False
 
+    # Token logits come from the frozen target head (weight tying), so the draft
+    # carries no lm_head of its own to publish.
+    trains_draft_lm_head = False
+
     def _build_draft_config(self, spec_model_path, target_hf_config):
         config_path = (
             os.path.join(spec_model_path, "config.json") if spec_model_path else None

verl_speco/backends/eagle3_trainer_backend.py

@@ -622,6 +622,13 @@ def _apply_coverage_mask_to_loss_mask(
 
 
 class Eagle3TrainerBackend:
+    # Hot-publish contract: which target-seeded drafter tensors this backend owns
+    # and therefore has to ship to the rollout engine. The EAGLE-3 draft has its
+    # own lm_head over the draft vocabulary, but seeds the embedding from the
+    # target and freezes it.
+    trains_draft_lm_head = True
+    trains_draft_embeddings = False
+
     def __init__(self, config, target_model_config):
         self.config = config
         self.target_model_config = target_model_config

verl_speco/backends/peagle_trainer_backend.py

@@ -203,6 +203,11 @@ def model_type(self):
     # P-EAGLE trains on full local sequences and does not implement the SP loss.
     supports_ulysses_sp = False
 
+    # Unlike EAGLE-3, P-EAGLE fine-tunes the draft embedding instead of freezing
+    # the target-seeded copy (speculators sets embed_requires_grad=True), so hot
+    # publish has to carry it along with the draft's own lm_head.
+    trains_draft_embeddings = True
+
     def _training_cfg(self):
         return self.config.rollout.drafter.training
 

verl_speco/trainer/base_trainer.py

@@ -1257,7 +1257,12 @@ def _is_frozen_publish_param(self, name: str) -> bool:
 
         if any(frozen_name in name for frozen_name in self._frozen_param_names):
             return True
-        return name == "embed_tokens.weight" or name.endswith(".embed_tokens.weight")
+        if name == "embed_tokens.weight" or name.endswith(".embed_tokens.weight"):
+            # Most backends seed the draft embedding from the target and freeze it,
+            # so the engine already holds the same rows. Backends that fine-tune it
+            # (P-EAGLE) must publish it or the engine keeps the seed forever.
+            return not bool(getattr(self.backend, "trains_draft_embeddings", False))
+        return False
 
     def _get_trainable_state_dict(self) -> dict[str, torch.Tensor]:
         """Get floating state dict entries excluding weights shared with the target model."""
@@ -1280,10 +1285,11 @@ def _get_trainable_state_dict(self) -> dict[str, torch.Tensor]:
                     f"Skipping non-floating drafter state: {name}, dtype={param.dtype}"
                 )
                 continue
-            # EAGLE3 trains and publishes its own lm_head; other backends skip lm_head by default.
+            # Backends whose draft owns an lm_head (EAGLE-3, P-EAGLE) publish it;
+            # the block drafters read logits off the target head, so theirs stays out.
             if self._is_frozen_publish_param(name) or (
                 "lm_head.weight" in name
-                and getattr(self.backend, "model_type", None) != "eagle3"
+                and not bool(getattr(self.backend, "trains_draft_lm_head", False))
             ):
                 logger.debug(f"Skipping frozen parameter: {name}")
                 continue