fix: handle gemma 3 weights prefix during hf conversion (#2156)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Bhimraj Yadav <[email protected]>

Author: 3 people

litgpt/scripts/convert_hf_checkpoint.py

@@ -18,7 +18,6 @@
 
 from litgpt.config import Config
 from litgpt.utils import (
-    _TRANSFORMERS_GREATER_EQUAL_4_52,
     extend_checkpoint_dir,
     incremental_save,
     lazy_load,
@@ -292,15 +291,6 @@ def copy_weights_gemma_2(
                 pbar.update(progress_per_file)
 
 
-GEMMA3_LANGUAGE_MODEL_PREFIX = "model.language_model" if _TRANSFORMERS_GREATER_EQUAL_4_52 else "language_model.model"
-
-GEMMA3_VISION_MODEL_PREFIX = "model.vision_tower" if _TRANSFORMERS_GREATER_EQUAL_4_52 else "vision_tower"
-
-GEMMA3_MM_PROJECTOR_PREFIX = (
-    "model.multi_modal_projector" if _TRANSFORMERS_GREATER_EQUAL_4_52 else "multi_modal_projector"
-)
-
-
 def copy_weights_gemma_3(
     qkv_weights: Dict[int, List[Optional[NotYetLoadedTensor]]],
     state_dict: Dict[str, torch.Tensor],
@@ -312,6 +302,22 @@ def copy_weights_gemma_3(
     debug_mode: Optional[bool] = False,
     config: Optional[Config] = None,
 ) -> None:
+    GEMMA3_LANGUAGE_MODEL_PREFIX = (
+        "model.language_model"
+        if any(k.startswith("model.language_model") for k in hf_weights)
+        else "language_model.model"
+    )
+
+    GEMMA3_VISION_MODEL_PREFIX = (
+        "model.vision_tower" if any(k.startswith("model.vision_tower") for k in hf_weights) else "vision_tower"
+    )
+
+    GEMMA3_MM_PROJECTOR_PREFIX = (
+        "model.multi_modal_projector"
+        if any(k.startswith("model.multi_modal_projector") for k in hf_weights)
+        else "multi_modal_projector"
+    )
+
     weight_map = {
         "model.embed_tokens.weight": "transformer.wte.weight",
         "model.layers.{}.self_attn.q_proj.weight": None,

litgpt/utils.py

@@ -47,7 +47,6 @@
 _LITDATA_AVAILABLE = RequirementCache("litdata")
 _LITSERVE_AVAILABLE = RequirementCache("litserve")
 _JINJA2_AVAILABLE = RequirementCache("jinja2")
-_TRANSFORMERS_GREATER_EQUAL_4_52 = RequirementCache("transformers>=4.52.0")
 _SAFETENSORS_AVAILABLE = RequirementCache("safetensors")
 _HF_TRANSFER_AVAILABLE = RequirementCache("hf_transfer")