File size: 10,497 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 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 | import os
import logging
from multiprocessing import Process
import random
import pytest
from types import SimpleNamespace
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 vllm.forward_context import set_forward_context
from dinfer.model import LLaDAMoeModelLM, 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
moe_model_path = '/mnt/infra/dulun.dl/models/LLaDA-MoE/fusemoe/step45567_converted_hf_fusemoe'
# moe_model_path = '/data/dulun/models/llada-moe-sft/llada-moe-sft-model/7bA1b_anneal_19t_500B_further_8k_anneal_train_4k_ep3_v8p5/step45567_converted_hf_fusemoe/'
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
gpu_id = 1
device = torch.device(gpu_id)
decoder = ThresholdParallelDecoder(0, threshold=0.9, mask_id=156895, eos_id=156892, use_float64=True)
h_decoder = HierarchyDecoder(0, threshold=0.9, mask_id=156895, eos_id=156892, low_threshold=0.4)
tokenizer = AutoTokenizer.from_pretrained(moe_model_path, trust_remote_code=True, local_files_only=True)
input_ids = get_prompts(tokenizer, mask_id=156895, device=device)
model = None
@pytest.fixture(scope="session", autouse=True)
def init_vllm_dist(worker_id):
torch.cuda.set_device(gpu_id)
from vllm import distributed
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '37977'
distributed.init_distributed_environment(1, 0, 'env://', 0, 'nccl')
distributed.initialize_model_parallel(1, backend='nccl')
print("[Loading model]")
# setup EP
parallel_config = ParallelConfig(enable_expert_parallel = True)
with set_current_vllm_config(VllmConfig(parallel_config = parallel_config)):
model_config = AutoConfig.from_pretrained(moe_model_path, trust_remote_code=True, local_files_only=True)
global model
model = LLaDAMoeModelLM(config=model_config).eval()
model.load_weights(moe_model_path, torch_dtype=torch.bfloat16)
model = model.to(device)
yield
distributed.destroy_model_parallel()
distributed.destroy_distributed_environment()
def test_llada_moe_hierarchy():
# Test block-wise hierarchical diffusion MOE-LLM without KV-cache
print('Test block-wise hierarchical diffusion MOE-LLM without KV-cache')
dllm = BlockWiseDiffusionLLM(model, h_decoder, BlockIteratorFactory(), early_stop=True)
vllm_config = get_current_vllm_config()
with set_forward_context(None, vllm_config):
res = dllm.generate(input_ids, gen_length=128, block_length=32)
res1, nfe = generate_hierarchy(model, input_ids, gen_length=128, block_length=32, threshold=0.9, mask_id=156895, eos_id=156892,decoding='hierarchy_fast_v2',
low_threshold=0.4, remask_threshold=0.4)
res1 = res1[res1 != 156892]
assert res.shape[1] == len(res1)
res1 = res1.to(res.device)
assert torch.all(res == res1)
def test_llada_moe_blockwise():
# Test generation without cache.
print('Test block-wise diffusion MOE-LLM without KV-cache')
dllm = BlockWiseDiffusionLLM(model, decoder, BlockIteratorFactory(), early_stop=True)
vllm_config = get_current_vllm_config()
with set_forward_context(None, vllm_config):
res = dllm.generate(input_ids, gen_length=128, block_length=32)
res1, nfe = generate(model, input_ids, gen_length=128, block_length=32, threshold=0.9, mask_id=156895, eos_id=156892)
res2, nfe = generate_merge(model, input_ids, None, gen_length=128, block_length=32, threshold=0.9, mask_id=156895, eos_id=156892, parallel_decoding='threshold', early_stop=False,)
res1 = res1[res1 != 156892]
res2 = res2[res2 != 156892]
assert res.shape[1] == len(res1)
assert res.shape[1] == len(res2)
res1 = res1.to(res.device)
res2 = res2.to(res.device)
assert torch.all(res == res1)
assert torch.all(res == res2)
def test_llada_moe_batching():
# Test generation without cache with batch size == 2.
dllm = BlockWiseDiffusionLLM(model, decoder, BlockIteratorFactory(), early_stop=True)
print('Test block-wise diffusion MOE-LLM without KV-cache and batch size == 2')
input_ids2 = get_prompts(tokenizer, mask_id=156895, device=device, num=2)
vllm_config = get_current_vllm_config()
with set_forward_context(None, vllm_config):
res2 = dllm.generate(input_ids2, gen_length=128, block_length=32)
assert res2.shape[0] == 2
def test_llada_moe_itersmooth():
# Test generation with iteration smooth without kv-cache.
print('Test block-wise diffusion MOE-LLM with iteration smooth without kv-cache')
dllm = IterSmoothDiffusionLLM(model, decoder, BlockIteratorFactory(), early_stop=True)
dllm1 = IterSmoothDiffusionLLM_test(model, decoder, BlockIteratorFactory(), early_stop=True)
vllm_config = get_current_vllm_config()
with set_forward_context(None, vllm_config):
res = dllm.generate(input_ids, gen_length=128, block_length=32)
res1 = dllm1.generate(input_ids, gen_length=128, block_length=32)
assert dllm.num_forwards == dllm1.num_forwards
assert dllm.cache_updates == 0
assert res.shape[1] == res1.shape[1]
res1 = res1.to(res.device)
assert torch.all(res == res1)
def test_llada_moe_dual_cache():
# Test generation with dual cache
print('Test block-wise diffusion MOE-LLM with dual KV-cache')
dllm = BlockWiseDiffusionLLM(model, decoder, BlockIteratorFactory(), early_stop=True, cache_factory=KVCacheFactory('dual'))
vllm_config = get_current_vllm_config()
with set_forward_context(None, vllm_config):
res = dllm.generate(input_ids, gen_length=256, block_length=32)
res1, nfe = generate_with_dual_cache(model, input_ids, gen_length=256, block_length=32, threshold=0.9, mask_id=156895, eos_id=156892)
res1 = res1[res1 != 156892]
assert res.shape[1] == len(res1)
res1 = res1.to(res.device)
assert torch.all(res == res1)
def test_llada_moe_dual_cache_batching():
# Test generation with dual cache with batch size == 2
print('Test block-wise diffusion MOE-LLM with dual KV-cache and batch size == 2')
dllm = BlockWiseDiffusionLLM(model, decoder, BlockIteratorFactory(), early_stop=True, cache_factory=KVCacheFactory('dual'))
input_ids2 = get_prompts(tokenizer, mask_id=156895, device=device, num=2)
vllm_config = get_current_vllm_config()
with set_forward_context(None, vllm_config):
res2 = dllm.generate(input_ids2, gen_length=256, block_length=32)
assert res2.shape[0] == 2
def test_llada_moe_itersmooth_cache():
# Test generation with iteration smooth with kv-cache.
print('Test block-wise diffusion MOE-LLM with iteration smooth with kv-cache')
dllm = IterSmoothDiffusionLLM(model, decoder, BlockIteratorFactory(), early_stop=True, cache_factory=KVCacheFactory('dual'))
dllm1 = IterSmoothDiffusionLLM_test(model, decoder, BlockIteratorFactory(), early_stop=True, cache_factory=KVCacheFactory('dual'))
vllm_config = get_current_vllm_config()
with set_forward_context(None, vllm_config):
res = dllm.generate(input_ids, gen_length=128, block_length=32)
res1 = dllm1.generate(input_ids, gen_length=128, block_length=32)
assert dllm.num_forwards == dllm1.num_forwards
assert dllm.cache_updates > 0
assert dllm.cache_updates == dllm1.cache_updates
assert res.shape[1] == res1.shape[1]
res1 = res1.to(res.device)
assert torch.all(res == res1)
def test_llada_moe_itersmooth_vicinity_cache():
# Test generation with iteration smooth and vicinity cache update.
print('Test block-wise diffusion MOE-LLM with iteration smooth with vicinity cache update')
dllm = IterSmoothWithVicinityCacheDiffusionLLM(model, decoder, BlockIteratorFactory(), early_stop=True, cache_factory=KVCacheFactory('dual'))
dllm1 = IterSmoothWithVicinityCacheDiffusionLLM_test(model, decoder, BlockIteratorFactory(), early_stop=True, cache_factory=KVCacheFactory('dual'))
vllm_config = get_current_vllm_config()
with set_forward_context(None, vllm_config):
res = dllm.generate(input_ids, gen_length=128, block_length=32)
res1 = dllm1.generate(input_ids, gen_length=128, block_length=32)
assert dllm.num_forwards == dllm1.num_forwards
assert dllm.cache_updates > 0
assert dllm.cache_updates == dllm1.cache_updates
assert res.shape[1] == res1.shape[1]
res1 = res1.to(res.device)
assert torch.all(res == res1)
|