Eagle3 Training

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]
@@ -236,6 +236,9 @@ select = [
 "src/speculators/convert/**/*.py" = [
     "BLE001", # allow catching Exception for conversion errors
 ]
+"scripts/**/*.py" = [
+    "INP001",
+]
 
 "src/speculators/data_generation/**/*.py" = [
     "S106", # false positives for chat template tokens

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