''' This file is inspired by the code from https://github.com/NVlabs/Fast-dLLM ''' import accelerate import torch import random import torch.nn.functional as F from datasets import Dataset from tqdm import tqdm, trange import accelerate import random import numpy as np import json import time import datasets import json import time import datasets import os from transformers import AutoTokenizer, AutoConfig import torch.multiprocessing as mp from multiprocessing import Process from lm_eval.api.model import LM from lm_eval.__main__ import cli_evaluate from lm_eval.api.model import LM from lm_eval.api.registry import register_model from dinfer.model import LLaDAMoeModelLM, LLaDAModelLM, LLaDA2MoeModelLM from dinfer import BlockIteratorFactory, KVCacheFactory from dinfer import ThresholdParallelDecoder,CreditThresholdParallelDecoder, HierarchyDecoder, BlockWiseDiffusionLLM, IterSmoothDiffusionLLM, VicinityCacheDiffusionLLM, IterSmoothWithVicinityCacheDiffusionLLM, BlockDiffusionLLM from vllm import distributed from vllm.config import VllmConfig, set_current_vllm_config, get_current_vllm_config from vllm.config import ParallelConfig from dataclasses import dataclass datasets.config.HF_DATASETS_TRUST_REMOTE_CODE = True datasets.config.DOWNLOAD_TIMEOUT = 180 os.environ['TOKENIZERS_PARALLELISM'] = 'false' bucket_size = 32 used_buckets = [] def warmup_cudagraph(rank, device, dllm, gen_len, block_length, batch_size, vocab_size): if rank==0: print('warmup') print(used_buckets) iterator = tqdm(used_buckets) else: iterator = used_buckets offset = 0 for i in iterator: input_ids = torch.randint(0, vocab_size, (batch_size, i - gen_len+offset), dtype=torch.long, device=device) dllm.generate(input_ids, gen_length=gen_len, block_length=block_length) def cut_eos(data, eos_id=156892): eos_indices = (data[0] == eos_id).nonzero(as_tuple=True)[0] if eos_indices.numel() > 0: first_eos_idx = eos_indices[0].item() return data[:, :first_eos_idx] else: return data @ torch.no_grad() def run_benchmark(world_size, rank, gpu_id, tokenizer, args): print('started', world_size, rank, gpu_id, args) torch.cuda.set_device(gpu_id) device = torch.device(gpu_id) all_input_ids, padded_gen_lens = args.all_input_ids, args.padded_gen_lens block_length=args.block_length # print() from vllm import distributed os.environ['MASTER_ADDR'] = 'localhost' os.environ['MASTER_PORT'] = str(args.master_port+args.port_offset) distributed.init_distributed_environment(world_size, rank, 'env://', rank, 'nccl') distributed.initialize_model_parallel(args.tp_size, backend='nccl') print("[Loading model]") # setup EP parallel_config = ParallelConfig(enable_expert_parallel = True) with set_current_vllm_config(VllmConfig(parallel_config = parallel_config)): vllm_config = get_current_vllm_config() print("EP Enabled:", vllm_config.parallel_config.enable_expert_parallel) model_config = AutoConfig.from_pretrained(args.model_name, trust_remote_code=True) if 'llada_moe' == args.model_type: model = LLaDAMoeModelLM(config=model_config).eval() model.load_weights(args.model_name, torch_dtype=torch.bfloat16) print('llada_moe') elif 'llada2' == args.model_type: model = LLaDA2MoeModelLM(config=model_config).eval() model.load_weights(args.model_name, torch_dtype=torch.bfloat16) elif 'llada' == args.model_type: model = LLaDAModelLM.from_pretrained(args.model_name, torch_dtype=torch.bfloat16, init_device=device).eval() else: raise ValueError('model type not supported') if args.tp_size>1 and args.use_tp: print('enabling tp') model.tensor_parallel(args.tp_size) x = torch.arange(50+args.gen_len, dtype=torch.long, device=device).unsqueeze(0) model = model.to(device) out = model(x, use_cache=False) out = model(x, use_cache=True) if args.use_compile: if args.use_cudagraph: model.forward = torch.compile(model.forward, mode='reduce-overhead', fullgraph=False, dynamic=True) else: model.forward = torch.compile(model.forward, fullgraph=False, dynamic=True) if args.parallel_decoding == 'threshold': if args.use_credit: decoder = CreditThresholdParallelDecoder(temperature=0, threshold=args.threshold, mask_id=args.mask_id, eos_id=args.eos_id) else: decoder = ThresholdParallelDecoder(temperature=0, threshold=args.threshold, mask_id=args.mask_id, eos_id=args.eos_id) else: decoder = HierarchyDecoder(temperature=0, threshold=args.threshold, low_threshold=args.low_threshold, mask_id=args.mask_id, eos_id=args.eos_id) use_sw = args.prefix_look > 0 or args.after_look > 0 or args.warmup_times > 0 if args.cache == 'prefix' or args.cache == 'dual': cache_factory=KVCacheFactory(args.cache, is_bd_model=args.use_bd) else: cache_factory=None if not args.use_bd: if args.cont_weight>0: if use_sw: print("IterSmoothWithVicinityCacheDiffusionLLM") dllm = IterSmoothWithVicinityCacheDiffusionLLM(model, decoder, BlockIteratorFactory(start_block_align=True), cache_factory=cache_factory, early_stop=True, cont_weight=args.cont_weight, prefix_look=args.prefix_look, after_look=args.after_look, warmup_steps=args.warmup_times) else: print("IterSmoothDiffusionLLM") dllm = IterSmoothDiffusionLLM(model, decoder, BlockIteratorFactory(start_block_align=True), cache_factory=cache_factory, early_stop=True, cont_weight=args.cont_weight) else: if use_sw: print("VicinityCacheDiffusionLLM") dllm = VicinityCacheDiffusionLLM(model, decoder, BlockIteratorFactory(start_block_align=True), cache_factory=cache_factory, early_stop=True,prefix_look=args.prefix_look, after_look=args.after_look, warmup_steps=args.warmup_times) else: print("BlockWiseDiffusionLLM") dllm = BlockWiseDiffusionLLM(model, decoder, BlockIteratorFactory(start_block_align=True), cache_factory=cache_factory, early_stop=True, use_shift=args.use_shift) else: print("BlockDiffusionLLM") dllm = BlockDiffusionLLM(model, decoder, BlockIteratorFactory(start_block_align=True), cache_factory=cache_factory, early_stop=True) warmup_cudagraph(rank, device, dllm, args.gen_len, block_length, args.batch_size, args.vocab_size) for wi in range(1): outputs = [] total_forward = 0 if rank==0: iterator = trange(0, len(all_input_ids), args.batch_size) else: iterator = range(0, len(all_input_ids), args.batch_size) start = time.time() tpfs = [] tpss = [] fpss = [] total_token = 0 token_numbers = [] for i in iterator: input_ids = all_input_ids[i:i+args.batch_size] max_length = 0 min_padded_length = 10000 for j, seq in enumerate(input_ids): if seq.shape[1] > max_length: max_length = seq.shape[1] min_padded_length = padded_gen_lens[i+j] batch_input_ids= torch.zeros((len(input_ids), max_length), dtype=torch.long, device=device).fill_(args.mask_id) for j in range(len(input_ids)): batch_input_ids[j, :input_ids[j].shape[1]] = input_ids[j].to(device) input_ids = batch_input_ids padded_gen_len = padded_gen_lens[i] inner_start = time.time() prev_forwards = dllm.num_forwards out = dllm.generate(input_ids, gen_length=min_padded_length, block_length=block_length) nfe = dllm.num_forwards - prev_forwards inner_stop = time.time() sample_time = inner_stop - inner_start for j in range(input_ids.shape[0]): outputs.append(out[j].unsqueeze(0)) total_forward += nfe batch_token_number = 0 for j in range(input_ids.shape[0]): token_number = int((out[j]!=156892).sum() - all_input_ids[i+j].shape[1]) batch_token_number += token_number token_numbers.append(token_number) tpf = batch_token_number/nfe/args.batch_size tps = batch_token_number/sample_time fps = nfe/sample_time if rank == 0: print(f'[iter {i:4d}]nfe={nfe:4d}, token number={batch_token_number:4d}, fps={fps:4.2f},tpf={tpf:2.2f}, tps={tps:4.2f}') if wi==0 and i<5: for j in range(input_ids.shape[0]): answer = cut_eos(out[j, all_input_ids[i+j].shape[1]:].unsqueeze(0))[0] # print(answer) print(f'generated text {j}: {tokenizer.decode(answer, skip_special_tokens=False)}') tpfs.append(tpf) tpss.append(tps) fpss.append(fps) total_token += token_number total_token = total_token stop = time.time() answers = [] if rank==0: for i in trange(len(outputs)): out = outputs[i] answer = (tokenizer.decode(out[0, all_input_ids[i].shape[1]:], skip_special_tokens=True)) answers.append(answer) print(f'Forward: {total_forward}, Time: {stop-start}, FPS: {total_forward/(stop-start)}({np.mean(fpss)}), TPS: {total_token/(stop-start)}({np.mean(tpss)}), TPF: {total_token/total_forward}({np.mean(tpfs)})') filename = args.save_path with open (filename, 'w') as f: for i in range(len(answers)): answer = answers[i] json.dump({'answer': answer, 'generated_length': token_numbers[i], 'tpf':tpfs[i//args.batch_size], 'tps':tpss[i//args.batch_size], 'fps':fpss[i//args.batch_size], }, f) f.write('\n') print(f'Forward: {total_forward}, Time: {stop-start}, FPS: {total_forward/(stop-start)}({np.mean(fpss)}), TPS: {total_token/(stop-start)}({np.mean(tpss)}), TPF: {total_token/total_forward}({np.mean(tpfs)})') with open(args.speed_path, 'a+') as f: print( args.config, args.parallel_decoding, args.threshold, args.prefix_look, args.batch_size, args.block_length, total_forward, stop-start, total_token / len(all_input_ids), total_forward/(stop-start), total_token/(stop-start), total_token/total_forward, sum(padded_gen_lens)/total_forward, np.mean(fpss), np.mean(tpss), np.mean(tpfs), file=f) return @dataclass class EvalConfig: model_name: str = '' gpu: str = '0,1,2,3' batch_size: int = 1 gen_len: int = 1024 prefix_look: int = 0 after_look: int = 0 block_length: int = 64 threshold: float = 0.9 warmup_times: int = 0 low_threshold: float = 0.3 cont_weight: float = 0 parallel_decoding: str = 'threshold' use_credit: bool = False cache: str = '' use_tp: bool = False save_path: str = '' config: int = 0 tp_size: int = 1 port_offset: int = 0 all_input_ids = None padded_gen_lens = None use_cudagraph: bool = False use_compile: bool = True use_bd: bool = False use_shift: bool = False model_type: str = 'llada' vocab_size: int = 156896 master_port: int = 23456 mask_id: int = 156895 eos_id: int = 156892 save_dir: str = './res' save_samples: bool = False speed_path: str = '' def set_seed(seed): torch.manual_seed(seed) random.seed(seed) np.random.seed(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False @register_model("dInfer_eval") class DInferEvalHarness(LM): def __init__( self, model_path='', device="cuda", mask_id=126336, eos_id=126081, max_length=4096, batch_size=2, mc_num=128, is_check_greedy=True, gen_length=1024, block_length=1024, save_dir=None, show_speed=False, parallel_decoding="threshold", threshold: float=0.9, cache: str="", warmup_times: int=0, low_threshold: float=0.3, cont_weight: float=0, use_credit: bool=False, tp_size: int=1, parallel = 'dp', use_compile = True, master_port = 23456, use_cudagraph = True, gpus = '0,1,2,3', use_bd = False, prefix_look = 0, after_look = 0, use_shift = False, model_type = 'llada', save_samples = False, **kwargs ): super().__init__() self.model_path = model_path self.mask_id = mask_id self.eos_id = eos_id self.mc_num = mc_num self.batch_size = int(batch_size) assert mc_num % self.batch_size == 0 self.sampling_eps = 0. self.max_length = max_length self.is_check_greedy = is_check_greedy self.gen_length = gen_length self.block_length = block_length self.save_dir = save_dir self.show_speed = show_speed self.parallel_decoding = parallel_decoding self.threshold = threshold self.cache = cache self.warmup_times = warmup_times self.low_threshold = low_threshold self.cont_weight = cont_weight self.use_credit = use_credit self.master_port = master_port self.tp_size = tp_size self.use_compile = use_compile self.parallel = parallel self.use_cudagraph = use_cudagraph self.gpus = gpus self.prefix_look = prefix_look self.after_look = after_look self.use_bd = use_bd self.kwargs = kwargs self.use_shift = use_shift self.model_type = model_type self.save_samples = save_samples if self.model_type == 'llada_moe': self.mask_id = 156895 self.eos_id = 156892 self.vocab_size = 156896 self.is_moe = True elif self.model_type == 'llada2': self.mask_id = 156895 self.eos_id = 156892 self.vocab_size = 156896 self.is_moe = True elif self.model_type == 'llada': self.vocab_size = 126464 self.is_moe = False else: raise ValueError('model type not supported') accelerator = accelerate.Accelerator() if accelerator.num_processes > 1: self.accelerator = accelerate.Accelerator() self._rank = self.accelerator.local_process_index self._world_size = self.accelerator.num_processes else: self.accelerator = None model_kwargs = {} if self.accelerator is not None: model_kwargs.update({'device_map': {'': f'{self.accelerator.device}'}}) # set decoder if parallel_decoding == "threshold": if use_credit: decoder = CreditThresholdParallelDecoder(temperature=0, threshold=threshold, mask_id=self.mask_id, eos_id=self.eos_id) else: decoder = ThresholdParallelDecoder(temperature=0, threshold=threshold, mask_id=self.mask_id, eos_id=self.eos_id) else: decoder = HierarchyDecoder(temperature=0, threshold=threshold, low_threshold=low_threshold, mask_id=self.mask_id, eos_id=self.eos_id) if parallel == 'dp': self.device= torch.device(device) self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) os.environ['MASTER_ADDR'] = 'localhost' os.environ['MASTER_PORT'] = str(self.master_port + self.rank) distributed.init_distributed_environment(1, 0, 'env://', 0, 'nccl') distributed.initialize_model_parallel(1, backend='nccl') parallel_config = ParallelConfig(enable_expert_parallel = True) with set_current_vllm_config(VllmConfig(parallel_config = parallel_config)): vllm_config = get_current_vllm_config() print("EP Enabled:", vllm_config.parallel_config.enable_expert_parallel) config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) # load model if self.model_type == 'llada_moe': self.model = LLaDAMoeModelLM(config=config).eval() self.model.load_weights(self.model_path, torch_dtype=torch.bfloat16) elif self.model_type == 'llada2': self.model = LLaDA2MoeModelLM(config=config).eval() self.model.load_weights(self.model_path, torch_dtype=torch.bfloat16) elif self.model_type == 'llada': self.model = LLaDAModelLM.from_pretrained(model_path, torch_dtype=torch.bfloat16, init_device=str(device)).eval() else: raise ValueError('model type not supported') self.vllm_config = vllm_config if self.accelerator is not None: self.model = self.accelerator.prepare(self.model) self.device = torch.device(f'{self.accelerator.device}') else: self.model = self.model.to(self.device) if self.use_compile: # compile model if self.use_cudagraph: self.model.forward = torch.compile(self.model.forward, fullgraph=False, dynamic=True, mode='reduce-overhead') else: self.model.forward = torch.compile(self.model.forward, fullgraph=False, dynamic=True) if self.cache == 'prefix' or self.cache == 'dual': cache_factory=KVCacheFactory(self.cache, is_bd_model=self.use_bd) else: cache_factory=None use_sw = self.cache != '' and (self.prefix_look > 0 or self.after_look > 0 or self.warmup_times > 0 ) if not self.use_bd: if self.cont_weight>0: if use_sw: self.dllm = IterSmoothWithVicinityCacheDiffusionLLM(self.model, decoder, BlockIteratorFactory(start_block_align=True), cache_factory=cache_factory, early_stop=True, cont_weight=self.cont_weight, prefix_look=self.prefix_look, after_look=self.after_look, warmup_steps=self.warmup_times) else: self.dllm = IterSmoothDiffusionLLM(self.model, decoder, BlockIteratorFactory(start_block_align=True), cache_factory=cache_factory, early_stop=True, cont_weight=self.cont_weight) else: if use_sw: self.dllm = VicinityCacheDiffusionLLM(self.model, decoder, BlockIteratorFactory(start_block_align=True), cache_factory=cache_factory, early_stop=True, prefix_look=self.prefix_look, after_look=self.after_look, warmup_steps=self.warmup_times) else: self.dllm = BlockWiseDiffusionLLM(self.model, decoder, BlockIteratorFactory(start_block_align=True), cache_factory=cache_factory, early_stop=True, use_shift=self.use_shift) else: self.dllm = BlockDiffusionLLM(self.model, decoder, BlockIteratorFactory(start_block_align=True), cache_factory=cache_factory, early_stop=True) elif parallel == 'tp': self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) else: raise NotImplementedError(parallel) @property def rank(self): return self._rank @property def world_size(self): return self._world_size @property def tokenizer_name(self) -> str: return self.model_path def apply_chat_template(self, chat_history, **kwargs) -> str: if "tokenize" not in kwargs: kwargs["tokenize"] = False return self.tokenizer.apply_chat_template(chat_history, **kwargs) def _forward_process(self, batch, prompt_index): b, l = batch.shape target_len = (l - prompt_index.sum()).item() k = torch.randint(1, target_len + 1, (), device=batch.device) x = torch.round(torch.linspace(float(k), k + (b - 1) * (target_len / b), steps=b, device=batch.device)).long() x = ((x - 1) % target_len) + 1 assert x.min() >= 1 and x.max() <= target_len indices = torch.arange(target_len, device=batch.device).repeat(b, 1) is_mask = indices < x.unsqueeze(1) for i in range(b): is_mask[i] = is_mask[i][torch.randperm(target_len)] is_mask = torch.cat((torch.zeros(b, prompt_index.sum(), dtype=torch.bool, device=batch.device), is_mask), dim=1) noisy_batch = torch.where(is_mask, self.mask_id, batch) return noisy_batch, (x / target_len).unsqueeze(1).repeat(1, l) @torch.no_grad() def get_logits(self, batch, prompt_index): if self.cfg > 0.: assert len(prompt_index) == batch.shape[1] prompt_index = prompt_index.unsqueeze(0).repeat(batch.shape[0], 1) un_batch = batch.clone() un_batch[prompt_index] = self.mask_id batch = torch.cat([batch, un_batch]) logits = self.model(batch).logits if self.cfg > 0.: logits, un_logits = torch.chunk(logits, 2, dim=0) logits = un_logits + (self.cfg + 1) * (logits - un_logits) return logits[:, :batch.shape[1]] @torch.no_grad() def get_loglikelihood(self, prefix, target): seq = torch.concatenate([prefix, target])[None, :] seq = seq.repeat((self.batch_size, 1)).to(self.device) prompt_index = torch.arange(seq.shape[1], device=self.device) < len(prefix) loss_acc = [] for _ in range(self.mc_num // self.batch_size): perturbed_seq, p_mask = self._forward_process(seq, prompt_index) mask_indices = perturbed_seq == self.mask_id logits = self.get_logits(perturbed_seq, prompt_index) loss = F.cross_entropy(logits[mask_indices], seq[mask_indices], reduction='none') / p_mask[mask_indices] loss = loss.sum() / self.batch_size loss_acc.append(loss.item()) return - sum(loss_acc) / len(loss_acc) @torch.no_grad() def suffix_greedy_prediction(self, prefix, target): if not self.is_check_greedy: return False seq = torch.full((1, len(prefix) + len(target)), self.mask_id, device=self.device) prompt_index = torch.arange(seq.shape[1], device=self.device) < len(prefix) prefix, target = prefix.to(self.device), target.to(self.device) seq[0, :len(prefix)] = prefix for i in range(len(target)): mask_index = (seq == self.mask_id) logits = self.get_logits(seq, prompt_index)[mask_index] x0 = torch.argmax(logits, dim=-1) p = torch.softmax(logits.to(torch.float32), dim=-1) confidence = torch.gather(p, dim=-1, index=torch.unsqueeze(x0, -1)).squeeze(dim=-1) _, index = torch.sort(confidence, descending=True) x0[index[1:]] = self.mask_id seq[mask_index] = x0.clone() correct = target == seq[0, len(prefix):] correct = torch.all(correct) return correct def _encode_pair(self, context, continuation): n_spaces = len(context) - len(context.rstrip()) if n_spaces > 0: continuation = context[-n_spaces:] + continuation context = context[:-n_spaces] whole_enc = self.tokenizer(context + continuation)["input_ids"] context_enc = self.tokenizer(context)["input_ids"] context_enc_len = len(context_enc) continuation_enc = whole_enc[context_enc_len:] return context_enc, continuation_enc def loglikelihood(self, requests): def _tokenize(e): prefix, target = self._encode_pair(e["prefix"], e["target"]) return { "prefix_text": e["prefix"], "target_text": e["target"], "prefix": prefix, "target": target, } ds = [] ds = [{"prefix": req.args[0], "target": req.args[1]} for req in requests] ds = Dataset.from_list(ds) ds = ds.map(_tokenize) ds = ds.with_format("torch") prompt_len = [len(x["prefix"]) + len(x["target"]) for x in ds] assert max(prompt_len) <= 4096 out = [] with torch.no_grad(): for elem in tqdm(ds, desc="Computing likelihood..."): prefix = elem["prefix"] target = elem["target"] ll = self.get_loglikelihood(prefix, target) is_target_greedy_dec = self.suffix_greedy_prediction(prefix, target) out.append((ll, 1.0 if is_target_greedy_dec else 0.0)) torch.cuda.empty_cache() return out def loglikelihood_rolling(self, requests): raise NotImplementedError def generate_until(self, requests): if self.save_dir is not None: os.makedirs(self.save_dir, exist_ok=True) self.save_path = os.path.join(self.save_dir, f'rank_{self.rank}.jsonl') print(f"save_path: {self.save_path}") self.speed_path = os.path.join(self.save_dir, f'results.txt') def get_bucket_length(length): bucket_length = bucket_size*(length//bucket_size) if bucket_length not in used_buckets: used_buckets.append(bucket_length) return bucket_length def load_inputs(prompts, tokenizer): all_input_ids = [] for id, prompt in enumerate(prompts): input_ids = tokenizer(prompt.args[0])['input_ids'] input_ids = torch.tensor(input_ids).unsqueeze(0) all_input_ids.append(input_ids) return all_input_ids def cal_bucket_len(gen_len, all_input_ids): max_prompt_length = 0 padded_gen_lens = [] for i in range(len(all_input_ids)): input_ids = all_input_ids[i] if input_ids.shape[1] > max_prompt_length: max_prompt_length = input_ids.shape[1] padded_length = get_bucket_length(input_ids.shape[1]+gen_len) padded_gen_lens.append(padded_length - input_ids.shape[1]) return padded_gen_lens all_input_ids = load_inputs(requests, self.tokenizer) padded_gen_lens = cal_bucket_len(self.gen_length, all_input_ids) answers = [] outputs = [] total_forward = 0 start = time.time() tpfs = [] tpss = [] fpss = [] total_token = 0 token_numbers = [] if self.parallel == 'dp': with set_current_vllm_config(self.vllm_config): if self.use_cudagraph and self.use_cudagraph: warmup_cudagraph(self.rank, self.device, self.dllm, self.gen_length, self.block_length, self.batch_size, self.vocab_size) for i, req in enumerate(tqdm(requests, desc="Generating...")): input_ids = all_input_ids[i] padded_gen_len = padded_gen_lens[i] inner_start = time.time() input_ids = input_ids.to(self.device) prev_forwards = self.dllm.num_forwards out = self.dllm.generate(input_ids, gen_length=padded_gen_len,block_length=self.block_length) nfe = self.dllm.num_forwards - prev_forwards inner_stop = time.time() sample_time = inner_stop - inner_start outputs.append(out) answer = (self.tokenizer.decode(out[0, all_input_ids[i].shape[1]:], skip_special_tokens=True)) answers.append(answer) total_forward += nfe token_number = out.shape[1] - input_ids.shape[1] token_numbers.append(token_number) tpf = token_number/nfe tps = token_number/sample_time fps = nfe/sample_time if self.rank == 0: print(f'iter={i}, fps={fps}, nfe={nfe}') tpfs.append(tpf) tpss.append(tps) fpss.append(fps) total_token += token_number total_token = total_token stop = time.time() print(f'Forward: {total_forward}, Time: {stop-start}, FPS: {total_forward/(stop-start)}({np.mean(fpss)}), TPS: {total_token/(stop-start)}({np.mean(tpss)}), TPF: {total_token/total_forward}({np.mean(tpfs)})') if self.show_speed and self.save_dir is not None: with open (self.save_dir+f'/rank{self.rank}_results.jsonl', 'w', encoding='utf-8') as file: data={'rank':f'rank{self.rank}', 'forward per second': np.mean(fpss), 'tokens per second': np.mean(tpss), 'tokens per forward': np.mean(tpfs), 'average generated length': total_token / len(all_input_ids) } file.write(json.dumps(data, ensure_ascii=False) + '\n') elif self.parallel == 'tp': procs = [] answers = [] gpus = [int(gpu) for gpu in self.gpus.split(';')] args = {"gpu": gpus, "batch_size": self.batch_size, "model_name": self.model_path, "gen_len": self.gen_length, "block_length": self.block_length, "prefix_look": self.prefix_look, "after_look": self.after_look, "warmup_times": self.warmup_times, "low_threshold": self.low_threshold, "threshold": self.threshold, "cont_weight": self.cont_weight, "use_credit": self.use_credit, "cache": self.cache, "parallel_decoding": self.parallel_decoding, "tp_size": self.tp_size, "save_path": self.save_path, "use_cudagraph": self.use_cudagraph, "use_compile": self.use_compile,"use_bd": self.use_bd, "use_shift": self.use_shift, "model_type": self.model_type, "vocab_size": self.vocab_size, "mask_id": self.mask_id, "eos_id": self.eos_id, "save_dir": self.save_dir, "save_samples": self.save_samples, "speed_path": self.speed_path} args = EvalConfig(**args) args.tp_size = len(gpus) args.master_port = self.master_port args.use_tp = args.tp_size > 1 args.port_offset = gpus[0] args.all_input_ids = all_input_ids args.padded_gen_lens = padded_gen_lens if len(gpus) == 1: run_benchmark(1, 0, gpus[0], self.tokenizer, args) else: for i, gpu in enumerate(gpus): ctx = mp.get_context('spawn') p = ctx.Process(target=run_benchmark, args=(len(gpus), i, gpu, self.tokenizer, args)) p.daemon = True procs.append(p) p.start() for p in procs: p.join() answers = [] with open(self.save_path, 'r') as f: for line in f : answers.append(json.loads(line)["answer"]) if not self.save_samples: os.remove(self.save_path) return answers if __name__ == "__main__": set_seed(1234) cli_evaluate()