多GPU

假设您已经阅读了注入教程并对如何注入模型有了基本了解。在本教程中，我们将向您展示如何使用KTransformers在多个GPU上运行模型。

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

首先，对于多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

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

如何充分利用多GPU的显存

当你有多个GPU时，可以通过将更多的权重移到GPU上来充分利用每个GPU的VRAM。

例如，对于 DeepSeekV2-Chat，我们可以将专家的权重移动到 GPU 上。

例如，两个GPU的yaml是：

- match:
    name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp\\.experts$"
  replace:
    class: ktransformers.operators.experts.KTransformersExperts
    kwargs:
      generate_device: "cpu"
      generate_op:  "KExpertsCPU"
      out_device: "cuda:0"
  recursive: False

但是我们在cuda:0上有额外的60GB VRAM，我们可以将第4到第8层的专家移动到cuda:0。

# Add new rule before old rule.
- match:
    name: "^model\\.layers\\.([4-8])\\.mlp\\.experts$" # inject experts in layer 4~8 as marlin expert
  replace:
    class: ktransformers.operators.experts.KTransformersExperts  
    kwargs:
      generate_device: "cuda:0"
      generate_op:  "KExpertsMarlin"
  recursive: False

- match:
    name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp\\.experts$"
  replace:
    class: ktransformers.operators.experts.KTransformersExperts     
    kwargs:
      generate_device: "cpu"
      generate_op:  "KExpertsCPU"
      out_device: "cuda:0"
  recursive: False 

根据您的需要调整层范围。请注意：

• 每个移动到GPU的专家的加载速度将显著减慢。
• 如果您想将专家移动到GPU，必须关闭cuda图。
• 对于DeepSeek-R1/V3，每个移动到GPU的专家将消耗大约6GB的显存。
• 在yaml中，第一个匹配的规则将被应用。例如，如果您有两个匹配同一层的规则，则仅第一个规则的替换有效。