Skip to content

Commit cb87136

Browse files
wooway777wooway777
committed
feat: add DeepSeek V2 model
1 parent d6da380 commit cb87136

8 files changed

Lines changed: 647 additions & 0 deletions

‎csrc/models/deepseek_v2/deepseek_v2_attention.cpp‎

Lines changed: 183 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,183 @@
1+
#include "deepseek_v2_attention.hpp"
2+
3+
#include "../../global_state/global_state.hpp"
4+
#include "../../utils.hpp"
5+
#include "infinicore/ops.hpp"
6+
#include "infinicore/ops/broadcast_to.hpp"
7+
#include "infinicore/ops/cat.hpp"
8+
#include "infinicore/ops/pad.hpp"
9+
10+
#include <cmath>
11+
#include <stdexcept>
12+
13+
namespace infinilm::models::deepseek_v2 {
14+
namespace {
15+
16+
float yarn_get_mscale(float scale, float mscale) {
17+
if (scale <= 1.0f) {
18+
return 1.0f;
19+
}
20+
return 0.1f * mscale * std::log(scale) + 1.0f;
21+
}
22+
23+
} // namespace
24+
25+
DeepseekV2Attention::DeepseekV2Attention(std::shared_ptr<infinilm::config::ModelConfig> model_config,
26+
size_t layer_idx,
27+
const infinicore::Device &device) {
28+
layer_idx_ = layer_idx;
29+
hidden_size_ = model_config->get<size_t>("hidden_size");
30+
qk_nope_head_dim_ = model_config->get<size_t>("qk_nope_head_dim");
31+
qk_rope_head_dim_ = model_config->get<size_t>("qk_rope_head_dim");
32+
q_head_dim_ = qk_nope_head_dim_ + qk_rope_head_dim_;
33+
v_head_dim_ = model_config->get<size_t>("v_head_dim");
34+
35+
const auto &dtype{model_config->get_dtype()};
36+
const size_t total_num_heads = model_config->get<size_t>("num_attention_heads");
37+
const size_t kv_lora_rank = model_config->get<size_t>("kv_lora_rank");
38+
const bool attention_bias = model_config->get_or<bool>("attention_bias", false);
39+
const double rms_norm_eps = model_config->get<double>("rms_norm_eps");
40+
41+
const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info();
42+
const int tp_rank = rank_info.tp_rank;
43+
const int tp_size = rank_info.tp_size;
44+
if ((total_num_heads < static_cast<size_t>(tp_size)) || (total_num_heads % static_cast<size_t>(tp_size) != 0)) {
45+
throw std::runtime_error("DeepseekV2Attention: num_attention_heads must be divisible by tp_size");
46+
}
47+
num_attention_heads_ = total_num_heads / static_cast<size_t>(tp_size);
48+
attention_backend_ = infinilm::global_state::get_infinilm_config().attention_backend;
49+
50+
auto quantization_method = model_config->get_quantization_method();
51+
INFINICORE_NN_MODULE_INIT(q_proj, hidden_size_, total_num_heads * q_head_dim_, quantization_method, false, dtype, device, tp_rank, tp_size);
52+
INFINICORE_NN_MODULE_INIT(kv_a_proj_with_mqa, hidden_size_, kv_lora_rank + qk_rope_head_dim_, attention_bias, dtype, device);
53+
INFINICORE_NN_MODULE_INIT(kv_a_layernorm, kv_lora_rank, rms_norm_eps, dtype, device);
54+
INFINICORE_NN_MODULE_INIT(kv_b_proj, kv_lora_rank, total_num_heads * (qk_nope_head_dim_ + v_head_dim_), quantization_method, false, dtype, device, tp_rank, tp_size);
55+
INFINICORE_NN_MODULE_INIT(o_proj, total_num_heads * v_head_dim_, hidden_size_, quantization_method, attention_bias, dtype, device, tp_rank, tp_size, rank_info.comm);
56+
57+
const size_t max_position_embeddings = model_config->get<size_t>("max_position_embeddings");
58+
const double rope_theta = model_config->get<double>("rope_theta");
59+
rotary_emb_ = std::make_shared<infinicore::nn::RoPE>(
60+
qk_rope_head_dim_, qk_rope_head_dim_, max_position_embeddings, rope_theta,
61+
infinicore::nn::RoPE::Algo::GPT_J, dtype, device, nullptr);
62+
63+
softmax_scale_ = 1.0f / std::sqrt(static_cast<float>(q_head_dim_));
64+
auto &config_json = model_config->get_config_json();
65+
if (config_json.contains("rope_scaling") && config_json["rope_scaling"].is_object()) {
66+
const auto &rope_scaling = config_json["rope_scaling"];
67+
const float mscale_all_dim = rope_scaling.value("mscale_all_dim", 0.0f);
68+
if (mscale_all_dim != 0.0f) {
69+
const float scaling_factor = rope_scaling.value("factor", 1.0f);
70+
const float mscale = yarn_get_mscale(scaling_factor, mscale_all_dim);
71+
softmax_scale_ *= mscale * mscale;
72+
}
73+
}
74+
75+
attn_ = std::make_shared<infinilm::layers::attention::AttentionLayer>(
76+
num_attention_heads_, q_head_dim_, softmax_scale_, num_attention_heads_, layer_idx_,
77+
kv_cache_k_scale_, kv_cache_v_scale_, attention_backend_);
78+
infinilm::layers::attention::init_kv_cache_quant_params(
79+
[this](const std::string &n, infinicore::nn::Parameter p) { this->register_parameter(n, std::move(p)); },
80+
device, kv_cache_k_scale_, kv_cache_v_scale_);
81+
}
82+
83+
infinicore::Tensor DeepseekV2Attention::position_ids_for_rope_(const infinicore::Tensor &position_ids) const {
84+
auto pos_shape = position_ids->shape();
85+
if (pos_shape.size() == 2) {
86+
return position_ids->narrow({{0, 0, 1}})->contiguous()->view({pos_shape[1]});
87+
}
88+
if (pos_shape.size() == 1) {
89+
return position_ids->contiguous();
90+
}
91+
throw std::runtime_error("DeepseekV2Attention: unexpected position_ids shape");
92+
}
93+
94+
infinicore::Tensor DeepseekV2Attention::trim_value_padding_(const infinicore::Tensor &attn_output) const {
95+
const auto shape = attn_output->shape();
96+
const size_t batch_size = shape[0];
97+
const size_t seq_len = shape[1];
98+
return attn_output->view({batch_size, seq_len, num_attention_heads_, q_head_dim_})
99+
->narrow({{3, 0, v_head_dim_}})
100+
->contiguous()
101+
->view({batch_size, seq_len, num_attention_heads_ * v_head_dim_});
102+
}
103+
104+
infinicore::Tensor DeepseekV2Attention::forward(const infinicore::Tensor &positions,
105+
const infinicore::Tensor &hidden_states) const {
106+
if (::infinilm::backends::AttentionBackend::STATIC_ATTN == attention_backend_) {
107+
return forward_static_(positions, hidden_states);
108+
}
109+
return forward_paged_(positions, hidden_states);
110+
}
111+
112+
infinicore::Tensor DeepseekV2Attention::forward_static_(const infinicore::Tensor &position_ids,
113+
const infinicore::Tensor &hidden_states) const {
114+
auto shape = hidden_states->shape();
115+
const size_t batch_size = shape[0];
116+
const size_t seq_len = shape[1];
117+
auto hidden_states_mutable = hidden_states;
118+
119+
auto q = q_proj_->forward(hidden_states_mutable)->view({batch_size, seq_len, num_attention_heads_, q_head_dim_});
120+
auto q_nope = q->narrow({{3, 0, qk_nope_head_dim_}});
121+
auto q_pe = q->narrow({{3, qk_nope_head_dim_, qk_rope_head_dim_}})->contiguous();
122+
123+
auto compressed = kv_a_proj_with_mqa_->forward(hidden_states_mutable);
124+
auto compressed_kv = compressed->narrow({{2, 0, kv_a_layernorm_->normalized_shape()}})->contiguous();
125+
auto k_pe = compressed->narrow({{2, kv_a_layernorm_->normalized_shape(), qk_rope_head_dim_}})->contiguous();
126+
127+
auto kv_norm = kv_a_layernorm_->forward(compressed_kv);
128+
auto kv = kv_b_proj_->forward(kv_norm)->view({batch_size, seq_len, num_attention_heads_, qk_nope_head_dim_ + v_head_dim_});
129+
auto k_nope = kv->narrow({{3, 0, qk_nope_head_dim_}});
130+
auto value_states = kv->narrow({{3, qk_nope_head_dim_, v_head_dim_}})->contiguous();
131+
132+
auto pos_ids = position_ids_for_rope_(position_ids);
133+
q_pe = rotary_emb_->forward(q_pe, pos_ids, true);
134+
auto k_pe_broadcast = infinicore::op::broadcast_to(k_pe->view({batch_size, seq_len, 1, qk_rope_head_dim_}),
135+
{static_cast<int64_t>(batch_size), static_cast<int64_t>(seq_len), static_cast<int64_t>(num_attention_heads_), static_cast<int64_t>(qk_rope_head_dim_)});
136+
k_pe_broadcast = rotary_emb_->forward(k_pe_broadcast, pos_ids, true);
137+
138+
auto query_states = infinicore::op::cat({q_nope, q_pe}, 3);
139+
auto key_states = infinicore::op::cat({k_nope, k_pe_broadcast}, 3);
140+
auto value_padded = infinicore::op::pad(value_states, {0, static_cast<int>(q_head_dim_ - v_head_dim_)}, "constant", 0.0);
141+
142+
auto attn_output = attn_->forward(query_states, key_states, value_padded);
143+
auto trimmed_output = trim_value_padding_(attn_output);
144+
return o_proj_->forward(trimmed_output);
145+
}
146+
147+
infinicore::Tensor DeepseekV2Attention::forward_paged_(const infinicore::Tensor &position_ids,
148+
const infinicore::Tensor &hidden_states) const {
149+
auto shape = hidden_states->shape();
150+
const size_t batch_size = shape[0];
151+
const size_t seq_len = shape[1];
152+
ASSERT_EQ(batch_size, 1);
153+
auto hidden_states_mutable = hidden_states;
154+
155+
auto q = q_proj_->forward(hidden_states_mutable)->view({seq_len, num_attention_heads_, q_head_dim_});
156+
auto q_nope = q->narrow({{2, 0, qk_nope_head_dim_}});
157+
auto q_pe = q->narrow({{2, qk_nope_head_dim_, qk_rope_head_dim_}})->contiguous();
158+
159+
auto compressed = kv_a_proj_with_mqa_->forward(hidden_states_mutable)->view({seq_len, kv_a_layernorm_->normalized_shape() + qk_rope_head_dim_});
160+
auto compressed_kv = compressed->narrow({{1, 0, kv_a_layernorm_->normalized_shape()}})->contiguous();
161+
auto k_pe = compressed->narrow({{1, kv_a_layernorm_->normalized_shape(), qk_rope_head_dim_}})->contiguous();
162+
163+
auto kv_norm = kv_a_layernorm_->forward(compressed_kv);
164+
auto kv = kv_b_proj_->forward(kv_norm)->view({seq_len, num_attention_heads_, qk_nope_head_dim_ + v_head_dim_});
165+
auto k_nope = kv->narrow({{2, 0, qk_nope_head_dim_}});
166+
auto value_states = kv->narrow({{2, qk_nope_head_dim_, v_head_dim_}})->contiguous();
167+
168+
auto pos_ids = position_ids_for_rope_(position_ids);
169+
q_pe = rotary_emb_->forward(q_pe, pos_ids, true);
170+
auto k_pe_broadcast = infinicore::op::broadcast_to(k_pe->view({seq_len, 1, qk_rope_head_dim_}),
171+
{static_cast<int64_t>(seq_len), static_cast<int64_t>(num_attention_heads_), static_cast<int64_t>(qk_rope_head_dim_)});
172+
k_pe_broadcast = rotary_emb_->forward(k_pe_broadcast, pos_ids, true);
173+
174+
auto query_states = infinicore::op::cat({q_nope, q_pe}, 2);
175+
auto key_states = infinicore::op::cat({k_nope, k_pe_broadcast}, 2);
176+
auto value_padded = infinicore::op::pad(value_states, {0, static_cast<int>(q_head_dim_ - v_head_dim_)}, "constant", 0.0);
177+
178+
auto attn_output = attn_->forward(query_states, key_states, value_padded);
179+
auto trimmed_output = trim_value_padding_(attn_output);
180+
return o_proj_->forward(trimmed_output);
181+
}
182+
183+
} // namespace infinilm::models::deepseek_v2

‎csrc/models/deepseek_v2/deepseek_v2_attention.hpp‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
#pragma once
2+
3+
#include "../../config/model_config.hpp"
4+
#include "../../layers/attention/attention.hpp"
5+
#include "../../layers/linear/linear.hpp"
6+
#include "infinicore/nn/module.hpp"
7+
#include "infinicore/nn/rmsnorm.hpp"
8+
#include "infinicore/nn/rope.hpp"
9+
#include "infinicore/tensor.hpp"
10+
11+
#include <memory>
12+
13+
namespace infinilm::models::deepseek_v2 {
14+
15+
class DeepseekV2Attention : public infinicore::nn::Module {
16+
public:
17+
DeepseekV2Attention(std::shared_ptr<infinilm::config::ModelConfig> model_config,
18+
size_t layer_idx,
19+
const infinicore::Device &device);
20+
21+
infinicore::Tensor forward(const infinicore::Tensor &positions,
22+
const infinicore::Tensor &hidden_states) const;
23+
24+
private:
25+
infinicore::Tensor forward_static_(const infinicore::Tensor &positions,
26+
const infinicore::Tensor &hidden_states) const;
27+
infinicore::Tensor forward_paged_(const infinicore::Tensor &positions,
28+
const infinicore::Tensor &hidden_states) const;
29+
infinicore::Tensor trim_value_padding_(const infinicore::Tensor &attn_output) const;
30+
infinicore::Tensor position_ids_for_rope_(const infinicore::Tensor &position_ids) const;
31+
32+
size_t layer_idx_{0};
33+
size_t hidden_size_{0};
34+
size_t num_attention_heads_{0};
35+
size_t qk_nope_head_dim_{0};
36+
size_t qk_rope_head_dim_{0};
37+
size_t q_head_dim_{0};
38+
size_t v_head_dim_{0};
39+
float softmax_scale_{1.0f};
40+
infinilm::backends::AttentionBackend attention_backend_;
41+
42+
INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, q_proj);
43+
INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, kv_a_proj_with_mqa);
44+
INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, kv_a_layernorm);
45+
INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, kv_b_proj);
46+
INFINICORE_NN_MODULE(infinilm::layers::linear::RowParallelLinear, o_proj);
47+
48+
std::shared_ptr<infinicore::nn::RoPE> rotary_emb_;
49+
std::shared_ptr<infinilm::layers::attention::AttentionLayer> attn_;
50+
infinicore::nn::Parameter kv_cache_k_scale_;
51+
infinicore::nn::Parameter kv_cache_v_scale_;
52+
};
53+
54+
} // namespace infinilm::models::deepseek_v2

‎csrc/models/deepseek_v2/deepseek_v2_decoder_layer.cpp‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
#include "deepseek_v2_decoder_layer.hpp"
2+
3+
namespace infinilm::models::deepseek_v2 {
4+
5+
DeepseekV2DecoderLayer::DeepseekV2DecoderLayer(std::shared_ptr<infinilm::config::ModelConfig> model_config,
6+
size_t layer_idx,
7+
const infinicore::Device &device) {
8+
const auto &dtype{model_config->get_dtype()};
9+
const size_t hidden_size = model_config->get<size_t>("hidden_size");
10+
const double rms_norm_eps = model_config->get<double>("rms_norm_eps");
11+
INFINICORE_NN_MODULE_INIT(input_layernorm, hidden_size, rms_norm_eps, dtype, device);
12+
INFINICORE_NN_MODULE_INIT(post_attention_layernorm, hidden_size, rms_norm_eps, dtype, device);
13+
INFINICORE_NN_MODULE_INIT(self_attn, model_config, layer_idx, device);
14+
15+
const size_t first_k_dense_replace = model_config->get_or<size_t>("first_k_dense_replace", 0);
16+
const size_t moe_layer_freq = model_config->get_or<size_t>("moe_layer_freq", 1);
17+
use_moe_ = model_config->get_or<size_t>("n_routed_experts", 0) > 0
18+
&& layer_idx >= first_k_dense_replace
19+
&& (moe_layer_freq == 0 || layer_idx % moe_layer_freq == 0);
20+
if (use_moe_) {
21+
moe_mlp_ = this->register_module<DeepseekV2MoE>("mlp", model_config, device);
22+
} else {
23+
dense_mlp_ = this->register_module<DeepseekV2MLP>("mlp", model_config, device);
24+
}
25+
}
26+
27+
std::tuple<infinicore::Tensor, infinicore::Tensor>
28+
DeepseekV2DecoderLayer::forward(const infinicore::Tensor &positions,
29+
infinicore::Tensor &hidden_states,
30+
infinicore::Tensor &residual) const {
31+
input_layernorm_->forward_inplace(hidden_states, residual);
32+
hidden_states = self_attn_->forward(positions, hidden_states);
33+
post_attention_layernorm_->forward_inplace(hidden_states, residual);
34+
hidden_states = use_moe_ ? moe_mlp_->forward(hidden_states) : dense_mlp_->forward(hidden_states);
35+
return {hidden_states, residual};
36+
}
37+
38+
} // namespace infinilm::models::deepseek_v2

‎csrc/models/deepseek_v2/deepseek_v2_decoder_layer.hpp‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
#pragma once
2+
3+
#include "../../config/model_config.hpp"
4+
#include "deepseek_v2_attention.hpp"
5+
#include "deepseek_v2_moe.hpp"
6+
#include "infinicore/device.hpp"
7+
#include "infinicore/nn/module.hpp"
8+
#include "infinicore/nn/rmsnorm.hpp"
9+
#include "infinicore/tensor.hpp"
10+
11+
#include <memory>
12+
#include <tuple>
13+
14+
namespace infinilm::models::deepseek_v2 {
15+
16+
class DeepseekV2DecoderLayer : public infinicore::nn::Module {
17+
public:
18+
DeepseekV2DecoderLayer(std::shared_ptr<infinilm::config::ModelConfig> model_config,
19+
size_t layer_idx,
20+
const infinicore::Device &device);
21+
22+
std::tuple<infinicore::Tensor, infinicore::Tensor> forward(const infinicore::Tensor &positions,
23+
infinicore::Tensor &hidden_states,
24+
infinicore::Tensor &residual) const;
25+
26+
private:
27+
INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, input_layernorm);
28+
INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, post_attention_layernorm);
29+
INFINICORE_NN_MODULE(DeepseekV2Attention, self_attn);
30+
INFINICORE_NN_MODULE(DeepseekV2MLP, dense_mlp);
31+
INFINICORE_NN_MODULE(DeepseekV2MoE, moe_mlp);
32+
bool use_moe_{false};
33+
};
34+
35+
} // namespace infinilm::models::deepseek_v2

‎csrc/models/deepseek_v2/deepseek_v2_for_causal_lm.cpp‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
#include "deepseek_v2_for_causal_lm.hpp"
2+
3+
#include "../models_registry.hpp"
4+
#include "infinicore/ops.hpp"
5+
6+
#include <stdexcept>
7+
#include <string>
8+
9+
namespace infinilm::models::deepseek_v2 {
10+
11+
DeepseekV2Model::DeepseekV2Model(std::shared_ptr<infinilm::config::ModelConfig> model_config,
12+
const infinicore::Device &device) {
13+
const auto &dtype{model_config->get_dtype()};
14+
const size_t vocab_size = model_config->get<size_t>("vocab_size");
15+
const size_t hidden_size = model_config->get<size_t>("hidden_size");
16+
const size_t num_hidden_layers = model_config->get<size_t>("num_hidden_layers");
17+
const double rms_norm_eps = model_config->get<double>("rms_norm_eps");
18+
19+
INFINICORE_NN_MODULE_INIT(embed_tokens, vocab_size, hidden_size, std::nullopt, dtype, device);
20+
layers_.reserve(num_hidden_layers);
21+
for (size_t i = 0; i < num_hidden_layers; ++i) {
22+
layers_.push_back(this->register_module<DeepseekV2DecoderLayer>("layers." + std::to_string(i), model_config, i, device));
23+
}
24+
INFINICORE_NN_MODULE_INIT(norm, hidden_size, rms_norm_eps, dtype, device);
25+
}
26+
27+
infinicore::Tensor DeepseekV2Model::forward(const infinilm::InfinilmModel::Input &input) const {
28+
auto input_ids = input.input_ids.value();
29+
auto positions = input.position_ids.value();
30+
auto hidden_states = embed_tokens_->forward(input_ids);
31+
32+
infinicore::Tensor residual;
33+
for (const auto &layer : layers_) {
34+
layer->forward(positions, hidden_states, residual);
35+
}
36+
norm_->forward_inplace(hidden_states, residual);
37+
return hidden_states;
38+
}
39+
40+
DeepseekV2ForCausalLM::DeepseekV2ForCausalLM(std::shared_ptr<infinilm::config::ModelConfig> model_config,
41+
const infinicore::Device &device) {
42+
model_config_ = model_config;
43+
const auto &dtype{model_config->get_dtype()};
44+
const size_t hidden_size = model_config->get<size_t>("hidden_size");
45+
const size_t vocab_size = model_config->get<size_t>("vocab_size");
46+
INFINICORE_NN_MODULE_INIT(model, model_config, device);
47+
INFINICORE_NN_MODULE_INIT(lm_head, hidden_size, vocab_size, false, dtype, device);
48+
}
49+
50+
infinilm::InfinilmModel::Output DeepseekV2ForCausalLM::forward(const infinilm::InfinilmModel::Input &input) const {
51+
auto hidden_states = model_->forward(input);
52+
auto logits = lm_head_->forward(hidden_states);
53+
return {logits};
54+
}
55+
56+
std::shared_ptr<infinilm::config::ModelConfig> create_deepseek_v2_model_config(std::shared_ptr<infinilm::config::ModelConfig> model_config) {
57+
const std::string model_type = model_config->get<std::string>("model_type");
58+
if ("deepseek_v2" != model_type) {
59+
throw std::runtime_error("create_deepseek_v2_model_config: model_type is not deepseek_v2");
60+
}
61+
62+
auto &config_json = model_config->get_config_json();
63+
const size_t q_head_dim = config_json.at("qk_nope_head_dim").get<size_t>() + config_json.at("qk_rope_head_dim").get<size_t>();
64+
config_json["head_dim"] = q_head_dim;
65+
config_json["num_experts"] = config_json.value("n_routed_experts", 0);
66+
config_json["mlp_bias"] = false;
67+
if (!config_json.contains("attention_output_bias")) {
68+
config_json["attention_output_bias"] = config_json.value("attention_bias", false);
69+
}
70+
if (!config_json.contains("dtype") && config_json.contains("torch_dtype")) {
71+
config_json["dtype"] = config_json["torch_dtype"];
72+
}
73+
return model_config;
74+
}
75+
76+
} // namespace infinilm::models::deepseek_v2
77+
78+
namespace {
79+
INFINILM_REGISTER_CAUSAL_LM_MODEL(
80+
deepseek_v2,
81+
infinilm::models::deepseek_v2::DeepseekV2ForCausalLM,
82+
infinilm::models::deepseek_v2::create_deepseek_v2_model_config);
83+
} // namespace

0 commit comments

Comments
 (0)