File size: 8,021 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 | import torch
import numpy as np
from torch._C import dtype
import torch.nn.functional as F
import os
from transformers import AutoTokenizer, AutoModel, AutoConfig, AutoModelForCausalLM
import torch.distributed as dist
import time
import tqdm
from sglang.srt.server_args import ServerArgs
from sglang.srt.layers.moe import initialize_moe_config
from dinfer.model.modeling_llada2_moe_sglang import LLaDA2SGLangLM
from dinfer.decoding.diffusion_runner import ModelRunner
from dinfer.decoding import serving
from queue import Empty
from dinfer.decoding.serving import ServerGroup
from dinfer import BlockIteratorFactory, KVCacheFactory, BlockDiffusionLLM
from dinfer import ThresholdParallelDecoder,CreditThresholdParallelDecoder, HierarchyDecoder, BlockWiseDiffusionLLM, IterSmoothDiffusionLLM, VicinityCacheDiffusionLLM, IterSmoothWithVicinityCacheDiffusionLLM
import logging
import traceback
import json
from multiprocessing import Process
from pathlib import Path
import pytest
from dinfer.model import LLaDA2MoeModelLM
from dinfer import BlockIteratorFactory, KVCacheFactory, SamplingParams, DiffusionLLMServing
from dinfer import ThresholdParallelDecoder, BlockDiffusionLLMAttnmask, BlockDiffusionLLM
import difflib
import time
#model_path = '/mnt/dllm/luxiaocheng/moe-mini-v2-e256-1009-fp8-ml4-grouprouter-20T-mdmcpt-block-diffusion-bl32-4k-noshift-100B'
model_path = '/mnt/infra/dulun.dl/models/dllm-mini/block-diffusion-sft-2k-v2-full-bd/LLaDA2-mini-preview-ep4-v0'
#model_path = '/mnt/infra/dulun.dl/models/dllm-mini/block-diffusion-sft-2k-v2-full-bd/LLaDA2-mini-preview-ep4-v0'
dataset_path = '/ossfs/workspace/dumped_prompts'
dataset='openai_humaneval'
FILE_PATH = Path(__file__).resolve()
sample_path = FILE_PATH.with_name(f"{FILE_PATH.stem[:-8]}_sample.json")
model = None
gpu_id = 0
device = torch.device(gpu_id)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
decoder = ThresholdParallelDecoder(temperature=0, threshold=0.9, mask_id=156895, eos_id=156892)
def generate_test(*args, tout=10, **kwargs):
while True:
if len(args) >= 2:
dllm, device=args[0], args[1]
req_q = kwargs.get('req_q', None)
res = kwargs.get('res_q', None)
data = req_q.get()
if isinstance(data, str):
assert data == 'stop'
break
else:
input_ids, gen_len, block_len = data
raise RuntimeError
def init_sglang_dist():
torch.cuda.set_device(gpu_id)
from sglang.srt import distributed
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '40399'
distributed.init_distributed_environment(1, 0, 'env://', 0, 'nccl')
distributed.initialize_model_parallel(1, 1, 1, backend='nccl')
print("[Loading model]")
from sglang.srt.layers.dp_attention import initialize_dp_attention
model_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
server_args = ServerArgs(model_path=model_path, enable_dp_attention=True, trust_remote_code=True, tp_size=1, dp_size = 1, pp_size = 1)
try:
from sglang.srt.server_args import set_global_server_args_for_scheduler
except ImportError:
pass
else:
set_global_server_args_for_scheduler(server_args)
initialize_dp_attention(
server_args=server_args,
model_config=model_config,
)
initialize_moe_config(server_args)
model = LLaDA2SGLangLM(config=model_config, expert_map_path='.').eval()
torch.set_default_dtype(torch.bfloat16)
model.load_weights(model_path, device=device)
initialize_moe_config(server_args)
model = model.to(device)
max_length = 2048
model = ModelRunner(model, device, server_args=server_args, max_length=max_length)
return model
model = init_sglang_dist()
def run_bd(use_kvcache):
with open(sample_path, "r") as f:
samples = json.load(f)
ans = []
for sample in samples:
prompt = [sample['question']]
prompt[0] = '<role>SYSTEM</role>detailed thinking off<|role_end|><role>HUMAN</role>'+prompt[0]+'<|role_end|><role>ASSISTANT</role>'
input_ids = tokenizer(prompt)['input_ids']
input_ids = torch.tensor(input_ids).to(device)
if not use_kvcache:
dllm = BlockDiffusionLLMAttnmask(model, decoder, BlockIteratorFactory(use_block_diffusion=True), early_stop=True)
else:
dllm = BlockDiffusionLLM(model, decoder, BlockIteratorFactory(start_block_align=True, use_block_diffusion=True),
cache_factory=KVCacheFactory('prefix', is_bd_model=True, max_length=2048), early_stop=True,
maximum_unroll=4, expected_tpf=4, backend='sglang')
out = dllm.generate(input_ids, gen_length=256, block_length=32)
new_ans = tokenizer.decode(out[0, input_ids.shape[1]:], skip_special_tokens=True)
#assert(new_ans == sample['answer'])
ans.append(new_ans)
return ans
def run_bd_serving(use_kvcache):
with open(sample_path, "r") as f:
samples = json.load(f)
sample_params = SamplingParams(threshold=0.9, cache='prefix', temperature=0., early_stop=True, cont_weight=0, prefix_look=0,
after_look=0, warmup_steps=0, enable_torch_compile=True, mask_id=156895, eos_id=156892, parallel_decoding='threshold',
use_credit=False, use_bd=True, max_length=2048)
dllm_server = DiffusionLLMServing(model_path, model_type='llada2-mini', sample_params=sample_params, server_port=40567, num_gpus=1, dp_size=1, tpep_size=1, backend='sglang')
ans = []
for sample in samples:
prompt = [sample['question']]
prompt[0] = '<role>SYSTEM</role>detailed thinking off<|role_end|><role>HUMAN</role>'+prompt[0]+'<|role_end|><role>ASSISTANT</role>'
input_ids = tokenizer(prompt)['input_ids']
input_ids = torch.tensor(input_ids).to(device)
out = dllm_server.generate(input_ids, gen_length=256, block_length=32)
new_ans = tokenizer.decode(out[0, input_ids.shape[1]:], skip_special_tokens=True)
#assert(new_ans == sample['answer'])
ans.append(new_ans)
return ans
def run_bd_serving_error(use_kvcache):
with open(sample_path, "r") as f:
samples = json.load(f)
tout=10
sample_params = SamplingParams(threshold=0.9, cache='prefix', temperature=0., early_stop=True, cont_weight=0, prefix_look=0,
after_look=0, warmup_steps=0, enable_torch_compile=True, mask_id=156895, eos_id=156892, parallel_decoding='threshold',
use_credit=False, use_bd=True, max_length=2048)
dllm_server = DiffusionLLMServing(model_path, model_type='llada2-mini', sample_params=sample_params, server_port=40567, num_gpus=4, dp_size=1, tpep_size=4, backend='sglang',
timeout=tout)
ans = []
for sample in samples:
prompt = [sample['question']]
prompt[0] = '<role>SYSTEM</role>detailed thinking off<|role_end|><role>HUMAN</role>'+prompt[0]+'<|role_end|><role>ASSISTANT</role>'
input_ids = tokenizer(prompt)['input_ids']
input_ids = torch.tensor(input_ids).to(device)
try:
out = dllm_server.generate(input_ids, gen_length=128, block_length=128) # return the error type in server group
except Exception as e:
if isinstance(e, RuntimeError):
dllm_server.stop_serving()
return
else:
raise ValueError("The return should be RuntimeError")
new_ans = tokenizer.decode(out[0, input_ids.shape[1]:], skip_special_tokens=True)
ans.append(new_ans)
return ans
def test_bd():
ans_cache = run_bd(use_kvcache=True)
ans_serving = run_bd_serving(use_kvcache=True)
for i in range(len(ans_cache)):
assert(ans_cache[i] == ans_serving[i])
if __name__ == '__main__':
# ans = run_bd_serving(use_kvcache=True) # using serving to generate response.
model = init_sglang_dist() # test the init of sglang model.
test_bd()
# ans = run_bd_serving_error(use_kvcache=True) # with code with error capture. When timeout, the process will return runtime error.
|