DiffusionGemma

DiffusionGemma est un modèle de langage encodeur-décodeur à diffusion par blocs. Un encodeur causal lit le prompt propre (et les blocs générés précédemment) dans un cache KV, et un décodeur bidirectionnel débruite un « canevas » de taille fixe de canvas_length jetons en effectuant une attention croisée avec ce cache. La génération alterne une boucle externe autorégressive sur les canevas et une boucle interne de débruitage, où chaque étape échantillonne des jetons candidats, valide les plus confiants via BlockRefinementScheduler en mode de corruption uniforme, et rebruite le reste. Le modèle lui-même se trouve dans transformers sous la forme de DiffusionGemmaForBlockDiffusion ; le point de contrôle publié est google/diffusiongemma-26B-A4B-it.

Utilisation

import torch
from transformers import AutoProcessor, DiffusionGemmaForBlockDiffusion

from diffusers import BlockRefinementScheduler, DiffusionGemmaPipeline

model_id = "google/diffusiongemma-26B-A4B-it"
model = DiffusionGemmaForBlockDiffusion.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto")
processor = AutoProcessor.from_pretrained(model_id)
scheduler = BlockRefinementScheduler()

pipe = DiffusionGemmaPipeline(model=model, scheduler=scheduler, processor=processor)
pipe.model.model.decoder = torch.compile(pipe.model.model.decoder, mode="reduce-overhead", fullgraph=True)
output = pipe(
    prompt="Why is the sky blue?",
    gen_length=256,
    num_inference_steps=48,
    cache_implementation="static",
)
print(output.texts[0])

num_inference_steps est le nombre d’étapes de débruitage par canevas (48 correspond au point de contrôle publié) ; moins d’étapes est plus rapide mais de moindre qualité. cache_implementation="static" permet au décodeur d’être compilé avec torch.compile et cudagraphs (voir Cache statique et compilation) ; supprimez les deux pour une exécution plus simple avec un cache dynamique.

Pour les entrées multi-tours ou multimodales, passez une conversation brute messages au lieu de prompt. Il s’agit d’une liste de dictionnaires {"role", "content"} au format de chat habituel, que le processeur fait passer par son modèle de chat :

messages = [
    {"role": "user", "content": "Why is the sky blue?"},
]
# or with an image:
messages = [
    {
        "role": "user",
        "content": [
            {"type": "image", "image": image},
            {"type": "text", "text": "Describe this image."},
        ],
    },
]
output = pipe(messages=messages, gen_length=256)

Pour un tour utilisateur unique, vous pouvez omettre messages et passer une image avec le prompt ; le processeur la convertit automatiquement en entrées d’image du modèle.

Planificateurs

Le planificateur est le module d'échantillonnage qui débruite chaque canevas, et il est interchangeable : remplacez-le pour modifier la stratégie d'échantillonnage sans rien toucher d'autre. Trois planificateurs sont disponibles :

• BlockRefinementScheduler (par défaut) : engage les jetons les plus sûrs à chaque étape (ceux au-dessus de threshold, plus un quota par étape) et réintroduit du bruit sur le reste. editing_threshold permet en outre de rééditer les jetons déjà engagés.
• DiscreteDDIMScheduler : échantillonne chaque position à partir de la postérieure discrète exacte du processus de corruption uniforme (D3PM). Il ne comporte aucun paramètre et l'étape finale engage de manière déterministe les jetons prédits.
• EntropyBoundScheduler : engage les positions à plus faible entropie dont l’entropie conjointe reste sous entropy_bound, de sorte que des jetons approximativement indépendants soient acceptés ensemble. Il fait recuire sa température d’échantillonnage depuis t_max (0.8) à la première étape jusqu’à t_min (0.4) à la dernière, conformément à l’échantillonneur du point de contrôle publié.

from diffusers import DiscreteDDIMScheduler, EntropyBoundScheduler

pipe.scheduler = DiscreteDDIMScheduler()
# or: pipe.scheduler = EntropyBoundScheduler(entropy_bound=0.1)
output = pipe(prompt="Why is the sky blue?", gen_length=256, num_inference_steps=48)
print(output.texts[0])

Les paramètres d’échantillonnage spécifiques au planificateur (le threshold/top_k de raffinement des blocs, la borne d’entropie, …) sont définis dans la configuration du planificateur :

from diffusers import BlockRefinementScheduler

pipe.scheduler = BlockRefinementScheduler.from_config(pipe.scheduler.config, threshold=0.9)

EntropyBoundScheduler effectue un recuit de sa température d’échantillonnage (t_max/t_min) en interne au fil des étapes de débruitage ; DiscreteDDIMScheduler et BlockRefinementScheduler utilisent la temperature plate transmise au pipeline (0.0 pour un échantillonnage glouton).

Échantillonnage prédicteur-correcteur

DiscreteDDIMScheduler prend en charge le prédicteur-correcteur leave-one-out de Uniform Diffusion Models Revisited: Leave-One-Out Denoiser and Absorbing State Reformulation. Il affine le canevas à l’aide de corrector_steps balayages de Gibbs qui rééchantillonnent les positions les moins confiantes à partir de la conditionnelle à une coordonnée de la marginale bruitée, ce qui laisse cette marginale invariante et améliore la génération sans coût d’entraînement supplémentaire. Il fonctionne directement sur le point de contrôle publié : pour la diffusion uniforme, le débruiteur et la postérieure leave-one-out sont interchangeables sous forme fermée, de sorte que le correcteur récupère les quantités leave-one-out dont il a besoin sans aucun réentraînement.

Les balayages du correcteur sont intégrés au budget de num_inference_steps plutôt qu’ajoutés en sus : le pipeline exécute moins d’étapes de prédicteur et consacre les passages avant libérés aux correcteurs, de sorte que le nombre total de passages avant du modèle reste num_inference_steps et que le coût prédicteur-correcteur est identique à celui de l’échantillonnage ancestral simple.

from diffusers import DiscreteDDIMScheduler

pipe.scheduler = DiscreteDDIMScheduler(corrector_steps=2, corrector_k=12)
output = pipe(prompt="Why is the sky blue?", gen_length=256, num_inference_steps=48)
print(output.texts[0])

Adaptateurs PEFT

Le débruiteur est un modèle 🤗 Transformers, les adaptateurs sont donc chargés via son intégration native PEFT plutôt que via l’API load_lora_weights de diffusers. Comme cette intégration est indépendante du type d’adaptateur, les mêmes appels chargent des adaptateurs LoRA, DoRA ou tout autre adaptateur PEFT (par exemple la sortie de SFTTrainer de TRL). Gérez les adaptateurs directement sur le composant modèle :

pipe.model.load_adapter("path/to/adapter", adapter_name="sft")  # LoRA, DoRA, ...
pipe.model.set_adapter("sft")
output = pipe(prompt="Why is the sky blue?", gen_length=256)

pipe.model.disable_adapters()  # run the base model
pipe.model.delete_adapter("sft")

Les adaptateurs restent actifs et non fusionnés : DiffusionGemma lie les poids de base de l’encodeur et du décodeur, de sorte que la fusion d’un adaptateur dans ceux-ci corromprait les deux branches.

Cache statique et compilation

Le pipeline préremplit l'encodeur une fois par bloc dans un cache réutilisable (un DynamicCache par défaut). Passer cache_implementation="static" utilise à la place un StaticCache de forme fixe, dont les formes permettent de torch.compile le décodeur avec cudagraphs pour un gain de vitesse supplémentaire (le pipeline marque chaque étape et clone les logits afin que la mémoire cudagraph ne soit pas écrasée) ; c'est la configuration présentée dans Utilisation. Supprimez à la fois l'appel torch.compile et cache_implementation="static" pour une exécution plus simple avec un cache dynamique.

Arrêt adaptatif

Un bloc converge généralement avant que toutes les num_inference_steps ne soient consommées. Par défaut, le pipeline fige chaque exemple du lot une fois que sa prédiction argmax est stable pendant stability_threshold étapes et que son entropie moyenne par jeton passe sous confidence_threshold (0.005, la valeur utilisée par le point de contrôle publié). La boucle de débruitage se termine une fois que tous les exemples sont figés. Cela divise approximativement par deux le nombre de passages en avant du décodeur à qualité égale et constitue le levier de débit le plus important. Passez confidence_threshold=None pour toujours exécuter les num_inference_steps complètes :

output = pipe(prompt="Why is the sky blue?", gen_length=256, confidence_threshold=None)  # disable adaptive stopping

Callbacks

Les callbacks s'exécutent après chaque étape de débruitage. Passez callback_on_step_end_tensor_inputs pour sélectionner les tenseurs inclus dans callback_kwargs ; canvas (les jetons du bloc actuel) et logits sont disponibles. Retournez {"canvas": ...} depuis le callback pour remplacer le canvas.

def on_step_end(pipe, step, timestep, callback_kwargs):
    canvas = callback_kwargs["canvas"]
    # Inspect or modify `canvas` here.
    return {"canvas": canvas}


out = pipe(
    prompt="Why is the sky blue?",
    callback_on_step_end=on_step_end,
    callback_on_step_end_tensor_inputs=["canvas"],
)

DiffusionGemmaPipeline

>>> import torch
>>> from transformers import AutoProcessor, DiffusionGemmaForBlockDiffusion
>>> from diffusers import BlockRefinementScheduler, DiffusionGemmaPipeline

>>> model_id = "google/diffusiongemma-26B-A4B-it"
>>> model = DiffusionGemmaForBlockDiffusion.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto")
>>> processor = AutoProcessor.from_pretrained(model_id)
>>> scheduler = BlockRefinementScheduler()

>>> pipe = DiffusionGemmaPipeline(model=model, scheduler=scheduler, processor=processor)
>>> output = pipe(prompt="Why is the sky blue?", gen_length=256)
>>> print(output.texts[0])

DiffusionGemmaPipelineOutput