Views
No views yet
1import torch
2import transformers
3from transformers.cache_utils import DynamicCache
4# refer to https://github.com/iiiutch-ii/RemeDi/blob/main/RL-code
5from networks.block_llada.modelling_llada_bitowel import LLaDAUPMModelLM
6
7@torch.no_grad()
8def generate_block_diffusion(
9 model,
10 conv,
11 tokenizer,
12 device,
13 num_generations,
14 kv_cache=None,
15 steps: int = 32,
16 max_length = 1024,
17 block_size = 32,
18 mask_token_id = 126336,
19 eos_id = 126081,
20):
21 m = [conv]
22 prompts = tokenizer.apply_chat_template(m, add_generation_prompt=True, tokenize=False)
23 inputs = tokenizer(prompts, return_tensors='pt', padding=True, padding_side='left')
24 x_t = inputs['input_ids'].to(device)
25
26 attention_mask = inputs['attention_mask'].to(device)
27 prompt_len = attention_mask.sum(dim=1)
28 attn_bias = torch.where(
29 attention_mask + attention_mask.T > 0,
30 0, -torch.inf
31 )[None, None].repeat(x_t.shape[0], 1, 1, 1)
32
33 x_t = x_t.repeat(num_generations, 1)
34 prompt_len = prompt_len.repeat(num_generations)
35 attn_bias = attn_bias.repeat(num_generations, 1, 1, 1)
36 batch_size = x_t.shape[0]
37
38 position_ids = torch.arange(x_t.shape[1], device=x_t.device, dtype=torch.long).unsqueeze(0) - (1 - attention_mask).sum(dim=-1)
39 if kv_cache is None:
40 kv_cache = DynamicCache()
41
42 # cache prompt first
43 with torch.autocast(device_type="cuda", enabled=True, dtype=torch.bfloat16):
44 model(
45 x_t,
46 kv_cache=kv_cache,
47 update_kv_cache=True,
48 )
49
50 cur_blocks = 0
51 responses = [x_t]
52 is_eos_meet = torch.zeros((batch_size,), device=x_t.device, dtype=torch.bool)
53
54 while (cur_blocks * block_size) < max_length:
55 x_t = torch.full((batch_size, block_size), fill_value=mask_token_id, device=device, dtype=torch.long)
56
57 position_ids = torch.arange(
58 cur_blocks * block_size,
59 (cur_blocks + 1) * block_size,
60 device=x_t.device, dtype=torch.long).unsqueeze(0) + prompt_len.unsqueeze(1)
61
62 num_transfer_tokens = torch.tensor([block_size // steps for _ in range(steps)])
63 if block_size % steps != 0:
64 num_transfer_tokens[-block_size % steps:] += 1
65 # cumsum
66 num_transfer_tokens = num_transfer_tokens.cumsum(dim=0)
67
68 for i in range(steps):
69 mask_index = (x_t == mask_token_id)
70
71 with torch.autocast(device_type="cuda", enabled=True, dtype=torch.bfloat16):
72 out = model(
73 x_t,
74 position_ids=position_ids,
75 kv_cache=kv_cache,
76 )
77 logits = out.logits.to(torch.float32)
78 x0 = torch.argmax(logits, dim=-1) # b, l
79 x0 = torch.where(mask_index, x0, x_t)
80
81 upm_prob = logits.gather(dim=-1, index=x0.unsqueeze(-1)).squeeze(-1)
82 samples = torch.topk(upm_prob, k=num_transfer_tokens[i], dim=-1).indices
83
84 bs_idx = torch.arange(batch_size, dtype=samples.dtype).unsqueeze(1)
85 remask_index = torch.ones_like(x_t).bool()
86 remask_index[bs_idx, samples] = False
87
88 x_t = torch.where(remask_index, mask_token_id, x0)
89
90 responses.append(x_t.clone())
91 cur_blocks += 1
92 if is_eos_meet.all(): break
93
94 # update kv_cache
95 with torch.autocast(device_type="cuda", enabled=True, dtype=torch.bfloat16):
96 model(
97 x_t,
98 position_ids=position_ids,
99 kv_cache=kv_cache,
100 update_kv_cache=True,
101 )
102
103
104 response_tokens = torch.cat(responses, dim=1)
105 responses = []
106 responses_length = []
107 for i in range(batch_size):
108 if eos_id in response_tokens[i]:
109 eos_token_idx = (response_tokens[i] == eos_id).nonzero(as_tuple=True)[0][0].item()
110 resp_token = response_tokens[i, prompt_len[i]:eos_token_idx]
111 else:
112 resp_token = response_tokens[i, prompt_len[i]:]
113 responses.append(tokenizer.decode(resp_token, skip_special_tokens=True))
114 responses_length.append(resp_token.shape[0])
115
116 return responses
117
118def main(
119 ckpt_path = 'iiiutch/RemeDi-Instruct',
120 seed: int = 112,
121):
122 torch.manual_seed(seed)
123 device = 'cuda'
124
125 tokenizer = transformers.AutoTokenizer.from_pretrained(ckpt_path)
126
127 model = LLaDAUPMModelLM.from_pretrained(
128 ckpt_path,
129 torch_dtype=torch.bfloat16,
130 )
131 model.eval().requires_grad_(False).to(device)
132
133 conv = []
134 while True:
135 conv = []
136 print('=' * 20)
137 prompt = input("User: ").strip()
138 print('Assistant: ', end='')
139 conv = [{'role': 'user', 'content': prompt}]
140
141 inputs = generate_block_diffusion(
142 model,
143 conv,
144 tokenizer,
145 reward_fn=None,
146 device=device,
147 viz=True,
148 num_generations=1,
149 steps=32, max_length=1024, block_size=32,
150 )
151
152 conv.append({'role': 'assistant', 'content': inputs[0]})
153
154
155if __name__ == "__main__":
156 main()
157