ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add FastDiffusionModel slow path for text-diffusion models (DiffusionGemma) (#6158)

Browse source 
* 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 <danielhanchen@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-06-11 04:02:56 -07:00 • committed by GitHub
commit fbaf527787
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 352 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

301
unsloth/models/diffusion.py Normal file
View file

@ -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

51
unsloth/models/loader.py
View file

@ -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