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

<div align="center"> <img src="https://huggingface.co/sriq-ai/sriq-diffusiongemma-26B-A4B-v1.5/resolve/main/banner.png" width="100%" alt="SRIQ-DiffusionGemma-26B-A4B-V1.5 — compressed reasoning" /> </div> <p align="center"> <a href="https://sriq.org" target="_blank">SRIQ</a> | <a href="https://bench.sriq.org" target="_blank">Benchmarks</a> | <a href="https://huggingface.co/datasets/sriq-ai/sriq-sft-v1.5" target="_blank">Training data</a> | <a href="https://huggingface.co/kaivoss/diffusiongemma-26B-A4B-it-ALWAYS-THINK" target="_blank">Base model</a> | <a href="https://huggingface.co/unsloth/diffusiongemma-26B-A4B-it-GGUF" target="_blank">Base GGUF</a> | <a href="https://huggingface.co/sriq-ai/sriq-diffusiongemma-26B-A4B-v1.6" target="_blank">v1.6 (supersedes this)</a> </p>

⚠️ Superseded by v1.6

sriq-ai/sriq-diffusiongemma-26B-A4B-v1.6 replaces this release. It ships merged bf16 weights — a full checkpoint with no PEFT in the path, so vLLM serves it directly — and it is trained on sriq-sft-v1.6 (9,988 rows vs 1,424 here) with verified reasoning and answer checks, which addresses the reasoning-validity limitation described below.

This adapter remains up for reproducibility. Its headline CJK number (24.6%) is higher than v1.6's 18.1%, but that is not a regression: v1.6's easy rows carry ~25 reasoning tokens against v1.5's flat 349, so there is less text to be Chinese. Use v1.6.

A LoRA adapter for kaivoss/diffusiongemma-26B-A4B-it-ALWAYS-THINK, trained on sriq-ai/sriq-sft-v1.5 to move the model's chain-of-thought into compressed Simplified Chinese.

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 steps. forward takes no labels and returns no loss, so the training objective is written by hand (see Training below).

Results

Measured on 20 held-out prompts, 384 new tokens each, adapter loaded unmerged:

Every held-out row shifted, and generations got shorter — which is what "compressed" reasoning should look like. The held-out set is a seeded random 10% split; note that sriq-sft-v1.5 is ordered (v1.5 appended its long agentic-harness prompts at the end), so a tail split would hold out an unrepresentative slice.

Known limitation: reasoning validity

The adapter changes how the reasoning looks before it changes whether the reasoning is correct. On IMO 1990 P3 (n² | 2ⁿ+1) it reaches the correct n = 1, 3, but the intermediate steps inside the thought block contain invalid claims (e.g. p | p-1, which holds for no prime) and mangled superscripts. The base model produces cleaner, valid English proofs.

This is expected for a short LoRA run against a style dataset: 800 steps at r=16 is enough to copy surface form, not to teach the underlying logic. Two things are planned to address it:

• Future dataset revisions will supervise reasoning validity directly rather than style alone.
• Scaling the corpus to 10–20k rows (v1.5 is 1,424) to give the objective enough signal to learn the reasoning rather than imitate its shape.

Treat this adapter as a demonstration that the transfer works, not as a production reasoning model.

Usage

import torch
from peft import PeftModel
from transformers import AutoModelForImageTextToText, AutoTokenizer

BASE = "kaivoss/diffusiongemma-26B-A4B-it-ALWAYS-THINK"

tok = AutoTokenizer.from_pretrained(BASE)
model = AutoModelForImageTextToText.from_pretrained(
    BASE, dtype=torch.bfloat16, device_map={"": 0}
)
model = PeftModel.from_pretrained(model, "sriq-ai/sriq-diffusiongemma-26B-A4B-v1.5")
model.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)
input_ids = torch.tensor([ids], device="cuda")

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

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

Load this adapter unmerged. Encoder and decoder share tied base weights, and this adapter carries separate encoder and decoder deltas. merge_and_unload() would write both into the one shared tensor, so each half receives Δ_enc + Δ_dec — not the model that was trained.

This is not the Gemma4ClippableLinear problem quoted in unsloth's notebook: that class is built from Gemma4VisionConfig | Gemma4AudioConfig and lives in the vision/audio towers, which this adapter's target regex excludes. The projections here are plain nn.Linear.

A mergeable variant now exists: v1.6 was retrained with a tied LoRA and ships merged bf16 weights. Prior art for the same shape — see thesysdev/OUI-1, which trains a tied LoRA on decoder projections only, so the single delta propagates to both halves through the weight tying, and ships merged bf16 weights. That needs a retrain with tied targets, not a post-hoc merge of this adapter.

Needs ~52 GB of VRAM in bf16 on a single accelerator (A100 80GB / H100). The ~46 GB of MoE experts are fused 3D parameters, so 4-bit quantisation does not shrink them — load_in_4bit buys you nothing here.

Deployment

vLLM serves the base model through a dedicated block-diffusion path; SGLang does not support this architecture yet (see below). vLLM's docs do not cover LoRA for it either, so serving this adapter through it is untested — the command below is for the base model, and the transformers snippet above is the verified path for the adapter.

The clean fix is a merged bf16 checkpoint, which vLLM serves directly with no LoRA support required. That needs a tied-LoRA retrain (see the note under Usage); thesysdev/OUI-1 ships exactly that shape.

vLLM

Needs vLLM 0.24.0+ and the Gemma image:

docker pull vllm/vllm-openai:gemma

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

--max-num-seqs 4 matters: diffusion state buffers pre-allocate max_seqs × canvas_length × vocab_size tensors, so raising it costs VRAM quadratically in a way normal 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.

llama.cpp

Requires the diffusiongemma build and its llama-diffusion-cli runner; the standard conversion path does not cover this architecture. Prebuilt base GGUFs: unsloth/diffusiongemma-26B-A4B-it-GGUF.

Training

See FINETUNE.md for the full writeup — why this model has no loss function, the traps that cost the most time, and the measurements behind each decision. The training script is train.py.

Objective follows NVIDIA NeMo Automodel's DiffusionGemma recipe:

total = diffusion_CE + encoder_loss_weight * encoder_AR_CE

• diffusion CE — canvas positions corrupted with uniform-random vocab tokens at a per-example rate t ~ U(eps, 1). There is no [MASK] token.
• encoder AR CE — plain causal next-token CE on the encoder's logits over the clean full sequence.

The AR term is not optional. Runs with the diffusion term alone produced no style transfer whatsoever at 30 and 100 steps (mean CJK 5.1% and 3.3%, and only 2 of 14 rows, both of which were Chinese-language questions the base model already answered in Chinese).

Implementation notes that cost real time:

• The encoder must see the clean full sequence for the AR term. Keep the decoder block-causal via decoder_attention_mask, or it reads the very block it is denoising.
• forward already applies tanh(logits/cap)*cap. Applying the softcap again damps gradients.
• transformers 5.11 does not expose encoder_last_hidden_state from the top-level forward, so the AR term requires calling model.encoder / model.decoder directly. 5.17 added the field.