### 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