diff --git a/csrc/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp index 8f7f789b..dee3123c 100644 --- a/csrc/engine/compiler/paged_compiler.cpp +++ b/csrc/engine/compiler/paged_compiler.cpp @@ -62,6 +62,14 @@ void PagedCompiler::compile() { auto &forward_context = infinilm::global_state::get_forward_context(); const bool has_mamba_state = has_mamba_cache(forward_context); + const auto &model_config = model_->get_model_config(); + const size_t position_id_axes = model_config == nullptr + ? 1 + : model_config->get_or("position_id_axes", 1); + if (position_id_axes == 0) { + throw std::runtime_error("PagedCompiler: position_id_axes must be positive"); + } + size_t max_batch_size = *std::max_element(decode_batch_sizes_.begin(), decode_batch_sizes_.end()); compiled_map_decode_.clear(); block_tables_holder_ = infinicore::Tensor::empty( @@ -71,7 +79,11 @@ void PagedCompiler::compile() { auto make_decode_input = [&](size_t b) { InfinilmModel::Input input; input.input_ids = infinicore::Tensor::empty({1, b}, infinicore::DataType::I64, infinicore::context::getDevice()); - input.position_ids = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice()); + input.position_ids = infinicore::Tensor::empty( + position_id_axes > 1 + ? std::vector{position_id_axes, b} + : std::vector{b}, + infinicore::DataType::I64, infinicore::context::getDevice()); input.total_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::I32, infinicore::context::getDevice()); set_zeros(input.input_ids.value()); set_zeros(input.position_ids.value()); diff --git a/csrc/models/infinilm_model.hpp b/csrc/models/infinilm_model.hpp index ac994fd6..27275ee2 100644 --- a/csrc/models/infinilm_model.hpp +++ b/csrc/models/infinilm_model.hpp @@ -70,6 +70,9 @@ class InfinilmModel : public infinicore::nn::Module { virtual const cache::CacheConfig *get_cache_config() const { return cache_config_.get(); } + const std::shared_ptr &get_model_config() const { + return model_config_; + } void process_weights_after_loading(); void reset_runtime_state() const; diff --git a/csrc/models/qwen3_5/qwen3_5_attention.cpp b/csrc/models/qwen3_5/qwen3_5_attention.cpp index 3db8aba6..e7a47f4f 100644 --- a/csrc/models/qwen3_5/qwen3_5_attention.cpp +++ b/csrc/models/qwen3_5/qwen3_5_attention.cpp @@ -32,12 +32,15 @@ Qwen35Attention::Qwen35Attention(std::shared_ptr const engine::distributed::RankInfo &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); int tp_rank = infinilm::global_state::get_tensor_model_parallel_rank(); int tp_size = infinilm::global_state::get_tensor_model_parallel_world_size(); - if ((total_num_kv_heads < tp_size) || (0 != (total_num_kv_heads % tp_size))) { + if (total_num_kv_heads >= static_cast(tp_size) + && total_num_kv_heads % tp_size != 0) { throw std::runtime_error("infinilm::models::qwen3_5::Qwen35Attention: num_key_value_heads must be divisible by tp_size"); } num_attention_heads_ = total_num_heads / tp_size; - num_key_value_heads_ = total_num_kv_heads / tp_size; + num_key_value_heads_ = total_num_kv_heads < static_cast(tp_size) + ? 1 + : total_num_kv_heads / tp_size; auto quantization_method = model_config->get_quantization_method(); auto register_fn = [this](const std::string &n, infinicore::nn::Parameter p) { this->register_parameter(n, std::move(p)); }; diff --git a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp index aece19fd..72fe1a87 100644 --- a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp +++ b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp @@ -1,32 +1,33 @@ #include "qwen3_5_for_causal_lm.hpp" -#include "../../global_state/global_state.hpp" #include "../models_registry.hpp" -#include "../qwen3_next/qwen3_next_for_causal_lm.hpp" #include #include #include namespace infinilm::models::qwen3_5 { -Qwen35ForCausalLM::Qwen35ForCausalLM(std::shared_ptr model_config, - const infinicore::Device &device) { +Qwen35ForCausalLM::Qwen35ForCausalLM( + std::shared_ptr model_config, + const infinicore::Device &device) { model_config_ = model_config; - size_t hidden_size = model_config->get("hidden_size"); - size_t vocab_size = model_config->get("vocab_size"); - const auto &dtype{model_config->get_dtype()}; + const size_t hidden_size = model_config->get("hidden_size"); + const size_t vocab_size = model_config->get("vocab_size"); + const auto &dtype = model_config->get_dtype(); INFINICORE_NN_MODULE_INIT(model, model_config, device); - INFINICORE_NN_MODULE_INIT(lm_head, hidden_size, vocab_size, false, dtype, device); + INFINICORE_NN_MODULE_INIT( + lm_head, hidden_size, vocab_size, false, dtype, device); } -infinilm::InfinilmModel::Output Qwen35ForCausalLM::forward(const infinilm::InfinilmModel::Input &input) const { +InfinilmModel::Output Qwen35ForCausalLM::forward( + const InfinilmModel::Input &input) const { auto hidden_states = model_->forward(input); - auto logits = lm_head_->forward(hidden_states); - return {logits}; + return {lm_head_->forward(hidden_states)}; } -void Qwen35ForCausalLM::reset_cache(const cache::CacheConfig *cache_config) { +void Qwen35ForCausalLM::reset_cache( + const cache::CacheConfig *cache_config) { if (cache_config == nullptr) { cache_config_.reset(); } else { @@ -35,12 +36,7 @@ void Qwen35ForCausalLM::reset_cache(const cache::CacheConfig *cache_config) { model_->reset_cache(cache_config); } -std::shared_ptr create_qwen3_5_model_config(std::shared_ptr model_config) { - const std::string model_type = model_config->get("model_type"); - if ("qwen3_5" != model_type) { - throw std::runtime_error("infinilm::models::qwen3_5::create_qwen3_next_model_config: model_type is not qwen3_5"); - } - +std::shared_ptr prepare_qwen3_5_model_config(std::shared_ptr model_config) { nlohmann::json &config_json = model_config->get_config_json(); if (config_json.contains("text_config") && config_json["text_config"].is_object()) { const nlohmann::json &text_config_json = config_json["text_config"]; @@ -53,32 +49,56 @@ std::shared_ptr create_qwen3_5_model_config(std:: config_json["dtype"] = config_json["torch_dtype"]; } } + if (!config_json.contains("position_id_axes")) { + size_t position_id_axes = 1; + if (config_json.contains("rope_parameters") + && config_json["rope_parameters"].is_object()) { + const auto &rope_parameters = config_json["rope_parameters"]; + if (rope_parameters.contains("mrope_section") + && rope_parameters["mrope_section"].is_array() + && !rope_parameters["mrope_section"].empty()) { + position_id_axes = rope_parameters["mrope_section"].size(); + } + } + config_json["position_id_axes"] = position_id_axes; + } if (!config_json.contains("rope_theta") && config_json.contains("rope_parameters") && config_json["rope_parameters"].is_object() && config_json["rope_parameters"].contains("rope_theta")) { - // TODO: This is only a temporary loader shim. Qwen3.6 uses mRoPE, - // which needs proper support in InfiniCore instead of treating it as - // plain RoPE through a top-level rope_theta. + // Normalize the nested HuggingFace field for the Qwen3.5 attention module. config_json["rope_theta"] = config_json["rope_parameters"]["rope_theta"]; } if (!config_json.contains("partial_rotary_factor") && config_json.contains("rope_parameters") && config_json["rope_parameters"].is_object() && config_json["rope_parameters"].contains("partial_rotary_factor")) { config_json["partial_rotary_factor"] = config_json["rope_parameters"]["partial_rotary_factor"]; } if (!config_json.contains("layer_types")) { - size_t full_attention_interval = model_config->get("full_attention_interval"); - size_t num_hidden_layers = model_config->get("num_hidden_layers"); + const size_t full_attention_interval = model_config->get("full_attention_interval"); + if (full_attention_interval == 0) { + throw std::runtime_error("Qwen3.5 full_attention_interval must be positive"); + } + const size_t num_hidden_layers = model_config->get("num_hidden_layers"); std::vector layer_types; layer_types.reserve(num_hidden_layers); - for (size_t i = 0; i < num_hidden_layers; i++) { - layer_types.push_back(bool((i + 1) % full_attention_interval) ? "linear_attention" : "full_attention"); + for (size_t i = 0; i < num_hidden_layers; ++i) { + layer_types.push_back( + (i + 1) % full_attention_interval == 0 + ? "full_attention" + : "linear_attention"); } - config_json["layer_types"] = layer_types; + config_json["layer_types"] = std::move(layer_types); } - if (!config_json.contains("attention_bias")) { config_json["attention_bias"] = false; } return model_config; } +std::shared_ptr create_qwen3_5_model_config(std::shared_ptr model_config) { + const std::string model_type = model_config->get("model_type"); + if ("qwen3_5" != model_type) { + throw std::runtime_error("infinilm::models::qwen3_5::create_qwen3_5_model_config: model_type is not qwen3_5"); + } + return prepare_qwen3_5_model_config(model_config); +} + } // namespace infinilm::models::qwen3_5 namespace { diff --git a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp index a7c751e3..51211481 100644 --- a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp +++ b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp @@ -20,6 +20,9 @@ class Qwen35ForCausalLM : public InfinilmModel { INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, lm_head); }; +std::shared_ptr prepare_qwen3_5_model_config( + std::shared_ptr model_config); + std::shared_ptr create_qwen3_5_model_config(std::shared_ptr model_config); } // namespace infinilm::models::qwen3_5 diff --git a/csrc/models/qwen3_5/qwen3_5_model.cpp b/csrc/models/qwen3_5/qwen3_5_model.cpp index 82c8b0d4..0f66990a 100644 --- a/csrc/models/qwen3_5/qwen3_5_model.cpp +++ b/csrc/models/qwen3_5/qwen3_5_model.cpp @@ -31,8 +31,8 @@ std::vector tensor_to_i32_vector(const infinicore::Tensor &tensor) { } // namespace -Qwen35Model::Qwen35Model(std::shared_ptr model_config, - const infinicore::Device &device) +Qwen35ModelBase::Qwen35ModelBase(std::shared_ptr model_config, + const infinicore::Device &device) : model_config_(model_config) { const auto &dtype{model_config->get_dtype()}; nlohmann::json &config_json = model_config->get_config_json(); @@ -40,11 +40,10 @@ Qwen35Model::Qwen35Model(std::shared_ptr model_co if (config_json.contains("vision_config") && !config_json["vision_config"].is_null()) { INFINICORE_NN_MODULE_INIT(visual, config_json["vision_config"], dtype, device); } - INFINICORE_NN_MODULE_INIT(language_model, model_config, device); } -void Qwen35Model::replace_image_embeddings(infinicore::Tensor &inputs_embeds, - const InfinilmModel::Input &input) const { +void Qwen35ModelBase::replace_image_embeddings(infinicore::Tensor &inputs_embeds, + const InfinilmModel::Input &input) const { if (!input.pixel_values.has_value() || input.pixel_values->empty()) { return; } @@ -107,28 +106,19 @@ void Qwen35Model::replace_image_embeddings(infinicore::Tensor &inputs_embeds, } } -infinicore::Tensor Qwen35Model::forward(const InfinilmModel::Input &input) const { - if (input.pixel_values.has_value() && !input.pixel_values->empty()) { - auto inputs_embeds = language_model_->embed_tokens(input.input_ids.value()); - replace_image_embeddings(inputs_embeds, input); - return language_model_->forward_embeds(inputs_embeds, input.position_ids.value()); - } - return language_model_->forward(input); -} - -void Qwen35Model::reset_cache(const cache::CacheConfig *cache_config) { - if (nullptr == cache_config) { - return; - } - +void Qwen35ModelBase::reset_cache(const cache::CacheConfig *cache_config) { auto &forward_context = infinilm::global_state::get_forward_context(); forward_context.kv_cache_vec.clear(); forward_context.conv_state_vec.clear(); forward_context.ssm_state_vec.clear(); + if (nullptr == cache_config) { + return; + } const backends::AttentionBackend attention_backend = infinilm::global_state::get_infinilm_config().attention_backend; - auto cache_vectors = infinilm::models::qwen3_next::qwen3_next_allocate_cache_tensors(cache_config, model_config_, attention_backend); + auto cache_vectors = infinilm::models::qwen3_next::qwen3_next_allocate_cache_tensors( + cache_config, model_config_, attention_backend); forward_context.kv_cache_vec = std::move(cache_vectors.kv_cache_tensors); forward_context.conv_state_vec = std::move(cache_vectors.conv_state_tensors); forward_context.ssm_state_vec = std::move(cache_vectors.ssm_state_tensors); diff --git a/csrc/models/qwen3_5/qwen3_5_model.hpp b/csrc/models/qwen3_5/qwen3_5_model.hpp index bea1b78e..e9f2080b 100644 --- a/csrc/models/qwen3_5/qwen3_5_model.hpp +++ b/csrc/models/qwen3_5/qwen3_5_model.hpp @@ -10,24 +10,46 @@ namespace infinilm::models::qwen3_5 { using Qwen35LanguageModel = infinilm::layers::causal_lm_templates::TextModel; -class Qwen35Model : public infinicore::nn::Module { +class Qwen35ModelBase : public infinicore::nn::Module { public: - Qwen35Model(std::shared_ptr model_config, - const infinicore::Device &device); - - infinicore::Tensor forward(const InfinilmModel::Input &input) const; + Qwen35ModelBase(std::shared_ptr model_config, + const infinicore::Device &device); void reset_cache(const cache::CacheConfig *cache_config); -private: +protected: void replace_image_embeddings(infinicore::Tensor &inputs_embeds, const infinilm::InfinilmModel::Input &input) const; -protected: INFINICORE_NN_MODULE(Qwen35VisionModel, visual); - INFINICORE_NN_MODULE(Qwen35LanguageModel, language_model); - std::shared_ptr model_config_; }; +template +class Qwen35ModelTemplate : public Qwen35ModelBase { +public: + Qwen35ModelTemplate( + std::shared_ptr model_config, + const infinicore::Device &device) + : Qwen35ModelBase(model_config, device) { + language_model_ = this->register_module( + "language_model", model_config, device); + } + + infinicore::Tensor forward(const InfinilmModel::Input &input) const { + if (input.pixel_values.has_value() && !input.pixel_values->empty()) { + auto inputs_embeds = language_model_->embed_tokens(input.input_ids.value()); + replace_image_embeddings(inputs_embeds, input); + return language_model_->forward_embeds( + inputs_embeds, input.position_ids.value()); + } + return language_model_->forward(input); + } + +protected: + INFINICORE_NN_MODULE(LanguageModel, language_model); +}; + +using Qwen35Model = Qwen35ModelTemplate; + } // namespace infinilm::models::qwen3_5 diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.cpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.cpp new file mode 100644 index 00000000..aa513284 --- /dev/null +++ b/csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.cpp @@ -0,0 +1,71 @@ +#include "qwen3_5_moe_decoder_layer.hpp" + +#include "infinicore/ops.hpp" + +#include +#include +#include + +namespace infinilm::models::qwen3_5_moe { + +Qwen35MoeDecoderLayer::Qwen35MoeDecoderLayer( + std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device) + : layer_idx_(layer_idx) { + const auto &dtype = model_config->get_dtype(); + const size_t hidden_size = model_config->get("hidden_size"); + const double rms_norm_eps = model_config->get("rms_norm_eps"); + + INFINICORE_NN_MODULE_INIT(input_layernorm, hidden_size, rms_norm_eps, dtype, device); + INFINICORE_NN_MODULE_INIT(post_attention_layernorm, hidden_size, rms_norm_eps, dtype, device); + INFINICORE_NN_MODULE_INIT(mlp, model_config, layer_idx, device); + + const auto layer_types = model_config->get>("layer_types"); + layer_type_ = layer_types.at(layer_idx); + if ("linear_attention" == layer_type_) { + INFINICORE_NN_MODULE_INIT(linear_attn, model_config, layer_idx, device); + } else if ("full_attention" == layer_type_) { + INFINICORE_NN_MODULE_INIT(self_attn, model_config, layer_idx, device); + } else { + throw std::runtime_error( + "Qwen35MoeDecoderLayer: unsupported layer_type '" + layer_type_ + + "' for layer " + std::to_string(layer_idx)); + } +} + +std::tuple Qwen35MoeDecoderLayer::forward( + const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) { + input_layernorm_->forward_inplace(hidden_states, residual); + if ("linear_attention" == layer_type_) { + hidden_states = linear_attn_->forward(hidden_states); + } else { + hidden_states = self_attn_->forward(positions, hidden_states); + } + + post_attention_layernorm_->forward_inplace(hidden_states, residual); + hidden_states = mlp_->forward(hidden_states); + return std::make_tuple(hidden_states, residual); +} + +infinicore::Tensor Qwen35MoeDecoderLayer::forward( + const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states) { + auto residual = hidden_states; + hidden_states = input_layernorm_->forward(hidden_states); + if ("linear_attention" == layer_type_) { + hidden_states = linear_attn_->forward(hidden_states); + } else { + hidden_states = self_attn_->forward(positions, hidden_states); + } + hidden_states = infinicore::op::add(residual, hidden_states); + + residual = hidden_states; + hidden_states = post_attention_layernorm_->forward(hidden_states); + hidden_states = mlp_->forward(hidden_states); + return infinicore::op::add(residual, hidden_states); +} + +} // namespace infinilm::models::qwen3_5_moe diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.hpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.hpp new file mode 100644 index 00000000..50d279e0 --- /dev/null +++ b/csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.hpp @@ -0,0 +1,40 @@ +#pragma once + +#include "../qwen3_5/qwen3_5_attention.hpp" +#include "../qwen3_next/qwen3_next_gated_deltanet.hpp" +#include "../qwen3_next/qwen3_next_sparse_moe_block.hpp" + +#include +#include + +namespace infinilm::models::qwen3_5_moe { + +class Qwen35MoeDecoderLayer : public infinicore::nn::Module { +public: + Qwen35MoeDecoderLayer(std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device); + + std::tuple forward( + const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual); + + infinicore::Tensor forward(const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states); + + size_t layer_idx() const { return layer_idx_; } + +protected: + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, input_layernorm); + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, post_attention_layernorm); + INFINICORE_NN_MODULE(qwen3_5::Qwen35Attention, self_attn); + INFINICORE_NN_MODULE(qwen3_next::Qwen3NextGatedDeltaNet, linear_attn); + INFINICORE_NN_MODULE(qwen3_next::Qwen3NextSparseMoeBlock, mlp); + +private: + size_t layer_idx_; + std::string layer_type_; +}; + +} // namespace infinilm::models::qwen3_5_moe diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.cpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.cpp new file mode 100644 index 00000000..99dc16d9 --- /dev/null +++ b/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.cpp @@ -0,0 +1,34 @@ +#include "qwen3_5_moe_for_causal_lm.hpp" + +#include "../models_registry.hpp" +#include "../qwen3_5/qwen3_5_for_causal_lm.hpp" + +#include +#include + +namespace infinilm::models::qwen3_5_moe { + +std::shared_ptr create_qwen3_5_moe_model_config( + std::shared_ptr model_config) { + const std::string model_type = model_config->get("model_type"); + if ("qwen3_5_moe" != model_type) { + throw std::runtime_error( + "create_qwen3_5_moe_model_config: model_type is not qwen3_5_moe"); + } + + model_config = qwen3_5::prepare_qwen3_5_model_config(model_config); + auto &config_json = model_config->get_config_json(); + if (!config_json.contains("norm_topk_prob")) { + config_json["norm_topk_prob"] = true; + } + return model_config; +} + +} // namespace infinilm::models::qwen3_5_moe + +namespace { +INFINILM_REGISTER_CAUSAL_LM_MODEL( + qwen3_5_moe, + infinilm::models::qwen3_5_moe::Qwen35MoeForConditionalGeneration, + infinilm::models::qwen3_5_moe::create_qwen3_5_moe_model_config); +} // namespace diff --git a/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.hpp b/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.hpp new file mode 100644 index 00000000..099c544e --- /dev/null +++ b/csrc/models/qwen3_5_moe/qwen3_5_moe_for_causal_lm.hpp @@ -0,0 +1,33 @@ +#pragma once + +#include "../../layers/causal_lm_templates/text_causal_lm.hpp" +#include "../qwen3_5/qwen3_5_for_causal_lm.hpp" +#include "qwen3_5_moe_decoder_layer.hpp" + +#include + +namespace infinilm::models::qwen3_5_moe { + +using Qwen35MoeLanguageModel = infinilm::layers::causal_lm_templates::TextModel; +using Qwen35MoeModel = qwen3_5::Qwen35ModelTemplate; + +class Qwen35MoeForConditionalGeneration + : public infinilm::layers::causal_lm_templates::TextCausalLM { +public: + using Base = infinilm::layers::causal_lm_templates::TextCausalLM; + using Base::Base; + + void reset_cache(const cache::CacheConfig *cache_config) override { + if (cache_config == nullptr) { + this->cache_config_.reset(); + } else { + this->cache_config_ = cache_config->unique_copy(); + } + this->model().reset_cache(cache_config); + } +}; + +std::shared_ptr create_qwen3_5_moe_model_config( + std::shared_ptr model_config); + +} // namespace infinilm::models::qwen3_5_moe diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py index 00dae220..11bcdaa0 100644 --- a/python/infinilm/infer_engine.py +++ b/python/infinilm/infer_engine.py @@ -124,6 +124,27 @@ class GenerationConfig: stop_on_eos: bool = True +def _infer_position_id_axes(hf_config: dict) -> int: + text_config = hf_config.get("text_config", hf_config) + if not isinstance(text_config, dict): + return 1 + + explicit_axes = text_config.get( + "position_id_axes", hf_config.get("position_id_axes") + ) + if explicit_axes is not None: + axes = int(explicit_axes) + if axes < 1: + raise ValueError("position_id_axes must be positive") + return axes + + rope_parameters = text_config.get("rope_parameters") or {} + mrope_section = rope_parameters.get("mrope_section") + if isinstance(mrope_section, (list, tuple)) and mrope_section: + return len(mrope_section) + return 1 + + class InferEngine(_infinilm.InferEngine): def __init__( self, @@ -144,6 +165,11 @@ def __init__( self.hf_config = read_hf_config(model_path) self.hf_generation_config = read_hf_generation_config(model_path) self.hf_config["use_legacy_moe"] = bool(use_legacy_moe) + self.position_id_axes = _infer_position_id_axes(self.hf_config) + self.hf_config["position_id_axes"] = self.position_id_axes + text_config = self.hf_config.get("text_config") + if isinstance(text_config, dict): + text_config.setdefault("position_id_axes", self.position_id_axes) if device is None: device = infinicore.device() @@ -531,9 +557,13 @@ def generate( if self.enable_paged_attn: input_ids = input_ids.view([1, batch_size * seq_len]) + position_ids_list = ( + list(range(past_seq_len, past_seq_len + seq_len)) * batch_size + ) + if self.position_id_axes > 1: + position_ids_list = [position_ids_list] * self.position_id_axes position_ids = infinicore.from_list( - list(range(past_seq_len, past_seq_len + seq_len)) * batch_size, - dtype=infinicore.int64, + position_ids_list, dtype=infinicore.int64 ) if iter == 0: diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index 9bcefa7f..486df03b 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -985,6 +985,60 @@ def _remap_qwen3_next(state_dict, config): return state_dict +def _remap_qwen3_5_moe(state_dict, config): + """Adapt packed Qwen3.5-MoE experts to InfiniLM expert parameter names.""" + state_dict = _remap_qwen3_5(state_dict, config) + text_config = config.get("text_config", config) + expected_num_experts = text_config["num_experts"] + expected_intermediate_size = text_config["moe_intermediate_size"] + + remapped = {} + for key, tensor in state_dict.items(): + if key.endswith(".mlp.experts.gate_up_proj"): + if tensor.ndim != 3: + raise ValueError( + f"Expected packed gate_up_proj to be 3D, got {tensor.shape} for {key}" + ) + if tensor.shape[0] != expected_num_experts: + raise ValueError( + f"Expected {expected_num_experts} experts, got {tensor.shape[0]} for {key}" + ) + if tensor.shape[1] != expected_intermediate_size * 2: + raise ValueError( + f"Expected packed gate/up size {expected_intermediate_size * 2}, " + f"got {tensor.shape[1]} for {key}" + ) + + prefix = key[: -len("gate_up_proj")] + for expert_idx, expert_gate_up in enumerate(tensor.unbind(0)): + gate, up = expert_gate_up.chunk(2, dim=0) + expert_prefix = f"{prefix}{expert_idx}." + remapped[f"{expert_prefix}gate_proj.weight"] = gate + remapped[f"{expert_prefix}up_proj.weight"] = up + elif key.endswith(".mlp.experts.down_proj"): + if tensor.ndim != 3: + raise ValueError( + f"Expected packed down_proj to be 3D, got {tensor.shape} for {key}" + ) + if tensor.shape[0] != expected_num_experts: + raise ValueError( + f"Expected {expected_num_experts} experts, got {tensor.shape[0]} for {key}" + ) + if tensor.shape[2] != expected_intermediate_size: + raise ValueError( + f"Expected down projection input size {expected_intermediate_size}, " + f"got {tensor.shape[2]} for {key}" + ) + + prefix = key[: -len("down_proj")] + for expert_idx, expert_down in enumerate(tensor.unbind(0)): + remapped[f"{prefix}{expert_idx}.down_proj.weight"] = expert_down + else: + remapped[key] = tensor + + return remapped + + _WEIGHT_REMAPPER = { "glm4": _remap_glm4, "chatglm": _remap_chatglm, @@ -994,5 +1048,6 @@ def _remap_qwen3_next(state_dict, config): "videonsa": _remap_videonsa, "qwen3_5": _remap_qwen3_5, "ernie4_5_moe_vl": _remap_ernie4_5_moe_vl, + "qwen3_5_moe": _remap_qwen3_5_moe, "qwen3_next": _remap_qwen3_next, } diff --git a/python/infinilm/processors/__init__.py b/python/infinilm/processors/__init__.py index 3d5cad4d..f1a543b6 100644 --- a/python/infinilm/processors/__init__.py +++ b/python/infinilm/processors/__init__.py @@ -33,13 +33,20 @@ def from_pretrained(cls, model_dir_path: str, **kwargs) -> InfinilmProcessor: registered Processor. Falls back to the registered default processor for unregistered or standard architectures. """ - config = AutoConfig.from_pretrained(model_dir_path, trust_remote_code=True) - model_type = config.model_type.lower() raw_config_path = Path(model_dir_path) / "config.json" - architectures = [] + raw_config = {} if raw_config_path.exists(): with raw_config_path.open("r") as f: - architectures = json.load(f).get("architectures", []) or [] + raw_config = json.load(f) + + raw_model_type = str(raw_config.get("model_type", "")).lower() + if raw_model_type in {"qwen3_5", "qwen3_5_moe"}: + model_type = raw_model_type + else: + config = AutoConfig.from_pretrained(model_dir_path, trust_remote_code=True) + model_type = config.model_type.lower() + + architectures = raw_config.get("architectures", []) or [] if ( model_type == "qwen2_5_vl" and "VideoNSAForConditionalGeneration" in architectures diff --git a/python/infinilm/processors/qwen3_5_processor.py b/python/infinilm/processors/qwen3_5_processor.py index e550de5f..6b8fee90 100644 --- a/python/infinilm/processors/qwen3_5_processor.py +++ b/python/infinilm/processors/qwen3_5_processor.py @@ -10,6 +10,7 @@ from .processor import register_processor +@register_processor("qwen3_5_moe") @register_processor("qwen3_5") class Qwen35Processor(BasicLLMProcessor): def __init__(self, model_dir_path: str): diff --git a/test/models/qwen3_5_moe/test_adaptation.py b/test/models/qwen3_5_moe/test_adaptation.py new file mode 100644 index 00000000..3a0a43ba --- /dev/null +++ b/test/models/qwen3_5_moe/test_adaptation.py @@ -0,0 +1,74 @@ +import unittest + +import torch +from infinilm.infer_engine import _infer_position_id_axes +from infinilm.modeling_utils import _remap_qwen3_5_moe + + +class PositionIdAxesTest(unittest.TestCase): + def test_defaults_to_one_axis(self): + self.assertEqual(_infer_position_id_axes({"text_config": {}}), 1) + + def test_infers_axes_from_mrope_section(self): + config = {"text_config": {"rope_parameters": {"mrope_section": [11, 11, 10]}}} + self.assertEqual(_infer_position_id_axes(config), 3) + + def test_explicit_axes_take_precedence(self): + config = { + "position_id_axes": 2, + "text_config": { + "position_id_axes": 4, + "rope_parameters": {"mrope_section": [11, 11, 10]}, + }, + } + self.assertEqual(_infer_position_id_axes(config), 4) + + def test_rejects_non_positive_axes(self): + with self.assertRaisesRegex(ValueError, "must be positive"): + _infer_position_id_axes({"text_config": {"position_id_axes": 0}}) + + +class Qwen35MoeWeightRemapTest(unittest.TestCase): + def setUp(self): + self.config = { + "text_config": { + "linear_key_head_dim": 2, + "linear_num_key_heads": 1, + "num_experts": 2, + "moe_intermediate_size": 3, + } + } + + def test_splits_packed_expert_weights(self): + gate_up = torch.arange(2 * 6 * 4).reshape(2, 6, 4) + down = torch.arange(2 * 4 * 3).reshape(2, 4, 3) + state_dict = { + "model.language_model.layers.0.mlp.experts.gate_up_proj": gate_up, + "model.language_model.layers.0.mlp.experts.down_proj": down, + } + + remapped = _remap_qwen3_5_moe(state_dict, self.config) + + prefix = "model.language_model.layers.0.mlp.experts." + self.assertTrue( + torch.equal(remapped[f"{prefix}0.gate_proj.weight"], gate_up[0, :3]) + ) + self.assertTrue( + torch.equal(remapped[f"{prefix}0.up_proj.weight"], gate_up[0, 3:]) + ) + self.assertTrue(torch.equal(remapped[f"{prefix}1.down_proj.weight"], down[1])) + self.assertNotIn(f"{prefix}gate_up_proj", remapped) + self.assertNotIn(f"{prefix}down_proj", remapped) + + def test_rejects_wrong_expert_count(self): + state_dict = { + "model.language_model.layers.0.mlp.experts.gate_up_proj": torch.zeros( + 1, 6, 4 + ) + } + with self.assertRaisesRegex(ValueError, "Expected 2 experts"): + _remap_qwen3_5_moe(state_dict, self.config) + + +if __name__ == "__main__": + unittest.main()