Reading the Source: Decoding EAGLE-3 Training Data Formats in the Speculators Library

Introduction

In the sprawling, multi-week effort to deploy and optimize 1-trillion-parameter language models on 8×RTX PRO 6000 Blackwell GPUs, message [msg 2517] represents a quiet but pivotal moment. After a deep profiling campaign revealed that AllReduce communication was consuming 51.5% of decode time, the team pivoted to speculative decoding as a software-only optimization path. N-gram speculation had already been tested and found to be 9–26% slower than baseline — a failure that was both predictable (reasoning models generate novel thinking chains with little repetition) and instructive (MoE expert activation overhead during verification swamped any rare n-gram matches). The next candidate was EAGLE-3, a more sophisticated speculative decoding architecture that trains a lightweight draft model to predict hidden states rather than matching raw tokens.

Message [msg 2517] is the moment when the assistant transitions from high-level research into hands-on implementation. Having installed the speculators library (vLLM's official EAGLE-3 training framework) and verified that its core imports work, the assistant now needs to understand the exact data format that the training pipeline expects. The message consists of a single bash command — reading the data.py module from the installed speculators package — and its output showing the first ~60 lines of that file. On the surface, this is a mundane act of code reading. But in context, it is a critical reconnaissance step: the assistant is preparing to build a custom data generation pipeline that must produce data in precisely the format that speculators' Eagle3Dataset class expects, because the standard data generation path (which relies on vLLM internal worker APIs) is incompatible with the installed vLLM 0.16.

The Message in Full

[assistant] Let me read the key files I need to understand for the data format and training:
[bash] ssh root@10.1.230.174 'cat ~/ml-env/lib/python3.12/site-packages/speculators/train/data.py'
# ruff: noqa: ERA001
import math
import os
import random
from collections.abc import Callable
from pathlib import Path
from typing import Any

import torch
import torch.nn.functional as F  # noqa: N812
from torch.utils.data import Dataset

from speculators.train.noise_transforms import TransformTensors

BatchType = dict[str, Any]


def list_files(path):
    datapath = []
    for root, _directories, files in os.walk(path):
        for file in files:
            if not file.endswith("pt"):
       ...

The output is truncated — the command only captured the first screenful of a multi-hundred-line file. But even these opening lines reveal crucial structural information: the data format uses .pt (PyTorch serialized tensor) files, the dataset inherits from torch.utils.data.Dataset, and there is a list_files function that walks directories looking for .pt files. The presence of noise_transforms suggests data augmentation during training.

Why This Message Was Written

The assistant had just completed a multi-threaded research phase. Three parallel subagents had explored the SpecForge repo (SGLang's EAGLE-3 framework), the Speculators repo (vLLM's EAGLE-3 framework), and the existing AQ-MedAI/Kimi-K2-Instruct-eagle3 model on HuggingFace. The research revealed a critical obstacle: speculators' data generation pipeline (VllmHiddenStatesGenerator) uses internal vLLM worker extension APIs that changed between vLLM 0.15 and 0.16. The installed environment runs vLLM 0.16.0rc2, which means the standard data generation path would fail at runtime with API mismatches.

The assistant had already verified that speculators' training code imports cleanly — Eagle3DraftModel and the VllmHiddenStatesGenerator both imported without errors. But import-time success does not guarantee runtime compatibility. The data generation path monkey-patches into vLLM's internal worker processes to capture intermediate hidden states during prefill, and this is precisely where version incompatibilities manifest.

The user's directive was clear: "Start implementing the training scripts, on our existing machine with lowered numbers, then should be easy to port to much more expensive b300 machine for a simple train run." This meant the assistant needed to build a working pipeline now, on the existing 8×RTX PRO 6000 hardware, with the understanding that the production run would move to rented B300 NVL8 machines. The pipeline had to work with the current software stack — including vLLM 0.16 — which meant the standard speculators data generation path was off the table.

The assistant's strategy was to write a custom data generation script that captures hidden states from the running vLLM server using PyTorch hooks, bypassing speculators' internal vLLM integration entirely. But this custom script would need to produce data in the exact format that speculators' Eagle3Dataset expects. Hence the need to read data.py — the source code is the ground truth for the data contract.

The Reasoning Behind Reading Source Code Directly

There are several ways the assistant could have determined the data format. The speculators library has documentation, there are example scripts in the repository, and the data_generation_offline.py script in speculators shows how data is produced. But the assistant chose to read the training data module directly. This reflects a pragmatic engineering judgment: when building a custom pipeline that must interface with a third-party library, the source code is the most reliable specification. Documentation can be outdated, examples can be simplified, but the actual Dataset class defines the contract that the data must satisfy.

The assistant's comment — "Let me read the key files I need to understand for the data format and training" — reveals the deliberate, methodical approach. This is not random exploration; it is targeted reconnaissance. The assistant knows exactly what information is needed (the tensor shapes, the directory structure, the label format) and is going straight to the authoritative source.

What the Code Reveals

Even the truncated output of data.py reveals several important design decisions in the speculators training pipeline:

File format: Data is stored as .pt files — PyTorch's serialization format. This is efficient for tensor data and avoids the overhead of parsing text formats during training.
Directory walking: The list_files function recursively walks a directory tree looking for .pt files. This means the training data can be organized in subdirectories, which is useful when generating data in parallel across multiple GPUs or from multiple prompt sources.
Dataset abstraction: The dataset inherits from torch.utils.data.Dataset, the standard PyTorch dataset interface. This means it can be used with PyTorch's DataLoader for batching, shuffling, and multiprocess loading.
Noise transforms: The import of TransformTensors from speculators.train.noise_transforms indicates that the training pipeline applies data augmentation. This is common in EAGLE-3 training to improve robustness — adding small amounts of noise to the hidden states so the draft model learns to handle slight distribution shifts.
Batch type: The BatchType = dict[str, Any] type alias suggests that batches are dictionaries mapping string keys to tensors, which is the standard interface for HuggingFace-style training loops.

Assumptions Made

The assistant is making several assumptions in this message:

That the data format is stable across speculators versions. The installed version is 0.3.0, and the assistant is assuming that the Eagle3Dataset class in this version defines the format that will be expected. If the training pipeline is later run on a different machine with a different speculators version, the format might differ.

That reading the source code is sufficient to understand the format. The output is truncated — the assistant only sees the first ~60 lines of what is likely a 300+ line file. The critical parts (the __getitem__ method that defines what a single sample looks like, the __len__ method, the collate_fn if any) are further down in the file. The assistant would need to read the full file or run the dataset with sample data to verify understanding.

That the custom data generation approach is viable. The assistant is betting that capturing hidden states via PyTorch hooks on the running vLLM server will produce tensors that match what speculators' VllmHiddenStatesGenerator would have produced. This requires understanding exactly which layers' hidden states are needed (layers [2, 30, 58] for the K2 EAGLE-3 architecture), how they should be packaged, and what metadata (like attention masks) must be included.

That the training code does not depend on vLLM internals. The assistant has verified that Eagle3DraftModel imports cleanly, but the training loop might reference vLLM-specific configurations or utilities. This assumption will be tested when the training script actually runs.

Input Knowledge Required

To understand this message, the reader needs to know:

The EAGLE-3 architecture: EAGLE-3 is a speculative decoding framework where a lightweight draft model predicts future hidden states based on the target model's intermediate representations. It requires extracting hidden states from specific layers of the target model during prefill, then training the draft model to predict the hidden states of subsequent tokens.
The speculators library: This is vLLM's official library for training and deploying speculative decoding models, including EAGLE-3. It provides data generation, vocabulary mapping, model conversion, and training infrastructure.
The vLLM version constraint: Speculators 0.3.0 was designed for vLLM ≤0.15, while the installed environment uses vLLM 0.16.0rc2. The internal APIs for worker extensions changed between these versions, breaking the standard data generation path.
The hardware context: The training is being done on 8×RTX PRO 6000 Blackwell GPUs with PCIe interconnect, which limits data generation throughput. The production run will move to B300 NVL8 machines with NVLink for faster inter-GPU communication.
The model architecture: Kimi-K2.5 is a MoE (Mixture of Experts) model with approximately 1 trillion parameters. Its architecture includes a multimodal wrapper (model.language_model.model.layers instead of the standard model.model.layers), which complicates hidden state extraction.

Output Knowledge Created

This message creates knowledge in several forms:

Implicit knowledge for the assistant: The assistant now understands the data contract that the training pipeline expects. This enables the next step: writing a custom data generation script that produces .pt files in the expected format.

Documentation in the conversation: The output of the bash command is captured in the conversation log, serving as a reference for future steps. If the assistant needs to revisit the data format, it can scroll back to this message rather than re-reading the file.

A foundation for the pipeline: This message is the first concrete implementation step after the research phase. It transitions the conversation from "what should we build" to "how do we build it." The subsequent messages will build on this understanding to create the dataset preparation script, the hidden state extraction script, and the training orchestrator.

The Thinking Process Visible in the Message

The assistant's thought process is revealed in the framing comment: "Let me read the key files I need to understand for the data format and training." This indicates:

Goal-directed reading: The assistant is not browsing randomly. It has a specific information need — the data format — and is reading the file that defines it.
Prioritization: The assistant could have read any of several files (the training script, the data generation script, the model definition). It chose data.py first because the data format is the critical interface point for the custom pipeline.
Recognition of the bottleneck: The assistant understands that the custom data generation script is the hardest part of the pipeline (because it must interface with the running vLLM server), and getting the data format right is essential before writing that script.
Methodical approach: The assistant is building the pipeline piece by piece: first understand the data format, then write the data generation script, then the vocabulary mapping, then the training script, then the orchestrator. This is visible in the todo list from the previous message.

The Broader Significance

Message [msg 2517] exemplifies a pattern that recurs throughout the opencode session: the assistant repeatedly encounters API incompatibilities and version mismatches, and responds by reading source code and building custom workarounds. Earlier in the session, the assistant had to patch the vLLM GGUF loader to force-dequantize indexer weights, fix the Triton MLA attention backend for SM120 GPUs, and modify the speculators library to handle the Kimi-K2.5 multimodal wrapper architecture. Each time, the approach is the same: understand the contract (what the code expects), understand the reality (what the environment provides), and build a bridge between them.

This message also illustrates the tension between using existing frameworks and building custom solutions. The assistant initially tried to use speculators' standard data generation path, but the vLLM 0.16 incompatibility forced a custom approach. The custom approach is more work upfront but gives the assistant full control over the data generation process, which is valuable when debugging the inevitable issues that arise when training a draft model for a 1T-parameter MoE model on unfamiliar hardware.

Conclusion

Message [msg 2517] is a small but revealing moment in a complex engineering effort. On the surface, it is simply reading a file. But in context, it represents the transition from research to implementation, the recognition of a critical interface contract, and the beginning of a custom pipeline that will ultimately enable EAGLE-3 training for Kimi-K2.5. The assistant's methodical approach — understanding the data format before writing code, reading source code as ground truth, and documenting findings in the conversation — reflects the disciplined engineering practice required to deploy cutting-edge AI models on non-standard hardware configurations. The message is a reminder that in machine learning engineering, the most important code is often not the model architecture or the training loop, but the data pipeline that connects them.