README.md

---
license: apache-2.0
base_model:
- edwixx/diffusiongemma-26B-A4B-it-HERETIC-Uncensored
library_name: peft
pipeline_tag: text-generation
tags:
- diffusion
- diffusion-gemma
- grpo
- rlhf
- reasoning
- glm-5.1
- unsloth
- moe
datasets:
- Jackrong/GLM-5.1-Reasoning-1M-Cleaned
---

# DiffusionGemma-26B-A4B ByzantineSilk (GLM-5.1 GRPO)

<p align="center">
  <img src="https://raw.githubusercontent.com/unslothai/unsloth/main/images/unsloth%20logo%20black%20text.png" width="30%">
</p>

## Model Overview

**DiffusionGemma-26B-A4B ByzantineSilk** is a specialized block-diffusion language model fine-tuned using a novel **Two-Stage Masked-Diffusion SFT + GRPO (Group Relative Policy Optimization)** alignment pipeline.

Built on top of the 26.6B total / 4B active parameter Mixture-of-Experts (MoE) block diffusion architecture (`DiffusionGemmaForBlockDiffusion`), this model was trained on a GLM 5.1 dataset to perform non-autoregressive parallel canvas generation guided by LLM-as-a-Judge perplexity rewards.

## Quickstart
<a target="_blank" href="https://colab.research.google.com/drive/1S0E8Jha5H9UpEvzZx-1hLN-CcBmczeVR?usp=sharing">
  <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>
</a>

### Key Highlights
* **Architecture:** Block Diffusion MoE (26.6B Total / 4B Active Parameters, 256 Canvas Length).
* **Base Model:** `edwixx/diffusiongemma-26B-A4B-it-HERETIC-Uncensored` (Uncensored DiffusionGemma-26B).
* **Primary Dataset:** `Jackrong/GLM-5.1-Reasoning-1M-Cleaned`.
* **Evaluator Reward Model:** `GnLOLot/MiniCPM5-1B-Claude-Opus-Fable5-V2-Thinking` (Causal LM scoring target perplexity given diffusion reasoning traces).
* **Optimization Framework:** Unsloth FastModel with activation checkpointing for standard VRAM footprint.

---

## Technical Methodology & Accomplishments

### 1. Stage 1: Masked-Diffusion Pretraining
The model was first pre-trained using a masked-diffusion objective adapted for block diffusion canvas generation:
* **Corrupting Process ($x_t$):** Clean canvas sequences ($x_0$) of length 256 were corrupted at dynamic noise levels $t \sim U(0.1, 1.0)$ using uniform token corruption across content tokens.
* **Loss Function:** Cross-entropy over masked/corrupted token positions against target completion tokens ($x_0$).

$$L_{\text{SFT}} = -\sum_{i \in \text{corrupted}} \log P_\theta(x_{0, i} \mid x_t, \text{prompt})$$

### 2. Stage 2: Policy Gradient via GRPO Alignment
Following SFT pretraining, the model underwent **Group Relative Policy Optimization (GRPO)**:
1. **Sampling:** For each prompt, $N=4$ diverse candidates were generated with 48 denoising steps (`max_denoising_steps=48`).
2. **Perplexity Scoring:** A auxiliary evaluator (`MiniCPM5-1B-Claude-Opus-Fable5-V2`) evaluated trace quality by calculating target perplexity given the prompt and generated reasoning trace:
   
   $$\text{Reward}_i = -\text{PPL}_{\text{Evaluator}}(\text{Target} \mid \text{Prompt} + \text{Trace}_i)$$

3. **Advantage Calculation:** Standardized relative advantages were computed per candidate group:
   
   $$A_i = \frac{R_i - \bar{R}}{\sigma_R + \epsilon}$$

4. **Policy Gradient Step:** Backpropagated advantages directly into teacher-forced canvas token log-probabilities to align non-autoregressive canvas generation toward higher-reasoning quality traces.

### 3. VRAM Optimization & Infrastructure Breakthroughs
Fine-tuning a 26B-A4B MoE block-diffusion model alongside a secondary reward evaluator model required custom infrastructure fixes:
* **MoE Activation Checkpointing:** Fixed the 95 GB VRAM activation explosion in Hugging Face's `grouped_mm_experts_forward` by configuring `use_gradient_checkpointing=True` on `FastModel`.
* **Cache Management:** Implemented dynamic CUDA memory purging between sampling rollouts and backpropagation passes to ensure single-GPU execution compatibility (A100/H100 80GB+).

---

## How to Use

### Loading the Fine-Tuned Adapter via Unsloth / Hugging Face

```python
import torch
from unsloth import FastModel

# Load Base Model and Adapter
model, processor = FastModel.from_pretrained(
    model_name="dataopsnick/diffusiongemma-26B-A4B-it-ByzantineSilk",
    dtype=torch.bfloat16,
    device_map={"": 0},
)

# Enable fast inference mode (re-enables KV cache for sampling)
FastModel.for_inference(model)

# Prepare Prompt
messages = [{"role": "user", "content": "Explain the moral of the fable of the North Wind and the Sun."}]
inputs = processor.apply_chat_template(
    messages,
    tokenize=True,
    add_generation_prompt=True,
    return_tensors="pt"
).to("cuda")

# Generate via Denoising
generation_config = model.generation_config
generation_config.max_denoising_steps = 48
generation_config.max_new_tokens = model.config.canvas_length

with torch.no_grad():
    output = model.generate(input_ids=inputs, generation_config=generation_config)

# Decode Output
prompt_len = inputs.shape[1]
gen_ids = output.sequences[0][prompt_len:]
decoded_text = processor.tokenizer.decode(gen_ids.tolist(), skip_special_tokens=True)
print(decoded_text)
```

---

## Training Hyperparameters

| Hyperparameter | Pretraining (Stage 1) | GRPO Alignment (Stage 2) |
| :--- | :--- | :--- |
| **Learning Rate** | `1e-4` | `1e-6` |
| **Optimizer** | AdamW $(\beta_1=0.9, \beta_2=0.95)$ | AdamW |
| **Gradient Accumulation** | `4` | `1` |
| **Noise Threshold ($t_{lo}$)**| `0.1` | N/A |
| **Canvas Length** | `256` | `256` |
| **GRPO Group Size ($N$)**| N/A | `4` |
| **Denoise Steps** | N/A | `48` |
| **Precision** | `bfloat16` | `bfloat16` |

---

## Repository Artifacts

* `adapter_model.safetensors`: Trained LoRA adapter weights (r=64, alpha=128).
* `adapter_config.json`: PEFT configuration for target attention & MoE projections.
* `chat_template.jinja`: Jinja prompt formatting template for DiffusionGemma chat format.
* `processor_config.json` & `tokenizer.json`: Tokenizer configs tuned for canvas generation.

---

## Citation & Acknowledgments

* **Unsloth AI** for `FastModel` block-diffusion acceleration utilities.
* **Google DeepMind** for the base `DiffusionGemma` architecture.
* **Kassadin88** for the original `LM-5.1-1000000x` dataset.
* **Jackrong** for the cleaned `Jackrong/GLM5.1-Reasoning-1M-Cleaned` dataset.
* **edwixx** for the Heretic-abliterated `diffusiongemma-26B-A4B-it` pretrained model
* **Weidmann, Philipp Emanuel** for [`Heretic: Fully automatic censorship removal for language models`](https://github.com/p-e-w/heretic)
* **GnLOLot** for the `GnLOLot/MiniCPM5-1B-Claude-Opus-Fable5-V2-Thinking` VRAM-efficient eval model for LLM-as-Judge Critic