README.md

---
base_model: google/diffusiongemma-26B-A4B-it
library_name: peft
tags: [diffusion-lm, block-diffusion, lora, terminal-agent, tmax]
---

# DiffusionGemma tmax terminal-agent block-SFT — step 3000

LoRA adapter for `google/diffusiongemma-26B-A4B-it`, supervised-finetuned on Allen AI's
[tmax](https://huggingface.co/datasets/allenai/tmax-sft) 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 | 3000 |
| last logged train loss | 0.7370 |

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

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.