diff --git a/csrc/config/model_config.hpp b/csrc/config/model_config.hpp index dc0e89287..b27717863 100644 --- a/csrc/config/model_config.hpp +++ b/csrc/config/model_config.hpp @@ -7,11 +7,17 @@ #include #include #include +#include #include #include #include namespace infinilm::config { + +inline size_t json_size(const nlohmann::json &config, const char *key, size_t fallback = 0) { + return config.contains(key) ? config.at(key).get() : fallback; +} + class ModelConfig { // Model config is implemented using nlohmann/json and is primarily used for advanced configuration // beyond the standard model config. It is initialized via ModelConfig(const std::string& path) diff --git a/csrc/models/qwen3_vl/qwen3_vl_for_conditional_generation.cpp b/csrc/models/qwen3_vl/qwen3_vl_for_conditional_generation.cpp index 512fe3079..3f7c41130 100644 --- a/csrc/models/qwen3_vl/qwen3_vl_for_conditional_generation.cpp +++ b/csrc/models/qwen3_vl/qwen3_vl_for_conditional_generation.cpp @@ -1,6 +1,8 @@ #include "qwen3_vl_for_conditional_generation.hpp" + #include "../../global_state/global_state.hpp" #include "../models_registry.hpp" + #include #include @@ -10,14 +12,56 @@ Qwen3VLModel::Qwen3VLModel(std::shared_ptr model_ const infinicore::Device &device) { nlohmann::json &config_json = model_config->get_config_json(); nlohmann::json &text_config_json = config_json["text_config"]; - std::shared_ptr text_config = std::make_shared(text_config_json); + auto text_config = std::make_shared(text_config_json); + const auto &dtype{model_config->get_dtype()}; INFINICORE_NN_MODULE_INIT(language_model, text_config, device); + INFINICORE_NN_MODULE_INIT(visual, config_json["vision_config"], dtype, device); +} + +void Qwen3VLModel::replace_image_embeddings_(infinicore::Tensor inputs_embeds, + const infinicore::Tensor &image_bound, + const infinicore::Tensor &image_embeds) const { + auto bound_cpu = image_bound->to(infinicore::Device::cpu()); + const int64_t *bound = reinterpret_cast(bound_cpu->data()); + const size_t start = static_cast(bound[0]); + const size_t end = static_cast(bound[1]); + if (end < start || end > inputs_embeds->size(1)) { + throw std::runtime_error("Qwen3VLModel: invalid image_bound"); + } + const size_t image_tokens = end - start; + if (image_tokens != image_embeds->size(0)) { + throw std::runtime_error("Qwen3VLModel: image_bound and image features do not match"); + } + const size_t hidden_size = inputs_embeds->size(2); + inputs_embeds->narrow({{1, start, image_tokens}}) + ->copy_from(image_embeds->view({1, image_tokens, hidden_size})); } infinicore::Tensor Qwen3VLModel::forward(const infinilm::InfinilmModel::Input &input) const { - auto hidden_states = language_model_->forward(input); - return hidden_states; + auto input_ids = input.input_ids.value(); + if (input.pixel_values.has_value() && !input.pixel_values->empty()) { + if (!input.image_grid_thw.has_value() || !input.image_bound.has_value()) { + throw std::runtime_error("Qwen3VLModel: image_grid_thw and image_bound must be provided with pixel_values"); + } + if (input.pixel_values->size() != input.image_grid_thw->size() || input.pixel_values->size() != input.image_bound->size()) { + throw std::runtime_error("Qwen3VLModel: pixel_values, image_grid_thw and image_bound must have the same number of elements"); + } + auto inputs_embeds = language_model_->embed_tokens(input_ids); + auto input_offsets_cpu = input.input_offsets.value()->to(infinicore::Device::cpu()); + const int32_t *offsets = reinterpret_cast(input_offsets_cpu->data()); + auto req_ids = infinilm::global_state::get_forward_context().mm_metadata.image_req_ids; + for (size_t i = 0; i < input.pixel_values->size(); ++i) { + size_t req_id = req_ids.has_value() ? req_ids->at(i) : i; + auto image_embeds = visual_->forward(input.pixel_values->at(i), input.image_grid_thw->at(i)); + size_t start = static_cast(offsets[req_id]); + size_t len = static_cast(offsets[req_id + 1] - offsets[req_id]); + auto embeds_slice = inputs_embeds->narrow({{1, start, len}}); + replace_image_embeddings_(embeds_slice, input.image_bound->at(i), image_embeds); + } + return language_model_->forward_embeds(inputs_embeds, input.position_ids.value()); + } + return language_model_->forward(input); } Qwen3VLForConditionalGeneration::Qwen3VLForConditionalGeneration(std::shared_ptr model_config, @@ -66,9 +110,17 @@ std::shared_ptr create_qwen3_vl_model_config(std: nlohmann::json &config_json = model_config->get_config_json(); nlohmann::json &text_config_json = config_json["text_config"]; - if (!config_json.contains("torch_dtype")) { - std::string dtype = text_config_json["dtype"]; - config_json["torch_dtype"] = dtype; + if (!config_json.contains("torch_dtype") || config_json["torch_dtype"].is_null()) { + config_json["torch_dtype"] = text_config_json.value("dtype", "bfloat16"); + } + if (!config_json.contains("dtype") || config_json["dtype"].is_null()) { + config_json["dtype"] = text_config_json.value("dtype", "bfloat16"); + } + if (!text_config_json.contains("torch_dtype") || text_config_json["torch_dtype"].is_null()) { + text_config_json["torch_dtype"] = text_config_json.value("dtype", "bfloat16"); + } + if (!text_config_json.contains("model_type") || text_config_json["model_type"] == "qwen3_vl_text") { + text_config_json["model_type"] = "qwen3"; } return model_config; } diff --git a/csrc/models/qwen3_vl/qwen3_vl_for_conditional_generation.hpp b/csrc/models/qwen3_vl/qwen3_vl_for_conditional_generation.hpp index dae5bc39b..ecbe51c0d 100644 --- a/csrc/models/qwen3_vl/qwen3_vl_for_conditional_generation.hpp +++ b/csrc/models/qwen3_vl/qwen3_vl_for_conditional_generation.hpp @@ -1,6 +1,10 @@ #pragma once +#include "../../layers/linear/linear.hpp" #include "../../models/qwen3/qwen3_for_causal_lm.hpp" +#include "infinicore/nn/module.hpp" +#include "infinicore/tensor.hpp" +#include "qwen3_vl_vision.hpp" namespace infinilm::models::qwen3_vl { @@ -14,7 +18,12 @@ class Qwen3VLModel : public infinicore::nn::Module { infinicore::Tensor forward(const infinilm::InfinilmModel::Input &input) const; protected: + void replace_image_embeddings_(infinicore::Tensor inputs_embeds, + const infinicore::Tensor &image_bound, + const infinicore::Tensor &image_embeds) const; + INFINICORE_NN_MODULE(Qwen3VLTextModel, language_model); + INFINICORE_NN_MODULE(Qwen3VLVisionModel, visual); }; class Qwen3VLForConditionalGeneration : public InfinilmModel { diff --git a/csrc/models/qwen3_vl/qwen3_vl_position.cpp b/csrc/models/qwen3_vl/qwen3_vl_position.cpp new file mode 100644 index 000000000..41395d787 --- /dev/null +++ b/csrc/models/qwen3_vl/qwen3_vl_position.cpp @@ -0,0 +1,209 @@ +#include "qwen3_vl_position.hpp" + +#include "../../utils.hpp" + +#include +#include +#include +#include +#include +#include + +namespace infinilm::models::qwen3_vl { +namespace { + +size_t grid_num_patches(const int64_t *grid, size_t num_grids) { + size_t total_patches = 0; + for (size_t i = 0; i < num_grids; ++i) { + total_patches += static_cast(grid[i * 3]) + * static_cast(grid[i * 3 + 1]) + * static_cast(grid[i * 3 + 2]); + } + return total_patches; +} + +void validate_grid(size_t h, size_t w, size_t spatial_merge_size) { + if (spatial_merge_size == 0 || h % spatial_merge_size != 0 || w % spatial_merge_size != 0) { + throw std::runtime_error("Qwen3VLPositionBuilder: image grid must be divisible by spatial_merge_size"); + } +} + +float load_scalar(const std::byte *ptr, infinicore::DataType dtype) { + switch (dtype) { + case infinicore::DataType::F32: + return *reinterpret_cast(ptr); + case infinicore::DataType::F16: + return f16_to_f32(*reinterpret_cast(ptr)); + case infinicore::DataType::BF16: + return bf16_to_f32(*reinterpret_cast(ptr)); + default: + throw std::runtime_error("Qwen3VLPositionBuilder: unsupported pos_embed dtype"); + } +} + +} // namespace + +Qwen3VLPositionBuilder::Qwen3VLPositionBuilder(size_t hidden_size, + size_t spatial_merge_size, + size_t num_grid_per_side, + size_t num_heads, + const infinicore::DataType &dtype, + const infinicore::Device &device) + : hidden_size_(hidden_size), + spatial_merge_size_(spatial_merge_size), + num_grid_per_side_(num_grid_per_side), + num_heads_(num_heads), + dtype_(dtype), + device_(device) {} + +infinicore::Tensor Qwen3VLPositionBuilder::values_to_device_(const std::vector &values, + const infinicore::Shape &shape) const { + if (dtype_ == infinicore::DataType::F32) { + auto cpu = infinicore::Tensor::from_blob(const_cast(values.data()), shape, dtype_, infinicore::Device::cpu()); + return cpu->to(device_); + } + + std::vector packed(values.size()); + if (dtype_ == infinicore::DataType::BF16) { + for (size_t i = 0; i < values.size(); ++i) { + packed[i] = f32_to_bf16(values[i]); + } + } else if (dtype_ == infinicore::DataType::F16) { + for (size_t i = 0; i < values.size(); ++i) { + packed[i] = f32_to_f16(values[i]); + } + } else { + throw std::runtime_error("Qwen3VLPositionBuilder: unsupported dtype for generated position tables"); + } + + auto cpu = infinicore::Tensor::from_blob(packed.data(), shape, dtype_, infinicore::Device::cpu()); + return cpu->to(device_); +} + +infinicore::Tensor Qwen3VLPositionBuilder::position_embeddings(const infinicore::Tensor &image_grid_thw, + const infinicore::nn::Embedding &pos_embed) const { + auto grid_cpu = image_grid_thw->to(infinicore::Device::cpu()); + auto weight_cpu = pos_embed.weight()->to(infinicore::Device::cpu()); + const int64_t *grid = reinterpret_cast(grid_cpu->data()); + const size_t num_grids = grid_cpu->size(0); + const size_t total_patches = grid_num_patches(grid, num_grids); + const size_t num_positions = weight_cpu->size(0); + const size_t hidden_size = weight_cpu->size(1); + const size_t elem_size = weight_cpu->element_size(); + const auto dtype = weight_cpu->dtype(); + const auto *weight = weight_cpu->data(); + std::vector values(total_patches * hidden_size); + + size_t offset = 0; + for (size_t g = 0; g < num_grids; ++g) { + const size_t t = static_cast(grid[g * 3]); + const size_t h = static_cast(grid[g * 3 + 1]); + const size_t w = static_cast(grid[g * 3 + 2]); + validate_grid(h, w, spatial_merge_size_); + + for (size_t token = 0; token < t * h * w; ++token) { + const size_t frame_offset = token % (h * w); + const size_t merged_idx = frame_offset / (spatial_merge_size_ * spatial_merge_size_); + const size_t intra = frame_offset % (spatial_merge_size_ * spatial_merge_size_); + const size_t merged_h = merged_idx / (w / spatial_merge_size_); + const size_t merged_w = merged_idx % (w / spatial_merge_size_); + const size_t ph = merged_h * spatial_merge_size_ + intra / spatial_merge_size_; + const size_t pw = merged_w * spatial_merge_size_ + intra % spatial_merge_size_; + + const float h_pos = h > 1 ? static_cast(ph) * static_cast(num_grid_per_side_ - 1) / static_cast(h - 1) : 0.0f; + const float w_pos = w > 1 ? static_cast(pw) * static_cast(num_grid_per_side_ - 1) / static_cast(w - 1) : 0.0f; + const size_t h_floor = static_cast(std::floor(h_pos)); + const size_t w_floor = static_cast(std::floor(w_pos)); + const size_t h_ceil = std::min(h_floor + 1, num_grid_per_side_ - 1); + const size_t w_ceil = std::min(w_floor + 1, num_grid_per_side_ - 1); + const float dh = h_pos - static_cast(h_floor); + const float dw = w_pos - static_cast(w_floor); + + const size_t pos_ids[4] = { + h_floor * num_grid_per_side_ + w_floor, + h_floor * num_grid_per_side_ + w_ceil, + h_ceil * num_grid_per_side_ + w_floor, + h_ceil * num_grid_per_side_ + w_ceil, + }; + const float weights[4] = { + (1.0f - dh) * (1.0f - dw), + (1.0f - dh) * dw, + dh * (1.0f - dw), + dh * dw, + }; + + for (size_t k = 0; k < 4; ++k) { + if (pos_ids[k] >= num_positions) { + throw std::runtime_error("Qwen3VLPositionBuilder: generated position id is out of range"); + } + } + + for (size_t d = 0; d < hidden_size; ++d) { + float value = 0.0f; + for (size_t k = 0; k < 4; ++k) { + const auto *src = weight + (pos_ids[k] * hidden_size + d) * elem_size; + value += weights[k] * load_scalar(src, dtype); + } + values[offset * hidden_size + d] = value; + } + ++offset; + } + } + + return values_to_device_(values, {total_patches, hidden_size}); +} + +std::tuple +Qwen3VLPositionBuilder::rotary_embeddings(const infinicore::Tensor &image_grid_thw) const { + auto grid_cpu = image_grid_thw->to(infinicore::Device::cpu()); + const int64_t *grid = reinterpret_cast(grid_cpu->data()); + const size_t num_grids = grid_cpu->size(0); + const size_t total_patches = grid_num_patches(grid, num_grids); + + auto pos_cpu = infinicore::Tensor::empty({total_patches}, infinicore::DataType::I64, infinicore::Device::cpu()); + auto *pos = reinterpret_cast(pos_cpu->data()); + for (size_t i = 0; i < total_patches; ++i) { + pos[i] = static_cast(i); + } + + const size_t head_dim = hidden_size_ / num_heads_; + const size_t half_dim = head_dim / 2; + const size_t axis_dim = half_dim / 2; + std::vector sin_values(total_patches * half_dim); + std::vector cos_values(total_patches * half_dim); + + size_t offset = 0; + for (size_t g = 0; g < num_grids; ++g) { + const size_t t = static_cast(grid[g * 3]); + const size_t h = static_cast(grid[g * 3 + 1]); + const size_t w = static_cast(grid[g * 3 + 2]); + validate_grid(h, w, spatial_merge_size_); + + for (size_t token = 0; token < t * h * w; ++token) { + const size_t frame_offset = token % (h * w); + const size_t merged_idx = frame_offset / (spatial_merge_size_ * spatial_merge_size_); + const size_t intra = frame_offset % (spatial_merge_size_ * spatial_merge_size_); + const size_t merged_h = merged_idx / (w / spatial_merge_size_); + const size_t merged_w = merged_idx % (w / spatial_merge_size_); + const size_t py = merged_h * spatial_merge_size_ + intra / spatial_merge_size_; + const size_t px = merged_w * spatial_merge_size_ + intra % spatial_merge_size_; + + for (size_t d = 0; d < axis_dim; ++d) { + const float inv_freq = 1.0f / std::pow(10000.0f, static_cast(2 * d) / static_cast(half_dim)); + const float ay = static_cast(py) * inv_freq; + const float ax = static_cast(px) * inv_freq; + sin_values[offset * half_dim + d] = std::sin(ay); + cos_values[offset * half_dim + d] = std::cos(ay); + sin_values[offset * half_dim + axis_dim + d] = std::sin(ax); + cos_values[offset * half_dim + axis_dim + d] = std::cos(ax); + } + ++offset; + } + } + + return {pos_cpu->to(device_), + values_to_device_(sin_values, {total_patches, half_dim}), + values_to_device_(cos_values, {total_patches, half_dim})}; +} + +} // namespace infinilm::models::qwen3_vl diff --git a/csrc/models/qwen3_vl/qwen3_vl_position.hpp b/csrc/models/qwen3_vl/qwen3_vl_position.hpp new file mode 100644 index 000000000..750b49c8c --- /dev/null +++ b/csrc/models/qwen3_vl/qwen3_vl_position.hpp @@ -0,0 +1,38 @@ +#pragma once + +#include "infinicore/nn/embedding.hpp" +#include "infinicore/tensor.hpp" + +#include +#include + +namespace infinilm::models::qwen3_vl { + +class Qwen3VLPositionBuilder { +public: + Qwen3VLPositionBuilder(size_t hidden_size, + size_t spatial_merge_size, + size_t num_grid_per_side, + size_t num_heads, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor position_embeddings(const infinicore::Tensor &image_grid_thw, + const infinicore::nn::Embedding &pos_embed) const; + + std::tuple + rotary_embeddings(const infinicore::Tensor &image_grid_thw) const; + +private: + infinicore::Tensor values_to_device_(const std::vector &values, + const infinicore::Shape &shape) const; + + size_t hidden_size_; + size_t spatial_merge_size_; + size_t num_grid_per_side_; + size_t num_heads_; + infinicore::DataType dtype_; + infinicore::Device device_; +}; + +} // namespace infinilm::models::qwen3_vl diff --git a/csrc/models/qwen3_vl/qwen3_vl_vision.cpp b/csrc/models/qwen3_vl/qwen3_vl_vision.cpp new file mode 100644 index 000000000..9b5fea7e2 --- /dev/null +++ b/csrc/models/qwen3_vl/qwen3_vl_vision.cpp @@ -0,0 +1,43 @@ +#include "qwen3_vl_vision.hpp" + +#include "../../config/model_config.hpp" +#include "infinicore/ops.hpp" + +#include + +namespace infinilm::models::qwen3_vl { + +using infinilm::config::json_size; + +Qwen3VLVisionModel::Qwen3VLVisionModel(const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device) + : hidden_size_(json_size(config, "hidden_size")), + spatial_merge_size_(json_size(config, "spatial_merge_size", 2)), + num_grid_per_side_(static_cast(std::sqrt(static_cast(json_size(config, "num_position_embeddings"))))), + position_builder_(hidden_size_, + spatial_merge_size_, + num_grid_per_side_, + json_size(config, "num_heads"), + dtype, + device) { + INFINICORE_NN_MODULE_INIT(patch_embed, config, dtype, device); + INFINICORE_NN_MODULE_INIT(pos_embed, json_size(config, "num_position_embeddings"), hidden_size_, std::nullopt, dtype, device); + INFINICORE_NN_MODULE_VEC_INIT(blocks, json_size(config, "depth"), Qwen3VLVisionBlock, config, dtype, device); + INFINICORE_NN_MODULE_INIT(merger, config, false, dtype, device); + size_t deepstack_count = config.contains("deepstack_visual_indexes") ? config["deepstack_visual_indexes"].size() : 0; + INFINICORE_NN_MODULE_VEC_INIT(deepstack_merger_list, deepstack_count, Qwen3VLPatchMerger, config, true, dtype, device); +} + +infinicore::Tensor Qwen3VLVisionModel::forward(const infinicore::Tensor &pixel_values, + const infinicore::Tensor &image_grid_thw) const { + auto hidden_states = patch_embed_->forward(pixel_values); + hidden_states = infinicore::op::add(hidden_states, position_builder_.position_embeddings(image_grid_thw, *pos_embed_)); + auto [rotary_pos_ids, sin_table, cos_table] = position_builder_.rotary_embeddings(image_grid_thw); + for (const auto &block : blocks_) { + hidden_states = block->forward(hidden_states, rotary_pos_ids, sin_table, cos_table); + } + return merger_->forward(hidden_states); +} + +} // namespace infinilm::models::qwen3_vl diff --git a/csrc/models/qwen3_vl/qwen3_vl_vision.hpp b/csrc/models/qwen3_vl/qwen3_vl_vision.hpp new file mode 100644 index 000000000..914192ce9 --- /dev/null +++ b/csrc/models/qwen3_vl/qwen3_vl_vision.hpp @@ -0,0 +1,35 @@ +#pragma once + +#include "infinicore/nn/embedding.hpp" +#include "infinicore/nn/module.hpp" +#include "infinicore/tensor.hpp" +#include "qwen3_vl_position.hpp" +#include "qwen3_vl_vision_layers.hpp" + +#include + +namespace infinilm::models::qwen3_vl { + +class Qwen3VLVisionModel : public infinicore::nn::Module { +public: + Qwen3VLVisionModel(const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &pixel_values, + const infinicore::Tensor &image_grid_thw) const; + +private: + size_t hidden_size_; + size_t spatial_merge_size_; + size_t num_grid_per_side_; + Qwen3VLPositionBuilder position_builder_; + + INFINICORE_NN_MODULE(Qwen3VLPatchEmbed, patch_embed); + INFINICORE_NN_MODULE(infinicore::nn::Embedding, pos_embed); + INFINICORE_NN_MODULE_VEC(Qwen3VLVisionBlock, blocks); + INFINICORE_NN_MODULE(Qwen3VLPatchMerger, merger); + INFINICORE_NN_MODULE_VEC(Qwen3VLPatchMerger, deepstack_merger_list); +}; + +} // namespace infinilm::models::qwen3_vl diff --git a/csrc/models/qwen3_vl/qwen3_vl_vision_layers.cpp b/csrc/models/qwen3_vl/qwen3_vl_vision_layers.cpp new file mode 100644 index 000000000..f669fcdfc --- /dev/null +++ b/csrc/models/qwen3_vl/qwen3_vl_vision_layers.cpp @@ -0,0 +1,149 @@ +#include "qwen3_vl_vision_layers.hpp" + +#include "../../config/model_config.hpp" +#include "infinicore/nn/rope.hpp" +#include "infinicore/ops.hpp" +#include "infinicore/ops/mha.hpp" +#include "infinicore/ops/rope.hpp" + +#include +#include + +namespace infinilm::models::qwen3_vl { + +using infinilm::config::json_size; + +Qwen3VLPatchProjection::Qwen3VLPatchProjection(size_t out_features, + size_t in_channels, + size_t temporal_patch_size, + size_t patch_size, + const infinicore::DataType &dtype, + const infinicore::Device &device) + : patch_dim_(in_channels * temporal_patch_size * patch_size * patch_size) { + INFINICORE_NN_PARAMETER_INIT(weight, ({out_features, in_channels, temporal_patch_size, patch_size, patch_size}, dtype, device)); + INFINICORE_NN_PARAMETER_INIT(bias, ({out_features}, dtype, device)); +} + +infinicore::Tensor Qwen3VLPatchProjection::forward(const infinicore::Tensor &pixel_values) const { + auto input = const_cast(pixel_values); + auto weight_2d = static_cast(weight_)->view({weight_->size(0), patch_dim_}); + auto bias_tensor = static_cast(bias_); + return infinicore::op::linear(input, weight_2d, bias_tensor); +} + +Qwen3VLPatchEmbed::Qwen3VLPatchEmbed(const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device) { + INFINICORE_NN_MODULE_INIT(proj, + json_size(config, "hidden_size"), + json_size(config, "in_channels", 3), + json_size(config, "temporal_patch_size", 2), + json_size(config, "patch_size", 16), + dtype, + device); +} + +infinicore::Tensor Qwen3VLPatchEmbed::forward(const infinicore::Tensor &pixel_values) const { + return proj_->forward(pixel_values); +} + +Qwen3VLVisionMLP::Qwen3VLVisionMLP(const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device) { + size_t hidden_size = json_size(config, "hidden_size"); + size_t intermediate_size = json_size(config, "intermediate_size"); + INFINICORE_NN_MODULE_INIT(linear_fc1, hidden_size, intermediate_size, true, dtype, device); + INFINICORE_NN_MODULE_INIT(linear_fc2, intermediate_size, hidden_size, true, dtype, device); +} + +infinicore::Tensor Qwen3VLVisionMLP::forward(const infinicore::Tensor &hidden_states) const { + auto x = linear_fc1_->forward(const_cast(hidden_states)); + x = infinicore::op::gelu_tanh(x); + return linear_fc2_->forward(x); +} + +Qwen3VLVisionAttention::Qwen3VLVisionAttention(const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device) + : hidden_size_(json_size(config, "hidden_size")), + num_heads_(json_size(config, "num_heads")), + head_dim_(hidden_size_ / num_heads_), + scale_(1.0f / std::sqrt(static_cast(head_dim_))) { + INFINICORE_NN_MODULE_INIT(qkv, hidden_size_, hidden_size_ * 3, true, dtype, device); + INFINICORE_NN_MODULE_INIT(proj, hidden_size_, hidden_size_, true, dtype, device); +} + +infinicore::Tensor Qwen3VLVisionAttention::forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &position_ids, + const infinicore::Tensor &sin_table, + const infinicore::Tensor &cos_table) const { + size_t seq_len = hidden_states->size(0); + auto qkv = qkv_->forward(const_cast(hidden_states)) + ->view({seq_len, 3, num_heads_, head_dim_}); + auto q = qkv->narrow({{1, 0, 1}})->squeeze(1); + auto k = qkv->narrow({{1, 1, 1}})->squeeze(1); + auto v = qkv->narrow({{1, 2, 1}})->squeeze(1)->view({1, seq_len, num_heads_, head_dim_}); + + q = infinicore::op::rope(q, position_ids, sin_table, cos_table, infinicore::nn::RoPE::Algo::GPT_NEOX) + ->view({1, seq_len, num_heads_, head_dim_}); + k = infinicore::op::rope(k, position_ids, sin_table, cos_table, infinicore::nn::RoPE::Algo::GPT_NEOX) + ->view({1, seq_len, num_heads_, head_dim_}); + + auto attn_output = infinicore::op::mha(q, k, v, std::nullopt, scale_, false) + ->view({seq_len, hidden_size_}); + return proj_->forward(attn_output); +} + +Qwen3VLVisionBlock::Qwen3VLVisionBlock(const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device) { + size_t hidden_size = json_size(config, "hidden_size"); + INFINICORE_NN_MODULE_INIT(norm1, hidden_size, 1e-6, dtype, device); + INFINICORE_NN_MODULE_INIT(attn, config, dtype, device); + INFINICORE_NN_MODULE_INIT(norm2, hidden_size, 1e-6, dtype, device); + INFINICORE_NN_MODULE_INIT(mlp, config, dtype, device); +} + +infinicore::Tensor Qwen3VLVisionBlock::forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &position_ids, + const infinicore::Tensor &sin_table, + const infinicore::Tensor &cos_table) const { + auto residual = hidden_states; + auto x = norm1_->forward(hidden_states); + x = attn_->forward(x, position_ids, sin_table, cos_table); + x = infinicore::op::add(x, residual); + + residual = x; + x = norm2_->forward(x); + x = mlp_->forward(x); + return infinicore::op::add(x, residual); +} + +Qwen3VLPatchMerger::Qwen3VLPatchMerger(const nlohmann::json &config, + bool use_postshuffle_norm, + const infinicore::DataType &dtype, + const infinicore::Device &device) + : hidden_size_(json_size(config, "hidden_size")), + merged_size_(hidden_size_ * json_size(config, "spatial_merge_size", 2) * json_size(config, "spatial_merge_size", 2)), + use_postshuffle_norm_(use_postshuffle_norm) { + size_t out_hidden_size = json_size(config, "out_hidden_size"); + INFINICORE_NN_MODULE_INIT(norm, use_postshuffle_norm_ ? merged_size_ : hidden_size_, 1e-6, dtype, device); + INFINICORE_NN_MODULE_INIT(linear_fc1, merged_size_, merged_size_, true, dtype, device); + INFINICORE_NN_MODULE_INIT(linear_fc2, merged_size_, out_hidden_size, true, dtype, device); +} + +infinicore::Tensor Qwen3VLPatchMerger::forward(const infinicore::Tensor &hidden_states) const { + infinicore::Tensor x; + if (use_postshuffle_norm_) { + x = hidden_states->view({hidden_states->size(0) / 4, merged_size_}); + x = norm_->forward(x); + } else { + x = norm_->forward(hidden_states); + x = x->view({x->size(0) / 4, merged_size_}); + } + x = linear_fc1_->forward(x); + x = infinicore::op::gelu(x); + return linear_fc2_->forward(x); +} + +} // namespace infinilm::models::qwen3_vl diff --git a/csrc/models/qwen3_vl/qwen3_vl_vision_layers.hpp b/csrc/models/qwen3_vl/qwen3_vl_vision_layers.hpp new file mode 100644 index 000000000..d656e2468 --- /dev/null +++ b/csrc/models/qwen3_vl/qwen3_vl_vision_layers.hpp @@ -0,0 +1,110 @@ +#pragma once + +#include "../../layers/linear/linear.hpp" +#include "infinicore/nn/layer_norm.hpp" +#include "infinicore/nn/module.hpp" +#include "infinicore/tensor.hpp" + +#include + +namespace infinilm::models::qwen3_vl { + +class Qwen3VLPatchProjection : public infinicore::nn::Module { +public: + Qwen3VLPatchProjection(size_t out_features, + size_t in_channels, + size_t temporal_patch_size, + size_t patch_size, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &pixel_values) const; + +private: + size_t patch_dim_; + INFINICORE_NN_PARAMETER(weight); + INFINICORE_NN_PARAMETER(bias); +}; + +class Qwen3VLPatchEmbed : public infinicore::nn::Module { +public: + Qwen3VLPatchEmbed(const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &pixel_values) const; + +private: + INFINICORE_NN_MODULE(Qwen3VLPatchProjection, proj); +}; + +class Qwen3VLVisionMLP : public infinicore::nn::Module { +public: + Qwen3VLVisionMLP(const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + +private: + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, linear_fc1); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, linear_fc2); +}; + +class Qwen3VLVisionAttention : public infinicore::nn::Module { +public: + Qwen3VLVisionAttention(const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &position_ids, + const infinicore::Tensor &sin_table, + const infinicore::Tensor &cos_table) const; + +private: + size_t hidden_size_; + size_t num_heads_; + size_t head_dim_; + float scale_; + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, qkv); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, proj); +}; + +class Qwen3VLVisionBlock : public infinicore::nn::Module { +public: + Qwen3VLVisionBlock(const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &position_ids, + const infinicore::Tensor &sin_table, + const infinicore::Tensor &cos_table) const; + +private: + INFINICORE_NN_MODULE(infinicore::nn::LayerNorm, norm1); + INFINICORE_NN_MODULE(Qwen3VLVisionAttention, attn); + INFINICORE_NN_MODULE(infinicore::nn::LayerNorm, norm2); + INFINICORE_NN_MODULE(Qwen3VLVisionMLP, mlp); +}; + +class Qwen3VLPatchMerger : public infinicore::nn::Module { +public: + Qwen3VLPatchMerger(const nlohmann::json &config, + bool use_postshuffle_norm, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + +private: + size_t hidden_size_; + size_t merged_size_; + bool use_postshuffle_norm_; + INFINICORE_NN_MODULE(infinicore::nn::LayerNorm, norm); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, linear_fc1); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, linear_fc2); +}; + +} // namespace infinilm::models::qwen3_vl diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py index 3e874046e..57f8bad5d 100644 --- a/python/infinilm/infer_engine.py +++ b/python/infinilm/infer_engine.py @@ -26,6 +26,16 @@ def _apply_torch_dtype_defaults(config: dict) -> dict: return config +def _get_config_dtype(config_dict): + dtype = config_dict.get("torch_dtype") or config_dict.get("dtype") + if dtype is not None: + return dtype + text_config = config_dict.get("text_config") + if isinstance(text_config, dict): + return text_config.get("dtype") or text_config.get("torch_dtype") + return None + + def _normalize_videonsa_config(config_dict): model_type = config_dict.get("model_type") @@ -55,6 +65,11 @@ def read_hf_config(model_path): with open(config_path, "r") as f: config_dict = json.load(f) + if config_dict.get("torch_dtype") is None and config_dict.get("dtype") is None: + dtype = _get_config_dtype(config_dict) + if dtype is not None: + config_dict["torch_dtype"] = dtype + if "model_type" not in config_dict: raise ValueError( f"`model_type` is not specified in the config file `{config_path}`." diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index 8f84c9b97..2010597cc 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -196,6 +196,7 @@ def load_model_state_dict_by_file( already_loaded_keys = [] embed_tokens_torch_unscaled = None + qwen3_vl_embed_tokens_torch = None weights_processed = False remapper = _WEIGHT_REMAPPER.get(model_type) @@ -241,6 +242,10 @@ def load_model_state_dict_by_file( model_param["model.embed_tokens.weight"] = ( embed_tokens_torch_unscaled * float(scale_emb) ) + if "model.language_model.embed_tokens.weight" in model_param: + qwen3_vl_embed_tokens_torch = model_param[ + "model.language_model.embed_tokens.weight" + ] # --------------------------------------------------------- # # model_param_infini references torch.Tensor @@ -282,6 +287,10 @@ def load_model_state_dict_by_file( model_params["model.embed_tokens.weight"] = ( embed_tokens_torch_unscaled * float(scale_emb) ) + if "model.language_model.embed_tokens.weight" in model_params: + qwen3_vl_embed_tokens_torch = model_params[ + "model.language_model.embed_tokens.weight" + ].to(dtype=torch_dtype) model_param_infini = {} for key in model_params.keys(): @@ -301,12 +310,18 @@ def load_model_state_dict_by_file( # Handle tied weights: if lm_head.weight is missing, share embed_tokens.weight # Use unscaled weight for lm_head (C++ alpha handles dim_model_base scaling) if "lm_head.weight" in model_keys and "lm_head.weight" not in already_loaded_keys: - if embed_tokens_torch_unscaled is not None: - lm_head_tensor = infinicore.from_torch(embed_tokens_torch_unscaled) + tied_weight = ( + embed_tokens_torch_unscaled + if embed_tokens_torch_unscaled is not None + else qwen3_vl_embed_tokens_torch + ) + if tied_weight is not None: + lm_head_tensor = infinicore.from_torch(tied_weight) model.load_state_dict({"lm_head.weight": lm_head_tensor}, strict=False) already_loaded_keys.append("lm_head.weight") del lm_head_tensor embed_tokens_torch_unscaled = None + qwen3_vl_embed_tokens_torch = None gc.collect() check_parameters(model_keys, already_loaded_keys) diff --git a/python/infinilm/processors/qwen3_vl_processor.py b/python/infinilm/processors/qwen3_vl_processor.py new file mode 100644 index 000000000..1baa08666 --- /dev/null +++ b/python/infinilm/processors/qwen3_vl_processor.py @@ -0,0 +1,232 @@ +import torch +from transformers import AutoConfig, AutoProcessor +from typing_extensions import override + +from ..llm.scheduler import SchedulerOutput +from ..llm.static_scheduler import StaticSchedulerOutput +from .processor import register_processor +from .qwen3_5_processor import Qwen35Processor + + +@register_processor("qwen3_vl") +class Qwen3VLProcessor(Qwen35Processor): + def __init__(self, model_dir_path: str): + self.processor = AutoProcessor.from_pretrained( + model_dir_path, trust_remote_code=True + ) + self.tokenizer = self.processor.tokenizer + self.config = AutoConfig.from_pretrained(model_dir_path, trust_remote_code=True) + self.image_token_id = self.config.image_token_id + vision_config = getattr(self.config, "vision_config", None) + self.spatial_merge_size = int(getattr(vision_config, "spatial_merge_size", 2)) + text_config = getattr(self.config, "text_config", None) + dtype_name = getattr(text_config, "dtype", None) or getattr( + text_config, "torch_dtype", None + ) + if isinstance(dtype_name, torch.dtype): + self.pixel_values_dtype = dtype_name + elif dtype_name is not None: + self.pixel_values_dtype = getattr(torch, str(dtype_name)) + else: + self.pixel_values_dtype = torch.bfloat16 + + @override + def __call__( + self, + prompt, + images=None, + videos=None, + audios=None, + return_tensors: str = None, + **kwargs, + ) -> dict: + if not images and not videos and not audios: + return self.tokenizer(prompt, return_tensors=return_tensors, **kwargs) + + processor_kwargs = {"text": [prompt], "return_tensors": "pt", **kwargs} + if images: + processor_kwargs["images"] = images + if videos: + processor_kwargs["videos"] = videos + results = self.processor(**processor_kwargs) + self._append_image_bound(results) + return results + + @override + def apply_chat_template( + self, + conversation, + add_generation_prompt: bool = False, + tokenize: bool = True, + **kwargs, + ): + normalized = [] + for msg in conversation: + content = msg["content"] + if not isinstance(content, list): + normalized.append(msg) + continue + + normalized_content = [] + for item in content: + if item.get("type") == "text": + normalized_content.append(item) + elif item.get("type") == "image_url": + normalized_content.append( + {"type": "image", "image": item["image_url"]["url"]} + ) + else: + normalized_content.append(item) + normalized.append( + {"role": msg.get("role", "user"), "content": normalized_content} + ) + + return self.processor.apply_chat_template( + conversation=normalized, + add_generation_prompt=add_generation_prompt, + tokenize=tokenize, + **kwargs, + ) + + def _append_static_mm_inputs( + self, model_inputs: dict, req, num_cached: int + ) -> None: + processed_inputs = req.processed_inputs + if processed_inputs is None or "pixel_values" not in processed_inputs: + return + + import infinicore + + grid_thw = processed_inputs.get("image_grid_thw") + bounds = processed_inputs.get("image_bound") + if grid_thw is None or bounds is None: + raise RuntimeError( + "Qwen3-VL image input requires image_grid_thw and image_bound" + ) + + grids = self._as_grid_list(grid_thw) + + def split_pixel_values(pixel_value, grids): + if isinstance(pixel_value, list): + return [ + p if isinstance(p, torch.Tensor) else torch.as_tensor(p) + for p in pixel_value + ] + pixel_tensor = ( + pixel_value + if isinstance(pixel_value, torch.Tensor) + else torch.as_tensor(pixel_value) + ) + if len(grids) == 1: + return [pixel_tensor] + counts = [int(g[0].item() * g[1].item() * g[2].item()) for g in grids] + return list(torch.split(pixel_tensor, counts, dim=0)) + + pixels = split_pixel_values(processed_inputs["pixel_values"], grids) + bounds = bounds if isinstance(bounds, torch.Tensor) else torch.as_tensor(bounds) + if bounds.ndim == 3 and bounds.shape[0] == 1: + bounds = bounds.squeeze(0) + if bounds.ndim != 2 or bounds.shape[-1] != 2: + raise RuntimeError("Qwen3-VL image_bound must have shape [num_images, 2]") + if len(pixels) != len(grids) or len(pixels) != bounds.shape[0]: + raise RuntimeError("Qwen3-VL image input count mismatch") + + partial_cached = (bounds[:, 0] < num_cached) & (bounds[:, 1] > num_cached) + if partial_cached.any().item(): + raise RuntimeError("Qwen3-VL does not support partially cached image spans") + + pixel_values = [] + image_grid_thw = [] + image_bound = [] + for image_idx, (pixel, grid) in enumerate(zip(pixels, grids)): + if bounds[image_idx, 1].item() <= num_cached: + continue + pixel_values.append( + infinicore.from_torch( + ( + pixel + if isinstance(pixel, torch.Tensor) + else torch.as_tensor(pixel) + ).to(self.pixel_values_dtype) + ) + ) + image_grid_thw.append( + infinicore.from_torch( + grid if isinstance(grid, torch.Tensor) else torch.as_tensor(grid) + ) + ) + image_bound.append( + infinicore.from_torch((bounds[image_idx] - num_cached).to(torch.int64)) + ) + + if pixel_values: + model_inputs["pixel_values"] = pixel_values + model_inputs["image_grid_thw"] = image_grid_thw + model_inputs["image_bound"] = image_bound + model_inputs["image_req_ids"] = [0] * len(pixel_values) + + @override + def build_model_inputs( + self, + scheduler_output: SchedulerOutput | StaticSchedulerOutput, + temperature: float = 1.0, + top_p: float = 0.8, + top_k: int = 1, + **kwargs, + ) -> dict: + if isinstance(scheduler_output, StaticSchedulerOutput): + model_inputs = self._build_model_input_from_static_scheduler_output( + scheduler_output, temperature, top_p, top_k + ) + if scheduler_output.is_prefill: + self._append_static_mm_inputs( + model_inputs, + scheduler_output.scheduled_requests[0], + scheduler_output.prefix_hit_len, + ) + elif isinstance(scheduler_output, SchedulerOutput): + model_inputs = self._build_model_input_from_batch_scheduler_output( + scheduler_output, temperature, top_p, top_k + ) + self._append_qwen35_mm_inputs(model_inputs, scheduler_output) + else: + raise ValueError( + "scheduler_output must be an instance of SchedulerOutput or StaticSchedulerOutput" + ) + + return model_inputs + + @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, **kwargs + ): + mappings = [] + idx = 0 + image_idx = 0 + while idx < len(prompt_token_ids): + if prompt_token_ids[idx] != self.image_token_id: + idx += 1 + continue + start = idx + while ( + idx < len(prompt_token_ids) + and prompt_token_ids[idx] == self.image_token_id + ): + idx += 1 + mappings.append( + { + "start_index": start, + "end_index": idx, + "identifier": ( + image_ids[image_idx] + if image_ids and image_idx < len(image_ids) + else image_idx + ), + } + ) + image_idx += 1 + return mappings