sriq-diffusiongemma-26B-A4B-v1.6

Merged bf16 weights — a full model, not an adapter. A LoRA trained on sriq-ai/sriq-sft-v1.6 has been merged into kaivoss/diffusiongemma-26B-A4B-it-ALWAYS-THINK, so it loads and serves like any HF checkpoint with no PEFT in the path.

The model reasons in compressed Simplified Chinese and answers in the language it was asked in.

DiffusionGemma is not autoregressive: a causal encoder reads the prompt into a read-only KV cache, and a bidirectional decoder denoises a 256-token canvas over ~48 refinement passes.

Results

20 held-out prompts (seeded random 5% split of v1.6's 9,988 rows), 384 new tokens, 48 denoising steps.

Base scores zero on every row, so the shift is unambiguous.

Training step count matters a lot here — the same recipe at 200 steps reaches less than two thirds of the coverage:

Verification behaviour

v1.6 supervises traces that catch a real mistake and repair it (19.5% of rows). That shows up in generation — the model abandons a bad line of attack rather than rationalising it:

配方：3(x-1/3)^2+3(y-1/3)^2=3+2/3=11/3。
不对。重新检查x+y=3z。

The v1.5 model did the opposite: it produced steps containing p | p-1 (true for no prime) and stamped the result correct.

Known limitation: repetition on a minority of prompts

2 of 20 held-out rows collapse into a repeated phrase until the token budget runs out, e.g. 验1≤x,y,z≤100。 emitted a dozen times. The base model does this on 0 of 20, so finetuning introduced it.

Measured by zlib compression ratio (looped text compresses far harder than prose):

The mean is slightly higher than base, so output is not globally repetitive — 18 of 20 rows land between 0.35 and 0.61. It is a minority failure mode, not a general collapse, but it is real and you will hit it.

Reasoning validity is also not benchmarked. Only the language shift and this repetition rate were measured; no accuracy evaluation was run.

Usage

import torch
from transformers import AutoModelForImageTextToText, AutoTokenizer

M = "sriq-ai/sriq-diffusiongemma-26B-A4B-v1.6"
tok = AutoTokenizer.from_pretrained(M)
model = AutoModelForImageTextToText.from_pretrained(M, dtype=torch.bfloat16,
                                                    device_map={"": 0}).eval()

messages = [{"role": "user", "content": "求所有正整数 n，使得 n^2 | 2^n+1。"}]
ids = tok.apply_chat_template(messages, add_generation_prompt=True, tokenize=True,
                              return_dict=False)
inp = torch.tensor([ids], device="cuda")

gen = model.generation_config          # ships max_denoising_steps = 48
gen.max_new_tokens = 1536

with torch.no_grad():
    out = model.generate(input_ids=inp, generation_config=gen)
print(tok.decode(out.sequences[0, inp.shape[1]:], skip_special_tokens=True))

Needs ~52 GB VRAM in bf16 on one accelerator (A100 80GB / H100). The ~46 GB of MoE experts are fused 3D parameters, so load_in_4bit does not shrink them.

Deployment

An FP8 build is available at sriq-ai/sriq-diffusiongemma-26B-A4B-v1.6-FP8-dynamic — 27.2 GB instead of 51.6 GB, so it fits on one 40 GB card, with no measured quality loss (8.0% vs 5.9% mean CJK when both are served through the same vLLM build). It needs Hopper/Ada or newer. The commands below are for these bf16 weights.

Because these are merged weights, vLLM serves them directly — no LoRA support required. SGLang does not support this architecture yet (see below).

vLLM

docker pull vllm/vllm-openai:gemma

docker run -itd --name sriq-diffusiongemma \
    --ipc=host --network host --gpus all \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    vllm/vllm-openai:gemma \
        --model sriq-ai/sriq-diffusiongemma-26B-A4B-v1.6 \
        --max-model-len 262144 \
        --max-num-seqs 4 \
        --gpu-memory-utilization 0.85 \
        --host 0.0.0.0 --port 8000

Requires vLLM 0.24.0+. Keep --max-num-seqs low: diffusion state buffers pre-allocate max_seqs × canvas_length × vocab_size tensors, so it costs VRAM in a way ordinary autoregressive serving does not.

SGLang — not supported yet

SGLang cannot serve DiffusionGemma. --dllm-algorithm Gemma4Renoise parses, which makes it look supported, but the flag value is not a registered algorithm and the architecture is not a registered model:

RuntimeError: Unknown diffusion LLM: DiffusionGemmaForBlockDiffusion

Verified on lmsysorg/sglang:latest (v0.5.20). The dLLM registry in srt/dllm/config.py holds only LLaDA2MoeModelLM, SDARForCausalLM and SDARMoeForCausalLM — on the released image and on main — and srt/dllm/algorithm/ contains no renoise sampler. This is not specific to these weights: Google's own checkpoint declares the same architectures.

Support is in PR #34061, still open, and it adds bf16 serving only. The LMSYS cookbook page documents the command ahead of the code. Use vLLM until that lands.

Training

Full writeup in FINETUNE.md; scripts are train.py and merge.py. The unmerged adapter is in adapter/.

Objective follows NVIDIA NeMo Automodel's recipe:

total = diffusion_CE + encoder_loss_weight * encoder_AR_CE

The AR term is not optional — the diffusion term alone teaches denoising and nothing about the text distribution, and produces no style transfer at all.

The adapter is tied, and merged once

Encoder and decoder share tied base weights, which dictates the adapter's shape:

• Separate encoder/decoder adapters — correct unmerged, wrong merged: the shared tensor receives Δ_enc + Δ_dec.
• Decoder-only — merges exactly, but the encoder then has no trainable parameters, so the AR loss has no gradient path and contributes nothing.
• Tied (used here) — encoder and decoder share the same lora_A/lora_B, so both compute W+Δ during training and the AR loss reaches it. merge.py then writes Δ into the tied tensor once.

peft's merge_and_unload() would walk both modules and write Δ twice into the same storage. NVIDIA's own tools/merge_lora.py uses that path and has no tied-weight handling, so it is not safe for this architecture either.