File size: 7,026 Bytes
3a464db | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | import os
import logging
from multiprocessing import Process
import random
import torch
import torch.distributed as dist
from transformers import AutoTokenizer, AutoModel, AutoConfig
from vllm.config import CompilationConfig, ParallelConfig
from vllm.config import VllmConfig, set_current_vllm_config, get_current_vllm_config
from dinfer.model import LLaDAModelLM
from dinfer import BlockWiseDiffusionLLM, VicinityCacheDiffusionLLM, IterSmoothDiffusionLLM, IterSmoothWithVicinityCacheDiffusionLLM, BlockWiseDiffusionLLMWithSP
from dinfer import ThresholdParallelDecoder, HierarchyDecoder
from dinfer import DiffusionLLMServing, SamplingParams
from dinfer.model.modeling_llada_fastdllm import LLaDAModelLM as LLaDAModelLM_fastdllm
from dinfer.decoding.generate_fastdllm import generate, generate_with_prefix_cache, generate_with_dual_cache
from dinfer.decoding.generate_dist import generate as generate_sp
from dinfer.decoding.generate_uniform import BaseDiffusionIteration
from dinfer.decoding.generate_hierarchy import generate_hierarchy
from dinfer.decoding.utils import TokenArray, DistAlignedTokenArray, BlockIterator, BlockIteratorFactory, KVCacheFactory, gather_sequence_block, BlockLoc
from dinfer.decoding.utils import DiffusionKVCacheManager
from dinfer.decoding.generate_merge import generate_merge
os.environ['TOKENIZERS_PARALLELISM'] = 'false'
from test_generate import IterSmoothDiffusionLLM as IterSmoothDiffusionLLM_test
from test_generate import IterSmoothWithVicinityCacheDiffusionLLM as IterSmoothWithVicinityCacheDiffusionLLM_test
model_path = "/mnt/infra/myx/models/LLaDA-1.5/"
def get_prompts(tokenizer, mask_id, device, num=1):
prompt = "Lily can run 12 kilometers per hour for 4 hours. After that, she can run 6 kilometers per hour. How many kilometers can she run in 8 hours? "
m = [{"role": "user", "content": prompt}, ]
prompt = tokenizer.apply_chat_template(m, add_generation_prompt=True, tokenize=False)
input_ids1 = torch.tensor(tokenizer(prompt)['input_ids']).to(device).unsqueeze(0)
len1 = input_ids1.shape[1]
if num == 2:
prompt = "Lily can run 12 kilometers per hour for 4 hours. How many kilometers can she run in 4 hours? "
m = [{"role": "user", "content": prompt}, ]
prompt = tokenizer.apply_chat_template(m, add_generation_prompt=True, tokenize=False)
input_ids2 = torch.tensor(tokenizer(prompt)['input_ids']).to(device).unsqueeze(0)
len2 = input_ids2.shape[1]
ret = torch.zeros(2, max(len1, len2), dtype=input_ids1.dtype)
ret[0, 0:len1] = input_ids1
ret[1, 0:len2] = input_ids2
else:
ret = input_ids1
return ret
class SimulateBlockIterator:
""" This class simulates the block iterator in VicinityCacheDiffusionLLM.
"""
def __init__(self, x, block_length, mask_id):
self.x = x
self.iter = 0
self.block_length = block_length
self.mask_id = mask_id
def __iter__(self):
self.iter = 0
return self
def move_next(self):
current_block_start = self.x.prompt.shape[1] + self.iter * self.block_length
current_block_end = current_block_start + self.block_length
current_block_end = min(current_block_end, self.x.total_length)
# If all tokens have been decoded, move to the next block.
if torch.all(self.x[:, current_block_start:current_block_end] != self.mask_id):
self.iter += 1
def __next__(self):
self.move_next()
current_block_start = self.x.prompt.shape[1] + self.iter * self.block_length
if current_block_start >= self.x.total_length:
raise StopIteration
current_block_end = current_block_start + self.block_length
current_block_end = min(current_block_end, self.x.total_length)
return BlockLoc(current_block_start, current_block_end), self.x[:, current_block_start:current_block_end]
class SimulateBlockIteratorFactory:
def create(self, x, block_length):
return SimulateBlockIterator(x, block_length, 126336)
torch.cuda.set_device(0)
device = torch.device(0)
config = AutoConfig.from_pretrained(model_path, trust_remote_code=True, local_files_only=True)
config.flash_attention = True
config.train_max_sequence_length = 4096
model = LLaDAModelLM.from_pretrained(model_path, torch_dtype=torch.bfloat16, config=config).eval()
model.init_h2e_module()
model = model.to(device)
fastdllm_model = LLaDAModelLM_fastdllm.from_pretrained(model_path, torch_dtype=torch.bfloat16, config=config).eval()
fastdllm_model = fastdllm_model.to(device)
decoder = ThresholdParallelDecoder(0, threshold=0.9, use_float64=True)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, local_files_only=True)
input_ids = get_prompts(tokenizer, mask_id=126336, device=device)
batch_size = 1
input_ids = input_ids.clone().detach().to(device).repeat(batch_size, 1)
def test_sw_dual_cache():
print('Test sliding-window diffusion LLM with dual KV-cache')
dllm = VicinityCacheDiffusionLLM(model, decoder, SimulateBlockIteratorFactory(), KVCacheFactory('dual'))
res = dllm.generate(input_ids, gen_length=128, block_length=32)
res1, nfe = generate_with_dual_cache(fastdllm_model, input_ids, gen_length=128, block_length=32, threshold=0.9)
res1 = res1[res1 != 126081]
assert res.shape[1] == len(res1)
res1 = res1.to(res.device)
assert torch.all(res == res1)
def test_prefix_cache():
print('Test block-wise diffusion LLM with prefix KV-cache')
dllm = BlockWiseDiffusionLLM(model, decoder, BlockIteratorFactory(), early_stop=True, cache_factory=KVCacheFactory('prefix'))
res = dllm.generate(input_ids, gen_length=128, block_length=32)
res1, nfe = generate_with_prefix_cache(fastdllm_model, input_ids, gen_length=128, block_length=32, threshold=0.9)
res1 = res1[res1 != 126081]
assert res.shape[1] == len(res1)
res1 = res1.to(res.device)
assert torch.all(res == res1)
def test_dual_cache():
print('Test block-wise diffusion LLM with dual KV-cache')
dllm = BlockWiseDiffusionLLM(model, decoder, BlockIteratorFactory(), cache_factory=KVCacheFactory('dual'), early_stop=True)
res = dllm.generate(input_ids, gen_length=128, block_length=32)
res1, nfe = generate_with_dual_cache(fastdllm_model, input_ids, gen_length=128, block_length=32, threshold=0.9)
res1 = res1[res1 != 126081]
assert res.shape[1] == len(res1)
res1 = res1.to(res.device)
assert torch.all(res == res1)
def test_itersmooth():
print('Test block-wise diffusion LLM with dual KV-cache')
dllm = IterSmoothDiffusionLLM(model, decoder, BlockIteratorFactory(), cache_factory=KVCacheFactory('dual'), early_stop=True)
res = dllm.generate(input_ids, gen_length=128, block_length=32)
if __name__ == '__main__':
print("Start test iter smooth...")
test_itersmooth()
print("Start test sliding window with dual cache...")
test_sw_dual_cache()
print("Start test prefix cache...")
test_prefix_cache()
print("Start test dual cache...")
test_dual_cache()
|