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"]'