313 lines
14 KiB
C
313 lines
14 KiB
C
|
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
|
||
|
|
|
||
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
||
|
|
you may not use this file except in compliance with the License.
|
||
|
|
You may obtain a copy of the License at
|
||
|
|
|
||
|
|
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||
|
|
|
||
|
|
Unless required by applicable law or agreed to in writing, software
|
||
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
||
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
|
|
See the License for the specific language governing permissions and
|
||
|
|
limitations under the License.
|
||
|
|
==============================================================================*/
|
||
|
|
|
||
|
|
#pragma once
|
||
|
|
|
||
|
|
#include "core/framework/model/model_output.h"
|
||
|
|
#include "core/layers/common/lm_head.h"
|
||
|
|
#include "core/layers/common/qwen3_next_rms_norm.h"
|
||
|
|
#include "core/layers/common/rms_norm.h"
|
||
|
|
#include "core/layers/mlu/qwen3_5_decoder_layer.h"
|
||
|
|
#include "core/layers/qwen3_vision_layer.h"
|
||
|
|
#include "models/llm/llm_model_base.h"
|
||
|
|
#include "models/model_registry.h"
|
||
|
|
#include "models/vlm/qwen3_vl_base.h"
|
||
|
|
#include "processors/input_processor.h"
|
||
|
|
#include "processors/qwen2_vl_image_processor.h"
|
||
|
|
#include "qwen3_vl.h"
|
||
|
|
|
||
|
|
namespace xllm {
|
||
|
|
class Qwen3_5ModelImpl final
|
||
|
|
: public LlmModelImplBase<layer::Qwen3_5DecoderLayer> {
|
||
|
|
public:
|
||
|
|
Qwen3_5ModelImpl(const ModelContext& context)
|
||
|
|
: LlmModelImplBase<layer::Qwen3_5DecoderLayer>("qwen3_5",
|
||
|
|
context.get_model_args()) {
|
||
|
|
auto model_args = context.get_model_args();
|
||
|
|
auto options = context.get_tensor_options();
|
||
|
|
auto parallel_args = context.get_parallel_args();
|
||
|
|
dp_size_ = parallel_args.dp_size();
|
||
|
|
|
||
|
|
if (!mrope_section_.empty()) {
|
||
|
|
int64_t rotary_dim = static_cast<int64_t>(
|
||
|
|
model_args.head_dim() * model_args.partial_rotary_factor());
|
||
|
|
cos_sin_ = layer::rotary::get_concat_rotary_embedding(
|
||
|
|
rotary_dim,
|
||
|
|
model_args.max_position_embeddings(),
|
||
|
|
model_args.rope_theta(),
|
||
|
|
options);
|
||
|
|
}
|
||
|
|
|
||
|
|
layers_.reserve(model_args.n_layers());
|
||
|
|
rms_norm_ = register_module(
|
||
|
|
"norm",
|
||
|
|
layer::Qwen3NextRMSNorm(
|
||
|
|
model_args.hidden_size(), model_args.rms_norm_eps(), options));
|
||
|
|
embed_tokens_ =
|
||
|
|
register_module("embed_tokens", layer::WordEmbedding(context));
|
||
|
|
|
||
|
|
for (int32_t i = 0; i < model_args.n_layers(); i++) {
|
||
|
|
auto layer = layer::Qwen3_5DecoderLayer(context, i);
|
||
|
|
layers_.push_back(layer);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
void load_state_dict(const StateDict& state_dict) override {
|
||
|
|
embed_tokens_->load_state_dict(
|
||
|
|
state_dict.get_dict_with_prefix("embed_tokens."));
|
||
|
|
|
||
|
|
// call each layer's load_state_dict function
|
||
|
|
for (size_t i = 0; i < layers_.size(); i++) {
|
||
|
|
layers_[i]->load_state_dict(
|
||
|
|
state_dict.get_dict_with_prefix("layers." + std::to_string(i) + "."));
|
||
|
|
}
|
||
|
|
rms_norm_->load_state_dict(state_dict.get_dict_with_prefix("norm."));
|
||
|
|
}
|
||
|
|
|
||
|
|
std::pair<torch::Tensor, torch::Tensor> apply_mrope(
|
||
|
|
const torch::Tensor positions) override {
|
||
|
|
auto target_cos_sin = cos_sin_.index({positions});
|
||
|
|
auto target_cos_sin_chunks = target_cos_sin.chunk(/*chunks=*/2, /*dim=*/-1);
|
||
|
|
auto cos_pos = target_cos_sin_chunks[0].contiguous();
|
||
|
|
auto sin_pos = target_cos_sin_chunks[1].contiguous();
|
||
|
|
auto apply = [this](torch::Tensor x) {
|
||
|
|
auto freqs_t = x[0].clone();
|
||
|
|
int64_t mrop_length = static_cast<int64_t>(freqs_t.size(-1) / 2);
|
||
|
|
|
||
|
|
for (int32_t dim_idx = 1; dim_idx <= 2; ++dim_idx) {
|
||
|
|
int64_t offset = dim_idx;
|
||
|
|
int64_t section_len = mrope_section_[dim_idx];
|
||
|
|
int64_t length = section_len * 3;
|
||
|
|
|
||
|
|
auto idx_first_half = torch::arange(offset, length, 3, torch::kLong);
|
||
|
|
auto idx_second_half = torch::arange(
|
||
|
|
offset + mrop_length, length + mrop_length, 3, torch::kLong);
|
||
|
|
|
||
|
|
auto idx_tensor =
|
||
|
|
torch::cat({idx_first_half, idx_second_half}, 0).to(x.device());
|
||
|
|
auto src = x[dim_idx].index_select(-1, idx_tensor);
|
||
|
|
freqs_t.index_copy_(-1, idx_tensor, src);
|
||
|
|
}
|
||
|
|
return freqs_t;
|
||
|
|
};
|
||
|
|
cos_pos = apply(cos_pos.reshape({positions.size(0), -1, cos_pos.size(-1)}));
|
||
|
|
sin_pos = apply(sin_pos.reshape({positions.size(0), -1, sin_pos.size(-1)}));
|
||
|
|
return std::make_pair(cos_pos, sin_pos);
|
||
|
|
}
|
||
|
|
|
||
|
|
virtual ModelOutput forward(torch::Tensor tokens,
|
||
|
|
torch::Tensor positions,
|
||
|
|
std::vector<KVCache>& kv_caches,
|
||
|
|
const ModelInputParams& input_params) {
|
||
|
|
ModelInputParams& input_params_new =
|
||
|
|
const_cast<ModelInputParams&>(input_params);
|
||
|
|
std::vector<torch::Tensor> deep_stacks;
|
||
|
|
|
||
|
|
if (dp_size_ > 1) {
|
||
|
|
if (tokens.numel() == 0) {
|
||
|
|
tokens = torch::tensor({1}).to(torch::kInt32).to(tokens.device());
|
||
|
|
positions = torch::tensor({1}).to(torch::kInt32).to(positions.device());
|
||
|
|
}
|
||
|
|
auto& dp_token_nums = input_params_new.dp_global_token_nums;
|
||
|
|
std::replace(dp_token_nums.begin(), dp_token_nums.end(), 0, 1);
|
||
|
|
}
|
||
|
|
|
||
|
|
auto inputs_embeds = input_params.input_embedding;
|
||
|
|
torch::Tensor h;
|
||
|
|
if (inputs_embeds.defined()) {
|
||
|
|
h = inputs_embeds;
|
||
|
|
} else {
|
||
|
|
h = embed_tokens_(tokens);
|
||
|
|
}
|
||
|
|
|
||
|
|
if (!input_params_new.attn_metadata) {
|
||
|
|
input_params_new.attn_metadata =
|
||
|
|
std::make_shared<layer::AttentionMetadata>(
|
||
|
|
get_attention_metadata(input_params_new, h));
|
||
|
|
}
|
||
|
|
|
||
|
|
auto& attn_metadata = *(input_params_new.attn_metadata);
|
||
|
|
bool only_prefill =
|
||
|
|
(attn_metadata.is_prefill || attn_metadata.is_chunked_prefill);
|
||
|
|
if (positions.dim() == 2 && only_prefill && !mrope_section_.empty()) {
|
||
|
|
std::tie(attn_metadata.mrope_cos, attn_metadata.mrope_sin) =
|
||
|
|
apply_mrope(positions);
|
||
|
|
}
|
||
|
|
|
||
|
|
std::optional<torch::Tensor> residual;
|
||
|
|
for (size_t i = 0; i < layers_.size(); i++) {
|
||
|
|
auto& layer = layers_[i];
|
||
|
|
h = layer(h,
|
||
|
|
residual,
|
||
|
|
positions,
|
||
|
|
attn_metadata,
|
||
|
|
kv_caches[i],
|
||
|
|
input_params_new);
|
||
|
|
}
|
||
|
|
if (residual.has_value()) {
|
||
|
|
h = h + residual.value();
|
||
|
|
}
|
||
|
|
auto hidden_states = std::get<0>(rms_norm_(h));
|
||
|
|
return ModelOutput(hidden_states);
|
||
|
|
}
|
||
|
|
|
||
|
|
private:
|
||
|
|
int32_t dp_size_ = 1;
|
||
|
|
layer::Qwen3NextRMSNorm rms_norm_{nullptr};
|
||
|
|
layer::AttentionMetadata get_attention_metadata(
|
||
|
|
const ModelInputParams& params,
|
||
|
|
const torch::Tensor& h) {
|
||
|
|
auto attn_metadata = layer::AttentionMetadataBuilder::build(params, false);
|
||
|
|
// TODO: support linear attention
|
||
|
|
return attn_metadata;
|
||
|
|
}
|
||
|
|
};
|
||
|
|
TORCH_MODULE(Qwen3_5Model);
|
||
|
|
|
||
|
|
class Qwen3_5ForCausalLMImpl : public LlmForCausalLMImplBase<Qwen3_5Model> {
|
||
|
|
public:
|
||
|
|
Qwen3_5ForCausalLMImpl(const ModelContext& context)
|
||
|
|
: LlmForCausalLMImplBase<Qwen3_5Model>(context) {}
|
||
|
|
|
||
|
|
torch::Tensor pooler(const torch::Tensor& hidden_states,
|
||
|
|
const torch::Tensor& seleted_idxes) {
|
||
|
|
auto h = hidden_states;
|
||
|
|
if (seleted_idxes.defined()) {
|
||
|
|
h = h.index_select(/*dim=*/0, seleted_idxes);
|
||
|
|
}
|
||
|
|
namespace F = torch::nn::functional;
|
||
|
|
return F::normalize(h, F::NormalizeFuncOptions().p(2).dim(1));
|
||
|
|
}
|
||
|
|
};
|
||
|
|
TORCH_MODULE(Qwen3_5ForCausalLM);
|
||
|
|
|
||
|
|
using Qwen3_5ForConditionalGenerationImpl =
|
||
|
|
Qwen3VLForConditionalGenerationBase<Qwen3_VisionTransformer,
|
||
|
|
Qwen3_5ForCausalLM>;
|
||
|
|
TORCH_MODULE(Qwen3_5ForConditionalGeneration);
|
||
|
|
|
||
|
|
#define LOAD_QWEN3_5_COMMON_ARGS() \
|
||
|
|
LOAD_ARG_OR(model_type, "model_type", "qwen3_5"); \
|
||
|
|
LOAD_ARG_OR(dtype, "text_config.dtype", "bfloat16"); \
|
||
|
|
LOAD_ARG_OR(vocab_size, "text_config.vocab_size", 248320); \
|
||
|
|
LOAD_ARG_OR(hidden_size, "text_config.hidden_size", 5120); \
|
||
|
|
LOAD_ARG_OR(hidden_act, "text_config.hidden_act", "silu"); \
|
||
|
|
LOAD_ARG_OR(intermediate_size, "text_config.intermediate_size", 17408); \
|
||
|
|
LOAD_ARG_OR(n_layers, "text_config.num_hidden_layers", 64); \
|
||
|
|
LOAD_ARG_OR(n_heads, "text_config.num_attention_heads", 24); \
|
||
|
|
LOAD_ARG(n_kv_heads, "text_config.num_key_value_heads"); \
|
||
|
|
LOAD_ARG_OR( \
|
||
|
|
max_position_embeddings, "text_config.max_position_embeddings", 262144); \
|
||
|
|
LOAD_ARG_OR(rms_norm_eps, "text_config.rms_norm_eps", 1e-6); \
|
||
|
|
LOAD_ARG_OR(eos_token_id, "text_config.eos_token_id", 248044); \
|
||
|
|
LOAD_ARG_OR( \
|
||
|
|
rope_theta, "text_config.rope_parameters.rope_theta", 10000000.0f); \
|
||
|
|
LOAD_ARG_OR(head_dim, "text_config.head_dim", 256); \
|
||
|
|
LOAD_ARG_OR(tie_word_embeddings, "tie_word_embeddings", false); \
|
||
|
|
LOAD_ARG(layer_types, "text_config.layer_types"); \
|
||
|
|
LOAD_ARG_OR( \
|
||
|
|
linear_conv_kernel_dim, "text_config.linear_conv_kernel_dim", 4); \
|
||
|
|
LOAD_ARG_OR(linear_key_head_dim, "text_config.linear_key_head_dim", 128); \
|
||
|
|
LOAD_ARG_OR( \
|
||
|
|
linear_value_head_dim, "text_config.linear_value_head_dim", 128); \
|
||
|
|
LOAD_ARG_OR(linear_num_key_heads, "text_config.linear_num_key_heads", 16); \
|
||
|
|
LOAD_ARG_OR( \
|
||
|
|
linear_num_value_heads, "text_config.linear_num_value_heads", 48); \
|
||
|
|
LOAD_ARG_OR( \
|
||
|
|
full_attention_interval, "text_config.full_attention_interval", 4); \
|
||
|
|
LOAD_ARG_OR(attn_output_gate, "text_config.attn_output_gate", false); \
|
||
|
|
LOAD_ARG_OR( \
|
||
|
|
num_nextn_predict_layers, "text_config.mtp_num_hidden_layers", 0); \
|
||
|
|
LOAD_ARG_OR(attention_bias, "text_config.attention_bias", false); \
|
||
|
|
LOAD_ARG_OR(attention_dropout, "text_config.attention_dropout", 0.0f); \
|
||
|
|
LOAD_ARG_OR(initializer_range, "text_config.initializer_range", 0.02f); \
|
||
|
|
LOAD_ARG_OR( \
|
||
|
|
mlp_only_layers, "text_config.mlp_only_layers", std::vector<int32_t>()); \
|
||
|
|
LOAD_ARG(rope_scaling_mrope_section, \
|
||
|
|
"text_config.rope_parameters.mrope_section"); \
|
||
|
|
LOAD_ARG_OR(rope_scaling_rope_type, \
|
||
|
|
"text_config.rope_parameters.rope_type", \
|
||
|
|
"default"); \
|
||
|
|
LOAD_ARG_OR(partial_rotary_factor, \
|
||
|
|
"text_config.rope_parameters.partial_rotary_factor", \
|
||
|
|
0.25f)
|
||
|
|
|
||
|
|
#define LOAD_QWEN3_5_VISION_ARGS() \
|
||
|
|
LOAD_ARG_OR(image_token_id, "image_token_id", 248056); \
|
||
|
|
LOAD_ARG_OR(video_token_id, "video_token_id", 248057); \
|
||
|
|
LOAD_ARG_OR(vision_start_token_id, "vision_start_token_id", 248053); \
|
||
|
|
LOAD_ARG_OR(vision_end_token_id, "vision_end_token_id", 248054); \
|
||
|
|
LOAD_ARG(mm_deepstack_visual_indexes, \
|
||
|
|
"vision_config.deepstack_visual_indexes"); \
|
||
|
|
LOAD_ARG_OR(mm_num_hidden_layers, "vision_config.depth", 27); \
|
||
|
|
LOAD_ARG_OR(mm_hidden_act, "vision_config.hidden_act", "gelu_pytorch_tanh"); \
|
||
|
|
LOAD_ARG_OR(mm_hidden_size, "vision_config.hidden_size", 1152); \
|
||
|
|
LOAD_ARG_OR(mm_num_channels, "vision_config.in_channels", 3); \
|
||
|
|
LOAD_ARG_OR(mm_initializer_range, "vision_config.initializer_range", 0.02f); \
|
||
|
|
LOAD_ARG_OR(mm_intermediate_size, "vision_config.intermediate_size", 4304); \
|
||
|
|
LOAD_ARG_OR(mm_num_attention_heads, "vision_config.num_heads", 16); \
|
||
|
|
LOAD_ARG_OR(mm_num_position_embeddings, \
|
||
|
|
"vision_config.num_position_embeddings", \
|
||
|
|
2304); \
|
||
|
|
LOAD_ARG_OR(mm_projection_dim, "vision_config.out_hidden_size", 5120); \
|
||
|
|
LOAD_ARG_OR(mm_patch_size, "vision_config.patch_size", 16); \
|
||
|
|
LOAD_ARG_OR(mm_spatial_merge_size, "vision_config.spatial_merge_size", 2); \
|
||
|
|
LOAD_ARG_OR(mm_temporal_patch_size, "vision_config.temporal_patch_size", 2); \
|
||
|
|
LOAD_ARG_OR_FUNC(mm_head_dim, "head_dim", [&] { \
|
||
|
|
return args->mm_hidden_size() / args->mm_num_attention_heads(); \
|
||
|
|
}); \
|
||
|
|
LOAD_ARG_OR( \
|
||
|
|
rope_scaling_rope_type, "vision_config.rope_scaling.type", "mrope")
|
||
|
|
|
||
|
|
REGISTER_INPUT_PROCESSOR(qwen3_5, Qwen2_5_VLInputProcessor);
|
||
|
|
REGISTER_CAUSAL_VLM_MODEL(qwen3_5, Qwen3_5ForConditionalGeneration);
|
||
|
|
REGISTER_IMAGE_PROCESSOR(qwen3_5, Qwen2VLImageProcessor);
|
||
|
|
REGISTER_MODEL_ARGS(qwen3_5, [&] {
|
||
|
|
LOAD_QWEN3_5_COMMON_ARGS();
|
||
|
|
LOAD_QWEN3_5_VISION_ARGS();
|
||
|
|
SET_ARG(stop_token_ids, std::unordered_set<int32_t>({args->eos_token_id()}));
|
||
|
|
});
|
||
|
|
|
||
|
|
REGISTER_INPUT_PROCESSOR(qwen3_5_moe, Qwen2_5_VLInputProcessor);
|
||
|
|
REGISTER_CAUSAL_VLM_MODEL(qwen3_5_moe, Qwen3_5ForConditionalGeneration);
|
||
|
|
REGISTER_IMAGE_PROCESSOR(qwen3_5_moe, Qwen2VLImageProcessor);
|
||
|
|
REGISTER_MODEL_ARGS(qwen3_5_moe, [&] {
|
||
|
|
LOAD_QWEN3_5_COMMON_ARGS();
|
||
|
|
LOAD_QWEN3_5_VISION_ARGS();
|
||
|
|
|
||
|
|
LOAD_ARG_OR(decoder_sparse_step, "text_config.decoder_sparse_step", 1);
|
||
|
|
LOAD_ARG_OR(moe_intermediate_size, "text_config.moe_intermediate_size", 512);
|
||
|
|
LOAD_ARG_OR(num_experts, "text_config.num_experts", 512);
|
||
|
|
LOAD_ARG_OR(num_experts_per_tok, "text_config.num_experts_per_tok", 10);
|
||
|
|
LOAD_ARG_OR(shared_expert_intermediate_size,
|
||
|
|
"text_config.shared_expert_intermediate_size",
|
||
|
|
512);
|
||
|
|
LOAD_ARG_OR(norm_topk_prob, "text_config.norm_topk_prob", true);
|
||
|
|
|
||
|
|
LOAD_ARG_OR(
|
||
|
|
n_routed_experts, "text_config.n_routed_experts", args->num_experts());
|
||
|
|
SET_ARG(n_shared_experts,
|
||
|
|
args->shared_expert_intermediate_size() > 0 ? 1 : 0);
|
||
|
|
SET_ARG(scoring_func, "softmax");
|
||
|
|
SET_ARG(topk_method, "");
|
||
|
|
SET_ARG(n_group, -1);
|
||
|
|
SET_ARG(topk_group, 0);
|
||
|
|
SET_ARG(routed_scaling_factor, 1.0f);
|
||
|
|
|
||
|
|
SET_ARG(stop_token_ids, std::unordered_set<int32_t>({args->eos_token_id()}));
|
||
|
|
});
|
||
|
|
|
||
|
|
} // namespace xllm
|