[FIX] fbgemm version check (#3704)

* fbgemm version check

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* safer version check

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add check for torchvision-torch compatibility

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* refactor package check logic

* Remove logs and enforce torch

---------

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

unsloth/__init__.py

@@ -18,10 +18,18 @@ import os, re, subprocess, inspect, functools
 import numpy as np
 
 # Fix some issues before importing other packages
-from .import_fixes import fix_message_factory_issue
+from .import_fixes import (
+    fix_message_factory_issue,
+    check_fbgemm_gpu_version,
+    torchvision_compatibility_check,
+)
 
 fix_message_factory_issue()
+check_fbgemm_gpu_version()
+torchvision_compatibility_check()
 del fix_message_factory_issue
+del check_fbgemm_gpu_version
+del torchvision_compatibility_check
 
 # Check if modules that need patching are already imported
 critical_modules = ["trl", "transformers", "peft"]

unsloth/import_fixes.py

@@ -239,3 +239,57 @@ def patch_datasets():
             f"#### Unsloth: Using `datasets = {str(datasets_version)}` will cause recursion errors.\n"
             "Please downgrade datasets to `datasets==4.3.0"
         )
+
+
+def check_fbgemm_gpu_version():
+    if importlib.util.find_spec("fbgemm_gpu") is None:
+        return
+    fbgemm_gpu_version = importlib_version("fbgemm_gpu")
+    # We noticed some SegFault or bad alloc errors on lower versions of fbgemm_gpu.
+    if Version(fbgemm_gpu_version) < Version("1.4.0"):
+        raise ImportError(
+            f"Unsloth: fbgemm_gpu=={fbgemm_gpu_version} detected. It might cause unexpected issues like Segmentation Faults. Please uninstall the current one by doing `pip uninstall fbgemm-gpu` && `pip install fbgemm-gpu` to install fbgemm-gpu 1.4.0 or newer!"
+        )
+    elif UNSLOTH_ENABLE_LOGGING:
+        print(f"Unsloth: fbgemm_gpu=={fbgemm_gpu_version} detected.")
+
+
+def torchvision_compatibility_check():
+    if importlib.util.find_spec("torch") is None:
+        raise ImportError("Unsloth: torch not found. Please install torch first.")
+    if importlib.util.find_spec("torchvision") is None:
+        return
+    torch_version = importlib_version("torch")
+    torchvision_version = importlib_version("torchvision")
+
+    # Torch version -> minimum required torchvision version
+    # See https://pytorch.org/get-started/previous-versions/
+    TORCH_TORCHVISION_COMPAT = [
+        ("2.9.0", "0.24.0"),
+        ("2.8.0", "0.23.0"),
+        ("2.7.0", "0.22.0"),
+        ("2.6.0", "0.21.0"),
+        ("2.5.0", "0.20.0"),
+        ("2.4.0", "0.19.0"),
+    ]
+
+    required_torchvision = None
+    for min_torch, min_torchvision in TORCH_TORCHVISION_COMPAT:
+        if Version(torch_version) >= Version(min_torch):
+            required_torchvision = min_torchvision
+            break
+
+    if required_torchvision is None:
+        # Torch version not in compatibility table, skip check
+        return
+
+    if Version(torchvision_version) < Version(required_torchvision):
+        raise ImportError(
+            f"Unsloth: torch=={torch_version} requires torchvision>={required_torchvision}, "
+            f"but found torchvision=={torchvision_version}. "
+            f"Please refer to https://pytorch.org/get-started/previous-versions/ for more information."
+        )
+    elif UNSLOTH_ENABLE_LOGGING:
+        print(
+            f"Unsloth: torch=={torch_version} and torchvision=={torchvision_version} are compatible."
+        )