尧图网站建设 
人工智能大模型模型优化模型量化模型压缩【免费下载链接】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模块对 Dynamic Memory SparsificationDMS算法的完整落地实现。DMS 通过学习式驱逐策略压缩 Transformer 的 KV Cache在不明显损失精度的前提下将长序列推理成本大幅降低本仓库将该算法用于 Qwen3-8B 的适配retrofitting。读完本文你将掌握 DMS 的核心原理逐头驱逐策略、可微门控、知识蒸馏训练、paged KV Cache 推理、完整的训练/提取/评估命令与配置参数含义以及如何将 DMS 扩展到新的模型家族和数据集。为什么需要 DMS推理时扩展的 KV Cache 瓶颈推理时扩展Inference-Time Scaling通过生成更长的序列来换取更强的推理能力这在长链推理、深度思考类模型中已成为主流范式。但对于 Transformer 大语言模型生成成本往往被 KV Cache 的规模卡住序列越长每个头需要缓存的 Key/Value 越多显存占用与访问带宽都随之增长。DMSDynamic Memory Sparsification动态内存稀疏化的核心思路是学习一个逐注意力头的驱逐eviction策略在生成过程中决定哪些 KV Cache 条目该保留、哪些可以被丢弃从而在压缩缓存的同时尽量保住精度。该算法源自 Adrian Łańcucki、Konrad Staniszewski、Piotr Nawrot、Edoardo M. Ponti 的论文Inference-Time Hyper-Scaling with KV Cache CompressionNeurIPS 2025arXiv 2506.05345。Model-Optimizer 的experimental/dms目录给出了一个极简但已优化的参考实现既包含核心算法库dms也包含 Qwen3 模型的适配代码、训练/评估脚本与 YAML 配置可以直接在 8×H100 上复现 Qwen3-8B 的 DMS 适配流程。DMS 工作原理延迟驱逐 隐式信息合并与立即丢弃 token的朴素做法不同DMS 采用延迟驱逐delayed eviction驱逐决策并非在 token 写入缓存的那一刻就生效而是被推迟到后续时刻被标记为驱逐的槽位在写入新 token 时被就地覆盖相当于把新旧表示隐式地合并在一起从而保留关键信息而不是直接切断注意力路径。在训练阶段压缩比compression ratio会从 1× 逐步线性增大到目标值本仓库示例配置中为 16×README 中描述的典型目标是 8×同时使用知识蒸馏让学生模型带压缩的输出对齐未压缩教师模型的输出。通过这种课程式curriculum训练模型逐步学会在越来越激进的压缩下保持输出质量。驱逐决策的两种来源与两种粒度从源码结构看驱逐决策的生成集中在dms/core.py的prepare_attention_input中支持两种决策来源和两种决策粒度分别对应ARCHITECTURE.md的描述。决策来源一从 Key/Query 向量单神经元提取dms_separate_alphaFalse当不启用独立适配器时决策 logits 直接取自 query或 key向量最后一维的单个神经元乘上缩放系数并减去偏移量decision_logits ( query_states[:, ::gqa_factor, :, -1].clone() * dms_decision_scale - dms_initial_decision_offset )这一方式对应Dynamic Memory Compression: Retrofitting LLMs for Accelerated InferencearXiv 2403.09636第 3.1 节把每个 token 的是否保留信息编码进一个可训练/可用的标量位。注意此时该神经元位置会被置零后不再参与 RoPE 旋转代码中query_states[:, ::gqa_factor, :, -1] 0避免旋转引入位置耦合。决策来源二由隐藏状态经线性投影生成dms_separate_alphaTrue当启用独立适配器时决策 logits 由层归一化 线性投影从注意力前的隐藏状态计算得到decision_logits ( dms_proj_alpha_fn(dms_proj_alpha_norm_fn(pre_attn_norm_hidden_states)) * dms_decision_scale - dms_initial_decision_offset )这一方式对应论文第 3.2 节也是本仓库默认训练配置采用的方式。适配器由两层组成见models/qwen3/modeling_qwen3_dms.py中的Qwen3AttentionDMSdms_proj_alpha_norm一个Qwen3RMSNorm(hidden_size)dms_proj_alpha一个nn.Linear(hidden_size, num_key_value_heads)。投影输出形状为(batch, seq_len, num_kv_heads)再转置成(batch, num_kv_heads, seq_len)的逐 token 决策 logits。决策粒度逐头 vs 逐层通过dms_alpha_per配置项选择取值行为后果head默认每个注意力头独立决策各头的 KV Cache 长度可以不同压缩更精细layer层内所有头共享决策同一层所有头保留相同数量的 token实现更规整从models/qwen3/configuration_qwen3_dms.py可见dms_alpha_per只允许这两个取值非法值会在构造配置时直接断言失败。Qwen3ConfigDMS的完整 DMS 参数与默认值如下配置项默认值含义dms_alpha_scale100.0决策 logits 的缩放系数dms_initial_alpha_offset5.0决策 logits 的初始偏移dms_window_size512DMS 滑动窗口大小dms_paged_attention_block_size256paged cache 的块大小需小于窗口大小dms_cr8压缩比仅文档用途dms_disable_evictionFalse置 True 可将 DMS 模型退化为 vanilla 模型dms_separate_alphaFalse是否使用独立线性投影适配器dms_alpha_perhead驱逐决策粒度head/layerdms_tau0.1决策 logits 的 Gumbel 温度dms_compile_limit72torch.compile缓存上限dms_manual_inference_modeFalse是否手动切换 prefill/inference 模式dms_chunked_prefillNoneprefill 分块大小token 数dms_preallocate_for_tokens4096KV Cache 预分配长度配置层还做了两条硬性约束dms_paged_attention_block_size 0且dms_window_size dms_paged_attention_block_size。Gumbel 噪声与可微门控训练时决策 logits 需要可微地变成离散的保留/驱逐决策才能让梯度回传到适配器。实现位于dms/core.py的get_gating_with_noiselogits (gating_weights noise) / tau probs torch.nn.functional.sigmoid(logits) discretized (probs 0.5).to(probs.dtype) - probs.detach() probs即logits 加上 Gumbel 噪声后除以温度tau经 sigmoid 得到软概率离散决策用probs 0.5获得再通过 straight-through estimator离散值 - 概率.detach() 概率保证梯度可回传。推理阶段则直接使用硬阈值decisions (decision_logits 0)不再加噪声。噪声本身由step process_index grad_acc_step三个量做种子生成见dms/training/engine.py的get_dms_noise这也是训练可断点恢复的关键前提之一。注意力实现训练用 FlexAttention 软门控推理用 Flash Attention paged Cache注意力核心在dms/attention.py的dms_attention中分派dms_cache is None时走训练模式否则走推理模式。训练模式FlexAttention 软门控训练模式dms_attn_train_mode基于 PyTorch 的flex_attention通过score_mod实现滑动窗口 软驱逐门控。决策 logits 先做一次右移填充dms将决策 logits 相对移位 1然后取logsigmoid(-decision_logits)作为驱逐惩罚项dms_mask_values torch.nn.functional.logsigmoid(-decision_logits) 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即在滑动窗口内保持原始分数超出窗口的 key 会被加上一个由驱逐概率决定的惩罚项MASK_VALUE -50000.0用于因果掩码从而让注意力软性地被逐出窗口训练过程完全可微。推理模式Flash Attention paged KV Cache推理模式dms_attn_eval_mode使用flash_attn_with_kvcacheFlash Attention配合dms/cache_paged.py的 paged cache 实现缓存按块默认块大小 256管理每个 batch 元素通过 block table 记录已分配块每次追加新 token 时若该 token 落在滑动窗口内且滑动窗口末尾元素被标记为驱逐则用新值覆盖被驱逐的旧位置_aux_potential_eviction否则写入块内下一个空闲位置块的分配表按需动态增长类似 C 的std::vector避免一次性预留全部上下文长度推理阶段逐头序列长度可以不同这正是逐头驱逐得以实现的底层支撑。缓存层还提供Mode状态管理prefill 阶段与逐 token 生成阶段会切换模式见dms/core.py的dms_perform_chunked_prefill与dms/cache.py的DMSCache。由于这套 cache 是作为 Hugging FaceCache的替代品实现的训练得到的模型在 HF 生态中可以直接from_pretrained加载推理是drop-in replacement。损失函数与压缩比调度训练目标由两部分相加见dms/training/engine.py的distillation_forwardloss distil_loss dms_loss蒸馏损失前向 KL 散度distillation_loss计算学生与教师 logits 之间的前向 KLKLD(Student, Teacher)。实现上按process_vocab_using_chunk配置中为 4096对词表分块累加避免一次性物化整个词表的概率分布同时通过torch.utils.checkpoint.checkpointuse_reentrantFalse做激活重计算以省显存。只有attention_mask与eos_mask同时为真的位置参与蒸馏损失计算不预测掩码位置、不预测 EOS token。DMS 压缩损失dms_loss鼓励模型实际驱逐的比例对齐当前调度目标输出张量中携带dms_frac_closed各头平均的软驱逐比例训练日志中同时记录dms_target_cr与dms_cr 1 / clamp(1 - dms_closed_frac)供监控。压缩比调度dms_schedule按当前训练步线性拉升压缩比progress min(step / max_steps, 1.0) cr dms_initial_cr (dms_final_cr - dms_initial_cr) * progress target 1 - 1 / crinitial_cr通常为1.0不压缩final_cr为目标压缩比配置示例为16.0final_step为达到目标压缩比的步数配置示例为510未设置时用max_steps代替。实用特性一览特性说明8× 压缩README 表明该方法可在精度损失极小的情况下达到 8× 压缩Adapter 训练默认配方只训练驱逐适配器参数冻结基础权重成本远低于全量微调极短训练周期README 标明约 250 训练步8×H100 上约 4 小时即可适配 Qwen3-8BDrop-in 替换通过自定义 cache 支持各注意力头不同的序列长度可直接替换 Hugging Face 模型加载推理仓库 README 提供的训练耗时参考为Qwen3-8B8×H100约 4 小时注意仓库同时提供了final_cr16.0、final_step510、max_steps544的默认配置实际压缩比与步数可按需在 YAML 中调整。快速开始给 Qwen3-8B 加装 DMS环境与安装仓库设计为在 NVIDIA PyTorch 容器内运行docker pull nvcr.io/nvidia/pytorch:25.11-py3克隆后在experimental/dms目录下执行可编辑安装pyproject.toml声明了transformers4.57.3、datasets4.4.2、accelerate1.4.0、lm_eval[ruler]等依赖一次安装即可覆盖训练与评估git clone https://github.com/NVIDIA/Model-Optimizer cd Model-Optimizer/experimental/dms pip install -e .训练 DMS 适配器执行scripts/train.shbash scripts/train.sh configs/qwen3_8b.yamltrain.sh内部先做两件事先用单进程python -m models.qwen3.train --config $CONFIG --prepare-dataset-only准备数据集再用accelerate launch -m models.qwen3.train --config $CONFIG拉起分布式训练并预设PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True缓解显存碎片。训练会冻结原始 Qwen3-8B 权重仅训练 DMS 驱逐策略参数并使用知识蒸馏对齐未压缩教师模型的输出。GPU 数量决定有效 batch size配置是在 DGX H1008×H100 80GB上验证的若想在更小的算力预算下调试例如单张 RTX 5090可参考scripts/train_small.sh单卡 FSDP offload配套configs/qwen3_1.7b.yaml并额外提供forward_fn_kwargs.train_attn_kwargs.kernel_options来控制 FlexAttention 的 kernel 块尺寸例如BLOCK_M116、BLOCK_N116等。训练结束后学生模型会自动保存到outputs/qwen3_8b/student_model/对应train.py末尾的extract_student_model调用。配置文件详解训练入口models/qwen3/train.py从 YAML 读取三块配置model、dms、data、hf_trainer。以下为configs/qwen3_8b.yaml的完整内容model: name: Qwen/Qwen3-8B dtype: float32 dms: alpha_scale: 100.0 initial_alpha_offset: 5.0 window_size: 512 disable_eviction: false separate_alpha: true alpha_per: head tau: 0.1 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 hf_trainer: output_dir: outputs/qwen3_8b run_name: dms_qwen3_8b max_steps: 544 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 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_config: use_orig_params: true sync_module_states: true activation_checkpointing: true resume_from_checkpoint: # null fresh start, auto latest, or explicit path关键参数解读配置段参数含义modelname/dtype学生模型标识与加载精度float32加载、内部按需转bf16dmsseparate_alpha: true使用线性投影适配器产生决策对应论文 3.2 节dmsalpha_per: head逐头驱逐各头缓存长度可不同dmswindow_size: 512滑动窗口需大于 paged 块大小 256dmsfinal_cr / final_step压缩比调度目标与到达步数datablend数据集混合串格式为数据集名:权重逗号分隔datatrain_samples: 4000训练样本数datamax_length: 32768拼接后的最大序列长度长上下文datatokenizer_kwargs.enable_thinking是否启用 Qwen3 的 thinking 模式模板hf_trainerfsdp: full_shard offload全分片 参数 offload 的 FSDP 策略适配 8×H100 单机hf_trainerbf16: true训练使用 bf16 混合精度从源码看dms_cfg中的initial_cr、final_cr、final_step会被剔出模型构造参数build_combined_model中if k not in (initial_cr, final_cr, final_step)仅用于构造调度函数其余字段会以dms_前缀传给Qwen3ForCausalLMDMS.from_pretrained。此外学生模型加载时会zero_out_proj_alpha即把dms_proj_alpha权重清零后再开始训练让驱逐门在初始时全开。若学生与教师的model_name_or_path和dtype完全相同训练框架会自动复用同一个模型作为教师student_is_teacher优化节省一倍显存。从中间 checkpoint 提取学生模型可选如果想从某个中间 checkpoint 提取可直接部署的学生模型执行python -m models.qwen3.extract \ --config outputs/qwen3_8b/config.yaml \ --checkpoint outputs/qwen3_8b/checkpoint-238models/qwen3/extract.py会加载该 checkpoint 的状态字典pytorch_model.bin重建 combined model 后抽出学生权重输出默认在checkpoint/student_model也可用--output指定。被保存的模型目录包含bf16 权重、带auto_map的config.jsontrust_remote_codetrue加载、两份模型实现文件configuration_qwen3_dms.py与modeling_qwen3_dms.py以及 tokenizer。在 RULER 长上下文基准上评估评估脚本scripts/evaluate.sh基于 lm-eval-harnessbash scripts/evaluate.sh outputs/qwen3_8b/student_model该脚本使用niah_single_2RULER 的 needle-in-a-haystack 单针任务batch_size 2max_seq_lengths 32768并通过model_args传入dms_chunked_prefill4096。前提保存的模型依赖dms包提供注意力与缓存实现因此任何加载该模型进行推理或评估的环境都必须先pip install -e .安装dms。高级选项Chunked Prefill降低 prefill 峰值显存prefill 阶段把超长输入按固定 token 数分块逐段过 decoder 层见dms/core.py的dms_perform_chunked_prefill分块处理 mask、position_ids、cache_position 与 RoPE 位置嵌入最后拼接各块输出从而显著降低峰值显存。启用方式Qwen3ForCausalLMDMS.from_pretrained(..., dms_chunked_prefill4096)评估脚本中即以 4096 为默认分块大小。Cache 预分配避免生成期动态扩容paged cache 的 block table 在生成过程中会按需增长如果预先知道最大上下文长度可通过预分配消除运行时重分配的开销Qwen3ForCausalLMDMS.from_pretrained(..., dms_preallocate_for_tokens2048)该配置项默认值为 4096见Qwen3ConfigDMS与上述分块大小配合使用。训练中断恢复在 YAML 中设置hf_trainer: resume_from_checkpoint: autotrain.py的resolve_checkpoint会自动扫描输出目录下所有checkpoint-N目录并选取步数最大者也可显式指定路径hf_trainer: resume_from_checkpoint: outputs/qwen3_8b/checkpoint-300恢复之所以可靠源于三点设计见ARCHITECTURE.mdHugging Face Trainer 会恢复优化器状态、LR 调度器状态、训练步计数与 RNG 状态DMS 压缩比调度是当前步的确定性函数dms_schedule(step)Gumbel 噪声的种子由step process_index grad_acc_step决定恢复后与中断前完全一致。扩展到新的模型家族DMS 的核心算子prepare_attention_input、dms_attention、post_process_attention_output是模型无关的模型侧只需提供 Q/K/V 投影与所需归一化层。要支持新模型在models/下创建新目录models/new_model/ ├── configuration_new_model_dms.py # Config extending the base model config ├── extract.py # Checkpoint extraction ├── modeling_new_model_dms.py # Model with DMS attention └── train.py # Training entry point改造步骤以 Qwen3 为参照见models/qwen3/modeling_qwen3_dms.py扩展模型 config 类并加入 DMS 参数参照models/qwen3/configuration_qwen3_dms.py的Qwen3ConfigDMS包括对window_size block_size、alpha_per取值等的断言重写注意力 forward依次调用dms.core.prepare_attention_input与dms.attention.dms_attention最后用post_process_attention_output处理后投影在注意力层加入dms_proj_alpha与dms_proj_alpha_norm两层separate_alphaTrue时在configs/下添加对应 YAML 配置。新增训练数据集训练数据流水线位于dms/training/data.py支持数据集名:权重的混合blend机制并带有分阶段的磁盘缓存plain → tokenized → shuffled → concatenated。新增数据集只需两步为数据集定义filter_fn与extract_fn创建DatasetInfo实例。示例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控制是否总是从新样本开始拼接。仓库结构速览experimental/dms/ ├── configs # YAML 实验配置 │ ├── qwen3_8b.yaml │ └── qwen3_1.7b.yaml ├── dms # 核心 DMS 库pip install -e . │ ├── attention_prefill.py # 精确 prefill基于驱逐掩码 │ ├── attention.py # DMS 注意力训练 推理模式 │ ├── cache_paged.py # 分页缓存块式内存管理 │ ├── cache.py # KV CacheHF 包装 组合式 连续式 │ ├── core.py # 共享算子prepare_attention_input、门控、分块 prefill │ └── training │ ├── data.py # 数据流水线加载、混合、分词 │ └── engine.py # 蒸馏、模型配置、噪声、训练器状态 ├── ARCHITECTURE.md ├── example_inference.ipynb ├── models # 模型适配代码 │ └── qwen3 │ ├── configuration_qwen3_dms.py │ ├── extract.py │ ├── modeling_qwen3_dms.py │ └── train.py ├── scripts # 启动脚本 │ ├── evaluate.sh │ ├── train.sh │ └── train_small.sh └── pyproject.toml更多代码细节、高级选项与扩展指南见ARCHITECTURE.md推理接口示例见example_inference.ipynb。当前限制训练侧目前仅支持训练驱逐适配器、冻结基础权重。仓库说明这种训练方式在精度相当的前提下成本约比全量微调低两个数量级论文原始配方会更新全部模型权重仓库计划在近期支持。推理侧当前支持单次 prefill-then-generate 工作流。多轮对话中交错的prefill, generate, prefill, ...尚未优化独立序列之间必须重置 cache并存在一个用重复 prefill 模拟生成的慢速 fallback。相关细节可参阅example_inference.ipynb。引用如果 DMS 对你的工作有帮助可按仓库提供的 BibTeX 引用论文Inference-Time Hyper-Scaling with KV Cache CompressionAdrian Łańcucki 等人NeurIPS 2025逐头驱逐的单神经元决策方案则源自Dynamic Memory Compression: Retrofitting LLMs for Accelerated InferencearXiv 2403.09636。赞分享人工智能大模型模型优化模型量化模型压缩【免费下载链接】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 DMS 架构深度解析KV Cache 压缩的驱逐决策、训练机制与模型扩展实战指南Model Optimizer DMS 架构深度解析KV Cache 压缩的驱逐决策、训练机制与模型扩展实战指南 本篇技术指南以 Model Optimize人工智能大模型模型优化模型量化模型压缩LMCache KV Cache 压缩与解压缩实战通过 Controller 对 KV Cache 执行 CacheGen 压缩LMCache KV Cache 压缩与解压缩实战通过 Controller 对 KV Cache 执行 CacheGen 压缩 导读 本篇技术指南完整讲解人工智能大模型缓存抽象模型推理服务NVIDIA Model Optimizer Local-Hessian基于局部 Hessian 的 NVFP4 逐块权重缩放算法实战指南NVIDIA Model Optimizer Local Hessian基于局部 Hessian 的 NVFP4 逐块权重缩放算法实战指南 本篇技术指南深入讲人工智能大模型模型优化模型量化模型压缩创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 基于 DMS 的 KV Cache 压缩实战：为 Qwen3-8B 训练驱逐适配器（Model-Optimizer 实验模块）
人工智能大模型模型优化模型量化模型压缩【免费下载链接】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模块对 Dynamic Memory SparsificationDMS算法的完整落地实现。DMS 通过学习式驱逐策略压缩 Transformer 的 KV Cache在不明显损失精度的前提下将长序列推理成本大幅降低本仓库将该算法用于 Qwen3-8B 的适配retrofitting。读完本文你将掌握 DMS 的核心原理逐头驱逐策略、可微门控、知识蒸馏训练、paged KV Cache 推理、完整的训练/提取/评估命令与配置参数含义以及如何将 DMS 扩展到新的模型家族和数据集。为什么需要 DMS推理时扩展的 KV Cache 瓶颈推理时扩展Inference-Time Scaling通过生成更长的序列来换取更强的推理能力这在长链推理、深度思考类模型中已成为主流范式。但对于 Transformer 大语言模型生成成本往往被 KV Cache 的规模卡住序列越长每个头需要缓存的 Key/Value 越多显存占用与访问带宽都随之增长。DMSDynamic Memory Sparsification动态内存稀疏化的核心思路是学习一个逐注意力头的驱逐eviction策略在生成过程中决定哪些 KV Cache 条目该保留、哪些可以被丢弃从而在压缩缓存的同时尽量保住精度。该算法源自 Adrian Łańcucki、Konrad Staniszewski、Piotr Nawrot、Edoardo M. Ponti 的论文Inference-Time Hyper-Scaling with KV Cache CompressionNeurIPS 2025arXiv 2506.05345。Model-Optimizer 的experimental/dms目录给出了一个极简但已优化的参考实现既包含核心算法库dms也包含 Qwen3 模型的适配代码、训练/评估脚本与 YAML 配置可以直接在 8×H100 上复现 Qwen3-8B 的 DMS 适配流程。DMS 工作原理延迟驱逐 隐式信息合并与立即丢弃 token的朴素做法不同DMS 采用延迟驱逐delayed eviction驱逐决策并非在 token 写入缓存的那一刻就生效而是被推迟到后续时刻被标记为驱逐的槽位在写入新 token 时被就地覆盖相当于把新旧表示隐式地合并在一起从而保留关键信息而不是直接切断注意力路径。在训练阶段压缩比compression ratio会从 1× 逐步线性增大到目标值本仓库示例配置中为 16×README 中描述的典型目标是 8×同时使用知识蒸馏让学生模型带压缩的输出对齐未压缩教师模型的输出。通过这种课程式curriculum训练模型逐步学会在越来越激进的压缩下保持输出质量。驱逐决策的两种来源与两种粒度从源码结构看驱逐决策的生成集中在dms/core.py的prepare_attention_input中支持两种决策来源和两种决策粒度分别对应ARCHITECTURE.md的描述。决策来源一从 Key/Query 向量单神经元提取dms_separate_alphaFalse当不启用独立适配器时决策 logits 直接取自 query或 key向量最后一维的单个神经元乘上缩放系数并减去偏移量decision_logits ( query_states[:, ::gqa_factor, :, -1].clone() * dms_decision_scale - dms_initial_decision_offset )这一方式对应Dynamic Memory Compression: Retrofitting LLMs for Accelerated InferencearXiv 2403.09636第 3.1 节把每个 token 的是否保留信息编码进一个可训练/可用的标量位。注意此时该神经元位置会被置零后不再参与 RoPE 旋转代码中query_states[:, ::gqa_factor, :, -1] 0避免旋转引入位置耦合。决策来源二由隐藏状态经线性投影生成dms_separate_alphaTrue当启用独立适配器时决策 logits 由层归一化 线性投影从注意力前的隐藏状态计算得到decision_logits ( dms_proj_alpha_fn(dms_proj_alpha_norm_fn(pre_attn_norm_hidden_states)) * dms_decision_scale - dms_initial_decision_offset )这一方式对应论文第 3.2 节也是本仓库默认训练配置采用的方式。适配器由两层组成见models/qwen3/modeling_qwen3_dms.py中的Qwen3AttentionDMSdms_proj_alpha_norm一个Qwen3RMSNorm(hidden_size)dms_proj_alpha一个nn.Linear(hidden_size, num_key_value_heads)。投影输出形状为(batch, seq_len, num_kv_heads)再转置成(batch, num_kv_heads, seq_len)的逐 token 决策 logits。决策粒度逐头 vs 逐层通过dms_alpha_per配置项选择取值行为后果head默认每个注意力头独立决策各头的 KV Cache 长度可以不同压缩更精细layer层内所有头共享决策同一层所有头保留相同数量的 token实现更规整从models/qwen3/configuration_qwen3_dms.py可见dms_alpha_per只允许这两个取值非法值会在构造配置时直接断言失败。Qwen3ConfigDMS的完整 DMS 参数与默认值如下配置项默认值含义dms_alpha_scale100.0决策 logits 的缩放系数dms_initial_alpha_offset5.0决策 logits 的初始偏移dms_window_size512DMS 滑动窗口大小dms_paged_attention_block_size256paged cache 的块大小需小于窗口大小dms_cr8压缩比仅文档用途dms_disable_evictionFalse置 True 可将 DMS 模型退化为 vanilla 模型dms_separate_alphaFalse是否使用独立线性投影适配器dms_alpha_perhead驱逐决策粒度head/layerdms_tau0.1决策 logits 的 Gumbel 温度dms_compile_limit72torch.compile缓存上限dms_manual_inference_modeFalse是否手动切换 prefill/inference 模式dms_chunked_prefillNoneprefill 分块大小token 数dms_preallocate_for_tokens4096KV Cache 预分配长度配置层还做了两条硬性约束dms_paged_attention_block_size 0且dms_window_size dms_paged_attention_block_size。Gumbel 噪声与可微门控训练时决策 logits 需要可微地变成离散的保留/驱逐决策才能让梯度回传到适配器。实现位于dms/core.py的get_gating_with_noiselogits (gating_weights noise) / tau probs torch.nn.functional.sigmoid(logits) discretized (probs 0.5).to(probs.dtype) - probs.detach() probs即logits 加上 Gumbel 噪声后除以温度tau经 sigmoid 得到软概率离散决策用probs 0.5获得再通过 straight-through estimator离散值 - 概率.detach() 概率保证梯度可回传。推理阶段则直接使用硬阈值decisions (decision_logits 0)不再加噪声。噪声本身由step process_index grad_acc_step三个量做种子生成见dms/training/engine.py的get_dms_noise这也是训练可断点恢复的关键前提之一。注意力实现训练用 FlexAttention 软门控推理用 Flash Attention paged Cache注意力核心在dms/attention.py的dms_attention中分派dms_cache is None时走训练模式否则走推理模式。训练模式FlexAttention 软门控训练模式dms_attn_train_mode基于 PyTorch 的flex_attention通过score_mod实现滑动窗口 软驱逐门控。决策 logits 先做一次右移填充dms将决策 logits 相对移位 1然后取logsigmoid(-decision_logits)作为驱逐惩罚项dms_mask_values torch.nn.functional.logsigmoid(-decision_logits) 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即在滑动窗口内保持原始分数超出窗口的 key 会被加上一个由驱逐概率决定的惩罚项MASK_VALUE -50000.0用于因果掩码从而让注意力软性地被逐出窗口训练过程完全可微。推理模式Flash Attention paged KV Cache推理模式dms_attn_eval_mode使用flash_attn_with_kvcacheFlash Attention配合dms/cache_paged.py的 paged cache 实现缓存按块默认块大小 256管理每个 batch 元素通过 block table 记录已分配块每次追加新 token 时若该 token 落在滑动窗口内且滑动窗口末尾元素被标记为驱逐则用新值覆盖被驱逐的旧位置_aux_potential_eviction否则写入块内下一个空闲位置块的分配表按需动态增长类似 C 的std::vector避免一次性预留全部上下文长度推理阶段逐头序列长度可以不同这正是逐头驱逐得以实现的底层支撑。缓存层还提供Mode状态管理prefill 阶段与逐 token 生成阶段会切换模式见dms/core.py的dms_perform_chunked_prefill与dms/cache.py的DMSCache。由于这套 cache 是作为 Hugging FaceCache的替代品实现的训练得到的模型在 HF 生态中可以直接from_pretrained加载推理是drop-in replacement。损失函数与压缩比调度训练目标由两部分相加见dms/training/engine.py的distillation_forwardloss distil_loss dms_loss蒸馏损失前向 KL 散度distillation_loss计算学生与教师 logits 之间的前向 KLKLD(Student, Teacher)。实现上按process_vocab_using_chunk配置中为 4096对词表分块累加避免一次性物化整个词表的概率分布同时通过torch.utils.checkpoint.checkpointuse_reentrantFalse做激活重计算以省显存。只有attention_mask与eos_mask同时为真的位置参与蒸馏损失计算不预测掩码位置、不预测 EOS token。DMS 压缩损失dms_loss鼓励模型实际驱逐的比例对齐当前调度目标输出张量中携带dms_frac_closed各头平均的软驱逐比例训练日志中同时记录dms_target_cr与dms_cr 1 / clamp(1 - dms_closed_frac)供监控。压缩比调度dms_schedule按当前训练步线性拉升压缩比progress min(step / max_steps, 1.0) cr dms_initial_cr (dms_final_cr - dms_initial_cr) * progress target 1 - 1 / crinitial_cr通常为1.0不压缩final_cr为目标压缩比配置示例为16.0final_step为达到目标压缩比的步数配置示例为510未设置时用max_steps代替。实用特性一览特性说明8× 压缩README 表明该方法可在精度损失极小的情况下达到 8× 压缩Adapter 训练默认配方只训练驱逐适配器参数冻结基础权重成本远低于全量微调极短训练周期README 标明约 250 训练步8×H100 上约 4 小时即可适配 Qwen3-8BDrop-in 替换通过自定义 cache 支持各注意力头不同的序列长度可直接替换 Hugging Face 模型加载推理仓库 README 提供的训练耗时参考为Qwen3-8B8×H100约 4 小时注意仓库同时提供了final_cr16.0、final_step510、max_steps544的默认配置实际压缩比与步数可按需在 YAML 中调整。快速开始给 Qwen3-8B 加装 DMS环境与安装仓库设计为在 NVIDIA PyTorch 容器内运行docker pull nvcr.io/nvidia/pytorch:25.11-py3克隆后在experimental/dms目录下执行可编辑安装pyproject.toml声明了transformers4.57.3、datasets4.4.2、accelerate1.4.0、lm_eval[ruler]等依赖一次安装即可覆盖训练与评估git clone https://github.com/NVIDIA/Model-Optimizer cd Model-Optimizer/experimental/dms pip install -e .训练 DMS 适配器执行scripts/train.shbash scripts/train.sh configs/qwen3_8b.yamltrain.sh内部先做两件事先用单进程python -m models.qwen3.train --config $CONFIG --prepare-dataset-only准备数据集再用accelerate launch -m models.qwen3.train --config $CONFIG拉起分布式训练并预设PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True缓解显存碎片。训练会冻结原始 Qwen3-8B 权重仅训练 DMS 驱逐策略参数并使用知识蒸馏对齐未压缩教师模型的输出。GPU 数量决定有效 batch size配置是在 DGX H1008×H100 80GB上验证的若想在更小的算力预算下调试例如单张 RTX 5090可参考scripts/train_small.sh单卡 FSDP offload配套configs/qwen3_1.7b.yaml并额外提供forward_fn_kwargs.train_attn_kwargs.kernel_options来控制 FlexAttention 的 kernel 块尺寸例如BLOCK_M116、BLOCK_N116等。训练结束后学生模型会自动保存到outputs/qwen3_8b/student_model/对应train.py末尾的extract_student_model调用。配置文件详解训练入口models/qwen3/train.py从 YAML 读取三块配置model、dms、data、hf_trainer。以下为configs/qwen3_8b.yaml的完整内容model: name: Qwen/Qwen3-8B dtype: float32 dms: alpha_scale: 100.0 initial_alpha_offset: 5.0 window_size: 512 disable_eviction: false separate_alpha: true alpha_per: head tau: 0.1 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 hf_trainer: output_dir: outputs/qwen3_8b run_name: dms_qwen3_8b max_steps: 544 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 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_config: use_orig_params: true sync_module_states: true activation_checkpointing: true resume_from_checkpoint: # null fresh start, auto latest, or explicit path关键参数解读配置段参数含义modelname/dtype学生模型标识与加载精度float32加载、内部按需转bf16dmsseparate_alpha: true使用线性投影适配器产生决策对应论文 3.2 节dmsalpha_per: head逐头驱逐各头缓存长度可不同dmswindow_size: 512滑动窗口需大于 paged 块大小 256dmsfinal_cr / final_step压缩比调度目标与到达步数datablend数据集混合串格式为数据集名:权重逗号分隔datatrain_samples: 4000训练样本数datamax_length: 32768拼接后的最大序列长度长上下文datatokenizer_kwargs.enable_thinking是否启用 Qwen3 的 thinking 模式模板hf_trainerfsdp: full_shard offload全分片 参数 offload 的 FSDP 策略适配 8×H100 单机hf_trainerbf16: true训练使用 bf16 混合精度从源码看dms_cfg中的initial_cr、final_cr、final_step会被剔出模型构造参数build_combined_model中if k not in (initial_cr, final_cr, final_step)仅用于构造调度函数其余字段会以dms_前缀传给Qwen3ForCausalLMDMS.from_pretrained。此外学生模型加载时会zero_out_proj_alpha即把dms_proj_alpha权重清零后再开始训练让驱逐门在初始时全开。若学生与教师的model_name_or_path和dtype完全相同训练框架会自动复用同一个模型作为教师student_is_teacher优化节省一倍显存。从中间 checkpoint 提取学生模型可选如果想从某个中间 checkpoint 提取可直接部署的学生模型执行python -m models.qwen3.extract \ --config outputs/qwen3_8b/config.yaml \ --checkpoint outputs/qwen3_8b/checkpoint-238models/qwen3/extract.py会加载该 checkpoint 的状态字典pytorch_model.bin重建 combined model 后抽出学生权重输出默认在checkpoint/student_model也可用--output指定。被保存的模型目录包含bf16 权重、带auto_map的config.jsontrust_remote_codetrue加载、两份模型实现文件configuration_qwen3_dms.py与modeling_qwen3_dms.py以及 tokenizer。在 RULER 长上下文基准上评估评估脚本scripts/evaluate.sh基于 lm-eval-harnessbash scripts/evaluate.sh outputs/qwen3_8b/student_model该脚本使用niah_single_2RULER 的 needle-in-a-haystack 单针任务batch_size 2max_seq_lengths 32768并通过model_args传入dms_chunked_prefill4096。前提保存的模型依赖dms包提供注意力与缓存实现因此任何加载该模型进行推理或评估的环境都必须先pip install -e .安装dms。高级选项Chunked Prefill降低 prefill 峰值显存prefill 阶段把超长输入按固定 token 数分块逐段过 decoder 层见dms/core.py的dms_perform_chunked_prefill分块处理 mask、position_ids、cache_position 与 RoPE 位置嵌入最后拼接各块输出从而显著降低峰值显存。启用方式Qwen3ForCausalLMDMS.from_pretrained(..., dms_chunked_prefill4096)评估脚本中即以 4096 为默认分块大小。Cache 预分配避免生成期动态扩容paged cache 的 block table 在生成过程中会按需增长如果预先知道最大上下文长度可通过预分配消除运行时重分配的开销Qwen3ForCausalLMDMS.from_pretrained(..., dms_preallocate_for_tokens2048)该配置项默认值为 4096见Qwen3ConfigDMS与上述分块大小配合使用。训练中断恢复在 YAML 中设置hf_trainer: resume_from_checkpoint: autotrain.py的resolve_checkpoint会自动扫描输出目录下所有checkpoint-N目录并选取步数最大者也可显式指定路径hf_trainer: resume_from_checkpoint: outputs/qwen3_8b/checkpoint-300恢复之所以可靠源于三点设计见ARCHITECTURE.mdHugging Face Trainer 会恢复优化器状态、LR 调度器状态、训练步计数与 RNG 状态DMS 压缩比调度是当前步的确定性函数dms_schedule(step)Gumbel 噪声的种子由step process_index grad_acc_step决定恢复后与中断前完全一致。扩展到新的模型家族DMS 的核心算子prepare_attention_input、dms_attention、post_process_attention_output是模型无关的模型侧只需提供 Q/K/V 投影与所需归一化层。要支持新模型在models/下创建新目录models/new_model/ ├── configuration_new_model_dms.py # Config extending the base model config ├── extract.py # Checkpoint extraction ├── modeling_new_model_dms.py # Model with DMS attention └── train.py # Training entry point改造步骤以 Qwen3 为参照见models/qwen3/modeling_qwen3_dms.py扩展模型 config 类并加入 DMS 参数参照models/qwen3/configuration_qwen3_dms.py的Qwen3ConfigDMS包括对window_size block_size、alpha_per取值等的断言重写注意力 forward依次调用dms.core.prepare_attention_input与dms.attention.dms_attention最后用post_process_attention_output处理后投影在注意力层加入dms_proj_alpha与dms_proj_alpha_norm两层separate_alphaTrue时在configs/下添加对应 YAML 配置。新增训练数据集训练数据流水线位于dms/training/data.py支持数据集名:权重的混合blend机制并带有分阶段的磁盘缓存plain → tokenized → shuffled → concatenated。新增数据集只需两步为数据集定义filter_fn与extract_fn创建DatasetInfo实例。示例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控制是否总是从新样本开始拼接。仓库结构速览experimental/dms/ ├── configs # YAML 实验配置 │ ├── qwen3_8b.yaml │ └── qwen3_1.7b.yaml ├── dms # 核心 DMS 库pip install -e . │ ├── attention_prefill.py # 精确 prefill基于驱逐掩码 │ ├── attention.py # DMS 注意力训练 推理模式 │ ├── cache_paged.py # 分页缓存块式内存管理 │ ├── cache.py # KV CacheHF 包装 组合式 连续式 │ ├── core.py # 共享算子prepare_attention_input、门控、分块 prefill │ └── training │ ├── data.py # 数据流水线加载、混合、分词 │ └── engine.py # 蒸馏、模型配置、噪声、训练器状态 ├── ARCHITECTURE.md ├── example_inference.ipynb ├── models # 模型适配代码 │ └── qwen3 │ ├── configuration_qwen3_dms.py │ ├── extract.py │ ├── modeling_qwen3_dms.py │ └── train.py ├── scripts # 启动脚本 │ ├── evaluate.sh │ ├── train.sh │ └── train_small.sh └── pyproject.toml更多代码细节、高级选项与扩展指南见ARCHITECTURE.md推理接口示例见example_inference.ipynb。当前限制训练侧目前仅支持训练驱逐适配器、冻结基础权重。仓库说明这种训练方式在精度相当的前提下成本约比全量微调低两个数量级论文原始配方会更新全部模型权重仓库计划在近期支持。推理侧当前支持单次 prefill-then-generate 工作流。多轮对话中交错的prefill, generate, prefill, ...尚未优化独立序列之间必须重置 cache并存在一个用重复 prefill 模拟生成的慢速 fallback。相关细节可参阅example_inference.ipynb。引用如果 DMS 对你的工作有帮助可按仓库提供的 BibTeX 引用论文Inference-Time Hyper-Scaling with KV Cache CompressionAdrian Łańcucki 等人NeurIPS 2025逐头驱逐的单神经元决策方案则源自Dynamic Memory Compression: Retrofitting LLMs for Accelerated InferencearXiv 2403.09636。赞分享人工智能大模型模型优化模型量化模型压缩【免费下载链接】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 DMS 架构深度解析KV Cache 压缩的驱逐决策、训练机制与模型扩展实战指南Model Optimizer DMS 架构深度解析KV Cache 压缩的驱逐决策、训练机制与模型扩展实战指南 本篇技术指南以 Model Optimize人工智能大模型模型优化模型量化模型压缩LMCache KV Cache 压缩与解压缩实战通过 Controller 对 KV Cache 执行 CacheGen 压缩LMCache KV Cache 压缩与解压缩实战通过 Controller 对 KV Cache 执行 CacheGen 压缩 导读 本篇技术指南完整讲解人工智能大模型缓存抽象模型推理服务NVIDIA Model Optimizer Local-Hessian基于局部 Hessian 的 NVFP4 逐块权重缩放算法实战指南NVIDIA Model Optimizer Local Hessian基于局部 Hessian 的 NVFP4 逐块权重缩放算法实战指南 本篇技术指南深入讲人工智能大模型模型优化模型量化模型压缩创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考