vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-07-09 15:58:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 29

Patch FalconH1RMSNorm to fix float64 compilation crash on Intel Arc DG2 (#6691)

Browse source 
* Patch FalconH1RMSNorm to fix float64 compilation on Intel Arc DG2

Fixes unslothai/unsloth#6555

Root cause: FalconH1RMSNorm.forward() does hidden_states.pow(2).mean().rsqrt()
with self.variance_epsilon being a Python float64. When torch.compile fuses
this pattern into the auto-generated Triton kernel
'triton_per_fused__to_copy_mean_mul_pow_rsqrt_*', the float64 epsilon causes
type promotion to double. Intel Arc DG2 GPUs do not support double precision
(Double type is not supported on this platform).

The existing patch_rms_layernorm() only patches LlamaRMSNorm, not the
separate FalconH1RMSNorm class in transformers.models.falcon_h1.

Fix: add Unsloth_FalconH1RMSNorm that delegates to fast_rms_layernorm
(@torch.compiler.disable, handles epsilon as tl.float32), and call the
patch in FastFalconH1Model.pre_patch() before model creation.

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

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

* Condense FalconH1RMSNorm patch comments

* [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>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
Anmol Mishra 2026-06-26 14:15:36 +05:30 • committed by GitHub
parent a9c8bcf0e1
commit 3e43ed7b4a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 16 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

16
unsloth/models/falcon_h1.py
View file

@ -37,6 +37,8 @@ try:
FalconH1DecoderLayer,
FalconH1Model,
FalconH1ForCausalLM,
FalconH1RMSNorm,
FalconH1RMSNormGated,
FalconHybridMambaAttentionDynamicCache,
)
except:
@ -677,6 +679,18 @@ def fix_prepare_inputs_for_generation(module):
module.prepare_inputs_for_generation = _fast_prepare_inputs_for_generation
class Unsloth_FalconH1RMSNorm(FalconH1RMSNorm):
"""fast_rms_layernorm (compiler-disabled, fp32 eps) avoids the float64 torch.compile RMSNorm kernel that fails on Intel Arc DG2 (issue #6555)."""
def forward(self, hidden_states):
return fast_rms_layernorm(self, hidden_states, gemma = False)
def patch_falcon_h1_rms_layernorm():
import transformers.models.falcon_h1.modeling_falcon_h1
transformers.models.falcon_h1.modeling_falcon_h1.FalconH1RMSNorm = Unsloth_FalconH1RMSNorm
class FastFalconH1Model(FastLlamaModel):
@staticmethod
def pre_patch():
@ -708,6 +722,8 @@ class FastFalconH1Model(FastLlamaModel):
transformers.models.falcon_h1.modeling_falcon_h1.FalconH1RotaryEmbedding = (
LlamaRotaryEmbedding
)
# Avoids the float64 RMSNorm compile kernel that fails on Intel Arc DG2 (issue #6555).
patch_falcon_h1_rms_layernorm()
return
@staticmethod