The Unified Prep Script: Orchestrating a 10× Dataset Scale-Up for EAGLE-3 Training

Introduction

In the sprawling, multi-day effort to deploy a speculative decoding draft model for the Kimi-K2.5 large language model, few moments are as quietly pivotal as message [msg 3663]. On its surface, the message is unremarkable: the assistant announces its intent to write a unified dataset preparation script, executes a write tool call, and receives confirmation that the file was saved. But this brief exchange represents the transition from planning to execution at a critical juncture — the moment when a 10× scale-up of the training dataset shifts from a theoretical plan to concrete code. Understanding why this message was written, what decisions it embodies, and what assumptions it encodes reveals the architecture of a complex ML pipeline and the reasoning processes of an AI agent managing it.

The Context: A Pipeline at a Crossroads

To appreciate message [msg 3663], one must understand the state of the project at that moment. The assistant had just resolved a devastating bug in the EAGLE-3 speculative decoding system. For days, the trained draft model had achieved a zero acceptance rate — every token it proposed was rejected by the target model. The root cause turned out to be a single flag mismatch: the SGLang server was started with --speculative-algorithm EAGLE instead of EAGLE3. This caused the target model to pass only the final-layer hidden states (7168 dimensions) to the draft model instead of the concatenated intermediate-layer states (21504 dimensions) that the draft model's fusion layer expected. The trained weights were silently bypassed, rendering them useless.

With the flag corrected, the draft model finally achieved an acceptance length of ~2.1 tokens — meaning it could predict about two tokens correctly before the target model had to step in. However, benchmarking revealed a sobering reality: even with this fix, the EAGLE-3 system achieved only 82.3 tokens per second, compared to the 90 tok/s non-speculative baseline. The draft model was actually slowing the system down by ~9%. The overhead of running the draft model and verifying its predictions outweighed the benefit of its modest acceptance rate.

The EAGLE-3 research paper's scaling curves offered a clear diagnosis: more training data was the primary lever for improving acceptance rates. The assistant had been training on approximately 10,000 samples. To achieve meaningful speedup, the dataset needed to scale by an order of magnitude — to roughly 100,000 samples.

This set the stage for a massive data pipeline effort. The user approved a custom agentic-heavy dataset mix and had the disk resized to 11TB to accommodate the hidden states. The assistant laid out a detailed six-phase plan in train_plan_v4.md ([msg 3661]), covering dataset download, response regeneration through Kimi-K2.5, hidden state extraction, vocabulary mapping, training, and deployment. The user then gave the instruction to execute: "Start a general agent to prep each dataset into the correct format... then merge the datasets and run hidden state extraction."

The Message: Writing the Orchestrator

Message [msg 3663] is the assistant's response to that instruction. It reads:

Now let me write the unified dataset prep script that each agent will use. Each dataset needs a prep function that outputs either prompts.jsonl (for inference) or tokenized_data.jsonl (for Kimi-native data).

The assistant then executes a write tool call to create /home/theuser/glm-kimi-sm120-rtx6000bw/eagle3-train/datasets/prep_all.py, and receives confirmation that the file was written successfully. The message also includes LSP diagnostics from an unrelated file (server_args_sm120.py), which are noise from the development environment and not relevant to the dataset prep work.

The brevity of this message is deceptive. In a single sentence, the assistant encapsulates a design decision that will govern how ten separate datasets, drawn from diverse sources across the Hugging Face ecosystem, are normalized into a common format for downstream processing. The message is the visible tip of a much larger iceberg of reasoning that happened in the assistant's internal deliberation.

WHY: The Reasoning Behind a Unified Script

The assistant's choice to write a single unified script rather than ten separate dataset-specific scripts reflects several layers of reasoning.

First, the principle of code reuse. Each of the ten datasets requires a similar workflow: download from Hugging Face, parse the dataset's specific format, extract the relevant fields (prompts or full conversations), and write them to a standardized output format. Ten separate scripts would mean ten copies of boilerplate for Hugging Face authentication, JSON handling, file I/O, and error handling. A unified script with one function per dataset is more maintainable and less error-prone.

Second, the need for parallel execution. The user explicitly asked to "start a general agent to prep each dataset." The assistant's architecture supports spawning parallel subagents via the task tool. A unified script means each subagent can call the same module with a different dataset name argument, rather than requiring ten separate entry points. This is an architectural decision that enables the parallel launch that follows in subsequent messages ([msg 3664] onward).

Third, the binary classification of datasets. The assistant's sentence reveals a crucial design insight: the ten datasets fall into exactly two categories. Some datasets (like crownelius/KimiK2.5-2000x and SWE-Factory/DeepSWE-Agent-Kimi-K2-Trajectories-2.8K) already contain text generated by Kimi models. For these, the responses already match the target model's token distribution, so they can be tokenized directly and written as tokenized_data.jsonl. The remaining datasets (like nebius/SWE-agent-trajectories, nvidia/OpenCodeInstruct, HuggingFaceH4/ultrachat_200k) contain prompts written by humans or generated by other models. For these, the responses must be regenerated through Kimi-K2.5 inference, so only the prompts are extracted and written as prompts.jsonl. This binary split is the core organizational principle of the entire data pipeline.

HOW: Design Decisions Embedded in the Approach

The message encodes several design decisions that shape the entire pipeline.

The output format decision. The assistant specifies two output formats: prompts.jsonl for inference-bound datasets and tokenized_data.jsonl for Kimi-native datasets. The .jsonl format (JSON Lines, one JSON object per line) is a standard choice for ML data pipelines because it is line-delimited, streaming-friendly, and easy to inspect. The tokenized_data.jsonl format was already established by the existing pipeline scripts (the assistant had just inspected 01b_generate_synthetic.py in [msg 3660] to confirm the format). By reusing this format, the assistant ensures backward compatibility with the existing extraction and training scripts.

The function-per-dataset architecture. Each dataset gets its own prep function within prep_all.py. This modular design acknowledges that each Hugging Face dataset has a unique schema. nebius/SWE-agent-trajectories contains multi-turn agent trajectories where the prompt is the user's initial issue description. nvidia/OpenCodeInstruct uses an input field. open-r1/Mixture-of-Thoughts stores conversations in a messages structure. A single monolithic function would be unwieldy; separate functions allow each dataset's idiosyncrasies to be handled cleanly while sharing common infrastructure.

The inference pipeline dependency. The choice to output prompts.jsonl for most datasets creates a downstream dependency: these prompts must be fed through the Kimi-K2.5 inference server. This means the assistant must kill the existing EAGLE-3 speculative decoding server and start a non-speculative baseline server for inference — exactly what happens in the following messages ([msg 3666]). The unified script thus defines the interface between the data preparation phase and the inference phase.

Assumptions Made by the Assistant

The message rests on several assumptions, some explicit and some implicit.

The tokenizer is available. The assistant assumes that the Kimi-K2.5 tokenizer is accessible at a known path on the remote machine. This is a reasonable assumption given that the model is already deployed, but it creates a dependency on the environment setup.

Hugging Face datasets are accessible. The script assumes that all ten datasets can be downloaded from Hugging Face without authentication issues or rate limiting. In practice, some datasets may require login tokens or may have access restrictions.

The output directories exist. The assistant had just created the datasets directory in [msg 3662], but the subdirectories for each dataset's prepared output would need to be created by the script or by the launch commands.

The binary classification is correct. The assistant assumes that datasets like SWE-Factory/DeepSWE-Agent-Kimi-K2-Trajectories-2.8K truly contain Kimi-generated text that can be used directly. If these datasets contain mixed-source text or if the "Kimi" label is inaccurate, the training data would be contaminated with out-of-distribution responses.

The prompts.jsonl format is sufficient. The assistant assumes that extracting only the prompt is sufficient for inference. But some datasets (like SWE-agent-trajectories) contain multi-turn interactions where context matters. The assistant's plan to use "the user's initial issue as prompt" may lose important conversational context.

Input Knowledge Required

To understand and write this message, the assistant needed a substantial body of knowledge:

• The EAGLE-3 architecture: How speculative decoding works, the role of hidden states, the importance of matching the target model's token distribution.
• The existing pipeline scripts: The format of tokenized_data.jsonl, the structure of 01b_generate_synthetic.py and 02b_extract_hidden_states_sglang.py, which the assistant had just read in <msg id=3659-3660>.
• The ten datasets: Their schemas, sizes, and content types — gathered from the dataset research phase ([msg 3655]).
• Hugging Face datasets library: How to load, filter, and iterate over datasets programmatically.
• The SGLang inference server: How to send prompts and receive responses, the concurrency settings (C=150-200), the context length limits.
• The tokenizer: How to tokenize text into input_ids and create loss_mask arrays for the training format.

Output Knowledge Created

This message creates a concrete artifact: the prep_all.py script. But it also creates several forms of knowledge:

• A reusable data preparation module that can be invoked by parallel agents for each dataset.
• A standardized interface between data preparation and inference: prompts.jsonl files that the inference pipeline can consume.
• A template for dataset handling that can be extended if additional datasets are added later.
• Documentation of the two-category classification of datasets, which clarifies the pipeline logic for anyone reading the code.

The Thinking Process Visible in the Message

The assistant's reasoning is visible in the structure of its single sentence. The phrase "Now let me write the unified dataset prep script that each agent will use" reveals the assistant's mental model of the execution plan. It envisions a multi-agent architecture where parallel subagents each handle one dataset, all sharing a common codebase. The word "unified" is a deliberate design choice — it signals that the assistant considered and rejected the alternative of separate scripts.

The second clause — "Each dataset needs a prep function that outputs either prompts.jsonl (for inference) or tokenized_data.jsonl (for Kimi-native data)" — shows the assistant performing a conceptual grouping of the ten datasets into two categories. This is a moment of abstraction: the assistant recognizes that despite the diversity of sources and schemas, all datasets can be handled by one of two workflows. This binary classification simplifies the pipeline architecture and reduces cognitive load.

The fact that the assistant writes the file immediately after stating its intent (rather than, say, writing it incrementally or asking for confirmation) indicates confidence in the design. The assistant has already gathered all necessary information — it inspected the existing scripts, confirmed the tokenizer path, and reviewed the dataset list. The write is an execution step, not an exploratory one.

What Follows

The messages immediately after [msg 3663] show the unified script being deployed. In [msg 3664], the assistant attempts to SCP the script to the remote machine but hits a "No such file or directory" error because the target directory doesn't exist on the remote host. It recovers in [msg 3665] by creating the directory first, then successfully transfers the file. In [msg 3666], the assistant kills the EAGLE-3 server to prepare for the inference pipeline. In [msg 3667], it creates the output directories and launches all ten dataset preps in parallel.

This sequence reveals the assistant's execution model: write the code, deploy it to the remote machine, prepare the environment, and launch. The unified script is the foundation that makes the parallel launch possible. Without it, the assistant would have to write, deploy, and manage ten separate scripts — a significantly more complex orchestration problem.

Conclusion

Message [msg 3663] is a study in concise decision-making. In a single sentence and a file write, the assistant establishes the architecture for a 10× dataset scale-up, encodes a binary classification of data sources, defines the interface between preparation and inference, and sets the stage for parallel execution. The brevity of the message belies the depth of reasoning behind it — the assistant's understanding of the pipeline's data flow, its recognition of common patterns across diverse datasets, and its architectural choices for scalability and maintainability. For anyone studying how AI agents manage complex ML pipelines, this message offers a window into the moment when planning crystallizes into code.