Skip to content

Fix broken wandb import crashing unsloth startup #4147

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
danielhanchen merged 2 commits into from
Mar 3, 2026
Merged

Fix broken wandb import crashing unsloth startup #4147

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
68a26ec
Fix broken wandb import crashing unsloth startup
danielhanchen Mar 3, 2026
54f5bdc
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Mar 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions unsloth/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -151,6 +151,7 @@
patch_vllm_for_notebooks,
patch_torchcodec_audio_decoder,
disable_torchcodec_if_broken,
disable_broken_wandb,
)

fix_xformers_performance_issue()
Expand All @@ -172,6 +173,7 @@
patch_vllm_for_notebooks()
patch_torchcodec_audio_decoder()
disable_torchcodec_if_broken()
disable_broken_wandb()
Copy link
Copy Markdown

@chatgpt-codex-connector chatgpt-codex-connector Bot Mar 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Invoke wandb guard before importing transformers

disable_broken_wandb() is executed after patch_enable_input_require_grads(), but that earlier patch unconditionally does from transformers import PreTrainedModel (unsloth/import_fixes.py), so on environments where importing transformers already touches wandb (the failure mode this commit targets), startup can still crash before the new guard runs. Moving the guard earlier in the init sequence is necessary to make the workaround effective for those installs.

Useful? React with 👍 / 👎.


del fix_xformers_performance_issue
del fix_vllm_aimv2_issue
Expand All @@ -191,6 +193,7 @@
del patch_vllm_for_notebooks
del patch_torchcodec_audio_decoder
del disable_torchcodec_if_broken
del disable_broken_wandb

# Torch 2.4 has including_emulation
if DEVICE_TYPE == "cuda":
Expand Down
32 changes: 32 additions & 0 deletions unsloth/import_fixes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1290,6 +1290,38 @@ def disable_torchcodec_if_broken():
pass


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.

This function tests if wandb can actually import and if not, patches
transformers' is_wandb_available() to return False.
"""
if importlib.util.find_spec("wandb") is None:
return # wandb not installed, nothing to do

try:
import wandb
except Exception:
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Mar 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Using a broad except Exception: can catch unintended errors like KeyboardInterrupt and hide bugs. It's better to catch more specific exceptions that you expect to occur during a failed import of a library, especially one with native extensions. For consistency with other similar functions in this file, like disable_torchcodec_if_broken, consider catching (ImportError, RuntimeError, OSError).

Suggested change
except Exception:
except (ImportError, RuntimeError, OSError):

# wandb is installed but broken - patch transformers to skip it
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"
)
try:
import transformers.integrations.integration_utils as tf_integration

tf_integration.is_wandb_available = lambda: False
except (ImportError, AttributeError):
pass
# Also set env var as fallback for any other code path
os.environ["WANDB_DISABLED"] = "true"


CAUSAL_CONV1D_BROKEN = False
_CAUSAL_CONV1D_PREFIX = "causal_conv1d"
_CAUSAL_CONV1D_BLOCKER_SENTINEL = "_unsloth_causal_conv1d_blocker"
Expand Down