Views
No views yet
pip install torch transformers accelerate1import torch
2import numpy as np
3import torch.nn.functional as F
4
5from transformers import AutoTokenizer, AutoModelForMaskedLM
6
7
8def add_gumbel_noise(logits, temperature):
9 if temperature == 0:
10 return logits
11 logits = logits.to(torch.float64)
12 noise = torch.rand_like(logits, dtype=torch.float64)
13 gumbel_noise = (- torch.log(noise)) ** temperature
14 return logits.exp() / gumbel_noise
15
16
17def get_num_transfer_tokens(mask_index, steps):
18 mask_num = mask_index.sum(dim=1, keepdim=True)
19 base = mask_num // steps
20 remainder = mask_num % steps
21 num_transfer_tokens = torch.zeros(mask_num.size(0), steps, device=mask_index.device, dtype=torch.int64) + base
22
23 for i in range(mask_num.size(0)):
24 num_transfer_tokens[i, :remainder[i]] += 1
25 return num_transfer_tokens
26
27
28@ torch.no_grad()
29def generate(model, prompt, steps=128, gen_length=128, block_length=64, temperature=0.0, cfg_scale=0., remasking='random'):
30 mask_id = tokenizer.mask_token_id
31 x = torch.full((1, prompt.shape[1] + gen_length), mask_id, dtype=torch.long).to(model.device)
32 x[:, :prompt.shape[1]] = prompt.clone()
33 prompt_index = (x != mask_id)
34
35 assert gen_length % block_length == 0
36 num_blocks = gen_length // block_length
37 assert steps % num_blocks == 0
38 steps = steps // num_blocks
39
40 for num_block in range(num_blocks):
41 block_mask_index = (x[:, prompt.shape[1] + num_block * block_length: prompt.shape[1] + (num_block + 1) * block_length:] == mask_id)
42 num_transfer_tokens = get_num_transfer_tokens(block_mask_index, steps)
43 for i in range(steps):
44 mask_index = (x == mask_id)
45 if cfg_scale > 0.:
46 un_x = x.clone()
47 un_x[prompt_index] = mask_id
48 x_ = torch.cat([x, un_x], dim=0)
49 logits = model(x_).logits
50 logits, un_logits = torch.chunk(logits, 2, dim=0)
51 logits = un_logits + (cfg_scale + 1) * (logits - un_logits)
52 else:
53 logits = model(x).logits
54
55 logits_with_noise = add_gumbel_noise(logits, temperature=temperature)
56 x0 = torch.argmax(logits_with_noise, dim=-1) # b, l
57
58 if remasking == 'low_confidence':
59 p = F.softmax(logits, dim=-1)
60 x0_p = torch.squeeze(
61 torch.gather(p, dim=-1, index=torch.unsqueeze(x0, -1)), -1) # b, l
62 elif remasking == 'random':
63 x0_p = torch.rand((x0.shape[0], x0.shape[1]), device=x0.device)
64 else:
65 raise NotImplementedError(remasking)
66
67 x0_p[:, prompt.shape[1] + (num_block + 1) * block_length:] = -np.inf
68
69 x0 = torch.where(mask_index, x0, x)
70 confidence = torch.where(mask_index, x0_p, -np.inf)
71
72 transfer_index = torch.zeros_like(x0, dtype=torch.bool, device=x0.device)
73 for j in range(confidence.shape[0]):
74 _, select_index = torch.topk(confidence[j], k=num_transfer_tokens[j, i])
75 transfer_index[j, select_index] = True
76 x[transfer_index] = x0[transfer_index]
77
78 return x
79
80
81device = 'cuda'
82model = AutoModelForMaskedLM.from_pretrained('dllm-hub/ModernBERT-base-chat-v0.1', dtype=torch.bfloat16).to(device).eval()
83tokenizer = AutoTokenizer.from_pretrained('dllm-hub/ModernBERT-base-chat-v0.1')
84
85prompt = "Lily can run 12 kilometers per hour for 4 hours. After that, she runs 6 kilometers per hour. How many kilometers can she run in 8 hours?"
86m = [
87 {"role": "system", "content": "You are a helpful AI assistant."},
88 {"role": "user", "content": prompt}
89]
90prompt = tokenizer.apply_chat_template(m, add_generation_prompt=True, tokenize=False)
91
92input_ids = tokenizer(prompt)['input_ids']
93input_ids = torch.tensor(input_ids).to(device).unsqueeze(0)
94
95text = generate(model, input_ids, steps=128, gen_length=128, block_length=64, temperature=0.0, cfg_scale=0.0, remasking='random')
96print(tokenizer.batch_decode(text[:, input_ids.shape[1]:], skip_special_tokens=False)[0])| Parameter | Description | Default |
|---|---|---|
max_new_tokens | Number of tokens to generate | 128 |
steps | Number of diffusion denoising iterations | 128 |
temperature | Sampling temperature; set to 0.0 for deterministic generation | 0.0 |
block_length | Token block size used during iterative denoising | 64 |
cfg_scale | Classifier-free guidance scale controlling instruction adherence (higher = more deterministic) | 0.0 |
remasking | Strategy for re-masking during each denoising step (random, none, or confidence) | random |
1python -u examples/bert/chat.py \
2 --model_name_or_path dllm-hub/ModernBERT-base-chat-v0.1 \
3 --chat True| LAMBADA | GSM8K | CEval | BBH | MATH | MMLU | Winogrande | HellaSwag | CMMLU | |
|---|---|---|---|---|---|---|---|---|---|
| ModernBERT-base-chat-v0.1 | 49.3 | 5.9 | 25.0 | 17.9 | 3.1 | 26.1 | 49.7 | 41.0 | 24.3 |
| ModernBERT-large-chat-v0.1 | 46.3 | 17.1 | 24.6 | 25.1 | 3.8 | 33.5 | 53.1 | 45.0 | 27.5 |
1bash examples/bert/eval.sh \
2 --model_name_or_path "dllm-hub/ModernBERT-base-chat-v0.1"1@misc{zhou2026dllm,
2 title={dLLM: Simple Diffusion Language Modeling},
3 author={Zhanhui Zhou and Lingjie Chen and Hanghang Tong and Dawn Song},
4 year={2026},
5 eprint={2602.22661},
6 archivePrefix={arXiv},
7 primaryClass={cs.CL},
8 url={https://arxiv.org/abs/2602.22661},
9}