eagle.mdx

A step-by-step guide for training an EAGLE speculative decoding drafter to accelerate LLM inference using NeMo AutoModel.

---

What is EAGLE Speculative Decoding?

Large language models generate text one token at a time — each token requires a full forward pass through the entire model. Speculative decoding speeds this up by pairing the large target model with a small, fast drafter model. The drafter guesses multiple tokens ahead; the target model then verifies them all in a single forward pass, accepting correct guesses and rejecting wrong ones. The output is mathematically identical to running the target model alone, but 2-3x faster.

EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency) is a family of speculative decoding methods. NeMo AutoModel supports three variants:

Variant	Recipe	Description
EAGLE-1	train_eagle1	Lightweight 1-layer draft transformer; learns to predict target hidden states + next token
EAGLE-2	train_eagle2	Same architecture as EAGLE-1 (alias recipe)
EAGLE-3	train_eagle3	Advanced drafter with test-time training (TTT) unroll and vocabulary mapping; best speed

The Task

We train an EAGLE-3 drafter for Llama 3.1 8B Instruct on the PerfectBlend dataset — a chat corpus whose assistant turns were generated by the same Llama 3.1 8B model, ensuring distribution alignment between training data and target.

After training, we serve the target + drafter together via SGLang for accelerated inference.

Guide Overview

Step	Description
Step 0	Environment setup
Step 1	Understand EAGLE architecture
Step 2	Prepare the training dataset
Step 3	Configure and launch EAGLE-3 training
Step 4	Monitor training and inspect checkpoints
Step 5	Serve with SGLang
Step 6	(Bonus) Train an EAGLE-1 drafter

Hardware Requirements

Setup	Target Model	GPUs	Training Time
MVP (quick test)	Llama 3.2 1B	1x A100 80 GB	~10 min (1 epoch, 1k samples)
Production	Llama 3.1 8B Instruct	8x A100 80 GB	~2 h (1 epoch, 200k samples)

On CUDA, the target model is loaded in BF16 and frozen during training (the CPU fallback uses FP32). Only the small drafter is trained; standard EAGLE-3 uses one fused decoder layer plus an auxiliary projection. GPU memory is therefore dominated by the target model size.

---

Step 0 — Environment Setup

This guide runs inside the NeMo AutoModel Docker container:

docker run -it --rm --gpus all --ipc=host --network host \
    -v $(pwd):/workspace \
    nvcr.io/nvidia/nemo-automodel:26.06.00

huggingface-cli login

cd /opt/Automodel

For SGLang serving (Step 5), install it in the same environment:

uv pip install "sglang>=0.5.9"

**SGLang / transformers version compatibility.** SGLang `0.5.9` pins `transformers==4.57.1`. The SGLang target backend (`target_model_backend: sglang`) and `serve_sglang` run in that environment without issue. The `--compare-hf` check in `smoke_sglang_target.py` additionally builds the HuggingFace target through `NeMoAutoModelForCausalLM`, which imports `AutoModelForMultimodalLM` (available only in `transformers` 5.x), so that single-process comparison cannot run as written under `transformers==4.57.1`. To compare the two backends, use a `transformers` build that satisfies both, or load the HuggingFace side with plain `transformers.AutoModelForCausalLM` wrapped in `HFEagle3TargetModel`.

---

Step 1 — Understand EAGLE Architecture

How EAGLE-3 Works

EAGLE-3 pairs a frozen target LLM with a small trainable drafter. During training, the drafter learns to predict what the target model would produce next, using a technique called test-time training (TTT) unroll:

Target (frozen)          Drafter (trainable)
┌──────────────┐         ┌──────────────┐
│ Llama 3.1 8B │ ──────> │  1-layer     │
│              │ hidden  │  transformer │
│ Full model   │ states  │  + fc fusion │
│ (frozen)     │         │  + lm_head   │
└──────────────┘         └──────────────┘
                               │
                         predict next token
                         + hidden states

Key components:

• Target model: The full LLM (e.g., Llama 3.1 8B), completely frozen. Provides hidden states from selected intermediate layers as auxiliary inputs to the drafter.
• Draft model: A shallow transformer (a single fused decoder layer) with:
  • A fusion layer (fc) that combines auxiliary hidden states from 3 target layers
  • Its own attention layers, MLP, and layer norm
  • A smaller output vocabulary (e.g., 8192 or 32000 tokens instead of 128k) to reduce compute
• TTT unroll: The drafter runs the configured number of autoregressive steps during training; the maintained Llama examples use 4. Loss weights decay as 0.8^i, teaching the drafter to make multi-step predictions — exactly what speculative decoding needs.

Supported Target Architectures

The target is resolved from its HF config.architectures string through the registry in components/speculative/eagle/registry.py. Most dense and MoE decoders reuse the config-driven Llama-style draft (a one-line registry entry); architectures whose RoPE or attention layout the shared draft cannot represent get a thin dedicated draft class.

Target	Draft
Llama, Phi-3, Qwen3 dense, Qwen3-MoE	shared Llama-style dense draft
gpt-oss (GptOssForCausalLM)	dedicated draft reproducing gpt-oss YaRN RoPE
DeepSeek-V3 (DeepseekV3ForCausalLM)	dedicated MLA draft
Gemma4 (Gemma4ForConditionalGeneration)	thin dedicated draft (see below)

Gemma4. Gemma4 ships as a multimodal Gemma4ForConditionalGeneration whose text decoder config is nested under config.text_config; the draft is built from that inner config automatically. Set target_force_hf: true so the stock HuggingFace text-capable model is loaded (the draft hooks the decoder under model.language_model.layers). The draft consumes only post-block hidden states, so it is the same Llama-style dense draft used elsewhere, with two Gemma quirks reconciled: the hidden_activation (GeGLU) key and Gemma4's nested per-attention-type rope_parameters, which is flattened to a standard full-rotary Llama RoPE on the global (full-attention) base. The saved checkpoint keeps the canonical architectures: ["LlamaEagle3DraftModel"] layout, so SGLang and vLLM serve it with their existing EAGLE-3 Llama head. See examples/speculative/eagle3/gemma4_e2b_eagle3.yaml.

EAGLE-3.1 Drafter Toggles

The same train_eagle3 recipe supports the EAGLE-3.1 drafter variant via two optional flags in recipe_args. Both default to false, so existing EAGLE-3 configs and checkpoints behave identically. Setting them applies the EAGLE-3.1 architectural changes from vllm-project/vllm#42764 to the Llama-style draft. The MLA-backbone community release lightseekorg/kimi-k2.6-eagle3.1-mla is a separate architecture (Eagle3DeepseekV2ForCausalLM) and is not produced by this recipe.

Flag	Effect
fc_norm	Apply an independent RMSNorm to each of the num_aux_hidden_states (3 by default) auxiliary hidden-state chunks before they are concatenated and projected by model.fc. Stored as an nn.ModuleList with on-disk keys model.fc_norm.0.weight, model.fc_norm.1.weight, ... matching vLLM's layout so checkpoints load directly.
norm_output	Route the existing final RMSNorm (model.norm) over the per-step hidden state returned by the drafter so the next TTT step (and lm_head) consume the post-norm state instead of the raw decoder output. Adds no parameters.

Together they remove the "attention drift" pattern (loss of focus on sink tokens at deeper speculation depths) reported by the EAGLE-3.1 paper and let the drafter behave more like a recurrently applied module than a stack of extra layers bolted onto the target.

recipe_args:
  # ... standard EAGLE-3 fields ...
  fc_norm: true
  norm_output: true

How EAGLE-1 Differs

EAGLE-1 is simpler: it uses a single transformer layer, predicts the full vocabulary, and trains with a combined loss of SmoothL1 (Huber) on hidden states (hidden_loss_weight) and cross-entropy on tokens (token_loss_weight). No TTT unroll, no vocabulary mapping.

---

Step 2 — Prepare the Training Dataset

Data format

EAGLE training accepts chat data in the OpenAI messages format—either JSONL files or HuggingFace datasets with a messages column—and automatically converts supported ShareGPT-style conversations rows to that format:

{"messages": [
  {"role": "system", "content": "You are a helpful assistant."},
  {"role": "user", "content": "What is the capital of France?"},
  {"role": "assistant", "content": "The capital of France is Paris."}
]}

Option A: Use a pre-regenerated dataset (recommended)

For best results, the assistant turns in your training data should come from the same model you'll use as the target at inference time. The PerfectBlend dataset already has answers regenerated by Llama 3.1 8B Instruct:

python -c "
from datasets import load_dataset
ds = load_dataset(
    'frankleeeee/PerfectBlend-Regenerated-Llama-3.1-8B-Instruct',
    split='train[:5]'
)
print(f'Columns: {ds.column_names}')
print(f'Sample conversation:')
for msg in ds[0]['conversations'][:3]:
    role = msg['role']
    text = msg['content'][:80]
    print(f'  [{role}] {text}...')
"

Expected output:

Columns: ['conversations']
Sample conversation:
  [system] You are a helpful assistant....
  [user] What are the main differences between Python 2 and Python 3?...
  [assistant] Here are the key differences between Python 2 and Python 3:

1. **P...

`ChatDataset` automatically converts supported ShareGPT-style rows from the `conversations` column to OpenAI-style `messages`, including PerfectBlend's `role`/`content` records. A manual conversion is only needed for a dataset with unsupported role names or a different conversation schema.

Option B: Regenerate answers from your target model

If you have a chat dataset whose answers were generated by a different model, you can regenerate them using your target. This aligns the training data distribution with the model the drafter will actually assist at inference time.

Step B.1 — Start the target server (in one shell):

python -m sglang.launch_server \
    --model-path meta-llama/Llama-3.1-8B-Instruct \
    --port 30000

Wait for Uvicorn running on http://0.0.0.0:30000 before proceeding.

Step B.2 — Regenerate (in another shell):

python -m nemo_automodel.components.speculative.regenerate \
    --input-data Aeala/ShareGPT_Vicuna_unfiltered \
    --output-dir ./regenerated/sharegpt_llama31_8b \
    --target-server http://localhost:30000/v1 \
    --model meta-llama/Llama-3.1-8B-Instruct \
    --concurrency 64 \
    --shard-size 1000

For each sample, the script:

1. Loads the conversation from the input dataset
2. Drops the trailing assistant turn, keeping the user prompt context
3. Calls the target server to generate a new assistant response
4. Saves the rebuilt conversation to parquet shards

The output directory contains parquet files with a messages column — ready for EAGLE training. The script is resumable: rerun with --resume to skip completed shards.

Flag	Default	Notes
--concurrency	32	In-flight requests; raise to saturate the server
--shard-size	1000	Rows per parquet file
--temperature	0.0	Greedy by default (recommended for EAGLE)
--max-new-tokens	1024	Cap per-answer length
--split	train	Supports HF slice syntax, e.g., train[:10000]

---

Step 3 — Configure and Launch EAGLE-3 Training

YAML config

Save the following as eagle3_llama8b.yaml:

recipe: TrainEagle3Recipe

dist_env:
  backend: nccl
  timeout_minutes: 60

recipe_args:
  target_model_name_or_path: meta-llama/Llama-3.1-8B-Instruct

  # Point to your training data (HF dataset id, local parquet dir, or JSONL)
  train_data_path: frankleeeee/PerfectBlend-Regenerated-Llama-3.1-8B-Instruct
  val_data_path: null

  # Slice to 200k samples for a ~2h training run
  train_split: "train[:200000]"
  val_split: null

  output_dir: ./outputs/eagle3_llama8b
  seq_length: 2048
  micro_batch_size: 1
  grad_accumulation_steps: 4   # effective batch = 8 GPUs * 1 * 4 = 32
  num_workers: 4
  num_epochs: 1

  # EAGLE-3 specific
  ttt_steps: 4                 # TTT unroll depth (higher = better but slower)
  draft_vocab_size: 32000      # smaller vocab = faster drafter

  # The drafter copies the target's embedding table at init; this flag
  # freezes those copied weights so only the draft transformer layers
  # and lm_head are trained.
  freeze_embeddings: true
  shuffle_seed: 42
  log_every_steps: 20
  max_grad_norm: 1.0

optimizer:
  lr: 2.0e-4
  betas: [0.9, 0.95]
  weight_decay: 0.0
  warmup_ratio: 0.05           # 5% warmup
  min_lr_ratio: 0.1

checkpoint:
  enabled: true
  checkpoint_dir: ./outputs/eagle3_llama8b/checkpoints
  # The recipe defaults to safetensors + consolidated; these lines are
  # shown explicitly for clarity but can be omitted.
  model_save_format: safetensors
  save_consolidated: true

Config field reference

Field	What It Does
target_model_name_or_path	HuggingFace model ID for the frozen target LLM
train_data_path	Path to chat data (HF dataset id, parquet dir, or JSONL)
train_split	Optional HF slice syntax to limit data size
seq_length	Context window length (1024 for quick tests, 2048 for production)
micro_batch_size	Per-GPU batch size
grad_accumulation_steps	Gradient accumulation for larger effective batches
ttt_steps	Required TTT unroll depth; maintained Llama examples use 4, and cost is linear per step
draft_vocab_size	Draft output vocabulary size; smaller = faster inference
freeze_embeddings	Freeze the embedding table copied from the target so only draft layers train (recommended true)
target_attn_implementation	Optional attention backend for the frozen target (e.g. sdpa); defaults to HF auto-select. Set sdpa if the target's FlashAttention path is broken on your build (e.g. the Qwen3 FA2 s_aux crash)
fc_norm	EAGLE-3.1: per-chunk independent RMSNorm (ModuleList) on auxiliary hidden states before the fc projection (default false)
norm_output	EAGLE-3.1: feed the post-model.norm hidden state into the next TTT step and lm_head (default false)
warmup_ratio	Fraction of total steps for LR warmup

Launch training

Multi-GPU (8x A100, production):

torchrun --nproc-per-node=8 \
    -m nemo_automodel.recipes.llm.train_eagle3 \
    -c eagle3_llama8b.yaml

Single-GPU (quick test with Llama 3.2 1B):

For a quick test, use the MVP config with Llama 3.2 1B and a small dataset:

python -m nemo_automodel.recipes.llm.train_eagle3 \
    -c examples/speculative/eagle3/llama_eagle3_mvp.yaml

For GPUs with FlashAttention support, add `draft_attn_implementation: flash_attention_2` to `recipe_args` for faster training. See [llama_eagle3_mvp_flash_attn.yaml](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/speculative/eagle3/llama_eagle3_mvp_flash_attn.yaml) for a complete example.

---

Step 4 — Monitor Training and Inspect Checkpoints

What to watch

At each logging interval, the recipe emits epoch, step, train_loss, train_acc, an optional train_tau_sim, and lr. The values below illustrate the log format:

epoch=0 step=20 train_loss=3.452100 train_acc=0.214500 train_tau_sim=0.3812 lr=8.000e-05
epoch=0 step=40 train_loss=2.897300 train_acc=0.267100 train_tau_sim=0.4198 lr=1.600e-04

Checkpoint layout

The shown 200k-sample, 8-GPU config does not enable step or epoch checkpoint cadence, so it saves only the final checkpoint after 6250 optimizer steps:

outputs/eagle3_llama8b/checkpoints/
  epoch_1_step_6250/
    model/                   # HF-style draft artifacts
    eagle_meta.pt            # EAGLE metadata and token mapping
    optim/                   # optimizer and scheduler state
    rng/                     # distributed RNG state
    config.yaml              # training config snapshot
  LATEST -> epoch_1_step_6250

Resume from checkpoint

To create intermediate checkpoints that can recover an interrupted run, first set a cadence in recipe_args, for example ckpt_every_steps: 1000. Then resume from the latest checkpoint:

checkpoint:
  restore_from: LATEST

Or point to a specific checkpoint:

checkpoint:
  restore_from: epoch_0_step_3000

---

Step 5 — Serve with SGLang

The serve_sglang helper resolves the current HF-style training checkpoint and launches the server in one command. It also converts older bare-weight checkpoints as a compatibility fallback.

Launch the server

python -m nemo_automodel.components.speculative.serve_sglang \
    --target meta-llama/Llama-3.1-8B-Instruct \
    --draft ./outputs/eagle3_llama8b/checkpoints/LATEST \
    --algorithm EAGLE3 \
    --num-steps 3 \
    --topk 1 \
    --num-draft-tokens 4 \
    --port 30000

For a current checkpoint, the helper:

1. Resolves the checkpoint's HF-style model/ directory directly
2. Rewrites the architecture name for SGLang compatibility when needed (LlamaEagle3DraftModel -> LlamaForCausalLMEagle3)
3. Uses an existing model/speculative_token_map.pt, or generates it from the sibling eagle_meta.pt for EAGLE-3
4. Launches SGLang with the correct speculative decoding flags

For a legacy checkpoint containing draft_model.pt and config.json, the helper exports the weights and rewritten config into model/. For EAGLE-3, it uses the legacy eagle3_meta.pt when present to generate model/speculative_token_map.pt.

Serving parameters

Flag	Default	Notes
--algorithm	EAGLE3	EAGLE3 for EAGLE-3 drafters, EAGLE for EAGLE-1/2
--num-steps	3	Speculative steps per draft iteration
--topk	1	Branching factor for tree search
--num-draft-tokens	4	Budget of draft tokens per branch
--dtype	bfloat16	Must match training dtype
--tp-size	1	Tensor parallelism (shards the target model only)
--print-only	—	Inspect the resolved command without launching

Pass extra SGLang flags after --:

python -m nemo_automodel.components.speculative.serve_sglang \
    --target meta-llama/Llama-3.1-8B-Instruct \
    --draft ./outputs/eagle3_llama8b/checkpoints/LATEST \
    --algorithm EAGLE3 \
    -- --enable-torch-compile --schedule-conservativeness 1.2

Smoke-test the server

Once you see Uvicorn running on http://0.0.0.0:30000, test it:

curl http://localhost:30000/generate \
    -H "Content-Type: application/json" \
    -d '{
      "text": "Hello, my name is",
      "sampling_params": {"max_new_tokens": 64}
    }'

Expected output:

{
  "text": "Hello, my name is Sarah and I am a 25-year-old software engineer...",
  "meta_info": {
    "prompt_tokens": 6,
    "completion_tokens": 64,
    "accept_length_per_step": 3.2
  }
}

The accept_length_per_step metric shows how many tokens the target model accepts per speculative step on average. Higher is better — a value of 3.0+ indicates the drafter is accurately predicting the target's behavior.

OpenAI-compatible endpoint

SGLang also exposes an OpenAI-compatible API:

curl http://localhost:30000/v1/chat/completions \
    -H "Content-Type: application/json" \
    -d '{
      "model": "meta-llama/Llama-3.1-8B-Instruct",
      "messages": [
        {"role": "user", "content": "Explain speculative decoding in one paragraph."}
      ],
      "max_tokens": 256
    }'

Expected output:

{
  "choices": [{
    "message": {
      "role": "assistant",
      "content": "Speculative decoding is a technique for accelerating autoregressive language model inference. It works by using a small, fast \"draft\" model to predict multiple future tokens, which are then verified in parallel by the larger \"target\" model in a single forward pass. Tokens that match the target model's predictions are accepted, while incorrect tokens are rejected and regenerated. Because verification is cheaper than sequential generation (it processes all candidate tokens simultaneously), the overall throughput increases significantly — typically 2-3x — while producing output that is mathematically identical to running the target model alone."
    }
  }],
  "usage": {
    "prompt_tokens": 14,
    "completion_tokens": 112
  }
}

---

Step 6 — (Bonus) Train an EAGLE-1 Drafter

EAGLE-1 is simpler and faster to train, making it a good starting point for experimentation. It uses a single transformer layer and trains with a combined hidden-state SmoothL1 (Huber) + token cross-entropy loss.

YAML config

Save as eagle1_llama8b.yaml:

recipe: TrainEagle1Recipe

dist_env:
  backend: nccl
  timeout_minutes: 30

recipe_args:
  target_model_name_or_path: meta-llama/Llama-3.1-8B-Instruct
  train_data_path: frankleeeee/PerfectBlend-Regenerated-Llama-3.1-8B-Instruct
  val_data_path: null
  train_split: "train[:200000]"
  val_split: null
  output_dir: ./outputs/eagle1_llama8b
  seq_length: 2048
  micro_batch_size: 1
  grad_accumulation_steps: 4
  num_workers: 4
  num_epochs: 1

  # EAGLE-1 specific
  draft_num_hidden_layers: 1    # number of draft transformer layers
  hidden_loss_weight: 1.0       # SmoothL1 (Huber) loss on hidden states
  token_loss_weight: 0.1        # cross-entropy loss on tokens

  freeze_embeddings: true
  shuffle_seed: 42
  log_every_steps: 10
  max_grad_norm: 1.0

optimizer:
  lr: 1.0e-4
  betas: [0.9, 0.95]
  weight_decay: 0.0

checkpoint:
  enabled: true
  checkpoint_dir: ./outputs/eagle1_llama8b/checkpoints
  # Defaults to safetensors + consolidated; can be omitted.
  model_save_format: safetensors
  save_consolidated: true

Launch

torchrun --nproc-per-node=8 \
    -m nemo_automodel.recipes.llm.train_eagle1 \
    -c eagle1_llama8b.yaml

Serve

Use --algorithm EAGLE (not EAGLE3) for EAGLE-1/2 drafters:

python -m nemo_automodel.components.speculative.serve_sglang \
    --target meta-llama/Llama-3.1-8B-Instruct \
    --draft ./outputs/eagle1_llama8b/checkpoints/LATEST \
    --algorithm EAGLE \
    --num-steps 3 --topk 1 --num-draft-tokens 4 \
    --port 30000

EAGLE-1 vs EAGLE-3

	EAGLE-1	EAGLE-3
Draft layers	1 (configurable)	1 fused decoder layer + auxiliary fc projection
Training objective	Hidden SmoothL1 (Huber) + token CE	TTT unroll with decay
Vocabulary	Full target vocab	Reduced (e.g., 8k-32k)
Training speed	Faster	Slower (due to TTT unroll)
Inference speedup	Good (2-2.5x)	Better (2.5-3x)
Best for	Quick experiments	Production deployment

---

Example Configs Reference

Config	Target	Variant	Notes
llama_eagle3_mvp.yaml	Llama 3.2 1B	EAGLE-3	Quick test, single GPU
llama_eagle3_mvp_flash_attn.yaml	Llama 3.2 1B	EAGLE-3	With FlashAttention-2
llama_eagle3_perfectblend.yaml	Llama 3.1 8B	EAGLE-3	Production config, 200k samples
llama_eagle3_1_perfectblend.yaml	Llama 3.1 8B	EAGLE-3.1	Production config with fc_norm + norm_output enabled
llama_eagle1_perfectblend.yaml	Llama 3.2 1B	EAGLE-1	Quick test, single GPU
llama_eagle2_perfectblend.yaml	Llama 3.2 1B	EAGLE-2	Same as EAGLE-1 (alias)

Troubleshooting

Symptom	Fix
OutOfMemoryError during training	Reduce seq_length (1024 instead of 2048) or micro_batch_size
Loss stays flat or NaN	Check max_grad_norm (default 1.0), reduce lr
SGLang model not found error	Ensure --algorithm matches the recipe (EAGLE3 for train_eagle3, EAGLE for train_eagle1/2)
dtype mismatch at serving	Pass --dtype bfloat16 to match training precision
conversations vs messages conversion error	Check that each turn uses a supported ShareGPT role and either role/content or from/value; ChatDataset converts supported rows automatically
Checkpoint resume fails	Use restore_from: LATEST or an existing exact subdirectory name such as epoch_1_step_6250
Low accept_length_per_step at serving	Train longer, use more data, or try regenerating answers with the target model (Option B in Step 2)