Train an EAGLE-3 draft head

A draft head learns to predict what the target model will say next. Its training data must therefore be the target’s own output, not human-written responses. This page builds that data, chat and agent traces, with Qwen3-8B itself, then fine-tunes the public EAGLE-3 head on it.

Configure models and hardware

Everything that changes between clusters lives in config.py.

config.py
                
# The model you serve (the verifier) and the public EAGLE-3 draft head you start from.
TARGET_MODEL = "Qwen/Qwen3-8B"
PUBLIC_DRAFTER = "RedHatAI/Qwen3-8B-speculator.eagle3"

# Number of tokens the draft head proposes per verification step.
NUM_SPECULATIVE_TOKENS = 3

# One GPU type for every GPU task. Qwen3-8B in bf16 needs about 17 GB, so any 40 GB+ card
# works. On H100 or H200, change this one line.
GPU = "L40s"

The pipeline starts from RedHatAI/Qwen3-8B-speculator.eagle3, a general-purpose EAGLE-3 head trained by Red Hat. Fine-tuning it on your traffic is cheaper than training a head from scratch and keeps its general drafting ability. Set init_from="" on train_drafter to train from scratch instead. The same task trains the DFlash and DSpark drafters with a different method.

One GPU image runs vLLM for data generation, serving, and benchmarks, and speculators for training:

config.py
                
                    
                
            
                
            

                
                    
                
            
# One GPU image runs vLLM (serving, data generation, benchmarks) and speculators (training).
# vLLM 0.33 and speculators 0.9 agree on torch 2.14 and transformers 5.16-5.19.
gpu_image = (
    flyte.Image.from_debian_base(name="specdec-gpu", python_version=(3, 12))
    .with_apt_packages("curl", "git")
    .with_pip_packages("flyteplugins-vllm==2.13.0", "unionai-reuse>=0.1.3")
    .with_pip_packages(
        f"vllm=={VLLM_VERSION}",
        f"speculators=={SPECULATORS_VERSION}",
        "transformers==5.19.0",
        "arctic-inference==0.3.0",
        "datasets>=4.0.0",
        "httpx>=0.27",
    )
    # Prebuilt FlashInfer kernels, so vLLM's sampler does not JIT-compile at startup.
    .with_pip_packages(
        f"flashinfer-jit-cache=={FLASHINFER_VERSION}",
        index_url="https://flashinfer.ai/whl/cu130",
    )
)

# Orchestration, data wrangling, and reports run on a small CPU image.
# The vLLM plugin is small and pulls in no vLLM; drivers import it to define serving apps.
cpu_image = flyte.Image.from_debian_base(name="specdec-cpu", python_version=(3, 12)).with_pip_packages(
    "flyteplugins-vllm==2.13.0", "datasets>=4.0.0", "httpx>=0.27"
)

vLLM 0.33 and speculators 0.9 can share one environment because they agree on torch 2.14 and transformers 5.16 to 5.19. Pin all three. A newer transformers release can break that overlap.

Resource requests are sized to the nodes they land on. A request above a node’s allocatable CPU or memory never schedules, even when the GPU is free:

config.py
                
# Requests sized to the node shapes they land on. The defaults fit AWS L40S nodes:
# g6e.2xlarge (1 GPU, ~7.9 allocatable vCPU, ~61 GiB) and g6e.12xlarge (4 GPUs, 48 vCPU,
# 384 GiB). CPU-only tasks fit a t3a.xlarge (~3.9 vCPU, ~14.5 GiB). Leave headroom below a
# node's allocatable capacity, or the pod never schedules.
ONE_GPU = flyte.Resources(cpu=6, memory="48Gi", gpu=f"{GPU}:1", disk="100Gi", shm="auto")
FOUR_GPUS = flyte.Resources(cpu=40, memory="320Gi", gpu=f"{GPU}:4", disk="200Gi", shm="auto")

# Training needs a vLLM server for hidden states and the trainer. With 4 GPUs, two serve and
# two train. With 1, both share one GPU: slower, but it fits a single 48 GB card.
TRAIN_NUM_GPUS = 1

TRAIN_NUM_GPUS picks the training layout. With 1, the hidden-state server and the trainer share one 48 GB GPU, which fits any cluster that can serve the model. With 4, the task takes a whole 4-GPU node: two GPUs serve hidden states and two train, which is several times faster.

Each stage gets an environment sized for its work:

config.py
                
                    
                
            
                
            

                
                    
                
            
data_env = flyte.TaskEnvironment(
    name="specdec-data",
    image=cpu_image,
    resources=flyte.Resources(cpu=2, memory="12Gi", disk="20Gi"),
)

# Data generation keeps one vLLM server warm per replica, so each shard skips the model load.
generate_env = flyte.TaskEnvironment(
    name="specdec-generate",
    image=gpu_image,
    resources=ONE_GPU,
    reusable=flyte.ReusePolicy(replicas=(1, 4), concurrency=1, idle_ttl=600, scaledown_ttl=300),
)

# Online training: vLLM serves hidden states while the trainer consumes them.
train_env = flyte.TaskEnvironment(
    name="specdec-train",
    image=gpu_image,
    resources=FOUR_GPUS if TRAIN_NUM_GPUS == 4 else ONE_GPU,
)

# Each benchmark task owns one GPU and one vLLM server configuration.
bench_env = flyte.TaskEnvironment(
    name="specdec-bench",
    image=gpu_image,
    resources=ONE_GPU,
)

# The driver only orchestrates. depends_on registers every environment it calls.
driver_env = flyte.TaskEnvironment(
    name="specdec-driver",
    image=cpu_image,
    resources=flyte.Resources(cpu=2, memory="4Gi"),
    depends_on=[data_env, generate_env, train_env, bench_env],
)

generate_env is a reusable environment. Up to four replicas stay alive between tasks, so the vLLM server a replica starts for its first shard serves every later shard too. Loading Qwen3-8B and capturing CUDA graphs takes several minutes, which would otherwise repeat for every shard.

Collect target-generated data

Sample prompts

prepare_prompts samples UltraChat conversations for general chat and ToolACE tasks for agent traffic, and splits each into shards:

data.py
                
                    
                
            
                
            

                
                    
                
            
@data_env.task(cache="auto", report=True)
async def prepare_prompts(num_chat: int, num_agent: int, num_shards: int) -> PromptShards:
    """Sample prompts and split them into shards for parallel generation."""

    async def write_shards(rows: list[dict], prefix: str) -> list[File]:
        files = []
        for i in range(num_shards):
            path = f"/tmp/{prefix}-{i}.jsonl"
            with open(path, "w") as f:
                for row in rows[i::num_shards]:
                    f.write(json.dumps(row) + "\n")
            files.append(await File.from_local(path))
        return files

    chat = [{"messages": m} for m in load_chat_prompts(num_chat)]
    agent = load_agent_tasks("train", num_agent)
    await flyte.report.replace.aio(data_reports.prompts_report(chat, agent, num_shards))
    await flyte.report.flush.aio()
    return PromptShards(
        chat=await write_shards(chat, "chat"),
        agent=await write_shards(agent, "agent"),
    )

ToolACE is split three ways in workloads.py. The first 60% is initial training data, the next 20% stands in for production traffic in the retraining loop, and the last 20% is held out for benchmarks. No benchmark prompt is ever trained on.

Regenerate chat responses

speculators regenerate-responses sends each conversation to Qwen3-8B and writes the response in tokenized form, with a loss mask over the assistant tokens:

data.py
                
                    
                
            
                
            

                
                    
                
            
@generate_env.task(cache="auto", report=True)
async def regenerate_chat(shard: File) -> Shard:
    """Replace the dataset's responses with Qwen3-8B's own, tokenized for training."""
    server = _get_server()
    local = await shard.download()
    out = "/tmp/regenerated.jsonl"
    subprocess.run(
        [
            "speculators", "regenerate-responses",
            "--endpoint", f"{server.base_url}/v1/chat/completions",
            "--dataset", local,
            "--outfile", out,
            "--max-tokens", "2048",
            "--concurrency", "64",
            "--sampling-params", json.dumps({"temperature": 0.7, **NO_THINKING}),
        ],
        check=True,
    )
    with open(out) as f:
        rows = [json.loads(line) for line in f]
    with open(local) as f:
        requested = sum(1 for _ in f)

    lengths = [sum(r["loss_mask"]) for r in rows]
    result = Shard(
        data=await File.from_local(out),
        kind="chat",
        rows=len(rows),
        tokens=sum(len(r["input_ids"]) for r in rows),
        target_tokens=sum(lengths),
        failed=max(requested - len({r.get("primary_id") for r in rows}), 0),
        lengths=lengths,
    )
    await flyte.report.replace.aio(data_reports.chat_shard_report(result, rows))
    await flyte.report.flush.aio()
    return result

The sampling temperature is 0.7, not 0. A head trained only on greedy outputs drafts poorly for requests served with sampling. Thinking is disabled because this tutorial’s traffic runs without it. If your application uses Qwen3’s thinking mode, regenerate with it on so the head learns to draft reasoning too.

Record agent traces

ToolACE rows carry tool schemas, a user request, and recorded tool results. The agent sends the request and the tools to Qwen3-8B, and when the model calls a tool, replays the recorded result:

workloads.py
                
                    
                
            
                
            

                
                    
                
            
Complete = Callable[[list[dict], list[dict]], Awaitable[dict]]

def http_complete(
    client: httpx.AsyncClient,
    base_url: str,
    *,
    api_key: str | None = None,
    model: str = TARGET_MODEL,
    temperature: float = 0.7,
    max_tokens: int = 1024,
) -> Complete:
    """An OpenAI-compatible chat completion call that returns the assistant message."""
    headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}

    async def complete(messages: list[dict], tools: list[dict]) -> dict:
        response = await client.post(
            f"{base_url}/v1/chat/completions",
            headers=headers,
            json={
                "model": model,
                "messages": messages,
                "tools": tools,
                "temperature": temperature,
                "max_tokens": max_tokens,
                **NO_THINKING,
            },
            timeout=300,
        )
        response.raise_for_status()
        return response.json()["choices"][0]["message"]

    return complete

async def run_agent(complete: Complete, task: dict, max_rounds: int = 3) -> dict:
    """Run one tool-calling episode and return it as a trace.

    The trace is the full conversation in OpenAI chat format plus the tool schemas, which
    is exactly what a draft model trains on.
    """
    messages: list[dict] = [
        {"role": "system", "content": AGENT_SYSTEM_PROMPT},
        {"role": "user", "content": task["request"]},
    ]
    replay = list(task["tool_results"])
    for _ in range(max_rounds):
        message = await complete(messages, task["tools"])
        assistant = {"role": "assistant", "content": message.get("content") or ""}
        if message.get("tool_calls"):
            assistant["tool_calls"] = [
                {"id": c["id"], "type": "function", "function": c["function"]}
                for c in message["tool_calls"]
            ]
        messages.append(assistant)
        if not message.get("tool_calls"):
            break
        # Replay the recorded tool output instead of calling a live API.
        for call in message["tool_calls"]:
            result = replay.pop(0) if replay else json.dumps({"status": "ok"})
            messages.append({"role": "tool", "tool_call_id": call["id"], "content": result})
    return {"messages": messages, "tools": task["tools"]}

Each episode becomes one trace: the full conversation in OpenAI chat format plus the tool schemas. Every assistant turn in it is target output, including the JSON tool-call arguments.

data.py
                
                    
                
            
                
            

                
                    
                
            
@generate_env.task(cache="auto", report=True)
async def collect_agent_traces(shard: File, concurrency: int = 32) -> Shard:
    """Run each agent task against Qwen3-8B and record the full episode."""
    import httpx
    from transformers import AutoTokenizer

    server = _get_server()
    with open(await shard.download()) as f:
        tasks = [json.loads(line) for line in f]

    semaphore = asyncio.Semaphore(concurrency)
    async with httpx.AsyncClient() as client:
        complete = http_complete(client, server.base_url)

        async def one(task: dict) -> dict | None:
            async with semaphore:
                try:
                    return await run_agent(complete, task)
                except Exception as e:  # A malformed tool schema should not sink the shard.
                    print(f"agent episode failed: {e}")
                    return None

        traces = [t for t in await asyncio.gather(*(one(t) for t in tasks)) if t]

    out = "/tmp/agent-traces.jsonl"
    with open(out, "w") as f:
        for trace in traces:
            f.write(json.dumps(trace) + "\n")

    tokenizer = AutoTokenizer.from_pretrained(TARGET_MODEL)
    stats = [data_reports.trace_stats(t, tokenizer) for t in traces]
    result = Shard(
        data=await File.from_local(out),
        kind="agent",
        rows=len(traces),
        tokens=sum(s["tokens"] for s in stats),
        target_tokens=sum(s["target_tokens"] for s in stats),
        tool_calls=sum(s["tool_calls"] for s in stats),
        failed=len(tasks) - len(traces),
        lengths=[s["target_tokens"] for s in stats],
    )
    await flyte.report.replace.aio(data_reports.agent_shard_report(result, traces, stats))
    await flyte.report.flush.aio()
    return result

Report what each shard produced

Each generation task returns a Shard: the data file plus counts the later stages and reports use, so nothing downstream has to re-read the file to know what is in it:

data.py
                
@dataclass
class Shard:
    """One shard of generated training data and what is in it."""

    data: File
    kind: str  # "chat" or "agent"
    rows: int = 0
    tokens: int = 0  # Tokens the trainer sees, prompt included.
    target_tokens: int = 0  # Tokens written by the target model: what the draft learns.
    tool_calls: int = 0
    failed: int = 0
    lengths: list[int] = field(default_factory=list)  # Target tokens per row.

@dataclass
class TrainingData:
    chat: list[Shard]
    agent: list[Shard]

    @property
    def files(self) -> list[File]:
        return [s.data for s in self.chat + self.agent]

Every data task also renders a report:

prepare_prompts shows the sampled prompts: counts, prompt lengths, and tools per agent task.
regenerate_chat shows each shard’s rows, target-written tokens, and response length distribution, with sample rows as the trainer sees them after the chat template.
collect_agent_traces shows episodes, tool-call rate, turns per episode, and the most called tools, with a rendered example episode.
build_dataset shows the final mix: chat against agent tokens, and per-shard totals.

The build_dataset report: 4,000 chat rows and 4,000 agent episodes, with target tokens by source and per example

Check these before you spend GPU hours on training. A shard with many failed episodes or a response length distribution piled up at --max-tokens means the data, not the trainer, needs fixing.

Fan out the shards

build_dataset runs the chat and agent shards concurrently:

data.py
                
@driver_env.task(report=True)
async def build_dataset(num_chat: int, num_agent: int, num_shards: int = 4) -> TrainingData:
    """Generate the full training set, every shard in parallel."""
    shards = await prepare_prompts(num_chat, num_agent, num_shards)
    with flyte.group("regenerate-chat"):
        chat = asyncio.gather(*(regenerate_chat(s) for s in shards.chat))
    with flyte.group("agent-traces"):
        agent = asyncio.gather(*(collect_agent_traces(s) for s in shards.agent))
    chat_shards, agent_shards = await asyncio.gather(chat, agent)
    data = TrainingData(chat=list(chat_shards), agent=list(agent_shards))

    await flyte.report.replace.aio(data_reports.dataset_report(data))
    await flyte.report.flush.aio()
    return data

prepare_prompts, regenerate_chat, and collect_agent_traces all use cache="auto". A second pipeline run with the same data sizes reuses every shard and goes straight to training.

Fine-tune the head in online mode

EAGLE-3 trains on the target’s hidden states from three layers. Precomputing them offline costs about 32 KB per token for an 8B model, which is terabytes for a dataset of a few hundred million tokens. In online mode, a vLLM server computes hidden states on demand and the trainer consumes and discards them.

The hidden-state server is ordinary vllm serve with two extra flags. One tells vLLM which of the target’s layers to capture. The other attaches a connector that writes them to a shared path for every request:

vllm_server.py
                
                    
                
            
                
            

                
                    
                
            
def hidden_state_args(model: str, path: str, target_layers: tuple[int, ...]) -> list[str]:
    """vLLM flags that write the target's hidden states to `path` for every request.

    The drafter reads `target_layers`: three for EAGLE-3 (early, middle, late), five for
    DFlash and DSpark. The final layer is captured too, for the training targets.
    """
    from transformers import AutoConfig

    n = AutoConfig.from_pretrained(model).num_hidden_layers
    layers = [*target_layers, n]
    speculative_config = {
        "method": "extract_hidden_states",
        "num_speculative_tokens": 1,
        "draft_model_config": {"hf_config": {"eagle_aux_hidden_state_layer_ids": layers}},
    }
    kv_transfer_config = {
        "kv_connector": "ExampleHiddenStatesConnector",
        "kv_role": "kv_producer",
        "kv_connector_extra_config": {"shared_storage_path": path},
    }
    return [
        "--speculative-config", json.dumps(speculative_config),
        "--kv-transfer-config", json.dumps(kv_transfer_config),
        # Enables /v1/chat/completions/render, which prepare-data uses to tokenize
        # raw conversations with the target's chat template.
        "--enable-scale-out",
    ]

What to train is a recipe. The EAGLE-3 recipe names the public head to start from, the three target layers it reads, its serving draft length, and its training flags:

recipes.py
                
EAGLE3 = Recipe(
    method="eagle3",
    init_from=PUBLIC_DRAFTER,
    target_layers=(2, 18, 33),
    num_speculative_tokens=3,
    lr=5e-5,
    epochs=2,
    # Training-time test: unroll the draft on its own predictions for 3 steps.
    train_args=("--ttt-steps", "3"),
)

train_drafter starts that server, then runs the trainer with torchrun against it. On /dev/shm, hidden states move between the two processes through memory and are deleted as soon as the trainer reads them:

train.py
                
                    
                
            
                
            

                
                    
                
            
@train_env.task(cache="auto", retries=1, report=True)
async def train_drafter(
    data: list[File],
    method: str = "eagle3",
    init_from: str | None = None,
    warm_start: Dir | None = None,
    epochs: int | None = None,
    lr: float | None = None,
    max_steps: int | None = None,
) -> Dir:
    """Fine-tune (or train from scratch) a draft model on target-generated data.

    `method` picks a recipe: "eagle3", "dflash", or "dspark". `warm_start` continues from a
    previous run's head. Otherwise `init_from` overrides the recipe's Hugging Face head; an
    empty string trains from scratch. Unset arguments come from the recipe.

    The task is cached: the same data and arguments return the drafter already trained, so
    the benchmark reuses what each family's pipeline trained.
    """
    recipe = RECIPES[method]
    init_from = recipe.init_from if init_from is None else init_from
    lr = recipe.lr if lr is None else lr
    epochs = recipe.epochs if epochs is None else epochs
    shared_gpu = VLLM_GPUS == TRAIN_GPUS

    # Hidden-state server: vLLM capturing exactly the layers this drafter reads.
    server = start_server(
        gpus=VLLM_GPUS,
        hidden_states_path=HIDDEN_STATES,
        target_layers=recipe.target_layers,
        extra_args=[
            # Rendering conversations that carry tool schemas needs the tool parser.
            *TOOL_ARGS,
            "--data-parallel-size", str(VLLM_REPLICAS),
            "--max-model-len", str(recipe.seq_len + 512),
            "--gpu-memory-utilization", str(recipe.vllm_memory if shared_gpu else 0.85),
        ],
    )
    start = "previous drafter" if warm_start else (init_from or "scratch")
    history = train_report.History(
        method=method,
        config={
            "method": method,
            "start from": start,
            "data files": str(len(data)),
            "epochs": str(epochs),
            "max steps": str(max_steps or "all"),
            "learning rate": str(lr),
            "target layers": ", ".join(map(str, recipe.target_layers)),
            "method flags": " ".join(recipe.train_args),
            "sequence length": str(recipe.seq_len),
            "GPUs": f"vLLM {VLLM_GPUS}, trainer {TRAIN_GPUS}",
        },
    )
    try:
        # Tokenize every file with the target's chat template. Regenerated chat files
        # are already tokenized; agent traces are rendered through the server.
        local = [_for_prepare_data(await f.download()) for f in data]
        prepare_data(local, server.base_url, "/tmp/prepared", recipe.seq_len)
        _check_prepared(local, "/tmp/prepared")

        method_args = list(recipe.train_args)
        if warm_start is not None or init_from:
            # A pretrained head fixes its own architecture and vocabulary.
            method_args = _without(method_args, ARCHITECTURE_FLAGS)
            source = await warm_start.download() if warm_start is not None else init_from
            draft_args = ["--from-pretrained", source]
        else:
            draft_args = []

        save_path = Path("/tmp/checkpoints")
        await run_and_chart(
            [
                "torchrun", "--standalone", "--nproc_per_node", str(TRAINERS),
                "-m", "speculators.train",
                "--speculator-type", method,
                "--verifier-name-or-path", TARGET_MODEL,
                *draft_args,
                "--target-layer-ids", *map(str, recipe.target_layers),
                *method_args,
                "--data-path", "/tmp/prepared",
                "--vllm-endpoint", f"{server.base_url}/v1",
                "--hidden-states-path", HIDDEN_STATES,
                "--on-missing", "generate",
                "--save-path", str(save_path),
                "--save-best",
                "--epochs", str(epochs),
                "--lr", str(lr),
                "--total-seq-len", str(recipe.seq_len),
                "--num-workers", "4",
                *(["--max-steps", str(max_steps)] if max_steps else []),
            ],
            # A wide console keeps each metric line on one line for the parser.
            env={
                **os.environ,
                "CUDA_VISIBLE_DEVICES": TRAIN_GPUS,
                # Sharing a GPU with vLLM leaves free memory fragmented; let allocations grow.
                "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
                "COLUMNS": "8000",
            },
            history=history,
        )
    finally:
        server.stop()

    # Upload only what vLLM needs to serve the head: config and weights.
    best = _best_checkpoint(save_path)
    out = Path("/tmp/drafter")
    out.mkdir()
    for f in best.iterdir():
        if f.suffix in {".json", ".safetensors", ".py"}:
            shutil.copy(f, out / f.name)
    history.checkpoint = best.name
    await flyte.report.replace.aio(train_report.render(history, done=True))
    await flyte.report.flush.aio()
    return await Dir.from_local(str(out))

The trainer prints one metric line per optimizer step. run_and_chart streams those lines, parses them, and refreshes the task’s report every 20 seconds:

train.py
                
                    
                
            
                
            

                
                    
                
            
async def run_and_chart(cmd: list[str], env: dict, history: train_report.History) -> None:
    """Run the trainer, parse its metric lines, and refresh the report as it trains."""
    process = await asyncio.create_subprocess_exec(
        *cmd,
        env=env,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.STDOUT,
        limit=16 * 1024 * 1024,  # Long metric and config lines can exceed asyncio's 64 KB.
    )
    recent: deque[str] = deque(maxlen=80)
    last_flush = 0.0
    async for raw in process.stdout:
        # The trainer's rich console pads each line to the terminal width; collapse it.
        line = re.sub(r" {3,}", "  ", raw.decode(errors="replace")).rstrip()
        print(line, flush=True)
        recent.append(line)
        if history.parse(line) and time.monotonic() - last_flush > 20:
            await flyte.report.replace.aio(train_report.render(history))
            await flyte.report.flush.aio()
            last_flush = time.monotonic()
    if await process.wait() != 0:
        errors = [l for l in recent if re.search(r"Error|error:|Traceback|OutOfMemory", l)]
        detail = "\n".join(errors[-10:] or list(recent)[-10:])
        raise RuntimeError(f"training failed with exit code {process.returncode}:\n{detail}")

The report charts training loss and draft accuracy at each position, tabulates validation metrics per epoch, and records the configuration. Position 0 is the first drafted token. Each later position counts only when every earlier one was accepted, so acceptance length is roughly 1 + acc0 + acc0 × acc1 + …. Watch position 0: if it does not rise above the public head’s starting point, more training will not help.

The EAGLE-3 training report from the full run: loss, accuracy at each draft position, and the kept checkpoint

Things to know about the training settings:

Two kinds of training data, one dataset. Regenerated chat files are already tokenized. Agent traces are raw conversations, which prepare-data renders through the vLLM server’s chat template endpoint, tool schemas included. prepare_data prepares each kind separately and merges them, and the task fails if fewer rows come out than went in:

train.py
                  
                      
                  
              
                  
              

                  
                      
                  
              
def prepare_data(paths: list[str], render_endpoint: str, out: str, seq_len: int) -> None:
      """Tokenize training files with the target's chat template into one dataset.
  
      Regenerated chat rows arrive tokenized; agent traces are rendered through the vLLM
      server. speculators 0.9 cannot concatenate the two kinds in one call (their token id
      types differ), so each kind is prepared on its own and the results merged.
      """
      from datasets import concatenate_datasets, load_from_disk
      from speculators.train.vocab_mapping import save_token_frequency_distribution
  
      groups = [g for g in ([p for p in paths if _is_tokenized(p)], [p for p in paths if not _is_tokenized(p)]) if g]
      parts = []
      for i, group in enumerate(groups):
          part = f"{out}-{i}"
          subprocess.run(
              [
                  "speculators", "prepare-data",
                  "--model", TARGET_MODEL,
                  *[arg for path in group for arg in ("--data", path)],
                  "--output", part,
                  "--seq-length", str(seq_len),
                  "--render-endpoint", render_endpoint,
              ],
              check=True,
          )
          parts.append(load_from_disk(part))
      # Cast every part to 64-bit token ids, then merge and shuffle.
      target = next((p.features for p in parts if "int64" in str(p.features["input_ids"])), parts[0].features)
      merged = concatenate_datasets([p.cast(target) for p in parts]).shuffle(seed=0)
      merged.save_to_disk(out)
      # Token frequencies pick the reduced draft vocabulary for drafters trained from scratch.
      save_token_frequency_distribution(merged, f"{out}/token_freq.pt")
  

Rendering conversations that carry tool schemas needs vLLM’s tool-calling flags on the hidden-state server. Without them, every agent trace fails to render, and without the row check, training would silently run on chat data alone.

The training-time-test depth is 3. The draft is unrolled on its own predictions for --ttt-steps steps. Match it to num_speculative_tokens at serving time, or raise it for math and code, where longer drafts get accepted.
The learning rate is low for fine-tuning. 5e-5 adapts the public head without erasing it. Training from scratch uses about 1e-4 and several epochs.
Only the config and weights are uploaded. The returned Dir is the format vLLM loads with --speculative-config, so the serving app and benchmarks consume it directly.

Run the pipeline

main generates the dataset and fine-tunes the head:

pipeline.py
                
@driver_env.task
async def main(
    num_chat: int = 4000,
    num_agent: int = 4000,
    num_shards: int = 4,
    max_steps: int | None = None,
) -> Dir:
    """Generate the dataset and fine-tune the public EAGLE-3 head on it."""
    dataset = await build_dataset(num_chat, num_agent, num_shards)
    return await train_drafter(data=dataset.files, method="eagle3", max_steps=max_steps)

pipeline.py
                
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main, **(SMOKE if "--smoke" in sys.argv else FULL))
    print(run.url)

Start with a smoke run, which uses 64 prompts per source and 20 training steps:

python pipeline.py --smoke

Then run the full pipeline:

python pipeline.py

train_drafter is cached along with the data, so rerunning with the same inputs returns the trained head immediately. Changing a training argument reruns only training. Benchmark draft models measures the head against the other families.