From 3e43ed7b4afc42fe2afa295dd64c97ea9de64d77 Mon Sep 17 00:00:00 2001 From: Anmol Mishra Date: Fri, 26 Jun 2026 14:15:36 +0530 Subject: [PATCH] Patch FalconH1RMSNorm to fix float64 compilation crash on Intel Arc DG2 (#6691) * 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 --- unsloth/models/falcon_h1.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/unsloth/models/falcon_h1.py b/unsloth/models/falcon_h1.py index b940c95c6..05bfd2ebb 100644 --- a/unsloth/models/falcon_h1.py +++ b/unsloth/models/falcon_h1.py @@ -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