Eagle3 Training (#143)

This pr introduces Eagle3 Model training into the speculators repo. The
implementation is specific to Eagle3 but designed in a way that enables
future generalization to other speculative decoding algorithms.

# Components
<img width="1418" height="734" alt="Eagle3 Training Components"
src="https://github.com/user-attachments/assets/418a7d1f-0078-412a-ae56-a6427b756a05"
/>

## Example training script (~`scripts/train_llama3_8b_drafter.py`~
`scripts/train.py`)
Shows how to setup and run training. ~Currently specific to the
`meta-llama/Llama-3.1-8B-Instruct` model but doesn't require many
changes to run with a different model. Just need to update
`
VERIFIER_MODEL_NAME_OR_PATH = "meta-llama/Llama-3.1-8B-Instruct"
HIDDEN_SIZE = 4096  # Must match the verifier model's hidden size
VERIFIER_VOCAB_SIZE = 128256 # Must match the verifier model's vocab
size
`~

**Update:** I've generalize the training script. It now has a required
cli arg `--verifier_name_or_path` and supports arbitrary verifier
models. Note: this uses
`LlamaConfig.from_pretrained(args.verifier_name_or_path)` under the
hood, which does work for non-llama models (e.g. a Qwen model) but
prints a warning and may not work for every type of verifier.

You will also need to pass in a dataset and `t2d` / `d2t` tensors which
correspond to the verifier you are using.

## Flex Attention
Files:
- `src/speculators/train/eagle3/attention.py`
- `tests/unit/train/test_eagle3_attention.py`

The training code uses Flex attention which provides substantial speed
ups and memory efficiency over the full dense attention operations.

Functions:
- create_combined_mask_mod(lengths, total_seq_len): This function
creates the mask function used by flex attention.
- extend_mask_for_draft_tokens(block_mask): Helper function to extend
the block mask without needed to check each new squares mask value
- block_mask_to_dense_attention_mask: Only used for debugging purposes
- flex_attention_forward: lightweight wrapper around flex attention call

## Data processing
<img width="4531" height="2384" alt="Eagle3 Data Flow"
src="https://github.com/user-attachments/assets/b972ef8c-92d4-4d46-969f-66d33f801ceb"
/>
Files:
- `src/speculators/train/data.py`

Data is currently expected in the format of 1 file per data sample. We
load these samples and perform a shift to align `input_ids,
hidden_states, loss_mask, verifier_last_hidden_state` correctly. We also
automatically collate these samples into batches. Rather than padding
and wasting compute on padded tokens, we instead concatenate the
sequences along the sequence dimension, keeping track of the boundaries
between sequences and setting the attention mask accordingly.

## Batch sampling
Files: 
- `src/speculators/train/distributed_batch_sampler.py`
- `src/speculators/train/data.py`

Due to hardware limitations, we set a maximum sequence length for each
batch. We would like each batch of data to be close in size this max
length, so that each batch has a similar number of tokens. The way we
achieve this is through the `MultipackDistributedBatchSamplerV2` taken
from prior work I did on
[instructlab/training](https://github.com/instructlab/training). This
class produces indices of files that when batched together come close to
reaching the max length without exceeding it. It also does this in a
distributed aware manner so that there is no overlap in the data each
rank sees.

To run the packing algorithm, we need to know the lengths of each sample
in the dataset. Unfortunately, this would require opening every file in
the dataset which is expensive, so instead we approximate the lengths
(`_compute_approx_lengths` in `data.py`) using the length of the first
sample and the relative file sizes of samples.

## `Eagle3DraftModel`
Files: 
- `src/speculators/train/eagle3/core.py`

The draft model itself. Sets up and loads verifier components, as well
as the draft layers / weights. Contains the model `forward()` pass
which:
- sets up the block mask for the batch
- computes the target logits using the attached `verifier_lm_head`.
Note: this is computed here for data storage efficiency reasons, as
otherwise we would need to save the full logits: `[seq_len, vocab_size]`
instead of the last layer hidden states: `[seq_len, hidden_size]` to
disk. The verifier `vocab_size` is often > 100k whereas `hidden_size`
might be around 4-8k.
- For each ttt step:
  - Embeds tokens
  - concatenates with hidden_states
  - applies decoder layers
  - computes logits
  - computes loss and step accuracy
  - prepares next step tokens
  - Updates block mask
 
## Layer definitions
Files:
- `src/speculators/train/eagle3/model_definitions.py`

Currently just contains model definitions for llama3 style draft models.
Supports `norm_before_residual=True or False`. Attempted to keep
modifications to the original llama models minimal.

## Distributed training via FSDP
Files:
- `src/speculators/train/utils.py`
- `src/speculators/train/checkpointer.py`
- `src/speculators/train/trainer.py` (`setup_model` fn)

Full support for FSDP training by initializing the training script with
`torchrun --nnodes --nproc_per_node=N` where `N` is the number of gpus.
Tested with `N=2,3,4, 8` and all work. FSDP training also enables
Automatic Mixed Precision (AMP) for improved performance.

`checkpointer.py` contains checkpointing logic for FSDP distributed
model weights (gather all weights on rank 0 before saving).

Note: the way distributed works in general is `N` copies of the script
are started and all run the same code but with some env variables
setting which lets each process know its rank. Then explicit
`dist.barrier()` calls or implicit calls within FSDP forward/backwards
hooks force each process to wait until they all reach the same point in
the code, before continuing. It is important that all ranks reach these
operations as it allows them to perform synchronized operations (such as
gathering, reducing, etc). However, we can also limit certain code to
only one rank (rank 0) so that we only log once, or save to checkpoint
once, using simple `if local_rank == 0` statements.

## Logging
Files:
- `src/speculators/train/logger.py`
- `scripts/train.py`: (setup logger calls at start of `main()`)
- `src/speculators/train/trainer.py` and other files: usage of
`metric_logger` and `root_logger`

Another implementation mostly copied from prior work I did on
[instructlab/training](https://github.com/instructlab/training). This
uses python's std library `logging` module and extends it to support
training metric logging. We can log a nested dict of metrics anywhere in
the codebase like so:
```python
# Setup once
import logging
metric_logger = logging.getLogger("speculators.metrics")

# Log call
metric_logger.info(
    {"train": {"loss": loss.item(), **acc_values}, "epoch": epoch},
    extra={"step": self.global_step},
)
```
And when the user runs the training script they can select one (or
multiple) of `tensorboard`, `wandb`, and `trackio` and the results will
be logged to the respective experiment tracker.

There is also a `root_logger` which can be used for regular update
logging and everything logged to either the `root_logger` or
`metric_logger` will be pretty-printed to console.

## `Trainer`
Files:
- `src/speculators/train/trainer.py`

The `Trainer` class is initialized with the model, data loaders, and a
config and:
- Sets up model / optimizer (loads weights and configures distributed if
needed)
- Contains the training and validation loops (`train_epoch` and
`val_epoch` respectively)
- And the overall training loop which alternatives between training,
validation, and saving checkpoints

Todos:
- [x] Eagle3Draft Model definition with TTT steps and loss calculations
- [x] Patched Decoder layer definitions
- [x] Simple data loading from sample files
- [x] FlexAttention masking and implementation
- [x] Loss Masking
- [x] Training loop
  - [x] Train data loader
  - [x] `loss.backward()` + optimizer steps
  - [x] Distributed loss reduction 
  - [x] Val data loader
  - [x] Metric collection/reporting
  - [x] Model checkpointing
- [x] Data batching
  - [x] Collate fn
  - [x] Batch sampler (dynamic batch size through sample packing)
  - [x] Distributed (rank) aware sampling
- [x] Distributed support
- [ ] ~Code relocation / merging with existing definitions (Currently
just have everything under `speculators/train` but this will need to
change)~ FUTURE PR
- [x] Verify correctness of key components (attention masking, data
token alignment, etc).
- [x] General testing

Essential todos (as of 10/22/2025):
- [x] Save checkpoints to safetensors format w/ required config info
- [ ] ~Implement save best or save last logic (currently saving every
epoch)~ FUTURE PR
- [x] Better Verifier `lm_head`, `embed_tokens` loading (requires #144)
- [x] `Eagle3DraftModel.__init__` signature cleanup/better configuration
- [ ] ~Config/argparsing for `scripts/train.py`~ FUTURE PR
- [x] Ensure flex attention impl works with `torch==2.9` and
`torch.compile`
- [x] Fix lint / quality / type errors and pass CI

---------

Signed-off-by: Fynn Schmitt-Ulms <[email protected]>
Co-authored-by: Brian Dellabetta <[email protected]>

Author: fynnsu

pyproject.toml

@@ -48,8 +48,10 @@ dependencies = [
     "pydantic>=2.0.0",
     "pydantic-settings>=2.0.0",
     "pyyaml>=6.0.0",
+    "rich",
     "safetensors",
     "torch",
+    "tqdm",
     "transformers",
     "typer-slim>=0.12.0",
 ]
@@ -102,6 +104,7 @@ dev = [
     "types-PyYAML~=6.0.1",
     "types-requests~=2.32.0",
     "types-toml",
+    "types-tqdm",
 
     # link checking
     "mkdocs-linkcheck~=1.0.6",
@@ -211,9 +214,6 @@ select = [
     "UP", # pyupgrade: automatically upgrades syntax for newer versions of Python
     "W", # Warning: provides warnings about potential issues in the code
     "YTT", # flake8-2020: identifies code that will break with future Python releases
-
-    # Code Documentation
-    "FIX", # flake8-fixme: detects FIXMEs and other temporary comments that should be resolved
 ]
 
 [tool.ruff.lint.extend-per-file-ignores]
@@ -251,6 +251,7 @@ select = [
     "PTH", # os.path is acceptable in data generation
 ]
 "scripts/**/*.py" = [
+    "INP001", # allow implicit namespace packages in scripts
     "PTH", # os.path is acceptable in scripts
 ]
 

scripts/TRAINING.md

@@ -0,0 +1,56 @@
+# Eagle3 Training
+
+`scripts/train.py` provides the main entry point for training Eagle3 models.
+
+## Running the training script
+
+To run in a multi-node distributed training setup with FSDP, the scripts should be launched with `torchrun`:
+```bash
+torchrun --nnodes=1 --nproc_per_node=<num_gpus>  scripts/train.py
+```
+
+For single GPU training (useful for debugging), the script can be run directly:
+```bash
+python scripts/train.py
+```
+
+## Arguments
+The scripts has one required argument: `--verifier-name-or-path`, which is the name or path of the verifier model to use.
+
+The scripts has the following optional arguments:
+- `--data-path`: The path to the data directory. Defaults to `./data`. The script will collect all `.pt` files in this directory or its subdirectories and use them as training data.
+- `--save-path`: The path to save the checkpoints. Defaults to `./checkpoints`. The script will create subdirectories for each epoch to save the model weights and optimizer states. e.g. `./checkpoints/0/`
+- `--epochs`: The number of epochs to train for. Defaults to 20.
+- `--lr`: The learning rate to use. Defaults to 1e-4.
+- `--no-resume-from-checkpoint`: If set, the script will not resume from the last checkpoint if it exists, and will instead start from scratch and overwrite existing checkpoints.
+- `--logger`: The logger to use. Defaults to empty string, which means no logging. Supported loggers are `trackio`, `wandb`, and `tensorboard`.
+- `--total-seq-len`: The total sequence length to use. Defaults to 8192.
+- `--data-format-version`: The version of the data format to use. Defaults to 1. The structure of the data to train on. `1` is the default and is the structure produced by Speculators generation scripts. `0` exists for backwards compatibility with the old data format.
+- `--log-dir`: The path to save the logs. Defaults to `./logs`.
+- `--run-name`: The name of the run. Defaults to None.
+- `--num-layers`: The number of layers to use. Defaults to 1.
+- `--d2t-path`: The path to the d2t tensor. Defaults to `d2t.npy`.
+- `--t2d-path`: The path to the t2d tensor. Defaults to `t2d.npy`.
+- `--ttt-steps`: The number of TTT steps to use. Defaults to 3.
+- `--ttt-step-loss-decay`: The loss decay factor to use for the TTT steps. Defaults to 1.0.
+
+## Example run command
+```bash
+torchrun --nnodes=1 --nproc_per_node=8 scripts/train.py \
+    --verifier-name-or-path "meta-llama/Llama-3.1-8B" \
+    --data-path "./data/llama-3.1-8b_sharegpt/gen/" \
+    --save-path "./checkpoints/llama-3.1-8b.eagle3" \
+    --epochs 10 \
+    --lr 1e-4 \
+    --no-resume-from-checkpoint \
+    --logger "tensorboard" \
+    --total-seq-len 8192 \
+    --data-format-version 1 \
+    --log-dir "./logs/llama-3.1-8b.eagle3" \
+    --run-name "llama-3.1-8b.eagle3" \
+    --num-layers 1 \
+    --d2t-path "./data/llama-3.1-8b_sharegpt/d2t.npy" \
+    --t2d-path "./data/llama-3.1-8b_sharegpt/t2d.npy" \
+    --ttt-steps 3 \
+    --ttt-step-loss-decay 1.0
+```

scripts/train.py

@@ -0,0 +1,213 @@
+import argparse
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+from transformers import LlamaConfig
+
+from speculators.config import SpeculatorsConfig, VerifierConfig
+from speculators.models.eagle3 import Eagle3SpeculatorConfig
+from speculators.proposals.greedy import GreedyTokenProposalConfig
+from speculators.train.data import (
+    Eagle3SampleFileDataset,
+    create_collate_fn,
+    split_files,
+    standardize_data_v0,
+    standardize_data_v1,
+)
+from speculators.train.distributed_batch_sampler import (
+    MultipackDistributedBatchSamplerV2,
+)
+from speculators.train.eagle3.core import Eagle3DraftModel
+from speculators.train.logger import setup_metric_logger, setup_root_logger
+from speculators.train.noise_transforms import AddUniformNoise
+from speculators.train.trainer import Trainer, TrainerConfig
+from speculators.train.utils import maybe_destroy_distributed, maybe_setup_distributed
+
+# DRAFTER MODEL HYPARAMETERS
+NORM_BEFORE_RESIDUAL = True
+
+# Dataloader
+NUM_WORKERS = 12
+PREFETCH_FACTOR = 4
+NOISE_STD = 0.05
+
+
+def setup_dataloader(
+    file_list: list[str],
+    world_size: int,
+    local_rank: int,
+    add_noise: bool = True,
+    data_format_version: int = 1,
+) -> DataLoader:
+    """Setup dataloader for training.
+    Args:
+        file_list: List of file paths to load data from.
+        world_size: Number of processes in the distributed training.
+        local_rank: Rank of the current process.
+        add_noise: Whether to add noise to the data.
+        data_format_version: Version of the data format. Default is 1.
+    Returns:
+        DataLoader: Dataloader for training.
+    """
+    if add_noise:
+        noise_transform = AddUniformNoise(
+            std=NOISE_STD, tensors=("hidden_states", "verifier_last_hidden_states")
+        )
+    else:
+        noise_transform = None
+
+    standardize_fn = (
+        standardize_data_v1 if data_format_version == 1 else standardize_data_v0
+    )
+
+    dataset = Eagle3SampleFileDataset(
+        file_list=file_list,
+        max_len=args.total_seq_len,
+        transform=noise_transform,
+        standardize_fn=standardize_fn,
+    )
+    batch_sampler = MultipackDistributedBatchSamplerV2(
+        batch_max_length=args.total_seq_len,
+        lengths=dataset.approx_lengths,
+        num_replicas=world_size,
+        rank=local_rank,
+    )
+    return DataLoader(
+        dataset,
+        batch_sampler=batch_sampler,
+        num_workers=NUM_WORKERS,
+        prefetch_factor=PREFETCH_FACTOR,
+        pin_memory=True,
+        collate_fn=create_collate_fn(args.total_seq_len),
+        persistent_workers=True,
+    )
+
+
+def main(args: argparse.Namespace):
+    # Setup logging
+    setup_root_logger()
+    setup_metric_logger(
+        loggers=args.logger, run_name=args.run_name, output_dir=args.log_dir
+    )
+
+    # Setup distributed training
+    local_rank, world_size, rank, is_distributed = maybe_setup_distributed()
+    device = torch.device(local_rank)
+
+    # Setup speculator config
+    llama_config = LlamaConfig.from_pretrained(args.verifier_name_or_path)
+    llama_config.num_hidden_layers = args.num_layers
+    llama_config.model_type = "llama"  # reset to llama (handles non-llama verifiers)
+    llama_config._attn_implementation = "simple_flex_attention"  # noqa: SLF001
+
+    # Load t2d and d2t tensors
+    d2t = torch.from_numpy(np.load(args.d2t_path)).to(device)
+    t2d = torch.from_numpy(np.load(args.t2d_path)).to(device)
+    draft_vocab_size = d2t.shape[0]
+
+    speculator_config = Eagle3SpeculatorConfig(
+        transformer_layer_config=llama_config,
+        draft_vocab_size=draft_vocab_size,
+        norm_before_residual=NORM_BEFORE_RESIDUAL,
+        speculators_config=SpeculatorsConfig(
+            algorithm="eagle3",
+            proposal_methods=[
+                GreedyTokenProposalConfig(
+                    proposal_type="greedy",
+                    speculative_tokens=args.ttt_steps,
+                )
+            ],
+            default_proposal_method="greedy",
+            verifier=VerifierConfig(
+                name_or_path=args.verifier_name_or_path,
+                architectures=["LlamaForCausalLM"],
+            ),
+        ),
+    )
+
+    # Setup draft model
+    draft_model = Eagle3DraftModel(config=speculator_config, t2d=t2d, d2t=d2t)
+
+    # Setup dataloaders
+    train_files, val_files = split_files(args.data_path, ratio=0.9)
+    train_loader = setup_dataloader(
+        train_files,
+        world_size,
+        local_rank,
+        add_noise=True,
+        data_format_version=args.data_format_version,
+    )
+    val_loader = setup_dataloader(
+        val_files,
+        world_size,
+        local_rank,
+        add_noise=False,
+        data_format_version=args.data_format_version,
+    )
+
+    # Setup trainer
+    trainer_config = TrainerConfig(
+        num_epochs=args.epochs,
+        save_path=args.save_path,
+        lr=args.lr,
+        resume_from_checkpoint=not args.no_resume_from_checkpoint,
+        is_distributed=is_distributed,
+        local_rank=local_rank,
+        train_call_kwargs={
+            "use_off_policy_tokens": False,
+            "ttt_steps": args.ttt_steps,
+            "ttt_step_loss_decay": args.ttt_step_loss_decay,
+        },
+        val_call_kwargs={
+            "use_off_policy_tokens": False,
+            "ttt_steps": args.ttt_steps,
+            "ttt_step_loss_decay": args.ttt_step_loss_decay,
+        },
+    )
+    trainer = Trainer(draft_model, trainer_config, train_loader, val_loader)
+
+    # Run training
+    trainer.run_training()
+
+    # Cleanup
+    maybe_destroy_distributed()
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--verifier-name-or-path", type=str, required=True)
+    parser.add_argument("--data-path", type=str, default="./data")
+    parser.add_argument("--save-path", type=str, default="./checkpoints")
+    parser.add_argument("--epochs", type=int, default=20)
+    parser.add_argument("--lr", type=float, default=1e-4)
+    parser.add_argument("--no-resume-from-checkpoint", action="store_true")
+    parser.add_argument(
+        "--logger",
+        type=str,
+        default="",
+        help="One of 'trackio', 'wandb', 'tensorboard' or comma separated list of them",
+    )
+    parser.add_argument("--total-seq-len", type=int, default=8192)
+    parser.add_argument("--data-format-version", type=int, default=1)
+    parser.add_argument("--log-dir", type=str, default="./logs")
+    parser.add_argument("--run-name", type=str, default=None)
+    parser.add_argument("--num-layers", type=int, default=1)
+    parser.add_argument("--d2t-path", type=str, default="d2t.npy")
+    parser.add_argument("--t2d-path", type=str, default="t2d.npy")
+    parser.add_argument("--ttt-steps", type=int, default=3)
+    parser.add_argument("--ttt-step-loss-decay", type=float, default=1.0)
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)
+
+
+# RUN WITH:
+# torchrun --nnodes=1 --nproc_per_node=<num_gpus>  scripts/train.py
+# for FSDP training
+# OR
+# python scripts/train.py
+# for single GPU training