人工智能大模型模型优化模型量化模型压缩【免费下载链接】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点击查看免费下载本篇技术指南以 Model-Optimizer 仓库experimental/dms模块的 ARCHITECTURE.md 为核心骨架系统讲解 Dynamic Memory SparsificationDMS算法的内部实现驱逐决策如何生成、注意力与损失函数如何设计、压缩比调度如何运作以及如何将该机制适配到新模型家族、接入新数据集并实现训练断点恢复。读完本文你将掌握 DMS 核心组件的源码级原理并能直接复用仓库中的配置与命令完成 Qwen3-8B 的 DMS 适配训练、抽取与评估。DMS 是什么从 KV Cache 瓶颈到学习式驱逐策略在大语言模型LLM的推理阶段生成成本往往受限于 Key-ValueKVCache 的规模——随着序列变长KV Cache 占用的显存持续增长。DMS 通过学习一个逐注意力头的驱逐策略eviction policy决定生成过程中哪些 KV Cache 条目需要保留从而在保持精度的同时压缩缓存。与传统做法立即丢弃 token不同DMS 会延迟驱逐决策通过隐式地合并表征来保留关键信息。训练阶段压缩比会从 1× 逐步提升到目标值如 8×、16×同时使用知识蒸馏knowledge distillation让带压缩的学生模型对齐未压缩教师模型的输出。该实现对应的论文为Inference-Time Hyper-Scaling with KV Cache CompressionŁańcucki 等NeurIPS 2025论文编号已在 README.md 中给出。从实践角度看DMS 的优势体现在在最小精度损失下实现8× 压缩默认训练方式仅训练驱逐适配器eviction adapters冻结基础模型权重训练成本比全量微调低约两个数量级适配 Qwen3-8B 仅需约 250 训练步8× H100 上约 4 小时通过一个支持各注意力头变长序列的自定义缓存可无缝替换 Hugging Face 模型。驱逐决策Eviction Decisions两种生成方式与两种粒度DMS 的驱逐决策即每个 token 对应的是否驱逐二值判定有两种计算来源通过dms_separate_alpha配置项切换方式一从 K/Q 向量的单个神经元提取dms_separate_alphaFalse驱逐决策直接提取自 key 或 query 向量中的单个神经元对应论文Dynamic Memory Compression: Retrofitting LLMs for Accelerated InferencearXiv:2403.09636第 3.1 节。此时不引入任何新参数实现最轻量。方式二由隐藏状态经线性投影适配器产生dms_separate_alphaTrue驱逐决策由一个可学习的线性投影adapter从隐藏状态产生对应Inference-Time Hyper-Scaling with KV Cache CompressionarXiv:2506.05345第 3.2 节也是当前默认训练方案。在源码 configuration_qwen3_dms.py 中dms_separate_alpha的默认值为False而在实际训练配置 qwen3_8b.yaml 中显式设为True。从模型实现看两种方式的差异体现在注意力层。在 modeling_qwen3_dms.py 的Qwen3AttentionDMS中if self.config.dms_separate_alpha: self.dms_proj_alpha_norm Qwen3RMSNorm(config.hidden_size, epsconfig.rms_norm_eps) self.dms_proj_alpha nn.Linear(config.hidden_size, self.num_key_value_heads, biasconfig.attention_bias) else: self.dms_proj_alpha_norm None self.dms_proj_alpha None即dms_separate_alphaTrue时每个注意力层会额外挂载dms_proj_alpha_normRMSNorm与dms_proj_alpha线性层两个参数这正对应 ARCHITECTURE.md 中为注意力层添加dms_proj_alpha与dms_proj_alpha_norm层的扩展步骤。在 dms/core.py 的prepare_attention_input中可以看到两种方式的决策 logits 计算if dms_proj_alpha_fn is None: # 方式一取 query 每个 KV 头的最后一个维度作为决策 logits decision_logits ( query_states[:, ::gqa_factor, :, -1].clone() * dms_decision_scale - dms_initial_decision_offset ) else: # 方式二线性投影 归一化再缩放与平移 decision_logits ( dms_proj_alpha_fn(dms_proj_alpha_norm_fn(pre_attn_norm_hidden_states)) * dms_decision_scale - dms_initial_decision_offset )决策粒度按头head还是按层layerdms_alpha_per控制驱逐决策的共享粒度取值语义含义head按注意力头独立决策各头的 KV Cache 长度可以不同灵活性最高layer按层共享决策层内所有头保留相同数量的 token实现更规整在prepare_attention_input中dms_alpha_per layer时取第一个 KV 头的决策结果广播到全部头head时则直接逐头广播。注意 configuration_qwen3_dms.py 中有断言约束dms_alpha_per只能是head或layer其他取值会直接报错。训练时的软门控Gumbel 噪声训练阶段决策 logits 会被注入 Gumbel 噪声以实现可微分的门控对应dms.core.get_gating_with_noise推理阶段则使用硬阈值logits 0 判定为保留。在 dms/core.py 中该函数实现为def get_gating_with_noise(gating_weights, noise, tau): logits (gating_weights noise) / tau probs torch.nn.functional.sigmoid(logits) discretized (probs 0.5).to(probs.dtype) - probs.detach() probs return probs, discretized, logits其中discretized采用直通估计straight-through estimator技巧前向使用硬离散值、反向传播按软概率走保证梯度可流通。tau是温度参数配置中取0.1。此外若dms_disable_eviction为真或处于 teacher 模式决策会被清零、logits 置为 -1000相当于退化为普通模型见 qwen3_8b.yaml 中disable_eviction: false的默认形态。注意力实现训练用 FlexAttention推理用 FlashAttention 分页缓存ARCHITECTURE.md 指出给定决策 logits 的 DMS 注意力实现位于 dms/attention.py见dms_attn_train_mode。其入口dms_attention根据dms_cache is None判断当前处于训练还是评估模式训练模式dms_attn_train_mode使用 PyTorch 的flex_attention通过score_mod回调实现软驱逐——在滑动窗口外的 token其注意力分数会被加上log_sigmoid(-decision_logits)的对数掩码项源码常量MASK_VALUE -50000.0从而让被驱逐的 token 对注意力输出只有近似为零的贡献但整个掩码过程仍是可微的def score_mod(score, b, h, q_idx, k_idx): causal q_idx k_idx within_sliding_window q_idx - k_idx window_size modified_score within_sliding_window * score (1 - within_sliding_window) * ( dms_mask_values[b, k_idx] score ) return (1 - causal) * MASK_VALUE causal * modified_score推理模式dms_attn_eval_mode使用 FlashAttention 的flash_attn_with_kvcache配合分页 KV CacheDMSPagedCacheLayer逐 token 生成prefill 阶段则走dms_run_prefill_flexdms/attention_prefill.py按驱逐掩码进行精确 prefill。缓存会根据当前模式在 prefill/inference 间自动切换见 dms/cache_paged.py 的enable_prefill_mode/disable_prefill_mode。值得注意的是dms/core.py 中的prepare_attention_input、post_process_attention_output以及dms_perform_chunked_prefill都标注了torch.compile()核心算子会经过 torch.compile 编译优化setup_compile_limit_for_dms还会将 dynamo 的cache_size_limit调整为 72 以容纳逐层编译缓存。损失函数前向 KL 蒸馏 DMS 压缩损失训练目标由两部分叠加而成详见 dms/training/engine.py 的distillation_forward最终loss distil_loss dms_loss蒸馏损失distillation_loss学生与教师 logits 之间的前向 KL 散度forward KL。为避免超大词表一次展开爆显存实现按process_vocab_using_chunk对词表分块累计配置中为4096并借助torch.utils.checkpoint.checkpoint做激活重计算以降低显存s_logp s_chunk.float() - s_lse[:, :, None] t_logp t_chunk.float() - t_lse[:, :, None] # Forward KL: KLD(Student, Teacher) token_kl token_kl (s_logp.exp() * (s_logp - t_logp)).sum(dim-1)DMS 压缩损失dms_loss鼓励模型的实际驱逐比例贴近调度器给定的目标驱逐比例。模型输出中的dms_frac_closed各层逐头平均的软驱逐 token 数见 dms/core.py 的run_decoder_layers聚合逻辑被用于该损失的构建同时训练日志会记录dms_target_cr、dms_cr实际压缩比 1 / clamp(1 - frac_closed)等指标。此外distillation_forward还会在no_grad下计算学生模型自身的 LM 交叉熵calc_lm_loss作为旁路观测指标detached_lm_loss仅用于监控不参与梯度。训练日志中的指标会经ModifiedTrainer.log做跨进程all_reduce平均后统一以gl_前缀上报。DMS 调度压缩比从 1.0 线性增长到 16.0ARCHITECTURE.md 指出压缩比按训练步数从initial_cr通常 1.0线性增长到final_cr如 16.0该逻辑在 dms/training/engine.py 的dms_schedule()中实现progress min(step / max_steps, 1.0) cr dms_initial_cr (dms_final_cr - dms_initial_cr) * progress frac 1 / cr target 1 - frac # what fraction of gates to closemax_steps优先取配置中的dms_final_stepqwen3_8b 配置为510否则取training_args.max_steps该配置为544留有余量。target即当前步应关闭驱逐的 gate 比例每一步训练都会以此为基准计算 DMS 压缩损失使模型在先易后难的调度下渐进学会更激进的压缩。高级选项一Chunked Prefill分块预填充长序列 prefill 阶段峰值显存较高。DMS 支持把输入序列切成固定大小的块逐块处理以降低峰值内存。启用方式单位token 数Qwen3ForCausalLMDMS.from_pretrained(..., dms_chunked_prefill4096)底层由 dms/core.py 的dms_perform_chunked_prefill驱动按num_chunks ceil(seq_len / chunk_size)切分 hidden states并同步切分attention_mask取每个 query 的最后一行、position_ids、cache_position与position_embeddings逐块过 decoder 层后拼接输出。评估脚本 scripts/evaluate.sh 在 lm-eval-harness 的model_args中即传入dms_chunked_prefill4096配合 32K 序列的 RULER 评测使用。高级选项二Cache Preallocation缓存预分配分页 KV Cache 采用逐注意力层、可动态扩容的块表block table类似 C 的std::vector生成过程中按需增长。若提前知道最大上下文长度可以预分配以避免运行期反复扩容Qwen3ForCausalLMDMS.from_pretrained(..., dms_preallocate_for_tokens2048)在 dms/cache_paged.py 的DMSPagedCacheLayer中可以看到与内存管理相关的默认参数block_size256、growth_factor1.5、accommodate_min_initial_context_length4096、max_context_length由dms_preallocate_for_tokens决定。块表扩容时按growth_factor比例一次性拼接新 blockexpand_blocks并把新 page id 加入free_page_ids空闲池缓存层还提供offload()/prefetch()支持 CPU offload 与回迁。configuration_qwen3_dms.py中dms_preallocate_for_tokens的默认值为4096且断言dms_window_size dms_paged_attention_block_size512 256 满足。将 DMS 适配到新模型家族RetrofittingARCHITECTURE.md 给出了为新模型族添加 DMS 支持的目录规范。以models/qwen3/为参照experimental/dms/models/qwen3 下实际包含configuration_qwen3_dms.py、extract.py、modeling_qwen3_dms.py、train.py新模型应在models/下新建目录models/new_model/ ├── configuration_new_model_dms.py # 继承基础模型 config 的 DMS 配置类 ├── extract.py # 检查点抽取 ├── modeling_new_model_dms.py # 带 DMS 注意力的模型 └── train.py # 训练入口模型特定代码需要完成四件事扩展 config 类继承基础模型的 config 并加入 DMS 参数参照 configuration_qwen3_dms.pyQwen3ConfigDMS(Qwen3Config)其中包含dms_alpha_scale100.0、dms_initial_alpha_offset5.0、dms_window_size512、dms_paged_attention_block_size256、dms_cr8、dms_disable_evictionFalse、dms_separate_alphaFalse、dms_alpha_perhead、dms_tau0.1、dms_compile_limit72、dms_manual_inference_modeFalse、dms_chunked_prefillNone、dms_preallocate_for_tokens4096等默认值并内置参数合法性断言。覆写注意力 forward调用dms.core.prepare_attention_input与dms.attention.dms_attentionQwen3 实现在 modeling_qwen3_dms.py 的Qwen3AttentionDMS.forward中将 q/k/v 投影、QK 归一化、旋转位置编码、dms_proj_alpha等作为输入传给 core 函数。为注意力层添加dms_proj_alpha与dms_proj_alpha_norm层当dms_separate_alphaTrue时如上文代码所示。在configs/下添加对应的 YAML 配置参照 configs/qwen3_8b.yaml。核心 DMS 算子prepare_attention_input、dms_attention、post_process_attention_output都是模型无关的模型特定代码只需提供 Q/K/V 投影与所需的各种归一化层作为输入即可复用。添加新训练数据集训练数据管线位于 dms/training/data.py接入新数据集只需两步为数据集定义filter_fn与extract_fn创建DatasetInfo实例。ARCHITECTURE.md 给出了完整示例def my_dataset_filter_fn(ds_elem): return ds_elem[quality_score] 0.8 def my_dataset_extract_fn(ds_elem): return { conversation: [ {role: user, content: ds_elem[prompt]}, {role: assistant, content: ds_elem[response]}, ] } MyNewDataset DatasetInfo( args(org/my-dataset,), kwargs{split: train}, filter_fnmy_dataset_filter_fn, extract_fnmy_dataset_extract_fn, )随后在 YAML 配置中按名称:权重的混合格式引用data: blend: MyNewDataset:0.5,OpenR1Math220k:0.5数据层还支持train_samples采样条数、max_length截断长度、concat_always_start_new是否强制从新样本起头拼接与tokenizer_kwargs如enable_thinking: true等配置项详见 qwen3_8b.yaml。训练断点恢复Checkpoint Resume在 YAML 配置的hf_trainer段设置以下字段即可从最新检查点恢复hf_trainer: resume_from_checkpoint: autoauto会自动检测输出目录下最新的checkpoint-N目录也可以指定显式路径hf_trainer: resume_from_checkpoint: outputs/qwen3_8b/checkpoint-300恢复能够成立得益于以下机制见 dms/training/engine.pyHugging FaceTrainer会恢复优化器状态、LR 调度器状态、训练步计数器与 RNG 状态DMS 调度是确定性的dms_schedule(step)只依赖当前步数恢复后压缩比目标自动衔接Gumbel 噪声由step process_index grad_acc_step三元组播种get_dms_noise内部用str_to_seed(f{process_index}_{grad_acc_step}_{current_step})手动重置 RNG保证恢复后的噪声序列与中断前完全一致训练可无缝继续。实战从安装到训练、抽取、评估环境准备仓库被设计为在 NVIDIA PyTorch 容器内运行README.md 推荐nvcr.io/nvidia/pytorch:25.11-py3。安装 DMS 包cd Model-Optimizer/experimental/dms pip install -e .单次安装即可覆盖训练与评估所需依赖含 lm-eval-harness。需要说明的是由于本文聚焦架构原理完整训练流程与仓库结构请参阅 README.md。训练 DMS 适配器GPU 数量决定有效 batch size。以下配置在 8× H100 80GB 的 DGX 节点上验证通过小显存调试可参考 scripts/train_small.shbash scripts/train.sh configs/qwen3_8b.yamlscripts/train.sh 先以单进程准备数据集--prepare-dataset-only再通过accelerate launch -m models.qwen3.train --config $CONFIG启动分布式训练。训练冻结 Qwen3-8B 原始权重、仅训练 DMS 驱逐策略参数单个 DGX H100 节点约 4 小时完成训练结束时学生模型保存到outputs/qwen3_8b/student_model/。从中间检查点抽取可选python -m models.qwen3.extract \ --config outputs/qwen3_8b/config.yaml \ --checkpoint outputs/qwen3_8b/checkpoint-238评估在 RULER 长上下文基准上评估scripts/evaluate.shbash scripts/evaluate.sh outputs/qwen3_8b/student_model注意前提条件保存的模型依赖dms包提供注意力与缓存实现因此任何加载该模型做推理或评估的环境都必须先pip install -e .。关键配置文件逐项解读configs/qwen3_8b.yaml 是理解全部 DMS 参数落地的权威样例model: name: Qwen/Qwen3-8B dtype: float32 # 模型权重加载精度 dms: alpha_scale: 100.0 # 决策 logits 缩放因子 initial_alpha_offset: 5.0 # 决策 logits 初始偏移 window_size: 512 # DMS 滑动窗口大小必须大于 block_size 256 disable_eviction: false # 关闭驱逐即退化为普通模型 separate_alpha: true # 使用独立线性投影适配器产生决策 alpha_per: head # 决策粒度head 或 layer tau: 0.1 # Gumbel 门控温度 initial_cr: 1.0 # 初始压缩比 final_cr: 16.0 # 最终压缩比训练末期目标 final_step: 510 # 压缩比线性增长结束的步数 data: blend: OpenR1Math220k:1.0 # 数据集混合名称:权重 train_samples: 4000 # 采样样本数 max_length: 32768 # 序列最大长度 concat_always_start_new: true # 拼接时总是从新样本开始 process_vocab_using_chunk: 4096 # 蒸馏损失按词表分块大小 tokenizer_kwargs: enable_thinking: true # 启用 Qwen3 思考模式 hf_trainer: output_dir: outputs/qwen3_8b run_name: dms_qwen3_8b max_steps: 544 # 总训练步数略大于 final_step per_device_train_batch_size: 1 gradient_accumulation_steps: 1 learning_rate: 3.0e-5 weight_decay: 0.0 warmup_steps: 0 lr_scheduler_type: constant save_strategy: steps save_steps: 34 # 每 34 步保存一个检查点 save_total_limit: 5 logging_strategy: steps logging_steps: 1 gradient_checkpointing: false tf32: false bf16: true save_safetensors: false adam_beta1: 0.9 adam_beta2: 0.95 max_grad_norm: 1.0 seed: 42 fsdp: full_shard offload # FSDP 全分片 参数 offload fsdp_config: use_orig_params: true sync_module_states: true activation_checkpointing: true resume_from_checkpoint: # null 全新开始auto 最新或显式路径值得强调的是训练中仅有dms_proj_alpha含_norm以及_unfrozen_dummy_param模式的参数开启梯度见engine.py的_configure_gradients这与冻结基础权重、只训练驱逐适配器的设计一致get_student_model还支持zero_out_proj_alpha将投影适配器零初始化后再训练。当前实现的边界与限制ARCHITECTURE.md 配套的 README.md 明确了本仓库当前支持范围使用时应留意仅支持驱逐适配器训练基础模型权重保持冻结论文原始配方会更新全部权重官方计划后续支持推理仅支持单次 prefill-then-generate 流程多轮对话的prefill, generate, prefill, ...交错尚未优化独立序列之间需要重置缓存且存在用重复 prefill 模拟生成的慢速回退路径详见 example_inference.ipynb分页缓存层对 beam search 显式抛出NotImplementedError见 dms/cache_paged.py 的reorder_cache。如需追溯论文出处README 中提供了对应的 BibTeX 引用条目可依据Inference-Time Hyper-Scaling with KV Cache CompressionNeurIPS 2025进一步核对算法细节。整体而言DMS 模块在 Model-Optimizer 仓库中提供了一个模型无关核心算子 模型特定薄封装的清晰分层前者集中于 dms/core、attention、cache、training后者落在 models/qwen3这种结构正是将其快速推广到其他模型家族的关键。赞分享人工智能大模型模型优化模型量化模型压缩【免费下载链接】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点击查看免费下载相关推荐TensorRT-Model-Optimizer深度学习模型的量化与压缩TensorRT Model Optimizer深度学习模型的量化与压缩 项目介绍 TensorRT Model Optimizer 是由 NVIDIA 开发人工智能大模型模型优化模型量化模型压缩LMCache KV Cache 压缩与解压缩实战通过 Controller 对 KV Cache 执行 CacheGen 压缩LMCache KV Cache 压缩与解压缩实战通过 Controller 对 KV Cache 执行 CacheGen 压缩 导读 本篇技术指南完整讲解人工智能大模型缓存抽象模型推理服务LMCache 外部 KV Cache 压缩实战通过 Cache Controller API 对请求 KV Cache 进行 CacheGen 压缩与解压LMCache 外部 KV Cache 压缩实战通过 Cache Controller API 对请求 KV Cache 进行 CacheGen 压缩与解压人工智能大模型缓存抽象模型推理服务创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考