fix(ROCm): remove dead code fix_rocm_triton_key_error by GoldenGrapeGentleman · Pull Request #4125 · unslothai/unsloth

GoldenGrapeGentleman · 2026-02-27T10:33:05Z

Remove fix_rocm_triton_key_error() from #3923 — the function is dead code with multiple issues:

Checks wrong path: function checks triton.runtime.triton_key (missing on ROCm), but torch._inductor uses triton.compiler.compiler.triton_key which exists and works on ROCm.
Invalid env var: sets TORCHINDUCTOR_DISABLE which does not exist in PyTorch (102 TORCHINDUCTOR_* vars exist, this is not one of them).
Set too late: TORCH_COMPILE_DISABLE is set after import torch; torch._dynamo reads it only once at import time, so it has no effect.
Harmful side-effect: the env var is inherited by subprocesses where it does take effect, needlessly disabling torch.compile on working ROCm setups.

torch.compile works correctly on Triton 3.4.0 / gfx1100 — verified with matmul, CE loss, RMSNorm, transformer forward+backward, and LoRA backward.

cc @danielhanchen

gemini-code-assist · 2026-02-27T10:33:20Z

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 addresses a critical issue where torch.compile was not being effectively disabled on ROCm systems lacking triton_key, leading to potential Inductor backend crashes. The primary change involves modifying the fix_rocm_triton_key_error function to directly set torch._dynamo.config.disable = True, ensuring that torch.compile is correctly turned off at runtime, while also maintaining the TORCH_COMPILE_DISABLE environment variable for any child processes.

Highlights

Corrected torch.compile disabling logic: The fix_rocm_triton_key_error function was updated to effectively disable torch.compile on ROCm when triton_key is missing, preventing backend crashes.
Direct torch._dynamo.config.disable setting: Implemented direct setting of torch._dynamo.config.disable = True as the reliable mechanism to disable torch.compile after torch has been imported, as environment variables are read only at import time.
Environment variable handling refined: Retained the TORCH_COMPILE_DISABLE environment variable for child processes and removed the ineffective TORCHINDUCTOR_DISABLE environment variable.

🧠 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/import_fixes.py
- Updated comments to clarify the necessity of direct torch._dynamo.config.disable setting.
- Removed the TORCHINDUCTOR_DISABLE environment variable.
- Added code to directly set torch._dynamo.config.disable = True.
- Modified the logger info message to reflect disabling "Dynamo" instead of "Inductor".

Activity

No human activity has occurred on this pull request yet.

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 correctly identifies and fixes an issue where torch.compile was not being disabled on ROCm when triton_key is missing. The change to directly set torch._dynamo.config.disable = True is the right approach. My feedback includes a suggestion to improve the accuracy of the logging to reflect whether the disable operation was successful.

gemini-code-assist · 2026-02-27T10:33:59Z

+    # Runtime disable — this is the only mechanism that works after torch
+    # has already been imported (env vars are read once at import time).
+    try:
+        torch._dynamo.config.disable = True
+    except Exception:
+        pass
+
    logger.info(
        "Unsloth: ROCm detected and Triton lacks triton_key; "
-        "disabling torch.compile/Inductor to avoid backend crash."
+        "disabling torch.compile/Dynamo to avoid backend crash."
    )


The log message indicating that torch.compile is disabled is currently unconditional. If setting torch._dynamo.config.disable = True fails for some reason (e.g., due to future changes in PyTorch's internal API), the log message would be misleading as it would suggest the operation was successful when it wasn't. It's better to move the logging inside the try block to ensure it's only printed upon successful disabling.

# Runtime disable — this is the only mechanism that works after torch # has already been imported (env vars are read once at import time). try: torch._dynamo.config.disable = True logger.info( "Unsloth: ROCm detected and Triton lacks triton_key; " "disabling torch.compile/Dynamo to avoid backend crash." ) except Exception: pass

The function (introduced in unslothai#3923) assumed that the absence of `triton.runtime.triton_key` on ROCm means torch.compile will crash. Investigation shows this is incorrect: 1. `triton.runtime.triton_key` was renamed/removed in the ROCm Triton fork — it does not exist at that path. However, `triton.compiler.compiler.triton_key` (the path torch._inductor actually imports) EXISTS and works correctly on ROCm. 2. Both call-sites in torch._inductor (codecache.py and async_compile.py) already wrap the import in try/except, so even a genuinely missing triton_key would be handled gracefully. 3. Comprehensive testing on ROCm 7.1 + Triton 3.4.0 + gfx1100 confirms torch.compile works correctly for matmul, cross-entropy, RMSNorm, multi-layer transformer forward+backward, and LoRA — all without triton.runtime.triton_key. The original code was also ineffective (environment variables set after torch import have no effect on torch._dynamo config), so removing it has zero behavioral change on existing installations. Supersedes the compile-disable portion of unslothai#3923.

danielhanchen

Thank you! This works great!

unslothai#4125) The function (introduced in unslothai#3923) assumed that the absence of `triton.runtime.triton_key` on ROCm means torch.compile will crash. Investigation shows this is incorrect: 1. `triton.runtime.triton_key` was renamed/removed in the ROCm Triton fork — it does not exist at that path. However, `triton.compiler.compiler.triton_key` (the path torch._inductor actually imports) EXISTS and works correctly on ROCm. 2. Both call-sites in torch._inductor (codecache.py and async_compile.py) already wrap the import in try/except, so even a genuinely missing triton_key would be handled gracefully. 3. Comprehensive testing on ROCm 7.1 + Triton 3.4.0 + gfx1100 confirms torch.compile works correctly for matmul, cross-entropy, RMSNorm, multi-layer transformer forward+backward, and LoRA — all without triton.runtime.triton_key. The original code was also ineffective (environment variables set after torch import have no effect on torch._dynamo config), so removing it has zero behavioral change on existing installations. Supersedes the compile-disable portion of unslothai#3923.

GoldenGrapeGentleman requested a review from danielhanchen as a code owner February 27, 2026 10:33

gemini-code-assist Bot reviewed Feb 27, 2026

View reviewed changes

GoldenGrapeGentleman force-pushed the fix/rocm-torch-compile-disable branch 2 times, most recently from 352ecba to 8574d28 Compare February 27, 2026 15:33

GoldenGrapeGentleman closed this Feb 28, 2026

GoldenGrapeGentleman reopened this Feb 28, 2026

GoldenGrapeGentleman force-pushed the fix/rocm-torch-compile-disable branch from 8574d28 to 1f6cda0 Compare February 28, 2026 08:10

GoldenGrapeGentleman changed the title ~~fix(ROCm): ensure torch.compile is actually disabled when triton_key is missing~~ fix(ROCm): remove fix_rocm_triton_key_error — based on a false premise Feb 28, 2026

GoldenGrapeGentleman changed the title ~~fix(ROCm): remove fix_rocm_triton_key_error — based on a false premise~~ fix(ROCm): remove dead code fix_rocm_triton_key_error Feb 28, 2026

danielhanchen approved these changes Mar 1, 2026

View reviewed changes

danielhanchen merged commit 876cc86 into unslothai:main Mar 1, 2026
1 check passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

fix(ROCm): remove dead code fix_rocm_triton_key_error#4125

fix(ROCm): remove dead code fix_rocm_triton_key_error#4125
danielhanchen merged 1 commit into
unslothai:mainfrom
GoldenGrapeGentleman:fix/rocm-torch-compile-disable

GoldenGrapeGentleman commented Feb 27, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot commented Feb 27, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Feb 27, 2026

Uh oh!

danielhanchen left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

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

Uh oh!

Uh oh!

gemini-code-assist Bot commented Feb 27, 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 27, 2026

Choose a reason for hiding this comment

Uh oh!

danielhanchen left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

GoldenGrapeGentleman commented Feb 27, 2026 •

edited

Loading