feat: add EAGLE-1 / EAGLE-2 drafter training backend by khazic · Pull Request #15 · verl-project/verl-SpeCo

khazic · 2026-07-15T06:15:40Z

What

Adds an EAGLE-1 / EAGLE-2 drafter training backend to the SpeCo overlay, closing #14. The Draft Model Support matrix now covers EAGLE-1, EAGLE-2, EAGLE-3, DFLASH and DSpark.

Select it with:

actor_rollout_ref.rollout.drafter.speculative_algorithm=EAGLE1   # or EAGLE2

Design

The training objective follows the EAGLE-1/2 formulation (single-step feature prediction), which is unchanged between the two algorithms; EAGLE-2 differs from EAGLE-1 only in the inference-time speculative tree policy (a dynamic, context-aware tree), which lives in the rollout engine, not in drafter training. Both map to vLLM's native eagle method.

The dense draft (verl_speco/models/eagle1/) fuses the token embedding with the target's last-layer hidden state through a single fc layer (hidden_size + target_hidden_size -> hidden_size), runs one standard decoder layer, and predicts the target's next-step hidden state. It carries no lm_head: token logits come from the frozen target head (weight tying), matching the original EAGLE design. The loss combines a SmoothL1 feature-regression term (weight 1.0) with a full-vocabulary soft cross-entropy distillation term (weight 0.1), plus optional EAGLE feature-noise augmentation.

Eagle1TrainerBackend reuses the EAGLE-3 online data-collection plumbing by reporting model_type == "eagle3", so base_trainer assembles the shifted inputs and the last_hidden_states regression target; only build_model and compute_loss differ. Old-logprob hidden collection captures a single final target layer for EAGLE-1/2 (versus EAGLE-3's low/mid/high triple). The soft cross-entropy reuses the EAGLE-3 helper, so non-finite logits are sanitized before log_softmax and masked positions cannot backprop NaN gradients.

Not a duplicate

Checked gh issue view 14 (open, unassigned, no competing work) and the open PRs (gh pr list). #13 (Separate Draft Model Training) and #10 (draft weight loading) are unrelated; no open PR touches EAGLE-1/2.

Testing

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

CPU contract tests (tests/integration/test_eagle1_backend_contract.py), 12 passed: draft forward shape, loss equivalence against the reference EAGLE-1 formula, the use_logits=False requirement, config-time rejection of use_logits=True, the Ulysses-SP opt-out, EAGLE-1/EAGLE-2 to eagle vLLM method mapping, and single-final-layer aux selection (including ignoring a stray multi-layer config). The full suite is 100 passed, 2 skipped, 2 failed; both failures pre-exist on main (a tests/heavy runner-env check and a stale metrics assertion) and are unrelated to this change.

pytest tests/integration/test_eagle1_backend_contract.py -q
# 12 passed

GPU hardware smoke (ci/eagle1_gpu_smoke.py) on a real target (Qwen3-4B, 1x A100): it runs the frozen target forward to produce genuine last-layer hidden states, assembles the batch exactly as base_trainer does for the EAGLE shifted path, builds the draft via Eagle1TrainerBackend.build_model, and runs 100 optimizer steps. The draft is cold-started, so the signal is convergence: SmoothL1 feature loss and soft-CE token loss fall and draft-vs-target top-1 agreement rises. EAGLE-2 produces identical numbers (same training path).

python ci/eagle1_gpu_smoke.py --target /path/to/Qwen3-4B --algorithm EAGLE1 --steps 100

EAGLE-1 GPU training log (Qwen3-4B target, cold start, 100 steps)

[smoke] loading target /path/to/Qwen3-4B
[smoke] batch seq_len=66 hidden=2560
[smoke] algorithm=EAGLE1 draft_layers=1 fc=5120->2560 trainable_params=114,040,320
[smoke] step   0  hidden_loss=1.4905  token_loss=12.1883  loss=2.7094  draft_vs_target_top1=0.1364
[smoke] step  10  hidden_loss=1.1124  token_loss=2.7478  loss=1.3872  draft_vs_target_top1=0.8485
[smoke] step  20  hidden_loss=0.9030  token_loss=2.2944  loss=1.1324  draft_vs_target_top1=0.8636
[smoke] step  30  hidden_loss=0.7947  token_loss=2.1330  loss=1.0080  draft_vs_target_top1=0.9394
[smoke] step  40  hidden_loss=0.7217  token_loss=2.0425  loss=0.9260  draft_vs_target_top1=0.9545
[smoke] step  50  hidden_loss=0.6683  token_loss=2.0004  loss=0.8683  draft_vs_target_top1=0.9394
[smoke] step  60  hidden_loss=0.6311  token_loss=1.9821  loss=0.8293  draft_vs_target_top1=0.9545
[smoke] step  70  hidden_loss=0.6082  token_loss=1.9735  loss=0.8055  draft_vs_target_top1=0.9545
[smoke] step  80  hidden_loss=0.5965  token_loss=1.9690  loss=0.7934  draft_vs_target_top1=0.9545
[smoke] step  90  hidden_loss=0.5918  token_loss=1.9644  loss=0.7883  draft_vs_target_top1=0.9545
[smoke] step  99  hidden_loss=0.5903  token_loss=1.9612  loss=0.7864  draft_vs_target_top1=0.9545
[smoke] DONE  hidden_loss 1.4905->0.5903  token_loss 12.1883->1.9612  top1=0.9545

EAGLE-2 GPU training log (Qwen3-4B target, cold start, 100 steps) — identical training path

[smoke] loading target /path/to/Qwen3-4B
[smoke] batch seq_len=66 hidden=2560
[smoke] algorithm=EAGLE2 draft_layers=1 fc=5120->2560 trainable_params=114,040,320
[smoke] step   0  hidden_loss=1.4905  token_loss=12.1883  loss=2.7094  draft_vs_target_top1=0.1364
[smoke] step  10  hidden_loss=1.1124  token_loss=2.7478  loss=1.3872  draft_vs_target_top1=0.8485
[smoke] step  20  hidden_loss=0.9030  token_loss=2.2944  loss=1.1324  draft_vs_target_top1=0.8636
[smoke] step  30  hidden_loss=0.7947  token_loss=2.1330  loss=1.0080  draft_vs_target_top1=0.9394
[smoke] step  40  hidden_loss=0.7217  token_loss=2.0425  loss=0.9260  draft_vs_target_top1=0.9545
[smoke] step  50  hidden_loss=0.6683  token_loss=2.0004  loss=0.8683  draft_vs_target_top1=0.9394
[smoke] step  60  hidden_loss=0.6311  token_loss=1.9821  loss=0.8293  draft_vs_target_top1=0.9545
[smoke] step  70  hidden_loss=0.6082  token_loss=1.9735  loss=0.8055  draft_vs_target_top1=0.9545
[smoke] step  80  hidden_loss=0.5965  token_loss=1.9690  loss=0.7934  draft_vs_target_top1=0.9545
[smoke] step  90  hidden_loss=0.5918  token_loss=1.9644  loss=0.7883  draft_vs_target_top1=0.9545
[smoke] step  99  hidden_loss=0.5903  token_loss=1.9612  loss=0.7864  draft_vs_target_top1=0.9545
[smoke] DONE  hidden_loss 1.4905->0.5903  token_loss 12.1883->1.9612  top1=0.9545

Known follow-ups

Ulysses sequence parallelism is not yet implemented for the EAGLE-1/2 loss; the backend opts out so it is forced off rather than aborting.
The draft uses eager attention (as in the reference implementation); a flex-attention path like EAGLE-3's could be added for long sequences.
SGLang rollout wiring for EAGLE-1/2 can follow; this PR targets the vLLM path.

Port the EAGLE-1/2 draft-training objective from NeMo AutoModel's v12 path into the SpeCo overlay. The dense draft fuses the token embedding with the target's last-layer hidden state through a single fc layer, runs one standard decoder layer, and predicts the target's next-step hidden state; token logits are produced by the frozen target head (weight tying). The loss combines a SmoothL1 feature regression term with a full-vocabulary soft cross-entropy distillation term, with optional EAGLE feature-noise augmentation. EAGLE-1 and EAGLE-2 share this training path; they differ only in the inference-time speculative tree policy (EAGLE-2 uses a dynamic tree), which lives in the rollout engine. Both map to vLLM's native eagle method. The backend reuses the EAGLE-3 online data-collection plumbing by reporting model_type == eagle3, so base_trainer assembles the shifted inputs and the last_hidden_states regression target; only build_model and compute_loss differ. Old-logprob hidden collection captures a single final target layer for EAGLE-1/2. AI assistance was used for this change. Signed-off-by: khazic <khazzz1c@gmail.com>

AI assistance was used for this change. Signed-off-by: khazic <khazzz1c@gmail.com>

Drives build_model + compute_loss + optimizer on a real target model with real target hidden states, per the repo convention of keeping hardware smokes under ci/. AI assistance was used for this change. Signed-off-by: khazic <khazzz1c@gmail.com>

Make --target required and use a temp dir for the cold-start drafter path instead of hardcoded absolute paths. AI assistance was used for this change. Signed-off-by: khazic <khazzz1c@gmail.com>

- compute_loss: reuse eagle3 _masked_soft_cross_entropy and sanitize hidden states before SmoothL1 so masked non-finite positions cannot backprop NaN gradients into the draft. - oldlogprob_layer_ids: resolve the EAGLE-1/2 single-final-layer selection before the generic multi-layer config lookup, so a stray eagle3-style aux-layer id set cannot silently pick the wrong distill target. - build_model: reject use_logits=True at config time instead of crash-looping on the first step (the eagle3 data path yields no last_hidden_states then). - draft config advertises the draft depth, not the target's, in config.json. - opt out of Ulysses SP via supports_ulysses_sp so base_trainer forces it off rather than aborting compute_loss under rollout_tp>1. AI assistance was used for this change. Signed-off-by: khazic <khazzz1c@gmail.com>

gemini-code-assist

Code Review

This pull request introduces the training backend and model architecture for EAGLE-1 and EAGLE-2 speculative drafting. It adds the Eagle1TrainerBackend which leverages the existing EAGLE-3 data plumbing but implements a single-step feature regression and soft cross-entropy distillation loss against a frozen target head. Additionally, it includes the LlamaForCausalLMEagle1 draft model, configuration files, integration tests, and a GPU smoke test. The reviewer feedback highlights several key improvements: dynamically resolving the target linear head to prevent crashes when a wrapper module is used, avoiding CUDA synchronization bottlenecks in debug logging by checking the log level before calling .item(), supporting Hugging Face Hub IDs when loading pretrained weights, and logging the correct speculative algorithm name instead of the underlying model type when disabling Ulysses sequence parallelism.

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 · 2026-07-15T06:18:29Z

+        predicted_logits = self.target_model(predicted_hidden).float()
+        with torch.no_grad():
+            target_logits = self.target_model(last_hidden_states).float()
+            target_probs = torch.softmax(target_logits, dim=-1)


To ensure robustness across both unit tests (where self.target_model is directly a nn.Linear head) and real training runs (where self.target_model is a wrapper module containing a .fc attribute), we should resolve the actual linear head module dynamically. Calling self.target_model directly on hidden states might fail if the wrapper does not implement a forward method or expects different arguments.

Suggested change

predicted_logits = self.target_model(predicted_hidden).float()

with torch.no_grad():

target_logits = self.target_model(last_hidden_states).float()

target_probs = torch.softmax(target_logits, dim=-1)

target_head = getattr(self.target_model, "fc", self.target_model)

predicted_logits = target_head(predicted_hidden).float()

with torch.no_grad():

target_logits = target_head(last_hidden_states).float()

target_probs = torch.softmax(target_logits, dim=-1)

The eagle3 backend it mirrors calls self.target_model(...) directly (eagle3_trainer_backend.py:964,988), and both possible heads (_SyncedTargetHead and TargetHead) are nn.Modules with a forward. The GPU smoke exercises the real TargetHead path this way and reaches top1=0.95, so a direct call is correct. Keeping it consistent with eagle3.

gemini-code-assist · 2026-07-15T06:18:29Z

+        if num_tokens.detach().item() > 0:
+            with torch.no_grad():
+                draft_top1 = predicted_logits.argmax(dim=-1)
+                target_top1 = target_probs.argmax(dim=-1)
+                acc = ((draft_top1 == target_top1) & valid_mask).float().sum() / num_tokens.clamp_min(1)
+            logger.debug(
+                "[eagle1 loss] tokens=%s hidden_loss=%.6f token_loss=%.6f top1_acc=%.6f",
+                int(num_tokens.detach().cpu().item()),
+                float((total_local_vloss / num_tokens.clamp_min(1)).detach().cpu().item()),
+                float((total_local_ploss / num_tokens.clamp_min(1)).detach().cpu().item()),
+                float(acc.detach().cpu().item()),
+            )


Calling .item() on a GPU tensor triggers a host-device synchronization (CUDA sync), which is a known performance bottleneck in PyTorch. Since this block is only used for debug logging, we should wrap it in an if logger.isEnabledFor(logging.DEBUG): check. This completely avoids the synchronization overhead during normal training runs where the log level is set to INFO or higher.

Suggested change

if num_tokens.detach().item() > 0:

with torch.no_grad():

draft_top1 = predicted_logits.argmax(dim=-1)

target_top1 = target_probs.argmax(dim=-1)

acc = ((draft_top1 == target_top1) & valid_mask).float().sum() / num_tokens.clamp_min(1)

logger.debug(

"[eagle1 loss] tokens=%s hidden_loss=%.6f token_loss=%.6f top1_acc=%.6f",

int(num_tokens.detach().cpu().item()),

float((total_local_vloss / num_tokens.clamp_min(1)).detach().cpu().item()),

float((total_local_ploss / num_tokens.clamp_min(1)).detach().cpu().item()),

float(acc.detach().cpu().item()),

)

if logger.isEnabledFor(logging.DEBUG) and num_tokens.item() > 0:

with torch.no_grad():

draft_top1 = predicted_logits.argmax(dim=-1)

target_top1 = target_probs.argmax(dim=-1)

acc = ((draft_top1 == target_top1) & valid_mask).float().sum() / num_tokens.clamp_min(1)

logger.debug(

"[eagle1 loss] tokens=%s hidden_loss=%.6f token_loss=%.6f top1_acc=%.6f",

int(num_tokens.item()),

float((total_local_vloss / num_tokens.clamp_min(1)).item()),

float((total_local_ploss / num_tokens.clamp_min(1)).item()),

float(acc.item()),

)

Good catch, done: the per-step block is now gated on logger.isEnabledFor(logging.DEBUG) so the .item() syncs are skipped at INFO+ (84f0867).

gemini-code-assist · 2026-07-15T06:18:30Z

+        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 EAGLE-1/2 drafter weights")
+            drafter_module = LlamaForCausalLMEagle1.from_pretrained(spec_model_path, config=draft_config)
+        else:
+            drafter_module = LlamaForCausalLMEagle1(draft_config)


The current check os.path.exists(os.path.join(spec_model_path, "config.json")) only allows loading pretrained weights from a local directory. If a user specifies a Hugging Face Hub model ID (e.g., "some-org/some-model"), this check will evaluate to False, silently falling back to initializing a completely random model. We should enhance this check to support both local paths and Hugging Face Hub IDs while still avoiding the default placeholder path.

Suggested change

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 EAGLE-1/2 drafter weights")

drafter_module = LlamaForCausalLMEagle1.from_pretrained(spec_model_path, config=draft_config)

else:

drafter_module = LlamaForCausalLMEagle1(draft_config)

is_local_pretrained = spec_model_path and os.path.exists(os.path.join(spec_model_path, "config.json"))

is_hf_hub = spec_model_path and not os.path.exists(spec_model_path) and "/" in spec_model_path and "path/to" not in spec_model_path

if is_local_pretrained or is_hf_hub:

log_drafter_checkpoint_step(logger, spec_model_path, action="Loading EAGLE-1/2 drafter weights")

drafter_module = LlamaForCausalLMEagle1.from_pretrained(spec_model_path, config=draft_config)

else:

drafter_module = LlamaForCausalLMEagle1(draft_config)

Same behavior as the eagle3 backend, which also gates weight loading on a local path (eagle3_trainer_backend.py:644). Cold start is the common case here, and hub-id draft init is a separate enhancement that would apply to both backends; I'd rather not add the ad-hoc 'path/to' heuristic. Left as a follow-up.

gemini-code-assist · 2026-07-15T06:18:30Z

+                pred_logits = backend.target_model(predicted_hidden).float()
+                tgt_logits = backend.target_model(batch["last_hidden_states"]).float()


To align with the robust target head resolution in the backend and prevent potential crashes in the smoke test when backend.target_model is a wrapper module, we should dynamically resolve the linear head using getattr.

Suggested change

pred_logits = backend.target_model(predicted_hidden).float()

tgt_logits = backend.target_model(batch["last_hidden_states"]).float()

target_head = getattr(backend.target_model, "fc", backend.target_model)

pred_logits = target_head(predicted_hidden).float()

tgt_logits = target_head(batch["last_hidden_states"]).float()

backend.target_model is the real head built by build_model (TargetHead / _SyncedTargetHead), which is callable; the smoke already runs this path successfully. Same rationale as the backend comment.

gemini-code-assist · 2026-07-15T06:18:30Z

+                    "[Rank %s] Disable Ulysses SP for %s drafter training: "
+                    "rollout_tp=%s training_group_world_size=%s",
+                    self.rank,
+                    self.backend.model_type,


Since Eagle1TrainerBackend reports model_type == "eagle3" to reuse the online data-collection plumbing, logging self.backend.model_type here will incorrectly output "Disable Ulysses SP for eagle3 drafter training". This is highly confusing because EAGLE-3 actually supports Ulysses SP, whereas only EAGLE-1/2 does not. We should log the actual speculative algorithm name instead.

Suggested change

self.backend.model_type,

getattr(self.backend.config.rollout.drafter, "speculative_algorithm", self.backend.model_type),

Agreed, fixed: it now logs the actual speculative_algorithm (EAGLE1/EAGLE2) instead of the reused eagle3 model_type (84f0867).

- gate the per-step debug accuracy block on logger.isEnabledFor(DEBUG) so the .item() host-device syncs are skipped when training at INFO or higher. - log the actual speculative_algorithm (EAGLE1/EAGLE2) instead of the reused eagle3 model_type in the Ulysses-SP disable message. AI assistance was used for this change. Signed-off-by: khazic <khazzz1c@gmail.com>

…eagle2-drafter # Conflicts: # verl_speco/config/speco_trainer.yaml

…drafter Resolves conflicts with the Domino backend (verl-project#17), EAGLE-1/2 backend (verl-project#15), and the DSpark L1 target-hidden plumbing: the peagle and domino worker dispatch branches, vllm guardrails, and speco_base.yaml config blocks are all additive and kept side by side. The merge also brings in main's supports_ulysses_sp gate, which force-disables Ulysses SP for the P-EAGLE backend (supports_ulysses_sp=False) and enforces the unsliced-batch invariant the peagle assembly relies on. Signed-off-by: khazic <khazzz1c@gmail.com>

khazic added 5 commits July 15, 2026 13:48

test(eagle1): detach loss tensors in assertions to silence grad warning

7db4a2e

AI assistance was used for this change. Signed-off-by: khazic <khazzz1c@gmail.com>

test(eagle1): genericize smoke script paths for portability

284b43c

Make --target required and use a temp dir for the cold-start drafter path instead of hardcoded absolute paths. AI assistance was used for this change. Signed-off-by: khazic <khazzz1c@gmail.com>

khazic mentioned this pull request Jul 15, 2026

Offering to contribute EAGLE-1 and EAGLE-2 support #14

Closed

gemini-code-assist Bot reviewed Jul 15, 2026

View reviewed changes

khazic mentioned this pull request Jul 15, 2026

Add P-EAGLE and Domino drafter support #16

Closed

Merge remote-tracking branch 'upstream/main' into khazic/feat/eagle1-…

142fde5

…eagle2-drafter # Conflicts: # verl_speco/config/speco_trainer.yaml

tpx818 merged commit 3bbf1cd into verl-project:main Jul 16, 2026
1 check passed

Navigation Menu

Uh oh!

feat: add EAGLE-1 / EAGLE-2 drafter training backend - #15

feat: add EAGLE-1 / EAGLE-2 drafter training backend#15
tpx818 merged 7 commits into
verl-project:mainfrom
khazic:khazic/feat/eagle1-eagle2-drafter

khazic commented Jul 15, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Jul 15, 2026

Uh oh!

khazic Jul 15, 2026

Uh oh!

gemini-code-assist Bot Jul 15, 2026

Uh oh!

khazic Jul 15, 2026

Uh oh!

gemini-code-assist Bot Jul 15, 2026

Uh oh!

khazic Jul 15, 2026

Uh oh!

gemini-code-assist Bot Jul 15, 2026

Uh oh!

khazic Jul 15, 2026

Uh oh!

gemini-code-assist Bot Jul 15, 2026

Uh oh!

khazic Jul 15, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		pred_logits = backend.target_model(predicted_hidden).float()
		tgt_logits = backend.target_model(batch["last_hidden_states"]).float()

-                pred_logits = backend.target_model(predicted_hidden).float()
-                tgt_logits = backend.target_model(batch["last_hidden_states"]).float()
+                target_head = getattr(backend.target_model, "fc", backend.target_model)
+                pred_logits = target_head(predicted_hidden).float()
+                tgt_logits = target_head(batch["last_hidden_states"]).float()

	self.backend.model_type,
	getattr(self.backend.config.rollout.drafter, "speculative_algorithm", self.backend.model_type),

Uh oh!

Conversation

khazic commented Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What

Design

Not a duplicate

Testing

Known follow-ups

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

khazic Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

khazic Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

khazic Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

khazic Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

khazic Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

khazic commented Jul 15, 2026 •

edited

Loading