From ba19fdaef926da4ade91a43f305e79fdbe195f43 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 30 Jun 2025 07:15:48 -0700 Subject: [PATCH] Gemma 3N bug fixes (#2842) * Update vision.py * Bug fix * Update mapper.py * check SDPA for Mistral 3, Pixtral * Update vision.py * Versioning * Update rl_replacements.py * Update README.md * add model registry * move hf hub utils to unsloth/utils * refactor global model info dicts to dataclasses * fix dataclass init * fix llama registration * remove deprecated key function * start registry reog * add llama vision * quant types -> Enum * remap literal quant types to QuantType Enum * add llama model registration * fix quant tag mapping * add qwen2.5 models to registry * add option to include original model in registry * handle quant types per model size * separate registration of base and instruct llama3.2 * add QwenQVQ to registry * add gemma3 to registry * add phi * add deepseek v3 * add deepseek r1 base * add deepseek r1 zero * add deepseek distill llama * add deepseek distill models * remove redundant code when constructing model names * add mistral small to registry * rename model registration methods * rename deepseek registration methods * refactor naming for mistral and phi * add global register models * refactor model registration tests for new registry apis * add model search method * remove deprecated registration api * add quant type test * add registry readme * make llama registration more specific * clear registry when executing individual model registration file * more registry readme updates * Update _auto_install.py * Llama4 * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Synthetic data * Update mapper.py * Xet and Synthetic * Update synthetic.py * Update loader.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update pyproject.toml * Delete .gitignore * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update _utils.py * Update pyproject.toml * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update chat_templates.py * Seasame force float16 / float32 * Fix Seasame * Update loader.py * Update vision.py * Update vision.py * Update vision.py * Update loader.py * is_multimodal * Update loader.py * Update loader.py * Update loader.py * Update loader.py * Update vision.py * Update vision.py * Update vision.py * UNSLOTH_DISABLE_STATIC_GENERATION * Update vision.py * Auto vision detection * Sesame * Whisper * Update loader.py * Update loader.py * Update loader.py * Update mapper.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update loader.py * Update loader.py * Update loader.py * Update loader.py * Update _utils.py * Update rl.py * versioning * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * logging * Update pyproject.toml * Update rl.py * versioning * Update rl.py * Update rl.py * Update rl_replacements.py * Update rl_replacements.py * Update rl.py * Update rl_replacements.py * Update rl_replacements.py * logits / temperature * Update rl_replacements.py * Update pyproject.toml * Update rl_replacements.py * Update rl_replacements.py * Debugging only * Update llama.py * Update llama.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Generic efficient GRPO * Update rl_replacements.py * Update rl_replacements.py * Remove debugging * Update rl_replacements.py * Update rl_replacements.py * Update vision.py * Update llama.py * Update rl_replacements.py * versioning * Update _utils.py * Update vision.py * Update mapper.py * Update loader.py * Update mapper.py * Update vision.py * Update loader.py * Update vision.py * Update loader.py * Update _utils.py * Update vision.py * gradient checkpointing * Gemma 3N fixes * Update loader.py * Versioning * Gemma 3N fixes * Update vision.py * Update vision.py * Update loader.py * Update vision.py --------- Co-authored-by: Jack Shi Wei Lun <87535974+jackswl@users.noreply.github.com> Co-authored-by: jeromeku Co-authored-by: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> --- pyproject.toml | 4 ++-- unsloth/models/_utils.py | 2 +- unsloth/models/llama.py | 21 +++++++++++++------- unsloth/models/loader.py | 17 +++++++++++++++- unsloth/models/vision.py | 43 ++++++++++++++++++++++++++++++++++------ 5 files changed, 70 insertions(+), 17 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index a19d68f94a..e1e021d5ed 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,7 @@ triton = [ ] huggingface = [ - "unsloth_zoo>=2025.6.6", + "unsloth_zoo>=2025.6.7", "packaging", "tyro", "transformers>=4.51.3,!=4.47.0,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3", @@ -381,7 +381,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3", ] colab-new = [ - "unsloth_zoo>=2025.6.6", + "unsloth_zoo>=2025.6.7", "packaging", "tyro", "transformers>=4.51.3,!=4.47.0,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3", diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index a31bf7b604..154b437525 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "2025.6.8" +__version__ = "2025.6.9" __all__ = [ "SUPPORTS_BFLOAT16", diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index bc46ba177b..d0ff413925 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -909,12 +909,7 @@ def LlamaModel_fast_forward( mask = self. GA_mask if use_static_mask else dynamic_GA_mask pass - try: - is_gradient_checkpointing_layer = isinstance(decoder_layer, GradientCheckpointingLayer) - except: - is_gradient_checkpointing_layer = False - - if gradient_checkpointing and not is_gradient_checkpointing_layer: + if gradient_checkpointing and not isinstance(decoder_layer, GradientCheckpointingLayer): def create_custom_forward(module): def custom_forward(*inputs): return module(*inputs, past_key_value, output_attentions, padding_mask = padding_mask, position_embeddings = position_embeddings) @@ -2019,7 +2014,7 @@ class FastLlamaModel: f" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\n"\\ f"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\n"\\ f"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size:,}\\n"\\ - f' "-____-" Trainable parameters = {get_model_param_count(model, trainable_only=True):,}/{get_model_param_count(model):,} ({get_model_param_count(model, trainable_only=True)/get_model_param_count(model)*100:.2f}% trained)' + f' "-____-" Trainable parameters = {get_model_param_count(model, trainable_only=True):,} of {get_model_param_count(model):,} ({get_model_param_count(model, trainable_only=True)/get_model_param_count(model)*100:.2f}% trained)' logger.warning(debug_info) import gc for _ in range(3): @@ -2842,6 +2837,12 @@ class FastLlamaModel: m = m.model _for_inference(m) + # Since transformers 4.53, must turn off explicitly + for module in model.modules(): + if hasattr(module, "gradient_checkpointing"): + module.gradient_checkpointing = False + pass + # Also disable training for embeddings for NEFTune if hasattr(model, "get_input_embeddings"): embeddings = model.get_input_embeddings() @@ -2880,6 +2881,12 @@ class FastLlamaModel: m = m.model _for_training(m) + # Since transformers 4.53, must turn on explicitly + for module in model.modules(): + if hasattr(module, "gradient_checkpointing"): + module.gradient_checkpointing = use_gradient_checkpointing + pass + # Also re-enable training for embeddings for NEFTune if hasattr(model, "get_input_embeddings"): embeddings = model.get_input_embeddings() diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 70dd896514..a95a54b59d 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -547,25 +547,40 @@ class FastModel(FastBaseModel): lowered_model_name = model_name.lower() LATEST = '\nPlease use transformers via `pip install --no-deps git+https://github.com/huggingface/transformers.git`' NIGHTLY = '\nPlease use nightly transformers via pip install --upgrade "transformers>=4.49.0"`' + # Pixtral if "pixtral" in lowered_model_name and transformers_version < Version("4.49.0"): raise RuntimeError("Unsloth: Pixtral only works on transformers >= 4.49.0." + LATEST) + # Qwen 2.5 elif "qwen2.5" in lowered_model_name and transformers_version < Version("4.49.0"): raise RuntimeError("Unsloth: Qwen 2.5 only works on transformers >= 4.49.0." + LATEST) + # Gemma 3 elif "gemma-3" in lowered_model_name and transformers_version < Version("4.50.0.dev0"): raise RuntimeError("Unsloth: Gemma 3 only works on transformers >= 4.50.0." + NIGHTLY) + # Cohere elif "c4ai-command-a-03-2025" in lowered_model_name and transformers_version < Version("4.50.0.dev0"): raise RuntimeError("Unsloth: Cohere's Command model only works on transformers >= 4.50.0." + NIGHTLY) + # Sesame elif "csm-1b" in lowered_model_name: os.environ["UNSLOTH_DISABLE_STATIC_GENERATION"] = "1" # Sesame fails - os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"] = "torch.float16;if name.endswith(('_proj', 'fc1', 'fc2', 'codebook', 'head')): module.to(torch.float16)" + os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"] = \ + "all;torch.float32;torch.float16;"\ + "if name.endswith(('_proj', 'fc1', 'fc2', 'codebook', 'head')): module.to(torch.float16);" + # Granite 4 elif 'granite-4' in lowered_model_name: # granite-4 rms norms are stored as 16 bit, but we upcast os.environ["UNSLOTH_UPCAST_LAYERNORM"] = "1" os.environ["UNSLOTH_DISABLE_STATIC_GENERATION"] = "1" + # Olmo 2 elif "olmo-2" in lowered_model_name and transformers_version < Version("4.50.0.dev0"): raise RuntimeError("Unsloth: OLMo-2 only works on transformers >= 4.50.0." + NIGHTLY) + # Gemma 3N elif "gemma-3n" in lowered_model_name: os.environ["UNSLOTH_DISABLE_STATIC_GENERATION"] = "1" + os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"] = \ + "float16;torch.float16;torch.float16;"\ + "if name.endswith(('.conv')): module;"\ + "from unsloth_zoo.temporary_patches.gemma3n import patch_Gemma3nConvNormAct_forward; patch_Gemma3nConvNormAct_forward()" + if transformers_version < Version("4.53.0"): raise RuntimeError("Unsloth: Gemma 3N only works on transformers >= 4.53.0" + LATEST) else: diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 3137792dac..fb0a54489a 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -350,11 +350,23 @@ class FastBaseModel: correct_dtype = None if os.environ.get("UNSLOTH_FORCE_CUSTOM_DTYPE", "") != "": custom_datatype = os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"] - assert custom_datatype.count(";") == 1 - bnb_compute_dtype, custom_datatype = custom_datatype.split(";", 1) - dtype = torch.float32 - bnb_compute_dtype = eval(bnb_compute_dtype) - correct_dtype = bnb_compute_dtype + assert custom_datatype.count(";") >= 4 + checker, _dtype, _bnb_compute_dtype, _custom_datatype, execute_code = custom_datatype.split(";", 4) + + # Allow custom dtypes on all runs + allow_all_runs = (checker == "all") + # Allow only on float16 datatypes + allow_float16_runs = (checker == "float16" and dtype == torch.float16) + + if allow_all_runs or allow_float16_runs: + dtype = eval(_dtype) + bnb_compute_dtype = eval(_bnb_compute_dtype) + correct_dtype = bnb_compute_dtype + custom_datatype = _custom_datatype + # Execute code as well + if len(execute_code.strip()) != 0: + exec(execute_code) + pass pass # Stop SDPA for some archs like Pixtral / Mistral3 @@ -423,8 +435,15 @@ class FastBaseModel: # Edit data-types if custom_datatype is not None: - for name, module in model.named_modules(): + for jj, (name, module) in enumerate(model.named_modules()): exec(custom_datatype) + pass + # Clear deleted GPU items + for _ in range(3): + gc.collect() + if DEVICE_TYPE == "cuda": torch.cuda.empty_cache() + elif DEVICE_TYPE == "xpu": torch.xpu.empty_cache() + pass pass # Counteract saved tokenizers @@ -713,6 +732,12 @@ class FastBaseModel: m = m.model _for_inference(m) + # Since transformers 4.53, must turn off explicitly + for module in model.modules(): + if hasattr(module, "gradient_checkpointing"): + module.gradient_checkpointing = False + pass + # Also disable training for embeddings for NEFTune if hasattr(model, "get_input_embeddings"): embeddings = model.get_input_embeddings() @@ -755,6 +780,12 @@ class FastBaseModel: m = m.model _for_training(m) + # Since transformers 4.53, must turn on explicitly + for module in model.modules(): + if hasattr(module, "gradient_checkpointing"): + module.gradient_checkpointing = True + pass + # Also re-enable training for embeddings for NEFTune if hasattr(model, "get_input_embeddings"): embeddings = model.get_input_embeddings()