Views
No views yet

1VLLM_USE_V2_MODEL_RUNNER=1
2vllm serve RedHatAI/diffusiongemma-26B-A4B-it-FP8-dynamic \
3 --trust-remote-code \
4 --max-num-seqs 4 \
5 --hf-overrides '{"diffusion_sampler": "entropy_bound", "diffusion_entropy_bound": 0.1}' \
6 --default-chat-template-kwargs '{"enable_thinking": true}'1"""
2Quantize DiffusionGemma model to FP8 using LLM Compressor v0.11.0
3
4Model: google/diffusiongemma-26B-A4B-it
5- Total parameters: ~25.8B
6- Expert parameters: 22.8B (88.4%)
7- Non-expert parameters: 3.0B (11.6%)
8
9Note: This will require a local update to transformers to support the model definition.
10"""
11
12import torch
13from compressed_tensors.offload import dispatch_model
14from transformers import AutoProcessor
15from transformers.models.diffusion_gemma import DiffusionGemmaForBlockDiffusion
16
17from llmcompressor import oneshot
18from llmcompressor.modeling.diffusion_gemma4 import ( # noqa: F401
19 CalibrationDiffusionGemmaTextExperts,
20)
21from llmcompressor.modifiers.quantization import QuantizationModifier
22
23# Load model
24MODEL_ID = "google/diffusiongemma-26B-A4B-it"
25model = DiffusionGemmaForBlockDiffusion.from_pretrained(
26 MODEL_ID, dtype="auto", trust_remote_code=True
27)
28processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
29
30# CalibrationDiffusionGemmaTextExperts replaces the original
31# DiffusionGemmaTextExperts class during calibration to:
32# 1. Linearize the 3D expert tensors into individual nn.Linear modules
33# 2. Ensure all experts are properly calibrated, even those not activated
34# for certain tokens during calibration
35
36# Configure the quantization scheme
37# FP8 Dynamic for all Linear layers
38recipe = QuantizationModifier(
39 targets="Linear",
40 scheme="FP8_DYNAMIC",
41 ignore=[
42 "lm_head",
43 "re:.*embed.*",
44 "re:.*router",
45 "re:.*vision_tower.*",
46 "re:.*self_conditioning.*",
47 ],
48)
49
50oneshot(
51 model=model,
52 recipe=recipe
53)
54
55
56# Test sample generation
57print("========== SAMPLE GENERATION ==============")
58dispatch_model(model)
59
60# "The reason the sky is blue is because" + chat template
61input_ids = torch.tensor(
62 [[
63 2, 105, 2364, 107, 818, 3282, 506, 7217, 563, 3730, 563,
64 1547, 106, 107, 105, 4368, 107
65 ]]
66).to(model.device)
67
68output = model.generate(
69 input_ids,
70 max_new_tokens=100,
71 max_denoising_steps=48,
72)
73print(processor.tokenizer.decode(output[0]))
74print("==========================================\n\n")
75
76# Save to disk in compressed-tensors format
77SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-FP8-Dynamic"
78model.save_pretrained(SAVE_DIR)
79processor.save_pretrained(SAVE_DIR)| Benchmark | google/diffusiongemma-26B-A4B-it | RedHatAI/diffusiongemma-26B-A4B-it-FP8-dynamic | Recovery (%) |
|---|---|---|---|
| AIME 2025 | 0.437 | 0.423 | 96.8% |
| GPQA Diamond | 0.641 | 0.657 | 102.5% |
| IFEval | 0.879 | 0.862 | 98.1% |
| GSM8K | 0.943 | 0.942 | 99.9% |
| MMLU 0-Shot | 0.539 | 0.505 | 93.7% |
| Thinking | |||
| AIME 2025 | 0.650 | 0.660 | 101.5% |
| GPQA Diamond | 0.698 | 0.689 | 98.7% |
| GSM8K | 0.951 | 0.952 | 100.1% |