Fix broken wandb import crashing unsloth startup#4147
Conversation
When wandb is installed but broken (e.g., wandb < 0.19.11 with protobuf >= 6.0), the import chain unsloth -> trl -> transformers -> is_wandb_available() -> import wandb crashes with: ImportError: cannot import name 'Imports' from 'wandb.proto.wandb_telemetry_pb2' This happens because transformers' is_wandb_available() has no try/except around `import wandb`. The error propagates up and kills `from unsloth import FastLanguageModel` even though wandb is optional. Add disable_broken_wandb() following the same pattern as disable_torchcodec_if_broken(). It proactively tries importing wandb during early init, and if the import fails, patches is_wandb_available() to return False and sets WANDB_DISABLED=true.
for more information, see https://pre-commit.ci
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a robust mechanism to prevent Highlights
Changelog
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a fix to prevent a startup crash caused by a broken wandb installation. The approach involves proactively checking if wandb can be imported and gracefully disabling it if the import fails. This is a solid defensive programming practice that improves the robustness of the library. The implementation is clean and consistent with existing patterns for handling similar issues with other optional dependencies. I have one suggestion to make the exception handling more specific, which will further improve the code's robustness.
|
|
||
| try: | ||
| import wandb | ||
| except Exception: |
There was a problem hiding this comment.
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).
| except Exception: | |
| except (ImportError, RuntimeError, OSError): |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 54f5bdc40e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| patch_vllm_for_notebooks() | ||
| patch_torchcodec_audio_decoder() | ||
| disable_torchcodec_if_broken() | ||
| disable_broken_wandb() |
There was a problem hiding this comment.
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 👍 / 👎.
) trl/trainer/callbacks.py imports is_wandb_available from accelerate.utils, not from transformers. The original fix in #4147 only patched the transformers version, so `from trl import GRPOTrainer` still crashed via the callbacks.py -> accelerate -> wandb path. Must patch both the source module (accelerate.utils.imports) AND the re-export namespace (accelerate.utils) since Python's `from accelerate.utils import X` reads from the latter, which holds its own cached reference.
* Fix broken wandb import crashing unsloth startup When wandb is installed but broken (e.g., wandb < 0.19.11 with protobuf >= 6.0), the import chain unsloth -> trl -> transformers -> is_wandb_available() -> import wandb crashes with: ImportError: cannot import name 'Imports' from 'wandb.proto.wandb_telemetry_pb2' This happens because transformers' is_wandb_available() has no try/except around `import wandb`. The error propagates up and kills `from unsloth import FastLanguageModel` even though wandb is optional. Add disable_broken_wandb() following the same pattern as disable_torchcodec_if_broken(). It proactively tries importing wandb during early init, and if the import fails, patches is_wandb_available() to return False and sets WANDB_DISABLED=true. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
…slothai#4148) trl/trainer/callbacks.py imports is_wandb_available from accelerate.utils, not from transformers. The original fix in unslothai#4147 only patched the transformers version, so `from trl import GRPOTrainer` still crashed via the callbacks.py -> accelerate -> wandb path. Must patch both the source module (accelerate.utils.imports) AND the re-export namespace (accelerate.utils) since Python's `from accelerate.utils import X` reads from the latter, which holds its own cached reference.
1 participant
Summary
disable_broken_wandb()toimport_fixes.pythat proactively tests if wandb can import, and if not, patchesis_wandb_available()to returnFalseand setsWANDB_DISABLED=true__init__.pyalongside the existingdisable_torchcodec_if_broken()callProblem
When wandb is installed but broken (old wandb < 0.19.11 with protobuf >= 6.0),
from unsloth import FastLanguageModelcrashes with:The import chain is:
unsloth/__init__.py->tokenizer_utils.py->import trl.trainer.sft_trainerfrom transformers import is_wandb_availableis_wandb_available()doesimport wandbwith no try/exceptv6/proto stubs for protobuf 6.xThis commonly happens on Colab where wandb comes pre-installed but may be an older version, while protobuf gets pulled to v6+ by vllm/grpcio dependencies.
wandb is entirely optional -- notebooks use
report_to="none"-- but a broken install kills the entire import.Root cause
wandb/proto/wandb_telemetry_pb2.pydispatches tov3/,v4/,v5/, orv6/proto files based ongoogle.protobuf.__version__[0]. wandb versions before 0.19.11 do not ship av6/directory, so with protobuf >= 6.0 the dispatch falls through all branches and theImportsclass is never defined. This was fixed upstream in wandb/wandb#9612.Test plan
from unsloth import FastLanguageModelworks,is_wandb_available()returnsTrue