diff --git a/csrc/models/qwen3_moe/qwen3_moe_experts.cpp b/csrc/models/qwen3_moe/qwen3_moe_experts.cpp index d08b3bf7..2692db2c 100644 --- a/csrc/models/qwen3_moe/qwen3_moe_experts.cpp +++ b/csrc/models/qwen3_moe/qwen3_moe_experts.cpp @@ -1,21 +1,54 @@ #include "qwen3_moe_experts.hpp" +#include "../../config/model_config.hpp" +#include "../../global_state/global_state.hpp" #include "infinicore/ops.hpp" +#include #include namespace infinilm::models::qwen3_moe { Qwen3MoeExperts::Qwen3MoeExperts(std::shared_ptr model_config, const infinicore::Device &device) { + const auto &dtype = model_config->get_dtype(); + const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); + const int tp_rank = rank_info.tp_rank; + const int tp_size = rank_info.tp_size; num_experts_ = model_config->get("num_experts"); num_experts_per_tok_ = model_config->get("num_experts_per_tok"); + hidden_size_ = model_config->get("hidden_size"); + const size_t intermediate_size = model_config->get("moe_intermediate_size"); ASSERT((num_experts_ > 0) && (num_experts_per_tok_ > 0) && (num_experts_per_tok_ <= num_experts_)); - - for (size_t i = 0; i < num_experts_; ++i) { - experts_.push_back(this->register_module(std::to_string(i), model_config, device)); + ASSERT(intermediate_size % static_cast(tp_size) == 0); + intermediate_size_per_partition_ = intermediate_size / static_cast(tp_size); + + INFINICORE_NN_PARAMETER_INIT(w1, ({num_experts_, 2 * intermediate_size_per_partition_, hidden_size_}, dtype, device)); + INFINICORE_NN_PARAMETER_INIT(w2, ({num_experts_, hidden_size_, intermediate_size_per_partition_}, dtype, device)); + + for (size_t expert = 0; expert < num_experts_; ++expert) { + auto gate_weight = w1_ + ->narrow({{0, expert, 1}, {1, 0, intermediate_size_per_partition_}}) + ->squeeze(0); + auto up_weight = w1_ + ->narrow({{0, expert, 1}, {1, intermediate_size_per_partition_, intermediate_size_per_partition_}}) + ->squeeze(0); + auto down_weight = w2_ + ->narrow({{0, expert, 1}}) + ->squeeze(0); + + const std::string prefix = std::to_string(expert) + "."; + this->register_parameter( + prefix + "gate_proj.weight", + infinicore::nn::Parameter(gate_weight, 0, tp_rank, tp_size)); + this->register_parameter( + prefix + "up_proj.weight", + infinicore::nn::Parameter(up_weight, 0, tp_rank, tp_size)); + this->register_parameter( + prefix + "down_proj.weight", + infinicore::nn::Parameter(down_weight, 1, tp_rank, tp_size)); } } @@ -23,42 +56,10 @@ infinicore::Tensor Qwen3MoeExperts::forward(const infinicore::Tensor &hidden_sta const infinicore::Tensor &top_k_index, const infinicore::Tensor &top_k_weights) const { ASSERT(hidden_states->ndim() == 2); + ASSERT(top_k_index->ndim() == 2 && top_k_weights->ndim() == 2); - auto top_k_weights_cpu = top_k_weights->to(infinicore::Device::Type::CPU); - auto top_k_index_cpu = top_k_index->to(infinicore::Device::Type::CPU); - - int *top_k_index_ptr = reinterpret_cast(top_k_index_cpu->data()); - float *top_k_weights_ptr = reinterpret_cast(top_k_weights_cpu->data()); - - size_t ntoken = hidden_states->shape()[0]; - int index; - float score; - - auto final_hidden_states = infinicore::Tensor::empty(hidden_states->shape(), hidden_states->dtype(), hidden_states->device()); - for (size_t itok = 0; itok < ntoken; ++itok) { - auto hidden_states_i = hidden_states->narrow({{0, itok, 1}}); - const size_t route_row = itok * num_experts_per_tok_; - - infinicore::Tensor final_hidden_states_i; - for (size_t k = 0; k < num_experts_per_tok_; ++k) { - index = top_k_index_ptr[route_row + k]; - score = top_k_weights_ptr[route_row + k]; - - ASSERT(index >= 0 && static_cast(index) < num_experts_); - - experts_[index]->set_alpha(score); - auto expert_out = experts_[index]->forward(hidden_states_i); - - if (k == 0) { - final_hidden_states_i = expert_out; - } else { - infinicore::op::add_(final_hidden_states_i, final_hidden_states_i, expert_out); - } - } - - final_hidden_states->narrow({{0, itok, 1}})->copy_from(final_hidden_states_i); - } - return final_hidden_states; + return infinicore::op::fused_moe(hidden_states, top_k_index, top_k_weights, w1_, w2_, std::nullopt, std::nullopt, + infinicore::op::FusedMoeActivation::Swiglu); } } // namespace infinilm::models::qwen3_moe diff --git a/csrc/models/qwen3_moe/qwen3_moe_experts.hpp b/csrc/models/qwen3_moe/qwen3_moe_experts.hpp index 2f470ca8..0d3c00d3 100644 --- a/csrc/models/qwen3_moe/qwen3_moe_experts.hpp +++ b/csrc/models/qwen3_moe/qwen3_moe_experts.hpp @@ -1,15 +1,17 @@ #pragma once -#include "../../layers/moe/legacy/moe_mlp.hpp" #include "infinicore/nn/module.hpp" +#include "infinicore/nn/parameter.hpp" #include "infinicore/tensor.hpp" #include #include -namespace infinilm::models::qwen3_moe { +namespace infinilm::config { +class ModelConfig; +} -using Qwen3MoeMLP = infinilm::layers::moe::legacy::MoeMLP; +namespace infinilm::models::qwen3_moe { class Qwen3MoeExperts : public infinicore::nn::Module { public: @@ -21,9 +23,13 @@ class Qwen3MoeExperts : public infinicore::nn::Module { const infinicore::Tensor &top_k_weights) const; protected: - INFINICORE_NN_MODULE_VEC(Qwen3MoeMLP, experts); + INFINICORE_NN_PARAMETER(w1); + INFINICORE_NN_PARAMETER(w2); + size_t num_experts_per_tok_{0}; size_t num_experts_{0}; + size_t hidden_size_{0}; + size_t intermediate_size_per_partition_{0}; }; } // namespace infinilm::models::qwen3_moe diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index 8f84c9b9..ce84141d 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -59,8 +59,11 @@ def _is_internal_moe_packed_weight(key: str) -> bool: # InfiniLM registers packed MoE parameters internally. HF checkpoints # provide per-expert gate/up/down weights instead, so these packed tensors # are expected missing keys during non-strict checkpoint loading. - return key.endswith(".mlp.experts.w13_weight") or key.endswith( - ".mlp.experts.w2_weight" + return ( + key.endswith(".mlp.experts.w13_weight") + or key.endswith(".mlp.experts.w2_weight") + or key.endswith(".mlp.experts.w1") + or key.endswith(".mlp.experts.w2") )