diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index f6441b8d2662..a57efb324b5b 100755 --- a/convert_hf_to_gguf.py +++ b/convert_hf_to_gguf.py @@ -2630,6 +2630,147 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter yield from super().modify_tensors(data_torch, name, bid) +@ModelBase.register("FalconOCRForCausalLM") +class FalconOCRModel(TextModel): + model_arch = gguf.MODEL_ARCH.FALCON_OCR + + def set_vocab(self): + self._set_vocab_gpt2() + # this model does not actually use the chat template, but we need to make sure to avoid any additional formatting + self.gguf_writer.add_chat_template("{% for m in messages %}{{ m['content'] + '\\n' }}{% endfor %}") + + def set_gguf_parameters(self): + super().set_gguf_parameters() + hparams = self.hparams + + self.gguf_writer.add_context_length(hparams["max_seq_len"]) + self.gguf_writer.add_feed_forward_length(hparams["ffn_dim"]) + + # head_dim (64) differs from n_embd/n_heads (768/16=48) + self.gguf_writer.add_key_length(hparams["head_dim"]) + self.gguf_writer.add_value_length(hparams["head_dim"]) + + self.gguf_writer.add_layer_norm_rms_eps(hparams.get("norm_eps", 1e-5)) + self.gguf_writer.add_rope_freq_base(hparams.get("rope_theta", 10000)) + self.gguf_writer.add_rope_dimension_count(hparams["head_dim"] // 2) + self.gguf_writer.add_add_bos_token(False) + + # important: because "golden" rope must be applied to fit Q shape, + # we must force number of KV heads to be the same as number of Q heads + self.gguf_writer.add_head_count_kv(hparams["n_heads"]) # not n_kv_heads + + def tensor_force_quant(self, name, new_name, bid, n_dims): + if "freqs" in name: + return gguf.GGMLQuantizationType.F32 + return super().tensor_force_quant(name, new_name, bid, n_dims) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if "img_projector" in name: + return + + if name == "freqs_cis_golden": + # original shape: [n_heads, rope_dim // 2, 2] + # permute to [2, n_heads, rope_dim//2] so h-freqs and w-freqs are contiguous, + # then flatten to [2, n_heads * rope_dim//2] + # ggml loads this as ne[0]=n_heads*rope_dim//2, ne[1]=2 + data_torch = data_torch.permute(2, 0, 1).contiguous().reshape(2, -1) + # ggml_rope_ext computes theta = pos_int / freq_factor (freq_base=1.0) + # pos_int is fixed-point: pos_int = actual_pos * 1e6 + # golden rope needs theta = freqs_actual * actual_pos = freqs_actual * pos_int / 1e6 + # => freq_factor = 1e6 / freqs_actual + data_torch = 1e6 / data_torch + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FREQS), data_torch) + return + + # Deinterleave fused w13 into separate gate (even rows) and up (odd rows) + if "feed_forward.w13" in name: + gate = data_torch[0::2, :] + up = data_torch[1::2, :] + yield from super().modify_tensors(gate, name.replace("w13", "w1"), bid) + yield from super().modify_tensors(up, name.replace("w13", "w3"), bid) + return + + # Unfused w1 needs sqrt(2) scaling to match reference numerics + if "feed_forward.w1" in name: + data_torch = data_torch * math.sqrt(2.0) + yield from super().modify_tensors(data_torch, name, bid) + return + + yield from super().modify_tensors(data_torch, name, bid) + + +@ModelBase.register("FalconOCRForCausalLM") +class FalconOCRMmprojModel(MmprojModel): + has_vision_encoder = True + + # Important: Falcon OCR model does not actually have a vision encoder, + # the image patches are projected directly to the text embedding space and fed to the text encoder. + # we only use the mmproj to store the image projector weights and preprocessor config + + def __init__(self, dir_model: Path, *args, **kwargs): + # Inject synthetic vision_config / text_config for MmprojModel base class + hparams = ModelBase.load_hparams(dir_model, False) + hparams["text_config"] = {"hidden_size": hparams["dim"]} + hparams["vision_config"] = { + "hidden_size": hparams["dim"], + "patch_size": hparams["spatial_patch_size"], + "image_size": 1024, + "intermediate_size": hparams["dim"], + "num_attention_heads": 1, + "num_hidden_layers": 0, # no actual vision encoder layers + } + super().__init__(dir_model, *args, hparams=hparams, **kwargs) + + def set_gguf_parameters(self): + self.gguf_writer.add_file_type(self.ftype) + self.gguf_writer.add_clip_has_vision_encoder(True) + self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.FALCON_OCR) + self.gguf_writer.add_vision_projection_dim(self.n_embd_text) + self.gguf_writer.add_vision_patch_size(self.global_config["spatial_patch_size"]) + self.gguf_writer.add_vision_image_size(1024) + self.gguf_writer.add_vision_embedding_length(self.global_config["dim"]) + self.gguf_writer.add_vision_image_mean([0.5, 0.5, 0.5]) + self.gguf_writer.add_vision_image_std([0.5, 0.5, 0.5]) + self.gguf_writer.add_vision_min_pixels(64 * 64) + self.gguf_writer.add_vision_max_pixels(1024 * 1024) + self.gguf_writer.add_vision_head_count(1) + self.gguf_writer.add_vision_feed_forward_length(1) + self.gguf_writer.add_vision_block_count(0) + self.gguf_writer.add_vision_attention_layernorm_eps(1e-5) + + def tensor_force_quant(self, name, new_name, bid, n_dims): + if "tok_embeddings" in name: + return gguf.GGMLQuantizationType.F32 + return super().tensor_force_quant(name, new_name, bid, n_dims) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if name == "img_projector.weight": + # The HF linear weight [n_embd, patch_dim] has patch_dim = H*W*C with C + # fastest (from einops rearrange). Rearrange to PyTorch conv2d format + # [C_out, C_in, kH, kW] so ggml_conv_2d can use it directly. The planar + # image data fed to ggml matches this convention. + ps = self.global_config["spatial_patch_size"] # 16 + ch = self.global_config["channel_size"] # 3 + n_embd = data_torch.shape[0] + w = data_torch.reshape(n_embd, ps, ps, ch) # [n_embd, H, W, C] + w = w.permute(0, 3, 1, 2).contiguous() # [n_embd, C, H, W] + w = w.reshape(data_torch.shape) # flatten back to 2-D + yield (self.format_tensor_name(gguf.MODEL_TENSOR.V_MMPROJ, bid=0), w) + return + + if name == "tok_embeddings.weight": + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True) + prefix_str = "<|image_cls|><|image_reg_1|><|image_reg_2|><|image_reg_3|><|image_reg_4|>" + ids = tokenizer.encode(prefix_str, add_special_tokens=False) + prefix_embd = data_torch[ids].contiguous() + logger.info(f"Extracted {len(ids)} prefix embeddings (token IDs: {ids})") + yield (self.format_tensor_name(gguf.MODEL_TENSOR.V_TOK_IMG_BEGIN, suffix=""), prefix_embd) + return + + return + + @ModelBase.register("GPTBigCodeForCausalLM") class StarCoderModel(TextModel): model_arch = gguf.MODEL_ARCH.STARCODER diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index c5297a2f440f..4b30595d6ac7 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -494,6 +494,7 @@ class MODEL_ARCH(IntEnum): LLAMA_EMBED = auto() MAINCODER = auto() KIMI_LINEAR = auto() + FALCON_OCR = auto() class VISION_PROJECTOR_TYPE(IntEnum): @@ -980,6 +981,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.LLAMA_EMBED: "llama-embed", MODEL_ARCH.MAINCODER: "maincoder", MODEL_ARCH.KIMI_LINEAR: "kimi-linear", + MODEL_ARCH.FALCON_OCR: "falcon-ocr", } VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = { @@ -3877,6 +3879,18 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_DOWN_SHEXP, MODEL_TENSOR.FFN_UP_SHEXP, ], + MODEL_ARCH.FALCON_OCR: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_SINKS, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], # TODO } @@ -4137,7 +4151,8 @@ class VisionProjectorType: GLM4V = "glm4v" YOUTUVL = "youtuvl" NEMOTRON_V2_VL = "nemotron_v2_vl" - HUNYUANOCR = "hunyuanocr" + HUNYUANOCR = "hunyuanocr" + FALCON_OCR = "falcon-ocr" # Items here are (block size, type size) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 01a9b236000b..c5b41c2c0647 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -230,6 +230,7 @@ class TensorNameMap: "layers.{bid}.attn.Wqkv", # modern-bert "model.layers.{bid}.self_attn.language_expert_query_key_value", # cogvlm "model.layers.{bid}.linear_attn.in_proj_qkv", # qwen3.5 + "layers.{bid}.attention.wqkv", # falcon-ocr ), # Attention query @@ -357,6 +358,7 @@ class TensorNameMap: MODEL_TENSOR.ATTN_SINKS: ( "model.layers.{bid}.self_attn.sinks", # openai-moe "model.layers.{bid}.self_attn.attention_sink_bias", # mimov2 + "layers.{bid}.attention.sinks", # falcon_ocr ), MODEL_TENSOR.ATTN_GATE: ( diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 121c21fed957..495aeecf3ccb 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -67,6 +67,7 @@ add_library(llama models/exaone.cpp models/exaone4.cpp models/falcon-h1.cpp + models/falcon-ocr.cpp models/falcon.cpp models/gemma-embedding.cpp models/gemma.cpp diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 6904b9c1a645..7b7a59fccead 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -63,6 +63,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_MAMBA2, "mamba2" }, { LLM_ARCH_JAMBA, "jamba" }, { LLM_ARCH_FALCON_H1, "falcon-h1" }, + { LLM_ARCH_FALCON_OCR, "falcon-ocr" }, { LLM_ARCH_XVERSE, "xverse" }, { LLM_ARCH_COMMAND_R, "command-r" }, { LLM_ARCH_COHERE2, "cohere2" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index c4aabab7e0cf..135c8b7a455f 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -67,6 +67,7 @@ enum llm_arch { LLM_ARCH_MAMBA2, LLM_ARCH_JAMBA, LLM_ARCH_FALCON_H1, + LLM_ARCH_FALCON_OCR, LLM_ARCH_XVERSE, LLM_ARCH_COMMAND_R, LLM_ARCH_COHERE2, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index ee0c29235cd2..3a8a748c4476 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2070,7 +2070,7 @@ void llama_context::output_reorder() { // uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { - if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_KIMI_LINEAR || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE) { + if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_KIMI_LINEAR || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_FALCON_OCR) { return std::max(n_tokens * 40, 32u * model.n_tensors()); } uint32_t res = std::max(1024u, 8u*model.n_tensors()); diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 8e2b6ab8e7e1..dc6efeadbc18 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1201,8 +1201,9 @@ ggml_tensor * llm_graph_context::build_ffn( if (down) { cur = build_lora_mm(down, cur); - if (arch == LLM_ARCH_GLM4 || arch == LLM_ARCH_GLM4_MOE || arch == LLM_ARCH_JAIS2) { + if (arch == LLM_ARCH_GLM4 || arch == LLM_ARCH_GLM4_MOE || arch == LLM_ARCH_JAIS2 || arch == LLM_ARCH_FALCON_OCR) { // GLM4, GLM4_MOE, and JAIS2 seem to have numerical issues with half-precision accumulators + // Falcon-OCR's ReLU^2 activation produces values > 65504 (FP16 max), causing overflow. ggml_mul_mat_set_prec(cur, GGML_PREC_F32); } } diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 09102f549c8e..1ff1cbeed19a 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1088,7 +1088,8 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & bool llama_kv_cache::get_can_shift() const { // Step35 uses per-layer RoPE dims; K-shift assumes a single global n_rot. - if (model.arch == LLM_ARCH_STEP35) { + // falcon-ocr uses custom version of M-RoPE + if (model.arch == LLM_ARCH_STEP35 || model.arch == LLM_ARCH_FALCON_OCR) { return false; } if (hparams.n_pos_per_embd() > 1) { diff --git a/src/llama-model.cpp b/src/llama-model.cpp index d2ffc1f45f41..bcf0093fc6c4 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2936,6 +2936,15 @@ void llama_model::load_hparams(llama_model_loader & ml) { default: type = LLM_TYPE_UNKNOWN; } } break; + case LLM_ARCH_FALCON_OCR: + { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + + switch (hparams.n_layer) { + case 22: type = LLM_TYPE_SMALL; break; + default: type = LLM_TYPE_UNKNOWN; + } + } break; default: throw std::runtime_error("unsupported model architecture: " + arch_name()); } @@ -7969,6 +7978,34 @@ bool llama_model::load_tensors(llama_model_loader & ml) { layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0); layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); + } + } break; + case LLM_ARCH_FALCON_OCR: + { + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0); + + // note: the model doesn't actually use GQA due to "golden" rope enforcing Q dimension + const int64_t n_head_kv_ratio = 2; + const int64_t n_embd_qkv = (n_embd_head_k * n_head) + + n_embd_k_gqa / n_head_kv_ratio + + n_embd_v_gqa / n_head_kv_ratio; + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_head * n_rot/2, 2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0)); + + layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, n_embd_qkv}, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0); + + layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, i), {n_head}, TENSOR_NOT_REQUIRED); + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, 0); layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); @@ -9262,6 +9299,10 @@ ggml_cgraph * llama_model::build_graph(const llm_graph_params & params) const { { llm = std::make_unique(*this, params); } break; + case LLM_ARCH_FALCON_OCR: + { + llm = std::make_unique(*this, params); + } break; default: GGML_ABORT("fatal error"); } @@ -9518,6 +9559,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_QWEN2VL: case LLM_ARCH_PADDLEOCR: + case LLM_ARCH_FALCON_OCR: // note: falcon-ocr uses a variant of m-rope return LLAMA_ROPE_TYPE_MROPE; case LLM_ARCH_QWEN3VL: case LLM_ARCH_QWEN3VLMOE: diff --git a/src/models/falcon-ocr.cpp b/src/models/falcon-ocr.cpp new file mode 100644 index 000000000000..0a7fcd99085d --- /dev/null +++ b/src/models/falcon-ocr.cpp @@ -0,0 +1,221 @@ +#include "models.h" + +#include + +static ggml_tensor * rope_golden_axis( + ggml_context * ctx0, + ggml_tensor * cur, // [n_embd/2] + ggml_tensor * freqs, // [n_embd/4] + ggml_tensor * pos // [n_token] +) { + auto n_dim = cur->ne[0]; + return ggml_rope_ext( + ctx0, cur, pos, freqs, + n_dim, 0, 0, + 1.0f, // freq_base (ignored because we provide freqs directly) + 1.0f, // freq_scale + 0.0f, 1.0f, 0.0f, 0.0f + ); +} + +static ggml_tensor * rope_falcon( + ggml_context * ctx0, + ggml_tensor * cur, + ggml_tensor * freqs, + ggml_tensor * pos, + float freq_base, + float freq_scale +) { + // falcon-ocr style RoPE: + // - first half of head_dim rotates as normal (1D temporal RoPE) + // - second half of head_dim rotates with "golden" RoPE (2D RoPE with freq_h and freq_w) + // theta = freq_h * pos_h + freq_w * pos_w + + // the tricks for "golden" rope are: + // - we decompose it into 2 rotations: first rotate by freq_h * pos_h, then rotate by freq_w * pos_w + // - instead of rotating per-head, we rotate the whole n_embd_half (because "golden" freq different for each head, but ggml_rope only supports one set of freqs broadcasted across all heads) + + const int64_t n_dim = cur->ne[0]; + const int64_t n_head = cur->ne[1]; + const int64_t n_pos = cur->ne[2]; + + GGML_ASSERT(pos->type == GGML_TYPE_I32); + GGML_ASSERT(pos->ne[0] == n_pos * 4); // must be m-rope format + ggml_tensor * pos_t = ggml_view_1d(ctx0, pos, n_pos, 0); + ggml_tensor * pos_y = ggml_view_1d(ctx0, pos, n_pos, ggml_row_size(pos->type, n_pos)); + ggml_tensor * pos_x = ggml_view_1d(ctx0, pos, n_pos, ggml_row_size(pos->type, n_pos * 2)); + + // first half + ggml_tensor * first; + { + first = ggml_view_3d(ctx0, cur, + n_dim/2, n_head, n_pos, + cur->nb[1], + cur->nb[2], + 0); + first = ggml_rope_ext( + ctx0, first, pos_t, nullptr, + n_dim/2, GGML_ROPE_TYPE_NORMAL, + 0, + freq_base, + freq_scale, + 0.0f, 1.0f, 0.0f, 0.0f + ); + } + + // second half + ggml_tensor * second; + { + const int64_t n_embd_half = n_dim * n_head / 2; + // printf("shape of cur: %d x %d x %d\n", (int)cur->ne[0], (int)cur->ne[1], (int)cur->ne[2]); + // printf("shape of freqs: %d x %d x %d\n", (int)freqs->ne[0], (int)freqs->ne[1], (int)freqs->ne[2]); + // printf("n_embd_half: %d\n", (int)n_embd_half); + // freqs shape: ne[0]=n_head*n_rot/2, ne[1]=2 + // layout: all h-freqs contiguous first, then all w-freqs + GGML_ASSERT(freqs->type == GGML_TYPE_F32); + GGML_ASSERT(freqs->ne[0] == n_embd_half / 2 && freqs->ne[1] == 2); + // n_embd_half/2 = n_head * n_rot/2 (matches conversion: permute(2,0,1).reshape(2,-1)) + ggml_tensor * freqs_h = ggml_view_1d(ctx0, freqs, n_embd_half / 2, 0); + ggml_tensor * freqs_w = ggml_view_1d(ctx0, freqs, n_embd_half / 2, ggml_row_size(freqs->type, n_embd_half / 2)); + + second = ggml_view_3d(ctx0, cur, + n_dim/2, n_head, n_pos, + cur->nb[1], + cur->nb[2], + ggml_row_size(cur->type, n_dim/2)); + + // flatten head dim; keep n_pos on ne[2] so ggml_rope_ext sees a->ne[2] == pos->ne[0] + second = ggml_cont(ctx0, second); + second = ggml_reshape_3d(ctx0, second, n_embd_half, 1, n_pos); + + // apply each axis sequentially + second = rope_golden_axis(ctx0, second, freqs_w, pos_x); + second = rope_golden_axis(ctx0, second, freqs_h, pos_y); + + // unflatten head dim + second = ggml_reshape_3d(ctx0, second, n_dim/2, n_head, n_pos); + } + + cur = ggml_concat(ctx0, first, second, 0); + return cur; +} + +llm_build_falcon_ocr::llm_build_falcon_ocr(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { + const int64_t n_embd_head = hparams.n_embd_head_v(); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + GGML_ASSERT(n_embd_head == n_rot * 2); + + ggml_tensor * cur; + ggml_tensor * inpL; + + inpL = build_inp_embd(model.tok_embd); + + int sections[4]; + std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections); + + // inp_pos - contains the positions + ggml_tensor * inp_pos = build_inp_pos(); + + auto * inp_attn = build_attn_inp_kv(); + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + for (int il = 0; il < n_layer; ++il) { + // Parameterless RMSNorm (no learned weight) + cur = ggml_rms_norm(ctx0, inpL, hparams.f_norm_rms_eps); + cb(cur, "attn_norm", il); + + { + // note: model doesn't actually use GQA due to "golden" rope enforcing Q dimension + const int64_t n_head_kv_ratio = 2; + const int64_t n_head_kv = n_head / n_head_kv_ratio; + const int64_t n_embd_q = n_embd_head * n_head; + const int64_t n_embd_k = n_embd_head * n_head_kv; + const int64_t n_embd_v = n_embd_head * n_head_kv; + + cur = build_lora_mm(model.layers[il].wqkv, cur); + cb(cur, "wqkv", il); + + ggml_tensor * Qcur = ggml_view_3d(ctx0, cur, + n_embd_head, n_head, n_tokens, n_embd_head * sizeof(float), + cur->nb[1], ggml_row_size(cur->type, n_embd_q)); + ggml_tensor * Kcur = ggml_view_3d(ctx0, cur, + n_embd_head, n_head_kv, n_tokens, n_embd_head * sizeof(float), + cur->nb[1], ggml_row_size(cur->type, n_embd_k)); + ggml_tensor * Vcur = ggml_view_3d(ctx0, cur, + n_embd_head, n_head_kv, n_tokens, n_embd_head * sizeof(float), + cur->nb[1], ggml_row_size(cur->type, n_embd_v)); + + // Parameterless QK-norm (before RoPE) + Qcur = ggml_rms_norm(ctx0, Qcur, hparams.f_norm_rms_eps); + cb(Qcur, "Qcur_normed", il); + + Kcur = ggml_rms_norm(ctx0, Kcur, hparams.f_norm_rms_eps); + cb(Kcur, "Kcur_normed", il); + + // repeat K and V to match shape of Q (required by rope_falcon) + Kcur = ggml_repeat(ctx0, Kcur, Qcur); + Vcur = ggml_repeat(ctx0, Vcur, Qcur); + + // rope + Qcur = rope_falcon(ctx0, Qcur, model.layers[il].rope_freqs, inp_pos, freq_base, freq_scale); + Kcur = rope_falcon(ctx0, Kcur, model.layers[il].rope_freqs, inp_pos, freq_base, freq_scale); + + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + + cur = build_attn(inp_attn, + model.layers[il].wo, NULL, + Qcur, Kcur, Vcur, nullptr, model.layers[il].attn_sinks, nullptr, 1.0f/sqrtf(float(n_embd_head)), il); + } + + if (il == n_layer - 1 && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inpL = ggml_get_rows(ctx0, inpL, inp_out_ids); + } + + ggml_tensor * sa_out = ggml_add(ctx0, cur, inpL); + cb(sa_out, "sa_out", il); + + // Parameterless pre-FFN RMSNorm + cur = ggml_rms_norm(ctx0, sa_out, hparams.f_norm_rms_eps); + cb(cur, "ffn_norm", il); + + // Squared ReLU gating FFN + { + cur = build_ffn(cur, + model.layers[il].ffn_up, NULL, NULL, + model.layers[il].ffn_gate, NULL, NULL, + model.layers[il].ffn_down, NULL, NULL, + NULL, + LLM_FFN_RELU_SQR, LLM_FFN_PAR, il); + cb(cur, "ffn_out", il); + } + + cur = ggml_add(ctx0, cur, sa_out); + + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + + inpL = cur; + } + cur = inpL; + + // Final RMSNorm (with learned weight) + cur = build_norm(cur, + model.output_norm, NULL, + LLM_NORM_RMS, -1); + + cb(cur, "result_norm", -1); + + res->t_embd = cur; + + cur = build_lora_mm(model.output, cur); + + cb(cur, "result_output", -1); + + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} diff --git a/src/models/models.h b/src/models/models.h index a6682ebb287d..c495d194acb6 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -235,6 +235,10 @@ struct llm_build_falcon_h1 : public llm_build_mamba_base { llm_build_falcon_h1(const llama_model & model, const llm_graph_params & params); }; +struct llm_build_falcon_ocr : public llm_graph_context { + llm_build_falcon_ocr(const llama_model & model, const llm_graph_params & params); +}; + struct llm_build_gemma2_iswa : public llm_graph_context { llm_build_gemma2_iswa(const llama_model & model, const llm_graph_params & params); }; diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index 3bafde178de2..77cbb32937c1 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -18,6 +18,7 @@ add_library(mtmd models/cogvlm.cpp models/conformer.cpp models/dotsocr.cpp + models/falcon-ocr.cpp models/gemma4a.cpp models/gemma4v.cpp models/glm4v.cpp diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index 17cb703f7fbb..8481c4accaaa 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -293,6 +293,7 @@ enum projector_type { PROJECTOR_TYPE_KIMIK25, PROJECTOR_TYPE_NEMOTRON_V2_VL, PROJECTOR_TYPE_HUNYUANOCR, + PROJECTOR_TYPE_FALCON_OCR, PROJECTOR_TYPE_UNKNOWN, }; @@ -338,6 +339,7 @@ static std::map PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_KIMIK25, "kimik25"}, { PROJECTOR_TYPE_NEMOTRON_V2_VL, "nemotron_v2_vl"}, { PROJECTOR_TYPE_HUNYUANOCR, "hunyuanocr"}, + { PROJECTOR_TYPE_FALCON_OCR, "falcon-ocr"}, }; static projector_type clip_projector_type_from_string(const std::string & str) { diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index f0e8786b6601..b9de6b8edcb3 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -903,6 +903,10 @@ static ggml_cgraph * clip_image_build_graph(clip_ctx * ctx, const clip_image_f32 { builder = std::make_unique(ctx, img); } break; + case PROJECTOR_TYPE_FALCON_OCR: + { + builder = std::make_unique(ctx, img); + } break; case PROJECTOR_TYPE_KIMIK25: { builder = std::make_unique(ctx, img); @@ -1436,6 +1440,13 @@ struct clip_model_loader { hparams.set_warmup_n_tokens(28*28); // avoid OOM on warmup } break; + case PROJECTOR_TYPE_FALCON_OCR: + { + hparams.n_merge = 1; + get_u32(KEY_IMAGE_MIN_PIXELS, hparams.image_min_pixels); + get_u32(KEY_IMAGE_MAX_PIXELS, hparams.image_max_pixels); + hparams.set_warmup_n_tokens(16*16); + } break; case PROJECTOR_TYPE_DEEPSEEKOCR: { hparams.patch_size = 16; @@ -2340,6 +2351,11 @@ struct clip_model_loader { layer.conv_pw2_b = get_tensor(string_format(TN_CONV_PW2, prefix, il, "bias")); } } break; + case PROJECTOR_TYPE_FALCON_OCR: + { + model.mm_0_w = get_tensor(string_format(TN_LLAVA_PROJ, 0, "weight")); + model.mm_img_begin = get_tensor(TN_TOK_IMG_BEGIN); + } break; default: GGML_ASSERT(false && "unknown projector type"); } @@ -2800,6 +2816,7 @@ int clip_n_output_tokens_x(const struct clip_ctx * ctx, struct clip_image_f32 * case PROJECTOR_TYPE_YOUTUVL: return (img->nx / params.patch_size) / 2; case PROJECTOR_TYPE_STEP3VL: + case PROJECTOR_TYPE_FALCON_OCR: return img->nx / (params.patch_size * params.n_merge); default: break; @@ -2819,6 +2836,7 @@ int clip_n_output_tokens_y(const struct clip_ctx * ctx, struct clip_image_f32 * case PROJECTOR_TYPE_YOUTUVL: return (img->ny / params.patch_size) / 2; case PROJECTOR_TYPE_STEP3VL: + case PROJECTOR_TYPE_FALCON_OCR: return img->ny / (params.patch_size * params.n_merge); default: break; @@ -3023,6 +3041,10 @@ int clip_n_output_tokens(const struct clip_ctx * ctx, struct clip_image_f32 * im } n_patches = n; } break; + case PROJECTOR_TYPE_FALCON_OCR: + { + n_patches += clip_get_n_boi(ctx); // add number of BOI tokens + } break; default: GGML_ABORT("unsupported projector type"); } @@ -3463,6 +3485,7 @@ bool clip_image_batch_encode(clip_ctx * ctx, const int n_threads, const clip_ima case PROJECTOR_TYPE_PHI4: case PROJECTOR_TYPE_COGVLM: case PROJECTOR_TYPE_HUNYUANOCR: + case PROJECTOR_TYPE_FALCON_OCR: { // do nothing } break; @@ -3702,6 +3725,8 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { return ctx->model.hparams.projection_dim; case PROJECTOR_TYPE_GLM4V: return ctx->model.mm_ffn_down_w->ne[1]; + case PROJECTOR_TYPE_FALCON_OCR: + return ctx->model.mm_0_w->ne[1]; default: GGML_ABORT("Unknown projector type"); } @@ -3747,6 +3772,13 @@ bool clip_has_whisper_encoder(const struct clip_ctx * ctx) { } } +uint32_t clip_get_n_boi(const struct clip_ctx * ctx) { + if (ctx->proj_type() == PROJECTOR_TYPE_FALCON_OCR) { + return ctx->model.mm_img_begin->ne[1]; + } + return 0; +} + bool clip_encode_float_image (struct clip_ctx * ctx, int n_threads, float * img, int h, int w, float * vec) { clip_image_f32 clip_img; clip_img.buf.resize(h * w * 3); diff --git a/tools/mtmd/clip.h b/tools/mtmd/clip.h index a859b38658d3..6658583f42f6 100644 --- a/tools/mtmd/clip.h +++ b/tools/mtmd/clip.h @@ -116,3 +116,4 @@ void clip_image_f32_batch_add_mel(struct clip_image_f32_batch * batch, int n_mel bool clip_has_vision_encoder(const struct clip_ctx * ctx); bool clip_has_audio_encoder(const struct clip_ctx * ctx); bool clip_has_whisper_encoder(const struct clip_ctx * ctx); +uint32_t clip_get_n_boi(const struct clip_ctx * ctx); diff --git a/tools/mtmd/models/falcon-ocr.cpp b/tools/mtmd/models/falcon-ocr.cpp new file mode 100644 index 000000000000..eaa7eea0046a --- /dev/null +++ b/tools/mtmd/models/falcon-ocr.cpp @@ -0,0 +1,28 @@ +#include "models.h" + +// Important: Falcon OCR model does not actually have a vision encoder, +// the image patches are projected directly to the text embedding space and fed to the text encoder. +// we only use the mmproj to store the image projector weights and preprocessor config + +ggml_cgraph * clip_graph_falcon_ocr::build() { + ggml_tensor * inp_raw = build_inp_raw(); + + const int ps = patch_size; + const int pw = img.nx / ps; + const int ph = img.ny / ps; + const int n_patch = pw * ph; + + ggml_tensor * proj_w = ggml_reshape_4d(ctx0, model.mm_0_w, ps, ps, 3, n_embd); + + ggml_tensor * cur = ggml_conv_2d(ctx0, proj_w, inp_raw, ps, ps, 0, 0, 1, 1); + + // conv2d output [OW, OH, OC, 1] -> [n_embd, n_patch] + cur = ggml_reshape_2d(ctx0, cur, n_patch, n_embd); + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + + // prepend the BOI tokens (multiple tokens) + cur = ggml_concat(ctx0, model.mm_img_begin, cur, 1); + + ggml_build_forward_expand(gf, cur); + return gf; +} diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index 03d99e15b054..6078e6cc2c1c 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -163,3 +163,8 @@ struct clip_graph_kimik25 : clip_graph { ggml_tensor * resize_position_embeddings_3d(uint32_t interpolation_mode); }; + +struct clip_graph_falcon_ocr : clip_graph { + clip_graph_falcon_ocr(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} + ggml_cgraph * build() override; +}; diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index a56d3b35b484..98d15d037b56 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -36,8 +36,9 @@ struct mtmd_bitmap { struct mtmd_image_tokens { uint32_t nx; // number of tokens in x direction uint32_t ny; // number of tokens in y direction + uint32_t n_boi = 0; // number of begin-of-image tokens, used by falcon-ocr bool use_mrope_pos = false; // use M-RoPE position counting (the whole image is 1 temporal position) - uint32_t n_tokens() const { return nx * ny; } + uint32_t n_tokens() const { return nx * ny + n_boi; } clip_image_f32_batch batch_f32; // preprocessed image patches std::string id; // optional user-defined ID, useful for KV cache tracking @@ -45,6 +46,7 @@ struct mtmd_image_tokens { return mtmd_image_tokens{ nx, ny, + n_boi, use_mrope_pos, batch_f32.clone(), id @@ -436,6 +438,10 @@ struct mtmd_context { img_end = "<｜hy_place▁holder▁no▁101｜>"; image_preproc = std::make_unique(ctx_v); } break; + case PROJECTOR_TYPE_FALCON_OCR: + { + image_preproc = std::make_unique(ctx_v); + } break; default: throw std::runtime_error(string_format("%s: unexpected vision projector type %d\n", __func__, proj)); } @@ -793,6 +799,12 @@ struct mtmd_tokenizer { LOG_DBG("image_tokens->ny = %d\n", image_tokens->ny); LOG_DBG("batch_f32 size = %d\n", (int)image_tokens->batch_f32.entries.size()); + // used by falcon-ocr + auto n_boi = clip_get_n_boi(ctx->ctx_v); + if (n_boi > 0) { + image_tokens->n_boi = n_boi; + } + mtmd_input_chunk chunk{ MTMD_INPUT_CHUNK_TYPE_IMAGE, {}, // text tokens @@ -1025,6 +1037,7 @@ bool mtmd_decode_use_non_causal(mtmd_context * ctx, const mtmd_input_chunk * chu switch (proj_type) { case PROJECTOR_TYPE_GEMMA3: case PROJECTOR_TYPE_GEMMA4V: + case PROJECTOR_TYPE_FALCON_OCR: return true; default: return false; @@ -1042,6 +1055,7 @@ bool mtmd_decode_use_mrope(mtmd_context * ctx) { case PROJECTOR_TYPE_QWEN3VL: case PROJECTOR_TYPE_GLM4V: case PROJECTOR_TYPE_PADDLEOCR: + case PROJECTOR_TYPE_FALCON_OCR: // note: falcon-ocr uses a variant of m-rope return true; default: return false; @@ -1251,9 +1265,25 @@ size_t mtmd_image_tokens_get_ny(const mtmd_image_tokens * image_tokens) { mtmd_decoder_pos mtmd_image_tokens_get_decoder_pos(const mtmd_image_tokens * image_tokens, size_t i) { mtmd_decoder_pos pos; - pos.t = 0; - pos.x = i % image_tokens->nx; - pos.y = i / image_tokens->nx; + auto n_boi = image_tokens->n_boi; + if (n_boi > 0) { + // falcon-ocr style + if (i < n_boi) { + pos.t = i; + pos.x = 0; + pos.y = 0; + } else { + size_t idx = i - n_boi; + pos.t = n_boi; // all image tokens share the same temporal pos + pos.x = idx % image_tokens->nx; + pos.y = idx / image_tokens->nx; + } + } else { + // m-rope style + pos.t = 0; + pos.x = i % image_tokens->nx; + pos.y = i / image_tokens->nx; + } return pos; } @@ -1267,6 +1297,10 @@ llama_pos mtmd_image_tokens_get_n_pos(const mtmd_image_tokens * image_tokens) { // t is omitted as we don't support video input return std::max(image_tokens->nx, image_tokens->ny); } + if (image_tokens->n_boi > 0) { + // for falcon-ocr, temporal dimension = n_boi + 1 (for all image tokens) + return image_tokens->n_boi + 1; + } return image_tokens->n_tokens(); }