From fbaf5277870d79651c7787d3a25099d5cb0ff39d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 11 Jun 2026 04:02:56 -0700 Subject: [PATCH] Add FastDiffusionModel slow path for text-diffusion models (DiffusionGemma) (#6158) * Add FastDiffusionModel slow path for text-diffusion models (DiffusionGemma) Text-diffusion models (DiffusionGemma) use a block-diffusion generate loop and a novel backbone, so Unsloth's autoregressive kernel/compile patching does not apply. FastDiffusionModel loads the unmodified HuggingFace model (outputs stay bit-identical to transformers) and adds only the safe conveniences: 4bit/8bit loading, PEFT LoRA (attention + dense MLP; the fused 3D MoE experts are noted as a follow-up), the (model, tokenizer) API, and for_inference/for_training. FastModel.from_pretrained auto-routes diffusion model_types to this path, including a fallback that aliases the legacy "diffusion_gemma" config to the "diffusion_gemma4" classes current transformers ships. New file unsloth/models/diffusion.py; loader.py adds the dispatch and the slow-path-aware get_peft_model / for_inference / for_training. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * FastDiffusionModel: use unsloth's standard dtype resolution Match the FastModel loader instead of the one-line bf16/fp16 pick: SUPPORTS_BFLOAT16 = is_bfloat16_supported(), default to bf16 when supported else fp16, downgrade an explicit bf16 on unsupported hardware with a warning, and assert a valid dtype. --------- Co-authored-by: danielhanchen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/models/diffusion.py | 301 ++++++++++++++++++++++++++++++++++++ unsloth/models/loader.py | 51 ++++++ 2 files changed, 352 insertions(+) create mode 100644 unsloth/models/diffusion.py diff --git a/unsloth/models/diffusion.py b/unsloth/models/diffusion.py new file mode 100644 index 0000000000..12596b432e --- /dev/null +++ b/unsloth/models/diffusion.py @@ -0,0 +1,301 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +FastDiffusionModel: a transformers-only slow path for text-diffusion models (e.g. DiffusionGemma). + +These models use a block-diffusion sampling loop (custom generate) and a novel backbone, so we skip +Unsloth's autoregressive kernel/compile patching and load the unmodified HF model (outputs stay +bit-identical to transformers), keeping only the safe conveniences: 4bit/8bit loading, PEFT LoRA, the +(model, tokenizer) API, and for_inference/for_training. Extend DIFFUSION_MODEL_TYPES as more land. +""" + +import os +import torch +from transformers import AutoConfig, AutoProcessor, AutoTokenizer + +from ._utils import is_bfloat16_supported +from .llama import logger + +__all__ = ["FastDiffusionModel", "DIFFUSION_MODEL_TYPES", "is_diffusion_model_type"] + +# transformers model_type strings routed to this slow path +DIFFUSION_MODEL_TYPES = ("diffusion_gemma", "diffusion_gemma4") + +# Default LoRA targets: standard nn.Linear modules in the shared Gemma-4 backbone. The 128 MoE experts +# are fused 3D Parameters (gate_up_proj/down_proj), not nn.Linear, so PEFT LoRA cannot target them. +DIFFUSION_LORA_TARGETS = [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", # attention + "gate_proj", + "up_proj", + "down_proj", # dense (non-expert) MLP +] + +# Vision tower uses a custom Linear with the same suffix names; exclude it so only the text path is wrapped. +DIFFUSION_LORA_EXCLUDE = r".*(vision_tower|embed_vision).*" + + +def is_diffusion_model_type(model_types): + """model_types: str or iterable -> True if any is a known diffusion model_type.""" + if isinstance(model_types, str): + model_types = (model_types,) + return any(mt in DIFFUSION_MODEL_TYPES for mt in model_types) + + +def _resolve_diffusion_model_class(config): + """Resolve the HF model class for a diffusion checkpoint from config.architectures.""" + import transformers + + archs = getattr(config, "architectures", None) or [] + for arch in archs: + cls = getattr(transformers, arch, None) + if cls is not None: + return cls + # Fallbacks across naming revisions. + for name in ( + "DiffusionGemmaForBlockDiffusion", + "DiffusionGemma4ModelForBlockDiffusion", + "DiffusionGemma4ForBlockDiffusion", + ): + cls = getattr(transformers, name, None) + if cls is not None: + return cls + raise RuntimeError( + f"Unsloth: could not resolve a diffusion model class from architectures={archs}. " + "Ensure you have the transformers build that ships the DiffusionGemma implementation." + ) + + +def _load_diffusion_config(model_name, token, trust_remote_code, revision, local_files_only): + """Load the config, aliasing the legacy ``diffusion_gemma`` model_type to the ``diffusion_gemma4`` + classes current transformers ships. AutoConfig raises on the legacy type; catch that, rewrite the + type/arch names in-memory, and rebuild.""" + try: + return AutoConfig.from_pretrained( + model_name, + token = token, + trust_remote_code = trust_remote_code, + revision = revision, + local_files_only = local_files_only, + ) + except ValueError as e: + if "diffusion_gemma" not in str(e): + raise + import json + from transformers.utils import cached_file + + cfg_path = cached_file( + model_name, + "config.json", + token = token, + revision = revision, + local_files_only = local_files_only, + ) + with open(cfg_path, encoding = "utf-8") as f: + cd = json.load(f) + cd["model_type"] = "diffusion_gemma4" + cd.setdefault("architectures", ["DiffusionGemma4ModelForBlockDiffusion"]) + if isinstance(cd.get("text_config"), dict): + cd["text_config"]["model_type"] = "diffusion_gemma4_text" + if isinstance(cd.get("vision_config"), dict): + cd["vision_config"]["model_type"] = "diffusion_gemma4_vision" + from transformers import DiffusionGemma4Config + + return DiffusionGemma4Config.from_dict(cd) + + +class FastDiffusionModel: + """transformers-only slow path for text-diffusion models.""" + + @staticmethod + def from_pretrained( + model_name = "google/diffusiongemma-26B-A4B-it", + max_seq_length = None, # API-compat; diffusion uses canvas_length + dtype = None, + load_in_4bit = False, + load_in_8bit = False, + load_in_16bit = False, + full_finetuning = False, + token = None, + device_map = "auto", + trust_remote_code = False, + attn_implementation = "eager", # exact match with the reference golden logits + revision = None, + return_tokenizer = True, + **kwargs, + ): + SUPPORTS_BFLOAT16 = is_bfloat16_supported() + if dtype is None: + dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16 + elif dtype == torch.bfloat16 and not SUPPORTS_BFLOAT16: + logger.warning_once("Device does not support bfloat16. Will change to float16.") + dtype = torch.float16 + assert dtype in (torch.float16, torch.bfloat16, torch.float32) + + # Honor an explicit local_files_only; else fall back to the offline env vars. + local_files_only = kwargs.pop("local_files_only", None) + if local_files_only is None: + local_files_only = ( + os.environ.get("HF_HUB_OFFLINE", "0") == "1" + or os.environ.get("TRANSFORMERS_OFFLINE", "0") == "1" + ) + config = _load_diffusion_config( + model_name, + token, + trust_remote_code, + revision, + local_files_only, + ) + model_type = getattr(config, "model_type", None) + if not is_diffusion_model_type(model_type): + raise RuntimeError( + f"Unsloth: FastDiffusionModel only supports diffusion model_types {DIFFUSION_MODEL_TYPES}, " + f"got '{model_type}'. Use FastModel/FastLanguageModel for autoregressive models." + ) + + model_cls = _resolve_diffusion_model_class(config) + + load_kwargs = dict( + dtype = dtype, + device_map = device_map, + token = token, + trust_remote_code = trust_remote_code, + attn_implementation = attn_implementation, + revision = revision, + local_files_only = local_files_only, + ) + + # Optional bitsandbytes quant. The MoE experts (3D Parameters) are not nn.Linear so bnb skips + # them; only attention + dense MLP Linears quantize, lm_head/embeddings stay full precision. + if load_in_4bit or load_in_8bit: + from transformers import BitsAndBytesConfig + if load_in_4bit: + qcfg = BitsAndBytesConfig( + load_in_4bit = True, + bnb_4bit_use_double_quant = True, + bnb_4bit_quant_type = "nf4", + bnb_4bit_compute_dtype = dtype, + llm_int8_skip_modules = [ + "lm_head", + "embed_tokens", + "experts", + "self_conditioning", + "router", + ], + ) + else: + qcfg = BitsAndBytesConfig(load_in_8bit = True) + load_kwargs["quantization_config"] = qcfg + + print(f"==(( Unsloth: FastDiffusionModel (slow / transformers-only path) ))==") + print(f" Model: {model_name} | class: {model_cls.__name__} | model_type: {model_type}") + print( + f" dtype: {dtype} | 4bit: {load_in_4bit} | 8bit: {load_in_8bit} | attn: {attn_implementation}" + ) + + model = model_cls.from_pretrained(model_name, **load_kwargs).eval() + # Mark before any early return so get_peft_model/for_* route to the slow path. + model._unsloth_slow_diffusion = True + + if not return_tokenizer: + return model, None + + # Prefer the processor (chat template + tokenizer); fall back to a bare tokenizer. Returned as + # "tokenizer" to match the Unsloth (model, tokenizer) contract. + try: + tokenizer = AutoProcessor.from_pretrained( + model_name, + token = token, + trust_remote_code = trust_remote_code, + revision = revision, + local_files_only = local_files_only, + ) + except Exception: + tokenizer = AutoTokenizer.from_pretrained( + model_name, + token = token, + trust_remote_code = trust_remote_code, + revision = revision, + local_files_only = local_files_only, + ) + + return model, tokenizer + + @staticmethod + def get_peft_model( + model, + r = 16, + target_modules = None, + lora_alpha = 16, + lora_dropout = 0.0, + bias = "none", + use_gradient_checkpointing = True, + random_state = 3407, + task_type = None, + **kwargs, + ): + """Attach a PEFT LoRA to the diffusion backbone (attention + dense MLP). No fused kernels.""" + from peft import LoraConfig, get_peft_model as peft_get_peft_model + + if target_modules is None: + target_modules = DIFFUSION_LORA_TARGETS + + lora_kwargs = dict( + r = r, + lora_alpha = lora_alpha, + lora_dropout = lora_dropout, + bias = bias, + target_modules = target_modules, + task_type = task_type, # None: diffusion has no standard CAUSAL_LM head + **{k: v for k, v in kwargs.items() if k in ("modules_to_save", "init_lora_weights")}, + ) + # Exclude the vision tower's custom (non-Linear) modules that share suffix names. + exclude = kwargs.get("exclude_modules", DIFFUSION_LORA_EXCLUDE) + try: + lora_config = LoraConfig(exclude_modules = exclude, **lora_kwargs) + except TypeError: + # Older PEFT without exclude_modules: scope the target to the text decoder by regex. + lora_kwargs["target_modules"] = ( + r".*model\.decoder\.layers\.\d+\.(self_attn\.[qkvo]_proj|mlp\.(gate|up|down)_proj)" + ) + lora_config = LoraConfig(**lora_kwargs) + if use_gradient_checkpointing: + model.gradient_checkpointing_enable() + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + + model = peft_get_peft_model(model, lora_config) + model._unsloth_slow_diffusion = True + try: + model.print_trainable_parameters() + except Exception: + pass + return model + + @staticmethod + def for_inference(model): + model.eval() + for _, m in model.named_modules(): + if hasattr(m, "gradient_checkpointing"): + m.gradient_checkpointing = False + return model + + @staticmethod + def for_training(model, use_gradient_checkpointing = True): + model.train() + if use_gradient_checkpointing and hasattr(model, "gradient_checkpointing_enable"): + model.gradient_checkpointing_enable() + return model diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 0ee9e49ba1..4dc3046928 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -829,6 +829,7 @@ from ..kernels import ( post_patch_loss_function, ) from .vision import FastBaseModel +from .diffusion import FastDiffusionModel, is_diffusion_model_type from transformers import ( AutoModelForCausalLM, ) @@ -846,6 +847,25 @@ class FastModel(FastBaseModel): model = _prepare_model_for_qat(model, qat_scheme) return model + @staticmethod + def get_peft_model(model, *args, **kwargs): + # Route text-diffusion models (slow path) to the transformers-only PEFT helper. + if getattr(model, "_unsloth_slow_diffusion", False): + return FastDiffusionModel.get_peft_model(model, *args, **kwargs) + return FastBaseModel.get_peft_model(model, *args, **kwargs) + + @staticmethod + def for_inference(model): + if getattr(model, "_unsloth_slow_diffusion", False): + return FastDiffusionModel.for_inference(model) + return FastBaseModel.for_inference(model) + + @staticmethod + def for_training(model, use_gradient_checkpointing = True): + if getattr(model, "_unsloth_slow_diffusion", False): + return FastDiffusionModel.for_training(model, use_gradient_checkpointing) + return FastBaseModel.for_training(model, use_gradient_checkpointing) + @staticmethod def from_pretrained( model_name = "unsloth/Llama-3.2-11B-Vision-Instruct-bnb-4bit", @@ -1065,6 +1085,24 @@ class FastModel(FastBaseModel): local_files_only = True kwargs["local_files_only"] = True + # Text-diffusion slow-path dispatch, factored so both the normal route (below) and the + # legacy-config fallback (in the AutoConfig except handler) share one call site. + def _dispatch_diffusion(): + return FastDiffusionModel.from_pretrained( + model_name = model_name, + max_seq_length = max_seq_length, + dtype = dtype, + load_in_4bit = load_in_4bit, + load_in_8bit = load_in_8bit, + load_in_16bit = load_in_16bit, + full_finetuning = full_finetuning, + token = token, + device_map = device_map, + trust_remote_code = trust_remote_code, + revision = revision, + **kwargs, + ) + try: model_config = AutoConfig.from_pretrained( model_name, @@ -1078,6 +1116,12 @@ class FastModel(FastBaseModel): raise except Exception as error: autoconfig_error = str(error) + # Legacy text-diffusion configs use model_type "diffusion_gemma", which current + # transformers does not register by name (it ships "diffusion_gemma4"). AutoConfig + # raises before we can dispatch; route straight to the diffusion slow path, whose + # loader aliases the legacy type to the gemma4 classes. + if "diffusion_gemma" in autoconfig_error and is_diffusion_model_type("diffusion_gemma"): + return _dispatch_diffusion() if "architecture" in autoconfig_error: if "qwen3_5" in autoconfig_error: raise ImportError( @@ -1126,6 +1170,13 @@ class FastModel(FastBaseModel): ) model_types_all = ",".join(model_types) + "," + # ---- Text-diffusion models (e.g. DiffusionGemma) take a transformers-only slow path. ---- + # These use a custom block-diffusion `generate` and a novel backbone, so we skip Unsloth's + # autoregressive kernel/compile patching and load the unmodified HF model (bit-identical to + # naive transformers), keeping only 4bit/8bit + PEFT LoRA conveniences. + if is_diffusion_model_type(model_types): + return _dispatch_diffusion() + # Save model types and loading method lowered_model_name = model_name.lower() string = os.environ.get("UNSLOTH_MODEL_NAME", "") + model_types_all