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

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

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 LoRA adapter only (load onto the base model with PEFT). 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-adapter")
model = AutoModelForMaskedLM.from_pretrained(
    "LLM-OS-Models2/Qwen3-0.6B-diffusion-bd3lm-justgrpo-run2v2-mixed-lora-adapter", 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.