From 9ba2f3a917b8df94ebac33428a74ac23c0a93728 Mon Sep 17 00:00:00 2001 From: CearX <56916338+CearX@users.noreply.github.com> Date: Sun, 12 Jul 2026 23:09:40 +0800 Subject: [PATCH] feat(ernie-vl): add ERNIE 4.5 multimodal inference --- csrc/engine/infer_engine.cpp | 3 + csrc/engine/rank_worker.hpp | 6 + .../attention/backends/attention_layer.cpp | 5 + .../ernie4_5_moe_vl/ernie4_5_attention.cpp | 268 ++++++++ .../ernie4_5_moe_vl/ernie4_5_attention.hpp | 52 ++ .../ernie4_5_decoder_layer.cpp | 53 ++ .../ernie4_5_decoder_layer.hpp | 39 ++ csrc/models/ernie4_5_moe_vl/ernie4_5_moe.cpp | 438 +++++++++++++ csrc/models/ernie4_5_moe_vl/ernie4_5_moe.hpp | 66 ++ .../ernie4_5_moe_vl_for_causal_lm.cpp | 222 +++++++ .../ernie4_5_moe_vl_for_causal_lm.hpp | 62 ++ .../ernie4_5_moe_vl/ernie4_5_resampler.cpp | 113 ++++ .../ernie4_5_moe_vl/ernie4_5_resampler.hpp | 44 ++ .../ernie4_5_moe_vl/ernie4_5_vision.cpp | 374 +++++++++++ .../ernie4_5_moe_vl/ernie4_5_vision.hpp | 111 ++++ csrc/models/infinilm_model.hpp | 6 + csrc/pybind11/engine/engine.hpp | 12 + examples/ernie_vl_correctness.py | 590 ++++++++++++++++++ examples/ernie_vl_llm_smoke.py | 161 +++++ examples/ernie_vl_mmmu_smoke.py | 528 ++++++++++++++++ examples/ernie_vl_verification.md | 164 +++++ python/infinilm/infer_engine.py | 141 ++++- .../infinilm/llm/model_runner/model_runner.py | 14 +- python/infinilm/modeling_utils.py | 27 + .../processors/ernie4_5_moe_vl_processor.py | 385 ++++++++++++ test/bench/backends/transformers.py | 101 ++- tools/compare_mmmu_results.py | 251 ++++++++ tools/mmmu_run_status.py | 101 +++ 28 files changed, 4322 insertions(+), 15 deletions(-) create mode 100644 csrc/models/ernie4_5_moe_vl/ernie4_5_attention.cpp create mode 100644 csrc/models/ernie4_5_moe_vl/ernie4_5_attention.hpp create mode 100644 csrc/models/ernie4_5_moe_vl/ernie4_5_decoder_layer.cpp create mode 100644 csrc/models/ernie4_5_moe_vl/ernie4_5_decoder_layer.hpp create mode 100644 csrc/models/ernie4_5_moe_vl/ernie4_5_moe.cpp create mode 100644 csrc/models/ernie4_5_moe_vl/ernie4_5_moe.hpp create mode 100644 csrc/models/ernie4_5_moe_vl/ernie4_5_moe_vl_for_causal_lm.cpp create mode 100644 csrc/models/ernie4_5_moe_vl/ernie4_5_moe_vl_for_causal_lm.hpp create mode 100644 csrc/models/ernie4_5_moe_vl/ernie4_5_resampler.cpp create mode 100644 csrc/models/ernie4_5_moe_vl/ernie4_5_resampler.hpp create mode 100644 csrc/models/ernie4_5_moe_vl/ernie4_5_vision.cpp create mode 100644 csrc/models/ernie4_5_moe_vl/ernie4_5_vision.hpp create mode 100644 examples/ernie_vl_correctness.py create mode 100644 examples/ernie_vl_llm_smoke.py create mode 100644 examples/ernie_vl_mmmu_smoke.py create mode 100644 examples/ernie_vl_verification.md create mode 100644 python/infinilm/processors/ernie4_5_moe_vl_processor.py create mode 100644 tools/compare_mmmu_results.py create mode 100644 tools/mmmu_run_status.py diff --git a/csrc/engine/infer_engine.cpp b/csrc/engine/infer_engine.cpp index a5221eb37..f7df16f2f 100644 --- a/csrc/engine/infer_engine.cpp +++ b/csrc/engine/infer_engine.cpp @@ -167,6 +167,7 @@ InferEngine::Input::to_model_input(infinicore::Device device) const { infinilm::InfinilmModel::Input input = { to_device(input_ids), // @todo: on device in the future to_device(position_ids), + to_device(token_type_ids), to_device(past_sequence_lengths), // @todo: on device in the future to_device(total_sequence_lengths), to_device(input_offsets), @@ -179,6 +180,8 @@ InferEngine::Input::to_model_input(infinicore::Device device) const { to_device_vec(image_bound), to_device_vec(tgt_sizes), to_device_vec(image_grid_thw), + to_device_vec(grid_thw), + to_device_vec(image_type_ids), image_req_ids, visual_token_ranges, to_device(target_hidden_states)}; diff --git a/csrc/engine/rank_worker.hpp b/csrc/engine/rank_worker.hpp index d396ef6f1..c6d9d0a99 100644 --- a/csrc/engine/rank_worker.hpp +++ b/csrc/engine/rank_worker.hpp @@ -40,6 +40,8 @@ class RankWorker { std::optional input_ids; /// Position IDs tensor of shape `[batch, seq_len]` or `[seq_len]`. std::optional position_ids; + /// Token modality IDs. ERNIE-VL uses 0 for text and non-zero for vision. + std::optional token_type_ids; /// Past Lengths of cached sequence for each request, of shape `[num_requests]`. std::optional past_sequence_lengths; /// ToTal Lengths for each request sequence, of shape `[num_requests]`. @@ -64,6 +66,10 @@ class RankWorker { std::optional> tgt_sizes; /// Qwen-style image grids. Vector of tensors shape: [3] with temporal, height, width. std::optional> image_grid_thw; + /// ERNIE-VL image/video grids. Each tensor has shape [num_items, 3]. + std::optional> grid_thw; + /// ERNIE-VL item type IDs, where 0=image and 1=video. + std::optional> image_type_ids; /// req_id for each pixel_values among a batch std::optional> image_req_ids; /// Flattened [start, end) visual token ranges in the packed language sequence. diff --git a/csrc/layers/attention/backends/attention_layer.cpp b/csrc/layers/attention/backends/attention_layer.cpp index fcaefa292..f08edb416 100644 --- a/csrc/layers/attention/backends/attention_layer.cpp +++ b/csrc/layers/attention/backends/attention_layer.cpp @@ -30,6 +30,11 @@ infinicore::Tensor AttentionLayer::forward(infinicore::Tensor &query, infinicore::Tensor &value) const { auto &forward_context = infinilm::global_state::get_forward_context(); auto &attn_metadata = forward_context.attn_metadata; + if (forward_context.kv_cache_vec.size() <= layer_idx_) { + throw std::runtime_error( + "AttentionLayer::forward requires an initialized KV cache; " + "pass a cache_config to InferEngine or call reset_cache before generation."); + } auto &kv_cache = forward_context.kv_cache_vec[layer_idx_]; return std::visit( diff --git a/csrc/models/ernie4_5_moe_vl/ernie4_5_attention.cpp b/csrc/models/ernie4_5_moe_vl/ernie4_5_attention.cpp new file mode 100644 index 000000000..768ba499d --- /dev/null +++ b/csrc/models/ernie4_5_moe_vl/ernie4_5_attention.cpp @@ -0,0 +1,268 @@ +#include "ernie4_5_attention.hpp" + +#include "../../global_state/global_state.hpp" +#include "../../layers/attention/attention.hpp" +#include "../../utils.hpp" + +#include +#include +#include +#include +#include + +namespace infinilm::models::ernie4_5_moe_vl { +namespace { + +uint16_t fp32_to_bf16(float value) { + uint32_t bits; + std::memcpy(&bits, &value, sizeof(bits)); + const uint32_t lsb = (bits >> 16) & 1U; + bits += 0x7FFFU + lsb; + return static_cast(bits >> 16); +} + +float bf16_to_fp32(uint16_t value) { + uint32_t bits = static_cast(value) << 16; + float out; + std::memcpy(&out, &bits, sizeof(out)); + return out; +} + +float read_float(const void *src, size_t index, infinicore::DataType dtype) { + if (dtype == infinicore::DataType::F32) { + return reinterpret_cast(src)[index]; + } + if (dtype == infinicore::DataType::BF16) { + return bf16_to_fp32(reinterpret_cast(src)[index]); + } + throw std::runtime_error("Ernie4_5Attention: only float32 and bfloat16 3D RoPE are supported"); +} + +void write_float(void *dst, size_t index, infinicore::DataType dtype, float value) { + if (dtype == infinicore::DataType::F32) { + reinterpret_cast(dst)[index] = value; + return; + } + if (dtype == infinicore::DataType::BF16) { + reinterpret_cast(dst)[index] = fp32_to_bf16(value); + return; + } + throw std::runtime_error("Ernie4_5Attention: only float32 and bfloat16 3D RoPE are supported"); +} + +} // namespace + +Ernie4_5Attention::Ernie4_5Attention(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"); + use_rope_3d_ = model_config->get_or("rope_3d", false); + freq_allocation_ = model_config->get_or("freq_allocation", 20); + rope_theta_ = model_config->get_or("rope_theta", 10000.0); + compression_ratio_ = model_config->get_or("compression_ratio", 1.0); + + const auto &dtype{model_config->get_dtype()}; + size_t total_num_heads = model_config->get("num_attention_heads"); + size_t total_num_kv_heads = model_config->get("num_key_value_heads"); + bool use_bias = model_config->get_or("attention_bias", false); + bool use_output_bias = model_config->get_or("attention_output_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(); + int tp_rank = infinilm::global_state::get_tensor_model_parallel_rank(); + int tp_size = infinilm::global_state::get_tensor_model_parallel_world_size(); + + if ((total_num_kv_heads < static_cast(tp_size)) || (total_num_kv_heads % static_cast(tp_size) != 0)) { + throw std::runtime_error("Ernie4_5Attention: 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 &n, infinicore::nn::Parameter p) { this->register_parameter(n, std::move(p)); }; + 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); + + rotary_emb_ = infinilm::layers::rotary_embedding::get_rope(model_config, device); + + 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_); + + infinilm::layers::attention::init_kv_cache_quant_params(register_fn, device, kv_cache_k_scale_, kv_cache_v_scale_); +} + +infinicore::Tensor Ernie4_5Attention::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 Ernie4_5Attention::position_ids_for_rope_(const infinicore::Tensor &position_ids) const { + auto pos_shape = position_ids->shape(); + if (pos_shape.size() == 3) { + auto text_axis = position_ids->narrow({{2, 0, 1}})->contiguous()->view({pos_shape[0], pos_shape[1]}); + return text_axis->narrow({{0, 0, 1}})->contiguous()->view({pos_shape[1]}); + } + if (pos_shape.size() == 2) { + return position_ids->narrow({{0, 0, 1}})->contiguous()->view({pos_shape[1]}); + } + if (pos_shape.size() == 1) { + return position_ids->contiguous(); + } + throw std::runtime_error("Ernie4_5Attention: unexpected position_ids shape"); +} + +bool Ernie4_5Attention::should_use_rope_3d_(const infinicore::Tensor &position_ids) const { + return use_rope_3d_ && position_ids->shape().size() == 3; +} + +infinicore::Tensor Ernie4_5Attention::apply_rope_3d_(const infinicore::Tensor &states, + const infinicore::Tensor &position_ids) const { + const auto state_shape = states->shape(); + if (state_shape.size() != 3 && state_shape.size() != 4) { + throw std::runtime_error("Ernie4_5Attention: 3D RoPE expects [S,H,D] or [B,S,H,D] states"); + } + ASSERT(head_dim_ % 2 == 0); + const size_t half_dim = head_dim_ / 2; + ASSERT(freq_allocation_ <= half_dim); + + const bool has_batch_dim = state_shape.size() == 4; + const size_t batch = has_batch_dim ? state_shape[0] : 1; + const size_t seq_len = has_batch_dim ? state_shape[1] : state_shape[0]; + const size_t num_heads = has_batch_dim ? state_shape[2] : state_shape[1]; + ASSERT_EQ(has_batch_dim ? state_shape[3] : state_shape[2], head_dim_); + + auto pos_cpu = position_ids->to(infinicore::Device::cpu()); + const auto pos_shape = pos_cpu->shape(); + ASSERT(pos_shape.size() == 3); + ASSERT_EQ(pos_shape[2], 3); + ASSERT(pos_shape[0] == batch || pos_shape[0] == 1); + ASSERT(pos_shape[1] >= seq_len); + + auto states_cpu = states->contiguous()->to(infinicore::Device::cpu()); + auto out_shape = has_batch_dim + ? std::vector{batch, num_heads, seq_len, head_dim_} + : states->shape(); + auto out_cpu = infinicore::Tensor::empty(out_shape, states->dtype(), infinicore::Device::cpu()); + const auto *pos = reinterpret_cast(pos_cpu->data()); + const void *src = states_cpu->data(); + void *dst = out_cpu->data(); + + std::vector inv_freq(half_dim); + for (size_t j = 0; j < half_dim; ++j) { + inv_freq[j] = 1.0f / std::pow(static_cast(rope_theta_), static_cast(2 * j) / static_cast(head_dim_)); + } + + const size_t hw_freq_end = half_dim - freq_allocation_; + for (size_t b = 0; b < batch; ++b) { + const size_t pos_batch = pos_shape[0] == 1 ? 0 : b; + for (size_t s = 0; s < seq_len; ++s) { + const int64_t pos_t = pos[(pos_batch * pos_shape[1] + s) * 3]; + const int64_t pos_h = pos[(pos_batch * pos_shape[1] + s) * 3 + 1]; + const int64_t pos_w = pos[(pos_batch * pos_shape[1] + s) * 3 + 2]; + for (size_t h = 0; h < num_heads; ++h) { + const size_t src_base = has_batch_dim + ? ((b * seq_len + s) * num_heads + h) * head_dim_ + : (s * num_heads + h) * head_dim_; + const size_t dst_base = has_batch_dim + ? ((b * num_heads + h) * seq_len + s) * head_dim_ + : src_base; + for (size_t j = 0; j < half_dim; ++j) { + int64_t position = pos_t; + if (j < hw_freq_end) { + position = (j % 2 == 0) ? pos_h : pos_w; + } + const float angle = (static_cast(position) / static_cast(compression_ratio_)) * inv_freq[j]; + const float sn = std::sin(angle); + const float cs = std::cos(angle); + const size_t src_even_idx = src_base + 2 * j; + const size_t src_odd_idx = src_even_idx + 1; + const size_t dst_even_idx = dst_base + 2 * j; + const size_t dst_odd_idx = dst_even_idx + 1; + const float x0 = read_float(src, src_even_idx, states->dtype()); + const float x1 = read_float(src, src_odd_idx, states->dtype()); + write_float(dst, dst_even_idx, states->dtype(), x0 * cs - x1 * sn); + write_float(dst, dst_odd_idx, states->dtype(), x1 * cs + x0 * sn); + } + } + } + } + + auto out = out_cpu->to(states->device()); + if (has_batch_dim) { + return out->permute({0, 2, 1, 3}); + } + return out; +} + +infinicore::Tensor Ernie4_5Attention::forward_static_(const infinicore::Tensor &position_ids, + const infinicore::Tensor &hidden_states) const { + auto hidden_states_mutable = hidden_states; + auto shape = hidden_states->shape(); + size_t batch_size = shape[0]; + size_t seq_len = shape[1]; + + auto [q, k, v] = qkv_proj_->forward_split(hidden_states_mutable); + + auto q_reshaped = q->view({batch_size, seq_len, num_attention_heads_, head_dim_}); + auto k_reshaped = k->view({batch_size, seq_len, num_key_value_heads_, head_dim_}); + auto v_reshaped = v->view({batch_size, seq_len, num_key_value_heads_, head_dim_}); + + infinicore::Tensor q_rope; + if (should_use_rope_3d_(position_ids)) { + q_rope = apply_rope_3d_(q_reshaped, position_ids); + k_reshaped = apply_rope_3d_(k_reshaped, position_ids); + } else { + auto pos_ids_for_rope = position_ids_for_rope_(position_ids); + q_rope = infinicore::Tensor::empty( + {batch_size, num_attention_heads_, seq_len, head_dim_}, + q_reshaped->dtype(), q_reshaped->device()) + ->permute({0, 2, 1, 3}); + rotary_emb_->forward(q_rope, q_reshaped, pos_ids_for_rope); + rotary_emb_->forward(k_reshaped, pos_ids_for_rope, true); + } + + auto attn_output = attn_->forward(q_rope, k_reshaped, v_reshaped); + return o_proj_->forward(attn_output); +} + +infinicore::Tensor Ernie4_5Attention::forward_paged_(const infinicore::Tensor &position_ids, + const infinicore::Tensor &hidden_states) const { + auto hidden_states_mutable = hidden_states; + auto shape = hidden_states->shape(); + size_t batch_size = shape[0]; + size_t seq_len = shape[1]; + ASSERT_EQ(batch_size, 1); + + auto [q, k, v] = qkv_proj_->forward_split(hidden_states_mutable); + + auto q_reshaped = q->view({seq_len, num_attention_heads_, head_dim_}); + auto k_reshaped = k->view({seq_len, num_key_value_heads_, head_dim_}); + auto v_reshaped = v->view({seq_len, num_key_value_heads_, head_dim_}); + + if (should_use_rope_3d_(position_ids)) { + q_reshaped = apply_rope_3d_(q_reshaped, position_ids); + k_reshaped = apply_rope_3d_(k_reshaped, position_ids); + } else { + auto pos_ids_for_rope = position_ids_for_rope_(position_ids); + rotary_emb_->forward(q_reshaped, pos_ids_for_rope, true); + rotary_emb_->forward(k_reshaped, pos_ids_for_rope, true); + } + + auto attn_output = attn_->forward(q_reshaped, k_reshaped, v_reshaped); + return o_proj_->forward(attn_output); +} + +} // namespace infinilm::models::ernie4_5_moe_vl diff --git a/csrc/models/ernie4_5_moe_vl/ernie4_5_attention.hpp b/csrc/models/ernie4_5_moe_vl/ernie4_5_attention.hpp new file mode 100644 index 000000000..11fca2e12 --- /dev/null +++ b/csrc/models/ernie4_5_moe_vl/ernie4_5_attention.hpp @@ -0,0 +1,52 @@ +#pragma once + +#include "../../layers/attention/attention.hpp" +#include "../../layers/linear/fused_linear.hpp" +#include "../../layers/rotary_embedding/rotary_embedding.hpp" +#include "infinicore/nn/module.hpp" +#include "infinicore/tensor.hpp" + +#include + +namespace infinilm::models::ernie4_5_moe_vl { + +class Ernie4_5Attention : public infinicore::nn::Module { +public: + Ernie4_5Attention(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; + +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; + infinicore::Tensor position_ids_for_rope_(const infinicore::Tensor &positions) const; + bool should_use_rope_3d_(const infinicore::Tensor &positions) const; + infinicore::Tensor apply_rope_3d_(const infinicore::Tensor &states, + const infinicore::Tensor &positions) const; + + 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_{0}; + size_t num_attention_heads_{0}; + size_t num_key_value_heads_{0}; + size_t hidden_size_{0}; + size_t head_dim_{0}; + size_t freq_allocation_{20}; + double rope_theta_{10000.0}; + double compression_ratio_{1.0}; + bool use_rope_3d_{false}; + + INFINICORE_NN_PARAMETER(kv_cache_k_scale); + INFINICORE_NN_PARAMETER(kv_cache_v_scale); +}; + +} // namespace infinilm::models::ernie4_5_moe_vl diff --git a/csrc/models/ernie4_5_moe_vl/ernie4_5_decoder_layer.cpp b/csrc/models/ernie4_5_moe_vl/ernie4_5_decoder_layer.cpp new file mode 100644 index 000000000..1cff2f2d1 --- /dev/null +++ b/csrc/models/ernie4_5_moe_vl/ernie4_5_decoder_layer.cpp @@ -0,0 +1,53 @@ +#include "ernie4_5_decoder_layer.hpp" + +#include "infinicore/ops.hpp" + +namespace infinilm::models::ernie4_5_moe_vl { + +Ernie4_5DecoderLayer::Ernie4_5DecoderLayer(std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device) + : layer_idx_(layer_idx) { + const auto &dtype{model_config->get_dtype()}; + size_t hidden_size = model_config->get("hidden_size"); + 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, device); + + if (layer_idx == 0) { + dense_mlp_ = this->register_module("mlp", model_config, device); + } else { + moe_mlp_ = this->register_module("mlp", model_config, device); + } +} + +std::tuple +Ernie4_5DecoderLayer::forward(const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + const infinicore::Tensor &token_type_ids) { + 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 = dense_mlp_ ? dense_mlp_->forward(hidden_states) : moe_mlp_->forward(hidden_states, token_type_ids); + return std::make_tuple(hidden_states, residual); +} + +infinicore::Tensor Ernie4_5DecoderLayer::forward(const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + const infinicore::Tensor &token_type_ids) { + 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 = dense_mlp_ ? dense_mlp_->forward(hidden_states) : moe_mlp_->forward(hidden_states, token_type_ids); + hidden_states = infinicore::op::add(residual, hidden_states); + return hidden_states; +} + +} // namespace infinilm::models::ernie4_5_moe_vl diff --git a/csrc/models/ernie4_5_moe_vl/ernie4_5_decoder_layer.hpp b/csrc/models/ernie4_5_moe_vl/ernie4_5_decoder_layer.hpp new file mode 100644 index 000000000..19f88b1b7 --- /dev/null +++ b/csrc/models/ernie4_5_moe_vl/ernie4_5_decoder_layer.hpp @@ -0,0 +1,39 @@ +#pragma once + +#include "../../layers/mlp/mlp.hpp" +#include "ernie4_5_attention.hpp" +#include "ernie4_5_moe.hpp" +#include "infinicore/nn/module.hpp" +#include "infinicore/nn/rmsnorm.hpp" +#include "infinicore/tensor.hpp" + +#include + +namespace infinilm::models::ernie4_5_moe_vl { + +class Ernie4_5DecoderLayer : public infinicore::nn::Module { +public: + Ernie4_5DecoderLayer(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 &token_type_ids = infinicore::Tensor()); + + infinicore::Tensor forward(const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + const infinicore::Tensor &token_type_ids = infinicore::Tensor()); + +private: + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, input_layernorm); + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, post_attention_layernorm); + INFINICORE_NN_MODULE(Ernie4_5Attention, self_attn); + + std::shared_ptr dense_mlp_; + std::shared_ptr moe_mlp_; + size_t layer_idx_{0}; +}; + +} // namespace infinilm::models::ernie4_5_moe_vl diff --git a/csrc/models/ernie4_5_moe_vl/ernie4_5_moe.cpp b/csrc/models/ernie4_5_moe_vl/ernie4_5_moe.cpp new file mode 100644 index 000000000..c048d5473 --- /dev/null +++ b/csrc/models/ernie4_5_moe_vl/ernie4_5_moe.cpp @@ -0,0 +1,438 @@ +#include "ernie4_5_moe.hpp" + +#include "../../config/model_config.hpp" +#include "../../utils.hpp" +#include "infinicore/ops.hpp" + +#include +#include +#include +#include +#include +#include +#include + +namespace infinilm::models::ernie4_5_moe_vl { +namespace { + +size_t json_first_size(const nlohmann::json &value, size_t fallback) { + if (value.is_array() && !value.empty()) { + return value.at(0).get(); + } + if (value.is_number_unsigned() || value.is_number_integer()) { + return value.get(); + } + return fallback; +} + +size_t json_size_at(const nlohmann::json &value, size_t index, size_t fallback) { + if (value.is_array() && index < value.size()) { + return value.at(index).get(); + } + if (index == 0 && (value.is_number_unsigned() || value.is_number_integer())) { + return value.get(); + } + return fallback; +} + +std::shared_ptr +clone_with_moe_intermediate(const std::shared_ptr &model_config, + size_t moe_intermediate_size) { + auto json = model_config->get_config_json(); + json["moe_intermediate_size"] = moe_intermediate_size; + return std::make_shared(json); +} + +size_t tensor_offset_1d(const infinicore::Tensor &tensor, size_t index) { + ASSERT(tensor->ndim() == 1); + const auto stride = tensor->stride(0); + ASSERT(stride >= 0); + return index * static_cast(stride); +} + +size_t tensor_offset_2d(const infinicore::Tensor &tensor, size_t row, size_t col) { + ASSERT(tensor->ndim() == 2); + const auto row_stride = tensor->stride(0); + const auto col_stride = tensor->stride(1); + ASSERT(row_stride >= 0 && col_stride >= 0); + return row * static_cast(row_stride) + col * static_cast(col_stride); +} + +float read_as_f32_offset(const infinicore::Tensor &tensor, size_t offset) { + const auto dtype = tensor->dtype(); + const auto *data = tensor->data(); + if (dtype == infinicore::DataType::F32) { + return reinterpret_cast(data)[offset]; + } + if (dtype == infinicore::DataType::F16) { + return f16_to_f32(reinterpret_cast(data)[offset]); + } + if (dtype == infinicore::DataType::BF16) { + return bf16_to_f32(reinterpret_cast(data)[offset]); + } + ASSERT(false); + return 0.0f; +} + +float read_as_f32_1d(const infinicore::Tensor &tensor, size_t index) { + return read_as_f32_offset(tensor, tensor_offset_1d(tensor, index)); +} + +float read_as_f32_2d(const infinicore::Tensor &tensor, size_t row, size_t col) { + return read_as_f32_offset(tensor, tensor_offset_2d(tensor, row, col)); +} + +int64_t read_as_i64_offset(const infinicore::Tensor &tensor, size_t offset) { + const auto dtype = tensor->dtype(); + const auto *data = tensor->data(); + if (dtype == infinicore::DataType::I64) { + return reinterpret_cast(data)[offset]; + } + if (dtype == infinicore::DataType::I32) { + return reinterpret_cast(data)[offset]; + } + ASSERT(false); + return 0; +} + +int64_t read_as_i64_1d(const infinicore::Tensor &tensor, size_t index) { + return read_as_i64_offset(tensor, tensor_offset_1d(tensor, index)); +} + +int64_t read_as_i64_2d(const infinicore::Tensor &tensor, size_t row, size_t col) { + return read_as_i64_offset(tensor, tensor_offset_2d(tensor, row, col)); +} + +} // namespace + +Ernie4_5MoeTopKRouter::Ernie4_5MoeTopKRouter(std::shared_ptr model_config, + const infinicore::Device &device) { + size_t hidden_size = model_config->get("hidden_size"); + const auto &json = model_config->get_config_json(); + const auto experts_json = json.value("moe_num_experts", nlohmann::json(64)); + num_experts_ = json_first_size(experts_json, 64); + num_expert_groups_ = experts_json.is_array() ? experts_json.size() : 1; + num_experts_per_tok_ = model_config->get_or("moe_k", 6); + norm_topk_prob_ = model_config->get_or("moe_norm_gate_logits", true); + use_correction_bias_ = model_config->get_or("moe_use_aux_free", false); + ASSERT((num_experts_ > 0) && (num_experts_per_tok_ > 0) && (num_experts_per_tok_ <= num_experts_)); + ASSERT(num_expert_groups_ >= 1); + ASSERT(num_expert_groups_ <= 2); + if (num_expert_groups_ > 1) { + ASSERT_EQ(json_size_at(experts_json, 1, num_experts_), num_experts_); + } + + const auto &dtype{model_config->get_dtype()}; + const bool use_gpu_router = std::getenv("INFINILM_ERNIE_GPU_ROUTER") != nullptr; + const auto gate_dtype = use_gpu_router ? dtype : infinicore::DataType::F32; + INFINICORE_NN_PARAMETER_INIT(weight, ({num_experts_, hidden_size}, gate_dtype, device)); + if (num_expert_groups_ > 1) { + INFINICORE_NN_PARAMETER_INIT(weight_1, ({num_experts_, hidden_size}, gate_dtype, device)); + } +} + +std::tuple +Ernie4_5MoeTopKRouter::forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &score_correction_bias, + const std::vector *token_groups) const { + ASSERT(hidden_states->ndim() == 2); + + size_t ntoken = hidden_states->shape()[0]; + size_t hidden_size = hidden_states->shape()[1]; + auto correction_bias_cpu = use_correction_bias_ + ? score_correction_bias->to(infinicore::Device::Type::CPU) + : infinicore::Tensor(); + if (use_correction_bias_) { + ASSERT(correction_bias_cpu); + ASSERT(correction_bias_cpu->numel() >= num_experts_); + } + if (token_groups != nullptr) { + ASSERT_EQ(token_groups->size(), ntoken); + } + + auto router_scores_cpu = infinicore::Tensor::empty({ntoken, num_experts_per_tok_}, infinicore::DataType::F32, infinicore::Device::Type::CPU); + auto router_indices_cpu = infinicore::Tensor::empty({ntoken, num_experts_per_tok_}, infinicore::DataType::I32, infinicore::Device::Type::CPU); + auto *router_scores_ptr = reinterpret_cast(router_scores_cpu->data()); + auto *router_indices_ptr = reinterpret_cast(router_indices_cpu->data()); + + std::vector logits(num_experts_); + std::vector probs(num_experts_); + std::vector expert_order(num_experts_); + std::iota(expert_order.begin(), expert_order.end(), 0); + + infinicore::Tensor router_logits_cpu; + infinicore::Tensor hidden_states_cpu; + infinicore::Tensor weight_cpu; + infinicore::Tensor weight_1_cpu; + const bool use_gpu_gate_linear = token_groups == nullptr && std::getenv("INFINILM_ERNIE_GPU_ROUTER") != nullptr; + if (use_gpu_gate_linear) { + ASSERT(weight_->ndim() == 2); + ASSERT(weight_->shape()[0] == num_experts_); + ASSERT(weight_->shape()[1] == hidden_size); + auto router_logits = infinicore::op::linear(hidden_states, weight_, std::nullopt, 1.0f); + router_logits_cpu = router_logits->to(infinicore::Device::Type::CPU); + ASSERT(router_logits_cpu->ndim() == 2); + ASSERT(router_logits_cpu->shape()[0] == ntoken); + ASSERT(router_logits_cpu->shape()[1] == num_experts_); + } else { + hidden_states_cpu = hidden_states->to(infinicore::Device::Type::CPU); + weight_cpu = weight_->to(infinicore::Device::Type::CPU); + weight_1_cpu = num_expert_groups_ > 1 ? weight_1_->to(infinicore::Device::Type::CPU) : infinicore::Tensor(); + ASSERT(weight_cpu->ndim() == 2); + ASSERT(weight_cpu->shape()[0] == num_experts_); + ASSERT(weight_cpu->shape()[1] == hidden_size); + if (num_expert_groups_ > 1) { + ASSERT(weight_1_cpu->ndim() == 2); + ASSERT(weight_1_cpu->shape()[0] == num_experts_); + ASSERT(weight_1_cpu->shape()[1] == hidden_size); + } + } + + for (size_t itok = 0; itok < ntoken; ++itok) { + const size_t expert_group = token_groups == nullptr ? 0 : (*token_groups)[itok]; + ASSERT(expert_group < num_expert_groups_); + + if (use_gpu_gate_linear) { + for (size_t iexpert = 0; iexpert < num_experts_; ++iexpert) { + logits[iexpert] = read_as_f32_2d(router_logits_cpu, itok, iexpert); + } + } else { + const auto &gate_weight = expert_group == 0 ? weight_cpu : weight_1_cpu; + + for (size_t iexpert = 0; iexpert < num_experts_; ++iexpert) { + float acc = 0.0f; + for (size_t ihidden = 0; ihidden < hidden_size; ++ihidden) { + acc += read_as_f32_2d(hidden_states_cpu, itok, ihidden) * read_as_f32_2d(gate_weight, iexpert, ihidden); + } + logits[iexpert] = acc; + } + } + + float max_logit = logits[0]; + for (size_t iexpert = 1; iexpert < num_experts_; ++iexpert) { + max_logit = std::max(max_logit, logits[iexpert]); + } + + float denom = 0.0f; + for (size_t iexpert = 0; iexpert < num_experts_; ++iexpert) { + probs[iexpert] = std::exp(logits[iexpert] - max_logit); + denom += probs[iexpert]; + } + if (!std::isfinite(max_logit) || !std::isfinite(denom) || denom <= 0.0f) { + SPDLOG_ERROR("ERNIE MoE router invalid CPU softmax: token={} max_logit={} denom={} hidden_dtype={} weight_dtype={}", + itok, + max_logit, + denom, + static_cast(use_gpu_gate_linear ? hidden_states->dtype() : hidden_states_cpu->dtype()), + static_cast(use_gpu_gate_linear ? weight_->dtype() : weight_cpu->dtype())); + for (size_t i = 0; i < std::min(num_experts_, 8); ++i) { + SPDLOG_ERROR("ERNIE MoE router cpu_logit[{},{}]={}", itok, i, logits[i]); + } + if (!use_gpu_gate_linear) { + for (size_t i = 0; i < std::min(hidden_states_cpu->shape()[1], 8); ++i) { + SPDLOG_ERROR("ERNIE MoE router hidden[{},{}]={}", itok, i, read_as_f32_2d(hidden_states_cpu, itok, i)); + } + } + } + ASSERT(denom > 0.0f); + for (float &prob : probs) { + prob /= denom; + } + + auto route_score = [&](size_t expert) { + float score = probs[expert]; + if (use_correction_bias_) { + score += correction_bias_cpu->ndim() == 1 + ? read_as_f32_1d(correction_bias_cpu, expert) + : read_as_f32_2d(correction_bias_cpu, expert_group, expert); + } + return score; + }; + std::iota(expert_order.begin(), expert_order.end(), 0); + std::partial_sort( + expert_order.begin(), + expert_order.begin() + num_experts_per_tok_, + expert_order.end(), + [&](size_t lhs, size_t rhs) { + const float lhs_score = route_score(lhs); + const float rhs_score = route_score(rhs); + if (lhs_score == rhs_score) { + return lhs < rhs; + } + return lhs_score > rhs_score; + }); + + float selected_sum = 0.0f; + for (size_t k = 0; k < num_experts_per_tok_; ++k) { + selected_sum += probs[expert_order[k]]; + } + if (norm_topk_prob_) { + selected_sum = std::max(selected_sum, 1e-12f); + } + + const size_t topk_offset = itok * num_experts_per_tok_; + for (size_t k = 0; k < num_experts_per_tok_; ++k) { + const size_t expert = expert_order[k]; + router_indices_ptr[topk_offset + k] = static_cast(expert + expert_group * num_experts_); + router_scores_ptr[topk_offset + k] = norm_topk_prob_ ? (probs[expert] / selected_sum) : probs[expert]; + } + } + + return std::make_tuple(router_scores_cpu->to(hidden_states->device()), + router_indices_cpu->to(hidden_states->device())); +} + +Ernie4_5MoeStatics::Ernie4_5MoeStatics(std::shared_ptr model_config, + const infinicore::Device &device) { + const auto &json = model_config->get_config_json(); + size_t num_experts = json_first_size(json.value("moe_num_experts", nlohmann::json(64)), 64); + size_t num_groups = 1; + if (json.contains("moe_num_experts") && json["moe_num_experts"].is_array()) { + num_groups = json["moe_num_experts"].size(); + } + INFINICORE_NN_PARAMETER_INIT(e_score_correction_bias, ({num_groups, num_experts}, infinicore::DataType::F32, device)); +} + +Ernie4_5TextMoeBlock::Ernie4_5TextMoeBlock(std::shared_ptr model_config, + const infinicore::Device &device) { + const auto &json = model_config->get_config_json(); + const auto experts_json = json.value("moe_num_experts", nlohmann::json(64)); + const auto intermediate_json = json.value("moe_intermediate_size", nlohmann::json(1536)); + num_experts_ = json_first_size(experts_json, 64); + num_expert_groups_ = experts_json.is_array() ? experts_json.size() : 1; + num_experts_per_tok_ = model_config->get_or("moe_k", 6); + size_t num_shared_experts = model_config->get_or("moe_num_shared_experts", 0); + ASSERT(num_expert_groups_ >= 1); + ASSERT(num_expert_groups_ <= 2); + for (size_t group = 1; group < num_expert_groups_; ++group) { + ASSERT_EQ(json_size_at(experts_json, group, num_experts_), num_experts_); + } + + INFINICORE_NN_MODULE_INIT(gate, model_config, device); + INFINICORE_NN_MODULE_INIT(moe_statics, model_config, device); + + experts_.reserve(num_experts_ * num_expert_groups_); + size_t expert_idx = 0; + for (size_t group = 0; group < num_expert_groups_; ++group) { + const size_t intermediate_size = json_size_at(intermediate_json, group, json_first_size(intermediate_json, 1536)); + auto expert_config = clone_with_moe_intermediate(model_config, intermediate_size); + for (size_t i = 0; i < num_experts_; ++i) { + experts_.push_back(this->register_module( + "experts." + std::to_string(expert_idx), expert_config, device)); + ++expert_idx; + } + } + + if (num_shared_experts > 0) { + const size_t intermediate_size = json_first_size(intermediate_json, 1536); + auto shared_config = clone_with_moe_intermediate(model_config, intermediate_size * num_shared_experts); + INFINICORE_NN_MODULE_INIT(shared_experts, shared_config, device); + } +} + +std::vector +Ernie4_5TextMoeBlock::token_groups_(const infinicore::Tensor &token_type_ids, + const std::vector &hidden_shape) const { + std::vector groups; + if (!token_type_ids || num_expert_groups_ <= 1) { + return groups; + } + + const bool restore_3d = hidden_shape.size() == 3; + const size_t batch = restore_3d ? hidden_shape[0] : 1; + const size_t seq_len = restore_3d ? hidden_shape[1] : hidden_shape[0]; + const size_t ntoken = batch * seq_len; + + auto token_types_cpu = token_type_ids->to(infinicore::Device::Type::CPU); + const auto tt_shape = token_types_cpu->shape(); + groups.assign(ntoken, 0); + bool any_multimodal = false; + + if (restore_3d && tt_shape.size() == 2) { + ASSERT_EQ(tt_shape[0], batch); + ASSERT(tt_shape[1] >= seq_len); + for (size_t b = 0; b < batch; ++b) { + for (size_t s = 0; s < seq_len; ++s) { + const int64_t token_type = read_as_i64_2d(token_types_cpu, b, s); + const size_t group = token_type == 0 ? 0 : 1; + groups[b * seq_len + s] = group; + any_multimodal = any_multimodal || group != 0; + } + } + } else { + ASSERT(token_types_cpu->numel() >= ntoken); + for (size_t i = 0; i < ntoken; ++i) { + const int64_t token_type = token_types_cpu->ndim() == 1 + ? read_as_i64_1d(token_types_cpu, i) + : read_as_i64_offset(token_types_cpu, i); + const size_t group = token_type == 0 ? 0 : 1; + groups[i] = group; + any_multimodal = any_multimodal || group != 0; + } + } + + if (!any_multimodal) { + groups.clear(); + } + return groups; +} + +infinicore::Tensor Ernie4_5TextMoeBlock::forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &token_type_ids) const { + ASSERT((hidden_states->ndim() == 2) || (hidden_states->ndim() == 3)); + + auto shape = hidden_states->shape(); + bool restore_3d = hidden_states->ndim() == 3; + auto hidden_states_2d = restore_3d ? hidden_states->view({shape[0] * shape[1], shape[2]}) : hidden_states; + + auto token_groups = token_groups_(token_type_ids, shape); + auto [routing_weights, selected_experts] = gate_->forward( + hidden_states_2d, + moe_statics_->e_score_correction_bias(), + token_groups.empty() ? nullptr : &token_groups); + auto routing_weights_cpu = routing_weights->to(infinicore::Device::Type::CPU); + auto selected_experts_cpu = selected_experts->to(infinicore::Device::Type::CPU); + + float *routing_weights_ptr = reinterpret_cast(routing_weights_cpu->data()); + int *selected_experts_ptr = reinterpret_cast(selected_experts_cpu->data()); + + size_t ntoken = hidden_states_2d->shape()[0]; + auto final_hidden_states = infinicore::Tensor::empty(hidden_states_2d->shape(), hidden_states_2d->dtype(), hidden_states_2d->device()); + infinicore::Tensor shared_out; + if (shared_experts_) { + shared_out = shared_experts_->forward(hidden_states_2d); + } + + for (size_t itok = 0; itok < ntoken; ++itok) { + auto hidden_states_i = hidden_states_2d->narrow({{0, itok, 1}}); + const size_t route_row = itok * num_experts_per_tok_; + + infinicore::Tensor final_hidden_states_i; + for (size_t k = 0; k < num_experts_per_tok_; ++k) { + int index = selected_experts_ptr[route_row + k]; + float score = routing_weights_ptr[route_row + k]; + ASSERT(index >= 0 && static_cast(index) < experts_.size()); + + experts_[index]->set_alpha(score); + auto expert_out = experts_[index]->forward(hidden_states_i); + if (k == 0) { + final_hidden_states_i = expert_out; + } else { + infinicore::op::add_(final_hidden_states_i, final_hidden_states_i, expert_out); + } + } + + if (shared_out) { + auto shared_i = shared_out->narrow({{0, itok, 1}}); + infinicore::op::add_(final_hidden_states_i, final_hidden_states_i, shared_i); + } + final_hidden_states->narrow({{0, itok, 1}})->copy_from(final_hidden_states_i); + } + + if (restore_3d) { + return final_hidden_states->view({shape[0], shape[1], shape[2]}); + } + return final_hidden_states; +} + +} // namespace infinilm::models::ernie4_5_moe_vl diff --git a/csrc/models/ernie4_5_moe_vl/ernie4_5_moe.hpp b/csrc/models/ernie4_5_moe_vl/ernie4_5_moe.hpp new file mode 100644 index 000000000..5531886ac --- /dev/null +++ b/csrc/models/ernie4_5_moe_vl/ernie4_5_moe.hpp @@ -0,0 +1,66 @@ +#pragma once + +#include "../../layers/moe/legacy/moe_mlp.hpp" +#include "infinicore/nn/module.hpp" +#include "infinicore/tensor.hpp" + +#include +#include +#include + +namespace infinilm::models::ernie4_5_moe_vl { + +class Ernie4_5MoeTopKRouter : public infinicore::nn::Module { +public: + Ernie4_5MoeTopKRouter(std::shared_ptr model_config, + const infinicore::Device &device); + + std::tuple forward( + const infinicore::Tensor &hidden_states, + const infinicore::Tensor &score_correction_bias, + const std::vector *token_groups = nullptr) const; + +private: + INFINICORE_NN_PARAMETER(weight); + INFINICORE_NN_PARAMETER(weight_1); + size_t num_experts_per_tok_{0}; + size_t num_experts_{0}; + size_t num_expert_groups_{1}; + bool norm_topk_prob_{true}; + bool use_correction_bias_{false}; +}; + +class Ernie4_5MoeStatics : public infinicore::nn::Module { +public: + Ernie4_5MoeStatics(std::shared_ptr model_config, + const infinicore::Device &device); + + infinicore::Tensor e_score_correction_bias() const { return e_score_correction_bias_; } + +private: + INFINICORE_NN_PARAMETER(e_score_correction_bias); +}; + +class Ernie4_5TextMoeBlock : public infinicore::nn::Module { +public: + Ernie4_5TextMoeBlock(std::shared_ptr model_config, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &token_type_ids = infinicore::Tensor()) const; + +private: + std::vector token_groups_(const infinicore::Tensor &token_type_ids, + const std::vector &hidden_shape) const; + + INFINICORE_NN_MODULE(Ernie4_5MoeTopKRouter, gate); + INFINICORE_NN_MODULE(Ernie4_5MoeStatics, moe_statics); + INFINICORE_NN_MODULE_VEC(infinilm::layers::moe::legacy::MoeMLP, experts); + INFINICORE_NN_MODULE(infinilm::layers::moe::legacy::MoeMLP, shared_experts); + + size_t num_experts_per_tok_{0}; + size_t num_experts_{0}; + size_t num_expert_groups_{1}; +}; + +} // namespace infinilm::models::ernie4_5_moe_vl diff --git a/csrc/models/ernie4_5_moe_vl/ernie4_5_moe_vl_for_causal_lm.cpp b/csrc/models/ernie4_5_moe_vl/ernie4_5_moe_vl_for_causal_lm.cpp new file mode 100644 index 000000000..982ff85b7 --- /dev/null +++ b/csrc/models/ernie4_5_moe_vl/ernie4_5_moe_vl_for_causal_lm.cpp @@ -0,0 +1,222 @@ +#include "ernie4_5_moe_vl_for_causal_lm.hpp" + +#include "../models_registry.hpp" +#include "infinicore/nn/rope.hpp" +#include "infinicore/ops.hpp" +#include "infinicore/ops/cat.hpp" + +#include +#include +#include + +namespace infinilm::models::ernie4_5_moe_vl { + +Ernie4_5Model::Ernie4_5Model(std::shared_ptr model_config, + const infinicore::Device &device) { + const auto &dtype{model_config->get_dtype()}; + size_t vocab_size = model_config->get("vocab_size"); + size_t hidden_size = model_config->get("hidden_size"); + size_t num_hidden_layers = model_config->get("num_hidden_layers"); + double rms_norm_eps = model_config->get("rms_norm_eps"); + + 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, device)); + } + + INFINICORE_NN_MODULE_INIT(norm, hidden_size, rms_norm_eps, dtype, device); + INFINICORE_NN_MODULE_INIT(resampler_model, model_config, device); +} + +infinicore::Tensor Ernie4_5Model::forward(const infinilm::InfinilmModel::Input &input) const { + auto input_ids = input.input_ids.value(); + auto positions = input.position_ids.value(); + auto hidden_states = embed_tokens_->forward(input_ids); + return forward_embeds(hidden_states, positions, input.token_type_ids.value_or(infinicore::Tensor())); +} + +infinicore::Tensor Ernie4_5Model::forward_embeds(const infinicore::Tensor &inputs_embeds, + const infinicore::Tensor &position_ids, + const infinicore::Tensor &token_type_ids) const { + auto hidden_states = inputs_embeds; + infinicore::Tensor residual; + for (const auto &layer : layers_) { + layer->forward(position_ids, hidden_states, residual, token_type_ids); + } + norm_->forward_inplace(hidden_states, residual); + return hidden_states; +} + +infinicore::Tensor Ernie4_5Model::embed_tokens(const infinicore::Tensor &input_ids) const { + return embed_tokens_->forward(input_ids); +} + +infinicore::Tensor Ernie4_5Model::resample_vision(const infinicore::Tensor &vision_features, + const infinicore::Tensor &grid_thw) const { + return resampler_model_->forward(vision_features, grid_thw); +} + +Ernie4_5MoeVLForCausalLM::Ernie4_5MoeVLForCausalLM(std::shared_ptr model_config, + const infinicore::Device &device) { + model_config_ = model_config; + + size_t hidden_size = model_config->get("hidden_size"); + size_t vocab_size = model_config->get("vocab_size"); + const auto &dtype{model_config->get_dtype()}; + im_patch_id_ = model_config->get_or("im_patch_id", 100295); + + INFINICORE_NN_MODULE_INIT(vision_model, model_config, device); + INFINICORE_NN_MODULE_INIT(model, model_config, device); + INFINICORE_NN_MODULE_INIT(lm_head, hidden_size, vocab_size, false, dtype, device); +} + +InfinilmModel::Output Ernie4_5MoeVLForCausalLM::forward(const Input &input) const { + if (input.pixel_values.has_value() && !input.pixel_values->empty()) { + auto inputs_embeds = build_multimodal_embeds_(input); + auto hidden_states = model_->forward_embeds( + inputs_embeds, + input.position_ids.value(), + input.token_type_ids.value_or(infinicore::Tensor())); + auto logits = lm_head_->forward(hidden_states); + return {logits}; + } + auto hidden_states = model_->forward(input); + auto logits = lm_head_->forward(hidden_states); + return {logits}; +} + +infinicore::Tensor Ernie4_5MoeVLForCausalLM::logits_from_hidden(const infinicore::Tensor &hidden_states) const { + return lm_head_->forward(const_cast(hidden_states)); +} + +infinicore::Tensor Ernie4_5MoeVLForCausalLM::concat_optional_tensors_( + const std::optional> &tensors, + int dim) const { + if (!tensors.has_value() || tensors->empty()) { + return infinicore::Tensor(); + } + if (tensors->size() == 1) { + return tensors->front(); + } + return infinicore::op::cat(*tensors, dim); +} + +infinicore::Tensor Ernie4_5MoeVLForCausalLM::replace_image_embeds_(const infinicore::Tensor &input_ids, + const infinicore::Tensor &inputs_embeds, + const infinicore::Tensor &image_features) const { + ASSERT(input_ids->ndim() == 2); + ASSERT(inputs_embeds->ndim() == 3); + ASSERT(image_features->ndim() == 2); + ASSERT_EQ(input_ids->shape()[0], inputs_embeds->shape()[0]); + ASSERT_EQ(input_ids->shape()[1], inputs_embeds->shape()[1]); + ASSERT_EQ(inputs_embeds->shape()[2], image_features->shape()[1]); + + auto ids_cpu = input_ids->to(infinicore::Device::cpu()); + auto embeds_cpu = inputs_embeds->to(infinicore::Device::cpu()); + auto features_cpu = image_features->to(infinicore::Device::cpu()); + auto out_cpu = infinicore::Tensor::empty(inputs_embeds->shape(), inputs_embeds->dtype(), infinicore::Device::cpu()); + out_cpu->copy_from(embeds_cpu); + + const auto *ids = reinterpret_cast(ids_cpu->data()); + const auto *features = reinterpret_cast(features_cpu->data()); + auto *out = reinterpret_cast(out_cpu->data()); + const size_t batch = input_ids->shape()[0]; + const size_t seq_len = input_ids->shape()[1]; + const size_t hidden = inputs_embeds->shape()[2]; + const size_t elem_size = inputs_embeds->element_size(); + const size_t row_bytes = hidden * elem_size; + + size_t image_row = 0; + for (size_t b = 0; b < batch; ++b) { + for (size_t s = 0; s < seq_len; ++s) { + if (ids[b * seq_len + s] != static_cast(im_patch_id_)) { + continue; + } + if (image_row >= image_features->shape()[0]) { + throw std::runtime_error("Ernie4_5MoeVLForCausalLM: fewer image features than image patch tokens"); + } + const size_t dst_row = (b * seq_len + s) * hidden; + std::memcpy(out + dst_row * elem_size, features + image_row * row_bytes, row_bytes); + ++image_row; + } + } + if (image_row != image_features->shape()[0]) { + throw std::runtime_error("Ernie4_5MoeVLForCausalLM: image feature count does not match image patch tokens"); + } + + return out_cpu->to(inputs_embeds->device()); +} + +infinicore::Tensor Ernie4_5MoeVLForCausalLM::build_multimodal_embeds_(const Input &input) const { + if (!input.input_ids.has_value() || !input.position_ids.has_value()) { + throw std::runtime_error("Ernie4_5MoeVLForCausalLM: input_ids and position_ids are required"); + } + auto images = concat_optional_tensors_(input.pixel_values, 0); + auto grid_thw = concat_optional_tensors_(input.grid_thw, 0); + if (!images || !grid_thw) { + throw std::runtime_error("Ernie4_5MoeVLForCausalLM: images and grid_thw are required for vision input"); + } + + auto input_ids = input.input_ids.value(); + auto inputs_embeds = model_->embed_tokens(input_ids); + auto image_features = vision_model_->forward(images, grid_thw); + image_features = model_->resample_vision(image_features, grid_thw); + return replace_image_embeds_(input_ids, inputs_embeds, image_features); +} + +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("create_ernie4_5_moe_vl_model_config: model_type is not ernie4_5_moe_vl"); + } + + nlohmann::json &config_json = model_config->get_config_json(); + if (!config_json.contains("head_dim")) { + config_json["head_dim"] = model_config->get("hidden_size") + / model_config->get("num_attention_heads"); + } + if (!config_json.contains("num_key_value_heads") || config_json["num_key_value_heads"].is_null()) { + config_json["num_key_value_heads"] = model_config->get("num_attention_heads"); + } + if (!config_json.contains("rope_theta")) { + config_json["rope_theta"] = 10000.0; + } + if (!config_json.contains("attention_bias")) { + config_json["attention_bias"] = config_json.value("use_bias", false); + } + if (!config_json.contains("attention_output_bias")) { + config_json["attention_output_bias"] = config_json.value("use_bias", false); + } + if (!config_json.contains("mlp_bias")) { + config_json["mlp_bias"] = config_json.value("use_bias", false); + } + if (!config_json.contains("moe_k")) { + config_json["moe_k"] = 6; + } + if (!config_json.contains("moe_num_shared_experts")) { + config_json["moe_num_shared_experts"] = 0; + } + if (!config_json.contains("moe_norm_gate_logits")) { + config_json["moe_norm_gate_logits"] = true; + } + if (!config_json.contains("moe_use_aux_free")) { + config_json["moe_use_aux_free"] = false; + } + + // ERNIE applies RoPE on adjacent even/odd pairs, matching InfiniCore's GPT-J layout. + model_config->set_rope_algo(infinicore::nn::RoPE::Algo::GPT_J); + + return model_config; +} + +} // namespace infinilm::models::ernie4_5_moe_vl + +namespace { +INFINILM_REGISTER_CAUSAL_LM_MODEL( + ernie4_5_moe_vl, + infinilm::models::ernie4_5_moe_vl::Ernie4_5MoeVLForCausalLM, + infinilm::models::ernie4_5_moe_vl::create_ernie4_5_moe_vl_model_config); +} // namespace diff --git a/csrc/models/ernie4_5_moe_vl/ernie4_5_moe_vl_for_causal_lm.hpp b/csrc/models/ernie4_5_moe_vl/ernie4_5_moe_vl_for_causal_lm.hpp new file mode 100644 index 000000000..38faf2836 --- /dev/null +++ b/csrc/models/ernie4_5_moe_vl/ernie4_5_moe_vl_for_causal_lm.hpp @@ -0,0 +1,62 @@ +#pragma once + +#include "../../layers/linear/linear.hpp" +#include "../infinilm_model.hpp" +#include "ernie4_5_decoder_layer.hpp" +#include "ernie4_5_resampler.hpp" +#include "ernie4_5_vision.hpp" +#include "infinicore/nn/embedding.hpp" +#include "infinicore/nn/rmsnorm.hpp" + +#include + +namespace infinilm::models::ernie4_5_moe_vl { + +class Ernie4_5Model : public infinicore::nn::Module { +public: + Ernie4_5Model(std::shared_ptr model_config, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinilm::InfinilmModel::Input &input) const; + infinicore::Tensor forward_embeds(const infinicore::Tensor &inputs_embeds, + const infinicore::Tensor &position_ids, + const infinicore::Tensor &token_type_ids = infinicore::Tensor()) const; + infinicore::Tensor embed_tokens(const infinicore::Tensor &input_ids) const; + infinicore::Tensor resample_vision(const infinicore::Tensor &vision_features, + const infinicore::Tensor &grid_thw) const; + +private: + INFINICORE_NN_MODULE(infinicore::nn::Embedding, embed_tokens); + INFINICORE_NN_MODULE_VEC(Ernie4_5DecoderLayer, layers); + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, norm); + INFINICORE_NN_MODULE(Ernie4_5Resampler, resampler_model); +}; + +class Ernie4_5MoeVLForCausalLM : public InfinilmModel { +public: + Ernie4_5MoeVLForCausalLM(std::shared_ptr model_config, + const infinicore::Device &device); + + Output forward(const Input &input) const override; + infinicore::Tensor logits_from_hidden(const infinicore::Tensor &hidden_states) const; + Ernie4_5Model &model() { return *model_; } + +private: + infinicore::Tensor build_multimodal_embeds_(const Input &input) const; + infinicore::Tensor replace_image_embeds_(const infinicore::Tensor &input_ids, + const infinicore::Tensor &inputs_embeds, + const infinicore::Tensor &image_features) const; + infinicore::Tensor concat_optional_tensors_(const std::optional> &tensors, + int dim) const; + + size_t im_patch_id_{100295}; + + INFINICORE_NN_MODULE(Ernie4_5VisionModel, vision_model); + INFINICORE_NN_MODULE(Ernie4_5Model, 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_moe_vl diff --git a/csrc/models/ernie4_5_moe_vl/ernie4_5_resampler.cpp b/csrc/models/ernie4_5_moe_vl/ernie4_5_resampler.cpp new file mode 100644 index 000000000..7b252172b --- /dev/null +++ b/csrc/models/ernie4_5_moe_vl/ernie4_5_resampler.cpp @@ -0,0 +1,113 @@ +#include "ernie4_5_resampler.hpp" + +#include "../../utils.hpp" +#include "infinicore/ops.hpp" + +#include +#include +#include + +namespace infinilm::models::ernie4_5_moe_vl { + +Ernie4_5Resampler::Ernie4_5Resampler(std::shared_ptr model_config, + const infinicore::Device &device) { + const auto &dtype{model_config->get_dtype()}; + in_dim_ = model_config->get_or("pixel_hidden_size", 1280); + out_dim_ = model_config->get("hidden_size"); + spatial_conv_size_ = model_config->get_or("spatial_conv_size", 2); + temporal_conv_size_ = model_config->get_or("temporal_conv_size", 2); + use_temporal_conv_ = model_config->get_or("use_temporal_conv", true); + + spatial_dim_ = in_dim_ * spatial_conv_size_ * spatial_conv_size_; + temporal_dim_ = spatial_dim_ * temporal_conv_size_; + + 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); + 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); + mlp_ = this->register_module("mlp", spatial_dim_, out_dim_, true, dtype, device); + after_norm_ = this->register_module("after_norm", out_dim_, model_config->get("rms_norm_eps"), dtype, device); +} + +infinicore::Tensor Ernie4_5Resampler::temporal_placeholder_(const infinicore::Tensor &x, + const infinicore::Tensor &grid_thw) const { + ASSERT(x->ndim() == 2); + ASSERT(x->shape()[1] == spatial_dim_); + ASSERT(temporal_conv_size_ == 2); + + auto grid_cpu = grid_thw->to(infinicore::Device::cpu()); + ASSERT(grid_cpu->ndim() == 2); + ASSERT(grid_cpu->shape()[1] == 3); + auto *grid_ptr = reinterpret_cast(grid_cpu->data()); + const size_t n_grid = grid_cpu->shape()[0]; + + size_t total_rows = 0; + for (size_t i = 0; i < n_grid; ++i) { + const size_t t = static_cast(grid_ptr[i * 3]); + const size_t h = static_cast(grid_ptr[i * 3 + 1]); + const size_t w = static_cast(grid_ptr[i * 3 + 2]); + ASSERT(h % spatial_conv_size_ == 0); + ASSERT(w % spatial_conv_size_ == 0); + const size_t spatial_size = (h * w) / (spatial_conv_size_ * spatial_conv_size_); + total_rows += ((t + 1) / 2) * spatial_size; + } + + auto out = infinicore::Tensor::empty({total_rows, temporal_dim_}, x->dtype(), x->device()); + + size_t input_base = 0; + size_t out_row = 0; + for (size_t i = 0; i < n_grid; ++i) { + const size_t t = static_cast(grid_ptr[i * 3]); + const size_t h = static_cast(grid_ptr[i * 3 + 1]); + const size_t w = static_cast(grid_ptr[i * 3 + 2]); + const size_t spatial_size = (h * w) / (spatial_conv_size_ * spatial_conv_size_); + + for (size_t temp_offset = 0; temp_offset < t; temp_offset += 2) { + const size_t temp_offset2 = (t == 1) ? 0 : std::min(temp_offset + 1, t - 1); + for (size_t row = 0; row < spatial_size; ++row) { + const size_t src1 = input_base + temp_offset * spatial_size + row; + const size_t src2 = input_base + temp_offset2 * spatial_size + row; + auto dst_first = out->narrow({{0, out_row, 1}, {1, 0, spatial_dim_}}); + auto dst_second = out->narrow({{0, out_row, 1}, {1, spatial_dim_, spatial_dim_}}); + dst_first->copy_from(x->narrow({{0, src1, 1}})); + dst_second->copy_from(x->narrow({{0, src2, 1}})); + ++out_row; + } + } + input_base += t * spatial_size; + } + + ASSERT_EQ(out_row, total_rows); + ASSERT_EQ(input_base, x->shape()[0]); + return out; +} + +infinicore::Tensor Ernie4_5Resampler::forward(const infinicore::Tensor &vision_features, + const infinicore::Tensor &grid_thw) const { + ASSERT(vision_features->ndim() == 2); + ASSERT(vision_features->shape()[1] == in_dim_); + const size_t merge = spatial_conv_size_ * spatial_conv_size_; + ASSERT(vision_features->shape()[0] % merge == 0); + + auto x = vision_features->view({vision_features->shape()[0] / merge, spatial_dim_}); + 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_) { + x = temporal_placeholder_(x, grid_thw); + 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); + x = after_norm_->forward(x); + return x; +} + +} // namespace infinilm::models::ernie4_5_moe_vl diff --git a/csrc/models/ernie4_5_moe_vl/ernie4_5_resampler.hpp b/csrc/models/ernie4_5_moe_vl/ernie4_5_resampler.hpp new file mode 100644 index 000000000..5da530ca5 --- /dev/null +++ b/csrc/models/ernie4_5_moe_vl/ernie4_5_resampler.hpp @@ -0,0 +1,44 @@ +#pragma once + +#include "../../config/model_config.hpp" +#include "../../layers/linear/linear.hpp" +#include "infinicore/nn/layer_norm.hpp" +#include "infinicore/nn/module.hpp" +#include "infinicore/nn/rmsnorm.hpp" +#include "infinicore/tensor.hpp" + +#include + +namespace infinilm::models::ernie4_5_moe_vl { + +class Ernie4_5Resampler : public infinicore::nn::Module { +public: + Ernie4_5Resampler(std::shared_ptr model_config, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &vision_features, + const infinicore::Tensor &grid_thw) const; + +private: + infinicore::Tensor temporal_placeholder_(const infinicore::Tensor &x, + const infinicore::Tensor &grid_thw) const; + + size_t in_dim_{0}; + size_t out_dim_{0}; + size_t spatial_conv_size_{2}; + size_t temporal_conv_size_{2}; + size_t spatial_dim_{0}; + size_t temporal_dim_{0}; + bool use_temporal_conv_{true}; + + INFINICORE_NN_MODULE(infinilm::nn::Linear, spatial_linear_0); + INFINICORE_NN_MODULE(infinilm::nn::Linear, spatial_linear_2); + INFINICORE_NN_MODULE(infinicore::nn::LayerNorm, spatial_linear_3); + INFINICORE_NN_MODULE(infinilm::nn::Linear, temporal_linear_0); + INFINICORE_NN_MODULE(infinilm::nn::Linear, temporal_linear_2); + INFINICORE_NN_MODULE(infinicore::nn::LayerNorm, temporal_linear_3); + INFINICORE_NN_MODULE(infinilm::nn::Linear, mlp); + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, after_norm); +}; + +} // namespace infinilm::models::ernie4_5_moe_vl diff --git a/csrc/models/ernie4_5_moe_vl/ernie4_5_vision.cpp b/csrc/models/ernie4_5_moe_vl/ernie4_5_vision.cpp new file mode 100644 index 000000000..536a28f49 --- /dev/null +++ b/csrc/models/ernie4_5_moe_vl/ernie4_5_vision.cpp @@ -0,0 +1,374 @@ +#include "ernie4_5_vision.hpp" + +#include "../../utils.hpp" +#include "infinicore/ops.hpp" +#include "infinicore/ops/cat.hpp" +#include "infinicore/ops/gelu.hpp" +#include "infinicore/ops/gelutanh.hpp" +#include "infinicore/ops/matmul.hpp" +#include "infinicore/ops/mha.hpp" +#include "infinicore/ops/quickgelu.hpp" +#include "infinicore/ops/relu.hpp" +#include "infinicore/ops/softmax.hpp" + +#include +#include +#include + +namespace infinilm::models::ernie4_5_moe_vl { +namespace { + +constexpr float kImageMean[3] = {0.48145466f, 0.4578275f, 0.40821073f}; +constexpr float kImageStd[3] = {0.26862954f, 0.26130258f, 0.27577711f}; +constexpr float kRescaleFactor = 0.00392156862745098f; +constexpr float kRopeTheta = 10000.0f; + +uint16_t fp32_to_bf16(float value) { + uint32_t bits; + std::memcpy(&bits, &value, sizeof(bits)); + const uint32_t lsb = (bits >> 16) & 1U; + bits += 0x7FFFU + lsb; + return static_cast(bits >> 16); +} + +float bf16_to_fp32(uint16_t value) { + uint32_t bits = static_cast(value) << 16; + float out; + std::memcpy(&out, &bits, sizeof(out)); + return out; +} + +void write_float(void *dst, size_t index, infinicore::DataType dtype, float value) { + if (dtype == infinicore::DataType::F32) { + reinterpret_cast(dst)[index] = value; + return; + } + if (dtype == infinicore::DataType::BF16) { + reinterpret_cast(dst)[index] = fp32_to_bf16(value); + return; + } + throw std::runtime_error("ERNIE vision: only float32 and bfloat16 host conversion are supported"); +} + +float read_float(const void *src, size_t index, infinicore::DataType dtype) { + if (dtype == infinicore::DataType::F32) { + return reinterpret_cast(src)[index]; + } + if (dtype == infinicore::DataType::BF16) { + return bf16_to_fp32(reinterpret_cast(src)[index]); + } + throw std::runtime_error("ERNIE vision: only float32 and bfloat16 host conversion are supported"); +} + +std::vector read_grid(const infinicore::Tensor &grid_thw) { + auto grid_cpu = grid_thw->to(infinicore::Device::cpu()); + ASSERT(grid_cpu->ndim() == 2); + ASSERT(grid_cpu->shape()[1] == 3); + const auto *grid_ptr = reinterpret_cast(grid_cpu->data()); + return std::vector(grid_ptr, grid_ptr + grid_cpu->numel()); +} + +std::vector> build_vision_pos_ids(const std::vector &grid, + size_t spatial_merge_size) { + std::vector> pos_ids; + for (size_t row = 0; row < grid.size(); row += 3) { + const int64_t t = grid[row]; + const int64_t h = grid[row + 1]; + const int64_t w = grid[row + 2]; + ASSERT(h % static_cast(spatial_merge_size) == 0); + ASSERT(w % static_cast(spatial_merge_size) == 0); + + for (int64_t ti = 0; ti < t; ++ti) { + for (int64_t hb = 0; hb < h; hb += static_cast(spatial_merge_size)) { + for (int64_t wb = 0; wb < w; wb += static_cast(spatial_merge_size)) { + for (int64_t ih = 0; ih < static_cast(spatial_merge_size); ++ih) { + for (int64_t iw = 0; iw < static_cast(spatial_merge_size); ++iw) { + pos_ids.emplace_back(hb + ih, wb + iw); + } + } + } + } + } + } + return pos_ids; +} + +std::vector segment_lengths(const std::vector &grid) { + std::vector lengths; + for (size_t row = 0; row < grid.size(); row += 3) { + const int64_t t = grid[row]; + const int64_t h = grid[row + 1]; + const int64_t w = grid[row + 2]; + for (int64_t ti = 0; ti < t; ++ti) { + lengths.push_back(static_cast(h * w)); + } + } + return lengths; +} + +} // namespace + +Ernie4_5VisionPatchEmbed::Ernie4_5VisionPatchEmbed(const nlohmann::json &vision_config, + const infinicore::DataType &dtype, + const infinicore::Device &device) { + patch_size_ = vision_config.value("patch_size", 14); + in_channels_ = vision_config.value("in_channels", vision_config.value("in_chans", 3)); + embed_dim_ = vision_config.value("embed_dim", vision_config.value("hidden_size", 1280)); + + INFINICORE_NN_MODULE_INIT(proj, in_channels_ * patch_size_ * patch_size_, embed_dim_, false, dtype, device); +} + +infinicore::Tensor Ernie4_5VisionPatchEmbed::forward(const infinicore::Tensor &hidden_states) const { + auto x = hidden_states; + return proj_->forward(x); +} + +Ernie4_5VisionAttention::Ernie4_5VisionAttention(const nlohmann::json &vision_config, + const infinicore::DataType &dtype, + const infinicore::Device &device) { + embed_dim_ = vision_config.value("embed_dim", vision_config.value("hidden_size", 1280)); + num_heads_ = vision_config.value("num_heads", 16); + spatial_merge_size_ = vision_config.value("spatial_merge_size", 2); + if (embed_dim_ % num_heads_ != 0) { + throw std::runtime_error("Ernie4_5VisionAttention: embed_dim must be divisible by num_heads"); + } + head_dim_ = embed_dim_ / num_heads_; + scale_ = 1.0f / std::sqrt(static_cast(head_dim_)); + + INFINICORE_NN_MODULE_INIT(qkv, embed_dim_, embed_dim_ * 3, true, dtype, device); + INFINICORE_NN_MODULE_INIT(proj, embed_dim_, embed_dim_, true, dtype, device); +} + +infinicore::Tensor Ernie4_5VisionAttention::apply_rotary_pos_emb_(const infinicore::Tensor &tensor, + const infinicore::Tensor &grid_thw) const { + ASSERT(tensor->ndim() == 3); + ASSERT_EQ(tensor->shape()[1], num_heads_); + ASSERT_EQ(tensor->shape()[2], head_dim_); + ASSERT(head_dim_ % 4 == 0); + + const auto grid = read_grid(grid_thw); + const auto pos_ids = build_vision_pos_ids(grid, spatial_merge_size_); + ASSERT_EQ(pos_ids.size(), tensor->shape()[0]); + + auto src_cpu = tensor->to(infinicore::Device::cpu()); + auto out_cpu = infinicore::Tensor::empty(tensor->shape(), tensor->dtype(), infinicore::Device::cpu()); + const void *src = src_cpu->data(); + void *dst = out_cpu->data(); + + const size_t seq_len = tensor->shape()[0]; + const size_t half_dim = head_dim_ / 2; + const size_t freq_dim = head_dim_ / 4; + std::vector inv_freq(freq_dim); + for (size_t i = 0; i < freq_dim; ++i) { + inv_freq[i] = 1.0f / std::pow(kRopeTheta, static_cast(i * 2) / static_cast(half_dim)); + } + + for (size_t s = 0; s < seq_len; ++s) { + for (size_t h = 0; h < num_heads_; ++h) { + const size_t base = (s * num_heads_ + h) * head_dim_; + for (size_t j = 0; j < half_dim; ++j) { + const size_t freq_idx = j % freq_dim; + const float pos = j < freq_dim + ? static_cast(pos_ids[s].first) + : static_cast(pos_ids[s].second); + const float angle = pos * inv_freq[freq_idx]; + const float c = std::cos(angle); + const float sn = std::sin(angle); + const float x1 = read_float(src, base + j, tensor->dtype()); + const float x2 = read_float(src, base + half_dim + j, tensor->dtype()); + write_float(dst, base + j, tensor->dtype(), x1 * c - x2 * sn); + write_float(dst, base + half_dim + j, tensor->dtype(), x2 * c + x1 * sn); + } + } + } + + return out_cpu->to(tensor->device()); +} + +infinicore::Tensor Ernie4_5VisionAttention::segmented_attention_(const infinicore::Tensor &q, + const infinicore::Tensor &k, + const infinicore::Tensor &v, + const infinicore::Tensor &grid_thw) const { + const auto grid = read_grid(grid_thw); + const auto lengths = segment_lengths(grid); + std::vector outputs; + outputs.reserve(lengths.size()); + + size_t offset = 0; + for (size_t len : lengths) { + auto q_flat = q->narrow({{0, offset, len}}) + ->permute({1, 0, 2}) + ->contiguous() + ->view({num_heads_, len, head_dim_}); + auto k_flat = k->narrow({{0, offset, len}}) + ->permute({1, 0, 2}) + ->contiguous() + ->view({num_heads_, len, head_dim_}); + auto v_flat = v->narrow({{0, offset, len}}) + ->permute({1, 0, 2}) + ->contiguous() + ->view({num_heads_, len, head_dim_}); + + auto attn_weights = infinicore::op::matmul(q_flat, k_flat->permute({0, 2, 1}), scale_); + infinicore::op::softmax_(attn_weights, attn_weights, -1); + auto out = infinicore::op::matmul(attn_weights, v_flat) + ->view({num_heads_, len, head_dim_}) + ->permute({1, 0, 2}) + ->contiguous(); + outputs.push_back(out); + offset += len; + } + ASSERT_EQ(offset, q->shape()[0]); + + if (outputs.size() == 1) { + return outputs[0]; + } + return infinicore::op::cat(outputs, 0); +} + +infinicore::Tensor Ernie4_5VisionAttention::forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &grid_thw) const { + auto x = hidden_states; + auto qkv_states = qkv_->forward(x)->view({hidden_states->shape()[0], 3, num_heads_, head_dim_}); + auto q = qkv_states->narrow({{1, 0, 1}})->squeeze(1)->contiguous(); + auto k = qkv_states->narrow({{1, 1, 1}})->squeeze(1)->contiguous(); + auto v = qkv_states->narrow({{1, 2, 1}})->squeeze(1)->contiguous(); + + q = apply_rotary_pos_emb_(q, grid_thw); + k = apply_rotary_pos_emb_(k, grid_thw); + + auto attn_output = segmented_attention_(q, k, v, grid_thw) + ->contiguous() + ->view({hidden_states->shape()[0], embed_dim_}); + return proj_->forward(attn_output); +} + +Ernie4_5VisionMLP::Ernie4_5VisionMLP(const nlohmann::json &vision_config, + const infinicore::DataType &dtype, + const infinicore::Device &device) + : hidden_act_(vision_config.value("hidden_act", "quick_gelu")) { + const size_t dim = vision_config.value("embed_dim", vision_config.value("hidden_size", 1280)); + const double mlp_ratio = vision_config.value("mlp_ratio", 4.0); + const size_t hidden_dim = static_cast(static_cast(dim) * mlp_ratio); + INFINICORE_NN_MODULE_INIT(fc1, dim, hidden_dim, true, dtype, device); + INFINICORE_NN_MODULE_INIT(fc2, hidden_dim, dim, true, dtype, device); +} + +infinicore::Tensor Ernie4_5VisionMLP::forward(const infinicore::Tensor &hidden_states) const { + auto x_in = hidden_states; + auto x = fc1_->forward(x_in); + if (hidden_act_ == "quick_gelu") { + x = infinicore::op::quick_gelu(x); + } else if (hidden_act_ == "gelu") { + x = infinicore::op::gelu(x); + } else if (hidden_act_ == "gelu_tanh") { + x = infinicore::op::gelu_tanh(x); + } else if (hidden_act_ == "relu") { + x = infinicore::op::relu(x); + } else { + throw std::runtime_error("Ernie4_5VisionMLP: unsupported activation " + hidden_act_); + } + return fc2_->forward(x); +} + +Ernie4_5VisionBlock::Ernie4_5VisionBlock(const nlohmann::json &vision_config, + const infinicore::DataType &dtype, + const infinicore::Device &device) { + const size_t hidden_size = vision_config.value("embed_dim", vision_config.value("hidden_size", 1280)); + INFINICORE_NN_MODULE_INIT(norm1, hidden_size, 1e-6, dtype, device); + INFINICORE_NN_MODULE_INIT(attn, vision_config, dtype, device); + INFINICORE_NN_MODULE_INIT(norm2, hidden_size, 1e-6, dtype, device); + INFINICORE_NN_MODULE_INIT(mlp, vision_config, dtype, device); +} + +infinicore::Tensor Ernie4_5VisionBlock::forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &grid_thw) const { + auto residual = hidden_states; + auto x = norm1_->forward(hidden_states); + x = attn_->forward(x, grid_thw); + x = infinicore::op::add(x, residual); + + residual = x; + x = norm2_->forward(x); + x = mlp_->forward(x); + return infinicore::op::add(x, residual); +} + +Ernie4_5VisionModel::Ernie4_5VisionModel(std::shared_ptr model_config, + const infinicore::Device &device) { + vision_config_ = model_config->get_config_json().value("vision_config", nlohmann::json::object()); + dtype_ = model_config->get_dtype(); + depth_ = vision_config_.value("depth", 32); + hidden_size_ = vision_config_.value("hidden_size", vision_config_.value("embed_dim", 1280)); + patch_size_ = vision_config_.value("patch_size", 14); + + INFINICORE_NN_MODULE_INIT(patch_embed, vision_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), vision_config_, dtype_, device)); + } + INFINICORE_NN_MODULE_INIT(ln, hidden_size_, 1e-6, dtype_, device); +} + +infinicore::Tensor Ernie4_5VisionModel::normalize_images_(const infinicore::Tensor &images) const { + if (images->dtype() != infinicore::DataType::U8 && images->dtype() != infinicore::DataType::BYTE) { + return images; + } + + ASSERT(images->ndim() == 2); + const size_t patch_width = images->shape()[1]; + const size_t channel_patch = patch_size_ * patch_size_; + ASSERT_EQ(patch_width, channel_patch * 3); + + auto images_cpu = images->to(infinicore::Device::cpu()); + auto out_cpu = infinicore::Tensor::empty(images->shape(), dtype_, infinicore::Device::cpu()); + const auto *src = reinterpret_cast(images_cpu->data()); + void *dst = out_cpu->data(); + + for (size_t i = 0; i < images->shape()[0]; ++i) { + for (size_t j = 0; j < patch_width; ++j) { + const size_t channel = j / channel_patch; + const float value = (static_cast(src[i * patch_width + j]) * kRescaleFactor - kImageMean[channel]) / kImageStd[channel]; + write_float(dst, i * patch_width + j, dtype_, value); + } + } + return out_cpu->to(images->device()); +} + +infinicore::Tensor Ernie4_5VisionModel::vision_grid_thw_(const infinicore::Tensor &grid_thw) const { + const auto grid = read_grid(grid_thw); + size_t rows = 0; + for (size_t i = 0; i < grid.size(); i += 3) { + rows += static_cast(grid[i]); + } + + auto out_cpu = infinicore::Tensor::empty({rows, 3}, infinicore::DataType::I64, infinicore::Device::cpu()); + auto *out = reinterpret_cast(out_cpu->data()); + size_t row = 0; + for (size_t i = 0; i < grid.size(); i += 3) { + const int64_t t = grid[i]; + const int64_t h = grid[i + 1]; + const int64_t w = grid[i + 2]; + for (int64_t ti = 0; ti < t; ++ti) { + out[row * 3] = 1; + out[row * 3 + 1] = h; + out[row * 3 + 2] = w; + ++row; + } + } + return out_cpu->to(grid_thw->device()); +} + +infinicore::Tensor Ernie4_5VisionModel::forward(const infinicore::Tensor &images, + const infinicore::Tensor &grid_thw) const { + auto vision_grid = vision_grid_thw_(grid_thw); + auto hidden_states = normalize_images_(images); + hidden_states = patch_embed_->forward(hidden_states); + + for (const auto &block : blocks_) { + hidden_states = block->forward(hidden_states, vision_grid); + } + return ln_->forward(hidden_states); +} + +} // namespace infinilm::models::ernie4_5_moe_vl diff --git a/csrc/models/ernie4_5_moe_vl/ernie4_5_vision.hpp b/csrc/models/ernie4_5_moe_vl/ernie4_5_vision.hpp new file mode 100644 index 000000000..a3c3153b5 --- /dev/null +++ b/csrc/models/ernie4_5_moe_vl/ernie4_5_vision.hpp @@ -0,0 +1,111 @@ +#pragma once + +#include "../../config/model_config.hpp" +#include "../../layers/linear/linear.hpp" +#include "infinicore/nn/layer_norm.hpp" +#include "infinicore/nn/module.hpp" +#include "infinicore/tensor.hpp" + +#include +#include + +namespace infinilm::models::ernie4_5_moe_vl { + +class Ernie4_5VisionPatchEmbed : public infinicore::nn::Module { +public: + Ernie4_5VisionPatchEmbed(const nlohmann::json &vision_config, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + +private: + size_t patch_size_{14}; + size_t in_channels_{3}; + size_t embed_dim_{1280}; + + INFINICORE_NN_MODULE(infinilm::nn::Linear, proj); +}; + +class Ernie4_5VisionAttention : public infinicore::nn::Module { +public: + Ernie4_5VisionAttention(const nlohmann::json &vision_config, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &grid_thw) const; + +private: + infinicore::Tensor apply_rotary_pos_emb_(const infinicore::Tensor &tensor, + const infinicore::Tensor &grid_thw) const; + infinicore::Tensor segmented_attention_(const infinicore::Tensor &q, + const infinicore::Tensor &k, + const infinicore::Tensor &v, + const infinicore::Tensor &grid_thw) const; + + size_t embed_dim_{1280}; + size_t num_heads_{16}; + size_t head_dim_{80}; + size_t spatial_merge_size_{2}; + float scale_{1.0f}; + + INFINICORE_NN_MODULE(infinilm::nn::Linear, qkv); + INFINICORE_NN_MODULE(infinilm::nn::Linear, proj); +}; + +class Ernie4_5VisionMLP : public infinicore::nn::Module { +public: + Ernie4_5VisionMLP(const nlohmann::json &vision_config, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + +private: + std::string hidden_act_; + + INFINICORE_NN_MODULE(infinilm::nn::Linear, fc1); + INFINICORE_NN_MODULE(infinilm::nn::Linear, fc2); +}; + +class Ernie4_5VisionBlock : public infinicore::nn::Module { +public: + Ernie4_5VisionBlock(const nlohmann::json &vision_config, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &grid_thw) const; + +private: + INFINICORE_NN_MODULE(infinicore::nn::LayerNorm, norm1); + INFINICORE_NN_MODULE(Ernie4_5VisionAttention, attn); + INFINICORE_NN_MODULE(infinicore::nn::LayerNorm, norm2); + INFINICORE_NN_MODULE(Ernie4_5VisionMLP, mlp); +}; + +class Ernie4_5VisionModel : public infinicore::nn::Module { +public: + Ernie4_5VisionModel(std::shared_ptr model_config, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &images, + const infinicore::Tensor &grid_thw) const; + +private: + infinicore::Tensor normalize_images_(const infinicore::Tensor &images) const; + infinicore::Tensor vision_grid_thw_(const infinicore::Tensor &grid_thw) const; + + nlohmann::json vision_config_; + infinicore::DataType dtype_{infinicore::DataType::BF16}; + size_t depth_{32}; + size_t hidden_size_{1280}; + size_t patch_size_{14}; + + INFINICORE_NN_MODULE(Ernie4_5VisionPatchEmbed, patch_embed); + INFINICORE_NN_MODULE_VEC(Ernie4_5VisionBlock, blocks); + INFINICORE_NN_MODULE(infinicore::nn::LayerNorm, ln); +}; + +} // namespace infinilm::models::ernie4_5_moe_vl diff --git a/csrc/models/infinilm_model.hpp b/csrc/models/infinilm_model.hpp index ac994fd6d..510d6efdc 100644 --- a/csrc/models/infinilm_model.hpp +++ b/csrc/models/infinilm_model.hpp @@ -22,6 +22,8 @@ class InfinilmModel : public infinicore::nn::Module { std::optional input_ids; /// Position IDs tensor of shape `[batch, seq_len]` or `[seq_len]`. std::optional position_ids; + /// Token modality IDs. ERNIE-VL uses 0 for text and non-zero for vision. + std::optional token_type_ids; /// Past Lengths of cached sequence for each request, of shape `[num_requests]`. std::optional past_sequence_lengths; /// ToTal Lengths for each request sequence, of shape `[num_requests]`. @@ -49,6 +51,10 @@ class InfinilmModel : public infinicore::nn::Module { std::optional> tgt_sizes; /// Qwen-style image grids. Vector of tensors shape: [3] with temporal, height, width. std::optional> image_grid_thw; + /// ERNIE-VL image/video grids. Each tensor has shape [num_items, 3]. + std::optional> grid_thw; + /// ERNIE-VL item type IDs, where 0=image and 1=video. + std::optional> image_type_ids; /// req_id for each pixel_values among a batch. std::optional> image_req_ids; /// Flattened [start, end) visual token ranges in the packed language sequence. diff --git a/csrc/pybind11/engine/engine.hpp b/csrc/pybind11/engine/engine.hpp index b73ce06ae..2762efe3d 100644 --- a/csrc/pybind11/engine/engine.hpp +++ b/csrc/pybind11/engine/engine.hpp @@ -134,6 +134,7 @@ inline void bind_infer_engine(py::module &m) { py::init([]( std::optional input_ids, std::optional position_ids, + std::optional token_type_ids, std::optional past_sequence_lengths, std::optional total_sequence_lengths, std::optional input_offsets, @@ -146,6 +147,8 @@ inline void bind_infer_engine(py::module &m) { std::optional> image_bound, std::optional> tgt_sizes, std::optional> image_grid_thw, + std::optional> grid_thw, + std::optional> image_type_ids, std::optional> image_req_ids, std::optional> visual_token_ranges, std::optional target_hidden_states, @@ -154,6 +157,7 @@ inline void bind_infer_engine(py::module &m) { InferEngine::Input input{ std::move(input_ids), std::move(position_ids), + std::move(token_type_ids), std::move(past_sequence_lengths), std::move(total_sequence_lengths), std::move(input_offsets), @@ -166,6 +170,8 @@ inline void bind_infer_engine(py::module &m) { std::move(image_bound), std::move(tgt_sizes), std::move(image_grid_thw), + std::move(grid_thw), + std::move(image_type_ids), std::move(image_req_ids), std::move(visual_token_ranges), std::move(target_hidden_states), @@ -205,6 +211,7 @@ inline void bind_infer_engine(py::module &m) { }), py::arg("input_ids") = std::nullopt, py::arg("position_ids") = std::nullopt, + py::arg("token_type_ids") = std::nullopt, py::arg("past_sequence_lengths") = std::nullopt, py::arg("total_sequence_lengths") = std::nullopt, py::arg("input_offsets") = std::nullopt, @@ -217,12 +224,15 @@ inline void bind_infer_engine(py::module &m) { py::arg("image_bound") = std::nullopt, py::arg("tgt_sizes") = std::nullopt, py::arg("image_grid_thw") = std::nullopt, + py::arg("grid_thw") = std::nullopt, + py::arg("image_type_ids") = std::nullopt, py::arg("image_req_ids") = std::nullopt, py::arg("visual_token_ranges") = std::nullopt, py::arg("target_hidden_states") = std::nullopt, py::arg("sample_all_positions") = false) .def_readwrite("input_ids", &InferEngine::Input::input_ids) .def_readwrite("position_ids", &InferEngine::Input::position_ids) + .def_readwrite("token_type_ids", &InferEngine::Input::token_type_ids) .def_readwrite("past_sequence_lengths", &InferEngine::Input::past_sequence_lengths) .def_readwrite("total_sequence_lengths", &InferEngine::Input::total_sequence_lengths) .def_readwrite("input_offsets", &InferEngine::Input::input_offsets) @@ -235,6 +245,8 @@ inline void bind_infer_engine(py::module &m) { .def_readwrite("image_bound", &InferEngine::Input::image_bound) .def_readwrite("tgt_sizes", &InferEngine::Input::tgt_sizes) .def_readwrite("image_grid_thw", &InferEngine::Input::image_grid_thw) + .def_readwrite("grid_thw", &InferEngine::Input::grid_thw) + .def_readwrite("image_type_ids", &InferEngine::Input::image_type_ids) .def_readwrite("image_req_ids", &InferEngine::Input::image_req_ids) .def_readwrite("visual_token_ranges", &InferEngine::Input::visual_token_ranges) .def_readwrite("target_hidden_states", &InferEngine::Input::target_hidden_states) diff --git a/examples/ernie_vl_correctness.py b/examples/ernie_vl_correctness.py new file mode 100644 index 000000000..6d59a81b3 --- /dev/null +++ b/examples/ernie_vl_correctness.py @@ -0,0 +1,590 @@ +import argparse +import gc +import json +import os +import subprocess +import sys +import tempfile +import time + +import infinicore +import numpy as np +import torch +from infinilm.cache import StaticKVCacheConfig +from infinilm.distributed import DistConfig +from infinilm.infer_engine import GenerationConfig, InferEngine +from infinilm.modeling_utils import load_model_state_dict_by_file +from transformers import AutoProcessor, AutoTokenizer + +EXPECTED_IDS = { + "text": [ + 3843, + 5971, + 94036, + 31282, + 5502, + 965, + 93956, + 5119, + 94111, + 6385, + 5188, + 1555, + 94035, + 8217, + 94110, + 586, + ], + "image": [38020, 432, 93938, 1981, 93968, 93927, 1505, 93937], + "video": [3843, 1510, 1386, 94001, 5434, 1187, 6350, 93956], +} + + +def build_video(): + video = np.zeros((2, 32, 32, 3), dtype=np.uint8) + video[0, :, :, 0] = 40 + video[0, :, :, 1] = 120 + video[0, :, :, 2] = 200 + video[1, :, :, 0] = 120 + video[1, :, :, 1] = 120 + video[1, :, :, 2] = 140 + return video + + +def apply_template(tokenizer, messages): + return tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + + +def build_inputs(case, model_path, processor, tokenizer, image_path): + if case == "text": + messages = [ + { + "role": "user", + "content": [{"type": "text", "text": "用一句话介绍你自己。"}], + } + ] + text = apply_template(tokenizer, messages) + inputs = processor(text=[text], return_tensors="pt") + return text, inputs + + if case == "image": + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this image briefly."}, + { + "type": "image_url", + "image_url": { + "url": image_path, + "image_width": 224, + "image_height": 224, + }, + }, + ], + } + ] + text = apply_template(tokenizer, messages) + image_inputs, video_inputs = processor.process_vision_info(messages) + inputs = processor( + text=[text], + images=image_inputs, + videos=video_inputs, + padding=True, + return_tensors="pt", + ) + return text, inputs + + if case == "video": + messages = [ + { + "role": "user", + "content": [ + {"type": "video", "video": "dummy_numpy_2x32x32x3"}, + {"type": "text", "text": "Describe the video briefly."}, + ], + } + ] + text = apply_template(tokenizer, messages) + inputs = processor(text=text, videos=[build_video()], return_tensors="pt") + return text, inputs + + raise ValueError(f"unknown case: {case}") + + +def infini_from_torch(tensor): + return infinicore.from_torch(tensor.contiguous()) + + +def build_hf_references(args, model_path, processor, tokenizer, image_path): + child_code = r""" +import json +import sys +import time + +import numpy as np +import torch +from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer + + +def build_video(): + video = np.zeros((2, 32, 32, 3), dtype=np.uint8) + video[0, :, :, 0] = 40 + video[0, :, :, 1] = 120 + video[0, :, :, 2] = 200 + video[1, :, :, 0] = 120 + video[1, :, :, 1] = 120 + video[1, :, :, 2] = 140 + return video + + +def apply_template(tokenizer, messages): + return tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + + +def build_inputs(case, processor, tokenizer, image_path): + if case == "text": + messages = [ + { + "role": "user", + "content": [{"type": "text", "text": "用一句话介绍你自己。"}], + } + ] + text = apply_template(tokenizer, messages) + inputs = processor(text=[text], return_tensors="pt") + return text, inputs + + if case == "image": + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this image briefly."}, + { + "type": "image_url", + "image_url": { + "url": image_path, + "image_width": 224, + "image_height": 224, + }, + }, + ], + } + ] + text = apply_template(tokenizer, messages) + image_inputs, video_inputs = processor.process_vision_info(messages) + inputs = processor( + text=[text], + images=image_inputs, + videos=video_inputs, + padding=True, + return_tensors="pt", + ) + return text, inputs + + if case == "video": + messages = [ + { + "role": "user", + "content": [ + {"type": "video", "video": "dummy_numpy_2x32x32x3"}, + {"type": "text", "text": "Describe the video briefly."}, + ], + } + ] + text = apply_template(tokenizer, messages) + inputs = processor(text=text, videos=[build_video()], return_tensors="pt") + return text, inputs + + raise ValueError(f"unknown case: {case}") + + +def torch_dtype_from_name(name): + if name == "bf16": + return torch.bfloat16 + if name == "fp16": + return torch.float16 + if name == "fp32": + return torch.float32 + raise ValueError(f"unsupported torch dtype: {name}") + + +def move_inputs_to_device(inputs, device): + moved = {} + for key, value in inputs.items(): + if torch.is_tensor(value): + moved[key] = value.to(device) + else: + moved[key] = value + return moved + + +def fix_hf_meta_helper_tensors(model): + patched = [] + vision_model = getattr(model, "vision_model", None) or getattr(model, "visual", None) + rotary = getattr(vision_model, "rotary_pos_emb", None) if vision_model is not None else None + inv_freq = getattr(rotary, "inv_freq", None) if rotary is not None else None + if inv_freq is not None and getattr(inv_freq, "device", None).type == "meta": + dim = int(inv_freq.numel() * 2) + theta = float(getattr(rotary, "theta", 10000.0)) + rotary.inv_freq = 1.0 / theta ** ( + torch.arange(start=0, end=dim, step=2, dtype=torch.float32) / dim + ) + patched.append("vision_rotary_inv_freq") + + for module in model.modules(): + experts_type_ids = getattr(module, "experts_type_ids", None) + if experts_type_ids is None or getattr(experts_type_ids, "device", None).type != "meta": + continue + config = getattr(module, "config", None) + moe_num_experts = getattr(config, "moe_num_experts", None) + if not isinstance(moe_num_experts, (list, tuple)): + continue + rebuilt = torch.zeros([sum(moe_num_experts)], dtype=torch.int64) + offset = 0 + masks = [] + for idx, expert_num in enumerate(moe_num_experts): + rebuilt[offset : offset + expert_num] = idx + offset += expert_num + module.experts_type_ids = rebuilt + for idx, _ in enumerate(moe_num_experts): + masks.append(module.experts_type_ids == idx) + module.experts_type_mask = masks + patched.append("experts_type_ids") + return patched + + +config_path, output_path = sys.argv[1], sys.argv[2] +with open(config_path, "r", encoding="utf-8") as f: + cfg = json.load(f) + +processor = AutoProcessor.from_pretrained( + cfg["model_path"], + trust_remote_code=True, + video_min_pixels=3136, + video_max_pixels=3136, +) +tokenizer = AutoTokenizer.from_pretrained(cfg["model_path"], trust_remote_code=True) +model = AutoModelForCausalLM.from_pretrained( + cfg["model_path"], + device_map=cfg["hf_device_map"], + torch_dtype=torch_dtype_from_name(cfg["hf_torch_dtype"]), + trust_remote_code=True, + low_cpu_mem_usage=True, + use_flash_attention=cfg["hf_use_flash_attention"], +) +model.eval() +patched_meta = fix_hf_meta_helper_tensors(model) +if patched_meta: + print(f"patched HF meta helper tensors: {patched_meta}", flush=True) +if hasattr(model, "add_image_preprocess"): + model.add_image_preprocess(processor) + +references = {} +details = [] +for case in cfg["cases"]: + text, inputs = build_inputs(case, processor, tokenizer, cfg["image_path"]) + prompt_len = int(inputs["input_ids"].shape[-1]) + model_inputs = move_inputs_to_device(inputs, model.device) + max_new_tokens = cfg["hf_max_new_tokens_by_case"][case] + start = time.time() + with torch.no_grad(): + generated_ids = model.generate( + inputs=model_inputs["input_ids"], + **model_inputs, + max_new_tokens=max_new_tokens, + do_sample=False, + use_cache=cfg["hf_use_cache"], + ) + elapsed = time.time() - start + token_ids = generated_ids[0][prompt_len:].detach().cpu().tolist() + item = { + "hf_reference": case, + "prompt_len": prompt_len, + "new_token_ids": token_ids, + "output_text": tokenizer.decode(token_ids, skip_special_tokens=True), + "elapsed_sec": elapsed, + } + references[case] = token_ids + details.append(item) + print( + json.dumps( + { + **item, + "elapsed_sec": round(elapsed, 3), + }, + ensure_ascii=False, + ), + flush=True, + ) + +with open(output_path, "w", encoding="utf-8") as f: + json.dump({"references": references, "details": details}, f, ensure_ascii=False) +""" + + config = { + "model_path": model_path, + "image_path": image_path, + "cases": args.cases, + "hf_device_map": args.hf_device_map, + "hf_torch_dtype": args.hf_torch_dtype, + "hf_use_cache": args.hf_use_cache, + "hf_use_flash_attention": args.hf_use_flash_attention, + "hf_max_new_tokens_by_case": { + case: args.hf_max_new_tokens or len(EXPECTED_IDS[case]) + for case in args.cases + }, + } + + with tempfile.TemporaryDirectory(prefix="ernie_hf_reference_") as tmp_dir: + script_path = os.path.join(tmp_dir, "run_hf_reference.py") + config_path = os.path.join(tmp_dir, "config.json") + output_path = os.path.join(tmp_dir, "reference.json") + with open(script_path, "w", encoding="utf-8") as f: + f.write(child_code) + with open(config_path, "w", encoding="utf-8") as f: + json.dump(config, f, ensure_ascii=False) + + subprocess.run( + [sys.executable, script_path, config_path, output_path], check=True + ) + with open(output_path, "r", encoding="utf-8") as f: + result = json.load(f) + + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + return result["references"] + + +def summarize_inputs(inputs): + summary = {} + for key, value in inputs.items(): + if torch.is_tensor(value): + item = {"shape": list(value.shape), "dtype": str(value.dtype)} + if key in {"grid_thw", "image_type_ids"}: + item["value"] = value.detach().cpu().tolist() + summary[key] = item + else: + summary[key] = str(type(value)) + return summary + + +def parse_tp_devices(value): + device_ids = [] + for item in value.split(","): + item = item.strip() + if not item: + continue + try: + device_ids.append(int(item)) + except ValueError as exc: + raise argparse.ArgumentTypeError( + f"invalid device id in --tp-devices: {item!r}" + ) from exc + + if not device_ids: + raise argparse.ArgumentTypeError("--tp-devices must contain at least one id") + if any(device_id < 0 for device_id in device_ids): + raise argparse.ArgumentTypeError("--tp-devices cannot contain negative ids") + return device_ids + + +def build_dist_config(args): + if args.tp < 1: + raise ValueError("--tp must be >= 1") + + if args.tp_devices is not None: + tp_device_ids = args.tp_devices + if args.tp != 1 and args.tp != len(tp_device_ids): + raise ValueError( + f"--tp ({args.tp}) must match --tp-devices length ({len(tp_device_ids)})" + ) + dist_config = DistConfig(tp_device_ids=tp_device_ids) + else: + tp_device_ids = list(range(args.tp)) + dist_config = DistConfig(args.tp) + + if args.device == "cuda": + device_count = torch.cuda.device_count() + if device_count < len(tp_device_ids): + raise ValueError( + f"tensor parallel needs {len(tp_device_ids)} CUDA device(s), " + f"but torch sees {device_count}" + ) + invalid_ids = [ + device_id for device_id in tp_device_ids if device_id >= device_count + ] + if invalid_ids: + raise ValueError( + f"--tp-devices contains unavailable CUDA device id(s): {invalid_ids}; " + f"torch sees device ids 0..{device_count - 1}" + ) + + return dist_config, tp_device_ids + + +def run_case(engine, processor, tokenizer, case, text, inputs, expected_ids): + max_new_tokens = len(expected_ids) + kwargs = {} + if inputs.get("position_ids") is not None: + kwargs["position_ids"] = infini_from_torch( + inputs["position_ids"].to(torch.int64) + ) + if inputs.get("token_type_ids") is not None: + kwargs["token_type_ids"] = infini_from_torch( + inputs["token_type_ids"].to(torch.int64) + ) + if inputs.get("images") is not None: + kwargs["images"] = infini_from_torch(inputs["images"].contiguous()) + kwargs["grid_thw"] = infini_from_torch(inputs["grid_thw"].to(torch.int64)) + kwargs["image_type_ids"] = infini_from_torch( + inputs["image_type_ids"].to(torch.int64) + ) + + input_ids = infini_from_torch(inputs["input_ids"].to(torch.int64)) + start = time.time() + output = engine.generate( + input_ids, + GenerationConfig( + max_new_tokens=max_new_tokens, + temperature=1.0, + top_k=1, + top_p=1.0, + ), + **kwargs, + ) + elapsed = time.time() - start + + token_ids = [] + for tensor in output: + token_ids.extend(np.array(tensor.to_numpy()).reshape(-1).astype(int).tolist()) + + return { + "case": case, + "prompt": text, + "prompt_len": int(inputs["input_ids"].shape[-1]), + "input_summary": summarize_inputs(inputs), + "expected_token_ids": expected_ids, + "new_token_ids": token_ids, + "match_expected": token_ids == expected_ids, + "output_text": tokenizer.decode(token_ids, skip_special_tokens=True), + "elapsed_sec": elapsed, + } + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--model", required=True) + parser.add_argument("--device", default="cuda") + parser.add_argument("--tp", type=int, default=1) + parser.add_argument("--tp-devices", type=parse_tp_devices, default=None) + parser.add_argument( + "--reference-mode", + choices=["expected", "hf"], + default="expected", + help="Use baked HF token baselines or run the HF model live first.", + ) + parser.add_argument("--hf-device-map", default="auto") + parser.add_argument( + "--hf-torch-dtype", choices=["bf16", "fp16", "fp32"], default="bf16" + ) + parser.add_argument("--hf-max-new-tokens", type=int, default=None) + parser.add_argument("--hf-use-cache", action="store_true") + parser.add_argument("--hf-use-flash-attention", action="store_true") + parser.add_argument("--cases", nargs="+", default=["text", "image", "video"]) + parser.add_argument("--image", default=None) + parser.add_argument("--max-cache-len", type=int, default=512) + parser.add_argument("--output-json", default=None) + return parser.parse_args() + + +def main(): + args = parse_args() + model_path = os.path.expanduser(args.model) + image_path = args.image or os.path.join(model_path, "benchmark.jpg") + dist_config, tp_device_ids = build_dist_config(args) + run_config = { + "device": args.device, + "tp_device_ids": tp_device_ids, + "dist_config": str(dist_config), + "max_cache_len": args.max_cache_len, + "reference_mode": args.reference_mode, + } + print(json.dumps({"run_config": run_config}, ensure_ascii=False), flush=True) + + processor = AutoProcessor.from_pretrained( + model_path, + trust_remote_code=True, + video_min_pixels=3136, + video_max_pixels=3136, + ) + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + for case in args.cases: + if case not in EXPECTED_IDS: + raise ValueError(f"unsupported case: {case}") + + if args.reference_mode == "hf": + expected_ids_by_case = build_hf_references( + args, model_path, processor, tokenizer, image_path + ) + else: + expected_ids_by_case = {case: EXPECTED_IDS[case] for case in args.cases} + + engine = InferEngine( + model_path, + device=infinicore.device(args.device, 0), + distributed_config=dist_config, + cache_config=StaticKVCacheConfig( + max_batch_size=1, max_cache_len=args.max_cache_len + ), + attention_backend="default", + ) + load_model_state_dict_by_file(engine, model_path, dtype=engine.dtype) + + results = [] + for case in args.cases: + text, inputs = build_inputs(case, model_path, processor, tokenizer, image_path) + engine.reset_cache( + StaticKVCacheConfig(max_batch_size=1, max_cache_len=args.max_cache_len) + ) + result = run_case( + engine, + processor, + tokenizer, + case, + text, + inputs, + expected_ids_by_case[case], + ) + results.append(result) + print( + json.dumps( + { + "case": result["case"], + "prompt_len": result["prompt_len"], + "new_token_ids": result["new_token_ids"], + "match_expected": result["match_expected"], + "output_text": result["output_text"], + "elapsed_sec": round(result["elapsed_sec"], 3), + }, + ensure_ascii=False, + ), + flush=True, + ) + + ok = all(item["match_expected"] for item in results) + report = {"ok": ok, "run_config": run_config, "results": results} + if args.output_json: + with open(args.output_json, "w", encoding="utf-8") as f: + json.dump(report, f, ensure_ascii=False, indent=2) + if not ok: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/examples/ernie_vl_llm_smoke.py b/examples/ernie_vl_llm_smoke.py new file mode 100644 index 000000000..7ab515f10 --- /dev/null +++ b/examples/ernie_vl_llm_smoke.py @@ -0,0 +1,161 @@ +import argparse +import json +import os +import time + +from infinilm.llm.llm import LLM +from infinilm.llm.sampling_params import SamplingParams +from PIL import Image + +EXPECTED_IDS = { + "static": [ + 3843, + 1510, + 1386, + 5434, + 38225, + 6554, + 93977, + 5119, + 94101, + 6554, + 2293, + 94035, + 93986, + 96101, + 94552, + 94397, + ], + "paged": [3843, 1510, 1386, 5434, 38225, 6554, 93977, 5119], +} + + +def prepare_image(image_path, output_path): + image = Image.open(image_path).convert("RGB").resize((224, 224)) + image.save(output_path) + return output_path + + +def build_messages(image_path): + return [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": image_path}}, + {"type": "text", "text": "请用一句中文简短描述这张图片。"}, + ], + } + ] + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--model", required=True) + parser.add_argument("--image", default=None) + parser.add_argument("--cache-type", choices=["static", "paged"], required=True) + parser.add_argument("--device", default="cuda") + parser.add_argument("--tp", type=int, default=1) + parser.add_argument("--max-cache-len", type=int, default=512) + parser.add_argument("--num-blocks", type=int, default=64) + parser.add_argument("--block-size", type=int, default=16) + parser.add_argument("--max-tokens", type=int, default=None) + parser.add_argument("--output-json", default=None) + return parser.parse_args() + + +def main(): + args = parse_args() + model_path = os.path.expanduser(args.model) + image_path = args.image or os.path.join(model_path, "benchmark.jpg") + resized_image = prepare_image(image_path, "/tmp/ernie_vl_llm_smoke_224.jpg") + expected_ids = EXPECTED_IDS[args.cache_type] + max_tokens = args.max_tokens or len(expected_ids) + + run_config = { + "cache_type": args.cache_type, + "device": args.device, + "tp": args.tp, + "max_cache_len": args.max_cache_len, + "num_blocks": args.num_blocks, + "block_size": args.block_size, + "max_tokens": max_tokens, + "image": resized_image, + } + print(json.dumps({"run_config": run_config}, ensure_ascii=False), flush=True) + + start = time.time() + llm = LLM( + model_path=model_path, + device=args.device, + dtype="bfloat16", + tensor_parallel_size=args.tp, + cache_type=args.cache_type, + max_batch_size=1, + max_tokens=max_tokens, + num_blocks=args.num_blocks, + block_size=args.block_size, + max_cache_len=args.max_cache_len, + temperature=1.0, + top_k=1, + top_p=1.0, + enable_graph=False, + attn_backend="default", + ) + init_sec = time.time() - start + + params = SamplingParams( + max_tokens=max_tokens, + temperature=1.0, + top_k=1, + top_p=1.0, + ignore_eos=True, + ) + gen_start = time.time() + outputs = llm.generate( + messages=build_messages(resized_image), + sampling_params=params, + use_tqdm=False, + ) + gen_sec = time.time() - gen_start + if len(outputs) != 1 or len(outputs[0].outputs) != 1: + raise RuntimeError(f"unexpected output shape: {outputs!r}") + + request = outputs[0] + completion = request.outputs[0] + token_ids = completion.token_ids or [] + result = { + "ok": token_ids == expected_ids, + "run_config": run_config, + "prompt_len": len(request.prompt_token_ids or []), + "expected_token_ids": expected_ids, + "new_token_ids": token_ids, + "output_text": completion.text, + "finish_reason": str(completion.finish_reason), + "init_sec": init_sec, + "elapsed_sec": gen_sec, + } + print( + json.dumps( + { + "cache_type": args.cache_type, + "prompt_len": result["prompt_len"], + "new_token_ids": result["new_token_ids"], + "match_expected": result["ok"], + "output_text": result["output_text"], + "init_sec": round(init_sec, 3), + "elapsed_sec": round(gen_sec, 3), + }, + ensure_ascii=False, + ), + flush=True, + ) + + if args.output_json: + with open(args.output_json, "w", encoding="utf-8") as f: + json.dump(result, f, ensure_ascii=False, indent=2) + if not result["ok"]: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/examples/ernie_vl_mmmu_smoke.py b/examples/ernie_vl_mmmu_smoke.py new file mode 100644 index 000000000..c69489500 --- /dev/null +++ b/examples/ernie_vl_mmmu_smoke.py @@ -0,0 +1,528 @@ +import argparse +import ast +import csv +import json +import os +import re +import tempfile +import time + +import numpy as np +import torch +from datasets import get_dataset_config_names, load_dataset +from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--model", required=True) + parser.add_argument("--backend", choices=["hf", "cpp"], default="hf") + parser.add_argument("--subjects", default="Accounting") + parser.add_argument("--split", choices=["dev", "validation"], default="validation") + parser.add_argument("--num-samples", type=int, default=2) + parser.add_argument("--max-new-tokens", type=int, default=64) + parser.add_argument("--image-size", type=int, default=224) + parser.add_argument( + "--prompt-style", + choices=["official", "direct", "ernie"], + default="official", + help="official follows MMMU's default prompt format; direct/ernie are experimental.", + ) + parser.add_argument("--max-cache-len", type=int, default=2048) + parser.add_argument("--device", default="cuda") + parser.add_argument("--tp", type=int, default=1) + parser.add_argument("--tp-devices", type=parse_tp_devices, default=None) + parser.add_argument( + "--torch-dtype", choices=["bf16", "fp16", "fp32"], default="bf16" + ) + parser.add_argument("--output-csv", required=True) + parser.add_argument("--output-json", default=None) + return parser.parse_args() + + +def parse_tp_devices(value): + if value is None or value == "": + return None + return [int(item.strip()) for item in value.split(",") if item.strip()] + + +def torch_dtype(name): + if name == "bf16": + return torch.bfloat16 + if name == "fp16": + return torch.float16 + return torch.float32 + + +def selected_subjects(subjects_arg): + if subjects_arg == "all": + return get_dataset_config_names("MMMU/MMMU") + return [item.strip() for item in subjects_arg.split(",") if item.strip()] + + +def parse_options(raw): + if isinstance(raw, list): + return raw + try: + value = ast.literal_eval(raw) + except Exception: + value = raw + if isinstance(value, list): + return [str(item) for item in value] + return [part.strip() for part in str(value).split("\n") if part.strip()] + + +def normalize_question(question): + text = question + for idx in range(1, 8): + text = text.replace(f"", "") + return " ".join(text.split()) + + +def extract_answer(text, choices=("A", "B", "C", "D"), index2ans=None): + output_upper = text.upper().strip() + answer_matches = re.findall(r"ANSWER\s*[:：]\s*([ABCD])\b", output_upper) + if answer_matches: + return answer_matches[-1] + final_matches = re.findall(r"FINAL\s+ANSWER\s*[:：]\s*([ABCD])\b", output_upper) + if final_matches: + return final_matches[-1] + response = " " + output_upper.strip(",.!?;:'") + " " + candidates = [] + positions = [] + for choice in choices: + for pattern in (f"({choice})", f" {choice} "): + pos = response.rfind(pattern) + if pos >= 0: + candidates.append(choice) + positions.append(pos) + if not candidates and index2ans and len(response.split()) > 5: + for choice, answer_text in index2ans.items(): + pos = response.lower().rfind(str(answer_text).lower()) + if pos >= 0: + candidates.append(choice) + positions.append(pos) + if candidates: + return candidates[int(np.argmax(positions))] + return "" + + +def extract_debug_answers(text, choices=("A", "B", "C", "D")): + output_upper = text.upper() + choice_class = "".join(choices) + boxed_matches = re.findall( + rf"(?:BOXED|\\BOXED)\s*\{{\s*([{choice_class}])\s*\}}", output_upper + ) + explicit_matches = re.findall( + rf"(?:ANSWER\s*(?:IS|:|：)|FINAL\s+ANSWER\s*(?:IS|:|：))\s*\(?([{choice_class}])\)?", + output_upper, + ) + return { + "boxed_answer": boxed_matches[-1] if boxed_matches else "", + "explicit_answer": explicit_matches[-1] if explicit_matches else "", + } + + +def build_messages(row, image_paths, image_size, prompt_style): + options = parse_options(row["options"]) + question = normalize_question(row["question"]) + if prompt_style == "official": + choices_text = "".join( + f"({chr(65 + idx)}) {choice}\n" for idx, choice in enumerate(options) + ) + text = ( + f"{question} {choices_text}" + "Answer with the option's letter from the given choices directly." + ) + elif prompt_style == "ernie": + choices_text = "\n".join( + f"{chr(65 + idx)}. {choice}" for idx, choice in enumerate(options) + ) + text = ( + "Please solve this MMMU multiple-choice problem using the provided image(s). " + "You may reason internally, but the final response must contain the answer " + "letter in the form 'Answer: A', 'Answer: B', 'Answer: C', or 'Answer: D'.\n\n" + f"Question: {question}\n{choices_text}" + ) + else: + choices_text = "\n".join( + f"{chr(65 + idx)}. {choice}" for idx, choice in enumerate(options) + ) + text = ( + "Answer the multiple-choice question. Respond with only the letter " + "A, B, C, or D.\n\n" + f"Question: {question}\n{choices_text}" + ) + content = [] + for path in image_paths: + content.append( + { + "type": "image_url", + "image_url": { + "url": path, + "image_width": image_size, + "image_height": image_size, + }, + } + ) + content.append({"type": "text", "text": text}) + return [{"role": "user", "content": content}] + + +def save_images(row, tmp_dir): + paths = [] + for idx in range(1, 8): + image = row.get(f"image_{idx}") + if image is None: + continue + path = os.path.join(tmp_dir, f"image_{idx}.png") + image.save(path) + paths.append(path) + return paths + + +def build_inputs(processor, tokenizer, row, tmp_dir, image_size, prompt_style): + image_paths = save_images(row, tmp_dir) + messages = build_messages(row, image_paths, image_size, prompt_style) + prompt = tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + image_inputs, video_inputs = processor.process_vision_info(messages) + inputs = processor( + text=[prompt], + images=image_inputs, + videos=video_inputs, + padding=True, + return_tensors="pt", + ) + return prompt, inputs, len(image_paths) + + +def tensor_to_infini(tensor): + import infinicore + + return infinicore.from_torch(tensor.contiguous()) + + +def fix_hf_meta_helper_tensors(model): + patched = [] + vision_model = getattr(model, "vision_model", None) or getattr( + model, "visual", None + ) + rotary = ( + getattr(vision_model, "rotary_pos_emb", None) + if vision_model is not None + else None + ) + inv_freq = getattr(rotary, "inv_freq", None) if rotary is not None else None + if inv_freq is not None and getattr(inv_freq, "device", None).type == "meta": + dim = int(inv_freq.numel() * 2) + theta = float(getattr(rotary, "theta", 10000.0)) + rotary.inv_freq = 1.0 / theta ** ( + torch.arange(start=0, end=dim, step=2, dtype=torch.float32) / dim + ) + patched.append("vision_rotary_inv_freq") + + for module in model.modules(): + experts_type_ids = getattr(module, "experts_type_ids", None) + if ( + experts_type_ids is None + or getattr(experts_type_ids, "device", None).type != "meta" + ): + continue + config = getattr(module, "config", None) + moe_num_experts = getattr(config, "moe_num_experts", None) + if not isinstance(moe_num_experts, (list, tuple)): + continue + rebuilt = torch.zeros([sum(moe_num_experts)], dtype=torch.int64) + offset = 0 + masks = [] + for idx, expert_num in enumerate(moe_num_experts): + rebuilt[offset : offset + expert_num] = idx + offset += expert_num + module.experts_type_ids = rebuilt + for idx, _ in enumerate(moe_num_experts): + masks.append(module.experts_type_ids == idx) + module.experts_type_mask = masks + patched.append("experts_type_ids") + return patched + + +class HFRunner: + def __init__(self, model_path, device, dtype_name): + self.processor = AutoProcessor.from_pretrained( + model_path, trust_remote_code=True + ) + self.tokenizer = AutoTokenizer.from_pretrained( + model_path, trust_remote_code=True + ) + self.model = AutoModelForCausalLM.from_pretrained( + model_path, + trust_remote_code=True, + torch_dtype=torch_dtype(dtype_name), + device_map=device, + low_cpu_mem_usage=True, + ) + self.model.eval() + patched_meta = fix_hf_meta_helper_tensors(self.model) + if patched_meta: + print(f"patched HF meta helper tensors: {patched_meta}", flush=True) + if hasattr(self.model, "add_image_preprocess"): + self.model.add_image_preprocess(self.processor) + + def generate(self, row, max_new_tokens, tmp_dir, image_size, prompt_style): + prompt, inputs, image_count = build_inputs( + self.processor, self.tokenizer, row, tmp_dir, image_size, prompt_style + ) + model_inputs = { + key: value.to(self.model.device) + for key, value in inputs.items() + if value is not None + } + if "attention_mask" not in model_inputs: + model_inputs["attention_mask"] = torch.ones_like(model_inputs["input_ids"]) + prompt_len = int(model_inputs["input_ids"].shape[-1]) + start = time.time() + with torch.no_grad(): + outputs = self.model.generate( + **model_inputs, + max_new_tokens=max_new_tokens, + do_sample=False, + use_cache=False, + ) + elapsed = time.time() - start + ids = outputs[0][prompt_len:].detach().cpu().tolist() + return { + "prompt_len": prompt_len, + "image_count": image_count, + "output_ids": ids, + "output_text": self.tokenizer.decode(ids, skip_special_tokens=True), + "elapsed_sec": elapsed, + } + + +class CppRunner: + def __init__(self, model_path, device, max_cache_len, tp=1, tp_devices=None): + import infinicore + from infinilm.cache import StaticKVCacheConfig + from infinilm.distributed import DistConfig + from infinilm.infer_engine import InferEngine + from infinilm.modeling_utils import load_model_state_dict_by_file + + self.processor = AutoProcessor.from_pretrained( + model_path, trust_remote_code=True + ) + self.tokenizer = AutoTokenizer.from_pretrained( + model_path, trust_remote_code=True + ) + self.max_cache_len = max_cache_len + if tp_devices is not None and len(tp_devices) != tp: + raise ValueError( + f"--tp-devices length {len(tp_devices)} does not match --tp {tp}" + ) + dist_config = ( + DistConfig(tp_device_ids=tp_devices) + if tp_devices is not None + else DistConfig(tp) + ) + self.engine = InferEngine( + model_path, + device=infinicore.device(device, 0), + distributed_config=dist_config, + cache_config=StaticKVCacheConfig( + max_batch_size=1, max_cache_len=max_cache_len + ), + attention_backend="default", + ) + load_model_state_dict_by_file(self.engine, model_path, dtype=self.engine.dtype) + + def generate(self, row, max_new_tokens, tmp_dir, image_size, prompt_style): + from infinilm.cache import StaticKVCacheConfig + from infinilm.infer_engine import GenerationConfig + + prompt, inputs, image_count = build_inputs( + self.processor, self.tokenizer, row, tmp_dir, image_size, prompt_style + ) + kwargs = {} + if inputs.get("position_ids") is not None: + kwargs["position_ids"] = tensor_to_infini( + inputs["position_ids"].to(torch.int64) + ) + if inputs.get("token_type_ids") is not None: + kwargs["token_type_ids"] = tensor_to_infini( + inputs["token_type_ids"].to(torch.int64) + ) + if inputs.get("images") is not None: + kwargs["images"] = tensor_to_infini(inputs["images"].contiguous()) + kwargs["grid_thw"] = tensor_to_infini(inputs["grid_thw"].to(torch.int64)) + kwargs["image_type_ids"] = tensor_to_infini( + inputs["image_type_ids"].to(torch.int64) + ) + + input_ids = tensor_to_infini(inputs["input_ids"].to(torch.int64)) + self.engine.reset_cache( + StaticKVCacheConfig(max_batch_size=1, max_cache_len=self.max_cache_len) + ) + start = time.time() + output = self.engine.generate( + input_ids, + GenerationConfig( + max_new_tokens=max_new_tokens, + temperature=1.0, + top_k=1, + top_p=1.0, + ), + **kwargs, + ) + elapsed = time.time() - start + ids = [] + for tensor in output: + ids.extend(np.array(tensor.to_numpy()).reshape(-1).astype(int).tolist()) + return { + "prompt_len": int(inputs["input_ids"].shape[-1]), + "image_count": image_count, + "output_ids": ids, + "output_text": self.tokenizer.decode(ids, skip_special_tokens=True), + "elapsed_sec": elapsed, + } + + +def main(): + args = parse_args() + if args.backend == "hf": + if args.tp != 1 or args.tp_devices is not None: + raise ValueError( + "--tp and --tp-devices only configure the InfiniLM cpp backend; " + "HF multi-GPU reference uses transformers device_map via --device auto." + ) + if args.device == "cuda": + print( + "warning: HF backend with --device cuda loads on one GPU; " + "use --device auto for multi-GPU device_map reference runs.", + flush=True, + ) + runner = HFRunner(args.model, args.device, args.torch_dtype) + else: + runner = CppRunner( + args.model, + args.device, + args.max_cache_len, + tp=args.tp, + tp_devices=args.tp_devices, + ) + + rows = [] + official_outputs = {} + total_correct = 0 + total_count = 0 + total_skipped = 0 + + for subject in selected_subjects(args.subjects): + dataset = load_dataset("MMMU/MMMU", subject, split=args.split) + subject_correct = 0 + subject_count = 0 + with tempfile.TemporaryDirectory(prefix=f"mmmu_{subject}_") as tmp_dir: + for row_index, row in enumerate(dataset): + if args.num_samples is not None and subject_count >= args.num_samples: + break + if row.get("question_type") != "multiple-choice": + total_skipped += 1 + continue + result = runner.generate( + row, + args.max_new_tokens, + tmp_dir, + args.image_size, + args.prompt_style, + ) + options = parse_options(row["options"]) + choices = [chr(65 + idx) for idx in range(len(options))] + index2ans = {choice: options[idx] for idx, choice in enumerate(choices)} + pred = extract_answer(result["output_text"], choices, index2ans) + debug_answers = extract_debug_answers(result["output_text"], choices) + gold = str(row["answer"]).strip().upper()[:1] + ok = pred == gold + subject_correct += int(ok) + total_correct += int(ok) + subject_count += 1 + total_count += 1 + item = { + "subject": subject, + "id": row.get("id", row_index), + "gold": gold, + "pred": pred, + "ok": int(ok), + "prompt_len": result["prompt_len"], + "image_count": result["image_count"], + "elapsed_sec": f"{result['elapsed_sec']:.3f}", + "new_tokens": len(result["output_ids"]), + "hit_max_new_tokens": int( + len(result["output_ids"]) >= args.max_new_tokens + ), + **debug_answers, + "output_ids": " ".join(str(x) for x in result["output_ids"]), + "output_text": result["output_text"].replace("\n", "\\n"), + } + rows.append(item) + official_outputs[str(row.get("id", row_index))] = ( + pred or result["output_text"] + ) + print(json.dumps(item, ensure_ascii=False), flush=True) + + acc = subject_correct / subject_count if subject_count else 0.0 + print( + f"SUBJECT {subject} {subject_correct}/{subject_count} {acc:.4f}", flush=True + ) + + overall = total_correct / total_count if total_count else 0.0 + with open(args.output_csv, "w", newline="", encoding="utf-8") as f: + fieldnames = [ + "subject", + "id", + "gold", + "pred", + "ok", + "prompt_len", + "image_count", + "elapsed_sec", + "new_tokens", + "hit_max_new_tokens", + "boxed_answer", + "explicit_answer", + "output_ids", + "output_text", + ] + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(rows) + report = { + "backend": args.backend, + "subjects": selected_subjects(args.subjects), + "split": args.split, + "prompt_style": args.prompt_style, + "image_size": args.image_size, + "tp": args.tp if args.backend == "cpp" else None, + "tp_devices": args.tp_devices if args.backend == "cpp" else None, + "correct": total_correct, + "total": total_count, + "skipped": total_skipped, + "accuracy": overall, + "csv": args.output_csv, + } + if args.output_json: + with open(args.output_json, "w", encoding="utf-8") as f: + json.dump( + { + **report, + "official_eval_only_outputs": official_outputs, + "rows": rows, + }, + f, + ensure_ascii=False, + indent=2, + ) + print(json.dumps(report, ensure_ascii=False), flush=True) + + +if __name__ == "__main__": + main() diff --git a/examples/ernie_vl_verification.md b/examples/ernie_vl_verification.md new file mode 100644 index 000000000..bf12bef59 --- /dev/null +++ b/examples/ernie_vl_verification.md @@ -0,0 +1,164 @@ +# ERNIE-4.5-VL verification + +This note contains portable verification commands for +`ERNIE-4.5-VL-28B-A3B-Thinking`. It deliberately uses environment variables +and repository-relative output paths so it can be shared without exposing +machine-specific accounts, hosts, credentials, or storage layouts. + +## Prerequisites + +Build and install InfiniCore and InfiniLM following their repository +instructions. Use a Python environment that provides PyTorch, Transformers, +and the dependencies required by the selected benchmark dataset. + +From the InfiniLM repository root, set the model directory and create a local +result directory: + +```bash +export MODEL_DIR="${MODEL_DIR:?Set MODEL_DIR to the downloaded model directory}" +mkdir -p artifacts/ernie_vl +``` + +The commands below assume four visible CUDA devices. Adjust `--tp` and +`--tp-devices` only when the model fits the selected hardware configuration. + +## Tensor-parallel correctness + +Run deterministic text, image, and video checks against the token IDs embedded +in the verification script: + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 \ +python examples/ernie_vl_correctness.py \ + --model "$MODEL_DIR" \ + --tp 4 \ + --tp-devices 0,1,2,3 \ + --cases text image video \ + --max-cache-len 512 \ + --reference-mode expected \ + --output-json artifacts/ernie_vl/correctness_tp4.json +``` + +Acceptance criteria: + +- the top-level `ok` field is `true`; +- `text`, `image`, and `video` all report `match_expected=true`; +- the process exits successfully. + +To generate a live Transformers reference in a separate process, run: + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 \ +python examples/ernie_vl_correctness.py \ + --model "$MODEL_DIR" \ + --tp 4 \ + --tp-devices 0,1,2,3 \ + --cases text image video \ + --max-cache-len 512 \ + --reference-mode hf \ + --hf-device-map auto \ + --hf-torch-dtype bf16 \ + --output-json artifacts/ernie_vl/correctness_tp4_hf.json +``` + +This mode requires enough memory for the Transformers reference as well as the +InfiniLM run. Its acceptance criteria are the same as the deterministic check. + +## Messages API smoke tests + +Static-cache smoke test: + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 \ +python examples/ernie_vl_llm_smoke.py \ + --model "$MODEL_DIR" \ + --tp 4 \ + --cache-type static \ + --output-json artifacts/ernie_vl/llm_static_tp4.json +``` + +Paged-cache smoke test: + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 \ +python examples/ernie_vl_llm_smoke.py \ + --model "$MODEL_DIR" \ + --tp 4 \ + --cache-type paged \ + --num-blocks 64 \ + --block-size 16 \ + --output-json artifacts/ernie_vl/llm_paged_tp4.json +``` + +Each output JSON should contain `ok=true`. + +## C-Eval and MMLU + +The standard benchmark entry point can evaluate the C++ backend. The following +commands run the complete validation split and write repository-local CSVs: + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 \ +python test/bench/test_benchmark.py \ + --device nvidia \ + --model "$MODEL_DIR" \ + --bench ceval \ + --subject all \ + --split val \ + --max-new-tokens 5 \ + --backend cpp \ + --tp 4 \ + --dtype bfloat16 \ + --output-csv artifacts/ernie_vl/ceval_val_tp4.csv + +CUDA_VISIBLE_DEVICES=0,1,2,3 \ +python test/bench/test_benchmark.py \ + --device nvidia \ + --model "$MODEL_DIR" \ + --bench mmlu \ + --subject all \ + --split val \ + --max-new-tokens 5 \ + --backend cpp \ + --tp 4 \ + --dtype bfloat16 \ + --output-csv artifacts/ernie_vl/mmlu_val_tp4.csv +``` + +Dataset access may require an existing local cache or outbound network access. +Report the sample count, accuracy, hardware, dtype, and tensor-parallel degree +together; results from different configurations are not directly comparable. + +## MMMU adaptation smoke test + +`ernie_vl_mmmu_smoke.py` is an adaptation check, not an official leaderboard +submission. A small C++ backend run can be launched with: + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 \ +python examples/ernie_vl_mmmu_smoke.py \ + --model "$MODEL_DIR" \ + --backend cpp \ + --subjects Accounting \ + --split validation \ + --num-samples 2 \ + --max-new-tokens 64 \ + --prompt-style official \ + --tp 4 \ + --tp-devices 0,1,2,3 \ + --output-csv artifacts/ernie_vl/mmmu_smoke_tp4.csv \ + --output-json artifacts/ernie_vl/mmmu_smoke_tp4.json +``` + +For a fair backend comparison, keep the dataset split, subjects, prompt style, +sample count, image size, and generation limit identical. + +## Publishing evidence + +Before attaching logs or screenshots to a pull request: + +- show the command, exit status, and final result summary; +- remove terminal history, login banners, hostnames, addresses, credentials, + cache locations, and machine-specific absolute paths; +- do not publish model access tokens or dataset credentials; +- label smoke-test and partial-dataset results as such. diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py index 3e874046e..08df6ca7a 100644 --- a/python/infinilm/infer_engine.py +++ b/python/infinilm/infer_engine.py @@ -10,6 +10,7 @@ from infinilm.lib import _infinilm from .exception_utils import handle_oom_and_exit +from .generation.utils import infini_to_numpy # noqa: F401 - registers Tensor.to_numpy from .modeling_utils import parse_dtype _MODEL_DEFAULTS = { @@ -214,6 +215,7 @@ def _build_input( input_ids, *, position_ids=None, + token_type_ids=None, past_kv_lengths=None, total_kv_lengths=None, input_offsets=None, @@ -223,9 +225,12 @@ def _build_input( mamba_init_state_indices=None, mamba_final_state_indices=None, pixel_values=None, + images=None, image_bound=None, tgt_sizes=None, image_grid_thw=None, + grid_thw=None, + image_type_ids=None, image_req_ids=None, visual_token_ranges=None, target_hidden_states=None, @@ -241,6 +246,7 @@ def unwrap_tensor(tensor): input_ids = unwrap_tensor(input_ids) position_ids = unwrap_tensor(position_ids) + token_type_ids = unwrap_tensor(token_type_ids) past_kv_lengths = unwrap_tensor(past_kv_lengths) total_kv_lengths = unwrap_tensor(total_kv_lengths) input_offsets = unwrap_tensor(input_offsets) @@ -260,10 +266,14 @@ def convert_tensor_list(tensor_list_): return None return [unwrap_tensor(tensor) for tensor in tensor_list_] + if pixel_values is None and images is not None: + pixel_values = images pixel_values = convert_tensor_list(pixel_values) image_bound = convert_tensor_list(image_bound) tgt_sizes = convert_tensor_list(tgt_sizes) image_grid_thw = convert_tensor_list(image_grid_thw) + grid_thw = convert_tensor_list(grid_thw) + image_type_ids = convert_tensor_list(image_type_ids) temperature = 1.0 if temperature is None else temperature top_k = 1 if top_k is None else top_k @@ -272,6 +282,7 @@ def convert_tensor_list(tensor_list_): return super().Input( input_ids, position_ids=position_ids, + token_type_ids=token_type_ids, past_sequence_lengths=past_kv_lengths, total_sequence_lengths=total_kv_lengths, input_offsets=input_offsets, @@ -284,6 +295,8 @@ def convert_tensor_list(tensor_list_): image_bound=image_bound, tgt_sizes=tgt_sizes, image_grid_thw=image_grid_thw, + grid_thw=grid_thw, + image_type_ids=image_type_ids, image_req_ids=image_req_ids, visual_token_ranges=visual_token_ranges, target_hidden_states=target_hidden_states, @@ -298,6 +311,7 @@ def forward( input_ids, *, position_ids=None, + token_type_ids=None, past_kv_lengths=None, total_kv_lengths=None, input_offsets=None, @@ -307,9 +321,12 @@ def forward( mamba_init_state_indices=None, mamba_final_state_indices=None, pixel_values=None, + images=None, image_bound=None, tgt_sizes=None, image_grid_thw=None, + grid_thw=None, + image_type_ids=None, image_req_ids=None, visual_token_ranges=None, target_hidden_states=None, @@ -323,6 +340,9 @@ def forward( position_ids = ( position_ids._underlying if position_ids is not None else None ) + token_type_ids = ( + token_type_ids._underlying if token_type_ids is not None else None + ) past_kv_lengths = ( past_kv_lengths._underlying if past_kv_lengths is not None else None ) @@ -359,10 +379,14 @@ def convert_tensor_list(tensor_list_): return None return [tensor._underlying for tensor in tensor_list_] + if pixel_values is None and images is not None: + pixel_values = images pixel_values = convert_tensor_list(pixel_values) image_bound = convert_tensor_list(image_bound) tgt_sizes = convert_tensor_list(tgt_sizes) image_grid_thw = convert_tensor_list(image_grid_thw) + grid_thw = convert_tensor_list(grid_thw) + image_type_ids = convert_tensor_list(image_type_ids) return infinicore.Tensor( super() @@ -370,6 +394,7 @@ def convert_tensor_list(tensor_list_): self._build_input( input_ids, position_ids=position_ids, + token_type_ids=token_type_ids, past_kv_lengths=past_kv_lengths, total_kv_lengths=total_kv_lengths, input_offsets=input_offsets, @@ -382,6 +407,8 @@ def convert_tensor_list(tensor_list_): image_bound=image_bound, tgt_sizes=tgt_sizes, image_grid_thw=image_grid_thw, + grid_thw=grid_thw, + image_type_ids=image_type_ids, image_req_ids=image_req_ids, visual_token_ranges=visual_token_ranges, target_hidden_states=target_hidden_states, @@ -401,6 +428,7 @@ def forward_raw( input_ids, *, position_ids=None, + token_type_ids=None, past_kv_lengths=None, total_kv_lengths=None, input_offsets=None, @@ -408,8 +436,12 @@ def forward_raw( block_tables=None, slot_mapping=None, pixel_values=None, + images=None, image_bound=None, tgt_sizes=None, + image_grid_thw=None, + grid_thw=None, + image_type_ids=None, image_req_ids=None, visual_token_ranges=None, target_hidden_states=None, @@ -423,6 +455,7 @@ def forward_raw( self._build_input( input_ids, position_ids=position_ids, + token_type_ids=token_type_ids, past_kv_lengths=past_kv_lengths, total_kv_lengths=total_kv_lengths, input_offsets=input_offsets, @@ -430,8 +463,12 @@ def forward_raw( block_tables=block_tables, slot_mapping=slot_mapping, pixel_values=pixel_values, + images=images, image_bound=image_bound, tgt_sizes=tgt_sizes, + image_grid_thw=image_grid_thw, + grid_thw=grid_thw, + image_type_ids=image_type_ids, image_req_ids=image_req_ids, visual_token_ranges=visual_token_ranges, target_hidden_states=target_hidden_states, @@ -455,9 +492,15 @@ def generate( input_ids, generation_config, *, + position_ids=None, + token_type_ids=None, pixel_values=None, + images=None, image_bound=None, tgt_sizes=None, + image_grid_thw=None, + grid_thw=None, + image_type_ids=None, _measure_and_log_time=False, ): eos_token_id = self.eos_token_id @@ -473,6 +516,33 @@ def generate( "When `batch_size > 1`, `max_new_tokens` must be specified." ) + if self.get_cache_config() is None: + raise ValueError( + "InferEngine.generate requires a KV cache config. " + "Pass StaticKVCacheConfig or PagedKVCacheConfig when constructing " + "InferEngine, or call reset_cache before generate()." + ) + + external_position_ids = position_ids + external_token_type_ids = token_type_ids + generated_position_ids = None + if external_position_ids is not None: + position_ids_np = external_position_ids.to_numpy() + if len(position_ids_np.shape) == 3: + max_positions = position_ids_np.max(axis=1, keepdims=True) + generated_position_ids = (max_positions + 1).tolist() + elif len(position_ids_np.shape) == 2: + generated_position_ids = (position_ids_np[:, -1:] + 1).tolist() + else: + generated_position_ids = [[int(position_ids_np.reshape([-1])[-1]) + 1]] + generated_token_type_ids = None + if external_token_type_ids is not None: + token_type_ids_np = external_token_type_ids.to_numpy() + batch_for_token_types = ( + token_type_ids_np.shape[0] if len(token_type_ids_np.shape) >= 2 else 1 + ) + generated_token_type_ids = [[0] for _ in range(batch_for_token_types)] + if _measure_and_log_time: time_measurements = [] @@ -508,6 +578,14 @@ def generate( dtype=infinicore.int32, ) + def position_ids_from_list(position_ids_list): + try: + return infinicore.from_list(position_ids_list, dtype=infinicore.int64) + except ValueError: + return infinicore.from_list_by_numpy( + position_ids_list, dtype=infinicore.int64 + ) + for iter in range(0, generation_config.max_new_tokens): if _measure_and_log_time: start_time = time.perf_counter() @@ -516,10 +594,17 @@ def generate( if self.enable_paged_attn: input_ids = input_ids.view([1, batch_size * seq_len]) - position_ids = infinicore.from_list( - list(range(past_seq_len, past_seq_len + seq_len)) * batch_size, - dtype=infinicore.int64, - ) + if external_position_ids is None: + iter_position_ids = infinicore.from_list( + list(range(past_seq_len, past_seq_len + seq_len)) * batch_size, + dtype=infinicore.int64, + ) + else: + iter_position_ids = ( + external_position_ids + if iter == 0 + else position_ids_from_list(generated_position_ids) + ) if iter == 0: slot_mapping_list = [] @@ -547,13 +632,20 @@ def generate( dtype=infinicore.int64, ) else: - position_ids = infinicore.from_list( - [ - list(range(past_seq_len, past_seq_len + seq_len)) - for _ in range(batch_size) - ], - dtype=infinicore.int64, - ) + if external_position_ids is None: + iter_position_ids = infinicore.from_list( + [ + list(range(past_seq_len, past_seq_len + seq_len)) + for _ in range(batch_size) + ], + dtype=infinicore.int64, + ) + else: + iter_position_ids = ( + external_position_ids + if iter == 0 + else position_ids_from_list(generated_position_ids) + ) slot_mapping = None @@ -586,7 +678,19 @@ def generate( output_id = self( input_ids=input_ids, pixel_values=pixel_values if iter == 0 else None, - position_ids=position_ids, + images=images if iter == 0 else None, + position_ids=iter_position_ids, + token_type_ids=( + external_token_type_ids + if iter == 0 + else ( + infinicore.from_list( + generated_token_type_ids, dtype=infinicore.int64 + ) + if generated_token_type_ids is not None + else None + ) + ), past_kv_lengths=past_kv_lengths, total_kv_lengths=total_kv_lengths, input_offsets=input_offsets, @@ -597,6 +701,9 @@ def generate( mamba_final_state_indices=mamba_final_state_indices, image_bound=image_bound if iter == 0 else None, tgt_sizes=tgt_sizes if iter == 0 else None, + image_grid_thw=image_grid_thw if iter == 0 else None, + grid_thw=grid_thw if iter == 0 else None, + image_type_ids=image_type_ids if iter == 0 else None, temperature=generation_config.temperature, top_k=generation_config.top_k, top_p=generation_config.top_p, @@ -616,6 +723,16 @@ def generate( input_ids = output_id.view([batch_size, 1]) past_seq_len = past_seq_len + seq_len + if generated_position_ids is not None: + if isinstance(generated_position_ids[0][0], list): + generated_position_ids = [ + [[value + 1 for value in generated_position_ids[b][0]]] + for b in range(len(generated_position_ids)) + ] + else: + generated_position_ids = [ + [row[0] + 1] for row in generated_position_ids + ] if _measure_and_log_time: end_time = time.perf_counter() diff --git a/python/infinilm/llm/model_runner/model_runner.py b/python/infinilm/llm/model_runner/model_runner.py index 45f1eaee1..fdf5c4d87 100644 --- a/python/infinilm/llm/model_runner/model_runner.py +++ b/python/infinilm/llm/model_runner/model_runner.py @@ -65,6 +65,8 @@ def __init__(self, config: EngineConfig): else: raise ValueError(f"Unsupported cache_type: {config.cache_type}") + attention_backend = self._resolve_attention_backend(config) + # Initialize model engine self.model_engine = InferEngine( model_path=config.model_path, @@ -76,7 +78,7 @@ def __init__(self, config: EngineConfig): ), cache_config=cache_config, enable_graph_compiling=config.enable_graph, - attention_backend=config.attn_backend, + attention_backend=attention_backend, use_mla=config.use_mla, weight_load_mode=config.weight_load_mode, skip_legacy_moe=config.skip_legacy_moe, @@ -128,6 +130,16 @@ def __init__(self, config: EngineConfig): self.kv_connector.register_kv_caches(kv_caches) + @staticmethod + def _resolve_attention_backend(config: EngineConfig) -> str: + if config.attn_backend != "default": + return config.attn_backend + if config.cache_type == "paged": + return "paged-attn" + if config.cache_type == "static": + return "static-attn" + return config.attn_backend + @property def model_type(self): return self.model_engine.model_type diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index 8f84c9b97..19561c67f 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -232,6 +232,13 @@ def load_model_state_dict_by_file( already_loaded_keys.extend(model_param.keys()) + if os.getenv("INFINILM_ERNIE_GPU_ROUTER"): + for key in list(model_param.keys()): + if key.endswith(".mlp.gate.weight") or key.endswith( + ".mlp.gate.weight_1" + ): + model_param[key] = model_param[key].to(dtype=torch.bfloat16) + # --------------------------------------------------------- # # Scale embed_tokens on torch side before converting # --------------------------------------------------------- # @@ -344,6 +351,11 @@ def load_model_state_dict_by_tensor( with safe_open(file_path, "pt", "cpu") as f: for name in f.keys(): tensor = f.get_tensor(name).to(dtype=torch_dtype) + if os.getenv("INFINILM_ERNIE_GPU_ROUTER") and ( + name.endswith(".mlp.gate.weight") + or name.endswith(".mlp.gate.weight_1") + ): + tensor = tensor.to(dtype=torch.bfloat16) if name == "model.embed_tokens.weight": embed_tokens_torch_unscaled = tensor @@ -361,6 +373,10 @@ def load_model_state_dict_by_tensor( for key in model_params.keys(): tensor = model_params[key].to(dtype=torch_dtype) + if os.getenv("INFINILM_ERNIE_GPU_ROUTER") and ( + key.endswith(".mlp.gate.weight") or key.endswith(".mlp.gate.weight_1") + ): + tensor = tensor.to(dtype=torch.bfloat16) if key == "model.embed_tokens.weight": embed_tokens_torch_unscaled = tensor if scale_emb != 1.0: @@ -705,6 +721,16 @@ def _remap_videonsa(state_dict, config=None): return state_dict +def _remap_ernie4_5_moe_vl(state_dict, config=None): + """Remap ERNIE text, multimodal MoE, vision tower, and VL resampler weights.""" + result = {} + for key, tensor in state_dict.items(): + if key.endswith(".mlp.gate.weight") or key.endswith(".mlp.gate.weight_1"): + tensor = tensor.t().contiguous() + result[key] = tensor + return result + + # Model type → remap function mapping def _remap_qwen3_5(state_dict, config): """Apply Qwen3.5-specific load-time weight fixes.""" @@ -751,4 +777,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_moe_vl_processor.py b/python/infinilm/processors/ernie4_5_moe_vl_processor.py new file mode 100644 index 000000000..2ff360352 --- /dev/null +++ b/python/infinilm/processors/ernie4_5_moe_vl_processor.py @@ -0,0 +1,385 @@ +from transformers import AutoProcessor +from typing_extensions import override + +from .processor import InfinilmProcessor, register_processor + + +@register_processor("ernie4_5_moe_vl") +class Ernie4_5MoeVLProcessor(InfinilmProcessor): + def __init__(self, model_dir_path: str): + self.processor = AutoProcessor.from_pretrained( + model_dir_path, trust_remote_code=True + ) + self.tokenizer = self.processor.tokenizer + + @override + def __call__( + self, + prompt=None, + images=None, + videos=None, + audios=None, + return_tensors: str = None, + **kwargs, + ) -> dict: + if prompt is None: + prompt = kwargs.pop("text", None) + if prompt is None: + raise ValueError("prompt or text must be provided") + if images is None and videos is None and audios is None: + return self.tokenizer(prompt, return_tensors=return_tensors, **kwargs) + return self.processor( + prompt, + images=images, + videos=videos, + return_tensors=return_tensors or "pt", + **kwargs, + ) + + @override + def apply_chat_template( + self, + conversation, + add_generation_prompt: bool = False, + tokenize: bool = True, + **kwargs, + ): + normalized = [] + for message in conversation: + role = message.get("role", "user") + content = message.get("content", "") + if not isinstance(content, list): + normalized.append({"role": role, "content": content}) + continue + + parts = [] + for item in content: + item_type = item.get("type") + if item_type == "text": + parts.append(item.get("text", "")) + elif item_type == "image_url": + parts.append( + self.processor.IMG_START + + "<|image@placeholder|>" + + self.processor.IMG_END + ) + elif item_type == "video_url": + parts.append( + self.processor.VID_START + + "<|video@placeholder|>" + + self.processor.VID_END + ) + else: + raise NotImplementedError( + "Only text, image_url and video_url inputs are supported" + ) + normalized.append({"role": role, "content": "".join(parts)}) + + return self.tokenizer.apply_chat_template( + conversation=normalized, + add_generation_prompt=add_generation_prompt, + tokenize=tokenize, + **kwargs, + ) + + @override + def build_model_inputs( + self, + scheduler_output, + temperature: float = 1.0, + top_p: float = 0.8, + top_k: int = 1, + **kwargs, + ) -> dict: + import infinicore + import torch + + if not scheduler_output.scheduled_requests: + raise RuntimeError( + "build_model_inputs called with empty scheduled_requests" + ) + + static_prefix_hit = getattr(scheduler_output, "prefix_hit_len", None) + if static_prefix_hit is not None: + req = scheduler_output.scheduled_requests[0] + processed = req.processed_inputs + has_vision_inputs = ( + processed is not None and processed.get("images") is not None + ) + + if scheduler_output.is_prefill: + req_tokens = req.get_input_tokens() + # Static cache does not track multimodal token bounds. Recompute + # vision prompts from the start so image feature replacement stays + # aligned with the patch tokens present in input_ids. + num_cached = 0 if has_vision_inputs else static_prefix_hit + input_tokens = req_tokens[num_cached:] + seq_len = len(req_tokens) + input_offsets = [0, len(input_tokens)] + + if processed is not None and "position_ids" in processed: + position_ids = ( + processed["position_ids"][:, num_cached:seq_len] + .to(torch.int64) + .contiguous() + ) + else: + position_ids = [ + list(range(num_cached, num_cached + len(input_tokens))) + ] + + token_type_ids = None + if processed is not None and "token_type_ids" in processed: + token_type_ids = ( + processed["token_type_ids"][0, num_cached:seq_len] + .to(torch.int64) + .tolist() + ) + else: + seq_len = req.get_total_length() + input_tokens = [ + req.generated_token_ids[-1] + if req.generated_token_ids + else req.prompt_token_ids[-1] + ] + num_cached = seq_len - 1 + input_offsets = [0, 1] + + if processed is not None and "position_ids" in processed: + max_position = processed["position_ids"][0].max(dim=0)[0] + generated = max_position + (seq_len - len(req.get_input_tokens())) + position_ids = generated.to(torch.int64).view(1, 1, 3).contiguous() + else: + position_ids = [[seq_len - 1]] + + token_type_ids = None + if processed is not None and "token_type_ids" in processed: + token_type_ids = [0] + + result = { + "input_ids": infinicore.from_list( + [input_tokens], dtype=infinicore.int64 + ), + "position_ids": ( + infinicore.from_torch(position_ids) + if torch.is_tensor(position_ids) + else infinicore.from_list(position_ids, dtype=infinicore.int64) + ), + "past_kv_lengths": infinicore.from_list( + [num_cached], dtype=infinicore.int32 + ), + "total_kv_lengths": infinicore.from_list( + [seq_len], dtype=infinicore.int32 + ), + "input_offsets": infinicore.from_list( + input_offsets, dtype=infinicore.int32 + ), + "cu_seqlens": infinicore.from_list( + [0, seq_len], dtype=infinicore.int32 + ), + "block_tables": None, + "slot_mapping": None, + "temperature": temperature, + "top_k": top_k, + "top_p": top_p, + } + if token_type_ids is not None: + result["token_type_ids"] = infinicore.from_list( + [token_type_ids], dtype=infinicore.int64 + ) + if scheduler_output.is_prefill and has_vision_inputs: + result["images"] = infinicore.from_torch( + processed["images"].contiguous() + ) + result["grid_thw"] = infinicore.from_torch( + processed["grid_thw"].to(torch.int64).contiguous() + ) + result["image_type_ids"] = infinicore.from_torch( + processed["image_type_ids"].to(torch.int64).contiguous() + ) + result["image_req_ids"] = [0] + return result + + tokens = [] + seq_lens = [] + seq_offsets = [0] + block_tables = [] + slot_mapping = [] + cached_lens = [] + position_ids = [] + token_type_ids = [] + cu_seqlens = [0] + mm_data = {} + + max_block_table_len = max( + len(req.block_table) for req in scheduler_output.scheduled_requests + ) + current_offset = 0 + + def append_mm_data(key, value): + if mm_data.get(key) is None: + mm_data[key] = [value] + else: + mm_data[key].append(value) + + for req_id, req in enumerate(scheduler_output.scheduled_requests): + processed = req.processed_inputs + has_vision_inputs = ( + processed is not None and processed.get("images") is not None + ) + + if scheduler_output.is_prefill: + num_cached = req.num_local_cached_tokens + + req_tokens = req.get_input_tokens() + tokens_to_compute = req_tokens[num_cached:] + tokens.extend(tokens_to_compute) + + compute_len = len(tokens_to_compute) + seq_len = len(req_tokens) + seq_lens.append(seq_len) + + current_offset += compute_len + seq_offsets.append(current_offset) + + slot_mapping.extend(req.slot_mapping) + cached_lens.append(num_cached) + + if processed is not None and "position_ids" in processed: + position_ids.extend( + processed["position_ids"][0, num_cached:seq_len].tolist() + ) + else: + position_ids.extend(range(num_cached, num_cached + compute_len)) + + if processed is not None and "token_type_ids" in processed: + token_type_ids.extend( + processed["token_type_ids"][0, num_cached:seq_len] + .to(torch.int64) + .tolist() + ) + + if has_vision_inputs and ( + num_cached == 0 + or req_tokens[num_cached:].count(self.processor.image_patch_id) > 0 + ): + append_mm_data( + "images", + infinicore.from_torch(processed["images"].contiguous()), + ) + append_mm_data( + "grid_thw", + infinicore.from_torch( + processed["grid_thw"].to(torch.int64).contiguous() + ), + ) + append_mm_data( + "image_type_ids", + infinicore.from_torch( + processed["image_type_ids"].to(torch.int64).contiguous() + ), + ) + append_mm_data("image_req_ids", req_id) + else: + num_cached = req.num_local_cached_tokens + seq_len = req.get_total_length() + last_token = req.generated_token_ids[-1] + tokens.append(last_token) + seq_lens.append(seq_len) + + current_offset += 1 + seq_offsets.append(current_offset) + + slot_mapping.extend(req.slot_mapping) + cached_lens.append(num_cached) + + if processed is not None and "position_ids" in processed: + max_position = processed["position_ids"][0].max(dim=0)[0] + generated = max_position + (seq_len - len(req.get_input_tokens())) + position_ids.append(generated.to(torch.int64).tolist()) + else: + position_ids.append(seq_len - 1) + + if processed is not None and "token_type_ids" in processed: + token_type_ids.append(0) + + padded_block_table = req.block_table + [-1] * ( + max_block_table_len - len(req.block_table) + ) + block_tables.append(padded_block_table) + cu_seqlens.append(cu_seqlens[-1] + seq_len) + + position_ids_payload = ( + [position_ids] + if position_ids and isinstance(position_ids[0], list) + else position_ids + ) + position_ids_value = ( + infinicore.from_torch( + torch.tensor(position_ids_payload, dtype=torch.int64).contiguous() + ) + if ( + position_ids_payload + and isinstance(position_ids_payload[0], list) + and position_ids_payload[0] + and isinstance(position_ids_payload[0][0], list) + ) + else infinicore.from_list(position_ids_payload, dtype=infinicore.int64) + ) + result = { + "input_ids": infinicore.from_list([tokens], dtype=infinicore.int64), + "position_ids": position_ids_value, + "past_kv_lengths": infinicore.from_list( + cached_lens, dtype=infinicore.int32 + ), + "total_kv_lengths": infinicore.from_list(seq_lens, dtype=infinicore.int32), + "input_offsets": infinicore.from_list(seq_offsets, dtype=infinicore.int32), + "cu_seqlens": infinicore.from_list(cu_seqlens, dtype=infinicore.int32), + "block_tables": infinicore.from_list(block_tables, dtype=infinicore.int32), + "slot_mapping": infinicore.from_list(slot_mapping, dtype=infinicore.int64), + "temperature": temperature, + "top_k": top_k, + "top_p": top_p, + **mm_data, + } + if token_type_ids: + result["token_type_ids"] = infinicore.from_list( + [token_type_ids], dtype=infinicore.int64 + ) + return result + + @override + def get_tokenizer(self): + return self.tokenizer + + @override + def get_mm_token_index_list( + self, prompt_token_ids, image_ids=None, video_ids=None, audio_ids=None, **kwargs + ): + mappings = [] + vision_ids = (image_ids or []) + (video_ids or []) + image_index = 0 + patch_id = self.processor.image_patch_id + idx = 0 + prompt_len = len(prompt_token_ids) + while idx < len(prompt_token_ids): + if prompt_token_ids[idx] != patch_id: + idx += 1 + continue + start = idx + while idx < len(prompt_token_ids) and prompt_token_ids[idx] == patch_id: + idx += 1 + identifier = ( + vision_ids[image_index] + if image_index < len(vision_ids) + else f"vision_{image_index}" + ) + mappings.append( + { + "start_index": start, + "end_index": prompt_len, + "identifier": identifier, + } + ) + image_index += 1 + return mappings diff --git a/test/bench/backends/transformers.py b/test/bench/backends/transformers.py index 2958d9da5..78d3f8d43 100644 --- a/test/bench/backends/transformers.py +++ b/test/bench/backends/transformers.py @@ -5,6 +5,52 @@ from .base import BaseBenchmark +def _fix_hf_meta_helper_tensors(model): + """Materialize ERNIE remote-code helper tensors left on the meta device.""" + import torch + + patched = [] + vision_model = getattr(model, "vision_model", None) or getattr( + model, "visual", None + ) + rotary = ( + getattr(vision_model, "rotary_pos_emb", None) + if vision_model is not None + else None + ) + inv_freq = getattr(rotary, "inv_freq", None) if rotary is not None else None + if inv_freq is not None and getattr(inv_freq, "device", None).type == "meta": + dim = int(inv_freq.numel() * 2) + theta = float(getattr(rotary, "theta", 10000.0)) + rotary.inv_freq = 1.0 / theta ** ( + torch.arange(start=0, end=dim, step=2, dtype=torch.float32) / dim + ) + patched.append("vision_rotary_inv_freq") + + for module in model.modules(): + experts_type_ids = getattr(module, "experts_type_ids", None) + if ( + experts_type_ids is None + or getattr(experts_type_ids, "device", None).type != "meta" + ): + continue + config = getattr(module, "config", None) + moe_num_experts = getattr(config, "moe_num_experts", None) + if not isinstance(moe_num_experts, (list, tuple)): + continue + rebuilt = torch.zeros([sum(moe_num_experts)], dtype=torch.int64) + offset = 0 + for idx, expert_num in enumerate(moe_num_experts): + rebuilt[offset : offset + expert_num] = idx + offset += expert_num + module.experts_type_ids = rebuilt + module.experts_type_mask = [ + module.experts_type_ids == idx for idx, _ in enumerate(moe_num_experts) + ] + patched.append("experts_type_ids") + return patched + + class TransformersBenchmark(BaseBenchmark): """Hugging Face Transformers backend.""" @@ -31,9 +77,18 @@ def __init__( with open(os.path.join(model_dir_path, "config.json"), "r") as f: self.config_dict = json.load(f) - self.tokenizer = transformers.AutoTokenizer.from_pretrained( - model_dir_path, trust_remote_code=True + self.use_processor_inputs = ( + self.config_dict.get("model_type") == "ernie4_5_moe_vl" ) + if self.use_processor_inputs: + self.processor = transformers.AutoProcessor.from_pretrained( + model_dir_path, trust_remote_code=True + ) + self.tokenizer = self.processor.tokenizer + else: + self.tokenizer = transformers.AutoTokenizer.from_pretrained( + model_dir_path, trust_remote_code=True + ) print("Loading model with Transformers backend...") load_kwargs = { @@ -65,6 +120,9 @@ def __init__( if tensor_parallel_size <= 1: self.model = self.model.to(self.device) self.model.eval() + patched_meta = _fix_hf_meta_helper_tensors(self.model) + if patched_meta: + print(f"Patched HF meta helper tensors: {patched_meta}") self.input_device = self.model.get_input_embeddings().weight.device print("Transformers model loaded successfully") @@ -86,6 +144,45 @@ def generate(self, *args, max_steps=500, topp_=1.0, topk_=1, temperature_=1.0): prompt = self.render_input_content(*args) print(prompt, end="", flush=True) + if self.use_processor_inputs: + inputs = self.processor( + prompt, return_tensors="pt", add_special_tokens=False + ) + model_inputs = { + key: value.to(self.input_device) + for key, value in inputs.items() + if value is not None + } + if "attention_mask" not in model_inputs: + model_inputs["attention_mask"] = torch.ones_like( + model_inputs["input_ids"] + ) + input_tokens = int(model_inputs["input_ids"].shape[-1]) + + self._synchronize() + start_time = time.perf_counter() + outputs = self.model.generate( + **model_inputs, + max_new_tokens=max_steps, + do_sample=temperature_ > 0, + temperature=temperature_, + top_k=topk_, + top_p=topp_, + eos_token_id=self.eos_token_id, + pad_token_id=self.tokenizer.pad_token_id or 2, + use_cache=False, + ) + self._synchronize() + + generated_ids = outputs[0][input_tokens:] + output_text = self.tokenizer.decode(generated_ids, skip_special_tokens=True) + return self.record_generation( + output_text, + input_tokens, + generated_ids.numel(), + start_time, + ) + tokens = self.encode_text(prompt) input_ids = torch.tensor([tokens], device=self.input_device) diff --git a/tools/compare_mmmu_results.py b/tools/compare_mmmu_results.py new file mode 100644 index 000000000..abaebc6b2 --- /dev/null +++ b/tools/compare_mmmu_results.py @@ -0,0 +1,251 @@ +#!/usr/bin/env python3 +"""Compare two ERNIE MMMU smoke JSON outputs.""" + +from __future__ import annotations + +import argparse +import json +import sys +import tarfile +from collections import defaultdict +from pathlib import Path +from typing import Any + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "--left", required=True, help="Left JSON path or .tar.gz archive" + ) + parser.add_argument( + "--right", required=True, help="Right JSON path or .tar.gz archive" + ) + parser.add_argument("--left-label", default="left") + parser.add_argument("--right-label", default="right") + parser.add_argument( + "--left-member", + default=None, + help="JSON member inside left archive. Auto-detected if omitted.", + ) + parser.add_argument( + "--right-member", + default=None, + help="JSON member inside right archive. Auto-detected if omitted.", + ) + parser.add_argument("--output", default=None, help="Optional Markdown output path") + return parser.parse_args() + + +def read_json(path_arg: str, member: str | None = None) -> dict[str, Any]: + path = Path(path_arg) + if path.suffix == ".json": + with path.open("r", encoding="utf-8") as f: + return json.load(f) + if path.suffix == ".log": + return read_log_rows(path) + if path.name.endswith((".tar.gz", ".tgz")): + with tarfile.open(path, "r:gz") as archive: + selected = select_json_member(archive, member) + with archive.extractfile(selected) as f: + if f is None: + raise FileNotFoundError(selected) + return json.loads(f.read().decode("utf-8")) + raise ValueError(f"Unsupported input path: {path}") + + +def read_log_rows(path: Path) -> dict[str, Any]: + rows = [] + for line in path.read_text(errors="replace").splitlines(): + if not line.startswith("{") or '"subject"' not in line: + continue + try: + row = json.loads(line) + except json.JSONDecodeError: + continue + rows.append(row) + stats = subject_stats(rows) + return { + "backend": "log_rows", + "subjects": list(stats.keys()), + "split": "", + "prompt_style": "", + "image_size": "", + "tp": "", + "tp_devices": "", + "correct": sum(item["correct"] for item in stats.values()), + "total": sum(item["total"] for item in stats.values()), + "skipped": "", + "accuracy": ( + sum(item["correct"] for item in stats.values()) + / sum(item["total"] for item in stats.values()) + if stats + else 0.0 + ), + "rows": rows, + } + + +def select_json_member(archive: tarfile.TarFile, member: str | None) -> str: + names = [item.name for item in archive.getmembers() if item.isfile()] + if member: + matches = [ + name for name in names if name == member or name.endswith("/" + member) + ] + if not matches: + raise FileNotFoundError(f"Archive member not found: {member}") + return matches[0] + json_names = [ + name + for name in names + if name.endswith(".json") and "mmmu" in Path(name).name.lower() + ] + if len(json_names) == 1: + return json_names[0] + preferred = [ + name + for name in json_names + if Path(name).name == "ernie_vl_mmmu_cpp_50_1024.json" + ] + if len(preferred) == 1: + return preferred[0] + raise ValueError( + "Could not auto-detect archive JSON member; pass --left-member/--right-member. " + f"Candidates: {json_names}" + ) + + +def rows_by_id(data: dict[str, Any]) -> dict[str, dict[str, Any]]: + rows = {} + for row in data.get("rows", []): + row_id = str(row.get("id", "")) + if row_id: + rows[row_id] = row + return rows + + +def subject_stats(rows: list[dict[str, Any]]) -> dict[str, dict[str, int]]: + stats: dict[str, dict[str, int]] = defaultdict(lambda: {"total": 0, "correct": 0}) + for row in rows: + subject = str(row.get("subject", "")) + stats[subject]["total"] += 1 + stats[subject]["correct"] += int(row.get("ok", 0)) + return dict(stats) + + +def cell(value: Any) -> str: + if value is None: + return "" + text = str(value).replace("\n", "\\n").replace("|", "\\|") + return text + + +def config_lines(label: str, data: dict[str, Any]) -> list[str]: + fields = [ + "backend", + "subjects", + "split", + "prompt_style", + "image_size", + "tp", + "tp_devices", + "correct", + "total", + "skipped", + "accuracy", + ] + lines = [f"### {label}", "", "| field | value |", "|---|---|"] + for field in fields: + lines.append(f"| {field} | {cell(data.get(field))} |") + return lines + + +def build_report( + left: dict[str, Any], + right: dict[str, Any], + left_label: str, + right_label: str, +) -> str: + left_rows = left.get("rows", []) + right_rows = right.get("rows", []) + left_by_id = rows_by_id(left) + right_by_id = rows_by_id(right) + common_ids = sorted(left_by_id.keys() & right_by_id.keys()) + left_only = sorted(left_by_id.keys() - right_by_id.keys()) + right_only = sorted(right_by_id.keys() - left_by_id.keys()) + + lines: list[str] = [ + "# MMMU Comparison", + "", + *config_lines(left_label, left), + "", + *config_lines(right_label, right), + "", + "## Subject Summary", + "", + "| subject | " + f"{left_label} correct/total | {right_label} correct/total | delta correct |", + "|---|---:|---:|---:|", + ] + left_stats = subject_stats(left_rows) + right_stats = subject_stats(right_rows) + for subject in sorted(set(left_stats) | set(right_stats)): + left_stat = left_stats.get(subject, {"correct": 0, "total": 0}) + right_stat = right_stats.get(subject, {"correct": 0, "total": 0}) + lines.append( + f"| {cell(subject)} | {left_stat['correct']}/{left_stat['total']} | " + f"{right_stat['correct']}/{right_stat['total']} | " + f"{right_stat['correct'] - left_stat['correct']} |" + ) + + lines.extend( + [ + "", + "## Row Diff", + "", + f"Common rows: {len(common_ids)}", + f"Left-only rows: {len(left_only)}", + f"Right-only rows: {len(right_only)}", + "", + "| id | subject | gold | " + f"{left_label} pred/ok | {right_label} pred/ok | pred same | ok same | " + f"{left_label} boxed/explicit | {right_label} boxed/explicit |", + "|---|---|---|---|---|---:|---:|---|---|", + ] + ) + for row_id in common_ids: + lrow = left_by_id[row_id] + rrow = right_by_id[row_id] + pred_same = lrow.get("pred") == rrow.get("pred") + ok_same = str(lrow.get("ok")) == str(rrow.get("ok")) + lines.append( + f"| {cell(row_id)} | {cell(lrow.get('subject') or rrow.get('subject'))} | " + f"{cell(lrow.get('gold') or rrow.get('gold'))} | " + f"{cell(lrow.get('pred'))}/{cell(lrow.get('ok'))} | " + f"{cell(rrow.get('pred'))}/{cell(rrow.get('ok'))} | " + f"{int(pred_same)} | {int(ok_same)} | " + f"{cell(lrow.get('boxed_answer'))}/{cell(lrow.get('explicit_answer'))} | " + f"{cell(rrow.get('boxed_answer'))}/{cell(rrow.get('explicit_answer'))} |" + ) + + if left_only: + lines.extend(["", "Left-only ids:", "", ", ".join(left_only)]) + if right_only: + lines.extend(["", "Right-only ids:", "", ", ".join(right_only)]) + lines.append("") + return "\n".join(lines) + + +def main() -> int: + args = parse_args() + left = read_json(args.left, args.left_member) + right = read_json(args.right, args.right_member) + report = build_report(left, right, args.left_label, args.right_label) + if args.output: + Path(args.output).write_text(report, encoding="utf-8") + else: + sys.stdout.write(report) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/mmmu_run_status.py b/tools/mmmu_run_status.py new file mode 100644 index 000000000..63e63b5c7 --- /dev/null +++ b/tools/mmmu_run_status.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python3 +"""Summarize an in-progress ERNIE MMMU smoke run.""" + +from __future__ import annotations + +import argparse +import json +import re +from collections import defaultdict +from pathlib import Path +from typing import Any + +ERROR_RE = re.compile( + r"OOM|out of memory|killed|NCCL|Xid|Traceback|RuntimeError|Exception", + re.IGNORECASE, +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "prefix", + help="Run output prefix, or log path ending in .log", + ) + parser.add_argument("--recent", type=int, default=5) + return parser.parse_args() + + +def log_path(prefix_arg: str) -> Path: + path = Path(prefix_arg) + if path.suffix == ".log": + return path + return path.with_suffix(".log") + + +def parse_log(path: Path) -> tuple[list[dict[str, Any]], list[str]]: + rows: list[dict[str, Any]] = [] + errors: list[str] = [] + if not path.exists(): + return rows, [f"log not found: {path}"] + for line in path.read_text(errors="replace").splitlines(): + if line.startswith("{") and '"subject"' in line: + try: + rows.append(json.loads(line)) + continue + except json.JSONDecodeError: + pass + if ERROR_RE.search(line): + errors.append(line[:1000]) + return rows, errors + + +def subject_stats(rows: list[dict[str, Any]]) -> dict[str, dict[str, int]]: + stats: dict[str, dict[str, int]] = defaultdict(lambda: {"total": 0, "correct": 0}) + for row in rows: + subject = str(row.get("subject", "")) + stats[subject]["total"] += 1 + stats[subject]["correct"] += int(row.get("ok", 0)) + return dict(stats) + + +def main() -> int: + args = parse_args() + log = log_path(args.prefix) + prefix = log.with_suffix("") + rows, errors = parse_log(log) + print(f"log={log}") + print(f"log_exists={log.exists()}") + print(f"log_bytes={log.stat().st_size if log.exists() else 0}") + print(f"json_exists={prefix.with_suffix('.json').exists()}") + print(f"csv_exists={prefix.with_suffix('.csv').exists()}") + print(f"rows={len(rows)}") + print( + f"subjects={json.dumps(subject_stats(rows), ensure_ascii=False, sort_keys=True)}" + ) + print(f"errors={len(errors)}") + for error in errors[-args.recent :]: + print(f"error: {error}") + print("recent_rows:") + for row in rows[-args.recent :]: + subset = { + key: row.get(key) + for key in [ + "subject", + "id", + "gold", + "pred", + "ok", + "new_tokens", + "hit_max_new_tokens", + "elapsed_sec", + "boxed_answer", + "explicit_answer", + ] + } + print(json.dumps(subset, ensure_ascii=False, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())