from collections import OrderedDict from quantize.int_linear import QuantLinear import torch from quantize.int_matmul import QuantMatMul from models.transformation import * from pprint import pprint from parallel_utils import map_layers_to_multi_gpus import os from tqdm import tqdm import torch.nn as nn import numpy as np from datautils import get_loaders from accelerate.utils import get_balanced_memory from accelerate import dispatch_model,infer_auto_device_map import copy def get_max_memory_map(ratio=0.95): max_memory = {} for i in range(torch.cuda.device_count()): total_mem = torch.cuda.get_device_properties(i).total_memory mem_in_gib = int(total_mem * ratio / (1024 ** 3)) # 转换为 GiB max_memory[i] = f"{mem_in_gib}GiB" return max_memory @torch.no_grad() def evaluate(lm, args, logger): results = {} assert not (args.multigpu and args.parallelize) if args.multigpu: if "llama" in args.net.lower() or "vicuna" in args.net.lower() or "qwen" in args.net.lower(): map_layers_to_multi_gpus(lm.model.model.layers) input_device = lm.model.model.layers[0].device output_device = lm.model.model.layers[-1].device assert input_device == output_device lm._device = input_device lm.model.model.embed_tokens.to(input_device) lm.model.model.norm.to(output_device) lm.model.lm_head.to(output_device) else: raise NotImplementedError("Only support for llama/vicuna/qwen multigpu now") elif args.parallelize: balanced_mem = get_balanced_memory( lm.model, max_memory=get_max_memory_map(0.95), no_split_module_classes=["LlamaDecoderLayer","QuantLlamaDecoderLayer","Qwen3MoeDecoderLayer"] ) logger.info(f"mem is {balanced_mem}") device_map = infer_auto_device_map( lm.model, max_memory=balanced_mem, no_split_module_classes=["LlamaDecoderLayer","QuantLlamaDecoderLayer","Qwen3MoeDecoderLayer"] ) lm.model = dispatch_model(lm.model,device_map=device_map) else: if "llama" in args.net.lower() or "vicuna" in args.net.lower() or "qwen" in args.net.lower(): lm.model = lm.model.to(lm.device) else: raise NotImplementedError("Only support for llama/vicuna/qwen single gpu now") if args.eval_ppl: logger.info(f"model seqlen is {lm.seqlen}") datasets = args.test_datasets.split(",") for dataset in datasets: # for dataset in ["wikitext2"]: cache_testloader = f'{args.cache_dir}/testloader_{args.net}_{dataset}_all.cache' if os.path.exists(cache_testloader): testloader = torch.load(cache_testloader) logger.info(f"load calibration from {cache_testloader}") else: dataloader, testloader = get_loaders( dataset, seed=args.seed, model=args.model, seqlen=lm.seqlen, args=args ) torch.save(testloader, cache_testloader) if "c4" in dataset: testenc = testloader else: testenc = testloader.input_ids nsamples = testenc.numel() // lm.seqlen use_cache = lm.model.config.use_cache lm.model.config.use_cache = False lm.model.eval() nlls = [] all_hidden_states = [] output_hidden_states = False for i in tqdm(range(nsamples)): batch = testenc[:, (i * lm.seqlen) : ((i + 1) * lm.seqlen)].to(lm.device) outputs = lm.model.model(batch,output_hidden_states=output_hidden_states) hidden_states = outputs[0] if hasattr(lm.model.lm_head,"bias") and lm.model.lm_head.bias is not None: lm.model.lm_head.bias = torch.nn.Parameter(lm.model.lm_head.bias.to(lm.model.lm_head.weight.device)) logits = lm.model.lm_head(hidden_states) shift_logits = logits[:, :-1, :] shift_labels = testenc[:, (i * lm.seqlen) : ((i + 1) * lm.seqlen)][ :, 1: ].to(lm.model.lm_head.weight.device) loss_fct = nn.CrossEntropyLoss() loss = loss_fct( shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ) neg_log_likelihood = loss.float() * lm.seqlen nlls.append(neg_log_likelihood) if i == args.limit: break ppl = torch.exp(torch.stack(nlls).sum() / (nsamples * lm.seqlen)) logger.info(f'{dataset} : {ppl.item()}') lm.model.config.use_cache = use_cache results[dataset] = ppl.item() if args.tasks != "": args.tasks = args.tasks.split(",") import lm_eval from lm_eval import utils as lm_eval_utils from lm_eval.api.registry import ALL_TASKS from lm_eval.models.huggingface import HFLM print(f"use lm_eval in {lm_eval}") task_manager = lm_eval.tasks.TaskManager(include_path="./datasets_local/lm_eval_configs/tasks", include_defaults=True) hflm = HFLM(pretrained=lm.model,tokenizer=lm.tokenizer, batch_size=args.lm_eval_batch_size) t_results = lm_eval.simple_evaluate(hflm, tasks=args.tasks, batch_size=args.lm_eval_batch_size,task_manager=task_manager)['results'] metric_vals = {task: round(result.get('acc_norm,none', result['acc,none']), 4) for task, result in t_results.items()} logger.info(metric_vals) pprint(metric_vals) if args.eval_ppl is True: metric_vals.update(results) reported_metric_vals = {} for k,v in metric_vals.items(): if "mmlu" in k: if k == "mmlu": reported_metric_vals[k] = v else: reported_metric_vals[k] = v import pandas as pd if os.path.exists(f"{args.output_dir}/results.csv"): df = pd.read_csv(f"{args.output_dir}/results.csv") new_df = pd.DataFrame(reported_metric_vals,index=[0]) df[new_df.columns] = new_df else: df = pd.DataFrame(reported_metric_vals,index=[0]) if args.eval_ppl: new_columns = ['wikitext2','c4'] else: new_columns = [] new_columns += args.tasks if len(args.tasks) >= 5: df["avg-5"] = df[['piqa','arc_easy', 'arc_challenge','hellaswag', 'winogrande']].mean(axis=1) if len(args.tasks) >= 6: df["avg-6"] = df[['piqa','arc_easy', 'arc_challenge','hellaswag', 'winogrande', 'boolq']].mean(axis=1) if 'mmlu' in args.tasks and len(args.tasks) >= 7: df["avg-7"] = df[['piqa','arc_easy', 'arc_challenge','hellaswag', 'winogrande', 'boolq','mmlu']].mean(axis=1) logger.info(df) df.to_csv(f"{args.output_dir}/results.csv",index=False) model = lm.model if "llama" in args.net.lower() or "vicuna" in args.net.lower() or "qwen" in args.net.lower(): model.model.embed_tokens = model.model.embed_tokens.cpu() model.model.norm = model.model.norm.cpu() else: raise ValueError("Only support for opt/llama/Llama-2/falcon/mixtral now") return results def get_slider_parameters(sub_layers, use_list=["scale","alpha","shift"]): params = [] print(f"get {use_list} parameters!") for sub_layer_idx in range(len(sub_layers)): for n, m in sub_layers[sub_layer_idx].named_parameters(): if any(n.find(t) > -1 for t in use_list) and not n.find('bound_factor') > -1: params.append(m) # print(n) # print(params) return iter(params) def get_lwc_parameters(sub_layers): params = [] print("get lwc parameters!") for sub_layer_idx in range(len(sub_layers)): for n, m in sub_layers[sub_layer_idx].named_parameters(): if n.find('bound_factor') > -1: params.append(m) return iter(params) def try_delete_object(object,logger,name=None): try: del object except Exception as e: logger.info(f"del tensor occurs {e}, skip!") else: logger.info(f"del tensor {name} successfully!") def cleanup_memory(verbos=True,logger=None) -> None: """Run GC and clear GPU memory.""" import gc import inspect caller_name = '' try: caller_name = f' (from {inspect.stack()[1].function})' except (ValueError, KeyError): pass def total_reserved_mem() -> int: return sum(torch.cuda.memory_reserved(device=i) for i in range(torch.cuda.device_count())) memory_before = total_reserved_mem() # gc.collect and empty cache are necessary to clean up GPU memory if the model was distributed gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() memory_after = total_reserved_mem() if verbos and logger: logger.info( f"GPU memory{caller_name}: {memory_before / (1024 ** 3):.2f} -> {memory_after / (1024 ** 3):.2f} GB" f" ({(memory_after - memory_before) / (1024 ** 3):.2f} GB)" ) def slider_state_dict(model, destination=None, prefix='', keep_vars=False): if destination is None: destination = OrderedDict() for name, param in model.named_parameters(): if name.find('smooth') > -1 or name.find('bound_factor') > -1 or name.find('lora_') > -1 or name.find('Q_') > -1 : destination[prefix + name] = param if keep_vars else param.detach() return destination def register_scales_and_zeros(model): for name, module in model.named_modules(): if isinstance(module, QuantLinear): module.weight_quantizer.register_scales_and_zeros()