Eagle3 模型生产

Speculators 目前支持 Eagle3 模型的训练。此功能可通过本目录中的脚本实现。 1. data_generation_offline.py: 使用 vLLM 生成训练数据（验证器隐藏状态）。注意：此脚本也会在数据尚未预处理的情况下进行预处理。 2. build_vocab_mapping.py: 使用 Token 频率分布文件构建 d2t（草稿到目标）和 t2d（目标到草稿）词汇表映射。 3. train.py: 使用训练数据和词汇表映射训练 Eagle3 模型。 4. (可选) gen_and_train.py: 一个方便的包装脚本，将上述脚本集成在一个命令中运行。.

目录

• 数据生成
  
  • 快速入门
    
  • 高级用法
    
  • 故障排除
    
• 词汇表映射
  
  • 快速入门
    
• 训练
  
  • 快速入门
    
  • 参数
    
  • 示例命令
    
• 端到端流水线
  
  • 概述
    
  • 先决条件
    
  • 用法
    

数据生成

scripts/data_generation_offline.py 提供了生成 Eagle3 模型训练数据的主要入口点。数据生成使用 vLLM，需要安装可选的 datagen 模块。

快速入门

使用 Llama 3.1 8B 从 ShareGPT 生成训练数据

python scripts/data_generation_offline.py \
    --target-model-path meta-llama/Llama-3.1-8B-Instruct \
    --train-data-path sharegpt \
    --output-dir ./training_data \
    --max-samples 5000

脚本通过 apply_chat_template 自动使用 tokenizer 内置的聊天模板。它将使用 vllm 为训练数据生成目标模型隐藏状态，并将它们与 input_ids 和 loss_mask 张量一起保存到磁盘，格式为 .pt 文件。

生成的示例数据请参见：https://hugging-face.cn/datasets/nm-testing/sharegpt_llama3_8b_hidden_states

高级用法

使用自定义设置和多 GPU

python scripts/data_generation_offline.py \
    --target-model-path meta-llama/Llama-3.1-70B-Instruct \
    --train-data-path ./my_data.jsonl \
    --seq-length 4096 \
    --cache-dir ./cache \
    --output-dir ./training_data \
    --layer-ids 2 28 54 \
    --tensor-parallel-size 4 \
    --batch-size 16 \
    --max-samples 10000

数据配置文件

脚本将在输出目录中生成一个 data_config.json 文件，其中包含用于生成数据时使用的配置以及有关数据生成过程的其他元数据。

示例文件

{
  "version": "2.0",
  "generated_at": "2025-12-03T16:03:02.471808+00:00",
  "speculators_version": "0.3.0",
  "reproducibility": {
    "command": "data_generation_offline.py --target-model-path meta-llama/Llama-3.1-8B-Instruct --train-data-path sharegpt --output-dir ./training_data --max-samples 5000",
    "package_versions": {
      "torch": "2.8.0+cu128",
      "vllm": "0.11.0",
      "transformers": "4.57.3",
      "speculators": "0.3.0"
    },
    "gpu": "NVIDIA H100 80GB HBM3"
  },
  "model": {
    "target_model_path": "meta-llama/Llama-3.1-8B-Instruct",
    "tensor_parallel_size": 1,
    "max_model_len": 2048,
    "gpu_memory_utilization": 0.8,
    "hidden_size": 4096
  },
  "data": {
    "train_data_path": "sharegpt",
    "seq_length": 2048,
    "max_samples": 5000,
    "num_samples": 5000,
    "seed": 0,
    "chat_template_note": "Uses tokenizer's built-in chat template"
  },
  "hidden_states": {
    "layer_ids": [
      2,
      16,
      29,
      31
    ],
    "description": "Layers selected for EAGLE3 fusion and target logits"
  },
  "generation": {
    "cache_dir": "/home/***/.cache/huggingface/datasets"
  },
  "format": {
    "file_pattern": "data_{idx}.pt",
    "data_format_version": 1,
    "schema": {
      "input_ids": {
        "dtype": "torch.long",
        "shape": "[seq_len]",
        "description": "Tokenized input sequence"
      },
      "hidden_states": {
        "dtype": "list[torch.bfloat16]",
        "shape": "list of [seq_len, 4096]",
        "num_tensors": 4,
        "description": "Hidden states from 4 layers"
      },
      "loss_mask": {
        "dtype": "torch.long",
        "shape": "[seq_len]",
        "description": "1 for assistant tokens to train on, 0 elsewhere"
      }
    }
  }
}

Token 频率文件

除了 data_config.json，数据生成步骤还将生成一个 token_freq.pt 文件，其中包含 Token 频率。如果未指定，Token 频率文件的默认位置是 ./token_freq.pt，即脚本运行的同一目录。这些频率将用于 d2t，即 draft-to-target 和 t2d，即 target-to-draft 词汇表映射。

数据集

内置数据集（可在 --train-data-path 参数中直接按名称使用）： - sharegpt - ShareGPT Vicuna 未过滤 - ultrachat - HuggingFace UltraChat 200k

或者，您可以通过在 --train-data-path 参数中传递 HuggingFace 数据集路径或本地 JSON/JSONL 文件路径来使用其他数据集。

缓存

HuggingFace 数据集使用基于指纹的缓存无效化自动缓存预处理。当以下情况发生时，缓存会自动更新：

• Tokenizer 更改
• 预处理参数更改（seq_length 等）
• 数据集更改

缓存位置

默认：~/.cache/huggingface/datasets (可选) 通过设置 HF_HUB_CACHE 环境变量使用自定义缓存目录。

# Example: Use custom cache directory
export HF_HUB_CACHE=/path/to/your/cache
python scripts/data_generation_offline.py ...

故障排除

1. 隐藏状态提取过程中内存不足

   • 减少 --batch-size
   • 减少 --seq-length
   • 增加 --tensor-parallel-size

2. 层索引越界

   • 检查模型的实际层数
   • 自动选择使用：[2, num_layers // 2, num_layers - 3]

3. 未找到助手回复跨度

   • 确保 tokenizer 具有聊天模板（支持 apply_chat_template）
   • 检查对话是否具有正确格式的助手回复（role/content 键）

4. 缓存无效化

   • 更改预处理参数时删除缓存目录
   • 确保运行之间 --seed 值一致以实现可复现性

词汇表映射

scripts/build_vocab_mapping.py 使用 Token 频率分布文件构建 d2t（草稿到目标）和 t2d（目标到草稿）词汇表映射。

快速入门

使用 Llama 3.1 8B 生成词汇表映射

通过手动指定 target-vocab-size

    python scripts/build_vocab_mapping.py \
        --token-freq-path ./token_freq.pt \
        --draft-vocab-size 32000 \
        --target-vocab-size 128256 \
        --output-path ./vocab_mapping

或使用 target-model-path 自动推断目标词汇表大小

    python scripts/build_vocab_mapping.py \
        --token-freq-path ./token_freq.pt \
        --draft-vocab-size 32000 \
        --target-model-path meta-llama/Llama-3.1-8B-Instruct \
        --output-path ./vocab_mapping

如果未指定，Token 频率文件的默认位置是 ./token_freq.pt。请确保 target-vocab-size 与验证器模型的词汇表大小完全匹配。完成后，此步骤将在磁盘上生成并保存 t2d.npy 和 d2t.npy 文件。

训练

scripts/train.py 提供了训练 Eagle3 模型的主要入口点。

快速入门

要在单节点多 GPU 分布式训练设置中使用 FSDP，应使用 torchrun 启动脚本。

torchrun --standalone --nproc_per_node=<num_gpus>  scripts/train.py

对于单 GPU 训练（用于调试），可以直接运行脚本。

python scripts/train.py

[!NOTE] 使用 CUDA_VISIBLE_DEVICES=<gpu_ids> 控制脚本可见的 GPU。

参数

脚本有一个必需参数：--verifier-name-or-path，即要使用的验证器模型的名称或路径。

脚本具有以下可选参数： - --data-path：数据目录的路径。默认为 ./data。脚本将收集此目录或其子目录中的所有 .pt 文件并将其用作训练数据。 - --save-path：保存检查点的路径。默认为 ./checkpoints。脚本将为每个 epoch 创建子目录以保存模型权重和优化器状态。例如 ./checkpoints/0/ - --epochs：训练的 epoch 数。默认为 20。 - --lr：要使用的学习率。默认为 1e-4。 - --no-resume-from-checkpoint：如果设置，脚本将不会从最后一个检查点恢复（如果存在），而是从头开始并覆盖现有检查点。 - --logger：要使用的 logger。默认为空字符串，表示不记录。支持的 loggers 有 trackio、wandb 和 tensorboard。 - --total-seq-len：要使用的总序列长度。默认为 8192。 - --data-format-version：要使用的数据格式版本。默认为 1。训练数据的结构。1 是默认值，是 Speculators 生成脚本产生的结构。0 用于与旧数据格式向后兼容。 - --log-dir：保存日志的路径。默认为 ./logs。 - --run-name：运行的名称。默认为 None。 - --num-layers：要使用的层数。默认为 1。 - --d2t-path：d2t 张量的路径。默认为 d2t.npy。 - --t2d-path：t2d 张量的路径。默认为 t2d.npy。 - --ttt-steps：要使用的 TTT 步数。默认为 3。 - --ttt-step-loss-decay：TTT 步数要使用的损失衰减因子。默认为 1.0。

示例命令

torchrun --nnodes=1 --nproc_per_node=8 scripts/train.py \
    --verifier-name-or-path "meta-llama/Llama-3.1-8B-Instruct" \
    --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/gen_and_train.py 可用于在一个命令中运行完整的流水线。它还确保每个脚本都使用正确的参数和依赖项运行。

它按顺序调用以下脚本： 1. scripts/data_generation_offline.py 2. scripts/build_vocab_mapping.py 3. scripts/train.py

使用 uv 为每个脚本生成临时环境。

先决条件

• Python 3.10+
• uv (pip install uv)

用法

[!IMPORTANT] 运行前请更新脚本文件本身中的脚本参数部分。

然后运行

python scripts/gen_and_train.py

[!NOTE] 您可以使用环境变量（如 CUDA_VISIBLE_DEVICES 和 HF_HOME）调用脚本来控制脚本的行为。默认情况下，脚本将使用所有可用的 GPU。

CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/gen_and_train.py