ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Also patch accelerate's is_wandb_available for trl callbacks path

Browse source 
trl/trainer/callbacks.py imports is_wandb_available from accelerate.utils,
not from transformers. Must patch both the source module
(accelerate.utils.imports) AND the re-export namespace (accelerate.utils)
since `from accelerate.utils import X` reads from the latter.
This commit is contained in:
Daniel Han 2026-03-03 15:39:19 +00:00
commit 242a093a67
1 changed files with 33 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

40
unsloth/import_fixes.py
View file

@ -1294,30 +1294,54 @@ def disable_broken_wandb():
"""Disable wandb if it's installed but cannot actually import.
wandb can fail to import when there's a protobuf version mismatch
(e.g., wandb < 0.19.11 with protobuf >= 6.0). This causes a cascading
import failure through trl -> transformers -> wandb that crashes
unsloth's import chain.
(e.g., wandb < 0.19.11 with protobuf >= 6.0). This causes cascading
import failures through trl -> transformers/accelerate -> wandb that
crash unsloth's import chain.
This function tests if wandb can actually import and if not, patches
transformers' is_wandb_available() to return False.
There are two separate is_wandb_available() functions used by trl:
- transformers.integrations.integration_utils.is_wandb_available
(used by most trl trainers)
- accelerate.utils.imports.is_wandb_available
(used by trl/trainer/callbacks.py)
Both must be patched, and broken wandb modules must be purged from
sys.modules to prevent cached partial imports from causing crashes.
"""
import sys as _sys
if importlib.util.find_spec("wandb") is None:
return # wandb not installed, nothing to do
try:
import wandb
except Exception:
# wandb is installed but broken - patch transformers to skip it
# wandb is installed but broken - clean up and disable
logger.info(
"Unsloth: wandb is installed but broken (likely a protobuf version mismatch). "
"Disabling wandb to prevent import errors. To fix, run: pip install --upgrade wandb"
)
_wandb_false = lambda: False
# Patch transformers' is_wandb_available (used by most trl trainers)
try:
import transformers.integrations.integration_utils as tf_integration
tf_integration.is_wandb_available = lambda: False
tf_integration.is_wandb_available = _wandb_false
except (ImportError, AttributeError):
pass
# Also set env var as fallback for any other code path
# Patch accelerate's is_wandb_available (used by trl/trainer/callbacks.py)
# Must patch both the source module AND the re-export namespace since
# `from accelerate.utils import is_wandb_available` reads from
# accelerate.utils, not accelerate.utils.imports
try:
import accelerate.utils.imports as acc_imports
acc_imports.is_wandb_available = _wandb_false
except (ImportError, AttributeError):
pass
try:
import accelerate.utils as acc_utils
acc_utils.is_wandb_available = _wandb_false
except (ImportError, AttributeError):
pass
# Set env var as additional fallback
os.environ["WANDB_DISABLED"] = "true"