lightseekorg/TorchSpec

# TorchSpec TorchSpec 是一个原生的 PyTorch speculative decoding 训练框架。我们引入了一种解耦的 speculative decoding draft model 训练方式，其中推理和训练完全分离，并通过 [Mooncake](https://github.com/kvcache-ai/Mooncake) store 将 hidden states 直接从推理引擎组流式传输到分布式训练 worker，允许双方独立扩缩容。 ## 采用情况 TorchSpec 已被生产推理平台和 vLLM 生态系统采用： - [DigitalOcean](https://www.digitalocean.com/blog/how-we-built-fastest-deepseek-minimax-qwen-on-blackwell-ultra) 使用 TorchSpec 在 DigitalOcean Serverless Inference 上为 MiniMax-M2.5 训练了 EAGLE3 draft model。 - [vLLM](https://vllm.ai/blog/2026-05-11-vllm-tops-artificial-analysis) 使用 TorchSpec 和 vLLM 训练了其 Artificial Analysis 排行榜工作中主推的自定义 EAGLE3 draft model。 - [CoreWeave](https://www.coreweave.com/blog/kimi-k2-7-code-now-available-on-serverless-inference-with-leading-benchmark-price-performance) 使用 TorchSpec 为 Kimi K2.7 Code 训练了 DFlash speculative decoding model，并将 D-PACE 支持向上游贡献。 - [fal](https://blog.fal.ai/how-we-achieved-1000-tok-s-and-16x-throughput-with-dspark-for-ideogram-v4-prompt-expander/) 使用 TorchSpec 为其 Ideogram V4 prompt expander 训练了 DSpark speculative decoding model，并报告 DSpark 带来了 16 倍的吞吐量提升。 - [腾讯混元](https://x.com/TencentHunyuan/status/2082447023626944936) 开源了 [**AngelSpec**](https://arxiv.org/abs/2607.25852)，采用 TorchSpec 进行训练，并在论文中重点介绍。 ## 🤗 发布的模型 使用 TorchSpec 训练的 draft model，可在 [LightSeek Foundation](https://huggingface.co/lightseekorg) Hugging Face 组织下获取： - [lightseekorg/kimi-k2.5-eagle3](https://huggingface.co/lightseekorg/kimi-k2.5-eagle3) - [lightseekorg/kimi-k2.5-eagle3-mla](https://huggingface.co/lightseekorg/kimi-k2.5-eagle3-mla) - [lightseekorg/kimi-k2.6-eagle3](https://huggingface.co/lightseekorg/kimi-k2.6-eagle3) - [lightseekorg/kimi-k2.6-eagle3-mla](https://huggingface.co/lightseekorg/kimi-k2.6-eagle3-mla) - [lightseekorg/kimi-k2.6-eagle3.1-mla](https://huggingface.co/lightseekorg/kimi-k2.6-eagle3.1-mla) 使用 TorchSpec 训练的 draft model，可从其他组织获取： - [Inferact/MiniMax-M3-EAGLE3](https://huggingface.co/Inferact/MiniMax-M3-EAGLE3) ## 🚀 博客 ## 博客与公告 - [EAGLE 3.1：通过 EAGLE 团队、vLLM 和 TorchSpec 的合作推进 Speculative Decoding](https://lightseek.org/blog/eagle-3-1.html)（2026 年 5 月） - [TorchSpec：大规模 Speculative Decoding 训练](https://pytorch.org/blog/torchspec-speculative-decoding-training-at-scale/)（PyTorch 博客，2026 年 3 月） - [TorchSpec：大规模 Speculative Decoding 训练](https://lightseek.org/blog/torchspec-speculative-decoding-training-at-scale.html)（2026 年 3 月） ## 目录 - [架构概述](#architecture-overview) - [推理后端支持](#inference-backend-support) - [快速开始](#quick-start) - [设置](#setup) - [示例](#examples) - [训练模式](#training-modes) - [Checkpoint 转换](#checkpoint-conversion) - [指标报告](#metrics-reporting) - [故障排除](#troubleshooting) ## 架构概述

TorchSpec 围绕解耦的训练 pipeline 构建： - **推理引擎** 使用推理引擎生成 target-model hidden states。 - **Mooncake store** 在推理和训练之间传输 tensor，无需在磁盘上具体化。 - **训练 worker** 使用流式传输的 hidden states 来训练 speculative decoding draft model。 这种分离使训练端专注于优化，同时让推理端可以为了提升 hidden-state 生成吞吐量而独立扩缩容。 ## 推理后端支持 TorchSpec 将 hidden states 从推理引擎流式传输到训练 worker。 | 后端 | 支持级别 | 状态 | |---------|--------------|--------| | [vLLM](https://github.com/vllm-project/vllm) | 一级支持 | 可用 | | [TokenSpeed](https://github.com/lightseekorg/tokenspeed) | 一级支持 | 进行中 | | [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM) | 一级支持 | 可用 | | [SGLang](https://github.com/sgl-project/sglang) | 尽力而为的社区支持 | 可用 | | [HuggingFace Transformers](https://github.com/huggingface/transformers) | 尽力而为的社区支持 | 可用 | ## 快速开始 在具有 4 个 GPU（2 个用于训练，2 个用于推理）的单节点上，为 Qwen3-8B 训练 Eagle3 draft model： ``` ./examples/qwen3-8b-single-node/run.sh ``` 直接从 CLI 覆盖配置值： ``` ./examples/qwen3-8b-single-node/run.sh training.learning_rate=5e-5 training.num_train_steps=500 ``` ## 设置 ### 快速设置 ``` # 使用 vLLM 安装 ./tools/build_conda.sh 1 vllm micromamba activate torchspec # 或使用 SGLang 安装 ./tools/build_conda.sh micromamba activate torchspec # 或从可编辑的源码 checkout 安装 TokenSpeed ./tools/build_conda.sh 1 tokenspeed micromamba activate torchspec ``` 若要安装到您当前的环境中，请执行： ``` ./tools/build_conda.sh current tokenspeed # or 'sglang', 'vllm', or 'both' ``` TokenSpeed 后端目前需要 Python 3.12 环境，因为其原生 kernel 依赖 wheel 不支持 Python 3.14。 可选：安装 Flash Attention 支持： ``` pip install -e ".[fa]" ``` ### 特定后端使用方法 **vLLM** ``` ./examples/qwen3-8b-single-node/run.sh configs/vllm_qwen3_8b.yaml ``` **SGLang** ``` ./examples/qwen3-8b-single-node/run.sh ``` **TensorRT-LLM** 在 TensorRT-LLM 镜像 (`docker/trtllm/v1.3.0rc18/Dockerfile`) 中运行，该镜像自带了针对 Mooncake hidden-state 捕获预补丁的 `tensorrt_llm`： ``` ./examples/qwen3-8b-single-node/run.sh configs/trtllm_qwen3_8b.yaml ``` 目前仅支持单节点 tensor parallelism（多节点 TP 尚未接入）。 TorchSpec 使用 vLLM 的 **Worker Extension** 机制挂钩到模型前向传播中，并直接在 worker 进程内捕获 hidden states，从而避免了提取过程中的 RPC 序列化开销。对于 SGLang，TorchSpec 对现有代码库应用补丁以启用 hidden-state 提取。对于 TensorRT-LLM，TorchSpec 在其原生的 **SaveHiddenStates** speculative mode 基础上构建，并应用了一个小补丁，将捕获的 aux + final hidden states 重定向到 Mooncake，而不是将它们写入磁盘。 ### 更新 SGLang 补丁 在 `_sglang` 检出中提交您的更改，然后使用 `./tools/update_sglang_patch.sh` 从锁定的基准提交重新生成 `patches/sglang//sglang.patch`。 ## 示例 | 示例 | 后端 | 模型 | |---------|---------|-------| | [hf-quickstart](examples/hf-quickstart/) | HuggingFace | Qwen3-8B | | [qwen3-8b-single-node](examples/qwen3-8b-single-node/) | 推理引擎 | Qwen3-8B | | [kimi-k25-2node-h200](examples/kimi-k25-2node-h200/) | 推理引擎 | Kimi-K2.5 | | [kimi-k25-3node-h100](examples/kimi-k25-3node-h100/) | 推理引擎 | Kimi-K2.5 | | [minimax-m25-5node-h200](examples/minimax-m25-5node-h200/) | 推理引擎 | MiniMax-M2.5 | 有关每个示例的更多详细信息，请参阅 [examples/README.md](examples/README.md)。 ## 训练模式 ### 离线重放训练 离线重放（[文档](docs/offline_training.md)）是一种通过从磁盘读取 hidden states 来将 target-output 生成分离出 draft-model 训练的训练模式。推荐在 1 个 GPU 上进行测试和开发。 ### 恢复与连续训练 两种模式都使用 `training.load_path`，但它们恢复的状态不同： | 目标 | `training.load_path` | `training.continual_training` | 恢复内容 | |------|----------------------|-------------------------------|--------------------| | 恢复中断的运行 | 必填 | `false`（默认） | 模型、optimizer、LR scheduler、RNG 和 step 元数据 | | 从现有权重开始新的运行 | 必填 | `true` | 仅模型权重 | 恢复相同的运行： ``` training: load_path: /path/to/old_run/checkpoints output_dir: /path/to/old_run ``` 从现有权重开始新的运行： ``` training: load_path: /path/to/old_run/checkpoints continual_training: true learning_rate: 1e-5 warmup_ratio: 0.01 num_epochs: 1 output_dir: /path/to/new_run ``` 从现有的 Hugging Face checkpoint 开始新的运行。`--input` 接受 Hugging Face Hub repo id（自动下载）、本地 HF/safetensors 目录或 `.safetensors` 文件，或 TorchSpec DCP checkpoint 目录： ``` python tools/convert_to_torchspec.py \ --input org/dflash-checkpoint \ --config torchspec/config/dspark_draft_config_qwen36_35b.json \ --output ./outputs/dspark_init ``` 然后从生成的初始化进行热启动： ``` training: load_path: ./outputs/dspark_init continual_training: true ``` 使用此技术，您还可以从预训练的 DFlash model 热启动 DSpark 训练运行，以加快收敛速度。 ## Checkpoint 转换 将 FSDP checkpoint 转换为 HuggingFace 格式： ``` python tools/convert_to_hf.py --input-dir ./outputs/my_experiment/iter_0010000/ ``` 词表裁剪（将 draft model `lm_head` 缩减到更小的 token 集，并输出 `d2t` 和 `t2d` 映射）可以在训练期间或在转换时应用。 - **预裁剪**：在您的训练 config 中设置 `draft_vocab_size`。Checkpoint 已包含裁剪后的 `lm_head` 和 `d2t`/`t2d` buffers，因此基本的转换命令就足够了。 - **后裁剪**：使用完整词表进行训练，然后在转换时连同用于计算 token 频率的代表性数据集一起传递 `--prune-vocab`。 ``` python tools/convert_to_hf.py \ --input-dir ./outputs/my_experiment/iter_0010000/ \ --prune-vocab \ --dataset-path Aeala/ShareGPT_Vicuna_unfiltered \ --draft-vocab-size 32000 \ --tokenizer Qwen/Qwen3-8B \ --chat-template qwen \ --prompt-key conversations ``` 传递 `--cache-dir ./cache` 以重复使用训练中分词后的数据集缓存。 ## 指标报告 默认情况下，W&B 日志记录通过 `report_to: none` 禁用。要启用它，请在您的 config 中设置 `report_to: wandb` 并提供您的 API key。 ## 故障排除 在诊断问题时设置 `TORCHSPEC_LOG_LEVEL=DEBUG` 以获取更详细的日志记录： ``` TORCHSPEC_LOG_LEVEL=DEBUG ./examples/qwen3-8b-single-node/run.sh ``` ### Mooncake SEGFAULT 当前的 Mooncake 版本存在一个仅限 TCP 主机的 bug，会导致 SEGFAULT 错误。在 [upstream issue](https://github.com/kvcache-ai/Mooncake/issues/1986) 修复之前，请设置 `MC_STORE_MEMCPY=0`。 ### RDMA 故障 如果您收到错误 `... libcudart symbols not found globally. Make sure PyTorch with CUDA is installed before using TileLang`，请重新安装 mooncake 以匹配您的系统 CUDA 版本，即如果您在 CUDA 13 上运行，mooncake 可能安装的是 cu12 wheel，请使用 `uv pip uninstall mooncake-transfer-engine && uv pip install mooncake-transfer-engine-cuda13` 重新安装。 ### 按级别文件记录日志 将 `TORCHSPEC_LOG_DIR` 设置为共享文件系统 (NFS) 上的绝对路径，以便为训练和推理双方的每个 Ray actor 启用按级别的日志文件： ``` export TORCHSPEC_LOG_DIR=/my_project/running_logs ``` 这将创建一个结构化的目录，每个 actor 对应一个文件，按角色和节点进行组织： ``` running_logs/ training/ 10.0.0.1/ training_g0_rank0_20260301_080012.log training_g0_rank1_20260301_080012.log 10.0.0.2/ training_g0_rank2_20260301_080013.log inference/ 10.0.0.1/ inference_g0_rank0_20260301_080014.log 10.0.0.2/ inference_g0_rank1_20260301_080015.log ``` 该路径必须是绝对路径，并且可从所有节点写入。如果 `TORCHSPEC_LOG_DIR` 未设置或不可写，则按级别的文件记录将保持禁用状态，Ray 会回退到 stdout/stderr 捕获。 | 问题 | 参考 | |-------|-----------| | 分布式运行卡住或失败，Ray actor 错误 | [docs/debugging_ray_jobs.md](docs/debugging_ray_jobs.md) | | Ray cluster 设置，actor 层次结构，placement groups | [docs/ray.md](docs/ray.md) | | Pipeline 瓶颈，步骤缓慢，吞吐量分析 | [docs/performance_metrics.md](docs/performance_metrics.md) |

标签：PyTorch, vLLM, 人工智能, 凭据扫描, 分布式训练, 大模型训练, 推测解码, 用户模式Hook绕过, 逆向工具