WIP bagel vision understanding patches

Author: francoishernandez

eole/bin/convert/HF_mappings.py

@@ -37,10 +37,45 @@
 MODEL_OVERRIDES = {
     "LlamaForCausalLM": {},  # default
     "MistralForCausalLM": {},
-    "Qwen2ForCausalLM": {
+    "Qwen2ForCausalLM": { # for bagel, but we need to add some conditions to keep supporting real qwen2...
+        "decoder_layer_prefix": "language_model.model.layers.",
+        "decoder.layer_norm.weight": "language_model.model.norm.weight",
+        "encoder_layer_prefix": "vit_model.vision_model.encoder.layers.",
+        "encoder.patch_conv.weight": "vit_model.vision_model.embeddings.patch_embedding.weight",
+        "encoder.patch_conv.bias": "vit_model.vision_model.embeddings.patch_embedding.bias",
+        "encoder.position_embeddings.weight": "vit_model.vision_model.embeddings.position_embedding.weight",
+        "encoder.post_layernorm.weight": "vit_model.vision_model.post_layernorm.weight",
+        "encoder.post_layernorm.bias": "vit_model.vision_model.post_layernorm.bias",
+        "tgt_emb.embeddings.weight": "language_model.model.embed_tokens.weight",
+        "generator.weight": "language_model.lm_head.weight",
+        # vision_adapter
+        "adapter.w_in.weight": "connector.fc1.weight",
+        "adapter.w_in.bias": "connector.fc1.bias",
+        "adapter.w_out.weight": "connector.fc2.weight",
+        "adapter.w_out.bias": "connector.fc2.bias",
+        "vit_pos_embed.pos_embed": "vit_pos_embed.pos_embed",
+        "decoder": {
+            ".self_attn.q_norm.": ".self_attn.q_norm.",
+            ".self_attn.k_norm.": ".self_attn.k_norm.",
+        },
+        "encoder": {
+            ".self_attn.linear_query.": ".self_attn.q_proj.",
+            ".self_attn.linear_keys.": ".self_attn.k_proj.",
+            ".self_attn.linear_values.": ".self_attn.v_proj.",
+            ".self_attn.final_linear.": ".self_attn.out_proj.",
+            ".mlp.gate_up_proj.": ".mlp.fc1.",
+            ".mlp.down_proj.": ".mlp.fc2.",
+            ".input_layernorm.": ".layer_norm1.",
+            ".post_attention_layernorm.": ".layer_norm2.",
+        },
         "config": {
             "add_qkvbias": True,
             "add_final_linear_bias": False,
+            # "ffn_layernorm": True,
+            "decoder": {
+                "query_norm": True,
+                "key_norm": True,
+            },
         }
     },
     "Qwen3ForCausalLM": {
@@ -353,6 +388,7 @@
         "Gemma2ForCausalLM": "gemma-rms",
         "M2M100ForConditionalGeneration": "standard",
         "Gemma3ForConditionalGeneration": "gemma-rms",
+        "Qwen2ForCausalLM": "rms",
     },
 )
 
@@ -386,6 +422,7 @@
         "Mistral3ForConditionalGeneration": VisionTransformerLMModelConfig,
         "Gemma3ForConditionalGeneration": VisionTransformerLMModelConfig,
         "M2M100ForConditionalGeneration": TransformerModelConfig,
+        "Qwen2ForCausalLM": VisionTransformerLMModelConfig,
     },
 )
 

eole/bin/convert/convert_HF.py

@@ -66,6 +66,7 @@ class HuggingfaceFiles:
     wmap_path: Optional[str] = None
     model_path: Optional[str] = None
     special_tokens_json: Optional[str] = None
+    vision_config_path: Optional[str] = None
 
     # Unified dictionary to cache loaded files
     _loaded_files: dict = field(default_factory=dict, init=False)
@@ -117,7 +118,7 @@ def download_file_from_hub(file_name, required=True):
 
         # Fetch required and optional files
         paths = {
-            "config_path": get_file_fn("config.json", required=True),
+            "config_path": get_file_fn("llm_config.json", required=False), # hard patch for bagel
             "tokenizer_config_json": get_file_fn("tokenizer_config.json", required=True),
             "generation_config_json": get_file_fn("generation_config.json", required=False),
             "tokenizer_model": get_file_fn("tokenizer.model", required=False)
@@ -126,8 +127,9 @@ def download_file_from_hub(file_name, required=True):
             "wmap_path": get_file_fn("model.safetensors.index.json", required=False)
             or get_file_fn("pytorch_model.bin.index.json", required=False),
             "model_path": get_file_fn("model.safetensors", required=False)
-            or get_file_fn("pytorch_model.bin", required=False),
+            or get_file_fn("pytorch_model.bin", required=False) or get_file_fn("ema.safetensors", required=False),
             "special_tokens_json": get_file_fn("special_tokens_map.json", required=False),
+            "vision_config_path": get_file_fn("vit_config.json", required=False),
         }
 
         return cls(**paths, model_dir=args.model_dir, token=args.token)
@@ -270,9 +272,13 @@ def build_config_dict(hf):
     arch = hf.arch
     print("Architecture: ", arch)
 
-    vision_config = config.get("vision_config", None)
-    other_config = config  # save what is not text/vision for later use
-    config = config.get("text_config", config)
+    vision_config = getattr(hf, "vision_config", None)
+    if vision_config is None:
+        vision_config = config.get("vision_config", None)
+        other_config = config  # save what is not text/vision for later use
+        config = config.get("text_config", config)
+
+    print("VISION_CONFIG:", vision_config)
 
     model_config = {}
     training_config = {}
@@ -289,6 +295,7 @@ def build_config_dict(hf):
         "transformer_ff_moe": config.get("moe_intermediate_size", None),
         "mlp_activation_fn": ACT_TABLE[arch],
         "layer_norm": LN_TABLE[arch],
+        # TODO: this can break encoder (e.g. bagel)
         "heads_kv": config.get("multi_query", False)
         or config.get(
             "num_key_value_heads",
@@ -351,6 +358,30 @@ def build_config_dict(hf):
         model_config["projector_activation_fn"] = other_config.get("projector_hidden_act", "gelu")
         model_config["spatial_merge_size"] = other_config.get("spatial_merge_size", None)
 
+    if arch == "Qwen2ForCausalLM":
+        model_config["adapter"] = "bagel"
+        model_config["encoder"] = {
+            "mlp_activation_fn": "gelu-tanh",  # no up_proj it seems
+            "hidden_size": vision_config.get("hidden_size", 1152),
+            # "image_size": vision_config["image_size"],
+            "image_size": 1024,
+            "patch_size": vision_config["patch_size"],
+            "heads": vision_config["num_attention_heads"],
+            "heads_kv": vision_config["num_attention_heads"],
+            "layers": 26, # 27 in config, but actually 26 in safetensors...
+            "transformer_ff": vision_config["intermediate_size"],
+            # siglip style learned position embeddings (like gemma3)
+            "position_encoding_type": PositionEncodingType.Learned,
+            "n_positions": (vision_config["image_size"] // vision_config["patch_size"]) ** 2,
+            "add_ffnbias": True,
+            "add_final_linear_bias": True,
+            "add_qkvbias": True,
+            "layer_norm": "standard",
+            "patch_conv_bias": True,
+            "layernorm_pre": False,  # implies post layernorm
+            "image_token_id": 151654,
+        }
+
     if arch == "Gemma3ForConditionalGeneration":
         if model_config.get("head_dim", None) is None:
             model_config["head_dim"] = 256  # src/transformers/models/gemma3/configuration_gemma3.py#L61

eole/decoders/transformer.py

@@ -110,6 +110,7 @@ def forward(self, layer_in, **kwargs):
         return_attn = kwargs.pop("return_attn", False)
         position_embeddings = kwargs.pop("position_embeddings", None)
 
+
         norm_layer_in = self.input_layernorm(layer_in)
 
         self_attn, attns = self.self_attn(
@@ -161,7 +162,8 @@ def forward(self, layer_in, **kwargs):
                 ctx_attn = 0
             ff_in = self.post_attention_layernorm(ctx_attn + self_attn + layer_in)
         # we apply residual with un-normed
-        layer_out = self.mlp(ff_in) + layer_in + self_attn + ctx_attn
+        MLP = self.mlp(ff_in)
+        layer_out = MLP + layer_in + self_attn + ctx_attn
 
         return layer_out, attns
 
@@ -311,7 +313,8 @@ def forward(self, emb, **kwargs):
         step = kwargs.pop("step", None)
         with_align = kwargs.pop("with_align", False)
         return_attn = with_align or kwargs.pop("return_attn", False)
-        position_embeddings = self.rope.update(emb.size(1), step=step)
+        positions = kwargs.pop("positions", None)
+        position_embeddings = self.rope.update(emb.size(1), step=step, positions=positions)
         if self.rope_local is not None:
             position_embeddings_local = self.rope_local.update(emb.size(1), step=step)
         else:
@@ -339,7 +342,7 @@ def forward(self, emb, **kwargs):
         # we need to adapt the mask for gemma3, TODO: find another condition?
         # SEEMS OK TO MASK IMAGES FOR LLAVA TOO ?
         if decoder_in is not None and attn_mask is not None:
-            attn_mask = self._update_causal_mask(attn_mask, decoder_in == image_token_id)
+            attn_mask = self._update_causal_mask(attn_mask, (decoder_in == image_token_id) | (decoder_in == 151652) | (decoder_in == 151653))
         if self.sliding_window > 0 and step >= self.sliding_window and attn_mask is not None:
             attn_mask = attn_mask[:, :, :, -self.sliding_window :]
 

eole/encoders/transformer.py

@@ -40,7 +40,7 @@ def __init__(
         )
         self.dropout = nn.Dropout(self.dropout_p)
         self.post_attention_layernorm = LayerNorm[encoder_config.layer_norm](
-            encoder_config.hidden_size, eps=encoder_config.norm_eps
+            encoder_config.hidden_size, eps=encoder_config.norm_eps, bias=True
         )
         self.mlp = MLP(
             encoder_config,

eole/encoders/vision.py

@@ -67,6 +67,15 @@ def position_ids_in_meshgrid(patch_embeds_list, max_width, flatten=True):
         return torch.stack(positions)
 
 
+
+# from bagel
+def get_flattened_position_ids_extrapolate(img_h, img_w, patch_size, max_num_patches_per_side):
+    num_patches_h, num_patches_w = img_h // patch_size, img_w // patch_size
+    coords_h = torch.arange(0, num_patches_h)
+    coords_w = torch.arange(0, num_patches_w)
+    pos_ids = (coords_h[:, None] * max_num_patches_per_side + coords_w).flatten()
+    return pos_ids
+
 def create_block_diagonal_mask(lengths, device):
     """
     Create a block diagonal mask based on sequence lengths.
@@ -88,6 +97,18 @@ def create_block_diagonal_mask(lengths, device):
     return mask.to(device)
 
 
+# grabbed from bagel repo
+
+def patchify(image, patch_size):
+    p = patch_size
+    c, h, w = image.shape
+    assert h % p == 0 and w % p == 0
+    image = image.reshape(c, h // p, p, w // p, p)
+    image = torch.einsum("chpwq->hwpqc", image)
+    image = image.reshape(-1, p**2 * c)
+    return image
+
+
 class VisionEncoder(nn.Module):
     def __init__(self, encoder_config, running_config=None):
         super(VisionEncoder, self).__init__()
@@ -99,12 +120,18 @@ def __init__(self, encoder_config, running_config=None):
             )
         else:
             self.rope = build_rope(encoder_config, mode="2d")
-        self.patch_conv = nn.Conv2d(
-            in_channels=encoder_config.num_channels,
-            out_channels=encoder_config.hidden_size,
-            kernel_size=encoder_config.patch_size,
-            stride=encoder_config.patch_size,
-            bias=encoder_config.patch_conv_bias,
+        # self.patch_conv = nn.Conv2d(
+        #     in_channels=encoder_config.num_channels,
+        #     out_channels=encoder_config.hidden_size,
+        #     kernel_size=encoder_config.patch_size,
+        #     stride=encoder_config.patch_size,
+        #     bias=encoder_config.patch_conv_bias,
+        # )
+        # linear patch conv for bagel
+        self.patch_conv = nn.Linear(
+            encoder_config.patch_size * encoder_config.patch_size * encoder_config.num_channels,
+            encoder_config.hidden_size,
+            bias=True,
         )
         if encoder_config.layernorm_pre:
             self.ln_pre = RMSNorm(encoder_config.hidden_size, eps=1e-5)
@@ -133,7 +160,8 @@ def from_config(cls, encoder_config, running_config=None):
 
     @property
     def max_patches_per_side(self):
-        return self.encoder_config.image_size // self.encoder_config.patch_size
+        return 70 # hardcoded bagel value
+        # return self.encoder_config.image_size // self.encoder_config.patch_size
 
     @property
     def device(self):
@@ -151,8 +179,10 @@ def forward(self, images):
         # TODO add as @property somewhere
         dtype = next(self.parameters()).dtype
 
+        pixel_values = [patchify(img, self.encoder_config.patch_size) for img in images]
+
         # pass images through initial convolution independently (because they may have different sizes)
-        patch_embeds_list = [self.patch_conv(img.to(dtype)) for img in images]
+        patch_embeds_list = [self.patch_conv(pv.to(dtype)) for pv in pixel_values]
 
         if self.ln_pre is not None:  # pixtral / mistral
             # flatten H+W then change to (H+W, C) and stack all images of ex
@@ -171,17 +201,32 @@ def forward(self, images):
             patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
             mask = None
 
+        patch_embeds = patch_embeds.transpose(1, 2)  # (N_img, Seqlen, D)
+
         # positional embeddings
-        positions = position_ids_in_meshgrid(
-            patch_embeds_list,
-            max_width=self.encoder_config.image_size // self.encoder_config.patch_size,
-            flatten=self.ln_pre is not None,  # dirty flag need to improve
-        ).to(self.device)
+        # positions = position_ids_in_meshgrid(
+        #     # patch_embeds_list,
+        #     images,
+        #     max_width=self.encoder_config.image_size // self.encoder_config.patch_size,
+        #     flatten=self.ln_pre is not None,  # dirty flag need to improve
+        # ).to(self.device)
+        positions = torch.cat([
+            get_flattened_position_ids_extrapolate(
+                img.shape[-2],
+                img.shape[-1],
+                self.encoder_config.patch_size,
+                self.max_patches_per_side,
+
+            )
+            for img in images
+            ], axis=0).unsqueeze(0).to(self.device)
+
         # TODO: make this cleaner
         if hasattr(self, "position_embeddings"):
             # this is only used for rope
             position_embeddings = None
-            patch_embeds += self.position_embeddings(positions)
+            pos_embeds = self.position_embeddings(positions)
+            patch_embeds += pos_embeds
         else:
             position_embeddings = self.rope.update(
                 patch_embeds.size(1),
@@ -197,7 +242,7 @@ def forward(self, images):
         if self.post_layernorm is not None:
             out = self.post_layernorm(out)
 
-        return out
+        return out, positions
 
 
 # Multi-Modal Projector
@@ -266,4 +311,5 @@ def from_config(cls, model_config, running_config=None):
 str2adapter = {
     "llava": VisionLanguageAdapter,
     "gemma3": Gemma3MultiModalProjector,
+    "bagel": VisionLanguageAdapter,
 }