mirror of
https://github.com/LostRuins/koboldcpp.git
synced 2026-09-02 02:51:17 +02:00
Deepseek 4: -sm tensor (#26490)
* DSV4: sm tensor * set coarser granularity for head splits * fix dspark * add model saving for dsv4 + allow dflash to return on specific device * add comment about dsv4 seq_rm * simplify * add shared expert delayed allreduce * remove special test for dsv4
This commit is contained in:
+67
-3
@@ -365,6 +365,8 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
|
||||
const llama_meta_device_get_split_state_userdata * ud = (const llama_meta_device_get_split_state_userdata *) userdata;
|
||||
const llama_hparams & hparams = ud->model->hparams;
|
||||
const std::string tensor_name = tensor->name;
|
||||
const bool is_dsv4 = ud->model->arch == LLM_ARCH_DEEPSEEK4 ||
|
||||
(ud->model->arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0);
|
||||
|
||||
static const std::regex pattern_q_weight ("blk\\.\\d*\\.attn_q.weight");
|
||||
static const std::regex pattern_kv_weight ("blk\\.\\d*\\.attn_(k|v).weight");
|
||||
@@ -374,9 +376,13 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
|
||||
static const std::regex pattern_qkv_bias ("blk\\.\\d*\\.attn_qkv.bias");
|
||||
static const std::regex pattern_qk_norm ("blk\\.\\d*\\.attn_(q|k)_norm\\.weight");
|
||||
static const std::regex pattern_kv_cache ("cache_(k|v)_l\\d*");
|
||||
static const std::regex pattern_dsv4_state ("dsv4_(csa|hca|lid)_state_(kv|score)_l\\d*");
|
||||
static const std::regex pattern_attn_sinks ("blk\\.\\d*\\.attn_sinks.weight");
|
||||
static const std::regex pattern_attn_out_weight ("blk\\.\\d*\\.attn_output.weight");
|
||||
static const std::regex pattern_attn_out_bias ("blk\\.\\d*\\.attn_output.bias");
|
||||
static const std::regex pattern_attn_out_a_weight("blk\\.\\d*\\.attn_output_a\\.weight");
|
||||
static const std::regex pattern_attn_out_b_weight("blk\\.\\d*\\.attn_output_b\\.weight");
|
||||
static const std::regex pattern_attn_q_b_weight ("blk\\.\\d*\\.attn_q_b\\.weight");
|
||||
static const std::regex pattern_attn_gate_weight("blk\\.\\d*\\.attn_gate.weight");
|
||||
|
||||
static const std::regex pattern_ssm_dt ("blk\\.\\d*\\.ssm_dt.bias");
|
||||
@@ -395,8 +401,11 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
|
||||
static const std::regex pattern_ffn_gate_bias ("blk\\.\\d*\\.ffn_gate(_exps)?.bias");
|
||||
static const std::regex pattern_ffn_gate_up_weight("blk\\.\\d*\\.ffn_gate_up(_exps)?.weight");
|
||||
static const std::regex pattern_ffn_down_weight ("blk\\.\\d*\\.ffn_down(_exps)?.weight");
|
||||
static const std::regex pattern_ffn_down_bias ("blk\\.\\d*\\.ffn_down.bias");
|
||||
static const std::regex pattern_ffn_down_exps_bias("blk\\.\\d*\\.ffn_down_exps.bias");
|
||||
static const std::regex pattern_ffn_down_bias ("blk\\.\\d*\\.ffn_down.bias");
|
||||
static const std::regex pattern_ffn_down_exps_bias ("blk\\.\\d*\\.ffn_down_exps.bias");
|
||||
static const std::regex pattern_ffn_up_shexp_weight ("blk\\.\\d*\\.ffn_up_shexp.weight");
|
||||
static const std::regex pattern_ffn_gate_shexp_weight ("blk\\.\\d*\\.ffn_gate_shexp.weight");
|
||||
static const std::regex pattern_ffn_down_shexp_weight ("blk\\.\\d*\\.ffn_down_shexp.weight");
|
||||
|
||||
static const std::regex pattern_output_weight("output\\.weight");
|
||||
static const std::regex pattern_output_bias ("output\\.bias");
|
||||
@@ -453,6 +462,32 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
|
||||
};
|
||||
|
||||
auto get_tensor_config = [&]() -> tensor_config {
|
||||
if (is_dsv4) {
|
||||
if (std::regex_match(tensor_name, pattern_kv_cache) ||
|
||||
std::regex_match(tensor_name, pattern_dsv4_state)) {
|
||||
return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED);
|
||||
}
|
||||
if (std::regex_match(tensor_name, pattern_attn_sinks)) {
|
||||
return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "attn_output_a.weight");
|
||||
}
|
||||
if (std::regex_match(tensor_name, pattern_attn_q_b_weight)) {
|
||||
return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output_a.weight");
|
||||
}
|
||||
if (std::regex_match(tensor_name, pattern_attn_out_a_weight)) {
|
||||
return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_2);
|
||||
}
|
||||
if (std::regex_match(tensor_name, pattern_attn_out_b_weight)) {
|
||||
return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0);
|
||||
}
|
||||
if (std::regex_match(tensor_name, pattern_ffn_up_shexp_weight) ||
|
||||
std::regex_match(tensor_name, pattern_ffn_gate_shexp_weight)) {
|
||||
return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "ffn_down_shexp.weight");
|
||||
}
|
||||
if (std::regex_match(tensor_name, pattern_ffn_down_shexp_weight)) {
|
||||
return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "ffn_down_shexp.weight");
|
||||
}
|
||||
}
|
||||
|
||||
// standard attention
|
||||
if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_kv_weight)) {
|
||||
return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output.weight", "ssm_out.weight");
|
||||
@@ -525,6 +560,9 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
|
||||
|
||||
// output
|
||||
if (std::regex_match(tensor_name, pattern_output_weight)) {
|
||||
if (is_dsv4) {
|
||||
return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED);
|
||||
}
|
||||
return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1);
|
||||
}
|
||||
if (std::regex_match(tensor_name, pattern_output_bias)) {
|
||||
@@ -649,8 +687,30 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
|
||||
const int64_t granularity_head = granularity_q / hparams.n_embd_head_k(il); // for tensors with one value per head
|
||||
if (std::regex_match(tensor_name, pattern_attn_sinks)) {
|
||||
GGML_ASSERT(segments.size() == 1);
|
||||
if (is_dsv4) {
|
||||
return {hparams.n_head(il) / hparams.dsv4_o_group_count};
|
||||
}
|
||||
return {granularity_head};
|
||||
}
|
||||
|
||||
if (is_dsv4) {
|
||||
if (std::regex_match(tensor_name, pattern_attn_q_b_weight)) {
|
||||
GGML_ASSERT(segments.size() == 1);
|
||||
// the grouped output projection requires each device to hold whole groups of heads
|
||||
const int64_t n_head_group = hparams.n_head(il) / hparams.dsv4_o_group_count;
|
||||
return {n_head_group * hparams.n_embd_head_k(il)};
|
||||
}
|
||||
if (std::regex_match(tensor_name, pattern_attn_out_a_weight)) {
|
||||
GGML_ASSERT(segments.size() == 1);
|
||||
return {1};
|
||||
}
|
||||
if (std::regex_match(tensor_name, pattern_attn_out_b_weight)) {
|
||||
GGML_ASSERT(segments.size() == 1);
|
||||
// the boundaries must align with wo_a's per-group split, so quant blocks must not straddle groups
|
||||
GGML_ASSERT(hparams.dsv4_o_lora_rank % blck_size == 0);
|
||||
return {hparams.dsv4_o_lora_rank};
|
||||
}
|
||||
}
|
||||
if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_q_bias)) {
|
||||
GGML_ASSERT(segments.size() == 1);
|
||||
// some models have Q gate tensors, for those cases the granularity needs to be doubled:
|
||||
@@ -687,7 +747,11 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
|
||||
// FFN
|
||||
if (std::regex_match(tensor_name, pattern_ffn_up_weight) || std::regex_match(tensor_name, pattern_ffn_up_bias) ||
|
||||
std::regex_match(tensor_name, pattern_ffn_gate_weight) || std::regex_match(tensor_name, pattern_ffn_gate_bias) ||
|
||||
std::regex_match(tensor_name, pattern_ffn_gate_up_weight) || std::regex_match(tensor_name, pattern_ffn_down_weight)) {
|
||||
std::regex_match(tensor_name, pattern_ffn_gate_up_weight) ||
|
||||
std::regex_match(tensor_name, pattern_ffn_down_weight) ||
|
||||
std::regex_match(tensor_name, pattern_ffn_up_shexp_weight) ||
|
||||
std::regex_match(tensor_name, pattern_ffn_gate_shexp_weight) ||
|
||||
std::regex_match(tensor_name, pattern_ffn_down_shexp_weight)) {
|
||||
const int64_t blck_size_perf = std::lcm(blck_size, 128);
|
||||
GGML_ASSERT(segments.size() == 1);
|
||||
return {blck_size_perf};
|
||||
|
||||
Reference in New Issue
Block a user