liger-kernel-nightly 0.6.2.dev20251005233549__tar.gz → 0.6.2.dev20251010092540__tar.gz

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.6.2.dev20251005233549
+Version: 0.6.2.dev20251010092540
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.6.2.dev20251005233549"
+version = "0.6.2.dev20251010092540"
 description = "Efficient Triton kernels for LLM Training"
 urls = { "Homepage" = "https://github.com/linkedin/Liger-Kernel" }
 readme = { file = "README.md", content-type = "text/markdown" }

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel/transformers/__init__.py

@@ -30,6 +30,7 @@ if TYPE_CHECKING:
     from liger_kernel.transformers.auto_model import AutoLigerKernelForCausalLM  # noqa: F401
     from liger_kernel.transformers.monkey_patch import _apply_liger_kernel  # noqa: F401
     from liger_kernel.transformers.monkey_patch import _apply_liger_kernel_to_instance  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_falcon_h1  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma2  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma3  # noqa: F401
@@ -91,6 +92,7 @@ def __getattr__(name: str):
     monkey_patch_symbols = {
         "_apply_liger_kernel",
         "_apply_liger_kernel_to_instance",
+        "apply_liger_kernel_to_falcon_h1",
         "apply_liger_kernel_to_gemma",
         "apply_liger_kernel_to_gemma2",
         "apply_liger_kernel_to_gemma3",
@@ -157,6 +159,7 @@ if _TRANSFORMERS_AVAILABLE:
             "AutoLigerKernelForCausalLM",
             "_apply_liger_kernel",
             "_apply_liger_kernel_to_instance",
+            "apply_liger_kernel_to_falcon_h1",
             "apply_liger_kernel_to_gemma",
             "apply_liger_kernel_to_gemma2",
             "apply_liger_kernel_to_gemma3",

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel/transformers/model/gemma.py

@@ -228,10 +228,11 @@ def lce_forward(
         )
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel/transformers/model/gemma2.py

@@ -252,8 +252,14 @@ def lce_forward(
             logits = logits * self.config.final_logit_softcapping
 
         loss = None
-        if labels is not None:
-            loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
+        if labels is not None or shift_labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                shift_labels=shift_labels,
+                vocab_size=self.vocab_size,
+                **kwargs,
+            )
 
     if not return_dict:
         output = (logits,) + outputs[1:]

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel/transformers/model/gemma3.py

@@ -119,8 +119,14 @@ def causal_forward(
             logits = logits / self.config.final_logit_softcapping
             logits = torch.tanh(logits)
             logits = logits * self.config.final_logit_softcapping
-        if labels is not None:
-            loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs)
+        if labels is not None or shift_labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                shift_labels=shift_labels,
+                vocab_size=self.vocab_size,
+                **loss_kwargs,
+            )
 
     if not return_dict:
         output = (logits,) + outputs[1:]
@@ -275,6 +281,25 @@ def multimodal_forward(
             # Flatten the tokens
             loss_fct = nn.CrossEntropyLoss()
 
+            flat_logits = shift_logits.view(-1, self.config.text_config.vocab_size)
+            flat_labels = shift_labels.view(-1).to(shift_logits.device)
+            loss = loss_fct(flat_logits, flat_labels)
+        elif shift_labels is not None:
+            # Upcast to float if we need to compute the loss to avoid potential precision issues
+            logits = logits.float()
+            shift_logits = logits[..., :-1, :]
+            if attention_mask is not None:
+                # we use the input attention mask to shift the logits and labels, because it is 2D.
+                # we also crop attn mask in case it is longer, which happens in PrefixTuning with peft
+                shift_attention_mask = attention_mask[:, -shift_logits.shape[1] :].to(logits.device)
+                shift_logits = shift_logits[shift_attention_mask.to(logits.device) != 0].contiguous()
+                shift_labels = shift_labels[shift_attention_mask.to(shift_labels.device) != 0].contiguous()
+            else:
+                shift_logits = shift_logits.contiguous()
+                shift_labels = shift_labels.contiguous()
+            # Flatten the tokens
+            loss_fct = nn.CrossEntropyLoss()
+
             flat_logits = shift_logits.view(-1, self.config.text_config.vocab_size)
             flat_labels = shift_labels.view(-1).to(shift_logits.device)
             loss = loss_fct(flat_logits, flat_labels)

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel/transformers/model/glm4.py

@@ -111,10 +111,11 @@ def lce_forward(
 
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel/transformers/model/glm4v.py

@@ -133,10 +133,11 @@ def lce_forward(
 
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel/transformers/model/glm4v_moe.py

@@ -134,10 +134,11 @@ def lce_forward(
 
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel/transformers/model/llama.py

@@ -248,10 +248,11 @@ def lce_forward(
 
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel/transformers/model/llama4.py

@@ -91,10 +91,11 @@ def lce_forward(
 
     else:  # if in inference mode materialize logits
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel/transformers/model/llava.py

@@ -313,9 +313,13 @@ def lce_forward(
 
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
-                logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **lm_kwargs
+                logits=logits,
+                labels=labels,
+                shift_labels=shift_labels,
+                vocab_size=self.config.text_config.vocab_size,
+                **lm_kwargs,
             )
 
     if not return_dict:

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel/transformers/model/mistral.py

@@ -115,10 +115,11 @@ def lce_forward(
         logits = self.lm_head(kept_hidden_states)
 
         loss = None
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel/transformers/model/mixtral.py

@@ -248,8 +248,14 @@ def lce_forward(
         logits = self.lm_head(kept_hidden_states)
 
         loss = None
-        if labels is not None:
-            loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
+        if labels is not None or shift_labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                shift_labels=shift_labels,
+                vocab_size=self.vocab_size,
+                **kwargs,
+            )
     aux_loss = None
     if output_router_logits:
         aux_loss = load_balancing_loss_func(

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel/transformers/model/mllama.py

@@ -239,10 +239,11 @@ def lce_forward(
 
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel/transformers/model/olmo2.py

@@ -111,10 +111,11 @@ def lce_forward(
 
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel/transformers/model/paligemma.py

@@ -379,6 +379,25 @@ def lce_forward(
             # Flatten the tokens
             loss_fct = CrossEntropyLoss()
 
+            flat_logits = shift_logits.view(-1, self.config.text_config.vocab_size)
+            flat_labels = shift_labels.view(-1).to(shift_logits.device)
+            loss = loss_fct(flat_logits, flat_labels)
+        elif shift_labels is not None:
+            # Upcast to float if we need to compute the loss to avoid potential precision issues
+            logits = logits.float()
+            shift_logits = logits[..., :-1, :]
+            if attention_mask is not None:
+                # we use the input attention mask to shift the logits and labels, because it is 2D.
+                # we also crop attn mask in case it is longer, which happens in PrefixTuning with peft
+                shift_attention_mask = attention_mask[:, -shift_logits.shape[1] :].to(logits.device)
+                shift_logits = shift_logits[shift_attention_mask.to(logits.device) != 0].contiguous()
+                shift_labels = shift_labels[shift_attention_mask.to(shift_labels.device) != 0].contiguous()
+            else:
+                shift_logits = shift_logits.contiguous()
+                shift_labels = shift_labels.contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+
             flat_logits = shift_logits.view(-1, self.config.text_config.vocab_size)
             flat_labels = shift_labels.view(-1).to(shift_logits.device)
             loss = loss_fct(flat_logits, flat_labels)

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel/transformers/model/phi3.py

@@ -91,10 +91,11 @@ def lce_forward(
 
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel/transformers/model/qwen2.py

@@ -228,10 +228,11 @@ def lce_forward(
 
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel/transformers/model/qwen2_5_vl.py

@@ -133,8 +133,13 @@ def lce_forward(
         logits = self.lm_head(hidden_states)
 
         loss = None
-        if labels is not None:
-            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size)
+        if labels is not None or shift_labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                shift_labels=shift_labels,
+                vocab_size=self.config.vocab_size,
+            )
 
     if not return_dict:
         output = (logits,) + outputs[1:]

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel/transformers/model/qwen2_vl.py

@@ -129,8 +129,13 @@ def lce_forward(
         logits = self.lm_head(hidden_states)
 
         loss = None
-        if labels is not None:
-            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size)
+        if labels is not None or shift_labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                shift_labels=shift_labels,
+                vocab_size=self.config.vocab_size,
+            )
 
     return Qwen2VLCausalLMOutputWithPast(
         loss=loss,

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel/transformers/model/qwen3.py

@@ -103,10 +103,11 @@ def lce_forward(
 
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel/transformers/model/qwen3_moe.py

@@ -107,8 +107,14 @@ def lce_forward(
         )
     else:  # if in inference model materialize logits
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
-            loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
+        if labels is not None or shift_labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                shift_labels=shift_labels,
+                vocab_size=self.vocab_size,
+                **kwargs,
+            )
 
     aux_loss = None
     if output_router_logits:

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel/transformers/model/smollm3.py

@@ -121,10 +121,11 @@ def lce_forward(
 
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel/transformers/monkey_patch.py

@@ -2107,6 +2107,72 @@ def apply_liger_kernel_to_internvl(
             logger.warning(f"{vision_model_name} is not supported by Liger kernel.")
 
 
+def apply_liger_kernel_to_falcon_h1(
+    rope: bool = True,
+    cross_entropy: bool = True,
+    fused_linear_cross_entropy: bool = False,
+    rms_norm: bool = True,
+    swiglu: bool = False,
+    model: PreTrainedModel = None,
+) -> None:
+    """
+    Apply Liger kernels to replace original implementation in HuggingFace Falcon-H1 models
+    Args:
+        rope (bool): Whether to apply Liger's rotary position embedding. Default is True.
+        cross_entropy (bool): Whether to apply Liger's cross entropy loss. Default is True.
+        fused_linear_cross_entropy (bool):
+            Whether to apply Liger's fused linear cross entropy loss. Default is False.
+            `cross_entropy` and `fused_linear_cross_entropy` cannot both be True.
+            If `fused_linear_cross_entropy` is True, the logits will not be materialized but more memory efficient.
+        rms_norm (bool): Whether to apply Liger's RMSNorm. Default is False.
+        swiglu (bool): Whether to apply Liger's SwiGLU MLP. Default is True.
+        model (PreTrainedModel): The model instance to apply Liger kernels to, if the model has already been
+        loaded. Default is None.
+    """
+
+    assert not (cross_entropy and fused_linear_cross_entropy), (
+        "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    )
+
+    from transformers.models.falcon_h1 import modeling_falcon_h1
+    from transformers.models.falcon_h1.modeling_falcon_h1 import FalconH1Model
+
+    if rope:
+        logger.info("Apply liger rotary pos emb.")
+        modeling_falcon_h1.apply_rotary_pos_emb = liger_rotary_pos_emb
+    if rms_norm:
+        logger.info("Apply liger RMSNorm")
+        modeling_falcon_h1.FalconH1RMSNorm = LigerRMSNorm
+    if swiglu:
+        raise NotImplementedError("LigerSwiGLUMLP is not available for Falcon-H1 models.")
+
+    if cross_entropy:
+        logger.info("Apply liger cross entropy")
+        from transformers.loss.loss_utils import nn
+
+        nn.functional.cross_entropy = liger_cross_entropy
+
+    if fused_linear_cross_entropy:
+        raise NotImplementedError("LigerFusedLinearCrossEntropy is not available for Falcon-H1 models.")
+
+    if model is not None:
+        # The model instance already exists, so we need to additionally patch the
+        # instance variables that reference already-instantiated modules (e.g. LlamaRMSNorm or LlamaMLP)
+
+        # get the base model from the model instance
+        base_model: FalconH1Model = getattr(model, model.base_model_prefix, model)
+
+        if rms_norm:
+            _patch_rms_norm_module(base_model.final_layernorm)
+
+        for decoder_layer in base_model.layers:
+            if swiglu:
+                _patch_swiglu_module(decoder_layer.mlp, LigerSwiGLUMLP)
+            if rms_norm:
+                _patch_rms_norm_module(decoder_layer.input_layernorm)
+                _patch_rms_norm_module(decoder_layer.pre_ff_layernorm)
+
+
 # Model type corresponds to the keys defined in transformers/models/auto/modeling_auto.py
 MODEL_TYPE_TO_APPLY_LIGER_FN = {
     "gemma": apply_liger_kernel_to_gemma,
@@ -2137,6 +2203,7 @@ MODEL_TYPE_TO_APPLY_LIGER_FN = {
     "smollm3": apply_liger_kernel_to_smollm3,
     "phi3": apply_liger_kernel_to_phi3,
     "paligemma": apply_liger_kernel_to_paligemma,
+    "falcon_h1": apply_liger_kernel_to_falcon_h1,
 }
 
 

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/src/liger_kernel_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.6.2.dev20251005233549
+Version: 0.6.2.dev20251010092540
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251010092540}/test/utils.py

@@ -14,7 +14,9 @@ from typing import Tuple
 import numpy as np
 import torch
 import torch.nn as nn
+import transformers
 
+from packaging import version
 from tokenizers import AddedToken
 from tokenizers import Tokenizer
 from tokenizers.models import BPE
@@ -149,18 +151,20 @@ class MiniModelConfig:
 
 def simple_collate_fn(data: List[Dict[str, Any]]):
     """A basic collate function to use for DataLoader"""
+    batch = {}
 
     input_ids = torch.stack([torch.tensor(item["input_ids"]) for item in data])
     attention_mask = torch.stack([torch.tensor(item["attention_mask"]) for item in data])
     labels = input_ids.clone()
+    batch["input_ids"] = input_ids
+    batch["attention_mask"] = attention_mask
+    batch["labels"] = labels
+    if version.parse("4.54.1") <= version.parse(transformers.__version__):
+        shift_labels = nn.functional.pad(labels, (0, 1), value=-100)
+        shift_labels = shift_labels[..., 1:].contiguous()
+        batch["shift_labels"] = shift_labels
 
-    return BatchEncoding(
-        {
-            "input_ids": input_ids,
-            "attention_mask": attention_mask,
-            "labels": labels,
-        }
-    )
+    return BatchEncoding(batch)
 
 
 def multimodal_collate_fn(data: List[Dict[str, Any]]):
@@ -174,6 +178,10 @@ def multimodal_collate_fn(data: List[Dict[str, Any]]):
 
     labels = input_ids.clone()
     batch["labels"] = labels
+    if version.parse("4.54.1") <= version.parse(transformers.__version__):
+        shift_labels = nn.functional.pad(labels, (0, 1), value=-100)
+        shift_labels = shift_labels[..., 1:].contiguous()
+        batch["shift_labels"] = shift_labels
 
     # Collate all other keys, e.g. pixel_values, attention_mask, image_grid_thw, etc
     for key in keys: