From 20e6882c99802072a32af6d41dbe1fe408fa5a04 Mon Sep 17 00:00:00 2001
From: Datta Nimmaturi <venkatadattasainimmaturi@gmail.com>
Date: Mon, 1 Dec 2025 13:14:07 +0000
Subject: [PATCH] seperate vllm fixes

---
 unsloth/__init__.py     |   6 ++-
 unsloth/import_fixes.py | 107 +++++++++++++++++++---------------------
 2 files changed, 56 insertions(+), 57 deletions(-)

diff --git a/unsloth/__init__.py b/unsloth/__init__.py
index 5ecc2a82cb..3b581fc38c 100644
--- a/unsloth/__init__.py
+++ b/unsloth/__init__.py
@@ -109,7 +109,8 @@ from unsloth_zoo.device_type import (
 # Fix other issues
 from .import_fixes import (
     fix_xformers_performance_issue,
-    patch_vllm_imports,
+    fix_vllm_aimv2_issue,
+    fix_vllm_guided_decoding_params,
     ignore_logger_messages,
     patch_ipykernel_hf_xet,
     patch_trackio,
@@ -117,7 +118,8 @@ from .import_fixes import (
 )
 
 fix_xformers_performance_issue()
-patch_vllm_imports()
+fix_vllm_aimv2_issue()
+fix_vllm_guided_decoding_params()
 ignore_logger_messages()
 patch_ipykernel_hf_xet()
 patch_trackio()
diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py
index 6a866dc5b2..eb5c65847f 100644
--- a/unsloth/import_fixes.py
+++ b/unsloth/import_fixes.py
@@ -114,65 +114,62 @@ def fix_xformers_performance_issue():
                 print(f"Unsloth: Failed patching Xformers with error = {str(e)}")
 
 
-def patch_vllm_imports():
+def fix_vllm_aimv2_issue():
     if importlib.util.find_spec("vllm") is None:
         return
-
-    def fix_vllm_aimv2_issue():
-        # ValueError: 'aimv2' is already used by a Transformers config, pick another name.
-        vllm_version = importlib_version("vllm")
-        if Version(vllm_version) < Version("0.10.1"):
-            vllm_version = importlib.util.find_spec("vllm").origin
-            vllm_version = os.path.split(vllm_version)[0]
-            ovis_config = (
-                Path(vllm_version) / "transformers_utils" / "configs" / "ovis.py"
-            )
-            try:
-                if ovis_config.exists():
-                    with open(ovis_config, "r+", encoding = "utf-8") as f:
-                        text = f.read()
-                        # See https://github.com/vllm-project/vllm-ascend/issues/2046
-                        if 'AutoConfig.register("aimv2", AIMv2Config)' in text:
-                            text = text.replace(
-                                'AutoConfig.register("aimv2", AIMv2Config)',
-                                "",
-                            )
-                            text = text.replace(
-                                """backbone_config.pop('model_type')
-                    backbone_config = AutoConfig.for_model(model_type,
-                                                        **backbone_config)""",
-                                """if model_type != "aimv2":
-                        backbone_config.pop('model_type')
-                        backbone_config = AutoConfig.for_model(model_type, **backbone_config)
-                    else:
-                        backbone_config = AIMv2Config(**backbone_config)""",
-                            )
-                            f.seek(0)
-                            f.write(text)
-                            f.truncate()
-                            if UNSLOTH_ENABLE_LOGGING:
-                                print(
-                                    "Unsloth: Patching vLLM to fix `'aimv2' is already used by a Transformers config, pick another name.`"
-                                )
-            except Exception as e:
-                if UNSLOTH_ENABLE_LOGGING:
-                    print(f"Unsloth: Failed patching vLLM with error = {str(e)}")
-
-    def fix_vllm_guided_decoding_params():
-        # GuidedDecodingParmas is renamed to StructuredOutputsParams in vLLM
-        # https://github.com/vllm-project/vllm/pull/22772/files
-        # trl still wants to use GuidedDecodingParams. This is a temporary patch till trl updates
-        import vllm
-
+    # ValueError: 'aimv2' is already used by a Transformers config, pick another name.
+    vllm_version = importlib_version("vllm")
+    if Version(vllm_version) < Version("0.10.1"):
+        vllm_version = importlib.util.find_spec("vllm").origin
+        vllm_version = os.path.split(vllm_version)[0]
+        ovis_config = (
+            Path(vllm_version) / "transformers_utils" / "configs" / "ovis.py"
+        )
         try:
-            from vllm.sampling_params import GuidedDecodingParams
-        except ImportError:
-            vllm.sampling_params.GuidedDecodingParams = (
-                vllm.sampling_params.StructuredOutputsParams
-            )
+            if ovis_config.exists():
+                with open(ovis_config, "r+", encoding = "utf-8") as f:
+                    text = f.read()
+                    # See https://github.com/vllm-project/vllm-ascend/issues/2046
+                    if 'AutoConfig.register("aimv2", AIMv2Config)' in text:
+                        text = text.replace(
+                            'AutoConfig.register("aimv2", AIMv2Config)',
+                            "",
+                        )
+                        text = text.replace(
+                            """backbone_config.pop('model_type')
+                backbone_config = AutoConfig.for_model(model_type,
+                                                    **backbone_config)""",
+                            """if model_type != "aimv2":
+                    backbone_config.pop('model_type')
+                    backbone_config = AutoConfig.for_model(model_type, **backbone_config)
+                else:
+                    backbone_config = AIMv2Config(**backbone_config)""",
+                        )
+                        f.seek(0)
+                        f.write(text)
+                        f.truncate()
+                        if UNSLOTH_ENABLE_LOGGING:
+                            print(
+                                "Unsloth: Patching vLLM to fix `'aimv2' is already used by a Transformers config, pick another name.`"
+                            )
+        except Exception as e:
+            if UNSLOTH_ENABLE_LOGGING:
+                print(f"Unsloth: Failed patching vLLM with error = {str(e)}")
 
-    fix_vllm_aimv2_issue()
-    fix_vllm_guided_decoding_params()
+def fix_vllm_guided_decoding_params():
+    if importlib.util.find_spec("vllm") is None:
+        return
+    # GuidedDecodingParmas is renamed to StructuredOutputsParams in vLLM
+    # https://github.com/vllm-project/vllm/pull/22772/files
+    # trl still wants to use GuidedDecodingParams. This is a temporary patch till trl updates
+    import vllm
+
+    try:
+        from vllm.sampling_params import GuidedDecodingParams
+    except ImportError:
+        vllm.sampling_params.GuidedDecodingParams = (
+            vllm.sampling_params.StructuredOutputsParams
+        )
 
 
 def ignore_logger_messages():