企业官网建设流程全解析

• 人工智能
• 大模型
• 模型优化
• 模型量化
• 模型压缩

本篇技术指南以 Model-Optimizer 仓库experimental/dms模块的 ARCHITECTURE.md 为核心骨架，系统讲解 Dynamic Memory Sparsification（DMS）算法的内部实现：驱逐决策如何生成、注意力与损失函数如何设计、压缩比调度如何运作，以及如何将该机制适配到新模型家族、接入新数据集并实现训练断点恢复。读完本文，你将掌握 DMS 核心组件的源码级原理，并能直接复用仓库中的配置与命令完成 Qwen3-8B 的 DMS 适配训练、抽取与评估。

DMS 是什么：从 KV Cache 瓶颈到学习式驱逐策略

在大语言模型（LLM）的推理阶段，生成成本往往受限于 Key-Value（KV）Cache 的规模——随着序列变长，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_alpha=False）

驱逐决策直接提取自 key 或 query 向量中的单个神经元，对应论文Dynamic Memory Compression: Retrofitting LLMs for Accelerated Inference（arXiv:2403.09636）第 3.1 节。此时不引入任何新参数，实现最轻量。

方式二：由隐藏状态经线性投影适配器产生（dms_separate_alpha=True）

驱逐决策由一个可学习的线性投影（adapter）从隐藏状态产生，对应Inference-Time Hyper-Scaling with KV Cache Compression（arXiv: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, eps=config.rms_norm_eps) self.dms_proj_alpha = nn.Linear(config.hidden_size, self.num_key_value_heads, bias=config.attention_bias) else: self.dms_proj_alpha_norm = None self.dms_proj_alpha = None

即dms_separate_alpha=True时，每个注意力层会额外挂载dms_proj_alpha_norm（RMSNorm）与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）还是按层（layer）

dms_alpha_per控制驱逐决策的共享粒度：

在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 Cache（DMSPagedCacheLayer）逐 token 生成；prefill 阶段则走dms_run_prefill_flex（dms/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）：

1. 蒸馏损失（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)

2. 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.0

ARCHITECTURE.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 close

max_steps优先取配置中的dms_final_step（qwen3_8b 配置为510），否则取training_args.max_steps（该配置为544，留有余量）。target即当前步应关闭（驱逐）的 gate 比例，每一步训练都会以此为基准计算 DMS 压缩损失，使模型在"先易后难"的调度下渐进学会更激进的压缩。

高级选项一：Chunked Prefill（分块预填充）

长序列 prefill 阶段峰值显存较高。DMS 支持把输入序列切成固定大小的块逐块处理，以降低峰值内存。启用方式（单位：token 数）：

Qwen3ForCausalLMDMS.from_pretrained(..., dms_chunked_prefill=4096)

底层由 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_prefill=4096，配合 32K 序列的 RULER 评测使用。

高级选项二：Cache Preallocation（缓存预分配）

分页 KV Cache 采用逐注意力层、可动态扩容的块表（block table）（类似 C++ 的std::vector），生成过程中按需增长。若提前知道最大上下文长度，可以预分配以避免运行期反复扩容：

Qwen3ForCausalLMDMS.from_pretrained(..., dms_preallocate_for_tokens=2048)

在 dms/cache_paged.py 的DMSPagedCacheLayer中可以看到与内存管理相关的默认参数：block_size=256、growth_factor=1.5、accommodate_min_initial_context_length=4096、max_context_length（由dms_preallocate_for_tokens决定）。块表扩容时按growth_factor比例一次性拼接新 block（expand_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_size（512 > 256 满足）。

将 DMS 适配到新模型家族（Retrofitting）

ARCHITECTURE.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 # 训练入口

模型特定代码需要完成四件事：

1. 扩展 config 类：继承基础模型的 config 并加入 DMS 参数，参照 configuration_qwen3_dms.py（Qwen3ConfigDMS(Qwen3Config)），其中包含dms_alpha_scale=100.0、dms_initial_alpha_offset=5.0、dms_window_size=512、dms_paged_attention_block_size=256、dms_cr=8、dms_disable_eviction=False、dms_separate_alpha=False、dms_alpha_per="head"、dms_tau=0.1、dms_compile_limit=72、dms_manual_inference_mode=False、dms_chunked_prefill=None、dms_preallocate_for_tokens=4096等默认值，并内置参数合法性断言。
2. 覆写注意力 forward，调用dms.core.prepare_attention_input与dms.attention.dms_attention（Qwen3 实现在 modeling_qwen3_dms.py 的Qwen3AttentionDMS.forward中，将 q/k/v 投影、QK 归一化、旋转位置编码、dms_proj_alpha等作为输入传给 core 函数）。
3. 为注意力层添加dms_proj_alpha与dms_proj_alpha_norm层（当dms_separate_alpha=True时，如上文代码所示）。
4. 在configs/下添加对应的 YAML 配置（参照 configs/qwen3_8b.yaml）。

核心 DMS 算子（prepare_attention_input、dms_attention、post_process_attention_output）都是模型无关的，模型特定代码只需提供 Q/K/V 投影与所需的各种归一化层作为输入即可复用。

添加新训练数据集

训练数据管线位于 dms/training/data.py，接入新数据集只需两步：

1. 为数据集定义filter_fn与extract_fn；
2. 创建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_fn=my_dataset_filter_fn, extract_fn=my_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: "auto"

"auto"会自动检测输出目录下最新的checkpoint-N目录；也可以指定显式路径：

hf_trainer: resume_from_checkpoint: outputs/qwen3_8b/checkpoint-300

恢复能够成立，得益于以下机制（见 dms/training/engine.py）：

• Hugging 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.sh）：

bash scripts/train.sh configs/qwen3_8b.yaml

scripts/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.sh）：

bash 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 Compression（NeurIPS 2025）进一步核对算法细节。整体而言，DMS 模块在 Model-Optimizer 仓库中提供了一个"模型无关核心算子 + 模型特定薄封装"的清晰分层：前者集中于 dms/（core、attention、cache、training），后者落在 models/qwen3，这种结构正是将其快速推广到其他模型家族的关键。

• 人工智能
• 大模型
• 模型优化
• 模型量化
• 模型压缩