feat: add P-EAGLE parallel-drafting drafter backend

[khazic]

What

Adds a P-EAGLE (parallel-drafting EAGLE) drafter training backend to the SpeCo overlay (part of #16). Select it with:

actor_rollout_ref.rollout.drafter.speculative_algorithm=PEAGLE

Design

Unlike EAGLE-3's sequential test-time-training recurrence, P-EAGLE predicts all num_depths tokens in a single parallel forward over a flat, COD-subsampled sequence, supervised by a count-normalized KL(target || draft) over the draft vocabulary (no per-depth decay). Logic follows the NeMo AutoModel reference (peagle_data.py / peagle_attention.py / peagle_draft.py / peagle_trainer.py, itself a port of speculators PR #480).

New verl_speco/models/peagle/:

• cod_sampling.py and peagle_mask.py: the COD geometric subsampling and the flex-attention COD block mask (verbatim ports). Depth 0 keeps every position; depth d keeps a down_sample_ratio**d fraction. The mask lets each element attend to the causal depth-0 context of its document plus earlier-or-equal depths of its own rollout.
• modeling_peagle.py: LlamaForCausalLMPeagle with a fused layer 0 ([embed, hidden] -> 2H attention), vanilla deep layers, a single learnable mask_hidden placeholder that substitutes for the target aux feature at masked (depth>=1) slots, and flex attention driven by the COD block mask.

PEagleTrainerBackend subclasses Eagle3TrainerBackend and reuses the aux + last-hidden collection, but reports model_type == "peagle" so base_trainer assembles the UNSHIFTED per-position batch P-EAGLE needs (aux f[p], token x[p], target distribution for x[p+1]). The frozen target head turns last_hidden_states into the full-vocab target logits, which are restricted to the draft vocab via t2d. Only build_model and compute_loss differ; preprocess/optimizer/target-head are inherited.

flex_attention is run eagerly: the COD block mask captures per-element index tensors and the compiled (max-autotune, dynamic) backward is unstable on it (illegal memory access); an opt-in VERL_PEAGLE_COMPILE_FLEX=1 keeps the compile path for later verification.

P-EAGLE targets vLLM's parallel-drafting runtime for serving; the vLLM config builder raises until that runtime is wired here, so this PR is training-side.

Not a duplicate

No open PR touches P-EAGLE (gh pr list); #13 and #10 are unrelated.

Testing

A human reviewed every changed line and ran the tests below.

CPU contract tests (tests/integration/test_peagle_backend_contract.py), 6 passed: COD sampling structure (depth-0 keeps all positions, deeper depths shrink geometrically), the flex-mask predicate (depth-0 causal context + own-rollout depth order), the draft-model modules (fc/mask_hidden/lm_head/selected_token_ids), the KL loss, the backend metadata, and the vLLM guardrail. The full suite is 94 passed, 2 skipped, 2 failed; both failures pre-exist on main and are unrelated.

pytest tests/integration/test_peagle_backend_contract.py -q
# 6 passed

GPU hardware smoke (ci/peagle_gpu_smoke.py) on a real target (Qwen3-4B, 1x A100): it collects the unshifted per-position data (3 aux layers + the final hidden + tokens + loss mask) from the frozen target, builds the P-EAGLE draft via PEagleTrainerBackend.build_model, and runs 150 optimizer steps through compute_loss (COD sampling -> flat 8-depth flex-attention forward -> count-normalized KL). The draft is cold-started, so the signal is convergence: the KL loss falls and draft-vs-target top-1 agreement rises across the 8 parallel depths.

python ci/peagle_gpu_smoke.py --target /path/to/Qwen3-4B --steps 150

P-EAGLE GPU training log (Qwen3-4B target, cold start, 8 depths, 150 steps)

[smoke] aux layers=[2, 18, 33] (of 36)
[smoke] batch seq_len=88 aux=7680
[smoke] draft_layers=2 num_depths=8 fc=7680->2560 trainable_params=1,015,175,680
[smoke] step   0  kl_loss=10.5380  draft_vs_target_top1=0.0000
[smoke] step  15  kl_loss=2.5845  draft_vs_target_top1=0.2822
[smoke] step  30  kl_loss=0.8606  draft_vs_target_top1=0.7108
[smoke] step  45  kl_loss=0.2523  draft_vs_target_top1=0.8432
[smoke] step  60  kl_loss=0.1088  draft_vs_target_top1=0.8850
[smoke] step  75  kl_loss=0.0611  draft_vs_target_top1=0.8979
[smoke] step  90  kl_loss=0.0370  draft_vs_target_top1=0.9059
[smoke] step 105  kl_loss=0.0239  draft_vs_target_top1=0.9439
[smoke] step 120  kl_loss=0.0216  draft_vs_target_top1=0.9199
[smoke] step 135  kl_loss=0.0258  draft_vs_target_top1=0.9059
[smoke] step 149  kl_loss=0.0165  draft_vs_target_top1=0.9053
[smoke] DONE  kl_loss 10.5380->0.0165  top1 0.0000->0.9053

Known follow-ups

• Serving: needs vLLM's parallel-drafting runtime (speculators PR #480 format; the checkpoint saves mask_hidden and the COD config for that runtime).
• The compiled flex-attention path is opt-in (VERL_PEAGLE_COMPILE_FLEX=1) pending a verified compiled backward.
• Sequence partitioning (Algorithm 1) for long-context memory is not ported yet; the single flat forward covers the common case.

[gemini-code-assist]

Code Review

This pull request introduces the P-EAGLE (parallel-drafting) drafter training backend, including its configuration, modeling with Conditional-On-Distribution (COD) sampling, custom flex-attention masking, integration into the base trainer, and associated tests. The review feedback identifies a critical issue where independent sequences in a batch are treated as a single document, causing cross-document attention leakage; to resolve this, individual sequence lengths must be preserved and passed in the batch. Additionally, several minor improvements are suggested, such as removing a redundant deepcopy call, using torch.full for cleaner tensor initialization, and replacing squeeze(0) with flatten() for safer tensor reshaping.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Because base_trainer.py concatenates all sequences in the batch into a single flat sequence of batch size 1, attention_mask[0] is a flat tensor of ones of length seq_len (the sum of all sequence lengths). As a result, row_length is computed as a single-element tensor [seq_len]. When passed to create_peagle_mask_mod, this causes document_ids to be a tensor of all 0s of length seq_len, meaning all independent sequences in the batch are treated as a single document. This allows tokens from independent prompts/documents to cross-attend to each other, leading to severe document leakage and training contamination. To fix this, retrieve the individual sequence lengths from the batch.

            row_length = batch.get(\"seq_lengths\")
            if row_length is None:
                row_length = attention_mask[b].sum().clamp_min(1).reshape(1).to(device)

[gemini-code-assist]

To support proper document isolation and prevent cross-document attention leakage in P-EAGLE training, we should preserve the individual sequence lengths of the concatenated chunks and pass them in the batch.

        elif self.backend.model_type == \"peagle\":
            batch[\"last_hidden_states\"] = last_hidden_states
            batch[\"seq_lengths\"] = torch.tensor([chunk.size(0) for chunk in input_id_chunks], dtype=torch.long, device=dev)

[gemini-code-assist]

Calling deepcopy on target_hf_config before calling .to_dict() is redundant because .to_dict() already creates and returns a new dictionary copy of the configuration parameters.

Suggested change

	cfg_dict = deepcopy(target_hf_config).to_dict()
	cfg_dict = target_hf_config.to_dict()

[gemini-code-assist]

Using torch.full is more direct and idiomatic for creating a tensor filled with a specific value than multiplying torch.ones by -1.

Suggested change

	-1 * torch.ones(total_seq_len - document_ids.shape[0], device=lengths.device, dtype=torch.long),
	torch.full((total_seq_len - document_ids.shape[0],), -1, device=lengths.device, dtype=torch.long),

[gemini-code-assist]

Using .flatten() or .view(-1) is safer than .squeeze(0) because it consistently returns a 1D tensor even if the input tensor has a batch size or sequence length of 1.

Suggested change

	loss_mask = loss_mask.squeeze(0)
	loss_mask = loss_mask.flatten()

[khazic]

Thanks for the review. Verified and fixed the high-priority one against the code path:

Cross-document leakage (real). base_trainer takes the flat non-block path for peagle (only dflash/dspark are block drafters), concatenating every document into one batch-1 sequence with an all-ones attention_mask, so the COD mask saw a single document and depth-0 positions leaked across documents. Fixed by passing per-chunk seq_lengths from base_trainer and building document_ids from them (fallback to single-document only when absent). Added a regression test asserting a cross-document depth-0 pair is masked while the merged-length variant leaks.

Also applied the three nits: dropped the redundant deepcopy before .to_dict() (and its now-unused import), torch.full for the padding fill, and .flatten() over .squeeze(0).

Pushed in 0acc2a0.

[khazic]

Pushed two fixes to this branch.

6691ac1 fix(peagle): apply the reference target-wrapper shift in batch assembly. The trainer here is a verbatim port of the reference _peagle_position_loss, but the reference feeds it shifted tensors: its target wrapper left-shifts input_ids, target logits, and loss_mask by one (_shift_left_with_zero), keeping only the aux hidden states unshifted. The previous assembly was fully unshifted, so it trained (aux[p], x[p]) -> P(x[p+1]) while the parallel-drafting runtime feeds (aux[t], x[t+1]) at inference (the newest sampled token has no hidden state yet). The assembly now mirrors the reference shift exactly: ids[1:1+T], aux[:T], last_hidden[1:1+T], loss_mask[1:1+T], and a contract test pins these slices.

8095358 fix(peagle): keep last_hidden_states and seq_lengths in the final batch. _prepare_training_batch rebuilds the batch dict after the Ulysses section and only re-attached the eagle3 keys, so the peagle last_hidden_states / seq_lengths set earlier were silently dropped and the real assembly path failed compute_loss's check on the first step. The GPU smoke drives the backend directly and never hit this path, which is why it passed. The new assembly contract test catches it.

Note: the training-side top1/loss from the earlier smoke were measured on the unshifted pairing, so they are not comparable to post-fix numbers; the shifted task is strictly harder (the draft no longer sees the hidden state of the position right before the predicted token). I will re-run the GPU smoke when cluster access is back.