教程：逐步注入操作符

作者: Azure-Tang

简而言之

本教程将指导您通过使用KTransformers框架将自定义运算符注入模型的过程。我们将以DeepSeekV2-Chat模型为例，逐步演示如何将自定义运算符注入模型。本教程将涵盖以下主题：

如何编写注入规则
- 理解模型的结构
多GPU
如何编写一个新的运算符并将其注入模型中

如何编写注入规则

Inject框架的注入规则的基本形式如下：

- match:
    name: "^model\\.layers\\..*\\.*$"  # Target module name
    class: torch.nn.Linear  # Target module
  replace:
    class: "default"
    kwargs:
      generate_device: "cuda:0"
      # your_op_param_1: 1234
      # your_op_param_2: 5678
  recursive: True

match：此字段标记匹配规则，可以以两种形式出现，名字和类。这两种匹配规则可以一起或单独出现；仅在满足两个条件时才匹配。
replace:
- class: 可以导入以替换目标模块的Python类。如果不需要替换，请设置为默认。
- kwargs: List of parameters needed for module initialization.
  - generate_device: 该模块的设备，可以设置为“cpu”、“cuda”、“cuda:1”等。
递归：是否递归地注入该模块的子模块，默认值为 True。

对于递归字段：某些模块包含多个子模块，例如自注意力模块通常包含q/k/v/o四个线性模块。如果我们替换自注意力模块但不希望内部线性模块被其他规则覆盖，请将此规则设置为False。

理解模型结构

以 deepseek-ai/DeepSeek-V2-Lite-Chat 为例，我们可以一步一步遵循上述规则来注入我们自定义的模块并运行它。KTransformers 提供了高度的灵活性，允许您替换/实验基本操作符。然而，这也要求用户清楚地理解他们正在运行的模型的结构。

幸运的是，了解模型的结构非常简单。在deepseek-ai/DeepSeek-V2-Lite主页上打开文件列表，您可以看到以下文件：

Inject-Struction

从 .saftensors 文件中，我们可以看到每层权重的名称，这与注入规则中的 match.name 属性对应。从 modeling_deepseek.py 文件中，我们可以看到每个模块类的具体实现，类名与注入规则中的 match.class 属性对应。

来自 .saftensors 和 modeling_deepseek.py 文件的 DeepSeekV2 模型结构如下：

Inject-Struction

支持的运算符及其对应的类别如下：

匹配	替换	后端	描述
线性	KTransformersLinear	KLinearMarlin	以Marlin为后端
		KLinearTorch	以pytorch作为后端
		KLinearCPUInfer	llamafile 作为后端
		KLinearFP8	Triton fp8_gemm 内核。需要GPU能够计算fp8数据
专家	KTransformers专家	KExpertsTorch	以pytorch作为后端
		KExpertsMarlin	Marlin 作为后端
		KExpertsCPU	以llamafile作为后端
注意	KDeepseekV2Attention	KDeepseekV2Attention	MLA 实现
MoE	KMistralSparseMoEBlock	KQwen2MoeSparseMoeBlock	Qwen2 的 MoE
	KDeepseekV2MoE	KDeepseekV2MoE	DeepseekV2的MoE
模型	KQwen2MoeModel	KQwen2MoeModel	Qwen2的模型
	KDeepseekV2Model	KDeepseekV2Model	DeepseekV2的模型
RoPE	旋转嵌入	旋转嵌入	RoPE模块
	YarnRotaryEmbedding	YarnRotaryEmbedding	RoPE模块

然后我们开始逐步注入自定义模块，我们的目标是：

用自定义的Marlin线性模块替换线性模块。
用自定义的基于吸收的MLA模块替换自注意力模块。
用自定义的专家模块替换专家模块。
用自定义的MoE模块替换MoE模块。
用自定义的RoPE模块替换RoPE模块。
为每个模块设置运行设备。

注入规则的完整实现可以在这里找到。

基于矩阵吸收的MLA注入

对于注意力模块的注入，我们只需要使用正则表达式来匹配在变换器中使用的模块名称，并用我们自己的MLA模块实现替换它们。YAML注入规则如下：

- match:
    name: "^model\\.layers\\..*\\.self_attn$"  # Regular expression
  replace:
    class: ktransformers.operators.attention.KDeepseekV2Attention # Optimized MLA implementation

正如您所看到的，YAML文件中的每条规则有两个部分：匹配和替换。匹配部分指定要替换的模块，而替换部分指定要注入到模型中的模块以及初始化关键字。

注入路由专家

对于路由专家（对应图中的exps），我们注入的模块是CPUInfer，它被包装在包装模块KTransformersExperts中。KTransformersExperts有多种实现，我们需要指定关键字来告诉包装模块我们想要使用哪种实现以及我们计划如何使用。

在变换器的源代码中，MoE是使用nn.ModuleList实现的。我们不希望KTransformers遍历列表中的所有子模块并逐个注入，因此在这个规则中，我们将recursive: False设置为防止递归注入此模块的子模块。YAML规则如下：

- match:
    name: "^model\\.layers\\..*\\.mlp\\.experts$"
  replace:
    class: ktransformers.operators.experts.KTransformersExperts     # Custom MoE kernel with expert parallelism
    kwargs:
      generate_device: "cpu"
      generate_op: "MLPCPUExperts"
      out_device: "cuda"
  recursive: False # Don't recursively inject submodules of this module

如果我们将Routed Experts作为一个自定义模块注入，则无法使用原始nn.ModuleList中的接口。因此，有必要修改FFN模块中的forward函数。最简单的方法是实现一个具有自定义forward函数的新模块并注入它。

- match:
    class: ktransformers.models.modeling_deepseek.DeepseekV2MoE
  replace:
    class: ktransformers.operators.experts.KDeepseekV2MoE     # MLP module with custom forward function

线性层的注入

对于剩余的线性层模块，我们旨在使用量化操作符来节省存储空间，同时提高性能。由于目前没有关于将 MLA 和量化结合使用的研究，我们不想将线性注入到 MLA 操作符中。因此，我们可以修改正则表达式，并在规则的匹配部分添加类型检查。只有同时匹配名称和类的模块才会被注入。我们还需要传递一些与 Routed Experts 注入类似的关键词。YAML 规则如下：

- match:
    name: "^model\\.layers\\.(?!.*self_attn).*$"  # Regular expression
    class: torch.nn.Linear  # Only match modules matching name and class simultaneously
  replace:
    class: ktransformers.operators.linear.KTransformersLinear  # Optimized kernel on quantized data types
    kwargs:
      generate_device: "cuda"
      generate_op: "QuantizedLinearMarlin"

带有预计算缓冲区的模块注入

为了避免在初始化注入的原始模型时占用资源，我们使用torch的元设备来初始化原始模型。RoPE模块在初始化时预先计算一些缓冲区，但在使用元设备时不会执行任何计算。因此，我们需要在加载模型时补偿缓冲区的计算。简单来说，我们将一个自定义模块注入到旋转嵌入模块中，该模块在加载时执行预计算。YAML规则如下：

- match:
    class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding
  replace:
    class: ktransformers.operators.RoPE.YarnRotaryEmbedding

指定模块的运行设备

最后，我们为所有模块设置了一个后备基本属性 generate_device：

- match:
    name: "^model\\.layers\\..*\\.|^lm_head"
  replace:
    class: "default"
    kwargs:
      generate_device: "cuda"
  
- match:
    name: "^model.embed_tokens"
  replace:
    class: "default"
    kwargs:
        generate_device: "cpu"

通过这两条规则，我们将所有之前未匹配的层（及其子模块）和 lm_head 放置在 cuda 上，将嵌入放置在 cpu 上。请注意，模块的属性将由它匹配的第一条规则决定。例如，如果您稍后在一个注入的模块中设置一个新的 replace.kwargs.generate_device，则之前设置的设备将优先考虑。如果您的计算机有多个显卡，您也可以将模型配置为使用多个显卡。

多GPU

如果您有多个GPU，您可以为每个模块设置不同的GPU。 DeepseekV2-Chat有60层，如果我们有2个GPU，我们可以将每个GPU分配30层。完整的多GPU规则示例在这里。

Inject-Struction

首先，对于多GPU，我们必须注入一个新的操作符 KDeepseekV2Model。并将层的划分设置为不同的GPU。对于我们的情况，我们必须设置 transfer_map 在 KDeepseekV2Model 操作符中如下：

- match:
    name: "^model$"
  replace:
    class: "ktransformers.operators.models.KDeepseekV2Model"
    kwargs:
      transfer_map: 
        30: "cuda:1"

我们必须为模型中的每个模块设置设备。

例如，对于 routed experts，一个 GPU 的 yaml 是：

- match:
    name: "^model\\.layers\\..*\\.mlp\\.experts$"
  replace:
    class: ktransformers.operators.experts.KTransformersExperts     # Custom MoE kernel with expert parallelism
    kwargs:
      generate_device: "cuda:0"
      generate_op: "MLPCUDAExperts"
      out_device: "cuda:0"
  recursive: False # Don't recursively inject submodules of this module

但对于两个GPU，我们需要为模型中的每个模块设置设备。

# allcate 0-29 layers‘s out_device to cuda:0
- match:
    name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp\\.experts$"
  replace:
    class: ktransformers.operators.experts.KTransformersExperts     # custom MoE Kernel with expert paralleism
    kwargs:
      generate_device: "cpu"
      generate_op:  "KExpertsCPU"
      out_device: "cuda:0"
  recursive: False # don't recursively inject submodules of this module

# allocate 30-59 layers‘s out_device to cuda:1
- match:
    name: "^model\\.layers\\.([345][0-9])\\.mlp\\.experts$"
  replace:
    class: ktransformers.operators.experts.KTransformersExperts     # custom MoE Kernel with expert paralleism
    kwargs:
      generate_device: "cpu"
      generate_op:  "KExpertsCPU"
      out_device: "cuda:1"
  recursive: False # don't recursively inject submodules of this module

对于其他模块，我们可以以相同的方式设置设备。

如何编写新的操作符并注入到模型中

在本节中，我们将解释如何编写可注入的运算符，以实现新的线性为例。

首先，所有可注入的操作符需要继承自BaseInjectedModule类，该类继承了我们的注入框架所需的一些属性。它的初始化函数需要符合以下基本格式：

class LinearTorchInject(BaseInjectedModule):
    def __init__(
        self,
        key: str,
        gguf_loader: GGUFLoader,
        config: PretrainedConfig,
        orig_module: nn.Module = None,
        generate_device: str = "cuda",
        **kwargs,
    ):
        super().__init__(key, gguf_loader, config, orig_module, generate_device, **kwargs)

如果用户有其他参数需要传递给这个类，它们也可以包含在初始化函数中，并在yaml文件中的kwargs参数中重新传递。例如，如果我们的操作员想要传递一个参数 my_param，初始化函数可以写成：

class LinearTorchInject(BaseInjectedModule):
    def __init__(
        self,
        key: str,
        gguf_loader: GGUFLoader,
        config: PretrainedConfig,
        orig_module: nn.Module = None,
        generate_device: str = "cuda",
        my_param: bool = True,
        **kwargs,
    ):
        super().__init__(key, gguf_loader, config, orig_module, generate_device, **kwargs)
        self.my_param = my_param

那么我们的注入规则可以写成：

- match: 
    name: "^model\\.layers\\..*$"  # Regular expression matches the module name.
    class: torch.nn.Linear  # Type restrictions can be added.
  replace:
    class: ktransformers.operators.linear.LinearTorchInject  # Inject module path
    kwargs: # Extra parameters
      generate_device: "cuda"
      my_param: True

对于线性模块，读取gguf文件中的权重也是必要的。我们提供了KLinearBase类，帮助用户从gguf文件中读取权重。用户只需继承并实现load、unload和forward函数。因此，一个完全可注入的线性类将如下所示：

class LinearTorchInject(BaseInjectedModule, KLinearBase):
    def __init__(
        self,
        key: str,
        gguf_loader: GGUFLoader,
        config: PretrainedConfig,
        orig_module: nn.Module = None,
        generate_device: str = "cuda",
        **kwargs,
    ):
        super().__init__(key, gguf_loader, config, orig_module, generate_device, **kwargs)
        KLinearBase.__init__(self)
        self.has_bias = False
        self.dtype = torch.get_default_dtype()
        self.w = None
        self.has_bias = False
    
    def load(self, w: dict | nn.Parameter | tuple | None = None, device: str|None = None):
        if device is None: device = self.device
        if w is None: w = self.load_weight(device=device)

        if isinstance(w, nn.Parameter):
            self.w = w.to(dtype=self.dtype).view(self.out_features, self.in_features).T
            self.has_bias = False
        elif isinstance(w, tuple):
            self.w = w[0].to(dtype=self.dtype).view(self.out_features, self.in_features).T
            self.bias = w[1].to(dtype=self.dtype)
            self.has_bias = True
        else:
            raise ValueError("Invalid weight type")
        self.w = self.w.to(device)
        if self.has_bias:
            self.bias = self.bias.to(device)

    def unload(self):
        if self.w is not None:
            self.w = None
        if self.has_bias:
            self.bias = None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        dtype = x.dtype
        out_device = x.device
        x = x.to(device=self.device, dtype=self.dtype)
        x = x @ self.w
        if self.has_bias:
            x = x + self.bias
        x = x.to(dtype=dtype, device=out_device)
        return x

请注意，self.load_weight 函数由 KLinearBase 类提供，以帮助用户将权重从 gguf 文件加载到模块中。KLinearBase 的实现细节可以在 GITHUB 上找到。

Ktransformers