diff --git a/csrc/models/ernie4_5_vl/ernie4_5_attention.cpp b/csrc/models/ernie4_5_vl/ernie4_5_attention.cpp new file mode 100644 index 00000000..ac92c3bc --- /dev/null +++ b/csrc/models/ernie4_5_vl/ernie4_5_attention.cpp @@ -0,0 +1,305 @@ +#include "ernie4_5_attention.hpp" + +#include "../../global_state/global_state.hpp" +#include "../../utils.hpp" +#include "infinicore/ops.hpp" +#include "infinicore/ops/rope.hpp" + +#include +#include +#include +#include +#include +#include + +namespace infinilm::models::ernie4_5_vl { +namespace { + +std::pair build_group_rope_cache(size_t max_seq_len, + size_t rotary_dim, + size_t group_pairs, + size_t first_pair_idx, + size_t pair_stride, + double theta, + const infinicore::DataType &dtype, + const infinicore::Device &device) { + const size_t numel = max_seq_len * group_pairs; + std::vector sin_data(numel); + std::vector cos_data(numel); + for (size_t pos = 0; pos < max_seq_len; ++pos) { + for (size_t group_idx = 0; group_idx < group_pairs; ++group_idx) { + const size_t pair_idx = first_pair_idx + group_idx * pair_stride; + const float inv_freq = 1.0f / std::pow(static_cast(theta), 2.0f * static_cast(pair_idx) / static_cast(rotary_dim)); + const float angle = static_cast(pos) * inv_freq; + const size_t offset = pos * group_pairs + group_idx; + sin_data[offset] = std::sin(angle); + cos_data[offset] = std::cos(angle); + } + } + + const auto cpu = infinicore::Device::cpu(); + auto sin_cache = infinicore::Tensor::empty({max_seq_len, group_pairs}, dtype, device); + auto cos_cache = infinicore::Tensor::empty({max_seq_len, group_pairs}, dtype, device); + if (dtype == infinicore::DataType::F32) { + auto sin_cpu = infinicore::Tensor::from_blob(sin_data.data(), {max_seq_len, group_pairs}, infinicore::DataType::F32, cpu); + auto cos_cpu = infinicore::Tensor::from_blob(cos_data.data(), {max_seq_len, group_pairs}, infinicore::DataType::F32, cpu); + sin_cache->copy_from(sin_cpu); + cos_cache->copy_from(cos_cpu); + return {sin_cache, cos_cache}; + } + if (dtype == infinicore::DataType::BF16) { + std::vector sin_bf16(numel); + std::vector cos_bf16(numel); + for (size_t i = 0; i < numel; ++i) { + sin_bf16[i] = f32_to_bf16(sin_data[i]); + cos_bf16[i] = f32_to_bf16(cos_data[i]); + } + auto sin_cpu = infinicore::Tensor::from_blob(sin_bf16.data(), {max_seq_len, group_pairs}, infinicore::DataType::BF16, cpu); + auto cos_cpu = infinicore::Tensor::from_blob(cos_bf16.data(), {max_seq_len, group_pairs}, infinicore::DataType::BF16, cpu); + sin_cache->copy_from(sin_cpu); + cos_cache->copy_from(cos_cpu); + return {sin_cache, cos_cache}; + } + if (dtype == infinicore::DataType::F16) { + std::vector sin_f16(numel); + std::vector cos_f16(numel); + for (size_t i = 0; i < numel; ++i) { + sin_f16[i] = f32_to_f16(sin_data[i]); + cos_f16[i] = f32_to_f16(cos_data[i]); + } + auto sin_cpu = infinicore::Tensor::from_blob(sin_f16.data(), {max_seq_len, group_pairs}, infinicore::DataType::F16, cpu); + auto cos_cpu = infinicore::Tensor::from_blob(cos_f16.data(), {max_seq_len, group_pairs}, infinicore::DataType::F16, cpu); + sin_cache->copy_from(sin_cpu); + cos_cache->copy_from(cos_cpu); + return {sin_cache, cos_cache}; + } + throw std::runtime_error("infinilm::models::ernie4_5_vl::Ernie45Attention: unsupported RoPE cache dtype"); +} + +infinicore::Tensor axis_positions_for_rope(const infinicore::Tensor &position_ids, size_t axis, bool has_batch_dim) { + const auto pos_shape = position_ids->shape(); + if (pos_shape.size() != 3 || pos_shape[2] < 3) { + throw std::runtime_error("infinilm::models::ernie4_5_vl::Ernie45Attention: ERNIE MRoPE expects [batch, seq, 3] position_ids"); + } + auto pos = position_ids->narrow({{0, 0, 1}, {2, axis, 1}})->contiguous(); + if (has_batch_dim) { + return pos->view({pos_shape[0], pos_shape[1]}); + } + return pos->view({pos_shape[1]}); +} + +void apply_grouped_rope_one(infinicore::Tensor &x, + const infinicore::Tensor &position_ids, + size_t axis, + size_t group_pairs, + size_t first_pair_idx, + size_t pair_stride, + const infinicore::Tensor &sin_cache, + const infinicore::Tensor &cos_cache) { + const size_t ndim = x->ndim(); + if (ndim != 3 && ndim != 4) { + throw std::runtime_error("infinilm::models::ernie4_5_vl::Ernie45Attention: ERNIE grouped RoPE expects 3D or 4D q/k"); + } + const size_t last_dim = ndim - 1; + auto group_shape = x->shape(); + group_shape[last_dim] = group_pairs * 2; + auto group = infinicore::Tensor::empty(group_shape, x->dtype(), x->device()); + + for (size_t group_idx = 0; group_idx < group_pairs; ++group_idx) { + const size_t src_pair_idx = first_pair_idx + group_idx * pair_stride; + group->narrow({{last_dim, 2 * group_idx, 2}})->copy_from(x->narrow({{last_dim, 2 * src_pair_idx, 2}})); + } + + auto positions = axis_positions_for_rope(position_ids, axis, ndim == 4); + infinicore::op::rope_(group, group, positions, sin_cache, cos_cache, infinicore::nn::RoPE::Algo::GPT_J); + + for (size_t group_idx = 0; group_idx < group_pairs; ++group_idx) { + const size_t dst_pair_idx = first_pair_idx + group_idx * pair_stride; + x->narrow({{last_dim, 2 * dst_pair_idx, 2}})->copy_from(group->narrow({{last_dim, 2 * group_idx, 2}})); + } +} + +void apply_ernie_grouped_mrope(infinicore::Tensor &q, + infinicore::Tensor &k, + const infinicore::Tensor &position_ids, + const std::vector §ion, + const infinicore::Tensor &sin_h, + const infinicore::Tensor &cos_h, + const infinicore::Tensor &sin_w, + const infinicore::Tensor &cos_w, + const infinicore::Tensor &sin_t, + const infinicore::Tensor &cos_t) { + const size_t h_pairs = static_cast(section[0]); + const size_t w_pairs = static_cast(section[1]); + const size_t t_pairs = static_cast(section[2]); + const size_t t_first_pair = h_pairs + w_pairs; + + apply_grouped_rope_one(q, position_ids, 1, h_pairs, 0, 2, sin_h, cos_h); + apply_grouped_rope_one(q, position_ids, 2, w_pairs, 1, 2, sin_w, cos_w); + apply_grouped_rope_one(q, position_ids, 0, t_pairs, t_first_pair, 1, sin_t, cos_t); + apply_grouped_rope_one(k, position_ids, 1, h_pairs, 0, 2, sin_h, cos_h); + apply_grouped_rope_one(k, position_ids, 2, w_pairs, 1, 2, sin_w, cos_w); + apply_grouped_rope_one(k, position_ids, 0, t_pairs, t_first_pair, 1, sin_t, cos_t); +} + +} // namespace + +std::shared_ptr build_ernie45_mrope_cache(std::shared_ptr model_config, + const infinicore::Device &device) { + auto cache = std::make_shared(); + const size_t rotary_dim = model_config->get_rotary_dim(); + const double rope_theta = model_config->get("rope_theta"); + const auto &config_json = model_config->get_config_json(); + if (config_json.contains("rope_parameters") && config_json["rope_parameters"].contains("mrope_section")) { + cache->section = config_json["rope_parameters"]["mrope_section"].get>(); + } else if (config_json.contains("rope_scaling") && config_json["rope_scaling"].contains("mrope_section")) { + cache->section = config_json["rope_scaling"]["mrope_section"].get>(); + } + if (cache->section.size() != 3 || static_cast(cache->section[0] + cache->section[1] + cache->section[2]) * 2 != rotary_dim) { + throw std::runtime_error("infinilm::models::ernie4_5_vl::Ernie45Attention: invalid mrope_section"); + } + + const auto &dtype = model_config->get_dtype(); + const size_t max_position_embeddings = model_config->get("max_position_embeddings"); + auto h_cache = build_group_rope_cache(max_position_embeddings, rotary_dim, static_cast(cache->section[0]), 0, 2, rope_theta, dtype, device); + auto w_cache = build_group_rope_cache(max_position_embeddings, rotary_dim, static_cast(cache->section[1]), 1, 2, rope_theta, dtype, device); + auto t_cache = build_group_rope_cache(max_position_embeddings, + rotary_dim, + static_cast(cache->section[2]), + static_cast(cache->section[0] + cache->section[1]), + 1, + rope_theta, + dtype, + device); + cache->sin_h = h_cache.first; + cache->cos_h = h_cache.second; + cache->sin_w = w_cache.first; + cache->cos_w = w_cache.second; + cache->sin_t = t_cache.first; + cache->cos_t = t_cache.second; + return cache; +} + +Ernie45Attention::Ernie45Attention(std::shared_ptr model_config, + size_t layer_idx, + std::shared_ptr mrope_cache, + const infinicore::Device &device) + : layer_idx_(layer_idx), + mrope_cache_(std::move(mrope_cache)) { + hidden_size_ = model_config->get("hidden_size"); + head_dim_ = model_config->get_head_dim(); + rotary_dim_ = model_config->get_rotary_dim(); + if (!mrope_cache_) { + throw std::runtime_error("infinilm::models::ernie4_5_vl::Ernie45Attention: mrope_cache is required"); + } + + 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("use_bias", false); + + 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 ((total_num_kv_heads < static_cast(tp_size)) || (0 != (total_num_kv_heads % static_cast(tp_size)))) { + throw std::runtime_error("infinilm::models::ernie4_5_vl::Ernie45Attention: num_key_value_heads must be divisible by tp_size"); + } + + num_attention_heads_ = total_num_heads / static_cast(tp_size); + num_key_value_heads_ = total_num_kv_heads / static_cast(tp_size); + + auto quantization_method = model_config->get_quantization_method(); + INFINICORE_NN_MODULE_INIT(q_proj, hidden_size_, total_num_heads * head_dim_, quantization_method, use_bias, dtype, device, tp_rank, tp_size); + INFINICORE_NN_MODULE_INIT(k_proj, hidden_size_, total_num_kv_heads * head_dim_, quantization_method, use_bias, dtype, device, tp_rank, tp_size); + INFINICORE_NN_MODULE_INIT(v_proj, hidden_size_, total_num_kv_heads * head_dim_, quantization_method, use_bias, dtype, device, tp_rank, tp_size); + INFINICORE_NN_MODULE_INIT(o_proj, total_num_heads * head_dim_, hidden_size_, quantization_method, use_bias, dtype, device, tp_rank, tp_size, rank_info.comm); + + rotary_emb_ = infinilm::layers::rotary_embedding::get_rope(model_config, device); + const float scaling = 1.0f / std::sqrt(static_cast(head_dim_)); + attn_ = std::make_shared(num_attention_heads_, head_dim_, scaling, num_key_value_heads_, layer_idx_, + kv_cache_k_scale_, kv_cache_v_scale_, attention_backend_); + + auto register_fn = [this](const std::string &n, infinicore::nn::Parameter p) { this->register_parameter(n, std::move(p)); }; + infinilm::layers::attention::init_kv_cache_quant_params(register_fn, device, kv_cache_k_scale_, kv_cache_v_scale_); +} + +infinicore::Tensor Ernie45Attention::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 Ernie45Attention::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 q = q_proj_->forward(hidden_states_mutable)->view({batch_size, seq_len, num_attention_heads_, head_dim_}); + auto k = k_proj_->forward(hidden_states_mutable)->view({batch_size, seq_len, num_key_value_heads_, head_dim_}); + auto v = v_proj_->forward(hidden_states_mutable)->view({batch_size, seq_len, num_key_value_heads_, head_dim_}); + + infinicore::Tensor pos_ids_for_rope = position_ids; + const auto pos_shape = position_ids->shape(); + if (pos_shape.size() == 2) { + pos_ids_for_rope = position_ids->narrow({{0, 0, 1}})->view({pos_shape[1]}); + } else if (pos_shape.size() == 3) { + pos_ids_for_rope = position_ids->narrow({{0, 0, 1}, {2, 0, 1}})->view({pos_shape[1]}); + } else if (pos_shape.size() != 1) { + throw std::runtime_error("infinilm::models::ernie4_5_vl::Ernie45Attention: unsupported position_ids shape"); + } + + auto q_rotary = q->narrow({{3, 0, rotary_dim_}}); + auto k_rotary = k->narrow({{3, 0, rotary_dim_}}); + if (pos_shape.size() == 3) { + apply_ernie_grouped_mrope(q, k, position_ids, mrope_cache_->section, mrope_cache_->sin_h, mrope_cache_->cos_h, mrope_cache_->sin_w, mrope_cache_->cos_w, mrope_cache_->sin_t, mrope_cache_->cos_t); + } else { + rotary_emb_->forward(q_rotary, pos_ids_for_rope, true); + rotary_emb_->forward(k_rotary, pos_ids_for_rope, true); + } + + auto attn_output = attn_->forward(q, k, v); + return o_proj_->forward(attn_output); +} + +infinicore::Tensor Ernie45Attention::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 q = q_proj_->forward(hidden_states_mutable)->view({seq_len, num_attention_heads_, head_dim_}); + auto k = k_proj_->forward(hidden_states_mutable)->view({seq_len, num_key_value_heads_, head_dim_}); + auto v = v_proj_->forward(hidden_states_mutable)->view({seq_len, num_key_value_heads_, head_dim_}); + + infinicore::Tensor pos_ids_for_rope = position_ids; + const auto pos_shape = position_ids->shape(); + if (pos_shape.size() == 2) { + pos_ids_for_rope = position_ids->narrow({{0, 0, 1}})->view({pos_shape[1]}); + } else if (pos_shape.size() == 3) { + pos_ids_for_rope = position_ids->narrow({{0, 0, 1}, {2, 0, 1}})->view({pos_shape[1]}); + } else if (pos_shape.size() != 1) { + throw std::runtime_error("infinilm::models::ernie4_5_vl::Ernie45Attention: unsupported position_ids shape"); + } + + auto q_rotary = q->narrow({{2, 0, rotary_dim_}}); + auto k_rotary = k->narrow({{2, 0, rotary_dim_}}); + if (pos_shape.size() == 3) { + apply_ernie_grouped_mrope(q, k, position_ids, mrope_cache_->section, mrope_cache_->sin_h, mrope_cache_->cos_h, mrope_cache_->sin_w, mrope_cache_->cos_w, mrope_cache_->sin_t, mrope_cache_->cos_t); + } else { + rotary_emb_->forward(q_rotary, pos_ids_for_rope, true); + rotary_emb_->forward(k_rotary, pos_ids_for_rope, true); + } + + auto attn_output = attn_->forward(q, k, v); + return o_proj_->forward(attn_output); +} + +} // namespace infinilm::models::ernie4_5_vl diff --git a/csrc/models/ernie4_5_vl/ernie4_5_attention.hpp b/csrc/models/ernie4_5_vl/ernie4_5_attention.hpp new file mode 100644 index 00000000..4cef9ac6 --- /dev/null +++ b/csrc/models/ernie4_5_vl/ernie4_5_attention.hpp @@ -0,0 +1,66 @@ +#pragma once + +#include "../../config/model_config.hpp" +#include "../../layers/attention/attention.hpp" +#include "../../layers/common_modules.hpp" +#include "../../layers/linear/linear.hpp" +#include "../../layers/rotary_embedding/rotary_embedding.hpp" +#include "infinicore/device.hpp" +#include "infinicore/nn/module.hpp" +#include "infinicore/tensor.hpp" + +#include +#include + +namespace infinilm::models::ernie4_5_vl { + +struct Ernie45MropeCache { + std::vector section{22, 22, 20}; + infinicore::Tensor sin_h; + infinicore::Tensor cos_h; + infinicore::Tensor sin_w; + infinicore::Tensor cos_w; + infinicore::Tensor sin_t; + infinicore::Tensor cos_t; +}; + +std::shared_ptr build_ernie45_mrope_cache(std::shared_ptr model_config, + const infinicore::Device &device); + +class Ernie45Attention : public infinicore::nn::Module { +public: + Ernie45Attention(std::shared_ptr model_config, + size_t layer_idx, + std::shared_ptr mrope_cache, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &positions, + const infinicore::Tensor &hidden_states) const; + +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; + + size_t layer_idx_{0}; + size_t hidden_size_{0}; + size_t head_dim_{0}; + size_t rotary_dim_{0}; + size_t num_attention_heads_{0}; + size_t num_key_value_heads_{0}; + std::shared_ptr mrope_cache_; + infinilm::backends::AttentionBackend attention_backend_; + + INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, q_proj); + INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, k_proj); + INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, v_proj); + INFINICORE_NN_MODULE(infinilm::layers::linear::RowParallelLinear, o_proj); + INFINICORE_NN_MODULE(infinicore::nn::RoPE, rotary_emb); + + infinicore::nn::Parameter kv_cache_k_scale_; + infinicore::nn::Parameter kv_cache_v_scale_; + std::shared_ptr attn_; +}; + +} // namespace infinilm::models::ernie4_5_vl diff --git a/csrc/models/ernie4_5_vl/ernie4_5_decoder_layer.cpp b/csrc/models/ernie4_5_vl/ernie4_5_decoder_layer.cpp new file mode 100644 index 00000000..75cee32a --- /dev/null +++ b/csrc/models/ernie4_5_vl/ernie4_5_decoder_layer.cpp @@ -0,0 +1,93 @@ +#include "ernie4_5_decoder_layer.hpp" + +#include "infinicore/ops.hpp" + +#include +#include + +namespace infinilm::models::ernie4_5_vl { +namespace { + +size_t min_json_size(const nlohmann::json &value, size_t fallback) { + if (value.is_array() && !value.empty()) { + size_t result = value.at(0).get(); + for (const auto &item : value) { + result = std::min(result, item.get()); + } + return result; + } + return value.is_number_unsigned() ? value.get() : fallback; +} + +size_t max_json_size(const nlohmann::json &value, size_t fallback) { + if (value.is_array() && !value.empty()) { + size_t result = value.at(0).get(); + for (const auto &item : value) { + result = std::max(result, item.get()); + } + return result; + } + return value.is_number_unsigned() ? value.get() : fallback; +} + +bool is_moe_layer(const nlohmann::json &config, size_t layer_idx) { + const size_t interval = config.value("moe_layer_interval", 1); + const size_t start = config.contains("moe_layer_start_index") ? min_json_size(config.at("moe_layer_start_index"), 0) : 0; + const size_t end = config.contains("moe_layer_end_index") ? max_json_size(config.at("moe_layer_end_index"), config.value("num_hidden_layers", 1) - 1) : config.value("num_hidden_layers", 1) - 1; + return config.value("use_moe", false) && interval > 0 && ((layer_idx + 1) % interval == 0) && layer_idx >= start && layer_idx <= end; +} + +} // namespace + +Ernie45DecoderLayer::Ernie45DecoderLayer(std::shared_ptr model_config, + size_t layer_idx, + std::shared_ptr mrope_cache, + const infinicore::Device &device) { + 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(self_attn, model_config, layer_idx, std::move(mrope_cache), device); + + use_moe_ = is_moe_layer(model_config->get_config_json(), layer_idx); + if (use_moe_) { + mlp_ = this->register_module("mlp", model_config, device); + // ERNIE also defines an unregistered mlp_text path for pure text tokens when multimodal + // experts are enabled. This first implementation routes all tokens through mlp. + } else { + mlp_ = this->register_module("mlp", model_config, device); + } +} + +std::tuple Ernie45DecoderLayer::forward(const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) { + input_layernorm_->forward_inplace(hidden_states, residual); + hidden_states = self_attn_->forward(positions, hidden_states); + + post_attention_layernorm_->forward_inplace(hidden_states, residual); + hidden_states = use_moe_ + ? std::static_pointer_cast(mlp_)->forward(hidden_states) + : std::static_pointer_cast(mlp_)->forward(hidden_states); + return {hidden_states, residual}; +} + +infinicore::Tensor Ernie45DecoderLayer::forward(const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states) { + auto residual = hidden_states; + hidden_states = input_layernorm_->forward(hidden_states); + 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 = use_moe_ + ? std::static_pointer_cast(mlp_)->forward(hidden_states) + : std::static_pointer_cast(mlp_)->forward(hidden_states); + hidden_states = infinicore::op::add(residual, hidden_states); + return hidden_states; +} + +} // namespace infinilm::models::ernie4_5_vl diff --git a/csrc/models/ernie4_5_vl/ernie4_5_decoder_layer.hpp b/csrc/models/ernie4_5_vl/ernie4_5_decoder_layer.hpp new file mode 100644 index 00000000..ae0f5767 --- /dev/null +++ b/csrc/models/ernie4_5_vl/ernie4_5_decoder_layer.hpp @@ -0,0 +1,38 @@ +#pragma once + +#include "../../layers/common_modules.hpp" +#include "ernie4_5_attention.hpp" +#include "ernie4_5_moe.hpp" +#include "infinicore/nn/module.hpp" +#include "infinicore/tensor.hpp" + +#include +#include + +namespace infinilm::models::ernie4_5_vl { + +class Ernie45DecoderLayer : public infinicore::nn::Module { +public: + Ernie45DecoderLayer(std::shared_ptr model_config, + size_t layer_idx, + std::shared_ptr mrope_cache, + 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); + +protected: + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, input_layernorm); + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, post_attention_layernorm); + INFINICORE_NN_MODULE(Ernie45Attention, self_attn); + std::shared_ptr mlp_; + +private: + bool use_moe_{false}; +}; + +} // namespace infinilm::models::ernie4_5_vl diff --git a/csrc/models/ernie4_5_vl/ernie4_5_for_causal_lm.cpp b/csrc/models/ernie4_5_vl/ernie4_5_for_causal_lm.cpp new file mode 100644 index 00000000..361721ba --- /dev/null +++ b/csrc/models/ernie4_5_vl/ernie4_5_for_causal_lm.cpp @@ -0,0 +1,136 @@ +#include "ernie4_5_for_causal_lm.hpp" + +#include "../../global_state/global_state.hpp" +#include "../models_registry.hpp" + +#include +#include + +namespace infinilm::models::ernie4_5_vl { +namespace { + +size_t get_first_size(const nlohmann::json &config, const char *key, size_t default_value) { + if (!config.contains(key) || config.at(key).is_null()) { + return default_value; + } + const auto &value = config.at(key); + if (value.is_array()) { + return value.empty() ? default_value : value.at(0).get(); + } + return value.get(); +} + +void normalize_ernie_config(nlohmann::json &config_json) { + if (!config_json.contains("dtype") && config_json.contains("torch_dtype")) { + config_json["dtype"] = config_json["torch_dtype"]; + } + if (!config_json.contains("head_dim")) { + config_json["head_dim"] = config_json["hidden_size"].get() / config_json["num_attention_heads"].get(); + } + if (!config_json.contains("partial_rotary_factor")) { + config_json["partial_rotary_factor"] = 1.0; + } + if (!config_json.contains("compression_ratio")) { + config_json["compression_ratio"] = 1.0; + } + if (!config_json.contains("use_flash_attention")) { + config_json["use_flash_attention"] = true; + } + if (!config_json.contains("attention_probs_dropout_prob")) { + config_json["attention_probs_dropout_prob"] = 0.0; + } + if (!config_json.contains("hidden_dropout_prob")) { + config_json["hidden_dropout_prob"] = 0.0; + } + if (!config_json.contains("mlp_bias")) { + config_json["mlp_bias"] = config_json.value("use_bias", false); + } + if (!config_json.contains("attention_bias")) { + config_json["attention_bias"] = config_json.value("use_bias", false); + } + if (!config_json.contains("norm_topk_prob")) { + config_json["norm_topk_prob"] = config_json.value("moe_norm_gate_logits", true); + } + if (!config_json.contains("num_experts")) { + config_json["num_experts"] = get_first_size(config_json, "moe_num_experts", 0); + } + if (!config_json.contains("num_experts_per_tok")) { + config_json["num_experts_per_tok"] = config_json.value("moe_k", 1); + } + if (!config_json.contains("use_moe")) { + const size_t num_experts = config_json.value("num_experts", 0); + config_json["use_moe"] = num_experts > 0; + } + if (!config_json.contains("moe_dropout_prob")) { + config_json["moe_dropout_prob"] = 0.0; + } + if (!config_json.contains("moe_reverse_token_drop")) { + config_json["moe_reverse_token_drop"] = false; + } + if (!config_json.contains("moe_group")) { + config_json["moe_group"] = "world"; + } + if (!config_json.contains("moe_all_to_all_dropout")) { + config_json["moe_all_to_all_dropout"] = 0.0; + } +} + +} // namespace + +Ernie45ForConditionalGeneration::Ernie45ForConditionalGeneration(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(); + auto &config_json = model_config->get_config_json(); + if (config_json.contains("vision_config") && config_json["vision_config"].is_object()) { + INFINICORE_NN_MODULE_INIT(vision_model, config_json["vision_config"], dtype, device); + } + + INFINICORE_NN_MODULE_INIT(model, model_config, device); + INFINICORE_NN_MODULE_INIT(lm_head, hidden_size, vocab_size, false, dtype, device); +} + +infinilm::InfinilmModel::Output Ernie45ForConditionalGeneration::forward(const infinilm::InfinilmModel::Input &input) const { + auto hidden_states = (input.pixel_values.has_value() && !input.pixel_values.value().empty()) + ? model_->forward(input, vision_model_.get()) + : model_->forward(input); + auto logits = lm_head_->forward(hidden_states); + return {logits}; +} + +void Ernie45ForConditionalGeneration::reset_cache(const cache::CacheConfig *cache_config) { + if (cache_config == nullptr) { + cache_config_.reset(); + return; + } + cache_config_ = cache_config->unique_copy(); + + 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(); + + const backends::AttentionBackend attention_backend = infinilm::global_state::get_infinilm_config().attention_backend; + forward_context.kv_cache_vec = std::move(default_allocate_kv_cache_tensors(cache_config, model_config_, attention_backend)); +} + +std::shared_ptr create_ernie4_5_moe_vl_model_config(std::shared_ptr model_config) { + const std::string model_type = model_config->get("model_type"); + if ("ernie4_5_moe_vl" != model_type) { + throw std::runtime_error("infinilm::models::ernie4_5_vl::create_ernie4_5_moe_vl_model_config: model_type is not ernie4_5_moe_vl"); + } + normalize_ernie_config(model_config->get_config_json()); + model_config->set_rope_algo(infinicore::nn::RoPE::Algo::GPT_J); + return model_config; +} + +} // namespace infinilm::models::ernie4_5_vl + +namespace { +INFINILM_REGISTER_CAUSAL_LM_MODEL( + ernie4_5_moe_vl, + infinilm::models::ernie4_5_vl::Ernie45ForConditionalGeneration, + infinilm::models::ernie4_5_vl::create_ernie4_5_moe_vl_model_config); +} // namespace diff --git a/csrc/models/ernie4_5_vl/ernie4_5_for_causal_lm.hpp b/csrc/models/ernie4_5_vl/ernie4_5_for_causal_lm.hpp new file mode 100644 index 00000000..dc19e03f --- /dev/null +++ b/csrc/models/ernie4_5_vl/ernie4_5_for_causal_lm.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include "ernie4_5_model.hpp" + +namespace infinilm::models::ernie4_5_vl { + +class Ernie45ForConditionalGeneration : public InfinilmModel { +public: + Ernie45ForConditionalGeneration(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(Ernie45VisionModel, vision_model); + INFINICORE_NN_MODULE(Ernie45Model, model); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, lm_head); +}; + +std::shared_ptr create_ernie4_5_moe_vl_model_config(std::shared_ptr model_config); + +} // namespace infinilm::models::ernie4_5_vl diff --git a/csrc/models/ernie4_5_vl/ernie4_5_model.cpp b/csrc/models/ernie4_5_vl/ernie4_5_model.cpp new file mode 100644 index 00000000..b51644b9 --- /dev/null +++ b/csrc/models/ernie4_5_vl/ernie4_5_model.cpp @@ -0,0 +1,133 @@ +#include "ernie4_5_model.hpp" + +#include "../../global_state/global_state.hpp" + +#include +#include +#include +#include +#include + +namespace infinilm::models::ernie4_5_vl { + +Ernie45Model::Ernie45Model(std::shared_ptr model_config, + const infinicore::Device &device) + : model_config_(model_config) { + const auto &dtype = model_config->get_dtype(); + auto &config_json = model_config->get_config_json(); + if (config_json.contains("vision_config") && config_json["vision_config"].is_object()) { + INFINICORE_NN_MODULE_INIT(resampler_model, config_json, dtype, device); + } + const size_t vocab_size = model_config->get("vocab_size"); + const size_t hidden_size = model_config->get("hidden_size"); + const size_t num_hidden_layers = model_config->get("num_hidden_layers"); + const double rms_norm_eps = model_config->get("rms_norm_eps"); + mrope_cache_ = build_ernie45_mrope_cache(model_config, device); + + INFINICORE_NN_MODULE_INIT(embed_tokens, vocab_size, hidden_size, std::nullopt, dtype, device); + layers_.reserve(num_hidden_layers); + for (size_t i = 0; i < num_hidden_layers; ++i) { + layers_.push_back(this->register_module("layers." + std::to_string(i), model_config, i, mrope_cache_, device)); + } + INFINICORE_NN_MODULE_INIT(norm, hidden_size, rms_norm_eps, dtype, device); +} + +void Ernie45Model::replace_embeddings(infinicore::Tensor inputs_embeds, + const infinicore::Tensor &vision_hidden, + const infinicore::Tensor &image_bound) const { + auto bounds_cpu = image_bound->to(infinicore::Device::cpu()); + auto bounds = reinterpret_cast(bounds_cpu->data()); + const int64_t start = bounds[0]; + const int64_t end = bounds[1]; + if (start < 0 || end < start) { + throw std::runtime_error("Ernie45Model: invalid image_bound"); + } + const size_t span = static_cast(end - start); + if (vision_hidden->size(0) != span) { + throw std::runtime_error("Ernie45Model: image feature length does not match image token span"); + } + + ASSERT_EQ(inputs_embeds->size(0), 1); + auto out_slice = inputs_embeds->squeeze(0); + out_slice->narrow({{0, static_cast(start), span}})->copy_from(vision_hidden); +} + +void Ernie45Model::apply_image_embeddings(infinicore::Tensor inputs_embeds, + const InfinilmModel::Input &input, + const Ernie45VisionModel &vision_model) const { + if (!input.image_bound.has_value() || !input.tgt_sizes.has_value()) { + throw std::runtime_error("Ernie45Model: image_bound and tgt_sizes/grid_thw must be provided with pixel_values"); + } + if (!resampler_model_) { + throw std::runtime_error("Ernie45Model: resampler_model is required for image input"); + } + const auto &pixel_values = input.pixel_values.value(); + const auto &image_bound = input.image_bound.value(); + const auto &grid_thw = input.tgt_sizes.value(); + if (pixel_values.size() != image_bound.size() || pixel_values.size() != grid_thw.size()) { + throw std::runtime_error("Ernie45Model: pixel_values, image_bound and grid_thw must have the same number of elements"); + } + auto &mm_metadata = global_state::get_forward_context().mm_metadata; + if (!mm_metadata.image_req_ids.has_value()) { + throw std::runtime_error("Ernie45Model: image_req_ids must be provided with pixel_values"); + } + const auto &image_req_ids = mm_metadata.image_req_ids.value(); + if (image_req_ids.size() != pixel_values.size()) { + throw std::runtime_error("Ernie45Model: image_req_ids count must match pixel_values count"); + } + if (!input.input_offsets.has_value()) { + throw std::runtime_error("Ernie45Model: input_offsets is required for multimodal replacement"); + } + + auto input_offsets_cpu = input.input_offsets.value()->to(infinicore::Device::cpu()); + auto *offsets = reinterpret_cast(input_offsets_cpu->data()); + std::vector visual_token_ranges; + for (size_t image_idx = 0; image_idx < pixel_values.size(); ++image_idx) { + const size_t req_id = image_req_ids[image_idx]; + auto bounds_cpu = image_bound[image_idx]->to(infinicore::Device::cpu()); + auto bounds = reinterpret_cast(bounds_cpu->data()); + const size_t packed_start = static_cast(offsets[req_id]) + static_cast(bounds[0]); + const size_t packed_end = static_cast(offsets[req_id]) + static_cast(bounds[1]); + visual_token_ranges.push_back(packed_start); + visual_token_ranges.push_back(packed_end); + auto req_embeds = inputs_embeds->narrow({{1, static_cast(offsets[req_id]), static_cast(offsets[req_id + 1] - offsets[req_id])}}); + auto image_features = vision_model.forward(pixel_values[image_idx], grid_thw[image_idx]); + auto vision_hidden = resampler_model_->forward(image_features, grid_thw[image_idx]); + replace_embeddings(req_embeds, vision_hidden, image_bound[image_idx]); + } + mm_metadata.visual_token_ranges = std::move(visual_token_ranges); +} + +infinicore::Tensor Ernie45Model::forward_embeds(infinicore::Tensor hidden_states, + const infinicore::Tensor &positions) const { + infinicore::Tensor residual; + for (auto &layer : layers_) { + layer->forward(positions, hidden_states, residual); + } + norm_->forward_inplace(hidden_states, residual); + return hidden_states; +} + +infinicore::Tensor Ernie45Model::forward(const InfinilmModel::Input &input) const { + return forward(input, nullptr); +} + +infinicore::Tensor Ernie45Model::forward(const InfinilmModel::Input &input, + const Ernie45VisionModel *vision_model) const { + auto input_ids = input.input_ids.value(); + auto positions = input.position_ids.value(); + auto hidden_states = embed_tokens_->forward(input_ids); + if (input.pixel_values.has_value() && !input.pixel_values.value().empty()) { + if (vision_model == nullptr) { + throw std::runtime_error("Ernie45Model: vision_model is required for image input"); + } + apply_image_embeddings(hidden_states, input, *vision_model); + } + return forward_embeds(hidden_states, positions); +} + +void Ernie45Model::reset_cache(const cache::CacheConfig *) { + // Cache tensors are allocated by Ernie45ForConditionalGeneration, which inherits InfinilmModel. +} + +} // namespace infinilm::models::ernie4_5_vl diff --git a/csrc/models/ernie4_5_vl/ernie4_5_model.hpp b/csrc/models/ernie4_5_vl/ernie4_5_model.hpp new file mode 100644 index 00000000..436a4ef6 --- /dev/null +++ b/csrc/models/ernie4_5_vl/ernie4_5_model.hpp @@ -0,0 +1,45 @@ +#pragma once + +#include "../../layers/common_modules.hpp" +#include "../infinilm_model.hpp" +#include "ernie4_5_decoder_layer.hpp" +#include "ernie4_5_vision.hpp" +#include "infinicore/nn/embedding.hpp" +#include "infinicore/nn/rmsnorm.hpp" + +#include + +namespace infinilm::models::ernie4_5_vl { + +class Ernie45Model : public infinicore::nn::Module { +public: + Ernie45Model(std::shared_ptr model_config, + const infinicore::Device &device); + + infinicore::Tensor forward(const InfinilmModel::Input &input) const; + infinicore::Tensor forward(const InfinilmModel::Input &input, + const Ernie45VisionModel *vision_model) const; + infinicore::Tensor forward_embeds(infinicore::Tensor hidden_states, + const infinicore::Tensor &positions) const; + + void reset_cache(const cache::CacheConfig *cache_config); + +protected: + INFINICORE_NN_MODULE(Ernie45ResamplerModel, resampler_model); + INFINICORE_NN_MODULE(infinicore::nn::Embedding, embed_tokens); + INFINICORE_NN_MODULE_VEC(Ernie45DecoderLayer, layers); + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, norm); + +private: + void replace_embeddings(infinicore::Tensor inputs_embeds, + const infinicore::Tensor &vision_hidden, + const infinicore::Tensor &image_bound) const; + void apply_image_embeddings(infinicore::Tensor inputs_embeds, + const InfinilmModel::Input &input, + const Ernie45VisionModel &vision_model) const; + + std::shared_ptr model_config_; + std::shared_ptr mrope_cache_; +}; + +} // namespace infinilm::models::ernie4_5_vl diff --git a/csrc/models/ernie4_5_vl/ernie4_5_moe.cpp b/csrc/models/ernie4_5_vl/ernie4_5_moe.cpp new file mode 100644 index 00000000..e3ee8e2d --- /dev/null +++ b/csrc/models/ernie4_5_vl/ernie4_5_moe.cpp @@ -0,0 +1,236 @@ +#include "ernie4_5_moe.hpp" + +#include "../../global_state/global_state.hpp" +#include "infinicore/ops.hpp" + +#include +#include +#include +#include +#include + +namespace infinilm::models::ernie4_5_vl { +namespace { + +std::vector get_size_list(const nlohmann::json &config, const char *key) { + std::vector values; + if (!config.contains(key) || config.at(key).is_null()) { + return values; + } + const auto &value = config.at(key); + if (value.is_array()) { + values.reserve(value.size()); + for (const auto &item : value) { + values.push_back(item.get()); + } + } else { + values.push_back(value.get()); + } + return values; +} + +} // namespace + +Ernie45TopKRouter::Ernie45TopKRouter(std::shared_ptr model_config, + const infinicore::Device &device) { + const size_t hidden_size = model_config->get("hidden_size"); + const size_t num_experts = model_config->get("num_experts"); + const auto expert_counts = get_size_list(model_config->get_config_json(), "moe_num_experts"); + const size_t num_router_groups = expert_counts.empty() ? 1 : expert_counts.size(); + num_experts_per_tok_ = model_config->get("num_experts_per_tok"); + norm_topk_prob_ = model_config->get("norm_topk_prob"); + + ASSERT((num_experts > 0) && (num_router_groups <= 2) && (num_experts_per_tok_ > 0) && (num_experts_per_tok_ <= num_experts * num_router_groups)); + + // ERNIE stores router weights as [hidden_size, group_num_experts], while Qwen3 stores + // [num_experts, hidden_size]. matmul(hidden, weight) matches ERNIE directly. + const auto &dtype = model_config->get_dtype(); + INFINICORE_NN_PARAMETER_INIT(weight, ({hidden_size, num_experts}, dtype, device)); + if (num_router_groups > 1) { + INFINICORE_NN_PARAMETER_INIT(weight_1, ({hidden_size, num_experts}, dtype, device)); + } +} + +std::tuple Ernie45TopKRouter::forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &correction_bias, + size_t group_idx) const { + ASSERT(hidden_states->ndim() == 2); + const auto router_weight = group_idx == 1 && weight_1_ ? static_cast(weight_1_) : static_cast(weight_); + ASSERT(hidden_states->dtype() == router_weight->dtype()); + auto router_logits = infinicore::op::matmul(hidden_states, router_weight); + return infinicore::op::moe_topk_softmax(router_logits, num_experts_per_tok_, norm_topk_prob_, 0.0f, correction_bias); +} + +Ernie45ExpertMLP::Ernie45ExpertMLP(std::shared_ptr model_config, + size_t intermediate_size, + const infinicore::Device &device) { + const auto &dtype = model_config->get_dtype(); + const size_t hidden_size = model_config->get("hidden_size"); + const bool use_bias = model_config->get_or("mlp_bias", false); + 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; + auto quantization_method = model_config->get_quantization_method(); + + INFINICORE_NN_MODULE_INIT(gate_proj, hidden_size, intermediate_size, quantization_method, use_bias, dtype, device, tp_rank, tp_size); + INFINICORE_NN_MODULE_INIT(up_proj, hidden_size, intermediate_size, quantization_method, use_bias, dtype, device, tp_rank, tp_size); + INFINICORE_NN_MODULE_INIT(down_proj, intermediate_size, hidden_size, quantization_method, use_bias, dtype, device, tp_rank, tp_size, rank_info.comm); +} + +infinicore::Tensor Ernie45ExpertMLP::forward(const infinicore::Tensor &hidden_states) const { + auto hidden_states_mutable = hidden_states; + auto gate = gate_proj_->forward(hidden_states_mutable); + auto up = up_proj_->forward(hidden_states_mutable); + auto intermediate = infinicore::op::swiglu(up, gate); + return down_proj_->forward(intermediate); +} + +Ernie45Experts::Ernie45Experts(std::shared_ptr model_config, + const infinicore::Device &device) { + const auto &config_json = model_config->get_config_json(); + auto expert_counts = get_size_list(config_json, "moe_num_experts"); + auto intermediate_sizes = get_size_list(config_json, "moe_intermediate_size"); + if (expert_counts.empty()) { + expert_counts.push_back(model_config->get("num_experts")); + } + if (intermediate_sizes.empty()) { + intermediate_sizes.push_back(model_config->get("moe_intermediate_size")); + } + if (intermediate_sizes.size() == 1 && expert_counts.size() > 1) { + intermediate_sizes.resize(expert_counts.size(), intermediate_sizes.front()); + } + if (expert_counts.size() != intermediate_sizes.size()) { + throw std::runtime_error("Ernie45Experts: moe_num_experts and moe_intermediate_size group counts differ"); + } + + const auto &dtype = model_config->get_dtype(); + const size_t hidden_size = model_config->get("hidden_size"); + const size_t text_intermediate_size = intermediate_sizes.front(); + const size_t vision_intermediate_size = intermediate_sizes.size() > 1 ? intermediate_sizes[1] : text_intermediate_size; + num_experts_per_tok_ = model_config->get("num_experts_per_tok"); + num_text_experts_ = expert_counts.front(); + num_vision_experts_ = expert_counts.size() > 1 ? expert_counts[1] : 0; + use_bias_ = model_config->get_or("mlp_bias", false); + + ASSERT((num_text_experts_ > 0) && (num_experts_per_tok_ > 0) && (num_experts_per_tok_ <= num_text_experts_)); + INFINICORE_NN_PARAMETER_INIT(w1, ({num_text_experts_, 2 * text_intermediate_size, hidden_size}, dtype, device)); + INFINICORE_NN_PARAMETER_INIT(w2, ({num_text_experts_, hidden_size, text_intermediate_size}, dtype, device)); + if (use_bias_) { + INFINICORE_NN_PARAMETER_INIT(b1, ({num_text_experts_, 2 * text_intermediate_size}, dtype, device)); + INFINICORE_NN_PARAMETER_INIT(b2, ({num_text_experts_, hidden_size}, dtype, device)); + } + if (num_vision_experts_ > 0) { + ASSERT(num_experts_per_tok_ <= num_vision_experts_); + INFINICORE_NN_PARAMETER_INIT(w1_1, ({num_vision_experts_, 2 * vision_intermediate_size, hidden_size}, dtype, device)); + INFINICORE_NN_PARAMETER_INIT(w2_1, ({num_vision_experts_, hidden_size, vision_intermediate_size}, dtype, device)); + if (use_bias_) { + INFINICORE_NN_PARAMETER_INIT(b1_1, ({num_vision_experts_, 2 * vision_intermediate_size}, dtype, device)); + INFINICORE_NN_PARAMETER_INIT(b2_1, ({num_vision_experts_, hidden_size}, dtype, device)); + } + } +} + +infinicore::Tensor Ernie45Experts::forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &top_k_index, + const infinicore::Tensor &top_k_weights, + size_t group_idx) const { + ASSERT(hidden_states->ndim() == 2); + ASSERT(top_k_index->ndim() == 2 && top_k_weights->ndim() == 2); + + const bool use_vision = group_idx == 1 && w1_1_ && w2_1_; + const auto w1 = use_vision ? static_cast(w1_1_) : static_cast(w1_); + const auto w2 = use_vision ? static_cast(w2_1_) : static_cast(w2_); + std::optional b1 = std::nullopt; + std::optional b2 = std::nullopt; + if (use_bias_) { + b1 = use_vision ? static_cast(b1_1_) : static_cast(b1_); + b2 = use_vision ? static_cast(b2_1_) : static_cast(b2_); + } + return infinicore::op::fused_moe(hidden_states, top_k_index, top_k_weights, w1, w2, b1, b2, + infinicore::op::FusedMoeActivation::Swiglu); +} + +Ernie45MoE::Ernie45MoE(std::shared_ptr model_config, + const infinicore::Device &device) { + num_experts_ = model_config->get("num_experts"); + const auto expert_counts = get_size_list(model_config->get_config_json(), "moe_num_experts"); + num_vision_experts_ = expert_counts.size() > 1 ? expert_counts[1] : 0; + INFINICORE_NN_MODULE_INIT(gate, model_config, device); + INFINICORE_NN_MODULE_INIT(experts, model_config, device); + + if (model_config->get_or("moe_use_aux_free", false)) { + size_t num_expert_groups = 1; + const auto &config_json = model_config->get_config_json(); + if (config_json.contains("moe_num_experts") && config_json["moe_num_experts"].is_array()) { + num_expert_groups = config_json["moe_num_experts"].size(); + } + const size_t num_experts = model_config->get("num_experts"); + e_score_correction_bias_ = infinicore::nn::Parameter({num_expert_groups, num_experts}, infinicore::DataType::F32, device); + this->register_parameter("moe_statics.e_score_correction_bias", e_score_correction_bias_); + } + + const size_t n_shared_experts = model_config->get_or("moe_num_shared_experts", 0); + has_shared_experts_ = n_shared_experts > 0; + if (has_shared_experts_) { + auto shared_config_json = model_config->get_config_json(); + auto intermediate_sizes = get_size_list(shared_config_json, "moe_intermediate_size"); + if (intermediate_sizes.empty()) { + throw std::runtime_error("Ernie45MoE: moe_intermediate_size is required for shared experts"); + } + shared_config_json["intermediate_size"] = intermediate_sizes.front() * n_shared_experts; + auto shared_config = std::make_shared(shared_config_json); + INFINICORE_NN_MODULE_INIT(shared_experts, shared_config, device); + } +} + +infinicore::Tensor Ernie45MoE::forward_group(const infinicore::Tensor &hidden_states_2d, size_t group_idx) const { + const size_t bias_row = group_idx == 1 ? 1 : 0; + auto correction_bias = e_score_correction_bias_ ? e_score_correction_bias_->narrow({{0, bias_row, 1}})->view({num_experts_}) : infinicore::Tensor(); + auto [routing_weights, selected_experts] = gate_->forward(hidden_states_2d, correction_bias, group_idx); + return experts_->forward(hidden_states_2d, selected_experts, routing_weights, group_idx); +} + +infinicore::Tensor Ernie45MoE::forward(const infinicore::Tensor &hidden_states) const { + ASSERT(hidden_states->ndim() == 3); + const auto shape = hidden_states->shape(); + auto hidden_states_reshaped = hidden_states->view({shape[0] * shape[1], shape[2]}); + + infinicore::Tensor final_hidden_states_2d; + const auto &mm_metadata = infinilm::global_state::get_forward_context().mm_metadata; + const bool has_visual_ranges = mm_metadata.visual_token_ranges.has_value() + && !mm_metadata.visual_token_ranges->empty() + && num_vision_experts_ > 0; + if (!has_visual_ranges) { + final_hidden_states_2d = forward_group(hidden_states_reshaped, 0); + } else { + final_hidden_states_2d = infinicore::Tensor::empty(hidden_states_reshaped->shape(), hidden_states_reshaped->dtype(), hidden_states_reshaped->device()); + const auto &ranges = mm_metadata.visual_token_ranges.value(); + const size_t total_tokens = hidden_states_reshaped->size(0); + size_t cursor = 0; + for (size_t i = 0; i + 1 < ranges.size(); i += 2) { + const size_t start = std::min(ranges[i], total_tokens); + const size_t end = std::min(ranges[i + 1], total_tokens); + if (cursor < start) { + const size_t len = start - cursor; + final_hidden_states_2d->narrow({{0, cursor, len}})->copy_from(forward_group(hidden_states_reshaped->narrow({{0, cursor, len}}), 0)); + } + if (start < end) { + const size_t len = end - start; + final_hidden_states_2d->narrow({{0, start, len}})->copy_from(forward_group(hidden_states_reshaped->narrow({{0, start, len}}), 1)); + } + cursor = std::max(cursor, end); + } + if (cursor < total_tokens) { + const size_t len = total_tokens - cursor; + final_hidden_states_2d->narrow({{0, cursor, len}})->copy_from(forward_group(hidden_states_reshaped->narrow({{0, cursor, len}}), 0)); + } + } + + auto final_hidden_states = final_hidden_states_2d->view(shape); + if (has_shared_experts_) { + final_hidden_states = infinicore::op::add(final_hidden_states, shared_experts_->forward(hidden_states)); + } + return final_hidden_states; +} + +} // namespace infinilm::models::ernie4_5_vl diff --git a/csrc/models/ernie4_5_vl/ernie4_5_moe.hpp b/csrc/models/ernie4_5_vl/ernie4_5_moe.hpp new file mode 100644 index 00000000..afc6fb95 --- /dev/null +++ b/csrc/models/ernie4_5_vl/ernie4_5_moe.hpp @@ -0,0 +1,89 @@ +#pragma once + +#include "../../layers/common_modules.hpp" +#include "../../layers/linear/linear.hpp" +#include "infinicore/nn/parameter.hpp" +#include "infinicore/tensor.hpp" + +#include +#include +#include + +namespace infinilm::models::ernie4_5_vl { + +class Ernie45TopKRouter : public infinicore::nn::Module { +public: + Ernie45TopKRouter(std::shared_ptr model_config, + const infinicore::Device &device); + + std::tuple forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &correction_bias = infinicore::Tensor(), + size_t group_idx = 0) const; + +protected: + INFINICORE_NN_PARAMETER(weight); + INFINICORE_NN_PARAMETER(weight_1); + size_t num_experts_per_tok_{0}; + bool norm_topk_prob_{false}; +}; + +class Ernie45ExpertMLP : public infinicore::nn::Module { +public: + Ernie45ExpertMLP(std::shared_ptr model_config, + size_t intermediate_size, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + void set_alpha(float alpha) { down_proj_->set_alpha(alpha); } + +protected: + INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, gate_proj); + INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, up_proj); + INFINICORE_NN_MODULE(infinilm::layers::linear::RowParallelLinear, down_proj); +}; + +class Ernie45Experts : public infinicore::nn::Module { +public: + Ernie45Experts(std::shared_ptr model_config, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &top_k_index, + const infinicore::Tensor &top_k_weights, + size_t group_idx = 0) const; + +protected: + INFINICORE_NN_PARAMETER(w1); + INFINICORE_NN_PARAMETER(w2); + INFINICORE_NN_PARAMETER(b1); + INFINICORE_NN_PARAMETER(b2); + INFINICORE_NN_PARAMETER(w1_1); + INFINICORE_NN_PARAMETER(w2_1); + INFINICORE_NN_PARAMETER(b1_1); + INFINICORE_NN_PARAMETER(b2_1); + size_t num_experts_per_tok_{0}; + size_t num_text_experts_{0}; + size_t num_vision_experts_{0}; + bool use_bias_{false}; +}; + +class Ernie45MoE : public infinicore::nn::Module { +public: + Ernie45MoE(std::shared_ptr model_config, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + +protected: + infinicore::Tensor forward_group(const infinicore::Tensor &hidden_states_2d, size_t group_idx) const; + + INFINICORE_NN_MODULE(Ernie45TopKRouter, gate); + INFINICORE_NN_MODULE(Ernie45Experts, experts); + INFINICORE_NN_MODULE(infinilm::layers::mlp::MLP, shared_experts); + infinicore::nn::Parameter e_score_correction_bias_; + size_t num_experts_{0}; + size_t num_vision_experts_{0}; + bool has_shared_experts_{false}; +}; + +} // namespace infinilm::models::ernie4_5_vl diff --git a/csrc/models/ernie4_5_vl/ernie4_5_vision.cpp b/csrc/models/ernie4_5_vl/ernie4_5_vision.cpp new file mode 100644 index 00000000..608dc201 --- /dev/null +++ b/csrc/models/ernie4_5_vl/ernie4_5_vision.cpp @@ -0,0 +1,308 @@ +#include "ernie4_5_vision.hpp" + +#include "infinicore/ops.hpp" +#include "infinicore/ops/cat.hpp" +#include "infinicore/ops/mha.hpp" + +#include +#include +#include +#include +#include + +namespace infinilm::models::ernie4_5_vl { +namespace { + +std::vector tensor_to_i64_vector(const infinicore::Tensor &tensor) { + auto cpu_tensor = tensor->to(infinicore::Device::cpu()); + std::vector values(cpu_tensor->numel()); + if (cpu_tensor->dtype() == infinicore::DataType::I64) { + const auto *ptr = reinterpret_cast(cpu_tensor->data()); + values.assign(ptr, ptr + cpu_tensor->numel()); + return values; + } + if (cpu_tensor->dtype() == infinicore::DataType::I32) { + const auto *ptr = reinterpret_cast(cpu_tensor->data()); + for (size_t i = 0; i < cpu_tensor->numel(); ++i) { + values[i] = static_cast(ptr[i]); + } + return values; + } + throw std::runtime_error("ERNIE 4.5 VL grid_thw must be int32 or int64"); +} + +infinicore::Tensor cat_or_single(std::vector tensors, int dim) { + if (tensors.empty()) { + throw std::runtime_error("ERNIE 4.5 VL internal cat received no tensors"); + } + if (tensors.size() == 1) { + return tensors[0]; + } + return infinicore::op::cat(std::move(tensors), dim); +} + +} // namespace + +Ernie45VisionLayerNorm::Ernie45VisionLayerNorm(size_t normalized_shape, + double eps, + const infinicore::DataType &dtype, + const infinicore::Device &device) + : eps_(eps) { + INFINICORE_NN_PARAMETER_INIT(weight, ({normalized_shape}, dtype, device)); + INFINICORE_NN_PARAMETER_INIT(bias, ({normalized_shape}, dtype, device)); +} + +infinicore::Tensor Ernie45VisionLayerNorm::forward(const infinicore::Tensor &hidden_states) const { + return infinicore::op::layer_norm(hidden_states, weight_, bias_, static_cast(eps_)); +} + +Ernie45VisionPatchEmbed::Ernie45VisionPatchEmbed(const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device) { + const size_t in_channels = config.value("in_channels", config.value("in_chans", 3)); + const size_t patch_size = config.value("patch_size", config.value("spatial_patch_size", 14)); + const size_t embed_dim = config.value("embed_dim", config.value("hidden_size", 1280)); + INFINICORE_NN_MODULE_INIT(proj, in_channels * patch_size * patch_size, embed_dim, false, dtype, device); +} + +infinicore::Tensor Ernie45VisionPatchEmbed::forward(const infinicore::Tensor &hidden_states) const { + auto hidden_states_mutable = hidden_states; + return proj_->forward(hidden_states_mutable); +} + +Ernie45VisionAttention::Ernie45VisionAttention(const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device) { + hidden_size_ = config.value("embed_dim", config.value("hidden_size", 1280)); + num_heads_ = config.value("num_heads", 16); + if (hidden_size_ % num_heads_ != 0) { + throw std::runtime_error("Ernie45VisionAttention: hidden_size must be divisible by num_heads"); + } + head_dim_ = hidden_size_ / num_heads_; + if (head_dim_ % 4 != 0) { + throw std::runtime_error("Ernie45VisionAttention: head_dim must be divisible by 4 for 2D RoPE"); + } + scale_ = 1.0f / std::sqrt(static_cast(head_dim_)); + const size_t axis_head_dim = head_dim_ / 2; + INFINICORE_NN_MODULE_INIT(rotary_emb, axis_head_dim, axis_head_dim, 8192, 10000.0, infinicore::nn::RoPE::Algo::GPT_NEOX, dtype, device); + INFINICORE_NN_MODULE_INIT(qkv, hidden_size_, hidden_size_ * 3, true, dtype, device); + INFINICORE_NN_MODULE_INIT(proj, hidden_size_, hidden_size_, true, dtype, device); +} + +infinicore::Tensor Ernie45VisionAttention::forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &row_position_ids, + const infinicore::Tensor &col_position_ids) const { + const size_t seq_len = hidden_states->size(0); + const size_t axis_head_dim = head_dim_ / 2; + const size_t axis_pair_dim = axis_head_dim / 2; + + auto hidden_mut = hidden_states; + auto qkv = qkv_->forward(hidden_mut)->view({seq_len, 3, num_heads_, head_dim_}); + auto q = qkv->narrow({{1, 0, 1}})->squeeze(1)->contiguous(); + auto k = qkv->narrow({{1, 1, 1}})->squeeze(1)->contiguous(); + auto v = qkv->narrow({{1, 2, 1}})->squeeze(1)->contiguous()->unsqueeze(0); + + auto apply_2d_rope = [&](const infinicore::Tensor &x) { + auto row = infinicore::Tensor::empty({seq_len, num_heads_, axis_head_dim}, x->dtype(), x->device()); + auto col = infinicore::Tensor::empty({seq_len, num_heads_, axis_head_dim}, x->dtype(), x->device()); + row->narrow({{2, 0, axis_pair_dim}})->copy_from(x->narrow({{2, 0, axis_pair_dim}})); + row->narrow({{2, axis_pair_dim, axis_pair_dim}})->copy_from(x->narrow({{2, axis_head_dim, axis_pair_dim}})); + col->narrow({{2, 0, axis_pair_dim}})->copy_from(x->narrow({{2, axis_pair_dim, axis_pair_dim}})); + col->narrow({{2, axis_pair_dim, axis_pair_dim}})->copy_from(x->narrow({{2, axis_head_dim + axis_pair_dim, axis_pair_dim}})); + + rotary_emb_->forward(row, row_position_ids, true); + rotary_emb_->forward(col, col_position_ids, true); + + x->narrow({{2, 0, axis_pair_dim}})->copy_from(row->narrow({{2, 0, axis_pair_dim}})); + x->narrow({{2, axis_head_dim, axis_pair_dim}})->copy_from(row->narrow({{2, axis_pair_dim, axis_pair_dim}})); + x->narrow({{2, axis_pair_dim, axis_pair_dim}})->copy_from(col->narrow({{2, 0, axis_pair_dim}})); + x->narrow({{2, axis_head_dim + axis_pair_dim, axis_pair_dim}})->copy_from(col->narrow({{2, axis_pair_dim, axis_pair_dim}})); + }; + apply_2d_rope(q); + apply_2d_rope(k); + + auto out = infinicore::op::mha(q->unsqueeze(0), k->unsqueeze(0), v, std::nullopt, scale_, false)->view({seq_len, hidden_size_}); + return proj_->forward(out); +} + +Ernie45VisionMLP::Ernie45VisionMLP(const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device) { + const size_t embed_dim = config.value("embed_dim", config.value("hidden_size", 1280)); + const size_t hidden_dim = static_cast(static_cast(embed_dim) * config.value("mlp_ratio", 4.0)); + INFINICORE_NN_MODULE_INIT(fc1, embed_dim, hidden_dim, true, dtype, device); + INFINICORE_NN_MODULE_INIT(fc2, hidden_dim, embed_dim, true, dtype, device); +} + +infinicore::Tensor Ernie45VisionMLP::forward(const infinicore::Tensor &hidden_states) const { + auto x = fc1_->forward(const_cast(hidden_states)); + x = infinicore::op::quick_gelu(x); + return fc2_->forward(x); +} + +Ernie45VisionBlock::Ernie45VisionBlock(const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device) { + const size_t embed_dim = config.value("embed_dim", config.value("hidden_size", 1280)); + INFINICORE_NN_MODULE_INIT(norm1, embed_dim, 1e-6, dtype, device); + INFINICORE_NN_MODULE_INIT(norm2, embed_dim, 1e-6, dtype, device); + INFINICORE_NN_MODULE_INIT(attn, config, dtype, device); + INFINICORE_NN_MODULE_INIT(mlp, config, dtype, device); +} + +infinicore::Tensor Ernie45VisionBlock::forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &row_position_ids, + const infinicore::Tensor &col_position_ids) const { + auto attn_out = attn_->forward(norm1_->forward(hidden_states), row_position_ids, col_position_ids); + auto hidden_after_attn = infinicore::op::add(hidden_states, attn_out); + auto mlp_out = mlp_->forward(norm2_->forward(hidden_after_attn)); + return infinicore::op::add(hidden_after_attn, mlp_out); +} + +Ernie45VisionModel::Ernie45VisionModel(const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device) { + const size_t depth = config.value("depth", 32); + const size_t embed_dim = config.value("embed_dim", config.value("hidden_size", 1280)); + spatial_merge_size_ = config.value("spatial_merge_size", 2); + INFINICORE_NN_MODULE_INIT(patch_embed, config, dtype, device); + blocks_.reserve(depth); + for (size_t i = 0; i < depth; ++i) { + blocks_.push_back(this->register_module("blocks." + std::to_string(i), config, dtype, device)); + } + INFINICORE_NN_MODULE_INIT(ln, embed_dim, 1e-6, dtype, device); +} + +infinicore::Tensor Ernie45VisionModel::build_rotary_position_ids(const infinicore::Tensor &grid_thw) const { + auto grid = tensor_to_i64_vector(grid_thw); + if (grid.size() % 3 != 0) { + throw std::runtime_error("Ernie45VisionModel: grid_thw must have shape [3] or [num_images, 3]"); + } + + size_t total_tokens = 0; + for (size_t i = 0; i < grid.size(); i += 3) { + total_tokens += static_cast(grid[i]) * static_cast(grid[i + 1]) * static_cast(grid[i + 2]); + } + + auto position_ids_cpu = infinicore::Tensor::empty({2, total_tokens}, infinicore::DataType::I64, infinicore::Device::cpu()); + auto *position_ids = reinterpret_cast(position_ids_cpu->data()); + size_t out_token = 0; + for (size_t i = 0; i < grid.size(); i += 3) { + const size_t grid_t = static_cast(grid[i]); + const size_t grid_h = static_cast(grid[i + 1]); + const size_t grid_w = static_cast(grid[i + 2]); + if (grid_h % spatial_merge_size_ != 0 || grid_w % spatial_merge_size_ != 0) { + throw std::runtime_error("Ernie45VisionModel: grid_h and grid_w must be divisible by spatial_merge_size"); + } + const size_t merged_h = grid_h / spatial_merge_size_; + const size_t merged_w = grid_w / spatial_merge_size_; + for (size_t t = 0; t < grid_t; ++t) { + (void)t; + for (size_t bh = 0; bh < merged_h; ++bh) { + for (size_t bw = 0; bw < merged_w; ++bw) { + for (size_t ih = 0; ih < spatial_merge_size_; ++ih) { + const size_t row = bh * spatial_merge_size_ + ih; + for (size_t iw = 0; iw < spatial_merge_size_; ++iw) { + const size_t col = bw * spatial_merge_size_ + iw; + position_ids[out_token] = static_cast(row); + position_ids[total_tokens + out_token] = static_cast(col); + ++out_token; + } + } + } + } + } + } + return position_ids_cpu->to(grid_thw->device()); +} + +infinicore::Tensor Ernie45VisionModel::forward(const infinicore::Tensor &pixel_values, + const infinicore::Tensor &grid_thw) const { + auto hidden_states = patch_embed_->forward(pixel_values); + auto position_ids = build_rotary_position_ids(grid_thw); + auto row_position_ids = position_ids->narrow({{0, 0, 1}})->view({position_ids->size(1)}); + auto col_position_ids = position_ids->narrow({{0, 1, 1}})->view({position_ids->size(1)}); + + for (auto &block : blocks_) { + hidden_states = block->forward(hidden_states, row_position_ids, col_position_ids); + } + return ln_->forward(hidden_states); +} + +Ernie45ResamplerModel::Ernie45ResamplerModel(const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device) { + const size_t pixel_hidden_size = config.value("pixel_hidden_size", 1280); + const size_t hidden_size = config.value("hidden_size", 2560); + spatial_conv_size_ = config.value("spatial_conv_size", 2); + temporal_conv_size_ = config.value("temporal_conv_size", 2); + const size_t spatial_dim = pixel_hidden_size * spatial_conv_size_ * spatial_conv_size_; + const size_t temporal_dim = spatial_dim * temporal_conv_size_; + use_temporal_conv_ = config.value("use_temporal_conv", true); + + spatial_linear_0_ = this->register_module("spatial_linear.0", spatial_dim, spatial_dim, true, dtype, device); + spatial_linear_2_ = this->register_module("spatial_linear.2", spatial_dim, spatial_dim, true, dtype, device); + spatial_linear_3_ = this->register_module("spatial_linear.3", spatial_dim, 1e-6, dtype, device); + if (use_temporal_conv_) { + temporal_linear_0_ = this->register_module("temporal_linear.0", temporal_dim, spatial_dim, true, dtype, device); + temporal_linear_2_ = this->register_module("temporal_linear.2", spatial_dim, spatial_dim, true, dtype, device); + temporal_linear_3_ = this->register_module("temporal_linear.3", spatial_dim, 1e-6, dtype, device); + } + INFINICORE_NN_MODULE_INIT(mlp, spatial_dim, hidden_size, true, dtype, device); + INFINICORE_NN_MODULE_INIT(after_norm, hidden_size, config.value("rms_norm_eps", 1e-5), dtype, device); +} + +infinicore::Tensor Ernie45ResamplerModel::forward(const infinicore::Tensor &image_features, + const infinicore::Tensor &grid_thw) const { + auto grid = tensor_to_i64_vector(grid_thw); + if (grid.size() != 3) { + throw std::runtime_error("Ernie45ResamplerModel: this fallback expects one image grid [3]"); + } + const size_t grid_t = static_cast(grid[0]); + const size_t grid_h = static_cast(grid[1]); + const size_t grid_w = static_cast(grid[2]); + const size_t spatial_group = spatial_conv_size_ * spatial_conv_size_; + if (grid_h % spatial_conv_size_ != 0 || grid_w % spatial_conv_size_ != 0) { + throw std::runtime_error("Ernie45ResamplerModel: grid_h/grid_w must be divisible by spatial_conv_size"); + } + if (image_features->size(0) != grid_t * grid_h * grid_w) { + throw std::runtime_error("Ernie45ResamplerModel: image feature length does not match grid_thw"); + } + + auto x = image_features->view({image_features->size(0) / spatial_group, image_features->size(1) * spatial_group}); + x = spatial_linear_0_->forward(x); + x = infinicore::op::gelu(x); + x = spatial_linear_2_->forward(x); + x = spatial_linear_3_->forward(x); + + if (use_temporal_conv_) { + if (temporal_conv_size_ != 2) { + throw std::runtime_error("Ernie45ResamplerModel: temporary temporal packing fallback only supports temporal_conv_size=2"); + } + const size_t spatial_tokens = (grid_h * grid_w) / spatial_group; + if (x->size(0) != grid_t * spatial_tokens) { + throw std::runtime_error("Ernie45ResamplerModel: spatial token length mismatch"); + } + + std::vector first_frames; + std::vector second_frames; + for (size_t t = 0; t < grid_t; t += 2) { + const size_t first = t; + const size_t second = (t + 1 < grid_t) ? (t + 1) : t; + first_frames.push_back(x->narrow({{0, first * spatial_tokens, spatial_tokens}})); + second_frames.push_back(x->narrow({{0, second * spatial_tokens, spatial_tokens}})); + } + auto x_first = cat_or_single(std::move(first_frames), 0); + auto x_second = cat_or_single(std::move(second_frames), 0); + x = infinicore::op::cat(std::vector{x_first, x_second}, 1); + x = temporal_linear_0_->forward(x); + x = infinicore::op::gelu(x); + x = temporal_linear_2_->forward(x); + x = temporal_linear_3_->forward(x); + } + + x = mlp_->forward(x); + return after_norm_->forward(x); +} + +} // namespace infinilm::models::ernie4_5_vl diff --git a/csrc/models/ernie4_5_vl/ernie4_5_vision.hpp b/csrc/models/ernie4_5_vl/ernie4_5_vision.hpp new file mode 100644 index 00000000..95134fbe --- /dev/null +++ b/csrc/models/ernie4_5_vl/ernie4_5_vision.hpp @@ -0,0 +1,132 @@ +#pragma once + +#include "../../layers/common_modules.hpp" +#include "infinicore/nn/layer_norm.hpp" +#include "infinicore/nn/module.hpp" +#include "infinicore/nn/rope.hpp" +#include "infinicore/tensor.hpp" +#include +#include + +namespace infinilm::models::ernie4_5_vl { +class Ernie45VisionLayerNorm : public infinicore::nn::Module { +public: + Ernie45VisionLayerNorm(size_t normalized_shape, + double eps, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + +private: + INFINICORE_NN_PARAMETER(weight); + INFINICORE_NN_PARAMETER(bias); + double eps_{1e-6}; +}; + +class Ernie45VisionPatchEmbed : public infinicore::nn::Module { +public: + Ernie45VisionPatchEmbed(const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + +private: + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, proj); +}; + +class Ernie45VisionAttention : public infinicore::nn::Module { +public: + Ernie45VisionAttention(const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &row_position_ids, + const infinicore::Tensor &col_position_ids) const; + +private: + size_t hidden_size_{1280}; + size_t num_heads_{16}; + size_t head_dim_{80}; + float scale_{1.0f}; + + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, qkv); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, proj); + INFINICORE_NN_MODULE(infinicore::nn::RoPE, rotary_emb); +}; + +class Ernie45VisionMLP : public infinicore::nn::Module { +public: + Ernie45VisionMLP(const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + +private: + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, fc1); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, fc2); +}; + +class Ernie45VisionBlock : public infinicore::nn::Module { +public: + Ernie45VisionBlock(const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &row_position_ids, + const infinicore::Tensor &col_position_ids) const; + +private: + INFINICORE_NN_MODULE(Ernie45VisionLayerNorm, norm1); + INFINICORE_NN_MODULE(Ernie45VisionLayerNorm, norm2); + INFINICORE_NN_MODULE(Ernie45VisionAttention, attn); + INFINICORE_NN_MODULE(Ernie45VisionMLP, mlp); +}; + +class Ernie45VisionModel : public infinicore::nn::Module { +public: + Ernie45VisionModel(const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &pixel_values, + const infinicore::Tensor &grid_thw) const; + +private: + infinicore::Tensor build_rotary_position_ids(const infinicore::Tensor &grid_thw) const; + + size_t spatial_merge_size_{2}; + + INFINICORE_NN_MODULE(Ernie45VisionPatchEmbed, patch_embed); + INFINICORE_NN_MODULE_VEC(Ernie45VisionBlock, blocks); + INFINICORE_NN_MODULE(Ernie45VisionLayerNorm, ln); +}; + +class Ernie45ResamplerModel : public infinicore::nn::Module { +public: + Ernie45ResamplerModel(const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &image_features, + const infinicore::Tensor &grid_thw) const; + +private: + std::shared_ptr spatial_linear_0_; + std::shared_ptr spatial_linear_2_; + std::shared_ptr spatial_linear_3_; + std::shared_ptr temporal_linear_0_; + std::shared_ptr temporal_linear_2_; + std::shared_ptr temporal_linear_3_; + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, mlp); + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, after_norm); + bool use_temporal_conv_{true}; + size_t spatial_conv_size_{2}; + size_t temporal_conv_size_{2}; +}; + +} // namespace infinilm::models::ernie4_5_vl diff --git a/examples/test_infer.py b/examples/test_infer.py index b0250042..dc4c72b2 100644 --- a/examples/test_infer.py +++ b/examples/test_infer.py @@ -6,6 +6,8 @@ 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], @@ -69,7 +71,9 @@ def test( for prompt in prompts ] if video_path is not None: - video_payload = decode_video_frames(video_path, video_num_frames) + 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}} diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index 8f84c9b9..c8946dc9 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -2,6 +2,7 @@ import glob import json import os +import re import time from typing import Dict, List, Optional, Union @@ -743,6 +744,137 @@ def _remap_qwen3_5(state_dict, config): return state_dict +def _remap_ernie4_5_moe_vl(state_dict, config=None): + """Apply ERNIE 4.5 VL load-time weight fixes. + + ERNIE checkpoints store router gate weights in fp32. They also store each + expert as separate gate/up/down projections; InfiniLM uses InfiniCore's + fused MoE op, so fuse the text expert weights before loading. + """ + target_dtype = torch.bfloat16 + hf_config = config or {} + for key in ("torch_dtype", "dtype"): + dtype_name = hf_config.get(key) + if dtype_name in ("float16", "float32", "bfloat16"): + target_dtype = { + "float16": torch.float16, + "float32": torch.float32, + "bfloat16": torch.bfloat16, + }[dtype_name] + break + + text_expert_count = hf_config.get("num_experts") + moe_num_experts = hf_config.get("moe_num_experts") + if isinstance(moe_num_experts, list) and len(moe_num_experts) > 0: + text_expert_count = int(moe_num_experts[0]) + elif text_expert_count is not None: + text_expert_count = int(text_expert_count) + + expert_re = re.compile( + r"^(?Pmodel\.layers\.\d+\.mlp\.experts)\." + r"(?P\d+)\." + r"(?Pgate_proj|up_proj|down_proj)\." + r"(?Pweight|bias)$" + ) + expert_parts = {} + remapped = {} + + for key, tensor in state_dict.items(): + match = expert_re.match(key) + if match is not None: + prefix = match.group("prefix") + expert = int(match.group("expert")) + proj = match.group("proj") + kind = match.group("kind") + expert_parts.setdefault(prefix, {}).setdefault(expert, {})[(proj, kind)] = ( + tensor + ) + continue + + if ( + key.endswith((".mlp.gate.weight", ".mlp.gate.weight_1")) + and tensor.is_floating_point() + ): + remapped[key] = tensor.to(dtype=target_dtype).contiguous() + else: + remapped[key] = tensor + + for prefix, experts in expert_parts.items(): + text_ids = sorted( + expert_id + for expert_id in experts + if text_expert_count is None or expert_id < text_expert_count + ) + vision_ids = [] + if text_expert_count is not None: + vision_ids = sorted( + expert_id for expert_id in experts if expert_id >= text_expert_count + ) + if not text_ids: + continue + + def fuse_expert_group(expert_ids): + w1_tensors = [] + w2_tensors = [] + b1_tensors = [] + b2_tensors = [] + has_all_bias = True + for expert_id in expert_ids: + parts = experts[expert_id] + gate = parts.get(("gate_proj", "weight")) + up = parts.get(("up_proj", "weight")) + down = parts.get(("down_proj", "weight")) + if gate is None or up is None or down is None: + raise KeyError( + f"Incomplete ERNIE MoE expert weights for {prefix}.{expert_id}" + ) + w1_tensors.append(torch.cat([gate, up], dim=0)) + w2_tensors.append(down) + + gate_bias = parts.get(("gate_proj", "bias")) + up_bias = parts.get(("up_proj", "bias")) + down_bias = parts.get(("down_proj", "bias")) + if gate_bias is None or up_bias is None or down_bias is None: + has_all_bias = False + else: + b1_tensors.append(torch.cat([gate_bias, up_bias], dim=0)) + b2_tensors.append(down_bias) + + fused = { + "w1": torch.stack(w1_tensors, dim=0) + .to(dtype=target_dtype) + .contiguous(), + "w2": torch.stack(w2_tensors, dim=0) + .to(dtype=target_dtype) + .contiguous(), + } + if has_all_bias: + fused["b1"] = ( + torch.stack(b1_tensors, dim=0).to(dtype=target_dtype).contiguous() + ) + fused["b2"] = ( + torch.stack(b2_tensors, dim=0).to(dtype=target_dtype).contiguous() + ) + return fused + + text_fused = fuse_expert_group(text_ids) + remapped[f"{prefix}.w1"] = text_fused["w1"] + remapped[f"{prefix}.w2"] = text_fused["w2"] + if "b1" in text_fused: + remapped[f"{prefix}.b1"] = text_fused["b1"] + remapped[f"{prefix}.b2"] = text_fused["b2"] + + if vision_ids: + vision_fused = fuse_expert_group(vision_ids) + remapped[f"{prefix}.w1_1"] = vision_fused["w1"] + remapped[f"{prefix}.w2_1"] = vision_fused["w2"] + if "b1" in vision_fused: + remapped[f"{prefix}.b1_1"] = vision_fused["b1"] + remapped[f"{prefix}.b2_1"] = vision_fused["b2"] + + return remapped + + _WEIGHT_REMAPPER = { "glm4": _remap_glm4, "chatglm": _remap_chatglm, @@ -751,4 +883,5 @@ def _remap_qwen3_5(state_dict, config): "mamba": _remap_mamba, "videonsa": _remap_videonsa, "qwen3_5": _remap_qwen3_5, + "ernie4_5_moe_vl": _remap_ernie4_5_moe_vl, } diff --git a/python/infinilm/processors/ernie4_5_vl_processor.py b/python/infinilm/processors/ernie4_5_vl_processor.py new file mode 100644 index 00000000..ccd98875 --- /dev/null +++ b/python/infinilm/processors/ernie4_5_vl_processor.py @@ -0,0 +1,448 @@ +import json +import os + +from transformers import AutoProcessor, AutoTokenizer +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 + +DEFAULT_VIDEO_NUM_FRAMES = 8 + + +@register_processor("ernie4_5_moe_vl") +class Ernie45VLProcessor(BasicLLMProcessor): + def __init__(self, model_dir_path: str): + self.pixel_values_dtype = None + self.im_patch_id = None + self.image_rescale_factor = 1.0 / 255.0 + self.image_mean = [0.48145466, 0.4578275, 0.40821073] + self.image_std = [0.26862954, 0.26130258, 0.27577711] + self.patch_size = 14 + + preprocessor_path = os.path.join(model_dir_path, "preprocessor_config.json") + if os.path.exists(preprocessor_path): + with open(preprocessor_path, "r") as f: + preprocessor_json = json.load(f) + self.image_rescale_factor = preprocessor_json.get( + "rescale_factor", self.image_rescale_factor + ) + self.image_mean = preprocessor_json.get("image_mean", self.image_mean) + self.image_std = preprocessor_json.get("image_std", self.image_std) + self.patch_size = int(preprocessor_json.get("patch_size", self.patch_size)) + + config_path = os.path.join(model_dir_path, "config.json") + if os.path.exists(config_path): + with open(config_path, "r") as f: + config_json = json.load(f) + self.im_patch_id = config_json.get("im_patch_id") + dtype_name = config_json.get("torch_dtype") or config_json.get("dtype") + if dtype_name is not None: + import torch + + self.pixel_values_dtype = getattr(torch, str(dtype_name)) + + try: + self.processor = AutoProcessor.from_pretrained( + model_dir_path, trust_remote_code=True + ) + self.tokenizer = self.processor.tokenizer + except Exception: + self.processor = None + self.tokenizer = AutoTokenizer.from_pretrained( + model_dir_path, trust_remote_code=True + ) + + if self.im_patch_id is None: + self.im_patch_id = self.tokenizer.convert_tokens_to_ids( + "<|IMAGE_PLACEHOLDER|>" + ) + self.im_patch_id = int(self.im_patch_id) + + def _normalize_videos(self, videos): + if not videos: + return videos + + normalized = [] + for video in videos: + if isinstance(video, tuple): + normalized.append(video) + continue + if isinstance(video, str): + from .videonsa_processor import decode_video_frames + + num_frames = int( + os.getenv( + "INFINILM_ERNIE_VIDEO_NUM_FRAMES", + os.getenv( + "INFINILM_VIDEONSA_VIDEO_NUM_FRAMES", + DEFAULT_VIDEO_NUM_FRAMES, + ), + ) + ) + video = decode_video_frames(video, num_frames) + if isinstance(video, list): + if not video: + continue + if hasattr(video[0], "convert"): + import numpy as np + + video = np.stack( + [np.array(frame.convert("RGB")) for frame in video], axis=0 + ) + normalized.append(video) + return normalized + + @override + def __call__( + self, + prompt, + images=None, + videos=None, + audios=None, + return_tensors: str = None, + **kwargs, + ) -> dict: + if not images and not videos and not audios: + return self.tokenizer( + prompt, return_tensors=return_tensors, add_special_tokens=False + ) + + if audios: + raise NotImplementedError("ERNIE 4.5 VL processor does not support audio") + if self.processor is None: + raise RuntimeError("ERNIE 4.5 VL multimodal processor is not available") + + videos = self._normalize_videos(videos) + processed = self.processor( + text=prompt, + images=images, + videos=videos, + return_tensors=return_tensors or "pt", + **kwargs, + ) + self._append_image_bound(processed) + return processed + + def _append_image_bound(self, processed_inputs: dict) -> None: + if "input_ids" not in processed_inputs: + return + + import torch + + input_ids = processed_inputs["input_ids"] + if isinstance(input_ids, torch.Tensor): + token_ids = ( + input_ids[0].tolist() if input_ids.ndim == 2 else input_ids.tolist() + ) + else: + token_ids = ( + input_ids[0] + if input_ids and isinstance(input_ids[0], list) + else input_ids + ) + + bounds = [] + i = 0 + while i < len(token_ids): + if int(token_ids[i]) != self.im_patch_id: + i += 1 + continue + start = i + while i < len(token_ids) and int(token_ids[i]) == self.im_patch_id: + i += 1 + bounds.append([start, i]) + + processed_inputs["image_bound"] = torch.tensor(bounds, dtype=torch.int64) + + @override + def apply_chat_template( + self, + conversation, + add_generation_prompt: bool = False, + tokenize: bool = True, + **kwargs, + ): + normalized_conversation = [] + for message in conversation: + content = message["content"] + if not isinstance(content, list): + normalized_conversation.append(message) + continue + + normalized_content = [] + for item in content: + item_type = item.get("type") + if item_type == "text": + normalized_content.append( + {"type": "text", "text": item.get("text", "")} + ) + elif item_type == "image_url": + normalized_content.append({"type": "image"}) + elif item_type == "video_url": + normalized_content.append({"type": "video"}) + else: + raise NotImplementedError( + f"Unsupported ERNIE 4.5 VL content type: {item_type}" + ) + + normalized_conversation.append( + {"role": message.get("role", "user"), "content": normalized_content} + ) + + return self.tokenizer.apply_chat_template( + conversation=normalized_conversation, + add_generation_prompt=add_generation_prompt, + tokenize=tokenize, + **kwargs, + ) + + @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: + if isinstance(scheduler_output, StaticSchedulerOutput): + model_inputs = self._build_model_input_from_static_scheduler_output( + scheduler_output, temperature, top_p, top_k + ) + self._append_ernie_position_ids(model_inputs, scheduler_output) + elif isinstance(scheduler_output, SchedulerOutput): + model_inputs = self._build_model_input_from_batch_scheduler_output( + scheduler_output, temperature, top_p, top_k + ) + self._append_ernie_mm_inputs(model_inputs, scheduler_output) + self._append_ernie_position_ids(model_inputs, scheduler_output) + else: + raise ValueError( + "scheduler_output must be an instance of SchedulerOutput or StaticSchedulerOutput" + ) + return model_inputs + + def _position_delta(self, req) -> int: + return int(getattr(req, "mrope_position_delta", 0)) + + def _update_position_delta(self, req) -> None: + import torch + + processed_inputs = req.processed_inputs + if processed_inputs is None or "position_ids" not in processed_inputs: + req.mrope_position_delta = 0 + return + + pos = processed_inputs["position_ids"] + pos = pos if isinstance(pos, torch.Tensor) else torch.as_tensor(pos) + if pos.ndim == 3 and pos.shape[0] == 1: + last_pos = int(pos[0, -1, 0].item()) + elif pos.ndim == 2: + last_pos = int(pos[-1, 0].item()) + else: + req.mrope_position_delta = 0 + return + req.mrope_position_delta = last_pos + 1 - len(req.get_input_tokens()) + + def _build_prefill_position_ids(self, req, num_cached: int, compute_len: int): + import torch + + processed_inputs = req.processed_inputs + if processed_inputs is None or "position_ids" not in processed_inputs: + pos = torch.arange(num_cached, num_cached + compute_len, dtype=torch.int64) + return pos + + pos = processed_inputs["position_ids"] + pos = pos if isinstance(pos, torch.Tensor) else torch.as_tensor(pos) + if pos.ndim == 3 and pos.shape[0] == 1: + return pos[:, num_cached : num_cached + compute_len, :].contiguous() + if pos.ndim == 2: + return pos[num_cached : num_cached + compute_len, :].contiguous() + raise RuntimeError("ERNIE 4.5 VL position_ids must have shape [1, seq, 3]") + + def _append_ernie_position_ids(self, model_inputs: dict, scheduler_output) -> None: + import infinicore + import torch + + position_chunks = [] + has_3d = False + for req in scheduler_output.scheduled_requests: + if scheduler_output.is_prefill: + self._update_position_delta(req) + num_cached = req.num_local_cached_tokens + compute_len = len(req.get_input_tokens()) - num_cached + pos = self._build_prefill_position_ids(req, num_cached, compute_len) + else: + current_position = ( + req.get_total_length() - 1 + self._position_delta(req) + ) + pos = torch.tensor([current_position], dtype=torch.int64) + has_3d = has_3d or pos.ndim == 3 + position_chunks.append(pos) + + if not position_chunks: + return + if has_3d: + normalized = [] + for pos in position_chunks: + if pos.ndim == 1: + pos = pos.view(1, -1, 1).expand(1, -1, 3) + elif pos.ndim == 2: + pos = pos.view(1, pos.shape[0], pos.shape[1]) + normalized.append(pos) + model_inputs["position_ids"] = infinicore.from_torch( + torch.cat(normalized, dim=1).contiguous() + ) + else: + model_inputs["position_ids"] = infinicore.from_torch( + torch.cat( + [pos.flatten() for pos in position_chunks], dim=0 + ).contiguous() + ) + + def _normalize_image_rows(self, image_rows): + import torch + + image_rows = image_rows.to(torch.float32) * float(self.image_rescale_factor) + repeat = self.patch_size * self.patch_size + mean = torch.tensor( + self.image_mean, dtype=torch.float32, device=image_rows.device + ) + std = torch.tensor( + self.image_std, dtype=torch.float32, device=image_rows.device + ) + mean = mean.repeat_interleave(repeat).view(1, -1) + std = std.repeat_interleave(repeat).view(1, -1) + return (image_rows - mean) / std + + def _append_ernie_mm_inputs( + self, model_inputs: dict, scheduler_output: SchedulerOutput + ) -> None: + import infinicore + import torch + + pixel_values = [] + image_bound = [] + grid_thw = [] + image_req_ids = [] + + for req_id, req in enumerate(scheduler_output.scheduled_requests): + processed_inputs = req.processed_inputs + if ( + not scheduler_output.is_prefill + or processed_inputs is None + or "images" not in processed_inputs + ): + continue + + bounds = processed_inputs.get("image_bound") + grids = processed_inputs.get("grid_thw") + if bounds is None or grids is None: + raise RuntimeError( + "ERNIE 4.5 VL multimodal input requires image_bound and grid_thw" + ) + + images = processed_inputs["images"] + images = ( + images if isinstance(images, torch.Tensor) else torch.as_tensor(images) + ) + grids = grids if isinstance(grids, torch.Tensor) else torch.as_tensor(grids) + bounds = ( + bounds if isinstance(bounds, torch.Tensor) else torch.as_tensor(bounds) + ) + if bounds.ndim == 3 and bounds.shape[0] == 1: + bounds = bounds.squeeze(0) + if grids.ndim == 1: + grids = grids.view(1, 3) + if bounds.ndim != 2 or bounds.shape[-1] != 2: + raise RuntimeError( + "ERNIE 4.5 VL image_bound must have shape [num_media, 2]" + ) + if grids.ndim != 2 or grids.shape[-1] != 3: + raise RuntimeError( + "ERNIE 4.5 VL grid_thw must have shape [num_media, 3]" + ) + if bounds.shape[0] != grids.shape[0]: + raise RuntimeError("ERNIE 4.5 VL media bounds and grids count mismatch") + + num_cached = req.num_local_cached_tokens + partial_cached = (bounds[:, 0] < num_cached) & (bounds[:, 1] > num_cached) + if partial_cached.any().item(): + raise RuntimeError( + "ERNIE 4.5 VL does not support partially cached multimodal spans" + ) + + row_offset = 0 + for image_idx in range(bounds.shape[0]): + grid = grids[image_idx].to(torch.int64) + rows = int(grid[0].item() * grid[1].item() * grid[2].item()) + image_rows = images[row_offset : row_offset + rows] + row_offset += rows + if bounds[image_idx, 1].item() <= num_cached: + continue + image_rows = self._normalize_image_rows(image_rows) + if self.pixel_values_dtype is not None: + image_rows = image_rows.to(self.pixel_values_dtype) + pixel_values.append(image_rows.contiguous()) + image_bound.append((bounds[image_idx] - num_cached).to(torch.int64)) + grid_thw.append(grid.contiguous()) + image_req_ids.append(req_id) + + if pixel_values: + model_inputs["pixel_values"] = [ + infinicore.from_torch(t) for t in pixel_values + ] + model_inputs["image_bound"] = [ + infinicore.from_torch(t) for t in image_bound + ] + model_inputs["tgt_sizes"] = [infinicore.from_torch(t) for t in grid_thw] + model_inputs["image_req_ids"] = image_req_ids + + @override + def get_mm_token_index_list( + self, prompt_token_ids, image_ids=None, video_ids=None, audio_ids=None, **kwargs + ): + mm_token_index_list = [] + image_ids = image_ids or [] + video_ids = video_ids or [] + if image_ids and video_ids: + raise NotImplementedError( + "ERNIE 4.5 VL cache mapping does not support mixed image and video inputs yet" + ) + media_ids = image_ids if image_ids else video_ids + + media_idx = 0 + i = 0 + while i < len(prompt_token_ids): + if int(prompt_token_ids[i]) != self.im_patch_id: + i += 1 + continue + + start = i + while ( + i < len(prompt_token_ids) + and int(prompt_token_ids[i]) == self.im_patch_id + ): + i += 1 + + if media_idx >= len(media_ids): + raise RuntimeError( + "The number of ERNIE 4.5 VL multimodal token spans exceeds media inputs" + ) + mm_token_index_list.append( + { + "start_index": start, + "end_index": i - 1, + "identifier": media_ids[media_idx], + } + ) + media_idx += 1 + + if media_idx != len(media_ids): + raise RuntimeError( + "The number of ERNIE 4.5 VL multimodal token spans does not match media inputs" + ) + return mm_token_index_list