Model-Optimizer 投机解码实战指南：用 Medusa 与 EAGLE 将 LLM 推理延迟降至最低

人工智能大模型模型优化模型量化模型压缩【免费下载链接】Model-OptimizerA unified library of SOTA model optimization techniques like quantization, distillation, pruning, neural architecture search, speculative decoding, etc. It compresses deep learning models for downstream deployment frameworks like TensorRT-LLM, TensorRT, vLLM, etc. to optimize inference speed.项目地址https://gitcode.com/GitHub_Trending/te/Model-Optimizer点击查看免费下载投机解码Speculative Decoding是目前在大语言模型LLM推理延迟敏感场景下最实用的加速手段之一。本文以 Model-Optimizer 的官方指南为主体结合仓库中modelopt.torch.speculative模块的源码、训练示例与单元测试完整讲解如何使用 Model-Optimizer 将一个普通的 Hugging Face 因果语言模型转换为支持 Medusa 或 EAGLE 的投机解码模型并完成微调、保存、恢复与部署。读完本文你将掌握mtsp.convert的完整用法、Medusa/EAGLE 的配置项语义、HF Trainer 微调流程以及投机解码背后的算法原理与权衡。Model-Optimizer 投机解码模块总览Model-Optimizer 的投机解码模块modelopt/torch/speculative通过注册表机制在原有模型结构上动态注入解码头或草稿模块让模型在每一个生成步骤中一次性产出多个 token从而显著降低自回归生成带来的串行开销、提升推理速度。目前该模块支持以下投机解码算法Medusa在目标模型上直接添加多个并行的解码头Medusa heads由模型自身同时预测多个未来 tokenEAGLE引入一个轻量级自回归解码器在隐藏状态feature层面预测未来 token精度与加速比通常更高。从 mode.py 的注册表可以看到模块内部通过SpeculativeDecodingModeRegistry统一管理多种模式描述符其中MedusaModeDescriptor与EagleModeDescriptor分别绑定各自的配置类MedusaConfig/EagleConfig以及转换/恢复入口函数。此外源码中还注册了DFlashModeDescriptor等扩展算法但本文严格围绕官方指南介绍的 Medusa 与 EAGLE 展开。使用该模块获得一个带投机解码能力的模型只需遵循三步工作流模型转换Convert通过mtsp.convert为目标模型添加 Medusa heads 或 EAGLE 模块微调Fine-tune仅训练解码头/EAGLE 模块基座模型推荐冻结保存与重载Checkpoint Re-load通过mto.save/mto.restore或 Hugging Face 的save_pretrained/from_pretrained持久化并恢复模型。模型转换mtsp.convert 详解转换是整个流程的入口。其核心 API 定义在 speculative_decoding.pydef convert(model: nn.Module, mode: ModeLike | dict[str, Any]) - nn.Module:model待转换的基座模型如 Hugging Face 的AutoModelForCausalLMmode可以是 mode 字符串、Mode对象或「mode config」的二元组列表若传dict则按{algorithm: ..., config: ...}解析。convert内部最终调用apply_mode(model, modemode, registrySpeculativeDecodingModeRegistry)完成动态模型改写。以官方指南中的示例为基础示例中mode分支示意实际运行时请直接传入明确的配置字典import torch from transformers import AutoModelForCausalLM, AutoTokenizer import modelopt.torch.speculative as mtsp # 用户自定义模型 model AutoModelForCausalLM.from_pretrained(TinyLlama/TinyLlama-1.1B-Chat-v1.0) tokenizer AutoTokenizer.from_pretrained(TinyLlama/TinyLlama-1.1B-Chat-v1.0) tokenizer.pad_token_id tokenizer.eos_token_id if mode medusa: # 配置并转换为 Medusa config { medusa_num_heads: 2, medusa_num_layers: 1, } elif mode eagle: config { eagle_num_layers: 1 } mtsp.convert(model, [(mode, config)])需要说明的是官方文档示例属于简化示意代码。对照当前仓库源码EagleConfig中并不存在eagle_num_layers字段——EAGLE 模块的层数是在eagle_architecture_config.num_hidden_layers中配置的见下文。更贴近当前 API 的写法如下# Medusa直接给出头部与 ResBlock 层数 mtsp.convert(model, [(medusa, {medusa_num_heads: 2, medusa_num_layers: 1})]) # EAGLE层数放在 eagle_architecture_config 中 mtsp.convert(model, [(eagle, {eagle_architecture_config: {num_hidden_layers: 1}})])MedusaConfig解码头参数MedusaConfig定义在 config.py只有两个核心字段字段默认值含义medusa_num_heads2添加到模型上的 Medusa 头数量即每个生成步骤并行预测的未来 token 数medusa_num_layers1Medusa 头内部使用的 ResBlock 层数转换时medusa/conversion.py 的convert_to_medusa_model会把基座模型注册进MedusaDMRegistry然后调用medusa_model.modify(medusa_num_heads..., medusa_num_layers...)完成结构注入。需要注意的是Medusa 转换不产生额外元数据metadata {}恢复时仅依赖配置即可重建结构。EagleConfig草稿模块参数EagleConfigconfig.py字段丰富除训练控制项外还包括架构、TTTTrain-Time Test与导出相关的配置字段默认值含义eagle_offlineFalse是否消费预计算的隐藏状态离线/流式训练由 recipe 从data.mode自动推导非用户可配eagle_architecture_config{}EAGLE 草稿模块的架构配置覆盖 default_config.py 中的默认值eagle_decoder_typellamaEAGLE 解码器类型可选llama或kimik2eagle_ttt_steps3训练中 Train-Time Test 的步数eagle_mix_hidden_statesFalse是否混合多个 TTT 步的隐藏状态降低训练成本的技巧eagle_loss_decay_factor0.9多层 EAGLE loss 的衰减因子eagle_self_logit_distillationTrue是否使用基座模型 logit 蒸馏eagle_hidden_state_distillationFalse是否使用特征隐藏状态蒸馏eagle_freeze_base_modelTrue训练期间是否冻结基座模型推荐冻结eagle_reuse_base_decoderFalse是否在 EAGLE 模块中复用基座解码器EAGLE-MTP 变体使用eagle_report_accTrue是否报告评估准确率eagle_use_torch_compileTrue是否对 EAGLE forward/loss 使用torch.compile加速训练eagle_enable_nvtxFalse是否启用 NVTX 范围以便 profileeagle_export_rope_scaling{rope_type: yarn, factor: 32.0, original_max_position_embeddings: 2048}导出时注入 HF config 的 rope_scaling训练用default导出时注入 YaRN 以支持长上下文eagle_base_loraFalse是否给基座模型添加 LoRA 适配器与 EAGLE 模块共训需peft与eagle_offlineTrue不兼容eagle_base_lora_*系列参数eagle_base_lora_rank默认 64、eagle_base_lora_alpha默认 16.0、eagle_base_lora_preservation_loss_weight默认 0.1、eagle_base_lora_warmup_steps默认 0、eagle_base_lora_logits_detach_prob默认 0.5 等用于控制 LoRA 共训的细节。EagleConfig 还内置了校验器_check_rope_scaling_consistency当eagle_export_rope_scaling被设置时若eagle_architecture_config.rope_scaling中的rope_type不是default会直接抛出ValueError避免导出期的 rope 覆盖逻辑与训练期配置冲突。EAGLE 默认架构default_config.pyEAGLE 草稿模块的默认架构在 eagle/default_config.py核心默认值如下可通过eagle_architecture_config覆盖hidden_act: silu、torch_dtype: bfloat16、position_embedding_type: rope、rope_theta: 10000num_hidden_layers: 1默认单层 EAGLE 头intermediate_size: 14336、num_attention_heads: 32、num_key_value_heads: 8rms_norm_eps: 1e-05、head_dim: 128、mlp_bias: False、attention_bias: Falseuse_input_layernorm_in_first_layer: True、use_last_layernorm: False、use_aux_hidden_state: False、use_mtp_layernorm: False仓库还在 config.py 中预置了两个开箱即用的配置常量可直接作为转换输入EAGLE3_DEFAULT_CFGEAGLE3 配置启用use_aux_hidden_state: True与use_last_layernorm: TrueEAGLE_MTP_DEFAULT_CFGEAGLE-MTP 配置额外启用eagle_reuse_base_decoder: True与use_mtp_layernorm: True。单元测试 test_hf_speculative.py 展示了以EAGLE3_DEFAULT_CFG为蓝本、通过更新eagle_architecture_config定制小模型架构后执行mtsp.convert的完整流程并断言转换结果类型为mtsp.plugins.HFEagleModel是验证转换、保存与恢复闭环的可靠参考。微调投机解码模型模型转换完成后需要对解码模块Medusa heads 或 EAGLE 模块进行微调。官方指南强调基座模型推荐冻结——这与EagleConfig.eagle_freeze_base_model默认True、HFMedusaModel.forward中freeze_base_model默认True的实现完全一致。使用 HF Trainer 微调指南给出的微调代码基于transformers.Trainer并配合mto.enable_huggingface_checkpointing()让 Model-Optimizer 的状态草稿模块权重等随 HF checkpoint 一起持久化import os from transformers import Trainer import modelopt.torch.opt as mto mto.enable_huggingface_checkpointing() trainer Trainer(modelmodel, processing_classtokenizer, argstraining_args, **data_module) trainer._move_model_to_device(model, trainer.args.device) trainer.train(resume_from_checkpointcheckpoint) trainer.save_state() trainer.save_model(path to the output directory)恢复已保存的投机解码模型同样简单model AutoModelForCausalLM.from_pretrained(path to the output directory)得益于enable_huggingface_checkpointingsave_pretrained会在输出目录生成modelopt_state.pth其中记录了 medusa heads / eagle 模块的权重与结构元数据而from_pretrained会自动将模型还原为投机解码模型——这一闭环在单元测试中被显式验证save_pretrained后检查modelopt_state.pth存在再通过AutoModelForCausalLM.from_pretrained恢复并断言类型为HFEagleModel见 test_hf_speculative.py。生产级训练入口仓库的 examples/speculative_decoding 提供了比指南更完整的端到端训练管线可作为实战参考安装与数据准备pip install -U nvidia-modelopt[hf] pip install -r examples/speculative_decoding/requirements.txt python examples/dataset/make_dataset.py -f examples/dataset/example_data_config.yaml --full-conversations一键工作流bash examples/speculative_decoding/train_eagle3_and_export.sh --base_model meta-llama/Llama-3.2-1B-Instruct即可完成「初始化草稿模型 → 微调 → MT-Bench 接受率评估 → 导出可部署 checkpoint」的最小闭环。在线训练小模型与草稿模型同驻显存./launch_train.sh \ --config ../../modelopt_recipes/general/speculative_decoding/eagle3.yaml \ model.model_name_or_pathmeta-llama/Llama-3.2-1B-Instruct \ data.data_pathinput_conversations/train.jsonl \ training.output_dirckpts/llama-3.2-1b-online离线训练大模型场景先落盘中间隐藏状态再只训练草稿模型显著降低显存占用# 用 TRT-LLM 或 HF 两个后端之一生成隐藏状态 python examples/speculative_decoding/collect_hidden_states/compute_hidden_states_trtllm.py \ --model $BASE_MODEL --input-file input_conversations/train.jsonl --output-dir $HIDDEN_STATES_DIR python examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py \ --model $BASE_MODEL --input-file input_conversations/train.jsonl --output-dir $HIDDEN_STATES_DIR # 离线训练 ./launch_train.sh \ --config ../../modelopt_recipes/general/speculative_decoding/eagle3.yaml \ model.model_name_or_pathmeta-llama/Llama-3.2-1B-Instruct \ data.offline_data_path$HIDDEN_STATES_DIR \ training.output_dirckpts/llama-3.2-1b-offline流式训练对于更大的基座模型可以通过共置的vllm serve实时产出基座隐藏状态并经 NIXL RDMA 流式发送给训练器免去磁盘落盘。训练数据的选择为提高投机解码接受率推荐使用由基座模型自身生成的对话作为训练数据数据合成使草稿模型的输出分布与基座模型对齐。训练 Recipeeagle3.yaml 参数详解modelopt_recipes/general/speculative_decoding/eagle3.yaml 是 EAGLE3 训练的官方 recipe所有默认训练设置都集中在这里可在命令行以 OmegaConf dotlist 覆盖。其结构分为四段modelmodel_name_or_path基座模型、trust_remote_code、use_fake_base_for_offlinedatadata_path对话数据、offline_data_path离线隐藏状态目录、draft_vocab_cache草稿词表映射d2t.pt、vlm_img_dir/vlm_processor多模态支持等modeonline/offline/streaming由offline_data_path/streaming_server_url字段自动推导不应手工设置training训练超参数eagleEagleConfig字段。training段常用参数参数默认值含义num_train_epochs1训练轮数per_device_train_batch_size1每设备 batch sizelearning_rate1.0e-4学习率warmup_steps1000warmup 步数training_seq_len2048训练序列长度logging_steps100日志频率save_steps8192checkpoint 保存频率cp_size1长上下文训练的上下文并行度可设为N开启estimate_ar/ar_validate_stepsfalse/-1是否在训练中评估接受率answer_only_lossfalse是否只对 assistant 回复计算 losseagle段默认值与EagleConfig对齐eagle_decoder_type: llama、eagle_ttt_steps: 3、eagle_self_logit_distillation: true、eagle_freeze_base_model: true、eagle_loss_decay_factor: 0.9、eagle_export_rope_scaling: {rope_type: yarn, factor: 32.0, original_max_position_embeddings: 2048}。如需调整草稿模块架构例如改为 2 层 EAGLE 头、intermediate_size 8192只需在 YAML 中覆盖eagle: eagle_architecture_config: num_hidden_layers: 2 intermediate_size: 8192另外仓库还支持草稿词表压缩先用examples/speculative_decoding/scripts/calibrate_draft_vocab.py统计训练集中最常出现的 token生成d2t.pt映射草稿 token → 目标 token再在 YAML 中设置eagle_architecture_config.draft_vocab_size与data.draft_vocab_cache即可用更小的词表训练与推理加快收敛并降低显存开销。评估、导出与部署训练完成后示例目录提供了完整的后续链路详见 examples/speculative_decoding/README.md评估仅限在线训练 checkpointMT-Bench 接受率python examples/speculative_decoding/scripts/ar_validate.py --model_path $ONLINE_CKPT导出Model-Optimizer checkpoint → 部署兼容格式python examples/speculative_decoding/scripts/export_hf_checkpoint.py --model_path $OUTPUT_DIR --export_path $EXPORT_PATH部署导出的 checkpoint 可直接用于 TRT-LLM、SGLang 或 vLLM。以 TRT-LLM 为例通过trtllm-serve配合extra-llm-api-config.yml中的speculative_config启用 EAGLE 草稿speculative_config: decoding_type: Eagle max_draft_len: 3 speculative_model_dir: draft_model_checkpoint支持矩阵来自 examples/speculative_decoding/README.mdLLAMA 2、LLAMA 3/3.1、Mistral、Phi 3、Qwen 1.5/2/2.5/3 均支持 Medusa 与 EAGLE1/2/3Kimi-K2.5/K2.6 支持 EAGLE3。训练得到的投机模块 checkpoint 同样可以继续接入 Model-Optimizer 的 PTQ/QAT 等量化流程。算法原理解读投机解码、Medusa 与 EAGLE从自回归解码到投机解码标准的 LLM 文本生成采用自回归autoregressive解码每个步骤只生成一个 token并追加到输入上下文后再进行下一次生成。这意味着要生成K个 token就需要K次串行的模型前向推理。对于 Transformer 这类大模型推理既慢又昂贵在延迟敏感应用中尤其明显。投机解码正是为了压缩这种串行开销而提出的通用范式。其典型流程是用一个更快的自回归「草稿模型」draft model先生成一段长度为K的候选序列——既可以是并行模型也可以串行调用草稿模型K次用更强大、规模更大的「目标模型」target model对草稿整体打分一次前向即可验证K个候选 token通过采样方案决定目标模型接受哪些草稿 token并在该过程中恢复目标模型的真实分布。如果草稿模型预测对了 α 个 token那么一次验证步骤就能接受并产出 α1 个 token——生成速度因此成倍提升。Medusa并行解码头方案Medusa 是投机解码中非常流行的实现路径。它不引入额外的独立草稿模型而是在目标模型上添加少量额外的解码头Medusa heads让这些头同时预测多个未来 token。生成过程中每个解码头在对应位置产出多个候选词这些候选被组合起来通过基于树tree的注意力机制一次性处理。最后用标准的接受方案从候选集合中挑选出最长的可接受前缀继续后续解码。由于「草稿模型」就是目标模型自身Medusa 能够保证输出分布与目标模型完全一致分布无损不会因为草稿模型的偏差而改变最终生成质量。其成本是模型结构上多出的若干解码头以及树形注意力带来的计算量。EAGLE隐藏状态级自回归EAGLE 与 Medusa 的核心区别在于预测的粒度Medusa 直接基于基座模型的隐藏状态预测未来 tokentoken 级别EAGLE 则用一个轻量级自回归解码器预测未来的隐藏状态hidden states再用预测出的隐藏状态去预测未来的 token。由于「在特征隐藏状态层面做自回归」比「在 token 层面做自回归」简单得多——语义特征的变化比离散 token 的分布更平滑、更易拟合——EAGLE 对未来 token 的预测通常比 Medusa 更准确因而在实际推理中能够取得更高的加速比。在 Model-Optimizer 的实现中EAGLE 草稿模块还支持 Train-Time Testeagle_ttt_steps、多步隐藏状态混合、隐藏状态蒸馏、LoRA 共训、导出期 YaRN rope 扩展等进阶能力使其在长上下文与多模态场景下同样可用。总结三条核心结论与进一步探索三步即可落地mtsp.convert完成结构注入 → HF Trainer基座冻结微调 →enable_huggingface_checkpointing配合save_pretrained/from_pretrained保存恢复这是官方指南给出的最小可用闭环且有 单元测试 背书。算法选择有明确权衡Medusa 结构简单、分布无损EAGLE 在隐藏状态层面建模草稿更准、加速比更高。两者在示例目录的支持矩阵中对主流开源模型均有良好覆盖。生产链路完整示例目录提供在线/离线/流式三种训练模式、接受率评估、导出以及 TRT-LLM / SGLang / vLLM 部署配置训练得到的草稿模块还可继续接入量化压缩流程。若需继续深入可阅读以下仓库文件官方指南原文docs/source/guides/5_speculative_decoding.rst转换 API 与模式注册modelopt/torch/speculative/speculative_decoding.py、modelopt/torch/speculative/mode.py配置类与预置常量modelopt/torch/speculative/config.pyEAGLE 默认架构modelopt/torch/speculative/eagle/default_config.pyMedusa / EAGLE 转换实现modelopt/torch/speculative/medusa/conversion.py、modelopt/torch/speculative/eagle/conversion.py端到端训练示例examples/speculative_decoding/README.md、examples/speculative_decoding/main.pyEAGLE3 训练 recipemodelopt_recipes/general/speculative_decoding/eagle3.yaml转换/保存/恢复测试tests/unit/torch/speculative/plugins/test_hf_speculative.py赞分享人工智能大模型模型优化模型量化模型压缩【免费下载链接】Model-OptimizerA unified library of SOTA model optimization techniques like quantization, distillation, pruning, neural architecture search, speculative decoding, etc. It compresses deep learning models for downstream deployment frameworks like TensorRT-LLM, TensorRT, vLLM, etc. to optimize inference speed.项目地址https://gitcode.com/GitHub_Trending/te/Model-Optimizer点击查看免费下载相关推荐vLLM推测解码技术大幅降低LLM推理延迟vLLM推测解码技术大幅降低LLM推理延迟 引言LLM推理的延迟瓶颈 你是否曾经历过等待AI模型生成回复的漫长时刻当使用大型语言模型LLM时推理延迟人工智能大模型模型推理服务推理引擎本地部署MediaMTX低延迟HLS终极优化指南如何将流媒体延迟降至1秒内MediaMTX低延迟HLS终极优化指南如何将流媒体延迟降至1秒内 还在为HLS直播的3 5秒延迟而烦恼吗想要在安防监控、在线教育等实时场景中实现秒级响应Escrcpy低延迟模式开启游戏操控延迟降低至10msEscrcpy低延迟模式开启游戏操控延迟降低至10ms 想要在电脑上流畅玩手机游戏却因为延迟问题而烦恼Escrcpy的低延迟模式能够将游戏操控延迟降低至1桌面应用移动开发创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考