Qwen3-0.6B-diffusion-bd3lm-justgrpo-run2v2-mixed-lora

Qwen3-0.6B-diffusion-bd3lm-justgrpo-run2v2-mixed-lora

JustGRPO-style RL (LoRA) applied to the tiny block-diffusion LM dllm-hub/Qwen3-0.6B-diffusion-bd3lm-v0.1. Trained on a round-robin mix of GSM8K + MATH + AceCode-Hard 21K (multi-domain RLVR) with verifiable rewards.

• Method & code: https://github.com/LLM-OS-Models/JustGRPO
• Papers this work builds on: JustGRPO (arXiv:2601.15165) · dLLM (arXiv:2602.22661) · BD3LM (arXiv:2503.09573)

How it was trained

• Rollouts: native block-diffusion sampling (block_size 32, temperature 1.0, 256 denoising steps, gen length 256) — AR-order rollout collapses on this base model, see the repo's ADAPTATION.md for the full analysis.
• Loss: exact autoregressive log-likelihood of each sampled token, computed in a single forward pass via the BD3LM [x0 || xt] concat-attention trick (mathematically identical to the per-token loop; verified to 2e-5 in fp32), weighted by GRPO group-normalized advantages with PPO-style clipping.
• LoRA r=128, alpha=64, dropout 0.05 on q/k/v/o/up/down/gate projections, lr 5e-5, 300 steps, 8 prompts x 16 rollouts per step, 1x H100.
• Rewards: per-domain: math +-1 exact match; code format + unit-test pass rate

Results

Evaluation in progress — see the GitHub repo's BENCHMARKS.md for the live table.

Usage

This repo contains the merged full model (base + LoRA). Generation uses block diffusion (NOT vanilla model.generate); the easiest path is the dllm sampler:

import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM
from dllm.core.samplers import BD3LMSampler, BD3LMSamplerConfig

tok = AutoTokenizer.from_pretrained("LLM-OS-Models2/Qwen3-0.6B-diffusion-bd3lm-justgrpo-run2v2-mixed-lora")
model = AutoModelForMaskedLM.from_pretrained(
    "LLM-OS-Models2/Qwen3-0.6B-diffusion-bd3lm-justgrpo-run2v2-mixed-lora", trust_remote_code=True, torch_dtype=torch.bfloat16
).cuda().eval()

sampler = BD3LMSampler(model=model, tokenizer=tok)
prompt = tok.apply_chat_template(
    [[{"role": "user", "content": "Natalia sold clips to 48 friends in April, "
       "then half as many in May. How many altogether?"}]],
    add_generation_prompt=True, tokenize=True)
seqs = sampler.sample(prompt, config=BD3LMSamplerConfig(
    max_new_tokens=256, steps=256, block_size=32, temperature=0.0))
print(tok.decode(seqs[0], skip_special_tokens=True))

To evaluate: use the dllm eval harness (--model a2d_bd3lm, max_new_tokens=256,steps=256,block_size=32) as described in https://github.com/LLM-OS-Models/JustGRPO.