在 DeepSpeed Inference V2 中接入新模型从 Parameter 到 Policy 的完整开发指南【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed在 DeepSpeed 的推理框架中deepspeed/inference/v2是一套围绕 ragged batching、Tensor Parallelism 与 KV Cache 块分配全新设计的高性能推理引擎。当需要接入一个尚未被官方支持的新模型例如新一代 Transformer 架构时最正统的扩展方式并不是去修改推理引擎本身而是按照model_implementations子系统的既定抽象为模型补齐“参数容器 模型实现 策略Policy”三个组件。本文以官方接入文档 AddingAModel.md 为主体骨架结合仓库中真实源码parameter_base.py、layer_container_base.py、inference_policy_base.py 与 inference_transformer_base.py做纵深解读并以内置的llama_v2真实实现为全程参照帮助你掌握如何为一个类 Transformer 模型开发三件套使其可直接被 V2 推理引擎加载并运行。阅读并动手完成后你将能够为任意类 Transformer 模型编写Parameter/LayerContainer含PARAM_MAPPING、编写继承DSTransformerModelBase的模型实现、编写InferenceV2PolicyContainerMap最终把模型接入 V2 引擎并享受其自动化的张量并行分片与参数融合优化。一、接入新模型的三大组件与代码落位官方文档开门见山地指出在 DeepSpeed Inference 中接入一个新模型需要开发三个相互关联的组件Containers容器描述模型包含哪些参数Model implementation模型实现描述模型应该如何被计算Policy策略负责把源模型的参数映射进容器并创建模型实现是三者之间的“组装层”。文档同时给出一个重要的前提假设如果你想接入的是一个“相对传统的 Transformer 风格模型”那么可以继承DSTransformerModelBase从而免费获得它提供的一整套工具方法。这里的“传统”通常意味着embedding → N 个标准 decoder/encoder 层self-attention MLP LayerNorm→ 输出 head 的堆叠结构而不是需要自定义稀疏注意力、非标准循环结构等特殊架构的模型。在当前的 model_implementations 目录 中这一“三件套”抽象已被多个真实模型采用每个模型一个子目录且均遵循container.pymodel.pypolicy.py的组织惯例llama_v2、mistral、falcon、phi、phi3、opt、qwen、qwen_v2、mixtral、qwen_v2_moe、exaone4、exaone4_5等。接入新模型时最直接的参考样板就是 llama_v2 这套实现。三个组件的核心类分别位于参数单元抽象parameter_base.py 中的ParameterBase参数容器抽象layer_container_base.py 中的LayerContainer策略与映射抽象inference_policy_base.py 中的InferenceV2Policy与ContainerMap模型基类inference_model_base.py 的DSInferenceModelBase以及 inference_transformer_base.py 的DSTransformerModelBase。小提示官方文档中的示例导入路径如from deepspeed.inference.module_implementations.parameter_base import ParameterBase是早期目录规划的写法当前仓库中这些模块已统一组织在deepspeed.inference.v2.model_implementations包下本文代码示例均以仓库当前实际路径为准。二、第一步用ParameterBase定义“融合参数”容器是“原模型参数”与“推理服务参数”之间的桥梁。但在定义容器之前必须先理解 DeepSpeed 推理中一个最小单元——Parameter的设计。官方文档概括了它的两个核心部分dependencies依赖来自源 checkpoint 的原始参数finalize 方法当所有依赖齐备后如何把它们转换/融合成真正的推理参数。其动机非常典型以 Llama 系列为例原始 checkpoint 中query、key、value三个投影是独立存储的但在推理时把三者沿dim0融合成一个更大的 QKV 投影通常能获得更高吞吐。这个“融合”动作正是用一个自定义Parameter来描述的文档给出的UnfusedQKVParameter示例经过导入路径修正后如下from deepspeed.inference.v2.model_implementations.parameter_base import ParameterBase class UnfusedQKVParameter(ParameterBase): query: torch.Tensor key: torch.Tensor value: torch.Tensor def finalize(self) - torch.Tensor: fused_param torch.cat([self.query, self.key, self.value], dim0) return self.inference_model.transform_qkv_param(fused_param)逐段拆解这个实现继承ParameterBase。基类会自动判断依赖何时被满足并在完成后把结果写回父级LayerContainer的对应槽位。需要特别注意的是ParameterBase并非普通类它由ParameterMetaclass驱动__call__阶段会为每个实例初始化dest_paramNone并把每个依赖的存储槽位初始化为None或空ParametrizedList。类上的类型注解 依赖声明。元类在__new__阶段解析注解凡是注解类型是torch.Tensor或ParametrizedList子类的字段都会被识别为“依赖”并自动替换成 property读 getter / 写 setter。因此类体里写query: torch.Tensor等价于声明“本参数依赖源 checkpoint 中一个叫query的张量”。由于原 Llama 模型天然是分离的 q/k/v融合参数就依次声明三个依赖。finalize方法。每个依赖一旦被 setter 赋值都会调用complete_component()把completed_components计数加一见 parameter_base.py当计数达到元类统计的n_dependencies时finalize被自动触发。它返回最终参数之后该ParameterBase实例会被父容器的finalization_callback替换为真正的张量并销毁你再也拿不到半成品。访问依赖与调用模型工具。finalize内可以通过self.{依赖名}如self.query拿到已填充的张量同时任何ParameterBase实例都持有一个self.inference_model弱引用——它提供了如何分片与转换参数的上下文。这里调用self.inference_model.transform_qkv_param正是利用了DSTransformerBase提供的能力该方法会按tp_rank/tp_size与 head 配置对 QKV 参数做张量并行分片随后交给底层 linear 模块实现执行量化等额外的形状转换或优化。从源码看transform_qkv_param先调用shard_qkv_param(...)完成按 head 的分片再委托self.qkv.transform_param(param)。2.1 内置的公共参数模板common_parameters由于 QKV/MLP/输出投影的融合模式在各类 Transformer 中高度重复仓库在 common_parameters 下预置了大量可直接复用且全部与DSTransformerBase兼容的实现接入新模型时优先查看这里文件提供的典型参数模板qkv_parameters.pyFusedQKVParameter原生已融合直接复制、UnfusedQKVParameterq/k/v 分离cat 融合、MegatronQKVParameter与GQAMegatronQKVParameter处理 Megatron[n_heads, 3, ...]头序布局及 GQA 分组布局转换attn_output_parameters.py注意力输出投影参数即文档示例中的AttentionOutputParametermlp_parameters.pyMLP 上下投影与门控参数模板norm_parameters.pyLayerNorm/RMSNorm 等归一化参数模板embedding_parameters.py、unembed_parameters.pyembedding / unembeddingLM Head参数模板invfreq_parameters.pyRoPEinv_freq等非张量数据的处理moe_parameters.pyMoE 路由/专家参数模板以真实的UnfusedQKVParameterqkv_parameters.py为例它与官方文档语义完全一致torch.cat([self.q_params, self.k_params, self.v_params], dim0)之后接self.inference_model.transform_qkv_param(fused_param)。2.2 变长依赖ParametrizedListParameterBase的依赖并不限于单一torch.Tensor。某些参数的数量取决于模型配置而非架构——典型例子是 MoE 层专家数量随模型规模变化逐个写expert_0、expert_1… 既笨拙又不可扩展。为此 parameter_base.py 提供了ParametrizedListclass MyParametrizedList(ParametrizedList): count_attr: str my_list_count其中my_list_count必须是推理模型实例上可访问的属性即self.inference_model.my_list_count它给出该列表的长度对列表使用整数索引experts[8]访问比命名成expert_8自然得多。该列表内部维护set_params计数当所有槽位都被填充完毕才回调父参数的complete_component()。若想省去手写子类可以直接使用工厂函数ParamList(attr)它会返回一个设置了指定count_attr的ParametrizedList子类。注意ParametrizedList不是一个普通的listtorch.cat(param_list)无法直接使用需要用torch.cat(tuple(param_list))包一层。三、第二步用LayerContainer组装参数并建立映射参数单元定义好之后下一步是把它们组合进LayerContainer。官方文档明确了两种容器分工Transformer 容器对应模型中的单层Transformer 的参数包括 FFN 全连接投影、QKV 投影、注意力输出投影、归一化等非 Transformer 容器存放其他一切——典型如 embedding 与 unembeddingLM Head参数。从源码看DSInferenceModelBase用_transformer每个 layer 一个容器的列表和_non_transformer一个容器来区分这两类参数二者在模型实例化时都是None要等 Policy 调用set_parameters(...)之后才会被填充。文档中简化版 Llama 的示例仅含 QKV 与注意力输出投影修正了语法后如下from deepspeed.inference.v2.model_implementations.layer_container_base import LayerContainer class ExampleContainer(LayerContainer): qkvw: UnfusedQKVParameter attn_o: AttentionOutputParameter PARAM_MAPPING { self_attn.q_proj.weight: qkvw.query, self_attn.k_proj.weight: qkvw.key, self_attn.v_proj.weight: qkvw.value, self_attn.o_proj.weight: attn_o.params, }同样有两个关键要素参数类型注解。每个注解对应模型实现中可以使用的“参数组”。在模型实现的 forward 中直接写container.qkvw就能拿到已经完成融合、分片与变换的 QKV 参数它实际上已被替换为InferenceParameter张量。之所以能做到这一点是因为LayerContainer同样由元类驱动LayerMetaclass它在__new__阶段收集本类与所有基类的注解并合并因此可以通过继承复用公共参数组为每个ParameterBase注解生成实例与统一的finalization_callback。PARAM_MAPPING字典。这是“源 checkpoint 参数名 → 容器内依赖”的显式路由表。它会被 Policy 利用在加载 checkpoint 时自动填充依赖因此文档强调它是显式映射写错、漏写都会被直接拦截。事实上LayerMetaclass在构建阶段会做大量静态校验见 layer_container_base.py映射目标必须写成参数名.依赖名且参数名必须是本容器注解中真实存在的参数依赖名必须真实存在于该参数的注解中同一个依赖不允许被多条映射规则命中重复映射直接抛ValueError容器中所有ParameterBase的所有依赖必须被映射覆盖完全否则会以“以下依赖未被映射”的明确错误拒绝若一个依赖目标是ParametrizedList则源名必须带且只能带一个通配符*如model.layers.*....元类会据此生成把源名中的数字索引解析到列表下标的路由 helper禁止把ParametrizedList与普通Tensor混在同一个映射规则里。在 checkpoint 加载期间容器会通过set_dependency(dep_name, dep_value)layer_container_base.py接收去除了前缀的依赖名先尝试精确匹配PARAM_MAPPING再尝试通配符*会被替换为正则.*与plist_helpers的列表索引解析。此外容器还提供direct_injection(name, tensor)用于直接注入张量以及两个非常有用的只读状态属性is_populated所有参数都已被 checkpoint 引擎填充与is_initialized在填充基础上进一步要求全部落在正确的加速设备上参数类型必须是InferenceParameter或显式None。文档建议当Transformer 容器与非 Transformer 容器都写好之后就可以进入模型实现环节了。四、第三步编写继承DSTransformerModelBase的模型实现DSTransformerModelBaseinference_transformer_base.py承担了绝大部分“分片与变换参数”的机械工作在其__init__中它按固定顺序调用make_norm_layer()、make_qkv_layer()、make_attn_layer()、make_attn_out_layer()、make_mlp_1_layer()、make_mlp_2_layer()、make_embedding_layer()与make_unembedding_layer()为模型组装出 embedding、QKV、self-attention、注意力输出投影、两层 MLP、归一化与 LM Head 的整套模块并借助modules.heuristics.instantiate_*依据引擎配置ragged state manager 的max_ragged_batch_size等选择具体 kernel 实现。即便如此接入者仍需要亲自完成官方文档点名的四个关键任务任务 1根据模型配置定义抽象属性DSTransformerModelBase用大量property abstractmethod把模型“尺寸”与“结构”参数化接入者必须全部实现。从源码归纳至少包括尺寸类num_layers层数、model_dimembedding 与残差维度、vocab_size含 padding 的词典大小、head_size每个注意力头维度、n_headsquery 头数、intermediate_dim未分片的中间投影维度对门控激活指第二个 MLP 层的输入维度结构类activation_dtype、mlp_activation_fnMLP 激活函数决定是否为门控、norm_type、positional_embedding_type位置编码类型与positional_embedding_config通常为 RoPE 的RotateHalfConfig。基类还基于这些属性派生了便捷工具n_heads_q_local/n_heads_kv_local给出按tp_rank/tp_size分片后的本地头数gated_mlp判断是否使用门控激活。需要特别留意n_heads_kv默认实现采用 MHA 形式return self.n_headsGQA 或 MQA 模型必须重写该属性。任务 2配置 embedding / unembedding 模块并实现其 forwardmake_embedding_layer()默认只做 dtype 转换不沿 channel 维度分片源码注释说明在支持非连续 all-gather 之前不会分片 embedding 参数make_unembedding_layer()假设 LM Head 之前存在一次归一化并对vocab维做分片sharded_unembed_dimtp_size 1时还会预分配通信用的 logits 缓冲。如果模型不符合这些默认假设例如无 pre-norm、embedding 需要特殊处理就应重写对应make_*方法并配套编写 embedding 与 unembedding 的 forward 计算。任务 3配置注意力与 KV Cache 行为基类为自回归密集注意力模式提供了一整套可覆盖的钩子make_attn_layer()用softmax_scale 1.0 / head_size**0.5、本地头数与位置编码信息构建DSSelfAttentionConfigkv_cache_config()返回KVCacheConfig其中cache_shape (num_layers, n_heads_kv_local, head_size)max_blocks_per_allocation_group由max_sequence_length与kv_block_size推算get_kv_requirements(...)/maybe_allocate_kv(...)负责估算并在需要时向state_manager申请 KV 块prepare_batch(...)在每次 forward 前构造 attention 相关的批元数据如build_atoms。若模型的注意力模式不是标准的自回归密集注意力例如有滑动窗口、稀疏模式文档与源码均明确要求重写这些默认实现。任务 4编写 Transformer 层的 forward最后是模型实现真正“计算如何发生”的部分以容器为输入按 归一化 → 注意力自注意力 输出投影 残差→ MLP两层投影 残差→ 输出 LM 计算 的顺序编写单层 forward以及全模型的逐层循环逻辑。注意 inference_model_base.py 中forward的接口约定它需要能被构图graphable因此不应依赖 Python 控制流编写逐 token 的逻辑。另外仓库还在同一文件中提供了 DSMoETransformerModelBaseMoE 版基类要求补充n_experts、n_top_k、normalize_expert_scores属性并用make_moe_layer()替代普通 MLP同时区分transform_moe_mlp_1_param与transform_moe_mlp_2_param因为同一 DSModule 同时持有两块专家参数时无法仅凭形状推断应做哪种变换。mixtral、qwen_v2_moe目录就是这一基类的最佳实践。五、第四步编写 Policy 与ContainerMap官方文档对 Policy 的定位是“组合层”InferenceV2Policy是直接传给推理引擎的对象负责把模型实现与容器组合成端到端解决方案。仓库中该抽象定义在 inference_policy_base.py其构造函数要求checkpoint_engine与inf_checkpoint_path二选一同时为None或同时给出都会抛ValueError分别对应两种建模型机制checkpoint_engine通用机制checkpoint 引擎逐个遍历模型参数交给 Policy再由模型实现完成分片/变换inf_checkpoint_path重新加载此前由 DeepSpeed 序列化保存的推理模型这种 checkpoint 不应跨模型后端配置混用源码以 TODO 注明该限制尚待代码强制。Policy 需要实现两个抽象方法5.1instantiate_model创建模型实例第一个抽象方法是创建前面定义好的模型。文档指出一般情况下只需直接调用模型构造函数并把引擎配置、张量并行通信对象、自定义模型配置三个参数传进去即可。这与仓库中llama_v2的实现完全吻合policy.pydef instantiate_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) - Llama2InferenceModel: return Llama2InferenceModel(configself._model_config, engine_configengine_config, base_mp_groupmp_group)5.2build_container_map建立 checkpoint 前缀 → 容器的路由第二个抽象方法是定义 checkpoint 参数如何映射到每个容器。上一节提过LayerContainer自己就能处理“checkpoint 参数 → 容器内依赖”的内部路由PARAM_MAPPING但为了找到“该参数到底属于哪一层、哪个容器”还需要ContainerMap这一层抽象。ContainerMapinference_policy_base.py通过把“checkpoint 前缀字符串”分类到“它所对应的容器类型”来完成映射核心 API 有三个set_transformer_params(prefixes, containers)注册 transformer 前缀与其对应的每层一个的容器列表入参必须是一个 list容器数与层数一一对应set_non_transformer_params(container)注册唯一的非 Transformer 容器set_unmapped_params(prefixes)登记“剩余”的、应当被忽略的前缀比如属于 RoPE 缓存等运行时数据、不属于权重的张量。在加载阶段ContainerMap.map_param(name, parameter)会对每个参数名依次做三件事先检查是否命中“忽略前缀”再检查是否命中 transformer 前缀——命中则剥离前缀并解析层号把剩余名字路由到第layer_idx个 transformer 容器的set_dependency否则尝试交给_non_transformer_params若连它也抛ValueError则包装成信息更明确的报错“找不到该参数对应的容器请复查 Containers/ContainerMap”。完成全部填充后调用validate()逐一校验非 Transformer 容器与每个 transformer 容器都已is_initialized。最简洁的build_container_map通常是遍历某个 PyTorch 模型的named_parameters()或遍历一个 checkpoint 的 state dict 后依前缀归类。官方文档示例中set_transformer_params(model.layers, transformer_containers)正是把model.layers前缀下的参数路由给每一层容器。下面是llama_v2的真实实现可作为模板policy.pydef build_container_map(self) - ContainerMap: map ContainerMap() transformer_containers [Llama2TransformerContainer(self.model) for _ in range(self.model.num_layers)] map.set_transformer_params([model.layers], transformer_containers) map.set_non_transformer_params(Llama2NonTransformerContainer(self.model)) map.set_unmapped_params( [fmodel.layers.{i}.self_attn.rotary_emb.inv_freq for i in range(self.model.num_layers)]) return map5.3 引擎如何驱动 Policy当引擎需要建模型时会调用InferenceV2Policy.build_model(engine_config, mp_group)inference_policy_base.py其内部流程为self.model self.instantiate_model(engine_config, mp_group)创建模型实现self.populate_model_parameters()调build_container_map()然后按上文两种机制之一填充参数若是 checkpoint 引擎则遍历checkpoint_engine.parameters()逐个map_param随后用 flat_model_helpers.py 的flatten_inference_model把容器参数拍平为一个连续 buffer 与其元数据若是序列化路径则按tp_rank/tp_size定位make_param_filename/make_metadata_filename生成的文件并restore_inference_modelcontainer_map.validate()校验完整性self.model.set_parameters(...)把 transformer / non-transformer 容器与拍平后的参数 buffer 交给模型。另外值得留意的是PolicyMeta元类会自动把每个非基类的 Policy 子类注册进模块级POLICIES字典——引擎侧据此按名字发现可用的模型策略。六、推荐阅读路径与调试建议动手实现时建议按以下顺序在仓库中“对照着抄”先读一套最完整的官方样例llama_v2/container.py看Llama2TransformerContainer/Llama2NonTransformerContainer如何写注解与PARAM_MAPPING、llama_v2/model.py看如何实现抽象属性与逐层 forward、llama_v2/policy.py看三行核心方法相似架构优先拷贝参数模板GQA 类Llama 2/3、Qwen2、Mistral 等从 common_parameters/qkv_parameters.py 选取UnfusedQKVParameter/GQAMegatronQKVParameterMegatron 风格布局参考MegatronQKVParameterMoE 架构参考 mixtral 与 qwen_v2_moe写完后自查容器完整性由于LayerMetaclass在类定义期就会强制“依赖全覆盖 无重复映射 目标存在”绝大多数映射笔误会在实例化容器的那一刻直接暴露加载期若报“Cannot find container for {name}”说明PARAM_MAPPING与ContainerMap的源名前缀不匹配——对照上述 llama 参考仔细核对model.layers.{i}.*之类的层前缀即可设备相关报错查is_initialized若容器校验提示参数不在正确设备检查你的finalize链路里参数是否都经过transform_*_param如transform_qkv_param、transform_mlp_1_param并落在InferenceParameter上。接入文档本身保存在 AddingAModel.md与各模型实现源码同目录非常适合作为“写代码时的案头手册”。整体而言这套“Parameter 声明融合 → LayerContainer 路由参数 → DSTransformerModelBase 计算 → Policy 组装”的分层设计把「checkpoint 参数布局的差异」和「计算图/分片/量化等推理细节」彻底解耦——新增模型时大部分工作集中在描述参数映射关系上而分片、融合、KV 缓存管理等繁琐部分则由框架基类自动接管。【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 在 DeepSpeed Inference V2 中接入新模型：从 Parameter 到 Policy 的完整开发指南
在 DeepSpeed Inference V2 中接入新模型从 Parameter 到 Policy 的完整开发指南【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed在 DeepSpeed 的推理框架中deepspeed/inference/v2是一套围绕 ragged batching、Tensor Parallelism 与 KV Cache 块分配全新设计的高性能推理引擎。当需要接入一个尚未被官方支持的新模型例如新一代 Transformer 架构时最正统的扩展方式并不是去修改推理引擎本身而是按照model_implementations子系统的既定抽象为模型补齐“参数容器 模型实现 策略Policy”三个组件。本文以官方接入文档 AddingAModel.md 为主体骨架结合仓库中真实源码parameter_base.py、layer_container_base.py、inference_policy_base.py 与 inference_transformer_base.py做纵深解读并以内置的llama_v2真实实现为全程参照帮助你掌握如何为一个类 Transformer 模型开发三件套使其可直接被 V2 推理引擎加载并运行。阅读并动手完成后你将能够为任意类 Transformer 模型编写Parameter/LayerContainer含PARAM_MAPPING、编写继承DSTransformerModelBase的模型实现、编写InferenceV2PolicyContainerMap最终把模型接入 V2 引擎并享受其自动化的张量并行分片与参数融合优化。一、接入新模型的三大组件与代码落位官方文档开门见山地指出在 DeepSpeed Inference 中接入一个新模型需要开发三个相互关联的组件Containers容器描述模型包含哪些参数Model implementation模型实现描述模型应该如何被计算Policy策略负责把源模型的参数映射进容器并创建模型实现是三者之间的“组装层”。文档同时给出一个重要的前提假设如果你想接入的是一个“相对传统的 Transformer 风格模型”那么可以继承DSTransformerModelBase从而免费获得它提供的一整套工具方法。这里的“传统”通常意味着embedding → N 个标准 decoder/encoder 层self-attention MLP LayerNorm→ 输出 head 的堆叠结构而不是需要自定义稀疏注意力、非标准循环结构等特殊架构的模型。在当前的 model_implementations 目录 中这一“三件套”抽象已被多个真实模型采用每个模型一个子目录且均遵循container.pymodel.pypolicy.py的组织惯例llama_v2、mistral、falcon、phi、phi3、opt、qwen、qwen_v2、mixtral、qwen_v2_moe、exaone4、exaone4_5等。接入新模型时最直接的参考样板就是 llama_v2 这套实现。三个组件的核心类分别位于参数单元抽象parameter_base.py 中的ParameterBase参数容器抽象layer_container_base.py 中的LayerContainer策略与映射抽象inference_policy_base.py 中的InferenceV2Policy与ContainerMap模型基类inference_model_base.py 的DSInferenceModelBase以及 inference_transformer_base.py 的DSTransformerModelBase。小提示官方文档中的示例导入路径如from deepspeed.inference.module_implementations.parameter_base import ParameterBase是早期目录规划的写法当前仓库中这些模块已统一组织在deepspeed.inference.v2.model_implementations包下本文代码示例均以仓库当前实际路径为准。二、第一步用ParameterBase定义“融合参数”容器是“原模型参数”与“推理服务参数”之间的桥梁。但在定义容器之前必须先理解 DeepSpeed 推理中一个最小单元——Parameter的设计。官方文档概括了它的两个核心部分dependencies依赖来自源 checkpoint 的原始参数finalize 方法当所有依赖齐备后如何把它们转换/融合成真正的推理参数。其动机非常典型以 Llama 系列为例原始 checkpoint 中query、key、value三个投影是独立存储的但在推理时把三者沿dim0融合成一个更大的 QKV 投影通常能获得更高吞吐。这个“融合”动作正是用一个自定义Parameter来描述的文档给出的UnfusedQKVParameter示例经过导入路径修正后如下from deepspeed.inference.v2.model_implementations.parameter_base import ParameterBase class UnfusedQKVParameter(ParameterBase): query: torch.Tensor key: torch.Tensor value: torch.Tensor def finalize(self) - torch.Tensor: fused_param torch.cat([self.query, self.key, self.value], dim0) return self.inference_model.transform_qkv_param(fused_param)逐段拆解这个实现继承ParameterBase。基类会自动判断依赖何时被满足并在完成后把结果写回父级LayerContainer的对应槽位。需要特别注意的是ParameterBase并非普通类它由ParameterMetaclass驱动__call__阶段会为每个实例初始化dest_paramNone并把每个依赖的存储槽位初始化为None或空ParametrizedList。类上的类型注解 依赖声明。元类在__new__阶段解析注解凡是注解类型是torch.Tensor或ParametrizedList子类的字段都会被识别为“依赖”并自动替换成 property读 getter / 写 setter。因此类体里写query: torch.Tensor等价于声明“本参数依赖源 checkpoint 中一个叫query的张量”。由于原 Llama 模型天然是分离的 q/k/v融合参数就依次声明三个依赖。finalize方法。每个依赖一旦被 setter 赋值都会调用complete_component()把completed_components计数加一见 parameter_base.py当计数达到元类统计的n_dependencies时finalize被自动触发。它返回最终参数之后该ParameterBase实例会被父容器的finalization_callback替换为真正的张量并销毁你再也拿不到半成品。访问依赖与调用模型工具。finalize内可以通过self.{依赖名}如self.query拿到已填充的张量同时任何ParameterBase实例都持有一个self.inference_model弱引用——它提供了如何分片与转换参数的上下文。这里调用self.inference_model.transform_qkv_param正是利用了DSTransformerBase提供的能力该方法会按tp_rank/tp_size与 head 配置对 QKV 参数做张量并行分片随后交给底层 linear 模块实现执行量化等额外的形状转换或优化。从源码看transform_qkv_param先调用shard_qkv_param(...)完成按 head 的分片再委托self.qkv.transform_param(param)。2.1 内置的公共参数模板common_parameters由于 QKV/MLP/输出投影的融合模式在各类 Transformer 中高度重复仓库在 common_parameters 下预置了大量可直接复用且全部与DSTransformerBase兼容的实现接入新模型时优先查看这里文件提供的典型参数模板qkv_parameters.pyFusedQKVParameter原生已融合直接复制、UnfusedQKVParameterq/k/v 分离cat 融合、MegatronQKVParameter与GQAMegatronQKVParameter处理 Megatron[n_heads, 3, ...]头序布局及 GQA 分组布局转换attn_output_parameters.py注意力输出投影参数即文档示例中的AttentionOutputParametermlp_parameters.pyMLP 上下投影与门控参数模板norm_parameters.pyLayerNorm/RMSNorm 等归一化参数模板embedding_parameters.py、unembed_parameters.pyembedding / unembeddingLM Head参数模板invfreq_parameters.pyRoPEinv_freq等非张量数据的处理moe_parameters.pyMoE 路由/专家参数模板以真实的UnfusedQKVParameterqkv_parameters.py为例它与官方文档语义完全一致torch.cat([self.q_params, self.k_params, self.v_params], dim0)之后接self.inference_model.transform_qkv_param(fused_param)。2.2 变长依赖ParametrizedListParameterBase的依赖并不限于单一torch.Tensor。某些参数的数量取决于模型配置而非架构——典型例子是 MoE 层专家数量随模型规模变化逐个写expert_0、expert_1… 既笨拙又不可扩展。为此 parameter_base.py 提供了ParametrizedListclass MyParametrizedList(ParametrizedList): count_attr: str my_list_count其中my_list_count必须是推理模型实例上可访问的属性即self.inference_model.my_list_count它给出该列表的长度对列表使用整数索引experts[8]访问比命名成expert_8自然得多。该列表内部维护set_params计数当所有槽位都被填充完毕才回调父参数的complete_component()。若想省去手写子类可以直接使用工厂函数ParamList(attr)它会返回一个设置了指定count_attr的ParametrizedList子类。注意ParametrizedList不是一个普通的listtorch.cat(param_list)无法直接使用需要用torch.cat(tuple(param_list))包一层。三、第二步用LayerContainer组装参数并建立映射参数单元定义好之后下一步是把它们组合进LayerContainer。官方文档明确了两种容器分工Transformer 容器对应模型中的单层Transformer 的参数包括 FFN 全连接投影、QKV 投影、注意力输出投影、归一化等非 Transformer 容器存放其他一切——典型如 embedding 与 unembeddingLM Head参数。从源码看DSInferenceModelBase用_transformer每个 layer 一个容器的列表和_non_transformer一个容器来区分这两类参数二者在模型实例化时都是None要等 Policy 调用set_parameters(...)之后才会被填充。文档中简化版 Llama 的示例仅含 QKV 与注意力输出投影修正了语法后如下from deepspeed.inference.v2.model_implementations.layer_container_base import LayerContainer class ExampleContainer(LayerContainer): qkvw: UnfusedQKVParameter attn_o: AttentionOutputParameter PARAM_MAPPING { self_attn.q_proj.weight: qkvw.query, self_attn.k_proj.weight: qkvw.key, self_attn.v_proj.weight: qkvw.value, self_attn.o_proj.weight: attn_o.params, }同样有两个关键要素参数类型注解。每个注解对应模型实现中可以使用的“参数组”。在模型实现的 forward 中直接写container.qkvw就能拿到已经完成融合、分片与变换的 QKV 参数它实际上已被替换为InferenceParameter张量。之所以能做到这一点是因为LayerContainer同样由元类驱动LayerMetaclass它在__new__阶段收集本类与所有基类的注解并合并因此可以通过继承复用公共参数组为每个ParameterBase注解生成实例与统一的finalization_callback。PARAM_MAPPING字典。这是“源 checkpoint 参数名 → 容器内依赖”的显式路由表。它会被 Policy 利用在加载 checkpoint 时自动填充依赖因此文档强调它是显式映射写错、漏写都会被直接拦截。事实上LayerMetaclass在构建阶段会做大量静态校验见 layer_container_base.py映射目标必须写成参数名.依赖名且参数名必须是本容器注解中真实存在的参数依赖名必须真实存在于该参数的注解中同一个依赖不允许被多条映射规则命中重复映射直接抛ValueError容器中所有ParameterBase的所有依赖必须被映射覆盖完全否则会以“以下依赖未被映射”的明确错误拒绝若一个依赖目标是ParametrizedList则源名必须带且只能带一个通配符*如model.layers.*....元类会据此生成把源名中的数字索引解析到列表下标的路由 helper禁止把ParametrizedList与普通Tensor混在同一个映射规则里。在 checkpoint 加载期间容器会通过set_dependency(dep_name, dep_value)layer_container_base.py接收去除了前缀的依赖名先尝试精确匹配PARAM_MAPPING再尝试通配符*会被替换为正则.*与plist_helpers的列表索引解析。此外容器还提供direct_injection(name, tensor)用于直接注入张量以及两个非常有用的只读状态属性is_populated所有参数都已被 checkpoint 引擎填充与is_initialized在填充基础上进一步要求全部落在正确的加速设备上参数类型必须是InferenceParameter或显式None。文档建议当Transformer 容器与非 Transformer 容器都写好之后就可以进入模型实现环节了。四、第三步编写继承DSTransformerModelBase的模型实现DSTransformerModelBaseinference_transformer_base.py承担了绝大部分“分片与变换参数”的机械工作在其__init__中它按固定顺序调用make_norm_layer()、make_qkv_layer()、make_attn_layer()、make_attn_out_layer()、make_mlp_1_layer()、make_mlp_2_layer()、make_embedding_layer()与make_unembedding_layer()为模型组装出 embedding、QKV、self-attention、注意力输出投影、两层 MLP、归一化与 LM Head 的整套模块并借助modules.heuristics.instantiate_*依据引擎配置ragged state manager 的max_ragged_batch_size等选择具体 kernel 实现。即便如此接入者仍需要亲自完成官方文档点名的四个关键任务任务 1根据模型配置定义抽象属性DSTransformerModelBase用大量property abstractmethod把模型“尺寸”与“结构”参数化接入者必须全部实现。从源码归纳至少包括尺寸类num_layers层数、model_dimembedding 与残差维度、vocab_size含 padding 的词典大小、head_size每个注意力头维度、n_headsquery 头数、intermediate_dim未分片的中间投影维度对门控激活指第二个 MLP 层的输入维度结构类activation_dtype、mlp_activation_fnMLP 激活函数决定是否为门控、norm_type、positional_embedding_type位置编码类型与positional_embedding_config通常为 RoPE 的RotateHalfConfig。基类还基于这些属性派生了便捷工具n_heads_q_local/n_heads_kv_local给出按tp_rank/tp_size分片后的本地头数gated_mlp判断是否使用门控激活。需要特别留意n_heads_kv默认实现采用 MHA 形式return self.n_headsGQA 或 MQA 模型必须重写该属性。任务 2配置 embedding / unembedding 模块并实现其 forwardmake_embedding_layer()默认只做 dtype 转换不沿 channel 维度分片源码注释说明在支持非连续 all-gather 之前不会分片 embedding 参数make_unembedding_layer()假设 LM Head 之前存在一次归一化并对vocab维做分片sharded_unembed_dimtp_size 1时还会预分配通信用的 logits 缓冲。如果模型不符合这些默认假设例如无 pre-norm、embedding 需要特殊处理就应重写对应make_*方法并配套编写 embedding 与 unembedding 的 forward 计算。任务 3配置注意力与 KV Cache 行为基类为自回归密集注意力模式提供了一整套可覆盖的钩子make_attn_layer()用softmax_scale 1.0 / head_size**0.5、本地头数与位置编码信息构建DSSelfAttentionConfigkv_cache_config()返回KVCacheConfig其中cache_shape (num_layers, n_heads_kv_local, head_size)max_blocks_per_allocation_group由max_sequence_length与kv_block_size推算get_kv_requirements(...)/maybe_allocate_kv(...)负责估算并在需要时向state_manager申请 KV 块prepare_batch(...)在每次 forward 前构造 attention 相关的批元数据如build_atoms。若模型的注意力模式不是标准的自回归密集注意力例如有滑动窗口、稀疏模式文档与源码均明确要求重写这些默认实现。任务 4编写 Transformer 层的 forward最后是模型实现真正“计算如何发生”的部分以容器为输入按 归一化 → 注意力自注意力 输出投影 残差→ MLP两层投影 残差→ 输出 LM 计算 的顺序编写单层 forward以及全模型的逐层循环逻辑。注意 inference_model_base.py 中forward的接口约定它需要能被构图graphable因此不应依赖 Python 控制流编写逐 token 的逻辑。另外仓库还在同一文件中提供了 DSMoETransformerModelBaseMoE 版基类要求补充n_experts、n_top_k、normalize_expert_scores属性并用make_moe_layer()替代普通 MLP同时区分transform_moe_mlp_1_param与transform_moe_mlp_2_param因为同一 DSModule 同时持有两块专家参数时无法仅凭形状推断应做哪种变换。mixtral、qwen_v2_moe目录就是这一基类的最佳实践。五、第四步编写 Policy 与ContainerMap官方文档对 Policy 的定位是“组合层”InferenceV2Policy是直接传给推理引擎的对象负责把模型实现与容器组合成端到端解决方案。仓库中该抽象定义在 inference_policy_base.py其构造函数要求checkpoint_engine与inf_checkpoint_path二选一同时为None或同时给出都会抛ValueError分别对应两种建模型机制checkpoint_engine通用机制checkpoint 引擎逐个遍历模型参数交给 Policy再由模型实现完成分片/变换inf_checkpoint_path重新加载此前由 DeepSpeed 序列化保存的推理模型这种 checkpoint 不应跨模型后端配置混用源码以 TODO 注明该限制尚待代码强制。Policy 需要实现两个抽象方法5.1instantiate_model创建模型实例第一个抽象方法是创建前面定义好的模型。文档指出一般情况下只需直接调用模型构造函数并把引擎配置、张量并行通信对象、自定义模型配置三个参数传进去即可。这与仓库中llama_v2的实现完全吻合policy.pydef instantiate_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) - Llama2InferenceModel: return Llama2InferenceModel(configself._model_config, engine_configengine_config, base_mp_groupmp_group)5.2build_container_map建立 checkpoint 前缀 → 容器的路由第二个抽象方法是定义 checkpoint 参数如何映射到每个容器。上一节提过LayerContainer自己就能处理“checkpoint 参数 → 容器内依赖”的内部路由PARAM_MAPPING但为了找到“该参数到底属于哪一层、哪个容器”还需要ContainerMap这一层抽象。ContainerMapinference_policy_base.py通过把“checkpoint 前缀字符串”分类到“它所对应的容器类型”来完成映射核心 API 有三个set_transformer_params(prefixes, containers)注册 transformer 前缀与其对应的每层一个的容器列表入参必须是一个 list容器数与层数一一对应set_non_transformer_params(container)注册唯一的非 Transformer 容器set_unmapped_params(prefixes)登记“剩余”的、应当被忽略的前缀比如属于 RoPE 缓存等运行时数据、不属于权重的张量。在加载阶段ContainerMap.map_param(name, parameter)会对每个参数名依次做三件事先检查是否命中“忽略前缀”再检查是否命中 transformer 前缀——命中则剥离前缀并解析层号把剩余名字路由到第layer_idx个 transformer 容器的set_dependency否则尝试交给_non_transformer_params若连它也抛ValueError则包装成信息更明确的报错“找不到该参数对应的容器请复查 Containers/ContainerMap”。完成全部填充后调用validate()逐一校验非 Transformer 容器与每个 transformer 容器都已is_initialized。最简洁的build_container_map通常是遍历某个 PyTorch 模型的named_parameters()或遍历一个 checkpoint 的 state dict 后依前缀归类。官方文档示例中set_transformer_params(model.layers, transformer_containers)正是把model.layers前缀下的参数路由给每一层容器。下面是llama_v2的真实实现可作为模板policy.pydef build_container_map(self) - ContainerMap: map ContainerMap() transformer_containers [Llama2TransformerContainer(self.model) for _ in range(self.model.num_layers)] map.set_transformer_params([model.layers], transformer_containers) map.set_non_transformer_params(Llama2NonTransformerContainer(self.model)) map.set_unmapped_params( [fmodel.layers.{i}.self_attn.rotary_emb.inv_freq for i in range(self.model.num_layers)]) return map5.3 引擎如何驱动 Policy当引擎需要建模型时会调用InferenceV2Policy.build_model(engine_config, mp_group)inference_policy_base.py其内部流程为self.model self.instantiate_model(engine_config, mp_group)创建模型实现self.populate_model_parameters()调build_container_map()然后按上文两种机制之一填充参数若是 checkpoint 引擎则遍历checkpoint_engine.parameters()逐个map_param随后用 flat_model_helpers.py 的flatten_inference_model把容器参数拍平为一个连续 buffer 与其元数据若是序列化路径则按tp_rank/tp_size定位make_param_filename/make_metadata_filename生成的文件并restore_inference_modelcontainer_map.validate()校验完整性self.model.set_parameters(...)把 transformer / non-transformer 容器与拍平后的参数 buffer 交给模型。另外值得留意的是PolicyMeta元类会自动把每个非基类的 Policy 子类注册进模块级POLICIES字典——引擎侧据此按名字发现可用的模型策略。六、推荐阅读路径与调试建议动手实现时建议按以下顺序在仓库中“对照着抄”先读一套最完整的官方样例llama_v2/container.py看Llama2TransformerContainer/Llama2NonTransformerContainer如何写注解与PARAM_MAPPING、llama_v2/model.py看如何实现抽象属性与逐层 forward、llama_v2/policy.py看三行核心方法相似架构优先拷贝参数模板GQA 类Llama 2/3、Qwen2、Mistral 等从 common_parameters/qkv_parameters.py 选取UnfusedQKVParameter/GQAMegatronQKVParameterMegatron 风格布局参考MegatronQKVParameterMoE 架构参考 mixtral 与 qwen_v2_moe写完后自查容器完整性由于LayerMetaclass在类定义期就会强制“依赖全覆盖 无重复映射 目标存在”绝大多数映射笔误会在实例化容器的那一刻直接暴露加载期若报“Cannot find container for {name}”说明PARAM_MAPPING与ContainerMap的源名前缀不匹配——对照上述 llama 参考仔细核对model.layers.{i}.*之类的层前缀即可设备相关报错查is_initialized若容器校验提示参数不在正确设备检查你的finalize链路里参数是否都经过transform_*_param如transform_qkv_param、transform_mlp_1_param并落在InferenceParameter上。接入文档本身保存在 AddingAModel.md与各模型实现源码同目录非常适合作为“写代码时的案头手册”。整体而言这套“Parameter 声明融合 → LayerContainer 路由参数 → DSTransformerModelBase 计算 → Policy 组装”的分层设计把「checkpoint 参数布局的差异」和「计算图/分片/量化等推理细节」彻底解耦——新增模型时大部分工作集中在描述参数映射关系上而分片、融合、KV 缓存管理等繁琐部分则由框架基类自动接管。【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考