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

DiffusionGemma tmax terminal-agent block-SFT — step 2000

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	2000
last logged train loss	0.8997

Usage

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-ckpt2000")
tok = AutoTokenizer.from_pretrained("Rhine-AI/diffusiongemma-26B-A4B-tmax-sft-ckpt2000")

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.