From 248fc5aea945997249bb00f08262006fabf4bf7e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 9 Feb 2026 07:39:26 -0800 Subject: [PATCH] Fix dtype mismatch in fp16 + 4-bit/8-bit LoRA training (#4005) * Fix dtype mismatch in fp16 + 4-bit/8-bit LoRA training Two fixes for training with dtype=torch.float16 and load_in_4bit=True: 1. fast_lora.py: fast_dequantize() returns tensors in quant_state.dtype (typically bfloat16 or float32), but activations may be float16. The subsequent matmul/addmm operations require matching dtypes. Add dtype casts after each fast_dequantize() call in LoRA_MLP.backward and LoRA_QKV.backward (5 locations total). 2. rl.py: TRL unconditionally casts trainable parameters to bfloat16 in the peft init block. When training with fp16=True, this causes GradScaler to crash since it requires float32 parameters. Make the cast conditional -- use float32 when fp16 is enabled, bfloat16 otherwise. This is a no-op for GRPOTrainer (whose peft init block is already removed by the existing regex), but fixes SFTTrainer and other TRL trainers. Tested with Llama-3.2-1B-Instruct 4-bit on both fp16 and bf16 training. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix fp16 + 4-bit LoRA: thread correct_dtype through post_patch Root cause: fast_dequantize returns tensors in quant_state.dtype, which for pre-quantized models is bfloat16 (from config.json). The post_patch methods in llama/gemma/gemma2 call patch_model_and_tokenizer without passing correct_dtype, so quant_state.dtype is never overridden to match the user's requested dtype. This causes a dtype mismatch crash in the backward pass when training with dtype=torch.float16. Fix: pass the user's dtype from from_pretrained through post_patch to patch_model_and_tokenizer as correct_dtype, matching the pattern already used by vision.py. Revert the 5 symptom-level dtype casts in fast_lora.py (upW, gateW, QW, KW, VW) since they are no longer needed with quant_state.dtype properly set at the source. Tested: fp16+4bit and bf16+4bit Llama-3.2-1B-Instruct 15-step SFT runs both complete successfully with similar losses (~1.558 vs ~1.563). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove TRL's unconditional bfloat16 cast instead of patching the dtype TRL 0.26.0+ hardcodes `param.data.to(torch.bfloat16)` for all trainable params in quantized models, citing the QLoRA paper recommendation. This is wrong: it ignores the user's requested dtype and breaks GradScaler when fp16=True. The block exists in sft_trainer, grpo_trainer, rloo_trainer, and reward_trainer (not dpo_trainer). Previous fix patched the cast to be dtype-conditional. This commit replaces the entire guard `if getattr(model, "is_loaded_in_4bit", ...) or getattr(model, "is_loaded_in_8bit", ...):` with `if False:` to disable the block entirely. Unsloth already handles adapter dtype via patch_model_and_tokenizer, making TRL's cast both unnecessary and harmful. For GRPOTrainer the enclosing peft init block is already removed by the regex above, making this a no-op for GRPO. --------- Co-authored-by: Daniel Hanchen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/models/gemma.py | 4 ++-- unsloth/models/gemma2.py | 4 ++-- unsloth/models/granite.py | 2 +- unsloth/models/llama.py | 8 +++++--- unsloth/models/rl.py | 12 ++++++++++++ 5 files changed, 22 insertions(+), 8 deletions(-) diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index 7173c03495..55a8c8697f 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -442,10 +442,10 @@ class FastGemmaModel(FastLlamaModel): return @staticmethod - def post_patch(model, tokenizer): + def post_patch(model, tokenizer, correct_dtype = None): # Gemma does not downcast RoPE model, tokenizer = patch_model_and_tokenizer( - model, tokenizer, downcast_rope = False + model, tokenizer, downcast_rope = False, correct_dtype = correct_dtype ) # Add 1 to weight diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py index 16d04955d3..03e77f6504 100644 --- a/unsloth/models/gemma2.py +++ b/unsloth/models/gemma2.py @@ -613,10 +613,10 @@ class FastGemma2Model(FastLlamaModel): return @staticmethod - def post_patch(model, tokenizer): + def post_patch(model, tokenizer, correct_dtype = None): # Gemma does not downcast RoPE model, tokenizer = patch_model_and_tokenizer( - model, tokenizer, downcast_rope = False + model, tokenizer, downcast_rope = False, correct_dtype = correct_dtype ) # Add 1 to weight diff --git a/unsloth/models/granite.py b/unsloth/models/granite.py index aae746aed1..168df90f4c 100644 --- a/unsloth/models/granite.py +++ b/unsloth/models/granite.py @@ -542,7 +542,7 @@ class FastGraniteModel(FastLlamaModel): return @staticmethod - def post_patch(model, tokenizer): + def post_patch(model, tokenizer, correct_dtype = None): # Torch.compile fails on embedding matrix?? # Workaround randomnly fixes it for torch versions < 2.2 model.model.embed_tokens = torch.nn.Embedding.from_pretrained( diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 61771e4567..c1e9110759 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2483,7 +2483,9 @@ class FastLlamaModel: ) model, tokenizer = patch_tokenizer(model, tokenizer) - model, tokenizer = model_patcher.post_patch(model, tokenizer) + model, tokenizer = model_patcher.post_patch( + model, tokenizer, correct_dtype = dtype + ) # Patch up QKV / O and MLP for idx, layer in enumerate(model.model.layers): @@ -2666,9 +2668,9 @@ class FastLlamaModel: return model, tokenizer @staticmethod - def post_patch(model, tokenizer): + def post_patch(model, tokenizer, correct_dtype = None): model, tokenizer = patch_model_and_tokenizer( - model, tokenizer, downcast_rope = True + model, tokenizer, downcast_rope = True, correct_dtype = correct_dtype ) return model, tokenizer diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 6cb12f6a12..67721d7531 100755 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -1309,6 +1309,18 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): flags = re.DOTALL, ) + # Remove TRL's unconditional bfloat16 cast of trainable params (added in + # TRL 0.26.0). TRL hardcodes bfloat16 for QLoRA per the original paper's + # recommendation, but this is wrong: it ignores the user's requested dtype + # and breaks GradScaler when training with fp16=True. Unsloth already + # handles adapter dtype correctly via patch_model_and_tokenizer, so the + # entire block is unnecessary. For GRPOTrainer the enclosing peft init + # block is already removed above, making this a no-op for GRPO. + RLTrainer_source = RLTrainer_source.replace( + 'if getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False):', + "if False:", + ) + if RLTrainer_name == "SFTTrainer": original_text = 'self._signature_columns = ["input_ids", "attention_mask", "completion_mask"]' new_text = 'self._signature_columns = ["input_ids", "attention_mask", "completion_mask","labels"]'