kvcache-ai / kvcache-ai/ktransformers

llama-factory deepseek-v3 sft AMXInt8 加载到44层 Segmentation fault

Open
#1,803 1 comment 0 reactions 1 assignee Claimed by @JimmyPeilinLi View on GitHub
Dominant language
Python
Stars
19.5k
Forks
1.6k
Avg merge
19h 32m
Merged PRs (30d)
27

Description

### Reminder

- [x] I have read the above rules and searched the existing issues.

### System Info

python 3.11
torch 2.9.0
transformer 4.51.3
ktransformers-0.4.4+cu128torch29fancy-cp311-cp311-linux_x86_64.whl

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2025 NVIDIA Corporation
Built on Fri_Feb_21_20:23:50_PST_2025
Cuda compilation tools, release 12.8, V12.8.93
Build cuda_12.8.r12.8/compiler.35583870_0

GPU NVIDIA H20-3e * 8

### Reproduction

https://blog.llamafactory.net/posts/ktransformers/

我按照这个blog尝试调deepseekv3 44层 Segmentation fault,之前尝试用transformer没支持的v3.2强行进行config映射按v3加载也是同样的在44层出问题,不过问题是 std::bad_alloc

```shell
GO INTO AMXInt8!!
loading model.layers.42.input_layernorm.weight to cuda:5
loading model.layers.42.post_attention_layernorm.weight to cuda:5
loading model.layers.43.self_attn.q_a_layernorm.weight to cuda:5
loading model.layers.43.self_attn.kv_a_layernorm.weight to cuda:5
loading model.layers.43.self_attn.kv_b_proj.weight to cuda:5
GO INTO AMXInt8!!
Segmentation fault (core dumped)
```

以下是我的yaml,试过blog提供的的单卡和多卡的yaml,问题都是在44层。
是依赖版本问题吗?

```yaml
- match:
name: "^model.embed_tokens"
replace:
class: "default"
kwargs:
generate_device: "cpu"
prefill_device: "cpu"

# === Rotary Embedding Replacement ===

# GPU 0: layers 0–7
- match:
name: "^model\\.layers\\.([0-7])\\."
class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding
replace:
class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3
kwargs:
generate_device: "cuda:0"
prefill_device: "cuda:0"

# GPU 1: layers 8–15
- match:
name: "^model\\.layers\\.(8|9|1[0-5])\\."
class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding
replace:
class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3
kwargs:
generate_device: "cuda:1"
prefill_device: "cuda:1"

# GPU 2: layers 16–23
- match:
name: "^model\\.layers\\.(1[6-9]|2[0-3])\\."
class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding
replace:
class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3
kwargs:
generate_device: "cuda:2"
prefill_device: "cuda:2"

# GPU 3: layers 24–31
- match:
name: "^model\\.layers\\.(2[4-9]|3[0-1])\\."
class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding
replace:
class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3
kwargs:
generate_device: "cuda:3"
prefill_device: "cuda:3"

# GPU 4: layers 32–39
- match:
name: "^model\\.layers\\.(3[2-9])\\."
class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding
replace:
class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3
kwargs:
generate_device: "cuda:4"
prefill_device: "cuda:4"

# GPU 5: layers 40–47
- match:
name: "^model\\.layers\\.(4[0-7])\\."
class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding
replace:
class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3
kwargs:
generate_device: "cuda:5"
prefill_device: "cuda:5"

# GPU 6: layers 48–55
- match:
name: "^model\\.layers\\.(4[8-9]|5[0-5])\\."
class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding
replace:
class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3
kwargs:
generate_device: "cuda:6"
prefill_device: "cuda:6"

# GPU 7: layers 56–60
- match:
name: "^model\\.layers\\.(5[6-9]|60)\\."
class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding
replace:
class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3
kwargs:
generate_device: "cuda:7"
prefill_device: "cuda:7"

# === Linear Layers Replacement (excluding self_attn.kv_b_proj) ===

# GPU 0: layers 0–7
- match:
name: "^model\\.layers\\.([0-7])\\.(?!self_attn\\.kv_b_proj).*$"
class: torch.nn.Linear
replace:
class: ktransformers.operators.linear.KTransformersLinear
kwargs:
generate_device: "cuda:0"
prefill_device: "cuda:0"
generate_op: "KLinearTorch"
prefill_op: "KLinearTorch"

# GPU 1: layers 8–15
- match:
name: "^model\\.layers\\.(8|9|1[0-5])\\.(?!self_attn\\.kv_b_proj).*$"
class: torch.nn.Linear
replace:
class: ktransformers.operators.linear.KTransformersLinear
kwargs:
generate_device: "cuda:1"
prefill_device: "cuda:1"
generate_op: "KLinearTorch"
prefill_op: "KLinearTorch"

# GPU 2: layers 16–23
- match:
name: "^model\\.layers\\.(1[6-9]|2[0-3])\\.(?!self_attn\\.kv_b_proj).*$"
class: torch.nn.Linear
replace:
class: ktransformers.operators.linear.KTransformersLinear
kwargs:
generate_device: "cuda:2"
prefill_device: "cuda:2"
generate_op: "KLinearTorch"
prefill_op: "KLinearTorch"

# GPU 3: layers 24–31
- match:
name: "^model\\.layers\\.(2[4-9]|3[0-1])\\.(?!self_attn\\.kv_b_proj).*$"
class: torch.nn.Linear
replace:
class: ktransformers.operators.linear.KTransformersLinear
kwargs:
generate_device: "cuda:3"
prefill_device: "cuda:3"
generate_op: "KLinearTorch"
prefill_op: "KLinearTorch"

# GPU 4: layers 32–39
- match:
name: "^model\\.layers\\.(3[2-9])\\.(?!self_attn\\.kv_b_proj).*$"
class: torch.nn.Linear
replace:
class: ktransformers.operators.linear.KTransformersLinear
kwargs:
generate_device: "cuda:4"
prefill_device: "cuda:4"
generate_op: "KLinearTorch"
prefill_op: "KLinearTorch"

# GPU 5: layers 40–47
- match:
name: "^model\\.layers\\.(4[0-7])\\.(?!self_attn\\.kv_b_proj).*$"
class: torch.nn.Linear
replace:
class: ktransformers.operators.linear.KTransformersLinear
kwargs:
generate_device: "cuda:5"
prefill_device: "cuda:5"
generate_op: "KLinearTorch"
prefill_op: "KLinearTorch"

# GPU 6: layers 48–55
- match:
name: "^model\\.layers\\.(4[8-9]|5[0-5])\\.(?!self_attn\\.kv_b_proj).*$"
class: torch.nn.Linear
replace:
class: ktransformers.operators.linear.KTransformersLinear
kwargs:
generate_device: "cuda:6"
prefill_device: "cuda:6"
generate_op: "KLinearTorch"
prefill_op: "KLinearTorch"

# GPU 7: layers 56–60
- match:
name: "^model\\.layers\\.(5[6-9]|60)\\.(?!self_attn\\.kv_b_proj).*$"
class: torch.nn.Linear
replace:
class: ktransformers.operators.linear.KTransformersLinear
kwargs:
generate_device: "cuda:7"
prefill_device: "cuda:7"
generate_op: "KLinearTorch"
prefill_op: "KLinearTorch"

# === MLP (MoE) Replacement ===

# GPU 0: layers 0–7
- match:
name: "^model\\.layers\\.([0-7])\\.mlp$"
class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE
replace:
class: ktransformers.operators.experts.KDeepseekV3MoE
kwargs:
generate_device: "cuda:0"
prefill_device: "cuda:0"

# GPU 1: layers 8–15
- match:
name: "^model\\.layers\\.(8|9|1[0-5])\\.mlp$"
class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE
replace:
class: ktransformers.operators.experts.KDeepseekV3MoE
kwargs:
generate_device: "cuda:1"
prefill_device: "cuda:1"

# GPU 2: layers 16–23
- match:
name: "^model\\.layers\\.(1[6-9]|2[0-3])\\.mlp$"
class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE
replace:
class: ktransformers.operators.experts.KDeepseekV3MoE
kwargs:
generate_device: "cuda:2"
prefill_device: "cuda:2"

# GPU 3: layers 24–31
- match:
name: "^model\\.layers\\.(2[4-9]|3[0-1])\\.mlp$"
class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE
replace:
class: ktransformers.operators.experts.KDeepseekV3MoE
kwargs:
generate_device: "cuda:3"
prefill_device: "cuda:3"

# GPU 4: layers 32–39
- match:
name: "^model\\.layers\\.(3[2-9])\\.mlp$"
class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE
replace:
class: ktransformers.operators.experts.KDeepseekV3MoE
kwargs:
generate_device: "cuda:4"
prefill_device: "cuda:4"

# GPU 5: layers 40–47
- match:
name: "^model\\.layers\\.(4[0-7])\\.mlp$"
class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE
replace:
class: ktransformers.operators.experts.KDeepseekV3MoE
kwargs:
generate_device: "cuda:5"
prefill_device: "cuda:5"

# GPU 6: layers 48–55
- match:
name: "^model\\.layers\\.(4[8-9]|5[0-5])\\.mlp$"
class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE
replace:
class: ktransformers.operators.experts.KDeepseekV3MoE
kwargs:
generate_device: "cuda:6"
prefill_device: "cuda:6"

# GPU 7: layers 56–60
- match:
name: "^model\\.layers\\.(5[6-9]|60)\\.mlp$"
class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE
replace:
class: ktransformers.operators.experts.KDeepseekV3MoE
kwargs:
generate_device: "cuda:7"
prefill_device: "cuda:7"

# === MLP Gate Replacement ===

# GPU 0: layers 0–7
- match:
name: "^model\\.layers\\.([0-7])\\.mlp\\.gate$"
class: ktransformers.models.modeling_deepseek_v3.MoEGate
replace:
class: ktransformers.operators.gate.KMoEGate
kwargs:
generate_device: "cuda:0"
prefill_device: "cuda:0"

# GPU 1: layers 8–15
- match:
name: "^model\\.layers\\.(8|9|1[0-5])\\.mlp\\.gate$"
class: ktransformers.models.modeling_deepseek_v3.MoEGate
replace:
class: ktransformers.operators.gate.KMoEGate
kwargs:
generate_device: "cuda:1"
prefill_device: "cuda:1"

# GPU 2: layers 16–23
- match:
name: "^model\\.layers\\.(1[6-9]|2[0-3])\\.mlp\\.gate$"
class: ktransformers.models.modeling_deepseek_v3.MoEGate
replace:
class: ktransformers.operators.gate.KMoEGate
kwargs:
generate_device: "cuda:2"
prefill_device: "cuda:2"

# GPU 3: layers 24–31
- match:
name: "^model\\.layers\\.(2[4-9]|3[0-1])\\.mlp\\.gate$"
class: ktransformers.models.modeling_deepseek_v3.MoEGate
replace:
class: ktransformers.operators.gate.KMoEGate
kwargs:
generate_device: "cuda:3"
prefill_device: "cuda:3"

# GPU 4: layers 32–39
- match:
name: "^model\\.layers\\.(3[2-9])\\.mlp\\.gate$"
class: ktransformers.models.modeling_deepseek_v3.MoEGate
replace:
class: ktransformers.operators.gate.KMoEGate
kwargs:
generate_device: "cuda:4"
prefill_device: "cuda:4"

# GPU 5: layers 40–47
- match:
name: "^model\\.layers\\.(4[0-7])\\.mlp\\.gate$"
class: ktransformers.models.modeling_deepseek_v3.MoEGate
replace:
class: ktransformers.operators.gate.KMoEGate
kwargs:
generate_device: "cuda:5"
prefill_device: "cuda:5"

# GPU 6: layers 48–55
- match:
name: "^model\\.layers\\.(4[8-9]|5[0-5])\\.mlp\\.gate$"
class: ktransformers.models.modeling_deepseek_v3.MoEGate
replace:
class: ktransformers.operators.gate.KMoEGate
kwargs:
generate_device: "cuda:6"
prefill_device: "cuda:6"

# GPU 7: layers 56–60
- match:
name: "^model\\.layers\\.(5[6-9]|60)\\.mlp\\.gate$"
class: ktransformers.models.modeling_deepseek_v3.MoEGate
replace:
class: ktransformers.operators.gate.KMoEGate
kwargs:
generate_device: "cuda:7"
prefill_device: "cuda:7"

# === MLP Experts Replacement (SFT配置) ===

# GPU 0: layers 0–7
- match:
name: "^model\\.layers\\.([0-7])\\.mlp\\.experts$"
replace:
class: ktransformers.operators.experts.KTransformersExperts
kwargs:
prefill_device: "cuda:0"
prefill_op: "KExpertsTorch"
generate_device: "cpu"
generate_op: "KSFTExpertsCPU"
out_device: "cuda:0"
backend: "AMXInt8" # 或 "AMXBF16" 或 "llamafile" (default)
recursive: False

# GPU 1: layers 8–15
- match:
name: "^model\\.layers\\.(8|9|1[0-5])\\.mlp\\.experts$"
replace:
class: ktransformers.operators.experts.KTransformersExperts
kwargs:
prefill_device: "cuda:1"
prefill_op: "KExpertsTorch"
generate_device: "cpu"
generate_op: "KSFTExpertsCPU"
out_device: "cuda:1"
backend: "AMXInt8"
recursive: False

# GPU 2: layers 16–23
- match:
name: "^model\\.layers\\.(1[6-9]|2[0-3])\\.mlp\\.experts$"
replace:
class: ktransformers.operators.experts.KTransformersExperts
kwargs:
prefill_device: "cuda:2"
prefill_op: "KExpertsTorch"
generate_device: "cpu"
generate_op: "KSFTExpertsCPU"
out_device: "cuda:2"
backend: "AMXInt8"
recursive: False

# GPU 3: layers 24–31
- match:
name: "^model\\.layers\\.(2[4-9]|3[0-1])\\.mlp\\.experts$"
replace:
class: ktransformers.operators.experts.KTransformersExperts
kwargs:
prefill_device: "cuda:3"
prefill_op: "KExpertsTorch"
generate_device: "cpu"
generate_op: "KSFTExpertsCPU"
out_device: "cuda:3"
backend: "AMXInt8"
recursive: False

# GPU 4: layers 32–39
- match:
name: "^model\\.layers\\.(3[2-9])\\.mlp\\.experts$"
replace:
class: ktransformers.operators.experts.KTransformersExperts
kwargs:
prefill_device: "cuda:4"
prefill_op: "KExpertsTorch"
generate_device: "cpu"
generate_op: "KSFTExpertsCPU"
out_device: "cuda:4"
backend: "AMXInt8"
recursive: False

# GPU 5: layers 40–47
- match:
name: "^model\\.layers\\.(4[0-7])\\.mlp\\.experts$"
replace:
class: ktransformers.operators.experts.KTransformersExperts
kwargs:
prefill_device: "cuda:5"
prefill_op: "KExpertsTorch"
generate_device: "cpu"
generate_op: "KSFTExpertsCPU"
out_device: "cuda:5"
backend: "AMXInt8"
recursive: False

# GPU 6: layers 48–55
- match:
name: "^model\\.layers\\.(4[8-9]|5[0-5])\\.mlp\\.experts$"
replace:
class: ktransformers.operators.experts.KTransformersExperts
kwargs:
prefill_device: "cuda:6"
prefill_op: "KExpertsTorch"
generate_device: "cpu"
generate_op: "KSFTExpertsCPU"
out_device: "cuda:6"
backend: "AMXInt8"
recursive: False

# GPU 7: layers 56–60
- match:
name: "^model\\.layers\\.(5[6-9]|60)\\.mlp\\.experts$"
replace:
class: ktransformers.operators.experts.KTransformersExperts
kwargs:
prefill_device: "cuda:7"
prefill_op: "KExpertsTorch"
generate_device: "cpu"
generate_op: "KSFTExpertsCPU"
out_device: "cuda:7"
backend: "AMXInt8"
recursive: False

# === Self-Attention Replacement ===

# GPU 0: layers 0–7
- match:
name: "^model\\.layers\\.([0-7])\\.self_attn$"
replace:
class: ktransformers.operators.attention.KDeepseekV2Attention
kwargs:
generate_device: "cuda:0"
prefill_device: "cuda:0"
absorb_for_prefill: False

# GPU 1: layers 8–15
- match:
name: "^model\\.layers\\.(8|9|1[0-5])\\.self_attn$"
replace:
class: ktransformers.operators.attention.KDeepseekV2Attention
kwargs:
generate_device: "cuda:1"
prefill_device: "cuda:1"
absorb_for_prefill: False

# GPU 2: layers 16–23
- match:
name: "^model\\.layers\\.(1[6-9]|2[0-3])\\.self_attn$"
replace:
class: ktransformers.operators.attention.KDeepseekV2Attention
kwargs:
generate_device: "cuda:2"
prefill_device: "cuda:2"
absorb_for_prefill: False

# GPU 3: layers 24–31
- match:
name: "^model\\.layers\\.(2[4-9]|3[0-1])\\.self_attn$"
replace:
class: ktransformers.operators.attention.KDeepseekV2Attention
kwargs:
generate_device: "cuda:3"
prefill_device: "cuda:3"
absorb_for_prefill: False

# GPU 4: layers 32–39
- match:
name: "^model\\.layers\\.(3[2-9])\\.self_attn$"
replace:
class: ktransformers.operators.attention.KDeepseekV2Attention
kwargs:
generate_device: "cuda:4"
prefill_device: "cuda:4"
absorb_for_prefill: False

# GPU 5: layers 40–47
- match:
name: "^model\\.layers\\.(4[0-7])\\.self_attn$"
replace:
class: ktransformers.operators.attention.KDeepseekV2Attention
kwargs:
generate_device: "cuda:5"
prefill_device: "cuda:5"
absorb_for_prefill: False

# GPU 6: layers 48–55
- match:
name: "^model\\.layers\\.(4[8-9]|5[0-5])\\.self_attn$"
replace:
class: ktransformers.operators.attention.KDeepseekV2Attention
kwargs:
generate_device: "cuda:6"
prefill_device: "cuda:6"
absorb_for_prefill: False

# GPU 7: layers 56–60
- match:
name: "^model\\.layers\\.(5[6-9]|60)\\.self_attn$"
replace:
class: ktransformers.operators.attention.KDeepseekV2Attention
kwargs:
generate_device: "cuda:7"
prefill_device: "cuda:7"
absorb_for_prefill: False

# === Overall Model Replacement with Transfer Map ===

- match:
name: "^model$"
replace:
class: "ktransformers.operators.models.KDeepseekV2Model"
kwargs:
per_layer_prefill_intput_threshold: 0 # 0 means close layer-wise prefill
transfer_map:
8: "cuda:1" # Layers 8+ on GPU 1
16: "cuda:2" # Layers 16+ on GPU 2
24: "cuda:3" # Layers 24+ on GPU 3
32: "cuda:4" # Layers 32+ on GPU 4
40: "cuda:5" # Layers 40+ on GPU 5
48: "cuda:6" # Layers 48+ on GPU 6
56: "cuda:7" # Layers 56+ on GPU 7

# === Default Catch-All for Other Modules ===

# GPU 0: layers 0–7
- match:
name: "^model\\.layers\\.([0-7])\\."
replace:
class: "default"
kwargs:
generate_device: "cuda:0"
prefill_device: "cuda:0"

# GPU 1: layers 8–15
- match:
name: "^model\\.layers\\.(8|9|1[0-5])\\."
replace:
class: "default"
kwargs:
generate_device: "cuda:1"
prefill_device: "cuda:1"

# GPU 2: layers 16–23
- match:
name: "^model\\.layers\\.(1[6-9]|2[0-3])\\."
replace:
class: "default"
kwargs:
generate_device: "cuda:2"
prefill_device: "cuda:2"

# GPU 3: layers 24–31
- match:
name: "^model\\.layers\\.(2[4-9]|3[0-1])\\."
replace:
class: "default"
kwargs:
generate_device: "cuda:3"
prefill_device: "cuda:3"

# GPU 4: layers 32–39
- match:
name: "^model\\.layers\\.(3[2-9])\\."
replace:
class: "default"
kwargs:
generate_device: "cuda:4"
prefill_device: "cuda:4"

# GPU 5: layers 40–47
- match:
name: "^model\\.layers\\.(4[0-7])\\."
replace:
class: "default"
kwargs:
generate_device: "cuda:5"
prefill_device: "cuda:5"

# GPU 6: layers 48–55
- match:
name: "^model\\.layers\\.(4[8-9]|5[0-5])\\."
replace:
class: "default"
kwargs:
generate_device: "cuda:6"
prefill_device: "cuda:6"

# GPU 7: layers 56–60
- match:
name: "^model\\.layers\\.(5[6-9]|60)\\."
replace:
class: "default"
kwargs:
generate_device: "cuda:7"
prefill_device: "cuda:7"

- match:
name: "^lm_head"
class: torch.nn.Linear
replace:
class: ktransformers.operators.linear.KTransformersLinear
kwargs:
generate_device: "cuda:7"
prefill_device: "cuda:7"
generate_op: "KLinearTorch"
prefill_op: "KLinearTorch"

# For final modules (model.norm), ensure they are on GPU 7
- match:
name: "^model\\.norm"
replace:
class: "default"
kwargs:
generate_device: "cuda:7"
prefill_device: "cuda:7"
```

### Others

_No response_

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.