RELAY-distilled Fast-dLLM v2 1.5B (c40m60, step 200)
A Fast-dLLM v2 student distilled from the two-pass RELAY teacher
brozonoyer/relay-fastdllm-v2-c40m60-relay-step200.
The student starts from Efficient-Large-Model/Fast_dLLM_v2_1.5B
and predicts without hidden-state carry (use_relay=false).
Single-pass refers to each student prediction step. A complete response still requires multiple block-diffusion decoding steps.
Training
- Data: 60,000 conversations, with 24,000 code examples from OpenCodeInstruct and 36,000 math examples from OpenMathInstruct-2.
- Objective: full-vocabulary forward KL against the frozen teacher's second-pass distribution, at masked target positions; distillation temperature 1.
- Teacher conditioning: nested masks with a default mask-rate gap of 0.1. The first pass runs on the noisier input; the second uses the less-masked input and detached, right-shifted RELAY hidden feedback.
- Training steps: 200; learning rate: 5e-6; cosine schedule; warmup ratio: 0.03.
- Effective batch size: 64 (8 devices, batch size 1, gradient accumulation 8).
- Diffusion block size: 32; checkpoint weights: BF16.
- Final reported training loss: 0.021416714957449585.
See the source repository for the distillation objective and generation implementation.
Loading and generation
This checkpoint uses custom block-diffusion model code included in this repository.
It was validated with Transformers 4.53.1 and PyTorch 2.14.0+cu130.
The model code also requires einops; a CUDA GPU is recommended for generation.
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "divelab/relay-fastdllm-v2-c40m60-distilled-step200"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
trust_remote_code=True,
torch_dtype=torch.bfloat16,
).cuda().eval()
prompt = tokenizer.apply_chat_template(
[{"role": "user", "content": "Write a Python function to merge two sorted lists."}],
add_generation_prompt=True,
return_tensors="pt",
).cuda()
with torch.inference_mode():
output = model.generate(
prompt,
tokenizer=tokenizer,
max_new_tokens=128,
block_size=32,
small_block_size=8,
threshold=0.85,
temperature=0,
use_block_cache=False,
)
print(tokenizer.decode(output[0, prompt.shape[1]:], skip_special_tokens=True))
Evaluation and limitations
This upload contains the final distilled student checkpoint, tokenizer, and custom model code. It does not include optimizer states or intermediate training checkpoints. Standalone generation should be evaluated separately from speculative decoding: accuracy measured with a target verifier describes the verified target output rather than this student's standalone accuracy.