Make bitsandbytes optional on ROCm and add bf16 helper by danielhanchen · Pull Request #4211 · unslothai/unsloth

danielhanchen · 2026-03-12T06:10:04Z

Replacement for #4000 after the original PR was closed during the branch-history incident.

Summary

This makes Unsloth more robust on environments where bitsandbytes is not available (notably ROCm), while keeping CUDA/NVIDIA behavior unchanged when bitsandbytes is installed.

Changes:

Guard bitsandbytes imports in several modules so works without bitsandbytes.
Avoid crashes by building type-tuples only from available classes.
Add a stable helper on HIP/XPU for backwards notebook compatibility.
Guard the vLLM aimv2 patch when vLLM package metadata is missing (module present but no dist-info).
GRPO: align mask/coef lengths in the loss path when left-padding creates a length mismatch.
PEFT compatibility: drop kwarg when running with older .

Testing

on the touched modules.
Validated in a ROCm notebook-suite environment (no bitsandbytes installed) where Unsloth notebooks need to import and train successfully.

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

gemini-code-assist · 2026-03-12T06:10:27Z

Summary of Changes

Hello, I'm Gemini Code Assist¹! 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 significantly enhances Unsloth's robustness and compatibility, particularly for environments like ROCm where bitsandbytes might not be present. The core changes involve making bitsandbytes an optional dependency through extensive conditional imports and type checks. Additionally, it improves compatibility with various versions of vLLM and PEFT, introduces a stable BF16 support helper for HIP/XPU, and addresses a specific padding-related issue in the GRPO loss calculation, leading to a more resilient and adaptable framework.

Highlights

Optional Bitsandbytes Dependency: Guarded bitsandbytes imports and related functional calls across multiple modules (kernels/utils.py, models/_utils.py, models/granite.py, save.py) to ensure Unsloth works robustly in environments where bitsandbytes is not available, such as ROCm.
Dynamic Type-Tuple Construction: Avoided crashes by dynamically building type-tuples for Linear4bit classes (_BNB_LINEAR_TYPES, _MERGE_LORA_LINEAR_TYPES) based on the runtime availability of bitsandbytes and peft modules.
BF16 Helper for HIP/XPU: Added a stable is_bf16_supported helper on HIP/XPU for backwards notebook compatibility, ensuring the function exists with a consistent signature.
vLLM Compatibility: Guarded the vLLM aimv2 patch by adding a try-except block around the importlib_version("vllm") call, preventing crashes when vLLM package metadata is missing.
GRPO Loss Path Alignment: Aligned mask and coefficient lengths in the GRPO loss path within masked_batch_mean to correctly handle scenarios where left-padding creates a length mismatch.
PEFT API Compatibility: Implemented logic to conditionally drop the ensure_weight_tying keyword argument when initializing LoraConfig if the installed PEFT version does not support it, ensuring broader compatibility.

Changelog

unsloth/init.py
- Added a stable is_bf16_supported helper for HIP/XPU for backwards compatibility.
unsloth/import_fixes.py
- Added a try-except block to guard the vLLM version check when package metadata is missing.
unsloth/kernels/utils.py
- Wrapped the bitsandbytes import in a try-except block to make it optional.
- Conditionalized the assignment of bitsandbytes functional library calls based on its availability.
unsloth/models/_utils.py
- Wrapped the bitsandbytes import in a try-except block to make it optional.
- Guarded the peft.tuners.lora.bnb import in patch_fast_lora to handle cases where it's unavailable.
unsloth/models/granite.py
- Wrapped bitsandbytes and peft Linear4bit imports in try-except blocks and created a conditional type tuple.
- Updated post_patch to conditionally check module types based on available bitsandbytes and peft Linear4bit classes.
unsloth/models/llama.py
- Added logic to conditionally remove the ensure_weight_tying argument for LoraConfig based on PEFT version compatibility.
unsloth/models/rl_replacements.py
- Aligned mask and coefficient lengths in masked_batch_mean to handle left-padding discrepancies.
unsloth/save.py
- Wrapped bitsandbytes and peft Linear4bit imports in try-except blocks.
- Created a conditional type tuple _MERGE_LORA_LINEAR_TYPES from available LoRA linear classes.
- Updated _merge_lora to conditionally check module types based on available LoRA linear classes.

Activity

This pull request serves as a replacement for a previously closed PR (Make bitsandbytes optional on ROCm and add bf16 helper #4000).
The changes have been tested with pytest on the touched modules.
Validation was performed in a ROCm notebook-suite environment, confirming that Unsloth notebooks successfully import and train without bitsandbytes installed.

Using Gemini Code Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request makes bitsandbytes an optional dependency, enhancing compatibility with environments like ROCm. The changes primarily involve adding try-except blocks around bitsandbytes imports and conditionally handling its components. This is a solid approach. My main feedback across several files is to replace broad except Exception: clauses with the more specific except ImportError: for import-related errors. This improves code robustness by not masking other potential issues. The other fixes for GRPO padding and PEFT compatibility appear correct.

gemini-code-assist · 2026-03-12T06:12:34Z

-    vllm_version = importlib_version("vllm")
+    try:
+        vllm_version = importlib_version("vllm")
+    except Exception:


Catching a broad Exception can hide unexpected errors. It's better to catch a more specific exception. importlib.metadata.version raises importlib.metadata.PackageNotFoundError when package metadata is not found, which is a subclass of ImportError. Using except ImportError: would be more specific and safer here, and it doesn't require adding a new import.

Suggested change

except Exception:

except ImportError:

gemini-code-assist · 2026-03-12T06:12:34Z

+    # https://github.com/bitsandbytes-foundation/bitsandbytes/pull/1330/files
+    HAS_CUDA_STREAM = Version(bnb.__version__) > Version("0.43.3")
+    get_ptr = bnb.functional.get_ptr
+except Exception:


It's better to catch a more specific exception than the general Exception. For a failed import, ImportError is the correct exception to catch. This prevents masking other potential issues during the setup of bnb.

Suggested change

except Exception:

except ImportError:

gemini-code-assist · 2026-03-12T06:12:34Z

-import bitsandbytes as bnb
+try:
+    import bitsandbytes as bnb
+except Exception:


Catching a broad Exception can hide unexpected errors. Since this block is intended to handle a failed import, it's better to catch the more specific ImportError.

Suggested change

except Exception:

except ImportError:

gemini-code-assist · 2026-03-12T06:12:34Z

-    import peft.tuners.lora.bnb
+    try:
+        import peft.tuners.lora.bnb
+    except Exception as e:


Catching a broad Exception can hide unexpected errors. Since this block is intended to handle a failed import of peft.tuners.lora.bnb, it's better to catch the more specific ImportError.

Suggested change

except Exception as e:

except ImportError as e:

gemini-code-assist · 2026-03-12T06:12:34Z

+try:
+    from bitsandbytes.nn import Linear4bit as Bnb_Linear4bit
+except Exception:
+    Bnb_Linear4bit = None
+
+try:
+    from peft.tuners.lora import Linear4bit as Peft_Linear4bit
+except Exception:
+    Peft_Linear4bit = None


It's better to catch a more specific exception than the general Exception. For failed imports, ImportError is the correct exception to catch. This applies to both try-except blocks here.

Suggested change

try:

from bitsandbytes.nn import Linear4bit as Bnb_Linear4bit

except Exception:

Bnb_Linear4bit = None

try:

from peft.tuners.lora import Linear4bit as Peft_Linear4bit

except Exception:

Peft_Linear4bit = None

try:

from bitsandbytes.nn import Linear4bit as Bnb_Linear4bit

except ImportError:

Bnb_Linear4bit = None

try:

from peft.tuners.lora import Linear4bit as Peft_Linear4bit

except ImportError:

Peft_Linear4bit = None

gemini-code-assist · 2026-03-12T06:12:34Z

+        except Exception:
+            arguments.pop("ensure_weight_tying", None)


Catching a broad Exception can hide unexpected errors. It's better to catch a more specific set of exceptions that you might expect from the introspection logic, such as ImportError, AttributeError, or TypeError. This makes the code more robust against unforeseen issues.

Suggested change

except Exception:

arguments.pop("ensure_weight_tying", None)

except (ImportError, AttributeError, TypeError, ValueError):

arguments.pop("ensure_weight_tying", None)

gemini-code-assist · 2026-03-12T06:12:34Z

+try:
+    from bitsandbytes.nn import Linear4bit as Bnb_Linear4bit
+except Exception:
+    Bnb_Linear4bit = None
+
+try:
+    from peft.tuners.lora import Linear4bit as Peft_Linear4bit
+except Exception:
+    Peft_Linear4bit = None


It's better to catch a more specific exception than the general Exception. For failed imports, ImportError is the correct exception to catch. This applies to both try-except blocks here.

Suggested change

try:

from bitsandbytes.nn import Linear4bit as Bnb_Linear4bit

except Exception:

Bnb_Linear4bit = None

try:

from peft.tuners.lora import Linear4bit as Peft_Linear4bit

except Exception:

Peft_Linear4bit = None

try:

from bitsandbytes.nn import Linear4bit as Bnb_Linear4bit

except ImportError:

Bnb_Linear4bit = None

try:

from peft.tuners.lora import Linear4bit as Peft_Linear4bit

except ImportError:

Peft_Linear4bit = None

# Conflicts: # unsloth/__init__.py # unsloth/models/_utils.py # unsloth/save.py

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

danielhanchen · 2026-05-31T11:48:55Z

AMD Support has been added into Unsloth with Bitsandbytes support.

danielhanchen and others added 2 commits February 8, 2026 16:20

Make bitsandbytes optional on ROCm and add bf16 helper

7ecb049

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

e8d8974

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

danielhanchen requested a review from mmathew23 as a code owner March 12, 2026 06:10

danielhanchen requested a review from Datta0 March 12, 2026 06:10

danielhanchen requested review from Datta0 and rolandtannous as code owners March 12, 2026 06:10

danielhanchen requested a review from mmathew23 March 12, 2026 06:10

danielhanchen requested a review from pluesclues as a code owner March 12, 2026 06:10

danielhanchen requested a review from rolandtannous March 12, 2026 06:10

gemini-code-assist Bot reviewed Mar 12, 2026

View reviewed changes

danielhanchen force-pushed the fix/amd-optional-bnb branch from d59efa4 to e8d8974 Compare March 12, 2026 06:15

danielhanchen force-pushed the main branch from 997f1a7 to bced783 Compare March 12, 2026 06:15

This comment was marked as low quality.

Sign in to view

shimmyshimmer and others added 2 commits May 31, 2026 09:50

Merge remote-tracking branch 'origin/main' into pr-4211

3e606ae

# Conflicts: # unsloth/__init__.py # unsloth/models/_utils.py # unsloth/save.py

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

c3d86a8

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

danielhanchen closed this May 31, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Make bitsandbytes optional on ROCm and add bf16 helper#4211

Make bitsandbytes optional on ROCm and add bf16 helper#4211
danielhanchen wants to merge 4 commits into
mainfrom
fix/amd-optional-bnb

danielhanchen commented Mar 12, 2026

Uh oh!

gemini-code-assist Bot commented Mar 12, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Mar 12, 2026

Uh oh!

gemini-code-assist Bot Mar 12, 2026

Uh oh!

gemini-code-assist Bot Mar 12, 2026

Uh oh!

gemini-code-assist Bot Mar 12, 2026

Uh oh!

gemini-code-assist Bot Mar 12, 2026

Uh oh!

gemini-code-assist Bot Mar 12, 2026

Uh oh!

gemini-code-assist Bot Mar 12, 2026

Uh oh!

This comment was marked as low quality.

danielhanchen commented May 31, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

danielhanchen commented Mar 12, 2026

Summary

Testing

Uh oh!

gemini-code-assist Bot commented Mar 12, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

This comment was marked as low quality.

danielhanchen commented May 31, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants