DiffusionGemma

DiffusionGemma là một mô hình ngôn ngữ encoder-decoder dạng block-diffusion. Một encoder causal đọc prompt sạch (và các block đã được tạo trước đó) vào KV cache, và một decoder hai chiều khử nhiễu một “canvas” có kích thước cố định gồm canvas_length token bằng cách cross-attend đến cache đó. Quá trình tạo ra xen kẽ một vòng lặp autoregressive bên ngoài trên các canvas với một vòng lặp khử nhiễu bên trong, trong đó mỗi bước lấy mẫu các token ứng viên, cam kết những token tự tin nhất thông qua BlockRefinementScheduler ở chế độ corruption đồng đều, và thêm nhiễu lại cho phần còn lại. Bản thân mô hình nằm trong transformers dưới dạng DiffusionGemmaForBlockDiffusion; checkpoint được phát hành là google/diffusiongemma-26B-A4B-it.

Cách sử dụng

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 là số bước khử nhiễu trên mỗi canvas (48 khớp với checkpoint đã phát hành); ít bước hơn sẽ nhanh hơn nhưng chất lượng thấp hơn. cache_implementation="static" cho phép bộ giải mã được torch.compile với cudagraphs (xem Static cache and compilation); bỏ cả hai để chạy với dynamic-cache đơn giản hơn.

Với đầu vào đa lượt hoặc đa phương thức, hãy truyền hội thoại messages thô thay vì prompt. Đây là một danh sách các dict {"role", "content"} theo định dạng chat thông thường, được bộ xử lý chạy qua mẫu chat của nó:

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)

Với một lượt của người dùng, bạn có thể bỏ qua messages và truyền image cùng với prompt; bộ xử lý sẽ tự động chuyển đổi nó thành đầu vào hình ảnh của mô hình.

Trình lập lịch

Trình lập lịch là trình lấy mẫu khử nhiễu cho mỗi canvas và có thể thay thế được: thay đổi nó để đổi chiến lược lấy mẫu mà không cần đụng chạm đến bất kỳ phần nào khác. Có ba trình lập lịch khả dụng:

• BlockRefinementScheduler (mặc định): cam kết các token tự tin nhất ở mỗi bước (trên threshold, cộng thêm một hạn mức đều cho mỗi bước) và thêm nhiễu lại phần còn lại. editing_threshold cho phép nó chỉnh sửa lại các token đã cam kết.
• DiscreteDDIMScheduler: lấy mẫu từng vị trí từ hậu nghiệm rời rạc chính xác của quá trình làm nhiễu đều (D3PM). Nó không có tham số, và bước cuối cùng cam kết một cách tất định các token được dự đoán.
• EntropyBoundScheduler: cam kết các vị trí có entropy thấp nhất mà entropy liên hợp của chúng vẫn nằm dưới entropy_bound, nhờ đó các token gần như độc lập được chấp nhận cùng nhau. Nó làm nguội nhiệt độ lấy mẫu từ t_max (0.8) ở bước đầu tiên xuống t_min (0.4) ở bước cuối, khớp với trình lấy mẫu của checkpoint đã phát hành.

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])

Các tùy chọn lấy mẫu cụ thể cho trình lập lịch (ngưỡng threshold/top_k của khối tinh chỉnh, ràng buộc entropy, …) được đặt trong cấu hình trình lập lịch:

from diffusers import BlockRefinementScheduler

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

EntropyBoundScheduler giảm dần nhiệt độ lấy mẫu (t_max/t_min) bên trong qua các bước khử nhiễu; DiscreteDDIMScheduler và BlockRefinementScheduler sử dụng temperature phẳng được truyền vào pipeline (0.0 cho greedy).

Mẫu hóa dự đoán-chỉnh sửa

DiscreteDDIMScheduler hỗ trợ bộ dự đoán-chỉnh sửa leave-one-out của Uniform Diffusion Models Revisited: Leave-One-Out Denoiser and Absorbing State Reformulation. Nó tinh chỉnh canvas bằng corrector_steps lượt quét Gibbs, trong đó lấy lại mẫu các vị trí kém tự tin nhất từ điều kiện một tọa độ của phân phối biên nhiễu, điều này giữ nguyên phân phối biên đó và cải thiện chất lượng sinh mà không phát sinh chi phí huấn luyện thêm. Nó hoạt động trực tiếp trên checkpoint đã phát hành: đối với uniform diffusion, bộ khử nhiễu và hậu nghiệm leave-one-out có thể thay thế cho nhau theo dạng đóng, do đó bộ chỉnh sửa khôi phục được các đại lượng leave-one-out cần thiết mà không cần huấn luyện lại.

Các bước corrector được gộp vào ngân sách num_inference_steps thay vì cộng thêm: pipeline thực hiện ít bước predictor hơn và dùng các lần forward được giải phóng cho corrector, do đó tổng số lần forward của model vẫn là num_inference_steps và chi phí predictor-corrector bằng với chi phí của phương pháp lấy mẫu tổ tiên thông thường.

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])

Bộ thích ứng PEFT

Denoiser là một mô hình 🤗 Transformers, do đó các adapter được nạp thông qua tích hợp PEFT bản địa của nó thay vì API load_lora_weights của diffusers. Vì tích hợp này không phụ thuộc vào loại adapter, các lệnh gọi giống nhau sẽ nạp LoRA, DoRA hoặc bất kỳ adapter PEFT nào khác (ví dụ: kết quả đầu ra của SFTTrainer của TRL). Hãy quản lý các adapter trực tiếp trên thành phần mô hình:

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")

Adapters vẫn ở trạng thái hoạt động và chưa được hợp nhất: DiffusionGemma liên kết trọng lượng cơ bản của encoder và decoder, do đó việc hợp nhất adapter vào chúng sẽ làm hỏng cả hai nhánh.

Bộ nhớ đệm tĩnh và biên dịch

Pipeline nạp trước encoder một lần cho mỗi khối vào một bộ nhớ đệm có thể tái sử dụng (mặc định là DynamicCache). Truyền cache_implementation="static" sẽ sử dụng StaticCache có hình dạng cố định thay thế, trong đó các hình dạng cho phép bạn torch.compile bộ giải mã với cudagraphs để tăng tốc độ hơn nữa (pipeline đánh dấu từng bước và nhân bản logits để bộ nhớ cudagraph không bị ghi đè); đây là cấu hình được hiển thị trong Sử dụng. Bỏ cả lệnh gọi torch.compile và cache_implementation="static" để chạy với bộ nhớ đệm động đơn giản hơn.

Dừng thích ứng

Một khối thường hội tụ trước khi dùng hết num_inference_steps, nên mặc định pipeline sẽ đóng băng mỗi ví dụ trong batch một khi dự đoán argmax của nó ổn định trong stability_threshold bước và entropy trung bình trên mỗi token của nó giảm xuống dưới confidence_threshold (0.005, giá trị được checkpoint đã phát hành sử dụng). Vòng lặp khử nhiễu kết thúc khi mọi ví dụ đều được đóng băng. Điều này giảm khoảng một nửa số lần chạy decoder ở cùng chất lượng và là đòn bẩy thông lượng lớn nhất. Truyền confidence_threshold=None để luôn chạy đủ num_inference_steps:

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

Callbacks

Callbacks được thực thi sau mỗi bước khử nhiễu. Truyền callback_on_step_end_tensor_inputs để chọn các tensor được bao gồm trong callback_kwargs; canvas (các token khối hiện tại) và logits có sẵn. Trả về {"canvas": ...} từ callback để thay thế 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