Systematic Code Archaeology: How an AI Assistant Mapped the EAGLE-3 Training Infrastructure

Introduction

In the landscape of large language model engineering, few tasks demand as much precision and depth as understanding an unfamiliar training library. Unlike using well-documented frameworks like Hugging Face Transformers or PyTorch Lightning, working with a research-grade codebase like the speculators library (version 0.2.0) requires a fundamentally different approach: systematic code archaeology. This article examines a remarkable sequence of ten messages in an opencode coding session where an AI assistant, tasked with exploring the speculators library's training API, executed a methodical reconnaissance of an unfamiliar codebase and synthesized its findings into a comprehensive, actionable analysis [7][10].

The session's broader context is itself impressive: the team had already set up an Ubuntu 24.04 ML environment with multiple RTX PRO 6000 Blackwell GPUs, resolved nightmarish flash-attn build issues requiring a secondary CUDA 12.8 toolkit, and deployed the GLM-5-NVFP4 model using SGLang. But the specific challenge addressed in this chunk was different: the user needed to rewrite a training script (04_train.py) to properly use the speculators library's native training API, and that required a complete understanding of how the library expected training to work.

The Opening Gambit: Defining the Exploration Mission

The chunk opens with the user issuing a detailed exploration request ([msg 0]). This is not a casual question—it is a structured reconnaissance mission with specific deliverables [7]. The user lists seven files to read, six directories to list, and seven key questions to answer. The questions reveal the user's mental model of what a training API should contain:

• What arguments does the Trainer class expect? — The orchestration layer
• How does Eagle3SpeculatorConfig work? — The configuration system
• How does the CLI orchestrate training? — The entry point
• What data format does training expect? — The data pipeline
• How are t2d/d2t mappings used? — The vocabulary bridge
• How are verifier weights extracted? — The model initialization
• What is the complete training loop? — The execution flow The user's request also contains assumptions that will prove partially incorrect. The path speculators/speculators/eagle3/ contains a double speculators/ that doesn't exist in the actual package structure. The user expects a dataset.py file when the actual file is data.py. And the user assumes the CLI has a training command—it does not. These are productive mistakes: they force the assistant to discover the actual structure through exploration rather than confirmation.

The First Round: Discovery Through Failure

The assistant's first response ([msg 1]) demonstrates a critical principle of systematic exploration: start broad, then narrow. The assistant issues four parallel SSH commands:

1. List the top-level directory — to understand the package's overall structure
2. List the train directory recursively — to discover all training-related files
3. List the assumed eagle3 path — which fails with "No such file or directory"
4. Read the train directory listing — already covered by command 2 The failure of the third command is actually the most informative result. It tells the assistant that the eagle3 model files are not where expected. This is a deliberate feature of the exploration strategy: by testing assumptions early, the assistant can correct course before investing effort in reading files from the wrong location. The successful commands reveal the package structure: a train/ directory with files like trainer.py, data.py, checkpointer.py, distributed_batch_sampler.py, logger.py, noise_transforms.py, and utils.py. The top-level directory shows config.py (13,567 bytes), model.py (26,789 bytes), and __main__.py (5,578 bytes). The assistant now has a map of the territory.

Course Correction: Finding the Actual Eagle3 Location

In the second round ([msg 2]), the assistant adapts to the failed assumption. Instead of retrying the wrong path, it uses find to locate the eagle3 directory:

find /root/ml-env/lib/python3.12/site-packages/speculators/ -type d | head -30

This reveals the correct path: speculators/models/eagle3/. The assistant then reads four key files in parallel:

• __main__.py — The CLI entry point, revealing that only a convert command exists (no train subcommand)
• trainer.py — The Trainer class and TrainerConfig
• data.py — The Eagle3SampleFileDataset and collation utilities
• noise_transforms.py — The data augmentation transforms This parallel read strategy is efficient but carries an implicit assumption: none of these reads depend on the results of the others. The assistant is making educated guesses about what's important based on filenames and the directory structure discovered in the previous round.

Drilling Deeper: The Model Architecture

The third round ([msg 3]) represents a shift from the training infrastructure to the model architecture itself. Having consumed the training loop and data pipeline, the assistant now needs the other half of the picture: what exactly is being trained?

The assistant issues seven parallel reads targeting:

1. The eagle3 directory listing — confirming the model files exist
2. The eagle3 __init__.py — revealing the public API exports Eagle3DraftModel and Eagle3SpeculatorConfig
3. The main config.py — the Pydantic-based configuration system
4. The main model.py — the base model classes
5. The checkpointer.py — checkpoint save/load logic
6. The train/utils.py — distributed training setup
7. The train/__init__.py — confirming the train module's exports This reading strategy reveals the assistant's mental model of the codebase's dependency chain. It's not reading randomly—it's following the logical flow: entry point → training loop → data pipeline → model architecture → configuration → persistence → infrastructure.

The Core Revealed: Eagle3DraftModel and Its Forward Pass

The fourth round ([msg 4]) drills into the heart of the EAGLE-3 architecture. The assistant reads four critical files:

• models/eagle3/config.py — The Eagle3SpeculatorConfig class with its transformer_layer_config, draft_vocab_size, norm_before_residual, target_hidden_size, and eagle_aux_hidden_state_layer_ids fields
• models/eagle3/core.py — The Eagle3DraftModel class with its multi-step autoregressive forward pass
• models/eagle3/attention.py — The flex-attention utilities for block-sparse attention masks
• models/eagle3/model_definitions.py — The custom LlamaDecoderEagle3FirstLayer that accepts 2× hidden_size input The forward pass documented in core.py is the most architecturally significant discovery. Unlike a standard language model that predicts one token at a time, Eagle3DraftModel runs a loop of ttt_steps iterations within a single forward call. At each step, it: 1. Embeds the current token IDs using the verifier's frozen embedding layer 2. Concatenates embeddings with verifier hidden states (creating a 2× hidden_size input) 3. Runs through decoder layers (the first layer handles the doubled dimension) 4. Computes logits via the trainable LM head 5. Computes KL divergence loss against the verifier's softmax targets 6. Generates draft tokens via argmax and maps them back to verifier vocabulary 7. Extends the attention mask for the next step This design is elegant but complex. The loss is a sum of KL divergences across all ttt_steps, weighted by ttt_step_loss_decay ** step. The model tracks both full accuracy (top-1 match) and conditional accuracy (accuracy conditioned on previous steps being correct). The use_off_policy_tokens flag enables teacher forcing by replacing predicted tokens with ground-truth shifted tokens.

The Supporting Cast: Vocabulary Mapping, Checkpointing, and Data Generation

The fifth and sixth rounds ([msg 5] and [msg 6]) fill in the remaining infrastructure. The assistant reads:

• vocab_mapping.py — The build_vocab_mappings_from_distribution() function that selects the top draft_vocab_size most frequent tokens and builds the t2d (boolean mask) and d2t (offset tensor) mappings
• distributed_batch_sampler.py — The MultipackDistributedBatchSamplerV2 for packing variable-length sequences across GPUs
• logger.py — The metric logging system supporting TensorBoard, Weights & Biases, and Trackio
• utils/loading.py — The load_model_layers() function for downloading verifier weights from Hugging Face Hub using safetensors shards
• data_generation/ — The VllmHiddenStatesGenerator and preprocessing pipeline that extracts hidden states from intermediate verifier layers during vLLM prefill The vocabulary mapping module is particularly important. The t2d tensor is a boolean mask of shape [verifier_vocab_size] where True indicates that the corresponding verifier token ID is included in the draft vocabulary. The d2t tensor stores offsets such that the verifier token ID can be computed as draft_idx + d2t[draft_idx]. This offset-based mapping is efficient: because selected tokens are sorted, consecutive draft indices map to consecutive verifier IDs with a constant offset.

The Final Pieces: Convert, Proposals, and Version

The seventh and eighth rounds ([msg 7] and [msg 8]) are cleanup reads. The assistant examines:

• The convert/ module — The convert_model() function and the eagle3-specific converter that packages trained draft models for inference deployment
• The proposals/greedy.py — The GreedyTokenProposalConfig that defines how the draft model generates candidate tokens during inference (single chain of speculative_tokens tokens)
• The models/__init__.py — Confirming the full model registry includes Eagle, Eagle3, Independent, and MLP speculators
• The version.py — Confirming the library is at version 0.2.0 These reads are less critical than the earlier ones, but they complete the picture. The assistant now knows that the library supports multiple speculative decoding algorithms (not just EAGLE-3), that the conversion pipeline exists for deployment, and that the proposal method is configurable.

The Synthesis: Message 9 as Architectural Reconstruction

The ninth message ([msg 9]) is the culmination of the entire exploration. It is not merely a summary of what was read—it is a complete architectural reconstruction of the EAGLE-3 training pipeline, organized not by file but by logical component. The assistant presents:

1. A complete package tree showing every module and submodule
2. The CLI entry point and its limitations (convert only, no train command)
3. The Eagle3SpeculatorConfig with all fields and a construction example
4. The Eagle3DraftModel constructor and forward pass, documented step by step
5. The custom first decoder layer and its dual-normalization design
6. The TrainerConfig and Trainer with setup sequence and training loop
7. The dataset pipeline including file format, standardization, shifting, and collation
8. Noise transforms for data augmentation
9. Distributed training with FSDP2 setup
10. Vocabulary mapping with t2d/d2t construction
11. Data generation with VllmHiddenStatesGenerator
12. Checkpointing with per-epoch save directories
13. A complete, executable training recipe (~70 lines of Python) The training recipe is the most valuable deliverable. It ties together every component discovered during the exploration into a working script that can be executed directly. The recipe demonstrates that the assistant's understanding is not just theoretical—it is complete enough to produce a functioning program.

The Architecture of Understanding

What makes this sequence of messages remarkable is not just the thoroughness of the exploration, but the structure of the understanding it produced. The assistant worked from the outside in:

1. Package structure → What modules exist
2. Entry points → How training is launched
3. Training loop → What happens during training
4. Data pipeline → What data looks like
5. Model architecture → What is being trained
6. Configuration → How everything is parameterized
7. Persistence → How progress is saved
8. Infrastructure → How distributed training works This is a textbook approach to understanding an unfamiliar codebase. Each layer of understanding builds on the previous one, and each file read answers questions that previous files left open. The assistant also made several critical discoveries that directly impacted the parent session's work: - The CLI has no training command — Training must be done programmatically - The data format is .pt files — Not raw text, but pre-processed tensors - The model requires t2d/d2t mappings — These must be built from token frequency distributions - The verifier weights are downloaded from Hugging Face Hub — The model is not self-contained - The forward pass is multi-step — ttt_steps controls the number of autoregressive steps within a single forward call

Conclusion

This chunk of the opencode session is a masterclass in systematic code exploration. The assistant began with a set of assumptions (some correct, some incorrect), tested them through parallel reads, corrected course when assumptions failed, and progressively built a complete mental model of a complex training library. The final synthesis in message 9 transformed scattered source code into actionable knowledge: a complete training recipe that the parent session could execute directly.

The exploration also reveals something important about the speculators library itself. It is a well-structured codebase with clear separation of concerns (model, train, data_generation, convert, proposals, utils), but it lacks user-facing documentation. There are no example scripts, no train CLI command, and no README explaining the training workflow. The library assumes the user will read the source code. This chunk is the story of someone accepting that assumption and acting on it—systematically, thoroughly, and with impressive results.

For anyone facing a similar challenge—understanding an unfamiliar, undocumented codebase—the strategy demonstrated here is worth emulating: start broad, test assumptions, correct course when wrong, drill deeper methodically, and synthesize everything into a working example. The result is not just understanding, but the ability to act on that understanding.## References

[1] "Drilling Down: The Systematic Exploration of EAGLE-3 Training Infrastructure" — Analysis of message 3, where the assistant reads the eagle3 model files, configuration system, and supporting infrastructure.

[2] "Completing the Speculators Exploration" — Analysis of message 8, the final data-gathering round before synthesis.

[3] "Mapping the Territory: The First Reconnaissance of the Speculators Training API" — Analysis of message 1, the initial exploration round.

[4] "The Final Piece of the Puzzle: Reading the Loading Utility and Data Generation Scripts" — Analysis of message 6, covering the data generation pipeline.

[5] "Navigating the Unknown: How an AI Assistant Corrects Course and Begins Mapping a Training Library" — Analysis of message 2, the course-correction round.

[6] "Systematic Code Archaeology: Reading the Remaining Support Files of the Speculators Library" — Analysis of message 5, covering vocab mapping, distributed batch sampling, and logging.

[7] "Decoding the Decoder: A Deep Dive into the EAGLE-3 Training API Exploration" — Analysis of message 0, the user's initial exploration request.

[8] "The Final Lap: Completing a Systematic Library Exploration" — Analysis of message 7, covering the convert and proposals modules.

[9] "The Pivot Point: Diving into the Eagle3 Model Architecture" — Analysis of message 4, the deep dive into Eagle3DraftModel.

[10] "The Architecture of EAGLE-3 Training: A Deep Dive into the Speculators Library" — Analysis of message 9, the complete architectural synthesis.