Peering into the Black Box: How One Debugging Message Uncovered the Hidden State Structure of an EAGLE-3 Draft Model

The Moment of Investigation

In the midst of a protracted debugging session spanning dozens of messages, a single command execution stands out as a pivotal moment of empirical discovery. The message at <msg id=4385> captures the assistant running a Python one-liner over SSH to inspect the contents of a training data file — specifically, the hidden states stored alongside token sequences during the EAGLE-3 training pipeline for the Kimi-K2.5 model. The output reveals something that had been implicit in the training code but never explicitly verified: the precise structure and dimensionality of the hidden states being fed into the draft model's fusion layer.

[assistant] So we have `hidden_states` as a list of tensors. Let me check how many hidden states and their shapes:
[bash] ssh root@10.1.230.174 '~/ml-env/bin/python3 -c "
import torch
d = torch.load(\"/data/eagle3/synth_100k/hidden_states/rows_0-2000/data_0.pt\", weights_only=False)
hs = d[\"hidden_states\"]
print(f\"Num hidden states: {len(hs)}\")
for i, h in enumerate(hs):
    print(f\"  hs[{i}]: shape={h.shape}, dtype={h.dtype}, norm={h.float().norm(dim=-1).mean():.4f}\")
print(f\"input_ids: shape={d[\"input_ids\"].shape}\")
print(f\"First 5 input_ids: {d[\"input_ids\"][:5].tolist()}\")
print(f\"loss_mask sum: {d[\"loss_mask\"].sum()}, total: {len(d[\"loss_mask\"])}\")
"'
Num hidden states: 4
  hs[0]: shape=torch.Size([1281, 7168]), dtype=torch.bfloat16, norm=1.0383
  hs[1]: shape=torch.Size([1281, 7168]), dtype=torch.bfloat16, norm=37.1650
  hs[2]: shape=torch.Size([1281, 7168]), dtype=torch.bfloat16, norm=174.3977
  hs[3]: shape=torch.Size([1281, 7168]), dtype=torch.bfloat16, norm=111.3359
input_ids: shape=torch.Size([1281])
First 5 input_ids: [163587, 2482, 163601, 7833, 398]
loss_mask sum: 1235, total: 1281

This seemingly simple inspection — four lines of Python wrapped in an SSH command — represents a critical turning point in the debugging effort. It is the moment the assistant moves from reasoning about the system in the abstract to actually examining the concrete data flowing through the pipeline. To understand why this message matters, we must trace the chain of reasoning that led to it, the decisions embedded within it, and the knowledge it produced.

The Debugging Context: Why This Message Was Written

The assistant had been engaged in a multi-hour effort to deploy and tune an EAGLE-3 speculative decoding system for the Kimi-K2.5 language model. The system uses a small "draft" model to predict multiple future tokens, which the large "target" model then verifies in parallel — a technique that can dramatically accelerate inference when the draft model's predictions are accurate.

The problem was stark: despite training the draft model to 74.7% validation accuracy, the deployed system was achieving only 54.8 tokens per second, far below the 90 tok/s baseline of the target model alone. Even more concerning, the "accept length" — the average number of draft tokens accepted before a rejection — was hovering around 1.8 out of 16 drafted tokens. This meant the draft model was essentially useless: the overhead of generating 16 draft tokens was not offset by meaningful acceptance.

The assistant had already discovered and fixed one critical bug: the --speculative-num-steps 1 argument was silently overriding the --speculative-num-draft-tokens 16 setting, reducing the draft chain to just 2 tokens. But fixing that only made performance worse (46.7 tok/s), confirming that the draft model itself was the bottleneck, not a configuration issue.

At this point, the user suggested a crucial diagnostic step: "Try with a training sample and see if accept rate looks correct. Possible we didn't wire in the model correctly." This prompted the assistant to bypass the entire SGLang serving infrastructure and test the draft model in isolation — a textbook debugging maneuver that eliminates layers of complexity to isolate the root cause.

The Reasoning Behind the Inspection

The assistant's decision to inspect the hidden states file was motivated by a specific hypothesis: the wiring between the training pipeline and the inference pipeline might be incorrect. The EAGLE-3 architecture works by concatenating hidden states from multiple layers of the target model (the "auxiliary hidden states") and feeding them through a fusion layer (fc) that projects them down to the model's hidden dimension. If the training pipeline stored hidden states in one format but SGLang expected another, the draft model would receive garbled input and produce garbage predictions.

The assistant had already discovered the structure of the training data in the preceding message ([msg 4384]): the .pt file contained a dictionary with input_ids, hidden_states (a list of tensors), and loss_mask. But the crucial details — how many hidden states, their shapes, their norms — remained unknown. The command in message 4385 was designed to answer exactly these questions.

The choice of metrics to print was deliberate. The shape reveals the dimensionality (7168 is the hidden size of Kimi-K2.5), and the number of tensors in the list reveals how many layers' outputs were captured. The norm values provide a fingerprint: the embedding output (before any transformer layers) typically has a norm near 1.0, while deeper layer outputs have much larger norms due to the accumulation of residual connections. The assistant was using these norms to identify which hidden state corresponded to which layer — a form of reverse engineering the training data format.

Assumptions Embedded in the Investigation

This message rests on several assumptions, some explicit and some implicit. The most important assumption is that the hidden states stored in the training data are the same hidden states that SGLang would capture during inference. If the training pipeline and the inference pipeline captured hidden states from different layers, or in a different order, the draft model would fail regardless of its training accuracy.

A second assumption is that the .pt file format is consistent across all training samples. The assistant loaded only the first file (data_0.pt) from the first shard (rows_0-2000), implicitly assuming it represents the general structure. This is a reasonable assumption for a well-engineered data pipeline, but it introduces a small risk of sampling bias.

A third assumption is that weights_only=False is safe. Modern PyTorch warns against loading pickled data without weights_only=True due to security risks, but in this debugging context — on a trusted machine — the assistant prioritized functionality over security.

The Knowledge Required to Interpret This Message

Understanding this message requires familiarity with several layers of technical knowledge. First, one must understand the EAGLE-3 architecture: it's a speculative decoding framework where a small draft model predicts tokens using hidden states captured from intermediate layers of the large target model. The draft model has a fusion layer (fc) that concatenates multiple hidden states and projects them to a single hidden dimension.

Second, one must understand the Kimi-K2.5 model architecture: it's a Mixture-of-Experts model based on DeepSeek-V2, with 60 transformer layers and a hidden dimension of 7168. The norms of the hidden states — 1.04 for the first, 37-174 for the subsequent ones — immediately signal that hs[0] is the embedding output (before any transformer computation amplifies the values), while hs[1] through hs[3] are outputs from deeper layers.

Third, one must understand the training data pipeline: the loss_mask distinguishes between prompt tokens (where loss is not computed) and response tokens (where loss is computed). The sum of 1235 out of 1281 total tokens means 46 prompt tokens and 1235 response tokens — consistent with the earlier metadata showing n_prompt=46 and n_resp=1235.

Fourth, one must understand the bfloat16 format and its implications for numerical precision during inference. The hidden states are stored in bfloat16, which has reduced precision compared to float32 but is standard for GPU inference.

The Knowledge Produced by This Message

The output of this command produces several critical pieces of knowledge:

There are exactly 4 hidden states per sample. This confirms that the training pipeline captures the embedding output plus three auxiliary hidden states from intermediate layers. The EAGLE-3 draft model's fc layer expects a concatenation of 3 * hidden_size = 21504 dimensions, which matches hs[1], hs[2], and hs[3] concatenated. But the presence of hs[0] (the embedding) raises an immediate question: is the training pipeline using cat([hs[0], hs[1], hs[2]]) (embedding + first two aux states) or cat([hs[1], hs[2], hs[3]]) (three aux states without embedding)?

The norms reveal the layer identity. The embedding output has a norm of approximately 1.0, while the auxiliary hidden states have norms of 37, 174, and 111. These correspond to layers 3, 31, and 59 (0-indexed) of the 60-layer Kimi-K2.5 model — the layers specified in the draft model's config as eagle_aux_hidden_state_layer_ids: [2, 30, 58] (which SGLang converts to 1-indexed: [3, 31, 59]).

The sequence is 1281 tokens long with 1235 response tokens. This confirms the training sample has a substantial response portion, which is necessary for the draft model to learn meaningful predictions.

The token IDs are consistent with the Kimi-K2.5 vocabulary. The first token 163587 is a special token (BOS or control token), followed by normal token IDs in the 1000-60000 range.

The Thinking Process Visible in the Message

The message reveals a methodical, hypothesis-driven approach to debugging. The assistant does not simply dump the file contents — it asks specific questions: how many hidden states, what are their shapes, what are their norms, what are the token IDs, what does the loss mask look like. Each question targets a specific potential failure mode.

The choice to print norms is particularly insightful. In transformer models, the scale of hidden states varies dramatically between layers. The embedding layer produces values with norm near 1.0 (since embeddings are typically initialized with unit variance), while deeper layers accumulate larger values through residual connections. By printing norms, the assistant can identify which hidden state is the embedding without needing to consult metadata — a form of numerical fingerprinting.

The assistant also prints the first 5 token IDs, which serves as a sanity check: are these real tokens from the Kimi-K2.5 vocabulary? The values (163587, 2482, 163601, 7833, 398) are all within the valid range of 0-163840, confirming the data is well-formed.

The loss_mask check — summing to 1235 out of 1281 — confirms that the mask correctly identifies 1235 response tokens. This is important because the draft model should only be trained on response tokens, not prompt tokens. If the mask were wrong, the training would have been corrupted from the start.

What This Message Enabled Next

This inspection directly enabled the next phase of debugging. In the following messages, the assistant uses this knowledge to build a standalone test that loads the draft model weights and runs them against the actual hidden states from training data. The discovery that there are 4 hidden states (not 3) leads to a critical realization: the training pipeline concatenates [embed_output, layer3, layer31] (taking the first 3 of 4 hidden states), while SGLang was passing [layer3, layer31, layer59] (the 3 auxiliary hidden states without the embedding). This mismatch — confirmed in subsequent messages — was the root cause of the poor acceptance rate.

The message at <msg id=4385> thus serves as the empirical foundation for the entire debugging effort. Without this inspection, the assistant would have continued to chase configuration issues or model quality problems, never realizing that the fundamental wiring between training and inference was wrong. It is a testament to the power of "look at the data" as a debugging strategy — a principle that applies as much to machine learning systems as to traditional software engineering.