Skip to content

ROCm: default GPT-OSS to BF16 and disable AITER #4021

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

Open
danielhanchen wants to merge 12 commits into
base: main
Choose a base branch
from
+141 −6
Open

ROCm: default GPT-OSS to BF16 and disable AITER #4021

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
4e1e24c
ROCm: default GPT-OSS to BF16 and disable AITER
danielhanchen Feb 10, 2026
3daaff6
ROCm: guard Trainer init patch against missing generated function
danielhanchen Feb 11, 2026
cff534d
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 11, 2026
24c7f2e
ROCm GPT-OSS: gate BF16 fallback by prequant capability
danielhanchen Feb 11, 2026
458af41
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 11, 2026
a211e8c
ROCm: trim unintended fast-inference fallback behaviors
danielhanchen Feb 11, 2026
b56ab6a
Refactor HIP GPT-OSS routing into shared loader helper
danielhanchen Feb 11, 2026
4f138ac
Move HIP GPT-OSS routing helper to loader footer
danielhanchen Feb 11, 2026
9956d1d
ROCm notebook stability: deepseek OCR hook + offline GGUF guard
danielhanchen Feb 11, 2026
28aa6c2
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 11, 2026
f0da826
Fix dequant global buffer dtype reuse across mixed precision
danielhanchen Feb 11, 2026
734649e
Remove redundant Deepseek OCR patch call from vision loader
danielhanchen Feb 11, 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
5 changes: 5 additions & 0 deletions unsloth/device_type.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
]

import torch
import os
import functools
import inspect
from unsloth_zoo.utils import Version
Expand Down Expand Up @@ -94,6 +95,10 @@ def get_device_count():
# HSA_STATUS_ERROR_EXCEPTION checks - sometimes AMD fails for BnB
ALLOW_BITSANDBYTES: bool = True
if DEVICE_TYPE == "hip":
# Disable AITER by default on ROCm to avoid JIT build locks and runtime faults.
# Users can override by explicitly setting env vars.
os.environ.setdefault("AITER_DISABLE", "1")
os.environ.setdefault("USE_ROCM_AITER_ROPE_BACKEND", "0")
try:
import bitsandbytes
except:
Expand Down
12 changes: 12 additions & 0 deletions unsloth/models/loader.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -275,6 +275,18 @@ def from_pretrained(
)
load_in_4bit = False

# AMD GPT-OSS: default to BF16 checkpoints to avoid MXFP4/prequant issues
if is_hip() and "gpt-oss" in model_name.lower() and not use_exact_model_name:
if not model_name.lower().endswith("-bf16"):
if "120b" in model_name.lower():
model_name = "unsloth/gpt-oss-120b-BF16"
else:
model_name = "unsloth/gpt-oss-20b-BF16"
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Feb 10, 2026

Choose a reason for hiding this comment

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

medium

The current logic for selecting the BF16 version of a GPT-OSS model assumes that any model not containing '120b' is a '20b' model. This is not robust and could lead to incorrectly loading the 20b model for other sizes (e.g., a hypothetical 60b model). I suggest making this logic more explicit by handling the '20b' case and adding a warning for any other unhandled sizes.

Suggested change
if "120b" in model_name.lower():
model_name = "unsloth/gpt-oss-120b-BF16"
else:
model_name = "unsloth/gpt-oss-20b-BF16"
if "120b" in model_name.lower():
model_name = "unsloth/gpt-oss-120b-BF16"
elif "20b" in model_name.lower():
model_name = "unsloth/gpt-oss-20b-BF16"
else:
logger.warning_once(
f"Unsloth: gpt-oss model '{model_name}' is not explicitly supported for BF16 defaulting on ROCm. "
f"Defaulting to 'unsloth/gpt-oss-20b-BF16', which might be incorrect."
)
model_name = "unsloth/gpt-oss-20b-BF16"

Copy link

@chatgpt-codex-connector chatgpt-codex-connector bot Feb 10, 2026

Choose a reason for hiding this comment

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

P1 Badge Preserve non-base GPT-OSS IDs on HIP

This rewrite forces any HIP model_name that contains gpt-oss into one of two base BF16 checkpoints, so with default use_exact_model_name=False it silently changes explicit variants (for example unsloth/gpt-oss-safeguard-20b/120b in unsloth/models/mapper.py) and custom fine-tuned IDs or paths into unsloth/gpt-oss-20b-BF16 or -120b-BF16. In those cases users load a different model than requested and can lose safeguard behavior without an error, so the remap should be restricted to the base GPT-OSS IDs (or fail clearly for unsupported variants).

Useful? React with 👍 / 👎.

load_in_4bit = False
load_in_8bit = False
load_in_fp8 = False
load_in_16bit = True

# Find FP8, BnB 4bit, other mapped names
old_model_name = model_name
fp8_mode = None
Expand Down