Verification at Scale: Confirming 17.3M Tokens of Hidden State Extraction for EAGLE-3 Training

Introduction

In the midst of a sprawling, multi-day effort to build an EAGLE-3 speculative decoding pipeline for the Kimi-K2.5 language model, one message stands as a quiet but critical milestone. Message [msg 3414] is a verification step—a brief, almost laconic check that the preceding hours of computation have produced valid results. The assistant writes:

10,000 samples, 17.3M tokens, 924 GB, 0 errors. Let me verify a sample:

Then follows a bash command that loads a single .pt file from the freshly extracted hidden states and prints its structure. The output reveals the data's anatomy: input_ids, loss_mask, and four hidden state tensors, each of shape [4096, 7168] in bfloat16. The statistics show a curious pattern—the first hidden state layer has modest min/max values (-0.42 to 1.06), while the middle layers exhibit startlingly wide ranges (-844 to 168, -644 to 133), and the final layer returns to a moderate scale (-33.5 to 28.0). All means hover near zero.

This message is not dramatic. It contains no breakthroughs, no debugging epiphanies, no architectural decisions. Yet it represents the culmination of an enormous amount of prior work: the tuning of SGLang to 90 tok/s single-stream performance, the development of a non-invasive server-side patch to capture intermediate hidden states during prefill, the careful disabling of radix cache and CUDA graphs to ensure correct extraction, the deletion of 828 GB of old vLLM-extracted states to free disk space, and the patient 2-3 hour wait for the extraction to complete across 10,000 samples. This message is the moment of truth—the instant when the assistant checks whether all that effort produced usable data.

The Reasoning and Motivation: Why Verify?

The motivation for this message is rooted in a fundamental principle of machine learning pipelines: garbage in, garbage out. The hidden state extraction is the data generation step for training an EAGLE-3 drafter—a lightweight auxiliary model that predicts multiple draft tokens in parallel for speculative decoding. If the extracted hidden states are corrupted, misaligned, or incorrectly formatted, the subsequent training run would produce a useless drafter, wasting hours of GPU time and potentially leading to incorrect conclusions about the approach's viability.

The assistant had already encountered this problem once before. In [msg 3389], the old vLLM-extracted hidden states (828 GB) were deleted to make room for the new SGLang extraction. Those old states had been used to train a previous EAGLE-3 drafter that achieved only a ~15% acceptance rate—far too low to be useful. The assistant had diagnosed that the issue was likely a hidden state mismatch between extraction and training, not insufficient data quantity. This time, the extraction method was different: instead of using vLLM's internal mechanisms, the assistant had developed a custom server-side patch for SGLang that captured hidden states at specific layers ([3, 31, 59] plus the final layer) during the prefill pass.

The verification in message [msg 3414] serves multiple purposes:

  1. Format validation: Confirming that the saved .pt files contain the expected keys (input_ids, hidden_states, loss_mask) with the correct shapes and dtypes.
  2. Semantic validation: Checking that the hidden state values are numerically reasonable—not all zeros, not all NaN, not all identical. The statistics (min, max, mean) provide a quick sanity check.
  3. Pipeline continuity: Ensuring the data format matches what the training script (04_train.py) expects, so the next step can proceed without surprises.
  4. Confidence building: After hours of unsupervised computation, a single verification provides psychological closure and confirms that the system behaved as intended.

How Decisions Were Made

The decisions visible in this message are largely implicit—they are the accumulated choices that made this verification possible and meaningful.

The decision to verify a single sample rather than run a comprehensive audit reflects a pragmatic trade-off. With 10,000 samples and 924 GB of data, a full validation would be prohibitively expensive. The assistant chooses a single representative sample (the first file in the first shard) and inspects its structure and statistics. This is a "spot check" that assumes the extraction pipeline is deterministic and consistent across samples—an assumption validated by the earlier test in [msg 3393], where a single request produced 19 tokens with a perfect match between input and dump.

The decision to use weights_only=False and map_location="cpu" in the torch.load() call shows careful handling of the data format. The weights_only=False parameter allows loading arbitrary pickle objects (not just PyTorch weights), which is necessary because the hidden state files contain dictionaries with tensor data rather than model state dicts. The map_location="cpu" ensures the tensors are loaded to CPU memory regardless of where they were saved, avoiding GPU memory pressure on the verification machine.

The decision to print min/max/mean statistics for each hidden state layer reveals what the assistant considers important: not just that the data exists, but that it has the right numerical characteristics. The extreme ranges in layers 1 and 2 (hs[1] and hs[2]) are notable—values spanning from -844 to 168 in bfloat16. This could indicate that these intermediate layers have genuinely large activations, or it could signal a numerical issue. The assistant does not flag this as a problem, suggesting that such ranges are expected for this model architecture (DeepSeek-based Kimi-K2.5 with MLA attention).

Assumptions and Potential Issues

Several assumptions underpin this verification:

Assumption 1: The sample is representative. By checking only the first file, the assistant assumes that the extraction pipeline is consistent across all 10,000 samples. This is reasonable given that the extraction script processes samples sequentially with the same server configuration, but it does not rule out edge cases—for example, very long sequences that get truncated to max_seq_len=4096 might produce different hidden state patterns, or sequences with special tokens might interact differently with the model.

Assumption 2: The layer indices are correct. The data_config.json (visible in [msg 3413]) shows layer_ids: [2, 30, 58, 60] and sglang_layers_to_capture: [3, 31, 59]. The mapping between SGLang's layer numbering (1-indexed?) and the stored indices requires careful alignment. If there's an off-by-one error, the training script would receive hidden states from the wrong layers, producing a drafter that doesn't match the model's actual computation graph.

Assumption 3: The loss mask is correct. The loss_mask tensor indicates which tokens should contribute to the training loss. If this mask is incorrectly computed (e.g., masking out tokens that should be trained on, or including padding tokens), the drafter training would be misdirected. The verification checks that the mask exists and has the right shape, but does not inspect its values.

Assumption 4: bfloat16 precision is sufficient. The hidden states are stored in bfloat16, which has 7 bits of mantissa precision. For training an EAGLE-3 drafter that predicts token-level features, this precision is typically adequate, but it's worth noting that the extreme value ranges (-844 to 168) could lead to quantization artifacts in bfloat16, especially for values near zero where the relative precision is lower.

A potential mistake: The assistant does not check for NaN or Inf values explicitly. The printed min/max values are finite, but a quick torch.isnan(hs).any() or torch.isinf(hs).any() would provide stronger assurance. In large-scale extractions, numerical issues can arise from GPU computation, especially with flash attention or reduced-precision arithmetic.

Input Knowledge Required

To understand this message, a reader needs knowledge of:

  1. EAGLE-3 speculative decoding: The technique of training a lightweight "drafter" model that predicts multiple future tokens in parallel, using hidden states from the base model as features. The drafter is trained on pairs of (hidden state, future token) extracted from the base model's forward pass.
  2. SGLang architecture: The serving framework that the assistant has configured with a custom patch to dump intermediate hidden states during the prefill (first forward pass) of each request. The patch captures states at layers [3, 31, 59] plus the final layer before the LM head.
  3. Kimi-K2.5 model architecture: A DeepSeek-based model with 7168-dimensional hidden states and MLA (Multi-head Latent Attention). The model uses 60 transformer layers (indices 0-59), and the selected layers for EAGLE-3 capture are early (3), middle (31), and late (59) to provide multi-scale representations.
  4. The speculators library format: The training pipeline expects data in a specific format: dictionaries with input_ids, hidden_states (a list of tensors, one per captured layer), and loss_mask. This format is defined by the speculators v1 data specification.
  5. The preceding debugging journey: The assistant had previously tried vLLM-based extraction, which produced 828 GB of hidden states that led to a drafter with only 15% acceptance rate. The pivot to SGLang was motivated by the hypothesis that vLLM's internal representation of hidden states differed from what the training script expected.

Output Knowledge Created

This message creates several forms of knowledge:

  1. Empirical confirmation that the SGLang-based extraction pipeline produces valid data at scale. The 10,000 samples, 17.3M tokens, 924 GB, and 0 errors demonstrate that the approach is reliable and production-ready.
  2. Numerical baseline data for the hidden state statistics at each captured layer. The observation that hs[0] (layer 3) has modest range (-0.42 to 1.06) while hs[1] (layer 31) and hs[2] (layer 59) have extreme ranges (-844 to 168, -644 to 133) provides insight into the model's internal dynamics. This could inform future decisions about quantization, normalization, or layer selection for EAGLE-3 training.
  3. Validation of the data format matching the speculators v1 specification. The keys (input_ids, hidden_states, loss_mask) and shapes ([4096, 7168]) confirm that the training script can consume this data without modification.
  4. A reference point for debugging: If the subsequent EAGLE-3 training fails, this verification provides a known-good starting point. The assistant can compare the data statistics against the training script's expectations to isolate the issue.

The Thinking Process Visible

The assistant's reasoning is compressed but visible in the structure of the verification:

  1. Summarize first: "10,000 samples, 17.3M tokens, 924 GB, 0 errors." This establishes the scope and confirms the extraction completed without failures. The 0 errors is particularly important—it means every request returned a valid response from the SGLang server.
  2. Then verify a sample: The assistant doesn't just trust the summary; it inspects actual data. The choice of data_0.pt from rows_0-2000 is the first file in the first shard—the most natural starting point.
  3. Print structural information: Keys, shapes, dtypes. This confirms the data format matches expectations.
  4. Print statistical information: Min, max, mean for each hidden state layer. This goes beyond structural validation to check that the values are numerically reasonable.
  5. Interpret the statistics implicitly: The assistant does not comment on the extreme ranges in hs[1] and hs[2], suggesting either that these are expected or that the assistant is reserving judgment until the training step reveals whether the data works. The near-zero means across all layers are a positive sign—they indicate the hidden states are properly normalized (zero-centered), which is typical for transformer hidden states. The thinking process also reveals what the assistant does not do: it does not compare the extracted hidden states against a known reference (e.g., running the model locally and comparing layer outputs), it does not verify that the loss mask correctly identifies padding tokens, and it does not check the alignment between input_ids and the original tokenized data. These omissions are reasonable given the scope of the verification—the assistant is doing a quick sanity check, not a formal audit.

Significance in the Broader Pipeline

This message sits at a critical juncture in the EAGLE-3 pipeline. The data generation phase is complete; the training phase is about to begin. The verification in [msg 3414] is the bridge between these two phases—the moment when the assistant confirms that the raw material for training is sound.

The significance of this verification is amplified by the history of the project. The previous attempt with vLLM had produced 828 GB of hidden states that led to a drafter with only 15% acceptance rate—too low to provide any speedup. The pivot to SGLang, the development of the server-side patch, the NCCL tuning to achieve 90 tok/s, and the careful configuration to disable caching all represent investments that could have been wasted if the extraction produced bad data. This verification is the first indication that those investments have paid off.

The 924 GB of extracted hidden states represent approximately 17.3 million tokens of training data, each token accompanied by hidden state vectors from four different layers of the model. This is a substantial dataset—comparable in scale to the training data used for the original EAGLE-3 paper. The 0 errors across 10,000 requests attest to the stability of the SGLang server and the extraction script.

Conclusion

Message [msg 3414] is a moment of validation in a complex engineering effort. It is the point where the assistant pauses, takes stock, and confirms that the data is real, correctly formatted, and numerically plausible before committing to the next expensive step. The verification is concise but thorough—it checks structure, statistics, and scale, providing confidence that the extraction pipeline worked as designed.

The message also reveals the assistant's engineering philosophy: trust but verify. The extraction log reported success, but the assistant still checks a sample. The numbers look right, but the assistant still inspects the tensors. This careful, empirical approach is what makes large-scale ML pipelines reliable—and it's why this brief verification message, for all its apparent simplicity, is a crucial milestone in the journey toward a working EAGLE-3 drafter for Kimi-K2.5.