You can deploy this model and start earning money today!

Rhine-AI/
diffusiongemma-26B-A4B-tmax-sft-ckpt4000

google/diffusiongemma-26B-A4B-it

adapter

Views

No views yet

DiffusionGemma tmax terminal-agent block-SFT — step 4000

LoRA adapter for google/diffusiongemma-26B-A4B-it, supervised-finetuned on Allen AI's tmax terminal-agent trajectories (config skill_tax_20260505_2.2k_combined_balanced_thinking_only_success).

Training objective — response-anchored block diffusion

Each trajectory is trimmed to <4096 tokens, always ending on an assistant turn. Per step we sample one assistant response, tile it into 256-token canvases, and sample one canvas:

Denoiser loss on the sampled canvas (the diffusion target).
Encoder AR loss over the entire prefix preceding that canvas (system prompt, user task, every prior assistant turn and tool output).
Everything after the canvas is dropped.

hyperparameter	value
LoRA rank / alpha	64 / 128
target modules	q,k,v,o,gate,up,down
canvas length	256
max completion / prompt	4096
learning rate	1.5e-5, cosine, 100 warmup
batch (per-device x accum)	1 x 8
step	4000
last logged train loss	0.6706

Usage

python

from peft import PeftModel
from transformers import AutoTokenizer, DiffusionGemmaForBlockDiffusion

base = DiffusionGemmaForBlockDiffusion.from_pretrained("google/diffusiongemma-26B-A4B-it", torch_dtype="bfloat16", device_map="cuda")
model = PeftModel.from_pretrained(base, "Rhine-AI/diffusiongemma-26B-A4B-tmax-sft-ckpt4000")
tok = AutoTokenizer.from_pretrained("Rhine-AI/diffusiongemma-26B-A4B-tmax-sft-ckpt4000")

Prompts are rendered with the tmax/Qwen chat template (not DiffGemma's, which drops the assistant reasoning) and tokenized with the DiffGemma tokenizer. <|im_end|> is ordinary text to this tokenizer, not an EOS token, so generation does not stop at a turn boundary — use a stop-string.