Fix Gemma3 NaN losses on ROCm by disabling torch.compile for RDNA GPUs by GoldenGrapeGentleman · Pull Request #4029 · unslothai/unsloth

GoldenGrapeGentleman · 2026-02-11T10:48:08Z

Summary

Gemma3 training produces NaN losses from the first step on all RDNA GPUs (gfx1100, gfx1151). The compiled forward path is numerically unstable on the ROCm/Triton backend, while the eager path trains correctly.

This adds a targeted compile disable for Gemma3 on HIP, following the same UNSLOTH_COMPILE_DISABLE pattern already used for Sesame/CSM models.

Reproduction

from unsloth import FastLanguageModel
from trl import SFTTrainer, SFTConfig
from datasets import Dataset

ds = Dataset.from_list([{"text": "Q: Hi\nA: Hello!"}] * 20)
m, t = FastLanguageModel.from_pretrained(
    "unsloth/gemma-3-1b-pt", max_seq_length=256, dtype=None, load_in_4bit=True,
)
m = FastLanguageModel.get_peft_model(m, r=16,
    target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"],
    use_gradient_checkpointing="unsloth",
)
trainer = SFTTrainer(model=m, tokenizer=t, train_dataset=ds,
    args=SFTConfig(output_dir="./out", max_steps=3, bf16=True,
        per_device_train_batch_size=2, dataset_text_field="text",
        max_seq_length=256, report_to="none"))
trainer.train()  # loss=nan on every step

Root Cause

The Triton/ROCm compiler backend generates numerically unstable code for Gemma3. Other architectures (Llama, Qwen, Mistral) compile and train correctly on the same GPU. The eager path for Gemma3 is numerically correct — disabling compile is sufficient.

The module-level constant unsloth_zoo.compiler.UNSLOTH_COMPILE_DISABLE must also be overridden at runtime because it is evaluated at import time, before the model-specific configuration in loader.py runs.

Testing

Tested on AMD Radeon PRO W7900 (gfx1100, ROCm 7.1, PyTorch 2.8.0):

Model	Before	After
Gemma-3-1B SFT (4bit bf16)	❌ NaN (all steps)	✅ loss=3.81
Llama-3.2-1B SFT (regression)	✅ loss=4.34	✅ loss=4.40
Qwen2.5-1.5B SFT (regression)	✅ loss=2.88	✅ (unaffected)

Also verified on a fresh container with no pre-existing compile cache.

Closes #3385

Related: PR #3588 (broader RDNA3 fix attempt, open since Nov 2025)

Co-authored-by: billishyahao [email protected]
Co-authored-by: yueyuan [email protected]

cc @danielhanchen @0xrushi @kyuz0

gemini-code-assist · 2026-02-11T10:48:24Z

Summary of Changes

Hello @GoldenGrapeGentleman, 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 resolves a critical numerical instability issue affecting Gemma3 model training on AMD RDNA GPUs with ROCm, where torch.compile led to NaN losses. By selectively disabling compilation for Gemma3 on HIP devices, the PR ensures stable and correct training, leveraging the numerically sound eager execution path without impacting other models or architectures.

Highlights

Fixes NaN losses for Gemma3 on ROCm/RDNA GPUs: Addresses a numerical instability issue where Gemma3 training produced NaN losses from the first step on RDNA GPUs (gfx1100, gfx1151) when using the compiled forward path.
Disables torch.compile for Gemma3 on HIP: Implements a targeted disable of torch.compile specifically for Gemma3 models when running on HIP (ROCm) devices, as the eager path is numerically stable.
Overrides UNSLOTH_COMPILE_DISABLE at runtime: Ensures the UNSLOTH_COMPILE_DISABLE flag is correctly set at runtime, as it's evaluated at import time before model-specific configurations.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

unsloth/models/loader.py
- Added conditional logic within the from_pretrained function to check for HIP devices when loading Gemma3 models.
- Set the UNSLOTH_COMPILE_DISABLE environment variable to "1".
- Imported unsloth_zoo.compiler and explicitly set unsloth_zoo.compiler.UNSLOTH_COMPILE_DISABLE to True to ensure the flag is active at runtime.

Activity

Testing was conducted on an AMD Radeon PRO W7900 (gfx1100, ROCm 7.1, PyTorch 2.8.0), confirming the fix for Gemma-3-1B SFT (loss=3.81) and verifying no regressions for Llama-3.2-1B and Qwen2.5-1.5B.
The fix was also verified on a fresh container to ensure no pre-existing compile cache interfered.
This pull request closes issue Training on ROCm (gfx1151, Strix Halo) results in NaN losses with Gemma3 fine-tuning #3385, which reported the NaN loss problem.
It is related to PR fix: unsloth fixes for gfx1151 #3588, which is a broader attempt at RDNA3 fixes.
The author has cc'd @danielhanchen, @0xrushi, and @kyuz0 for review.

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 introduces a targeted fix to disable torch.compile for Gemma3 models on ROCm platforms, addressing an issue with NaN losses during training on RDNA GPUs. The change correctly identifies the specific hardware (hip) and model type (gemma3) and applies a workaround by setting UNSLOTH_COMPILE_DISABLE and patching the unsloth_zoo.compiler module at runtime. The implementation is sound and directly addresses the root cause described. I have one minor suggestion to improve code style by moving the local import to the top of the file.

gemini-code-assist · 2026-02-11T10:50:05Z

unsloth/models/loader.py

+            # See https://github.com/unslothai/unsloth/issues/3385
+            if DEVICE_TYPE == "hip":
+                os.environ["UNSLOTH_COMPILE_DISABLE"] = "1"
+                import unsloth_zoo.compiler


To improve code style and adhere to PEP 8, it's better to have imports at the top of the file. Please move this import to the top of unsloth/models/loader.py and remove it from here.

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 18a6c6ad78

ℹ️ About Codex in GitHub

Codex has been enabled to automatically 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 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

chatgpt-codex-connector · 2026-02-11T10:51:23Z

unsloth/models/loader.py

+                os.environ["UNSLOTH_COMPILE_DISABLE"] = "1"
+                import unsloth_zoo.compiler
+
+                unsloth_zoo.compiler.UNSLOTH_COMPILE_DISABLE = True


Avoid persisting compile-disable globally after Gemma3 load

This assignment mutates a process-wide compiler flag and there is no corresponding reset path in FastLanguageModel.from_pretrained, so once a HIP Gemma3 load runs, later loads of other model families in the same Python process can inherit compile-disabled behavior unintentionally. That creates a silent cross-model regression (not just Gemma3) where subsequent calls into unsloth_compile_transformers may stay on eager paths and lose expected optimizations.

Useful? React with 👍 / 👎.

On RDNA GPUs (gfx1100, gfx1151, etc.), Gemma3's compiled forward path produces NaN losses from the first training step. The eager (uncompiled) path is numerically correct. Root cause: the Triton/ROCm compiler backend generates numerically unstable code for Gemma3's architecture. Other model architectures (Llama, Qwen) compile and train correctly on the same hardware. Fix: set UNSLOTH_COMPILE_DISABLE=1 for Gemma3 on HIP, following the same pattern used for Sesame/CSM models. The module-level constant in unsloth_zoo.compiler is also updated since it is evaluated at import time before loader.py's model-specific block runs. Tested on AMD Radeon PRO W7900 (gfx1100, ROCm 7.1, PyTorch 2.8.0): - Gemma-3-1B SFT 4bit bf16: NaN → loss=3.81 (fixed) - Llama-3.2-1B SFT (regression): loss=4.40 (unaffected) Closes unslothai#3385

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

GoldenGrapeGentleman mentioned this pull request Feb 11, 2026

Training on ROCm (gfx1151, Strix Halo) results in NaN losses with Gemma3 fine-tuning #3385

Open

gemini-code-assist bot reviewed Feb 11, 2026

View reviewed changes

chatgpt-codex-connector bot reviewed Feb 11, 2026

View reviewed changes

danielhanchen mentioned this pull request Feb 11, 2026

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

Open

GoldenGrapeGentleman and others added 2 commits February 12, 2026 02:57

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

18a6c6a

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

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Fix Gemma3 NaN losses on ROCm by disabling torch.compile for RDNA GPUs#4029

Fix Gemma3 NaN losses on ROCm by disabling torch.compile for RDNA GPUs#4029
GoldenGrapeGentleman wants to merge 2 commits intounslothai:mainfrom
GoldenGrapeGentleman:fix-gemma3-nan-on-rocm

GoldenGrapeGentleman commented Feb 11, 2026 •

edited

Loading

Uh oh!

gemini-code-assist bot commented Feb 11, 2026

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Feb 11, 2026

Uh oh!

chatgpt-codex-connector bot left a comment

Uh oh!

chatgpt-codex-connector bot Feb 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

GoldenGrapeGentleman commented Feb 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Reproduction

Root Cause

Testing

Uh oh!

gemini-code-assist bot commented Feb 11, 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 Feb 11, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector bot Feb 11, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

GoldenGrapeGentleman commented Feb 11, 2026 •

edited

Loading