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

Open In Colab

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):

Sampling: For each prompt, $N=4$ diverse candidates were generated with 48 denoising steps (max_denoising_steps=48).
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)$$
Advantage Calculation: Standardized relative advantages were computed per candidate group:

$$A_i = \frac{R_i - \bar{R}}{\sigma_R + \epsilon}$$
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
GnLOLot for the GnLOLot/MiniCPM5-1B-Claude-Opus-Fable5-V2-Thinking VRAM-efficient eval model for LLM-as-Judge Critic