基于1层Llama草稿模型实现2.16推测解码加速
加速比
- 在测试样本 3（Dolly Test Sample 3）中，系统完美打通并输出加速记录：
Sample | Match | Naive TPS | Spec TPS | Speedup
-----------------------------------------------------------------
3 | True | 3.95 | 8.53 | 2.16 x | 2.00 | 20.0 %
- Match = True 保证了文本级别 100% 逐字完全对齐。
- Naive TPS (自回归): 3.95
- Spec TPS (投机解码): 8.53
- 物理加速比 (Speedup): 2.16x 🚀
- 单步平均接受 Token 数: 2.00 个（草稿额外贡献了 1.0 个有效 Token 生成）。
全链路数量级压缩
1. 模型架构层：1/20 层数
|
标准 SmolLM2-360M |
本仓库 draft (Eaglet1) |
| Transformer layers |
20 |
1 |
| hidden_size |
960 |
960（不变） |
| vocab_size |
49152 |
49152（不变） |
embed_tokens 是最大权重块（49152×960 ≈ 47M），直接复用 teacher 权重，冻结不训练。
2. 隐状态压缩瓶颈：960 → 256
hidden_states ──► down_proj(960→256) ──► up_proj(256→960) ──► output
bottleneck_size=256 将每层隐状态压缩 3.75x，draft 只需要学习一个低维"信号摘要"。
3. LoRA 只训练极少数参数
| 超参 |
值 |
含义 |
lora_r |
8 |
秩，极低 |
lora_alpha |
16 |
alpha=2×r |
target_modules |
7 个 |
q/k/v/o_proj + gate/up/down_proj |
bias |
none |
不训 bias |
lora_dropout |
0.0 |
无 dropout |
| 总可训练参数 |
~几 MB |
只训 LoRA A/B 矩阵 |
4. 训练阶段数据/步骤压缩
| 参数 |
值 |
标准训练对比 |
num_train_steps |
600 |
完整预训练/微调通常数万步 |
num_train_samples |
1000 |
常规微调用数万条 |
batch_size |
1 |
常规 batch=32/64 |
max_seq_len |
256 |
常规 2048/4096 |
| 预估耗时 |
~4.3 min (GPU) |
数小时~数天 |
- 总 token 量：600步 × 1样本 × 256 tokens = 153,600 tokens（少于一个 epoch 的标准数据集）
5. 推理阶段 Speculative Decoding 树压缩
| 参数 |
值 |
作用 |
total_token |
32 |
整棵草稿树最多 32 个节点 |
depth |
5 |
树深度 5 层 |
top_k |
8 |
每步只展开 8 个候选 |
| 草稿调用次数 |
depth+1 = 6 次 |
标准自回归需 64 次 |
- 整棵树的候选数理论最大值：
1 + 8 + 8² + 8³ + 8⁴ + 8⁵ = 37273，但用 total_token=32 截断为最多 32 叶子节点，树的总 token 预算压缩了 >1000x。
6. 词表裁剪的物理降维
- 基于垂直数据集（如 Dolly）的实际词频分布，提取累计概率 $P \ge 0.95$ 的核心词集（在极速 Smoke 验证中被安全裁剪至 83 个核心高频词，或工业级泛化裁剪至约 20k），使得最终分类矩阵点积 Overhead 骤降 60% 以上。
- Teacher 完整前向 20层 × 256 tokens → baseline
- Draft 前向 1层 × (256 + 32) tokens → ≈1/20 层 × 1/8 长度
- LoRA 可训练参数 r=8, ~几 MB → 全参数微调的 <0.01%
- Spec tree 节点数 32 (top_k=8, depth=5) → 单 token 自回归的 ~1/2
- 训练总 token 153,600 → 标准微调的 ~1/100
技术栈
- Eagle推荐解码、LoRA 极致微调、BF16 混合精度、4-bit 量化目标模型（Target）以及词表裁剪（Vocab Trimming）
步骤
# source py3.10/bin/activate
PYTHONPATH=./ python evaluation/vocab_trim.py
python train/train_gpu_smoke.py --config-name gpu_smollm_smoke_config
python verify_effectiveness.py
code
运行
python train/train_gpu_smoke.py --config-name gpu_smollm_smoke_config
[smoke] Loading base model on cuda (4-bit NF4)...
Loading weights: 0%| | 0/290 [00:00<?, ?it/s]/opt/py10/lib/python3.10/site-packages/bitsandbytes/backends/cuda/ops.py:213: FutureWarning: _check_is_size will be removed in a future PyTorch release along with guard_size_oblivious. Use _check(i >= 0) instead.
torch._check_is_size(blocksize)
Loading weights: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 290/290 [00:00<00:00, 535.86it/s]
[smoke] base loaded in 2.8s
trainable params: 155,136 || all params: 59,508,096 || trainable%: 0.2607
[smoke] starting training: max_steps=600
[smoke] step= 0 loss=4.272652 lr=8.33e-07 (2.4s elapsed)
[smoke] step= 50 loss=4.505703 lr=4.25e-05 (14.4s elapsed)
[smoke] step= 100 loss=3.220418 lr=4.62e-05 (26.6s elapsed)
[smoke] step= 150 loss=3.591722 lr=4.16e-05 (38.7s elapsed)
[smoke] step= 200 loss=3.254315 lr=3.69e-05 (50.9s elapsed)
[smoke] step= 250 loss=2.473542 lr=3.23e-05 (63.0s elapsed)
[smoke] step= 300 loss=2.328086 lr=2.77e-05 (75.2s elapsed)
[smoke] step= 350 loss=2.481422 lr=2.31e-05 (87.4s elapsed)
[smoke] step= 400 loss=1.519966 lr=1.84e-05 (99.6s elapsed)
[smoke] step= 450 loss=2.222769 lr=1.38e-05 (111.8s elapsed)
[smoke] step= 500 loss=2.148213 lr=9.17e-06 (124.0s elapsed)
[smoke] step= 550 loss=2.185774 lr=4.54e-06 (136.3s elapsed)
[smoke] step= 599 loss=1.094599 lr=0.00e+00 (148.2s elapsed)
[smoke] training done in 148.2s (600 steps)
[smoke] loss summary: first10%=0.000000 last10%=0.000000 delta=+0.000000
[smoke] DONE