import torch import torch.nn as nn from models.int_llama_layer import QuantLlamaDecoderLayer import copy import math import os from tqdm import tqdm from train_utils import to_float,to_half from quantize.utils import get_slider_parameters, get_lwc_parameters, slider_state_dict from train_utils import to_dev,obtain_teacher_output,obtain_studnet_output,replace_ori_layer,init_model,model_to_inference_mode,SubLayer import time from transformers import get_scheduler from quantize.utils import cleanup_memory from torch.utils.data import DataLoader,Dataset import numpy as np import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DistributedSampler from tqdm import tqdm from quantize.utils import evaluate def setup_ddp(): dist.init_process_group( backend="nccl", init_method="env://", # 使用 torchrun 自动设置的环境变量 ) def cleanup_ddp(): dist.destroy_process_group() class Quant_dataset(Dataset): def __init__(self,aug_quant_inps=None,aug_fp_inps=None, aug_quant_targets=None,aug_fp_targets=None,samples_num=512,windows_num=1,args=None): """ In i-th round quant_inps: the output from (i-1)th quant model using quant_inps. aug_fp_inps: the output from (i-1)th fp16 model using fp_inps. fp_target: the output from i-th fp16 model using fp_inps. aug_quant_target: the output from i-th fp16 model using quant_inps. fp_inps ---------> [ fp16 model] ------------> fp_target quant_inps ---------> [ fp16 model] ------------> quant_target quant_inps -----> [quant model] -------------> out1 <-> [fp_target,quant_target] fp_inps -----> [quant model] -------------> out2 <->fp_target """ # self.quant_inps = quant_inps self.samples_num = samples_num self.windows_num = windows_num assert self.windows_num == aug_quant_inps.shape[1] assert self.samples_num == len(aug_quant_inps) self.aug_quant_inps = aug_quant_inps self.aug_fp_targets = aug_fp_targets if aug_fp_inps is not None: self.aug_fp_inps = aug_fp_inps else: self.aug_fp_inps = torch.ones(self.samples_num,self.windows_num) if aug_quant_targets is not None: self.aug_quant_targets = aug_quant_targets else: self.aug_quant_targets = torch.ones(self.samples_num,self.windows_num) def __len__(self): return self.samples_num def __getitem__(self, idx): return self.aug_quant_inps[idx],self.aug_fp_inps[idx],self.aug_quant_targets[idx],self.aug_fp_targets[idx] MB = 1024.0 * 1024.0 def train_one_round(r,epochs,sub_layers,layer_id_list,qdataset,cur_epochs,optimizer,lr_scheduler,attention_mask_batch,position_ids,position_embeddings,devs,args,logger,max_train_steps,init_quant_rate,fp16_type,global_start_time,total_epochs,loss_func,acts_round_idx): if args.use_ddp is True: rank = dist.get_rank() sub_layers = DDP(sub_layers,device_ids=[rank]) sampler = DistributedSampler(qdataset,shuffle=True) shuffle = None else: rank = 0 sampler = None shuffle = True qdataloader = DataLoader(qdataset,batch_size=args.batch_size,shuffle=shuffle,num_workers=0,pin_memory=True,sampler=sampler) if args.loss_type == "mean": if args.use_base_loss == "none": base_loss_num = 0 elif args.use_base_loss == "last": base_loss_num = 1 elif args.use_base_loss == "all": base_loss_num = args.last_round_inp_num else: raise NotImplementedError() Accumulated_loss_num = (int(args.use_fp_inp_loss) + int(args.use_quant_tar_loss) ) * args.last_round_inp_num + base_loss_num elif args.loss_type == "add": Accumulated_loss_num = 1 else: raise NotImplementedError("noly support mean and add!") logger.info(f"Accumulated_loss_num is {Accumulated_loss_num}") sub_layers.module.module.train() # 确保开启训练模型 for e in range(epochs): if args.use_ddp is True: sampler.set_epoch(e) start_time = time.time() epoch_losses = [] epoch_norms = [] epoch_losses_fp = [] epoch_losses_quant = [] epoch_losses_base = [] # import ipdb;ipdb.set_trace() for quant_input_list,fp_input_list,quant_tar_list,fp_tar_list in qdataloader: batch_loss = [0.0 for _ in range(quant_input_list.shape[1])] batch_base_loss = [0.0 for _ in range(quant_input_list.shape[1])] batch_fp_loss = [0.0 for _ in range(quant_input_list.shape[1])] batch_quant_loss = [0.0 for _ in range(quant_input_list.shape[1])] optimizer.zero_grad() # optimizer.zero_grad(set_to_none=True) # MOE节省显存 for w_idx in range(quant_input_list.shape[1]): quant_input,fp_input,quant_tar,fp_tar = quant_input_list[:,w_idx],fp_input_list[:,w_idx],quant_tar_list[:,w_idx],fp_tar_list[:,w_idx] inp_list = [quant_input] if args.use_fp_inp_loss is False else [quant_input,fp_input] for inp_idx,inp in enumerate(inp_list): if inp_idx == 0 and (args.use_base_loss == "none" or args.use_base_loss == "last" and w_idx != quant_input_list.shape[1] -1) and args.use_quant_tar_loss is False: continue loss_list = [] train_context = torch.cuda.amp.autocast(dtype=torch.bfloat16) with train_context: if args.use_ddp: out = sub_layers(inp) else: out = obtain_studnet_output( sub_layers,[args.quant_mode_layer_list[i] for i in layer_id_list], inp, attention_mask_batch[:len(inp)], position_ids,position_embeddings, args, devs=devs,return_gpu=True ) if inp_idx == 0: # out is quant_out if args.use_base_loss == "all" or (w_idx == quant_input_list.shape[1] -1 and args.use_base_loss == "last"): loss_base = loss_func(out, fp_tar.to(out.device)) loss_list.append(loss_base) batch_base_loss[w_idx] += loss_base.item() else: loss_base = torch.tensor(0.0) if args.use_quant_tar_loss is True: loss_quant = loss_func(out, quant_tar.to(out.device)) loss_list.append(loss_quant) batch_quant_loss[w_idx] += loss_quant.item() else: loss_quant = torch.tensor(0.0) else: # out is fp_out if args.use_fp_inp_loss is True: loss_fp = loss_func(out, fp_tar.to(out.device)) loss_list.append(loss_fp) batch_fp_loss[w_idx] += loss_fp.item() else: loss_fp = torch.tensor(0.0) loss = sum(loss_list) / Accumulated_loss_num loss.backward() batch_loss[w_idx] += loss.detach().item() if args.debug is True and not math.isfinite(loss.detach().item()): logger.info("Loss is NAN, stopping training") import ipdb;ipdb.set_trace() assert math.isfinite(loss.item()),"Loss is NAN, stopping training!" if args.grad_clip is not None: total_norm = torch.nn.utils.clip_grad_norm_(sub_layers.module.module.parameters(), max_norm=args.grad_clip) # logger.info(f"Gradient norm: {total_norm:.4f} Max norm: {args.grad_clip}") optimizer.step() epoch_norms.append(0.0) epoch_losses.append(batch_loss) epoch_losses_base.append(batch_base_loss) epoch_losses_fp.append(batch_fp_loss) epoch_losses_quant.append(batch_quant_loss) if args.use_lr_scheduler is True: lr_scheduler.step() current_memory = torch.cuda.memory_allocated() / MB max_memory = torch.cuda.max_memory_allocated() / MB cur_epochs += 1 epoch_mean_loss = torch.tensor(epoch_losses).mean(dim=0) # epoch_mean_norms = sum(epoch_norms) / len(epoch_norms) epoch_mean_loss_base = torch.tensor(epoch_losses_base).mean(dim=0) epoch_mean_loss_quant = torch.tensor(epoch_losses_quant).mean(dim=0) epoch_mean_loss_fp = torch.tensor(epoch_losses_fp).mean(dim=0) loss_str = "" for r_idx in range(quant_input_list.shape[1]): loss_str += f" loss r{acts_round_idx[r_idx]}:{epoch_mean_loss[r_idx]} " for r_idx in range(quant_input_list.shape[1]): loss_str += f" loss base r{acts_round_idx[r_idx]}:{epoch_mean_loss_base[r_idx]} " for r_idx in range(quant_input_list.shape[1]): loss_str += f" loss quant r{acts_round_idx[r_idx]}:{epoch_mean_loss_quant[r_idx]} " for r_idx in range(quant_input_list.shape[1]): loss_str += f" loss fp r{acts_round_idx[r_idx]}:{epoch_mean_loss_fp[r_idx]} " if rank == 0: logger.info( f"Round {r} epoch {e} {loss_str} max memory_allocated: {max_memory}MB current memory_allocated: {current_memory}MB epoch_time: {time.time() - start_time:.2f}s use_time:{(time.time() - global_start_time)/3600:.2f}h ETA: {(time.time() - global_start_time) / cur_epochs * (total_epochs-cur_epochs) / 3600:.2f}h" ) return sub_layers,cur_epochs def sliderquant( lm, args, dataloader, logger=None, teach_lm=None, ): logger.info("Starting ...") model = lm.model dev = lm.device use_cache = model.config.use_cache model.config.use_cache = False is_llama = False if args.use_ddp: rank = dist.get_rank() torch.cuda.set_device(rank) else: rank = 0 if "llama" in args.net.lower() or "vicuna" in args.net.lower() or "qwen" in args.net.lower(): is_llama = True layers = model.model.layers model.model.embed_tokens = model.model.embed_tokens.to(dev) model.model.norm = model.model.norm.to(dev) DecoderLayer = QuantLlamaDecoderLayer pairs = { "q_proj":"qkv", "o_proj":"out", "up_proj":"fc1", } if args.use_down_scale is True: pairs["down_proj"] = "fc2" layer_name_prefix = "model.layers" else: raise NotImplementedError("Only llama/qwen/vicuna are kept in this open-source snapshot.") # import ipdb;ipdb.set_trace() layers[0] = layers[0].to(dev) fp32_type = torch.float fp16_type = torch.bfloat16 if args.use_bfloat16 is True else torch.float16 act_dtype = fp16_type if args.fp16_act is True else fp32_type inps = torch.zeros( (args.nsamples, lm.seqlen, model.config.hidden_size), dtype=act_dtype, device="cpu" ) # import ipdb;ipdb.set_trace() cache = {"i": 0} # catch the first layer input class Catcher(nn.Module): def __init__(self, module): super().__init__() self.module = module self.is_llama = False def forward(self, inp, **kwargs): inps[cache["i"]] = inp.cpu() cache["i"] += 1 cache["attention_mask"] = kwargs["attention_mask"] cache["position_embeddings"] = kwargs["position_embeddings"] if self.is_llama: cache["position_ids"] = kwargs["position_ids"] raise ValueError layers[0] = Catcher(layers[0]) layers[0].is_llama = is_llama with torch.no_grad(): for batch in dataloader: if cache["i"] >= args.nsamples: break try: model(batch[0].to(dev)) except ValueError: pass # move embedding layer and first layer to cpu layers[0] = layers[0].module layers[0] = layers[0].cpu() 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 NotImplementedError("Only llama/qwen/vicuna are kept in this open-source snapshot.") # import ipdb;ipdb.set_trace() inps = inps.to("cpu") # same input of first layer for fp model and quant model cleanup_memory(logger=logger) attention_mask = cache["attention_mask"][:,:,:,:args.seqlen] if attention_mask is not None: attention_mask_batch = attention_mask.repeat(args.batch_size,1,1,1) else: logger.info( "No attention mask caught from the first layer." " Seems that model's attention works without a mask." ) attention_mask_batch = None if attention_mask is not None: infer_attention_mask = attention_mask.repeat(args.inference_batch_size,1,1,1) # import ipdb;ipdb.set_trace() if is_llama: position_ids = cache["position_ids"] position_embeddings = cache["position_embeddings"] else: position_ids = None position_embeddings = None if args.quant_mode in ["fp16"]: args.resume = None if args.resume: slider_parameters = torch.load(args.resume) else: slider_parameters = {} if args.train_resume is not None and args.test_mode is False: slider_parameters = torch.load(args.train_resume) args.resume = args.train_resume args.quant_layer_list = [int(layer_id) for layer_id in range(len(layers))] logger.info(f"these layer will quant:{args.quant_layer_list}") if args.use_lora is True: args.lora_layer_list = args.quant_layer_list # only when quant use lora else: args.lora_layer_list = [] logger.info(f"these layer will refine with lora:{args.lora_layer_list}") args.lora_iter_num_list = {layer_id:1 for layer_id in range(len(layers))} logger.info(f"each layer will refine with lora num iter:{args.lora_iter_num_list}") args.lora_r_list = {layer_id:args.lora_rank for layer_id in range(len(layers))} logger.info(f"each layer lora's r:{args.lora_r_list}") args.quant_mode_layer_list = { layer_id:(args.quant_mode if layer_id in args.quant_layer_list else "fp16") for layer_id in range(len(layers)) } logger.info(f"each layer quant mode:{args.quant_mode_layer_list}") if args.sliding_layer is None: args.sliding_layer = args.num_layer logger.info(f"sliding_layer:{args.sliding_layer}") init_quant_rate = args.quant_rate if args.quant_rate_list is None: args.quant_rate_list = np.linspace(0, 1, args.quant_step+1).tolist()[1:] logger.info(f"quant_step:{args.quant_step} quant_rate_list:{args.quant_rate_list} lora_quant:{args.lora_quant}") if args.test_mode is True: args.quant_rate = 1.0 if (args.resume is not None or args.train_resume is not None) and args.resume_layers_num is None: args.resume_layers_num = len(layers) # 模型初始化 model_attr = dict( is_llama=is_llama, pairs=pairs, layer_name_prefix=layer_name_prefix, slider_parameters=slider_parameters, dtype=fp32_type, ) init_model(config=lm.model.config,layers=layers,args=args,DecoderLayer=DecoderLayer,model_attr=model_attr,logger=logger,dev="cpu") logger.info("Model Initialized") num_update_steps_per_epoch = math.ceil(args.nsamples / args.batch_size) global_start_time = time.time() cur_epochs = 0 # 已经训练的epochs if args.fill_window_size is not None: args.fill_start_window_size = args.fill_window_size args.fill_end_window_size = args.fill_window_size if args.layer_windows_scheduler is not None: layer_windows_scheduler = [] for window_str in args.layer_windows_scheduler.split(","): layer_windows_scheduler.append([int(s) for s in window_str.split("-")]) num_round = len(layer_windows_scheduler) # import ipdb;ipdb.set_trace() assert layer_windows_scheduler[-1][-1] == len(layers) - 1 elif args.fill_window_size is not None: total_num_layers = len(layers) if args.fill_start_window_size is not None: start_layer_windows_scheduler = [list(range(i+1)) for i in range(args.fill_start_window_size)] start_len = (args.fill_start_window_size - args.sliding_layer) else: start_layer_windows_scheduler = [] start_len = 0 if args.fill_end_window_size is not None: end_start_layer_windows_scheduler = [list(range(total_num_layers-args.fill_end_window_size +i,total_num_layers)) for i in range(args.fill_end_window_size)] end_len = (args.fill_end_window_size - args.sliding_layer) else: end_start_layer_windows_scheduler = [] end_len = 0 mid_len = total_num_layers - start_len - end_len mid_round = math.ceil((mid_len - args.num_layer) / args.sliding_layer) + 1 mid_layer_windows_scheduler = [ [i for i in range(r * args.sliding_layer + start_len , min(r * args.sliding_layer + args.num_layer + start_len ,len(layers)))] for r in range(mid_round) ] layer_windows_scheduler = start_layer_windows_scheduler + mid_layer_windows_scheduler + end_start_layer_windows_scheduler num_round = len(layer_windows_scheduler) else: num_round = math.ceil((len(layers) - args.num_layer) / args.sliding_layer) + 1 layer_windows_scheduler = [ [i for i in range(r * args.sliding_layer, min(r * args.sliding_layer + args.num_layer,len(layers)))] for r in range(num_round) ] logger.info(f"layer_windows_scheduler is {layer_windows_scheduler}") if teach_lm is not None: teach_model = teach_lm.model teach_model.config.use_cache = False teach_layers = teach_model.model.layers teach_layers = teach_layers.to(dev) logger.info(f"Teacher model Initialized from {args.teach_model}") else: teach_layers = layers cleanup_memory(logger=logger) assert args.quant_step == len(args.quant_rate_list) if args.circular_aug: assert len(args.quant_rate_list) > 1 if args.littlt_bs_round is not None: littlt_bs_round = [int(i) for i in args.littlt_bs_round.split(",")] littlt_bs_round = [i if i >=0 else num_round+i for i in littlt_bs_round] else: littlt_bs_round = [] global_batch_size = args.batch_size for step,quant_rate in enumerate(args.quant_rate_list): if args.circular_aug is True and step+1 == len(args.quant_rate_list): windows_quant_inps[:,-1] = copy.deepcopy(inps) windows_fp_inps[:,-1] = copy.deepcopy(inps) else: windows_quant_inps = copy.deepcopy(inps).unsqueeze(1).repeat(1,args.last_round_inp_num,1, 1) # if None, not need cache. if True, need but had not been cached windows_fp_inps = copy.deepcopy(windows_quant_inps) # if None, not need cache. if True, need but had not been cached if step+1 == args.quant_step and args.debug is False: inps = None print("delete inps!") cleanup_memory(logger=logger) if args.use_quant_tar_loss: windows_quant_targets = copy.deepcopy(windows_quant_inps) # if None, not need cache. if True, need but had not been cached else: windows_quant_targets = None windows_fp_targets = copy.deepcopy(windows_fp_inps) args.quant_rate = quant_rate if args.low_memory is False: windows_quant_inps = windows_quant_inps.to(dev) windows_fp_inps = windows_fp_inps.to(dev) cleanup_memory(logger=logger) for r in range(num_round): if r in littlt_bs_round: args.batch_size = 1 else: args.batch_size = global_batch_size if args.test_mode is True: break layer_id_list = layer_windows_scheduler[r] logger.info(f"=== Step: {step+1}/{args.quant_step} Round: {r+1}/{num_round} ===") logger.info( f"=== Start quantize layer{layer_id_list[0]}-layer{layer_id_list[-1]} ===" ) if args.loss_function == "mse": loss_func = torch.nn.MSELoss() elif args.loss_function == "huber": delta = 0.1 + r/num_round * args.huber_loss_max loss_func = torch.nn.HuberLoss(delta=delta) else: raise NotImplementedError("only support mse and huber loss function") # del finished layers if args.low_cpu_memory is True and step+1 == len(args.quant_rate_list): for l_idx in range(layer_id_list[0]): if lm.model.model.layers[l_idx] is not None: lm.model.model.layers[l_idx] = torch.nn.Identity() # import ipdb;ipdb.set_trace() logger.info(f"del layer 0-{layer_id_list[0]-1}") sub_layers = layers[layer_id_list[0]:layer_id_list[-1]+1] teach_sub_layers = teach_layers[layer_id_list[0]:layer_id_list[-1]+1] cleanup_memory(logger=logger) logger.info(f"layer_id_list: {layer_id_list}") sub_layers = to_float(sub_layers,dtype=fp32_type) sub_layers = to_dev(sub_layers, [dev] * len(sub_layers)) #single gpu teach_sub_layers = to_float(teach_sub_layers,dtype=fp32_type) teach_sub_layers = to_dev(teach_sub_layers, [dev] * len(teach_sub_layers)) #single gpu acts_round_idx = list(range(r+1-args.last_round_inp_num,r+1)) logger.info(f"act_round_idx is {acts_round_idx}") # import ipdb;ipdb.set_trace() if r >= args.start_round: with torch.no_grad(): with torch.cuda.amp.autocast(dtype=fp16_type): for r_idx in range(args.last_round_inp_num): # get quant_target if args.use_quant_tar_loss: for j in tqdm(range(0,args.nsamples,args.inference_batch_size)): bs_local = min(args.inference_batch_size,args.nsamples-j) windows_quant_targets[j:j+bs_local,r_idx] = obtain_teacher_output( teach_sub_layers, windows_quant_inps[j:j+bs_local,r_idx], infer_attention_mask[:bs_local], position_ids, position_embeddings=position_embeddings, args=args, devs=[dev] * len(teach_sub_layers), ) logger.info(f"finish to obtain quant_target round {acts_round_idx[r_idx]} of full-precision model!") # get fp_target for j in tqdm(range(0,args.nsamples,args.inference_batch_size)): bs_local = min(args.inference_batch_size,args.nsamples-j) windows_fp_targets[j:j+bs_local,r_idx] = obtain_teacher_output( teach_sub_layers, windows_fp_inps[j:j+bs_local,r_idx], infer_attention_mask[:bs_local], position_ids, position_embeddings=position_embeddings, args=args, devs=[dev] * len(teach_sub_layers), ) logger.info(f"finish to obtain fp_target round {acts_round_idx[r_idx]} of full-precision model!") cleanup_memory(logger=logger) epochs = args.epochs // args.quant_step total_epochs = args.epochs*num_round if args.layers_assigned_gpu is not None: devs = [torch.device(f"cuda:{gpu}") for gpu in args.layers_assigned_gpu.split(",")] assert len(devs) == len(sub_layers), "layers_assigned_gpu number is not equal to layer number!" sub_layers = to_dev(sub_layers, devs) #mutil-gpu else: devs = [dev] * len(sub_layers) max_train_steps = epochs * num_update_steps_per_epoch lr_factor = args.batch_size logger.info(f"auto lr scale is {args.auto_lr_scale} lora_lr is {args.lora_lr*lr_factor} scale_lr is {args.scale_lr*lr_factor} lwc_lr is {args.lwc_lr*lr_factor}") params = [] if args.use_lora is True and (r not in littlt_bs_round or args.quant_mode == "lora_only"): params.append({"params":get_slider_parameters(sub_layers, ["lora_"]),"lr":args.lora_lr*lr_factor,"weight_decay":0.0}) if args.scale_lr > 0: params.append({"params":get_slider_parameters(sub_layers, ["scale"]),"lr":args.scale_lr*lr_factor,"weight_decay":0.0}) if args.lwc_lr > 0: params.append({"params":get_lwc_parameters(sub_layers),"lr":args.lwc_lr*lr_factor,"weight_decay":0.0}) optimizer = torch.optim.AdamW(params) lr_scheduler = get_scheduler( name="linear", optimizer=optimizer, num_warmup_steps=max_train_steps*args.warmup_ratio, num_training_steps=max_train_steps, ) # import ipdb;ipdb.set_trace() if args.use_ddp and r >= args.start_round: sub_layers = SubLayer(sub_layers,quant_mode_sub_layer_list=[args.quant_mode_layer_list[i] for i in layer_id_list], attention_mask=attention_mask_batch,position_ids=position_ids,position_embeddings=position_embeddings,args=args) cleanup_memory(logger=logger) if args.use_ddp and r >= args.start_round: sub_layers = sub_layers.cuda() if args.use_ddp: dist.barrier() # train loop if r < args.start_round: logger.info(f"round {r} skip because resume from disk!") qdataset = None else: qdataset = Quant_dataset(aug_quant_inps=windows_quant_inps,aug_fp_inps=windows_fp_inps,aug_quant_targets=windows_quant_targets,aug_fp_targets=windows_fp_targets,samples_num=args.nsamples,windows_num=args.last_round_inp_num,args=args) sub_layers,cur_epochs = train_one_round(r=r,epochs=epochs,sub_layers=sub_layers,layer_id_list=layer_id_list,attention_mask_batch=attention_mask_batch,cur_epochs=cur_epochs, position_ids=position_ids,position_embeddings=position_embeddings,devs=devs,args=args,logger=logger,max_train_steps=max_train_steps,optimizer=optimizer,lr_scheduler=lr_scheduler,qdataset=qdataset, init_quant_rate=init_quant_rate,fp16_type=fp16_type,global_start_time=global_start_time,total_epochs=total_epochs,loss_func=loss_func,acts_round_idx=acts_round_idx) if args.use_ddp and r >= args.start_round: sub_layers = sub_layers.module.module sub_layers = sub_layers.to("cpu") if args.use_ddp: dist.barrier() del optimizer,qdataset,lr_scheduler for r_idx in range(args.last_round_inp_num-1): windows_quant_inps[:,r_idx] = windows_quant_inps[:,r_idx+1] windows_fp_inps[:,r_idx] = windows_fp_inps[:,r_idx+1] cleanup_memory(logger=logger) sliding_layer = layer_windows_scheduler[min(r+1,num_round-1)][0] - layer_windows_scheduler[r][0] if r < num_round-1 and sliding_layer>0: sub_layers = to_dev(sub_layers, [dev] * len(sub_layers)) #single gpu with torch.no_grad(): with torch.cuda.amp.autocast(dtype=fp16_type): # get next fp_16 for j in tqdm(range(0,args.nsamples,args.inference_batch_size)): bs_local = min(args.inference_batch_size,args.nsamples-j) windows_fp_inps[j:j+bs_local,-1] = obtain_teacher_output( teach_sub_layers[:sliding_layer], windows_fp_inps[j:j+bs_local,-1], infer_attention_mask[:bs_local], position_ids, position_embeddings=position_embeddings, args=args, devs=[dev] * len(teach_sub_layers), ) logger.info(f"finish to obtain round {r+1} inps of full-precision model!") for j in tqdm(range(0,args.nsamples,args.inference_batch_size)): bs_local = min(args.inference_batch_size,args.nsamples-j) windows_quant_inps[j:j+bs_local,-1] = obtain_studnet_output( sub_layers[:sliding_layer], [args.quant_mode_layer_list[i] for i in layer_id_list[:sliding_layer]], windows_quant_inps[j:j+bs_local,-1], infer_attention_mask[:bs_local], position_ids, position_embeddings=position_embeddings, args=args, devs=[dev] * len(sub_layers), ) logger.info(f"finish to obtain round {r+1} inps of quant model!") with torch.no_grad(): for idx,i in enumerate(layer_id_list): qlayer = sub_layers[idx] qlayer.clear_temp_variable() if epochs>0: sub_layers[idx] = qlayer.to("cpu") slider_parameters[i] = slider_state_dict(qlayer) if args.use_ddp: if args.use_ddp and dist.get_rank() == 0: torch.save(slider_parameters, os.path.join(args.output_dir, f"slider_parameters.pth")) logger.info(f"save slider_parameters in layer{i} successfully!") else: torch.save(slider_parameters, os.path.join(args.output_dir, f"slider_parameters.pth")) logger.info(f"save slider_parameters in layer{i} successfully!") else: sub_layers[idx] = qlayer.to("cpu") del qlayer sub_layers = to_half(sub_layers,dtype=fp16_type) sub_layers = to_dev(sub_layers, ["cpu"] * len(sub_layers)) #single gpu teach_sub_layers = to_half(teach_sub_layers,dtype=fp16_type) teach_sub_layers = to_dev(teach_sub_layers, ["cpu"] * len(teach_sub_layers)) #single gpu replace_ori_layer(layers, sub_layers, layer_id_list, args) del sub_layers,teach_sub_layers cleanup_memory(logger=logger) if args.use_ddp: # 保证参数保存完整 dist.barrier() logger.info("Model quantization finished! start change model to inference mode!") args.quant_rate = 1.0 model_to_inference_mode(layers=layers,args=args,dtype=fp16_type,dev=dev) model.to(fp16_type) try: del quant_inps del fp_inps # del tmp_fp_inps except Exception as e: logger.info(f"del tensor occurs {e}, skip!") cleanup_memory(logger=logger) logger.info("Model is changed to inderence mode!") cleanup_memory(logger=logger) model.config.use_cache = use_cache return model,inps,infer_attention_mask,position_ids