diff --git a/csrc/models/granitemoehybrid/granitemoehybrid_allocate_kv_cache_tensors.cpp b/csrc/models/granitemoehybrid/granitemoehybrid_allocate_kv_cache_tensors.cpp new file mode 100644 index 000000000..d216599ae --- /dev/null +++ b/csrc/models/granitemoehybrid/granitemoehybrid_allocate_kv_cache_tensors.cpp @@ -0,0 +1,138 @@ +#include "granitemoehybrid_allocate_kv_cache_tensors.hpp" + +#include "../../global_state/global_state.hpp" + +#include "infinicore/context/context.hpp" + +#include +#include +#include +#include +#include + +namespace infinilm::models::granitemoehybrid { + +GraniteMoeHybridAllocatedCache granitemoehybrid_allocate_cache_tensors( + const cache::CacheConfig *cache_config, + const std::shared_ptr &text_config, + const backends::AttentionBackend &attention_backend) { + if (nullptr == cache_config) { + return {}; + } + if (nullptr == text_config) { + throw std::runtime_error("infinilm::models::granitemoehybrid::granitemoehybrid_allocate_cache_tensors: text_config is null"); + } + + const size_t num_hidden_layers = text_config->get("num_hidden_layers"); + const size_t head_dim = text_config->get_head_dim(); + const size_t num_key_value_heads = text_config->get("num_key_value_heads"); + + const size_t hidden_size = text_config->get("hidden_size"); + const size_t mamba_expand = text_config->get_or("mamba_expand", 2); + const size_t mamba_n_groups = text_config->get_or("mamba_n_groups", 1); + const size_t mamba_d_state = text_config->get("mamba_d_state"); + const size_t mamba_d_conv = text_config->get("mamba_d_conv"); + + const auto &dtype = text_config->get_dtype(); + const auto &kv_cache_dtype = text_config->get_kv_cache_dtype(); + const std::vector layer_types = text_config->get>("layer_types"); + + std::vector kv_cache_vec; + std::vector conv_state_vec; + kv_cache_vec.reserve(num_hidden_layers); + conv_state_vec.reserve(num_hidden_layers); + + const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); + const size_t local_mamba_groups = mamba_n_groups >= static_cast(rank_info.tp_size) + ? mamba_n_groups / rank_info.tp_size + : 1; + const size_t mamba_conv_dim = mamba_expand * hidden_size / rank_info.tp_size + 2 * local_mamba_groups * mamba_d_state; + + auto allocate_mamba_cache = [&](size_t pool_size) { + auto conv_state = infinicore::Tensor::zeros( + {pool_size, mamba_conv_dim, mamba_d_conv - 1}, + dtype, + rank_info.device); + kv_cache_vec.emplace_back(); + conv_state_vec.push_back(std::move(conv_state)); + }; + + auto allocate_static_attention_cache = [&](const cache::StaticKVCacheConfig &config) { + auto kv_cache = cache::StaticKVCache::create_layer_kv_cache( + head_dim, + head_dim, + num_key_value_heads, + num_key_value_heads, + text_config->get("max_position_embeddings"), + kv_cache_dtype, + config); + + kv_cache_vec.push_back(std::move(kv_cache)); + conv_state_vec.emplace_back(); + }; + + auto allocate_paged_attention_cache = [&](const cache::PagedKVCacheConfig &config) { + auto kv_cache = cache::PagedKVCache::create_layer_kv_cache( + head_dim, + head_dim, + num_key_value_heads, + num_key_value_heads, + kv_cache_dtype, + config); + + kv_cache_vec.push_back(std::move(kv_cache)); + conv_state_vec.emplace_back(); + }; + + switch (attention_backend) { + case backends::AttentionBackend::STATIC_ATTN: { + const auto *static_kv_cache_config = dynamic_cast(cache_config); + if (nullptr == static_kv_cache_config) { + throw std::runtime_error("infinilm::models::granitemoehybrid::granitemoehybrid_allocate_cache_tensors: invalid static kv cache config type"); + } + + const size_t mamba_pool_size = std::max(2, static_kv_cache_config->max_batch_size() + 1); + for (size_t layer_idx = 0; layer_idx < num_hidden_layers; ++layer_idx) { + const std::string &layer_type = layer_types[layer_idx]; + if ("mamba" == layer_type) { + allocate_mamba_cache(mamba_pool_size); + } else if ("attention" == layer_type) { + allocate_static_attention_cache(*static_kv_cache_config); + } else { + throw std::runtime_error("infinilm::models::granitemoehybrid::granitemoehybrid_allocate_cache_tensors: unsupported layer_type '" + layer_type + "' for layer " + std::to_string(layer_idx)); + } + } + break; + } + case backends::AttentionBackend::FLASH_ATTN: { + ; + } + case backends::AttentionBackend::PAGED_ATTN: { + const auto *paged_kv_cache_config = dynamic_cast(cache_config); + if (nullptr == paged_kv_cache_config) { + throw std::runtime_error("infinilm::models::granitemoehybrid::granitemoehybrid_allocate_cache_tensors: invalid paged kv cache config type"); + } + const size_t mamba_pool_size = std::max(2, paged_kv_cache_config->num_blocks() / 4); + + for (size_t layer_idx = 0; layer_idx < num_hidden_layers; ++layer_idx) { + const std::string &layer_type = layer_types[layer_idx]; + if ("mamba" == layer_type) { + allocate_mamba_cache(mamba_pool_size); + } else if ("attention" == layer_type) { + allocate_paged_attention_cache(*paged_kv_cache_config); + } else { + throw std::runtime_error("infinilm::models::granitemoehybrid::granitemoehybrid_allocate_cache_tensors: unsupported layer_type '" + layer_type + "' for layer " + std::to_string(layer_idx)); + } + } + break; + } + default: + throw std::runtime_error("infinilm::models::granitemoehybrid::granitemoehybrid_allocate_cache_tensors: unsupported attention backend " + std::to_string(static_cast(attention_backend))); + } + infinicore::context::syncStream(); + return GraniteMoeHybridAllocatedCache{ + std::move(kv_cache_vec), + std::move(conv_state_vec)}; +} + +} // namespace infinilm::models::granitemoehybrid diff --git a/csrc/models/granitemoehybrid/granitemoehybrid_allocate_kv_cache_tensors.hpp b/csrc/models/granitemoehybrid/granitemoehybrid_allocate_kv_cache_tensors.hpp new file mode 100644 index 000000000..816248ac4 --- /dev/null +++ b/csrc/models/granitemoehybrid/granitemoehybrid_allocate_kv_cache_tensors.hpp @@ -0,0 +1,22 @@ +#pragma once + +#include "../../backends/attention_backends.hpp" +#include "../../cache/kv_cache.hpp" +#include "../../config/model_config.hpp" + +#include +#include + +namespace infinilm::models::granitemoehybrid { + +struct GraniteMoeHybridAllocatedCache { + std::vector kv_cache_tensors; + std::vector conv_state_tensors; +}; + +GraniteMoeHybridAllocatedCache granitemoehybrid_allocate_cache_tensors( + const cache::CacheConfig *cache_config, + const std::shared_ptr &text_config, + const backends::AttentionBackend &attention_backend); + +} // namespace infinilm::models::granitemoehybrid diff --git a/csrc/models/granitemoehybrid/granitemoehybrid_attention.cpp b/csrc/models/granitemoehybrid/granitemoehybrid_attention.cpp new file mode 100644 index 000000000..24f7cd508 --- /dev/null +++ b/csrc/models/granitemoehybrid/granitemoehybrid_attention.cpp @@ -0,0 +1,177 @@ +#include "granitemoehybrid_attention.hpp" + +#include "../../global_state/global_state.hpp" +#include "../../layers/attention/attention.hpp" +#include "../../layers/rotary_embedding/rotary_embedding.hpp" +#include "../../utils.hpp" + +#include +#include +#include + +namespace infinilm::models::granitemoehybrid { + +GraniteMoeHybridAttention::GraniteMoeHybridAttention( + std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device) { + layer_idx_ = layer_idx; + hidden_size_ = model_config->get("hidden_size"); + head_dim_ = model_config->get_head_dim(); + + const auto &dtype{model_config->get_dtype()}; + const size_t total_num_heads = model_config->get("num_attention_heads"); + const size_t total_num_kv_heads = model_config->get("num_key_value_heads"); + const bool use_bias = model_config->get_or("attention_bias", false); + const bool use_output_bias = model_config->get_or("attention_output_bias", use_bias); + + attention_backend_ = infinilm::global_state::get_infinilm_config().attention_backend; + const engine::distributed::RankInfo &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; + if (tp_size <= 0 || total_num_heads % tp_size != 0) { + throw std::runtime_error( + "infinilm::models::granitemoehybrid::GraniteMoeHybridAttention: " + "num_attention_heads must be divisible by tp_size"); + } + if (total_num_kv_heads < static_cast(tp_size) || total_num_kv_heads % tp_size != 0) { + throw std::runtime_error( + "infinilm::models::granitemoehybrid::GraniteMoeHybridAttention: " + "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; + + auto quantization_method = model_config->get_quantization_method(); + auto register_fn = [this](const std::string &name, infinicore::nn::Parameter parameter) { + this->register_parameter(name, std::move(parameter)); + }; + + qkv_proj_ = std::make_shared( + hidden_size_, + head_dim_, + total_num_heads, + total_num_kv_heads, + "q_proj", + "k_proj", + "v_proj", + register_fn, + quantization_method, + use_bias, + dtype, + device, + rank_info); + o_proj_ = this->register_module( + "o_proj", + total_num_heads * head_dim_, + hidden_size_, + quantization_method, + use_output_bias, + dtype, + device, + tp_rank, + tp_size, + rank_info.comm); + o_proj_->set_alpha(model_config->get_or("residual_multiplier", 1.0f)); + + const std::string position_embedding_type = model_config->get_or("position_embedding_type", "nope"); + if ("rope" == position_embedding_type) { + rotary_emb_ = infinilm::layers::rotary_embedding::get_rope(model_config, device); + } + + const float attention_multiplier = model_config->get_or("attention_multiplier", 1.0f); + infinilm::layers::attention::init_kv_cache_quant_params( + register_fn, + device, + kv_cache_k_scale_, + kv_cache_v_scale_); + attn_ = std::make_shared( + num_attention_heads_, + head_dim_, + attention_multiplier, + num_key_value_heads_, + layer_idx_, + kv_cache_k_scale_, + kv_cache_v_scale_, + attention_backend_); +} + +infinicore::Tensor GraniteMoeHybridAttention::forward( + const infinicore::Tensor &positions, + const infinicore::Tensor &hidden_states) const { + if (::infinilm::backends::AttentionBackend::STATIC_ATTN == attention_backend_) { + return forward_static_(positions, hidden_states); + } + return forward_paged_(positions, hidden_states); +} + +infinicore::Tensor GraniteMoeHybridAttention::forward_static_( + const infinicore::Tensor &position_ids, + const infinicore::Tensor &hidden_states) const { + auto hidden_states_mutable = hidden_states; + const auto &shape = hidden_states->shape(); + const size_t batch_size = shape[0]; + const size_t seq_len = shape[1]; + + auto [query, key, value] = qkv_proj_->forward_split(hidden_states_mutable); + query = query->view({batch_size, seq_len, num_attention_heads_, head_dim_}); + key = key->view({batch_size, seq_len, num_key_value_heads_, head_dim_}); + value = value->view({batch_size, seq_len, num_key_value_heads_, head_dim_}); + + if (rotary_emb_) { + const auto &position_shape = position_ids->shape(); + infinicore::Tensor rope_positions; + if (position_shape.size() == 2) { + rope_positions = position_ids->narrow({{0, 0, 1}})->contiguous()->view({position_shape[1]}); + } else if (position_shape.size() == 1) { + rope_positions = position_ids->contiguous(); + } else { + throw std::runtime_error( + "infinilm::models::granitemoehybrid::GraniteMoeHybridAttention: " + "unexpected position_ids shape"); + } + + rotary_emb_->forward(query, rope_positions, true); + rotary_emb_->forward(key, rope_positions, true); + } + + auto attention_output = attn_->forward(query, key, value); + return o_proj_->forward(attention_output); +} + +infinicore::Tensor GraniteMoeHybridAttention::forward_paged_( + const infinicore::Tensor &position_ids, + const infinicore::Tensor &hidden_states) const { + auto hidden_states_mutable = hidden_states; + const auto &shape = hidden_states->shape(); + const size_t batch_size = shape[0]; + const size_t seq_len = shape[1]; + ASSERT_EQ(batch_size, 1); + + auto [query, key, value] = qkv_proj_->forward_split(hidden_states_mutable); + query = query->view({seq_len, num_attention_heads_, head_dim_}); + key = key->view({seq_len, num_key_value_heads_, head_dim_}); + value = value->view({seq_len, num_key_value_heads_, head_dim_}); + + if (rotary_emb_) { + const auto &position_shape = position_ids->shape(); + infinicore::Tensor rope_positions; + if (position_shape.size() == 2) { + rope_positions = position_ids->narrow({{0, 0, 1}})->view({position_shape[1]}); + } else if (position_shape.size() == 1) { + rope_positions = position_ids; + } else { + throw std::runtime_error( + "infinilm::models::granitemoehybrid::GraniteMoeHybridAttention: " + "unexpected position_ids shape"); + } + rotary_emb_->forward(query, rope_positions, true); + rotary_emb_->forward(key, rope_positions, true); + } + + auto attention_output = attn_->forward(query, key, value); + return o_proj_->forward(attention_output); +} + +} // namespace infinilm::models::granitemoehybrid diff --git a/csrc/models/granitemoehybrid/granitemoehybrid_attention.hpp b/csrc/models/granitemoehybrid/granitemoehybrid_attention.hpp new file mode 100644 index 000000000..42327b7fa --- /dev/null +++ b/csrc/models/granitemoehybrid/granitemoehybrid_attention.hpp @@ -0,0 +1,56 @@ +#pragma once + +#include "../../layers/common_modules.hpp" + +#include + +namespace infinilm::models::granitemoehybrid { + +class GraniteMoeHybridAttention : public infinicore::nn::Module { +public: + GraniteMoeHybridAttention(std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &positions, + const infinicore::Tensor &hidden_states) const; + + void process_weights_after_loading() override { + qkv_proj_->process_weights_after_loading(); + } + + void reset_runtime_state() const override { + qkv_proj_->reset_runtime_state(); + } + + size_t layer_idx() const { return layer_idx_; } + size_t num_heads() const { return num_attention_heads_; } + size_t num_kv_heads() const { return num_key_value_heads_; } + size_t head_dim() const { return head_dim_; } + size_t hidden_size() const { return hidden_size_; } + +private: + infinicore::Tensor forward_static_(const infinicore::Tensor &positions, + const infinicore::Tensor &hidden_states) const; + + infinicore::Tensor forward_paged_(const infinicore::Tensor &positions, + const infinicore::Tensor &hidden_states) const; + +protected: + std::shared_ptr qkv_proj_; + std::shared_ptr o_proj_; + std::shared_ptr rotary_emb_; + std::shared_ptr attn_; + ::infinilm::backends::AttentionBackend attention_backend_; + + size_t layer_idx_; + size_t num_attention_heads_; + size_t num_key_value_heads_; + size_t hidden_size_; + size_t head_dim_; + + INFINICORE_NN_PARAMETER(kv_cache_k_scale); + INFINICORE_NN_PARAMETER(kv_cache_v_scale); +}; + +} // namespace infinilm::models::granitemoehybrid diff --git a/csrc/models/granitemoehybrid/granitemoehybrid_decoderLayer.cpp b/csrc/models/granitemoehybrid/granitemoehybrid_decoderLayer.cpp new file mode 100644 index 000000000..eb112d2a2 --- /dev/null +++ b/csrc/models/granitemoehybrid/granitemoehybrid_decoderLayer.cpp @@ -0,0 +1,89 @@ +#include "granitemoehybrid_decoderLayer.hpp" + +#include "infinicore/ops.hpp" + +#include +#include +#include + +namespace infinilm::models::granitemoehybrid { + +GraniteMoeHybridDecoderLayer::GraniteMoeHybridDecoderLayer( + std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device) + : layer_idx_(layer_idx), + has_experts_(model_config->get_or("num_local_experts", 0) > 0) { + 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(shared_mlp, model_config, device); + if (has_experts_) { + INFINICORE_NN_MODULE_INIT(block_sparse_moe, model_config, device); + } + + const auto layer_types = model_config->get>("layer_types"); + layer_type_ = layer_types.at(layer_idx); + if ("mamba" == layer_type_) { + INFINICORE_NN_MODULE_INIT(mamba, model_config, layer_idx, device); + } else if ("attention" == layer_type_) { + INFINICORE_NN_MODULE_INIT(self_attn, model_config, layer_idx, device); + } else { + throw std::runtime_error( + "infinilm::models::granitemoehybrid::GraniteMoeHybridDecoderLayer: " + "unsupported layer_type '" + + layer_type_ + "' for layer " + std::to_string(layer_idx)); + } +} + +std::tuple GraniteMoeHybridDecoderLayer::forward( + const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const { + input_layernorm_->forward_inplace(hidden_states, residual); + if ("mamba" == layer_type_) { + hidden_states = mamba_->forward(hidden_states); + } else { + hidden_states = self_attn_->forward(positions, hidden_states); + } + + post_attention_layernorm_->forward_inplace(hidden_states, residual); + if (has_experts_) { + auto moe_hidden_states = block_sparse_moe_->forward(hidden_states); + auto shared_hidden_states = shared_mlp_->forward(hidden_states); + hidden_states = infinicore::op::add(moe_hidden_states, shared_hidden_states); + } else { + hidden_states = shared_mlp_->forward(hidden_states); + } + return {hidden_states, residual}; +} + +infinicore::Tensor GraniteMoeHybridDecoderLayer::forward( + const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states) const { + auto residual = hidden_states; + hidden_states = input_layernorm_->forward(hidden_states); + if ("mamba" == layer_type_) { + hidden_states = mamba_->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); + if (has_experts_) { + auto moe_hidden_states = block_sparse_moe_->forward(hidden_states); + auto shared_hidden_states = shared_mlp_->forward(hidden_states); + hidden_states = infinicore::op::add(moe_hidden_states, shared_hidden_states); + } else { + hidden_states = shared_mlp_->forward(hidden_states); + } + hidden_states = infinicore::op::add(residual, hidden_states); + return hidden_states; +} + +} // namespace infinilm::models::granitemoehybrid diff --git a/csrc/models/granitemoehybrid/granitemoehybrid_decoderLayer.hpp b/csrc/models/granitemoehybrid/granitemoehybrid_decoderLayer.hpp new file mode 100644 index 000000000..691aed417 --- /dev/null +++ b/csrc/models/granitemoehybrid/granitemoehybrid_decoderLayer.hpp @@ -0,0 +1,44 @@ +#pragma once + +#include "granitemoehybrid_attention.hpp" +#include "granitemoehybrid_mamba.hpp" +#include "granitemoehybrid_shared_mlp.hpp" +#include "granitemoehybrid_sparse_moe_block.hpp" + +#include +#include +#include + +namespace infinilm::models::granitemoehybrid { + +class GraniteMoeHybridDecoderLayer : public infinicore::nn::Module { +public: + GraniteMoeHybridDecoderLayer(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) const; + + infinicore::Tensor forward(const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states) const; + + 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(GraniteMoeHybridAttention, self_attn); + INFINICORE_NN_MODULE(GraniteMoeHybridMamba, mamba); + INFINICORE_NN_MODULE(GraniteMoeHybridSparseMoeBlock, block_sparse_moe); + INFINICORE_NN_MODULE(GraniteMoeHybridSharedMLP, shared_mlp); + +private: + size_t layer_idx_; + std::string layer_type_; + bool has_experts_; +}; + +} // namespace infinilm::models::granitemoehybrid diff --git a/csrc/models/granitemoehybrid/granitemoehybrid_for_causal_lm.cpp b/csrc/models/granitemoehybrid/granitemoehybrid_for_causal_lm.cpp new file mode 100644 index 000000000..079447b4a --- /dev/null +++ b/csrc/models/granitemoehybrid/granitemoehybrid_for_causal_lm.cpp @@ -0,0 +1,100 @@ +#include "granitemoehybrid_for_causal_lm.hpp" +#include "granitemoehybrid_allocate_kv_cache_tensors.hpp" + +#include "../../global_state/global_state.hpp" +#include "../models_registry.hpp" + +#include +#include +#include + +namespace infinilm::models::granitemoehybrid { + +GraniteMoeHybridForCausalLM::GraniteMoeHybridForCausalLM( + std::shared_ptr model_config, + const infinicore::Device &device) { + model_config_ = model_config; + 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); + + const float logits_scaling = model_config->get_or("logits_scaling", 1.0f); + if (logits_scaling <= 0.0f) { + throw std::runtime_error( + "infinilm::models::granitemoehybrid::GraniteMoeHybridForCausalLM: " + "logits_scaling must be greater than zero"); + } + lm_head_->set_alpha(1.0f / logits_scaling); +} + +infinilm::InfinilmModel::Output GraniteMoeHybridForCausalLM::forward( + const infinilm::InfinilmModel::Input &input) const { + auto hidden_states = model_->forward(input); + auto logits = lm_head_->forward(hidden_states); + return {logits}; +} + +void GraniteMoeHybridForCausalLM::reset_cache(const cache::CacheConfig *cache_config) { + reset_runtime_state(); + auto &forward_context = infinilm::global_state::get_forward_context(); + forward_context.conv_state_vec.clear(); + forward_context.ssm_state_vec.clear(); + if (nullptr == cache_config) { + InfinilmModel::reset_cache(nullptr); + return; + } + cache_config_ = cache_config->unique_copy(); + + forward_context.kv_cache_vec.clear(); + const auto attention_backend = infinilm::global_state::get_infinilm_config().attention_backend; + auto allocated_cache = granitemoehybrid_allocate_cache_tensors( + cache_config, + model_config_, + attention_backend); + forward_context.kv_cache_vec = std::move(allocated_cache.kv_cache_tensors); + forward_context.conv_state_vec = std::move(allocated_cache.conv_state_tensors); +} + +std::shared_ptr create_granitemoehybrid_model_config( + std::shared_ptr model_config) { + const std::string model_type = model_config->get("model_type"); + if ("granitemoehybrid" != model_type) { + throw std::runtime_error( + "infinilm::models::granitemoehybrid::create_granitemoehybrid_model_config: " + "model_type is not granitemoehybrid"); + } + + nlohmann::json &config_json = model_config->get_config_json(); + if (!config_json.contains("layer_types")) { + const size_t num_hidden_layers = model_config->get("num_hidden_layers"); + config_json["layer_types"] = std::vector(num_hidden_layers, "mamba"); + } + + if (!config_json.contains("attention_bias")) { + config_json["attention_bias"] = false; + } + + if (!config_json.contains("position_embedding_type") || config_json.at("position_embedding_type").is_null()) { + config_json["position_embedding_type"] = "nope"; + } + const std::string position_embedding_type = model_config->get("position_embedding_type"); + if ("rope" != position_embedding_type && "nope" != position_embedding_type) { + throw std::runtime_error( + "infinilm::models::granitemoehybrid::create_granitemoehybrid_model_config: " + "position_embedding_type must be either 'rope' or 'nope'"); + } + + return model_config; +} + +} // namespace infinilm::models::granitemoehybrid + +namespace { +INFINILM_REGISTER_CAUSAL_LM_MODEL( + granitemoehybrid, + infinilm::models::granitemoehybrid::GraniteMoeHybridForCausalLM, + infinilm::models::granitemoehybrid::create_granitemoehybrid_model_config); +} // namespace diff --git a/csrc/models/granitemoehybrid/granitemoehybrid_for_causal_lm.hpp b/csrc/models/granitemoehybrid/granitemoehybrid_for_causal_lm.hpp new file mode 100644 index 000000000..3f98cd018 --- /dev/null +++ b/csrc/models/granitemoehybrid/granitemoehybrid_for_causal_lm.hpp @@ -0,0 +1,27 @@ +#pragma once + +#include "granitemoehybrid_decoderLayer.hpp" +#include + +namespace infinilm::models::granitemoehybrid { + +using GraniteMoeHybridModel = infinilm::layers::causal_lm_templates::TextModel; + +class GraniteMoeHybridForCausalLM : public InfinilmModel { +public: + GraniteMoeHybridForCausalLM(std::shared_ptr model_config, + const infinicore::Device &device); + + Output forward(const Input &input) const override; + + void reset_cache(const cache::CacheConfig *cache_config) override; + +protected: + INFINICORE_NN_MODULE(GraniteMoeHybridModel, model); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, lm_head); +}; + +std::shared_ptr create_granitemoehybrid_model_config( + std::shared_ptr model_config); + +} // namespace infinilm::models::granitemoehybrid diff --git a/csrc/models/granitemoehybrid/granitemoehybrid_mamba.cpp b/csrc/models/granitemoehybrid/granitemoehybrid_mamba.cpp new file mode 100644 index 000000000..debd8d094 --- /dev/null +++ b/csrc/models/granitemoehybrid/granitemoehybrid_mamba.cpp @@ -0,0 +1,305 @@ +#include "granitemoehybrid_mamba.hpp" + +#include "../../global_state/global_state.hpp" + +#include "infinicore/ops/add.hpp" +#include "infinicore/ops/broadcast_to.hpp" +#include "infinicore/ops/causal_conv1d.hpp" +#include "infinicore/ops/distributed/allgather.hpp" +#include "infinicore/ops/mamba_selective_scan.hpp" +#include "infinicore/ops/mul.hpp" +#include "infinicore/ops/rms_norm.hpp" +#include "infinicore/ops/silu.hpp" + +#include +#include +#include + +namespace infinilm::models::granitemoehybrid { + +GraniteMoeHybridCausalConv1d::GraniteMoeHybridCausalConv1d( + 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 size_t mamba_expand = model_config->get_or("mamba_expand", 2); + const size_t mamba_n_groups = model_config->get_or("mamba_n_groups", 1); + const size_t mamba_n_heads = model_config->get("mamba_n_heads"); + const size_t mamba_d_state = model_config->get("mamba_d_state"); + const size_t mamba_d_conv = model_config->get("mamba_d_conv"); + const size_t intermediate_size = mamba_expand * hidden_size; + const size_t conv_dim = intermediate_size + 2 * mamba_n_groups * mamba_d_state; + use_bias_ = model_config->get_or("mamba_conv_bias", true); + + const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); + tp_size_ = rank_info.tp_size; + tp_rank_ = rank_info.tp_rank; + conv_kernel_dim_ = mamba_d_conv; + full_x_dim_ = intermediate_size; + full_bc_dim_ = mamba_n_groups * mamba_d_state; + local_x_dim_ = full_x_dim_ / tp_size_; + local_bc_dim_ = (mamba_n_groups >= tp_size_ ? mamba_n_groups / tp_size_ : 1) * mamba_d_state; + local_conv_dim_ = local_x_dim_ + 2 * local_bc_dim_; + bc_replicas_ = mamba_n_groups < tp_size_ ? tp_size_ / mamba_n_groups : 1; + + INFINICORE_NN_PARAMETER_INIT(weight, ({conv_dim, 1, mamba_d_conv}, dtype, device)); + if (use_bias_) { + INFINICORE_NN_PARAMETER_INIT(bias, ({conv_dim}, dtype, device)); + } +} + +void GraniteMoeHybridCausalConv1d::process_weights_after_loading() { + if (tp_size_ <= 1) { + return; + } + + const size_t expected_full_conv_dim = full_x_dim_ + 2 * full_bc_dim_; + const size_t bc_offset = (tp_rank_ / bc_replicas_) * local_bc_dim_; + const size_t src_x_offset = tp_rank_ * local_x_dim_; + const size_t src_b_offset = full_x_dim_ + bc_offset; + const size_t src_c_offset = full_x_dim_ + full_bc_dim_ + bc_offset; + + const size_t dst_x_offset = 0; + const size_t dst_b_offset = local_x_dim_; + const size_t dst_c_offset = local_x_dim_ + local_bc_dim_; + + if (weight_->size(0) != local_conv_dim_) { + if (weight_->shape() != infinicore::Shape{expected_full_conv_dim, 1, conv_kernel_dim_}) { + throw std::runtime_error("GraniteMoeHybridCausalConv1d: unexpected conv1d weight shape for TP slicing"); + } + + auto local_weight = infinicore::Tensor::empty( + {local_conv_dim_, 1, conv_kernel_dim_}, + weight_->dtype(), + weight_->device()); + + local_weight->narrow({{0, dst_x_offset, local_x_dim_}}) + ->copy_from(weight_->narrow({{0, src_x_offset, local_x_dim_}})); + local_weight->narrow({{0, dst_b_offset, local_bc_dim_}}) + ->copy_from(weight_->narrow({{0, src_b_offset, local_bc_dim_}})); + local_weight->narrow({{0, dst_c_offset, local_bc_dim_}}) + ->copy_from(weight_->narrow({{0, src_c_offset, local_bc_dim_}})); + + weight_ = infinicore::nn::Parameter(local_weight); + parameters_["weight"] = weight_; + } + + if (use_bias_ && bias_->size(0) != local_conv_dim_) { + if (bias_->shape() != infinicore::Shape{expected_full_conv_dim}) { + throw std::runtime_error("GraniteMoeHybridCausalConv1d: unexpected conv1d bias shape for TP slicing"); + } + + auto local_bias = infinicore::Tensor::empty( + {local_conv_dim_}, + bias_->dtype(), + bias_->device()); + + local_bias->narrow({{0, dst_x_offset, local_x_dim_}}) + ->copy_from(bias_->narrow({{0, src_x_offset, local_x_dim_}})); + local_bias->narrow({{0, dst_b_offset, local_bc_dim_}}) + ->copy_from(bias_->narrow({{0, src_b_offset, local_bc_dim_}})); + local_bias->narrow({{0, dst_c_offset, local_bc_dim_}}) + ->copy_from(bias_->narrow({{0, src_c_offset, local_bc_dim_}})); + + bias_ = infinicore::nn::Parameter(local_bias); + parameters_["bias"] = bias_; + } +} + +infinicore::Tensor GraniteMoeHybridCausalConv1d::forward( + const infinicore::Tensor &input) const { + auto &forward_context = infinilm::global_state::get_forward_context(); + auto &mamba_metadata = forward_context.mamba_metadata; + + std::optional bias = std::nullopt; + if (use_bias_) { + bias = bias_->narrow({{0, 0, local_conv_dim_}}); + } + + auto conv_out = infinicore::op::causal_conv1d( + input, + forward_context.conv_state_vec[layer_idx_], + weight_->narrow({{0, 0, local_conv_dim_}}), + bias, + mamba_metadata.input_offsets, + mamba_metadata.init_state_indices, + mamba_metadata.final_state_indices); + return infinicore::op::silu(conv_out); +} + +GraniteMoeHybridRMSNormGated::GraniteMoeHybridRMSNormGated( + std::shared_ptr model_config, + const infinicore::Device &device) { + const auto &dtype{model_config->get_dtype()}; + const size_t hidden_size = model_config->get("hidden_size"); + const size_t mamba_expand = model_config->get_or("mamba_expand", 2); + const size_t intermediate_size = mamba_expand * hidden_size; + eps_ = model_config->get("rms_norm_eps"); + + INFINICORE_NN_PARAMETER_INIT(weight, ({intermediate_size}, dtype, device)); +} + +infinicore::Tensor GraniteMoeHybridRMSNormGated::forward( + const infinicore::Tensor &hidden_states, + std::optional gate) const { + auto input = hidden_states; + if (gate.has_value()) { + input = infinicore::op::mul(input, infinicore::op::silu(gate.value())); + } + const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); + const size_t local_size = input->size(2); + if (rank_info.tp_size > 1) { + input = infinicore::op::distributed::allgather(input->contiguous(), rank_info.tp_size, rank_info.comm) + ->view({static_cast(rank_info.tp_size), input->size(0), input->size(1), local_size}) + ->permute({1, 2, 0, 3}) + ->contiguous() + ->view({input->size(0), input->size(1), local_size * rank_info.tp_size}); + } + auto output = infinicore::op::rms_norm(input, weight_, static_cast(eps_)); + if (rank_info.tp_size > 1) { + output = output->narrow({{2, rank_info.tp_rank * local_size, local_size}})->contiguous(); + } + return output; +} + +GraniteMoeHybridMamba::GraniteMoeHybridMamba( + std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device) { + const auto &dtype{model_config->get_dtype()}; + const size_t hidden_size = model_config->get("hidden_size"); + const size_t mamba_n_heads = model_config->get("mamba_n_heads"); + const size_t mamba_n_groups = model_config->get_or("mamba_n_groups", 1); + const size_t mamba_d_head = model_config->get("mamba_d_head"); + const size_t mamba_d_state = model_config->get("mamba_d_state"); + const size_t mamba_expand = model_config->get_or("mamba_expand", 2); + const size_t intermediate_size = mamba_expand * hidden_size; + const bool mamba_proj_bias = model_config->get_or("mamba_proj_bias", false); + if (intermediate_size != mamba_n_heads * mamba_d_head) { + throw std::runtime_error("GraniteMoeHybridMamba: intermediate_size must equal mamba_n_heads * mamba_d_head"); + } + + const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); + intermediate_size_ = intermediate_size / rank_info.tp_size; + num_heads_ = mamba_n_heads / rank_info.tp_size; + num_groups_ = mamba_n_groups >= static_cast(rank_info.tp_size) ? mamba_n_groups / rank_info.tp_size : 1; + head_dim_ = mamba_d_head; + state_size_ = mamba_d_state; + const size_t local_bc_size = num_groups_ * state_size_; + conv_dim_ = intermediate_size_ + 2 * local_bc_size; + const size_t local_projection_size = intermediate_size_ + conv_dim_ + num_heads_; + + in_proj_ = std::make_shared( + hidden_size, local_projection_size * rank_info.tp_size, + model_config->get_quantization_method(), mamba_proj_bias, dtype, device, + rank_info.tp_rank, rank_info.tp_size); + const std::vector splits = { + {"in_proj.gate", 0, intermediate_size_}, + {"in_proj.x", intermediate_size_, intermediate_size_}, + {"in_proj.B", 2 * intermediate_size_, local_bc_size, mamba_n_groups}, + {"in_proj.C", 2 * intermediate_size_ + local_bc_size, local_bc_size, mamba_n_groups}, + {"in_proj.dt", intermediate_size_ + conv_dim_, num_heads_}, + }; + for (auto ¶m : in_proj_->split_params(splits, rank_info.tp_rank, rank_info.tp_size, mamba_n_groups)) { + this->register_parameter(param.full_name, std::move(param.param)); + } + + INFINICORE_NN_MODULE_INIT(conv1d, model_config, layer_idx, device); + INFINICORE_NN_PARAMETER_INIT(dt_bias, ({mamba_n_heads}, dtype, device, 0, rank_info.tp_rank, rank_info.tp_size)); + INFINICORE_NN_PARAMETER_INIT(A_log, ({mamba_n_heads}, dtype, device, 0, rank_info.tp_rank, rank_info.tp_size)); + INFINICORE_NN_PARAMETER_INIT(D, ({mamba_n_heads}, dtype, device, 0, rank_info.tp_rank, rank_info.tp_size)); + INFINICORE_NN_MODULE_INIT(norm, model_config, device); + INFINICORE_NN_MODULE_INIT( + out_proj, intermediate_size, hidden_size, + model_config->get_quantization_method(), false, dtype, device, + rank_info.tp_rank, rank_info.tp_size, rank_info.comm); + if (mamba_proj_bias) { + out_proj_bias_ = infinicore::nn::Parameter({hidden_size}, dtype, device); + this->register_parameter("out_proj.bias", out_proj_bias_); + } + out_proj_->set_alpha(model_config->get_or("residual_multiplier", 1.0f)); +} + +infinicore::Tensor GraniteMoeHybridMamba::forward( + const infinicore::Tensor &hidden_states) const { + + const auto &hidden_shape = hidden_states->shape(); + const size_t batch_size = hidden_shape[0]; + const size_t seq_len = hidden_shape[1]; + auto projected_input = hidden_states; + auto projected_states = in_proj_->forward(projected_input); + + auto gate = projected_states->narrow({{2, 0, intermediate_size_}})->contiguous(); + auto conv_input = projected_states->narrow( + {{2, intermediate_size_, conv_dim_}}); + auto dt = projected_states->narrow( + {{2, intermediate_size_ + conv_dim_, num_heads_}}); + + auto conv_output = conv1d_->forward(conv_input); + + auto x = conv_output->narrow({{2, 0, intermediate_size_}})->contiguous(); + auto b = conv_output->narrow({{2, intermediate_size_, num_groups_ * state_size_}})->contiguous(); + auto c = conv_output->narrow({{2, intermediate_size_ + num_groups_ * state_size_, num_groups_ * state_size_}})->contiguous(); + + dt = infinicore::op::broadcast_to( + dt->view({batch_size, seq_len, num_heads_, 1}), + {static_cast(batch_size), + static_cast(seq_len), + static_cast(num_heads_), + static_cast(head_dim_)}) + ->view({batch_size, seq_len, intermediate_size_}); + + auto a_log = infinicore::op::broadcast_to( + A_log_->view({num_heads_, 1, 1}), + {static_cast(num_heads_), + static_cast(head_dim_), + static_cast(state_size_)}) + ->view({intermediate_size_, state_size_}); + auto d = infinicore::op::broadcast_to( + D_->view({num_heads_, 1}), + {static_cast(num_heads_), static_cast(head_dim_)}) + ->view({intermediate_size_}); + auto dt_bias = infinicore::op::broadcast_to( + dt_bias_->view({num_heads_, 1}), + {static_cast(num_heads_), static_cast(head_dim_)}) + ->view({intermediate_size_}); + + const infinicore::Shape ssm_state_shape{batch_size, intermediate_size_, state_size_}; + if (!ssm_state_ || ssm_state_->shape() != ssm_state_shape) { + ssm_state_ = infinicore::Tensor::zeros( + ssm_state_shape, + infinicore::DataType::F32, + hidden_states->device()); + } + + auto scan_output = infinicore::Tensor::empty(x->shape(), x->dtype(), x->device()); + const size_t group_size = intermediate_size_ / num_groups_; + for (size_t group = 0; group < num_groups_; ++group) { + auto group_state = ssm_state_->narrow({{1, group * group_size, group_size}})->contiguous(); + auto group_output = infinicore::op::mamba_selective_scan( + x->narrow({{2, group * group_size, group_size}})->contiguous(), + dt->narrow({{2, group * group_size, group_size}})->contiguous(), + b->narrow({{2, group * state_size_, state_size_}})->contiguous(), + c->narrow({{2, group * state_size_, state_size_}})->contiguous(), + a_log->narrow({{0, group * group_size, group_size}})->contiguous(), + d->narrow({{0, group * group_size, group_size}})->contiguous(), + gate->narrow({{2, group * group_size, group_size}})->contiguous(), + dt_bias->narrow({{0, group * group_size, group_size}})->contiguous(), + group_state); + scan_output->narrow({{2, group * group_size, group_size}})->copy_from(group_output); + if (num_groups_ > 1) { + ssm_state_->narrow({{1, group * group_size, group_size}})->copy_from(group_state); + } + } + + auto normalized = norm_->forward(scan_output); + auto output = out_proj_->forward(normalized); + if (out_proj_bias_) { + infinicore::op::add_(output, output, out_proj_bias_->as_strided(output->shape(), {0, 0, 1})); + } + return output; +} + +} // namespace infinilm::models::granitemoehybrid diff --git a/csrc/models/granitemoehybrid/granitemoehybrid_mamba.hpp b/csrc/models/granitemoehybrid/granitemoehybrid_mamba.hpp new file mode 100644 index 000000000..82594a390 --- /dev/null +++ b/csrc/models/granitemoehybrid/granitemoehybrid_mamba.hpp @@ -0,0 +1,86 @@ +#pragma once + +#include "../../layers/common_modules.hpp" + +#include +#include + +namespace infinilm::models::granitemoehybrid { + +class GraniteMoeHybridCausalConv1d : public infinicore::nn::Module { +public: + GraniteMoeHybridCausalConv1d( + std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &input) const; + void process_weights_after_loading() override; + +private: + INFINICORE_NN_PARAMETER(weight); + INFINICORE_NN_PARAMETER(bias); + size_t layer_idx_; + bool use_bias_; + size_t tp_size_; + size_t tp_rank_; + size_t conv_kernel_dim_; + size_t full_x_dim_; + size_t full_bc_dim_; + size_t local_x_dim_; + size_t local_bc_dim_; + size_t local_conv_dim_; + size_t bc_replicas_; +}; + +class GraniteMoeHybridRMSNormGated : public infinicore::nn::Module { +public: + GraniteMoeHybridRMSNormGated( + std::shared_ptr model_config, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states, + std::optional gate = std::nullopt) const; + +private: + INFINICORE_NN_PARAMETER(weight); + double eps_; +}; + +class GraniteMoeHybridMamba : public infinicore::nn::Module { +public: + GraniteMoeHybridMamba(std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + void process_weights_after_loading() override { + in_proj_->process_weights_after_loading(); + } + + void reset_runtime_state() const override { + in_proj_->reset_runtime_state(); + ssm_state_ = infinicore::Tensor(); + } + +private: + INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, in_proj); + INFINICORE_NN_MODULE(GraniteMoeHybridCausalConv1d, conv1d); + INFINICORE_NN_PARAMETER(dt_bias); + INFINICORE_NN_PARAMETER(A_log); + INFINICORE_NN_MODULE(GraniteMoeHybridRMSNormGated, norm); + INFINICORE_NN_PARAMETER(D); + INFINICORE_NN_MODULE(infinilm::layers::linear::RowParallelLinear, out_proj); + INFINICORE_NN_PARAMETER(out_proj_bias); + + size_t intermediate_size_; + size_t num_heads_; + size_t num_groups_; + size_t head_dim_; + size_t state_size_; + size_t conv_dim_; + + mutable infinicore::Tensor ssm_state_; +}; + +} // namespace infinilm::models::granitemoehybrid diff --git a/csrc/models/granitemoehybrid/granitemoehybrid_shared_mlp.cpp b/csrc/models/granitemoehybrid/granitemoehybrid_shared_mlp.cpp new file mode 100644 index 000000000..ca817f783 --- /dev/null +++ b/csrc/models/granitemoehybrid/granitemoehybrid_shared_mlp.cpp @@ -0,0 +1,65 @@ +#include "granitemoehybrid_shared_mlp.hpp" + +#include "../../global_state/global_state.hpp" +#include "infinicore/ops.hpp" + +#include +#include + +namespace infinilm::models::granitemoehybrid { + +GraniteMoeHybridSharedMLP::GraniteMoeHybridSharedMLP( + std::shared_ptr model_config, + const infinicore::Device &device) { + hidden_size_ = model_config->get("hidden_size"); + intermediate_size_ = model_config->get("shared_intermediate_size"); + + const std::string hidden_act = model_config->get_or("hidden_act", "silu"); + + const auto &dtype = model_config->get_dtype(); + const engine::distributed::RankInfo &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); + if (rank_info.tp_size <= 0 || intermediate_size_ % static_cast(rank_info.tp_size) != 0) { + throw std::runtime_error( + "infinilm::models::granitemoehybrid::GraniteMoeHybridSharedMLP: " + "shared_intermediate_size must be divisible by tp_size"); + } + + auto quantization_method = model_config->get_quantization_method(); + auto register_fn = + [this](const std::string &name, infinicore::nn::Parameter parameter) { + this->register_parameter(name, std::move(parameter)); + }; + input_linear_ = std::make_shared( + hidden_size_, + intermediate_size_, + "input_linear.gate", + "input_linear.up", + register_fn, + quantization_method, + false, + dtype, + device, + rank_info); + output_linear_ = this->register_module( + "output_linear", + intermediate_size_, + hidden_size_, + quantization_method, + false, + dtype, + device, + rank_info.tp_rank, + rank_info.tp_size, + rank_info.comm); + output_linear_->set_alpha(model_config->get_or("residual_multiplier", 1.0f)); +} + +infinicore::Tensor GraniteMoeHybridSharedMLP::forward( + const infinicore::Tensor &hidden_states) const { + auto input = hidden_states; + auto gate_up = input_linear_->forward(input); + auto activated = infinicore::op::silu_and_mul(gate_up); + return output_linear_->forward(activated); +} + +} // namespace infinilm::models::granitemoehybrid diff --git a/csrc/models/granitemoehybrid/granitemoehybrid_shared_mlp.hpp b/csrc/models/granitemoehybrid/granitemoehybrid_shared_mlp.hpp new file mode 100644 index 000000000..910976e5d --- /dev/null +++ b/csrc/models/granitemoehybrid/granitemoehybrid_shared_mlp.hpp @@ -0,0 +1,33 @@ +#pragma once + +#include "../../layers/common_modules.hpp" + +#include + +namespace infinilm::models::granitemoehybrid { + +class GraniteMoeHybridSharedMLP : public infinicore::nn::Module { +public: + GraniteMoeHybridSharedMLP( + std::shared_ptr model_config, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + + void process_weights_after_loading() override { + input_linear_->process_weights_after_loading(); + } + + void reset_runtime_state() const override { + input_linear_->reset_runtime_state(); + } + +private: + INFINICORE_NN_MODULE(infinilm::layers::linear::GateUpParallelLinear, input_linear); + INFINICORE_NN_MODULE(infinilm::layers::linear::RowParallelLinear, output_linear); + + size_t hidden_size_{0}; + size_t intermediate_size_{0}; +}; + +} // namespace infinilm::models::granitemoehybrid diff --git a/csrc/models/granitemoehybrid/granitemoehybrid_sparse_moe_block.cpp b/csrc/models/granitemoehybrid/granitemoehybrid_sparse_moe_block.cpp new file mode 100644 index 000000000..5550aca96 --- /dev/null +++ b/csrc/models/granitemoehybrid/granitemoehybrid_sparse_moe_block.cpp @@ -0,0 +1,206 @@ +#include "granitemoehybrid_sparse_moe_block.hpp" + +#include "../../global_state/global_state.hpp" +#include "infinicore/ops.hpp" +#include "infinicore/ops/mul_scalar.hpp" + +#include +#include + +namespace infinilm::models::granitemoehybrid { + +GraniteMoeHybridExpertMLP::GraniteMoeHybridExpertMLP( + std::shared_ptr model_config, + const infinicore::Device &device) { + const size_t hidden_size = model_config->get("hidden_size"); + const size_t intermediate_size = model_config->get("intermediate_size"); + const bool use_bias = model_config->get_or("mlp_bias", false); + const std::string hidden_act = model_config->get_or("hidden_act", "silu"); + + const auto &dtype = model_config->get_dtype(); + const engine::distributed::RankInfo &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); + if (rank_info.tp_size <= 0 || intermediate_size % static_cast(rank_info.tp_size) != 0) { + throw std::runtime_error( + "infinilm::models::granitemoehybrid::GraniteMoeHybridExpertMLP: " + "intermediate_size must be divisible by tp_size"); + } + + auto quantization_method = model_config->get_quantization_method(); + auto register_fn = + [this](const std::string &name, infinicore::nn::Parameter parameter) { + this->register_parameter(name, std::move(parameter)); + }; + input_linear_ = std::make_shared( + hidden_size, + intermediate_size, + "input_linear.gate", + "input_linear.up", + register_fn, + quantization_method, + use_bias, + dtype, + device, + rank_info); + output_linear_ = this->register_module( + "output_linear", + intermediate_size, + hidden_size, + quantization_method, + use_bias, + dtype, + device, + rank_info.tp_rank, + rank_info.tp_size, + rank_info.comm); +} + +infinicore::Tensor GraniteMoeHybridExpertMLP::forward( + const infinicore::Tensor &hidden_states) const { + auto input = hidden_states; + auto gate_up = input_linear_->forward(input); + auto activated = infinicore::op::silu_and_mul(gate_up); + return output_linear_->forward(activated); +} + +GraniteMoeHybridExperts::GraniteMoeHybridExperts( + std::shared_ptr model_config, + const infinicore::Device &device) { + num_experts_ = model_config->get("num_local_experts"); + num_experts_per_tok_ = model_config->get("num_experts_per_tok"); + residual_multiplier_ = model_config->get_or("residual_multiplier", 1.0f); + + if (num_experts_ == 0 || num_experts_per_tok_ == 0 + || num_experts_per_tok_ > num_experts_) { + throw std::runtime_error( + "infinilm::models::granitemoehybrid::GraniteMoeHybridExperts: " + "num_experts_per_tok must be in [1, num_local_experts]"); + } + + experts_.reserve(num_experts_); + for (size_t expert = 0; expert < num_experts_; ++expert) { + experts_.push_back( + this->register_module( + std::to_string(expert), model_config, device)); + } +} + +infinicore::Tensor GraniteMoeHybridExperts::forward( + const infinicore::Tensor &hidden_states, + const infinicore::Tensor &selected_experts, + const infinicore::Tensor &routing_weights) const { + if (hidden_states->ndim() != 2) { + throw std::runtime_error( + "infinilm::models::granitemoehybrid::GraniteMoeHybridExperts::forward: " + "hidden_states must have shape [num_tokens, hidden_size]"); + } + + auto selected_experts_cpu = selected_experts->to(infinicore::Device::Type::CPU); + auto routing_weights_cpu = routing_weights->to(infinicore::Device::Type::CPU); + const auto *selected_experts_ptr = reinterpret_cast(selected_experts_cpu->data()); + const auto *routing_weights_ptr = reinterpret_cast(routing_weights_cpu->data()); + + const size_t num_tokens = hidden_states->shape()[0]; + auto output = infinicore::Tensor::empty( + hidden_states->shape(), hidden_states->dtype(), hidden_states->device()); + for (size_t token = 0; token < num_tokens; ++token) { + auto token_input = hidden_states->narrow({{0, token, 1}}); + const size_t route_offset = token * num_experts_per_tok_; + infinicore::Tensor token_output; + + for (size_t route = 0; route < num_experts_per_tok_; ++route) { + const int expert = selected_experts_ptr[route_offset + route]; + if (expert < 0 || static_cast(expert) >= num_experts_) { + throw std::runtime_error( + "infinilm::models::granitemoehybrid::GraniteMoeHybridExperts::forward: " + "router selected an invalid expert index"); + } + + const float scale = routing_weights_ptr[route_offset + route] + * residual_multiplier_; + auto expert_output = experts_[expert]->forward(token_input); + expert_output = infinicore::op::mul_scalar(expert_output, scale); + if (route == 0) { + token_output = expert_output; + } else { + infinicore::op::add_( + token_output, token_output, expert_output); + } + } + output->narrow({{0, token, 1}})->copy_from(token_output); + } + return output; +} + +GraniteMoeHybridTopKRouter::GraniteMoeHybridTopKRouter( + std::shared_ptr model_config, + const infinicore::Device &device) { + const auto &dtype = model_config->get_dtype(); + const size_t hidden_size = model_config->get("hidden_size"); + const size_t num_experts = model_config->get("num_local_experts"); + num_experts_per_tok_ = model_config->get("num_experts_per_tok"); + + if (num_experts == 0 || num_experts_per_tok_ == 0 + || num_experts_per_tok_ > num_experts) { + throw std::runtime_error( + "infinilm::models::granitemoehybrid::GraniteMoeHybridTopKRouter: " + "num_experts_per_tok must be in [1, num_local_experts]"); + } + + INFINICORE_NN_MODULE_INIT( + layer, hidden_size, num_experts, false, dtype, device); +} + +std::tuple +GraniteMoeHybridTopKRouter::forward( + const infinicore::Tensor &hidden_states) const { + if (hidden_states->ndim() != 2) { + throw std::runtime_error( + "infinilm::models::granitemoehybrid::GraniteMoeHybridTopKRouter::forward: " + "hidden_states must have shape [num_tokens, hidden_size]"); + } + + const size_t num_tokens = hidden_states->shape()[0]; + auto input = hidden_states; + auto router_logits = layer_->forward(input); + auto router_scores = infinicore::Tensor::empty( + {num_tokens, num_experts_per_tok_}, + infinicore::DataType::F32, + hidden_states->device()); + auto router_indices = infinicore::Tensor::empty( + {num_tokens, num_experts_per_tok_}, + infinicore::DataType::I32, + hidden_states->device()); + + infinicore::op::topksoftmax( + router_scores, + router_indices, + router_logits, + num_experts_per_tok_, + true); + return {router_scores, router_indices}; +} + +GraniteMoeHybridSparseMoeBlock::GraniteMoeHybridSparseMoeBlock( + std::shared_ptr model_config, + const infinicore::Device &device) { + INFINICORE_NN_MODULE_INIT(router, model_config, device); + INFINICORE_NN_MODULE_INIT(experts, model_config, device); +} + +infinicore::Tensor GraniteMoeHybridSparseMoeBlock::forward( + const infinicore::Tensor &hidden_states) const { + if (hidden_states->ndim() != 3) { + throw std::runtime_error( + "infinilm::models::granitemoehybrid::GraniteMoeHybridSparseMoeBlock::forward: " + "hidden_states must have shape [batch_size, sequence_length, hidden_size]"); + } + + const auto &shape = hidden_states->shape(); + auto flat_hidden_states = hidden_states->view({shape[0] * shape[1], shape[2]}); + auto [routing_weights, selected_experts] = router_->forward(flat_hidden_states); + auto output = experts_->forward( + flat_hidden_states, selected_experts, routing_weights); + return output->view(shape); +} + +} // namespace infinilm::models::granitemoehybrid diff --git a/csrc/models/granitemoehybrid/granitemoehybrid_sparse_moe_block.hpp b/csrc/models/granitemoehybrid/granitemoehybrid_sparse_moe_block.hpp new file mode 100644 index 000000000..d0e3c1e1f --- /dev/null +++ b/csrc/models/granitemoehybrid/granitemoehybrid_sparse_moe_block.hpp @@ -0,0 +1,75 @@ +#pragma once + +#include "../../layers/common_modules.hpp" + +#include +#include + +namespace infinilm::models::granitemoehybrid { + +class GraniteMoeHybridExpertMLP : public infinicore::nn::Module { +public: + GraniteMoeHybridExpertMLP( + std::shared_ptr model_config, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + + void process_weights_after_loading() override { + input_linear_->process_weights_after_loading(); + } + + void reset_runtime_state() const override { + input_linear_->reset_runtime_state(); + } + +private: + INFINICORE_NN_MODULE(infinilm::layers::linear::GateUpParallelLinear, input_linear); + INFINICORE_NN_MODULE(infinilm::layers::linear::RowParallelLinear, output_linear); +}; + +class GraniteMoeHybridExperts : public infinicore::nn::Module { +public: + GraniteMoeHybridExperts( + std::shared_ptr model_config, + const infinicore::Device &device); + + infinicore::Tensor forward( + const infinicore::Tensor &hidden_states, + const infinicore::Tensor &selected_experts, + const infinicore::Tensor &routing_weights) const; + +private: + INFINICORE_NN_MODULE_VEC(GraniteMoeHybridExpertMLP, experts); + size_t num_experts_per_tok_{0}; + size_t num_experts_{0}; + float residual_multiplier_{1.0f}; +}; + +class GraniteMoeHybridTopKRouter : public infinicore::nn::Module { +public: + GraniteMoeHybridTopKRouter(std::shared_ptr model_config, + const infinicore::Device &device); + + std::tuple forward( + const infinicore::Tensor &hidden_states) const; + +private: + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, layer); + size_t num_experts_per_tok_{0}; +}; + +class GraniteMoeHybridSparseMoeBlock : public infinicore::nn::Module { +public: + GraniteMoeHybridSparseMoeBlock( + std::shared_ptr model_config, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + +protected: + INFINICORE_NN_MODULE(GraniteMoeHybridTopKRouter, router); + INFINICORE_NN_MODULE(GraniteMoeHybridExperts, experts); +}; + +} // namespace infinilm::models::granitemoehybrid diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py index 117d82f9a..46e6cb317 100644 --- a/python/infinilm/infer_engine.py +++ b/python/infinilm/infer_engine.py @@ -57,6 +57,9 @@ def model_uses_mamba_cache(config: dict) -> bool: return ( config.get("model_type") == "mamba" or llm_config.get("model_type") == "mamba" + or config.get("model_type") == "granitemoehybrid" + or llm_config.get("model_type") == "granitemoehybrid" + or "mamba" in layer_types or "linear_attention" in layer_types or all( key in llm_config diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index b210b0c56..403d3e104 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -18,7 +18,10 @@ def _get_scale_emb(model_path: str) -> float: raise FileNotFoundError(f"config.json not found at {config_path}") with open(config_path, "r") as f: config = json.load(f) - if config.get("model_type") not in ("fm9g", "minicpm"): + model_type = config.get("model_type") + if model_type in ("granite", "granitemoehybrid"): + return config.get("embedding_multiplier", 1.0) + if model_type not in ("fm9g", "minicpm"): return 1.0 return config.get("scale_emb", 1.0) @@ -162,7 +165,7 @@ def get_model_state_dict( load_state_dict(file_path, device=torch_device, dtype=torch_dtype) ) - # Apply scale_emb for fm9g models (embed_tokens uses lookup, not GEMM) + # Apply model-specific embedding scaling (embed_tokens uses lookup, not GEMM). scale_emb = _get_scale_emb(model_path) embed_tokens_unscaled = None if "model.embed_tokens.weight" in model_param: @@ -1052,6 +1055,128 @@ def _remap_qwen3_5_moe(state_dict, config): return remapped +def _remap_granitemoehybrid(state_dict, config=None): + """Unpack GraniteMoeHybrid Mamba, shared and routed-expert weights.""" + model_config = (config or {}).get("text_config", config or {}) + expected_num_experts = model_config.get("num_local_experts") + expected_intermediate_size = model_config.get("intermediate_size") + expected_shared_intermediate_size = model_config.get("shared_intermediate_size") + expert_weight_suffixes = ( + "input_linear.weight", + "output_linear.weight", + ) + + remapped = {} + for key, tensor in state_dict.items(): + if key.endswith((".mamba.in_proj.weight", ".mamba.in_proj.bias")): + intermediate_size = ( + model_config.get("mamba_expand", 2) * model_config["hidden_size"] + ) + bc_size = ( + model_config.get("mamba_n_groups", 1) * model_config["mamba_d_state"] + ) + sizes = ( + intermediate_size, + intermediate_size, + bc_size, + bc_size, + model_config["mamba_n_heads"], + ) + parameter_name = key.rsplit(".", 1)[1] + expected_ndim = 2 if parameter_name == "weight" else 1 + if tensor.ndim != expected_ndim or tensor.shape[0] != sum(sizes): + raise ValueError( + f"Expected GraniteMoeHybrid in_proj.{parameter_name} to " + f"have {expected_ndim} dimensions and {sum(sizes)} rows, " + f"got {tuple(tensor.shape)} for {key}" + ) + prefix = key[: -len(parameter_name)] + for name, part in zip( + ("gate", "x", "B", "C", "dt"), tensor.split(sizes, dim=0) + ): + remapped[f"{prefix}{name}.{parameter_name}"] = part.contiguous() + continue + + if key.endswith(".shared_mlp.input_linear.weight"): + if tensor.ndim != 2: + raise ValueError( + "Expected GraniteMoeHybrid shared input_linear.weight " + f"to be 2D, got {tensor.shape} for {key}" + ) + if tensor.shape[0] % 2 != 0: + raise ValueError( + "Expected GraniteMoeHybrid shared input_linear.weight " + f"output size to be even, got {tensor.shape[0]} for {key}" + ) + if ( + expected_shared_intermediate_size is not None + and tensor.shape[0] != 2 * expected_shared_intermediate_size + ): + raise ValueError( + "Expected GraniteMoeHybrid shared input_linear.weight " + f"output size {2 * expected_shared_intermediate_size}, " + f"got {tensor.shape[0]} for {key}" + ) + + gate, up = tensor.chunk(2, dim=0) + prefix = key[: -len("weight")] + remapped[f"{prefix}gate.weight"] = gate.contiguous() + remapped[f"{prefix}up.weight"] = up.contiguous() + continue + + matched_suffix = next( + ( + suffix + for suffix in expert_weight_suffixes + if key.endswith(f".block_sparse_moe.{suffix}") + ), + None, + ) + if matched_suffix is None: + remapped[key] = tensor + continue + + if tensor.ndim != 3: + raise ValueError( + f"Expected packed GraniteMoeHybrid {matched_suffix} to be 3D, " + f"got {tensor.shape} for {key}" + ) + if expected_num_experts is not None and tensor.shape[0] != expected_num_experts: + raise ValueError( + f"Expected {expected_num_experts} GraniteMoeHybrid experts, " + f"got {tensor.shape[0]} for {key}" + ) + + prefix = key[: -len(matched_suffix)] + for expert_idx, expert_weight in enumerate(tensor.unbind(0)): + expert_prefix = f"{prefix}experts.{expert_idx}." + if matched_suffix == "input_linear.weight": + if expert_weight.shape[0] % 2 != 0: + raise ValueError( + "Expected GraniteMoeHybrid expert input_linear.weight " + f"output size to be even, got {expert_weight.shape[0]} " + f"for {key}" + ) + if ( + expected_intermediate_size is not None + and expert_weight.shape[0] != 2 * expected_intermediate_size + ): + raise ValueError( + "Expected GraniteMoeHybrid expert input_linear.weight " + f"output size {2 * expected_intermediate_size}, got " + f"{expert_weight.shape[0]} for {key}" + ) + gate, up = expert_weight.chunk(2, dim=0) + remapped[f"{expert_prefix}input_linear.gate.weight"] = gate.contiguous() + remapped[f"{expert_prefix}input_linear.up.weight"] = up.contiguous() + else: + remapped[f"{expert_prefix}{matched_suffix}"] = ( + expert_weight.contiguous() + ) + + return remapped + + def _remap_kimi_k3(state_dict, config): """Adapt released Kimi-K3 KDA weights to the reference module layout.""" text_config = config.get("text_config", config) @@ -1082,5 +1207,6 @@ def _remap_kimi_k3(state_dict, config): "ernie4_5_moe_vl": _remap_ernie4_5_moe_vl, "qwen3_5_moe": _remap_qwen3_5_moe, "qwen3_next": _remap_qwen3_next, + "granitemoehybrid": _remap_granitemoehybrid, "kimi_k3": _remap_kimi_k3, } diff --git a/python/infinilm/processors/granitemoehybrid_processor.py b/python/infinilm/processors/granitemoehybrid_processor.py new file mode 100644 index 000000000..988a29b62 --- /dev/null +++ b/python/infinilm/processors/granitemoehybrid_processor.py @@ -0,0 +1,53 @@ +import infinicore +from typing_extensions import override + +from ..llm.scheduler import SchedulerOutput +from ..llm.static_scheduler import StaticSchedulerOutput +from .basic_llm_processor import BasicLLMProcessor +from .processor import register_processor + + +@register_processor("granitemoehybrid") +class GraniteMoeHybridProcessor(BasicLLMProcessor): + """Attach causal-convolution state rows to each Granite forward pass.""" + + @override + def build_model_inputs( + self, + scheduler_output: SchedulerOutput | StaticSchedulerOutput, + temperature: float = 1.0, + top_p: float = 0.8, + top_k: int = 1, + **kwargs, + ) -> dict: + model_inputs = super().build_model_inputs( + scheduler_output, + temperature, + top_p, + top_k, + **kwargs, + ) + + init_indices = [] + final_indices = [] + for request in scheduler_output.scheduled_requests: + if isinstance(scheduler_output, StaticSchedulerOutput): + state_index = 1 + else: + state_index = request.mamba_cache_index + if state_index is None: + raise RuntimeError( + f"Request {request.request_id} has no assigned mamba " + "cache index" + ) + + init_indices.append(0 if scheduler_output.is_prefill else state_index) + final_indices.append(state_index) + + model_inputs["mamba_init_state_indices"] = infinicore.from_list( + init_indices, dtype=infinicore.int32 + ) + model_inputs["mamba_final_state_indices"] = infinicore.from_list( + final_indices, dtype=infinicore.int32 + ) + return model_inputs diff --git a/test/models/granitemoehybrid/test_infer.py b/test/models/granitemoehybrid/test_infer.py new file mode 100644 index 000000000..827a6759d --- /dev/null +++ b/test/models/granitemoehybrid/test_infer.py @@ -0,0 +1,278 @@ +import gc +import logging +import os +import time + +from infinilm.base_config import BaseConfig +from infinilm.llm.llm import LLM +from infinilm.moe_config import configure_moe_ep_backend +from infinilm.processors.videonsa_processor import decode_video_frames + +DEFAULT_VIDEO_NUM_FRAMES = 8 + + +def test( + prompts: list[str], + model_path, + draft_model_path=None, + num_draft_tokens=4, + max_new_tokens=100, + device="cpu", + tp=1, + pp=1, + pp_stage=0, + master_addr="127.0.0.1", + master_port=29500, + moe_ep_backend="disabled", + ep=1, + enable_paged_attn=False, + enable_graph=False, + num_blocks=512, + block_size=256, + top_k=1, + top_p=1.0, + temperature=1.0, + attn_backend="default", + use_mla=False, + image_path=None, + video_path=None, + video_num_frames=None, + skip_load=False, + weight_load_mode="async", + use_legacy_moe=False, + enable_prefix_caching=True, + pre_transpose=False, +): + model_path = os.path.expanduser(model_path) + # ---------------------------------------------------------------------------- # + # Create Model + # ---------------------------------------------------------------------------- # + if enable_paged_attn and attn_backend == "default": + attn_backend = "paged-attn" + + import torch + from transformers import AutoModelForCausalLM, AutoTokenizer + + if top_k != 1: + raise ValueError("Hugging Face token comparison requires --top-k=1") + if image_path is not None or video_path is not None: + raise ValueError( + "The Granite Hugging Face comparison supports text prompts only" + ) + + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + print("Running Hugging Face (static cache, greedy decoding)...") + hf_model = AutoModelForCausalLM.from_pretrained( + model_path, + torch_dtype="auto", + trust_remote_code=True, + ).to(device) + hf_model.eval() + hf_results = [] + with torch.inference_mode(): + for prompt in prompts: + hf_inputs = tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + ) + prompt_token_ids = hf_inputs["input_ids"][0].tolist() + hf_inputs = {name: tensor.to(device) for name, tensor in hf_inputs.items()} + hf_sequences = hf_model.generate( + **hf_inputs, + do_sample=False, + max_new_tokens=max_new_tokens, + use_cache=True, + cache_implementation="static", + disable_compile=True, + ) + input_length = hf_inputs["input_ids"].shape[1] + hf_token_ids = hf_sequences[0, input_length:].cpu().tolist() + hf_results.append((prompt_token_ids, hf_token_ids)) + del hf_sequences, hf_inputs + del hf_model + gc.collect() + if device == "cuda": + torch.cuda.empty_cache() + + model = LLM( + model_path=model_path, + draft_model_path=draft_model_path, + num_draft_tokens=num_draft_tokens, + device=device, + tensor_parallel_size=tp, + pipeline_parallel_size=pp, + pipeline_parallel_stage=pp_stage, + master_addr=master_addr, + master_port=master_port, + moe_ep_backend=moe_ep_backend, + moe_ep_size=ep, + cache_type="paged" if enable_paged_attn else "static", + max_batch_size=len(prompts), + max_tokens=max_new_tokens, + num_blocks=num_blocks, + block_size=block_size, + temperature=temperature, + top_k=top_k, + top_p=top_p, + enable_graph=enable_graph, + attn_backend=attn_backend, + use_mla=use_mla, + skip_load=skip_load, + weight_load_mode=weight_load_mode, + use_legacy_moe=use_legacy_moe, + enable_prefix_caching=enable_prefix_caching, + pre_transpose=pre_transpose, + ) + + conversations = [ + [{"role": "user", "content": [{"type": "text", "text": prompt}]}] + for prompt in prompts + ] + if video_path is not None: + video_payload = decode_video_frames( + video_path, video_num_frames or DEFAULT_VIDEO_NUM_FRAMES + ) + for conversation in conversations: + conversation[0]["content"] = [ + {"type": "video_url", "video_url": {"url": video_payload}} + ] + conversation[0]["content"] + elif image_path is not None: + for conversation in conversations: + conversation[0]["content"] = [ + {"type": "image_url", "image_url": {"url": image_path}} + ] + conversation[0]["content"] + + t1 = time.time() + print("=================== start generate ====================") + + try: + outputs = model.chat( + messages=conversations, + ) + finally: + model.close() + t2 = time.time() + + all_matches = len(outputs) == len(hf_results) + for i, output in enumerate(outputs): + print(f"Resquest {i}:") + print("===Query===") + print(output.prompt) + print("===Response===") + print(output.outputs[0].text) + print("") + + prompt_token_ids, hf_token_ids = hf_results[i] + local_token_ids = output.outputs[0].token_ids + prompt_matches = output.prompt_token_ids == prompt_token_ids + mismatch = next( + ( + index + for index, (local_token, hf_token) in enumerate( + zip(local_token_ids, hf_token_ids) + ) + if local_token != hf_token + ), + None, + ) + if mismatch is None and len(local_token_ids) != len(hf_token_ids): + mismatch = min(len(local_token_ids), len(hf_token_ids)) + exact_match = prompt_matches and mismatch is None + all_matches = all_matches and exact_match + + print("=== Hugging Face ===") + print(f"token_ids: {hf_token_ids}") + print(tokenizer.decode(hf_token_ids, skip_special_tokens=True)) + print("=== InfiniLM ===") + print(f"token_ids: {local_token_ids}") + print("=== Comparison ===") + print(f"prompt token IDs match: {prompt_matches}") + print(f"generated token IDs match: {mismatch is None}") + if mismatch is not None: + hf_token = hf_token_ids[mismatch] if mismatch < len(hf_token_ids) else None + local_token = ( + local_token_ids[mismatch] if mismatch < len(local_token_ids) else None + ) + print( + f"first mismatch at generated token {mismatch}: " + f"Hugging Face={hf_token}, InfiniLM={local_token}" + ) + print(f"RESULT: {'PASS' if exact_match else 'FAIL'}") + print("") + + print( + f"total_time: {round((t2 - t1) * 1000, 2)} ms", + ) + + if not all_matches: + raise SystemExit(1) + + +if __name__ == "__main__": + cfg = BaseConfig() + logging.basicConfig( + level=getattr(logging, cfg.log_level.upper(), logging.INFO), + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + ) + if cfg.pp > 1 and cfg.node_rank > 0: + from infinilm.server.pipeline_worker import run_worker + + run_worker(cfg) + raise SystemExit(0) + + device_str = cfg.get_device_str(cfg.device) + + prompts = [cfg.prompt for _ in range(cfg.batch_size)] + + model_path = cfg.model + + max_new_tokens = cfg.max_new_tokens + + tp = cfg.tp + + enable_paged_attn = cfg.enable_paged_attn + + enable_graph = cfg.enable_graph + + if cfg.use_legacy_moe: + moe_ep_backend, ep = "disabled", 1 + else: + moe_ep_backend, ep = configure_moe_ep_backend( + cfg.tp, cfg.dp, cfg.ep, cfg.moe_ep_backend, cfg.model + ) + + test( + prompts, + model_path, + draft_model_path=cfg.draft_model, + num_draft_tokens=cfg.num_draft_tokens, + max_new_tokens=max_new_tokens, + device=device_str, + tp=tp, + pp=cfg.pp, + pp_stage=cfg.node_rank, + master_addr=cfg.master_addr, + master_port=cfg.master_port, + moe_ep_backend=moe_ep_backend, + ep=ep, + enable_paged_attn=enable_paged_attn, + enable_graph=enable_graph, + num_blocks=cfg.num_blocks, + block_size=cfg.block_size, + top_k=cfg.top_k, + top_p=cfg.top_p, + temperature=cfg.temperature, + attn_backend=cfg.attn, + use_mla=cfg.use_mla, + image_path=cfg.image, + video_path=cfg.video, + video_num_frames=cfg.video_num_frames, + skip_load=cfg.skip_load, + weight_load_mode=cfg.weight_load_mode, + use_legacy_moe=cfg.use_legacy_moe, + enable_prefix_caching=cfg.enable_prefix_caching, + pre_transpose=cfg.pre_transpose, + )