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

[khazic]

Problem

The hot-publish filter decides what to ship to the rollout engine by tensor name and by model_type:

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 self._is_frozen_publish_param(name) or (
    "lm_head.weight" in name
    and getattr(self.backend, "model_type", None) != "eagle3"
):
    continue

Both rules are wrong for P-EAGLE:

• Its draft owns an lm_head over the draft vocabulary (LlamaForCausalLMPeagle.lm_head), and model_type is "peagle", so the head is filtered out.
• It is the only backend that deliberately does not call freeze_embedding(). build_model seeds the embedding from the target and leaves it trainable (speculators sets embed_requires_grad=True), so the embedding is filtered out too.

Both tensors receive a gradient on every drafter step. The rollout engine therefore keeps the initial head and the target-seeded embedding for the whole run while training moves on, which is a silent train/serve divergence in the two most consequential tensors of the draft.

model_type is also not a usable proxy for this question in general: Eagle1TrainerBackend.model_type deliberately returns "eagle3" so it can reuse the EAGLE-3 data plumbing, so the check is already reading a value that does not mean what it looks like it means.

Fix

Replace the name/model_type heuristic with an explicit per-backend contract that each backend declares next to the freeze calls it actually makes:

backend	trains_draft_lm_head	trains_draft_embeddings
Eagle3TrainerBackend	True	False
Eagle1TrainerBackend	False (weight-tied to the target head)	False
PEagleTrainerBackend	True	True
DFlashTrainerBackend / DSpark / Domino	False	False

Offline export is unaffected: it goes through _get_full_export_state_dict, which already keeps frozen parameters and persistent buffers.

Scope

P-EAGLE rollout is not wired up yet (_speculative_method_from_drafter raises for PEAGLE and points users at training the drafter with drafter.enable=false), so this is a latent bug rather than a live one today. It becomes a live silent divergence the moment the parallel-drafting runtime lands, which is why it is worth closing now.

Validation

Ran a before/after repro on both main and this branch. It builds a tiny P-EAGLE draft, runs the real publish filter over its state dict with a real PEagleTrainerBackend, and reports which trained tensors reach the engine.

Before/after repro output

=================== before: main (333b754) ===================
[which backends freeze the draft embedding in build_model]
          Eagle3TrainerBackend     freezes
          Eagle1TrainerBackend     freezes
          PEagleTrainerBackend     TRAINS IT
          DFlashTrainerBackend     freezes
          DSparkTrainerBackend     freezes
          DominoTrainerBackend     freezes

[trained parameters on the P-EAGLE draft]
          embed_tokens.weight      requires_grad=True
          lm_head.weight           requires_grad=True
          fc.weight                requires_grad=True
          mask_hidden              requires_grad=True

[hot publish, P-EAGLE]
          trained tensors    : 24
          published tensors  : 22
          dropped tensors    : 2
                             - embed_tokens.weight
                             - lm_head.weight

[verdict]
          embed_tokens.weight      DROPPED
          lm_head.weight           DROPPED
          RESULT: trained P-EAGLE weights never reach the engine (bug present)

=================== after: this branch ===================
[which backends freeze the draft embedding in build_model]
          Eagle3TrainerBackend     freezes
          Eagle1TrainerBackend     freezes
          PEagleTrainerBackend     TRAINS IT
          DFlashTrainerBackend     freezes
          DSparkTrainerBackend     freezes
          DominoTrainerBackend     freezes

[trained parameters on the P-EAGLE draft]
          embed_tokens.weight      requires_grad=True
          lm_head.weight           requires_grad=True
          fc.weight                requires_grad=True
          mask_hidden              requires_grad=True

[hot publish, P-EAGLE]
          trained tensors    : 24
          published tensors  : 24
          dropped tensors    : 0

[verdict]
          embed_tokens.weight      published
          lm_head.weight           published
          RESULT: every trained P-EAGLE weight is published (fixed)

Tests

Two new tests in tests/integration/test_rollout_publish_contract.py:

• test_publish_state_filter_keeps_peagle_trained_head_and_embedding covers the P-EAGLE publish set.
• test_backend_publish_contract_matches_what_each_backend_trains pins the declared flags on every backend, so a new backend cannot silently inherit the wrong answer.

The three existing test_publish_state_filter_* tests build their backend as a SimpleNamespace, so they now declare the two flags explicitly. That is the point of the change: the stub has to state the contract it is exercising instead of relying on a model_type string.

Full CPU suite (tests/integration tests/compat tests/config tests/examples) run on both sides:

Test suite before/after

### BASELINE (origin/main, 333b754) ###
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[DSPARK-veomni-npu-veomni_lm_head_full]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[DFLASH-veomni-npu-veomni_lm_head_sparse]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[EAGLE3-veomni-cuda-veomni_lm_head_full]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[EAGLE1-fsdp-npu-engine_full_param]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[DOMINO-fsdp2-cuda-engine_full_param]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_transfer_waits_after_actor_update
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_async_publish_sets_pending_ref_and_waits_before_next_publish
FAILED tests/integration/test_dspark_trainer_backend.py::test_dspark_checkpoint_preserves_source_config_and_vllm_weight_names
FAILED tests/integration/test_verl_npu_vllm_compat.py::test_factory_fused_moe_survives_verl_npu_patch_import
9 failed, 245 passed, 2 warnings in 33.60s

### THIS BRANCH ###
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[DSPARK-veomni-npu-veomni_lm_head_full]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[DFLASH-veomni-npu-veomni_lm_head_sparse]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[EAGLE3-veomni-cuda-veomni_lm_head_full]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[EAGLE1-fsdp-npu-engine_full_param]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[DOMINO-fsdp2-cuda-engine_full_param]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_transfer_waits_after_actor_update
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_async_publish_sets_pending_ref_and_waits_before_next_publish
FAILED tests/integration/test_dspark_trainer_backend.py::test_dspark_checkpoint_preserves_source_config_and_vllm_weight_names
FAILED tests/integration/test_verl_npu_vllm_compat.py::test_factory_fused_moe_survives_verl_npu_patch_import
9 failed, 247 passed, 2 warnings in 20.50s

The same 9 tests fail on main and on this branch. They need optional dependencies (VeOmni, the NPU vLLM stack) that are absent in this environment, so they are pre-existing and unrelated. The +2 on this branch are the new tests above.