# 处理方法 - 把训练集SolRet_training的query和value做对比学习 # 参考代码 ```python ########################################################################### imports import os import json # os.environ['CUDA_VISIBLE_DEVICES'] = '0,1' import argparse import random import math from time import time import numpy as np import torch from torch.utils.data import DataLoader from torch.optim import AdamW from torch.optim.lr_scheduler import LambdaLR from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer from transformers.models.qwen2.modeling_qwen2 import Qwen2ForCausalLM as Qwen2ForCausalLMOrig from transformers.models.qwen2.modeling_qwen2 import QWEN2_INPUTS_DOCSTRING,_CONFIG_FOR_DOC from transformers.utils import add_start_docstrings_to_model_forward, replace_return_docstrings from typing import List, Optional, Tuple, Union from torch import nn from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss from torch.nn import functional as F from transformers.modeling_outputs import CausalLMOutputWithPast from datasets import load_dataset import deepspeed from peft import LoraConfig, get_peft_model, PeftModel from functools import partial import pynvml class print_time: def __init__(self, *desc): self.desc = desc def __enter__(self): print(*self.desc) self.t = time() def __exit__(self, type, value, traceback): print(f'{time()-self.t:.02f}s') # 指定 gpu def set_gpus(gpu): torch.cuda.set_device(gpu) def set_seed(seed): os.environ['PYTHONHASHSEED'] = str(seed) random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) def set_cuda(deterministic=True): if torch.cuda.is_available(): # 卷积算法确定 torch.backends.cudnn.deterministic = deterministic # 设置 torch.backends.cudnn.benchmark=True 将会让程序在开始时花费一点额外时间，为整个网络的每个卷积层搜索最适合它的卷积实现算法，进而实现网络的加速。适用场景是网络结构固定（不是动态变化的），网络的输入形状（包括 batch size，图片大小，输入的通道）是不变的，其实也就是一般情况下都比较适用。反之，如果卷积层的设置一直变化，将会导致程序不停地做优化，反而会耗费更多的时间。 torch.backends.cudnn.benchmark = not deterministic class Qwen2ForCausalLM(Qwen2ForCausalLMOrig): def __init__(self, config): super().__init__(config) retrieve_size=256 self.retrieve_proj = nn.Linear(config.hidden_size, retrieve_size, bias=False) @add_start_docstrings_to_model_forward(QWEN2_INPUTS_DOCSTRING) @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC) def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]: r""" Args: labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. Returns: Example: ```python >>> from transformers import AutoTokenizer, Qwen2ForCausalLM >>> model = Qwen2ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) >>> prompt = "Hey, are you conscious? Can you talk to me?" >>> inputs = tokenizer(prompt, return_tensors="pt") >>> # Generate >>> generate_ids = model.generate(inputs.input_ids, max_length=30) >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) ################################################################################################################# loss_fct = CrossEntropyLoss() # query_feats # value_feats with torch.no_grad(): outputs = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, cache_position=cache_position, ) hidden_states = outputs[0] input_feats = F.normalize( self.retrieve_proj(hidden_states[:, -1, :]), dim=-1 ) # 计算第一维度的一半 half_size = input_feats.shape[0] // 2 # 取出前一半 query_feats = input_feats[:half_size, :] # 取出后一半 value_feats = input_feats[half_size:, :] sim_t2q = torch.matmul( query_feats, value_feats.permute(1,0) ) sim_q2t = torch.matmul( value_feats, query_feats.permute(1,0) ) #n*n targets = torch.linspace(0, query_feats.size(0) - 1, query_feats.size(0), dtype=int).to( query_feats.device ) loss_itc = loss_fct(sim_t2q, targets) + loss_fct(sim_q2t, targets) # print(f'loss_itc\n{loss_itc}\n\nsim_t2q\n{sim_t2q}\n\nsim_q2t\n{sim_q2t}') return CausalLMOutputWithPast( loss=loss_itc, logits=None, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) # 送入模型前先tokenize def tokenize_batch(batch,tokenizer): # 前面半截是query，后面半截是value batch_list=[] batch_list.extend(batch['query']) batch_list.extend(batch['value']) batch_tokenize=tokenizer(batch_list,return_tensors='pt',padding='longest',return_token_type_ids=False) # 手动添加eos token batch_size=batch_tokenize['input_ids'].shape[0] eos_tokens = torch.full((batch_size, 1), tokenizer.eos_token_id, dtype=batch_tokenize['input_ids'].dtype) # 拼接EOS token到input_ids的末尾 batch_tokenize['input_ids'] = torch.cat([batch_tokenize['input_ids'], eos_tokens], dim=-1) # 对于attention_mask，我们也需要添加1来表明新添加的EOS token是有效的 attention_masks_eos = torch.ones((batch_size, 1), dtype=batch_tokenize['attention_mask'].dtype) batch_tokenize['attention_mask'] = torch.cat([batch_tokenize['attention_mask'], attention_masks_eos], dim=-1) return batch_tokenize # 保存模型，参数 def save_model_tokenizer_config_params(args,model_engine,tokenizer,training_params_save,epoch,step=None): if step is not None: dir_name=os.path.join(args.save_dir,f'epoch_{epoch}',f'step_{step}') else: dir_name=os.path.join(args.save_dir,f'epoch_{epoch}') model_engine.save_16bit_model(dir_name) tokenizer.save_pretrained(dir_name) model_engine.model.config.save_pretrained(dir_name) # 保存训练参数 with open(os.path.join(dir_name,'training_params_save.json'), 'w') as f: json.dump(training_params_save, f, indent=4) # 加载数据，并且选择刚好合适的大小 def load_training_dataset(args, file_folder_path): print('loading dataset \n') # 给定的路径是目录，从文件夹中加载parquet文件 if os.path.isdir(file_folder_path): parquet_files=[] for filepath,dirnames,filenames in os.walk(file_folder_path): for filename in filenames: if filename.endswith('parquet'): fullname = os.path.join(filepath, filename) parquet_files.append(fullname) ds = load_dataset("parquet", data_files=parquet_files)['train'] # 给定的路径是文件名，从文件中加载parquet文件 elif os.path.isfile(file_folder_path): ds = load_dataset("parquet", data_files=file_folder_path)['train'] # shuffle数据集 ds = ds.shuffle(seed=args.seed) return ds # 初始化模型 def initialize_model_tokenizer(args): tokenizer = AutoTokenizer.from_pretrained(args.model) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side='left' tokenizer.add_bos_token=False tokenizer.add_eos_token=False print('initializing model \n') config = AutoConfig.from_pretrained(args.model) if args.gradient_checkpointing: config.gradient_checkpointing = True config.use_cache = False model = Qwen2ForCausalLM.from_pretrained(args.model, config=config) # if hasattr(model, "enable_input_require_grads"): # model.enable_input_require_grads() # else: # def make_inputs_require_grad(module, input, output): # output.requires_grad_(True) # model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) model.train() def find_all_linear_names(model): cls = torch.nn.Linear lora_module_names = set() for name, module in model.named_modules(): if isinstance(module, cls): names = name.split('.') lora_module_names.add(names[0] if len(names) == 1 else names[-1]) if 'lm_head' in lora_module_names: # needed for 16-bit lora_module_names.remove('lm_head') return list(lora_module_names) ###############################################LoRA if LoRA: loraconfig = LoraConfig( r=128, lora_alpha=256, target_modules=find_all_linear_names(model), lora_dropout=0.05, bias="none", task_type="CAUSAL_LM", ) model = get_peft_model(model, loraconfig) # model=PeftModel.from_pretrained(model, args.save_dir) # TODO(enijkamp): we need to set this flag twice? #forget all activations during forward and recompute during the backward. if args.gradient_checkpointing: model.gradient_checkpointing_enable() return model,tokenizer def deepspeed_train(args,ds,model,tokenizer): set_seed(args.seed) set_cuda(deterministic=deterministic) print('initializing deepspeed \n') model_parameters = list(filter(lambda p: p.requires_grad, model.parameters())) optimizer=AdamW(model_parameters, lr=lr, betas=(0.9, 0.999),eps=1e-8,weight_decay=0.05) def _get_cosine_schedule_with_warmup_lr_lambda(current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_cycles: float): if current_step < num_warmup_steps: return float(current_step) / float(max(1, num_warmup_steps)) progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps)) return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))) lr_lambda = partial( _get_cosine_schedule_with_warmup_lr_lambda, num_warmup_steps=100, num_training_steps=num_epochs*steps, num_cycles=0.5, ) scheduler=LambdaLR(optimizer, lr_lambda) model_engine, optimizer, dataloader, lr_scheduler = deepspeed.initialize(config=args.deepspeed_config, model=model, model_parameters=model_parameters,training_data=ds,optimizer=optimizer, lr_scheduler=scheduler,) torch.cuda.empty_cache() print('starting training \n') training_params_save={ 'epoch':[], 'step':[], 'loss':[], 'lr':[], 'cost_time':[], 'total_cost_time':[], 'percent':[], } total_start_time=time() for epoch in range(1,num_epochs+1): for step, batch in enumerate(dataloader, start=1): batch=tokenize_batch(batch,tokenizer) start_time=time() batch.to(model_engine.device) loss = model_engine(input_ids=batch['input_ids'],attention_mask=batch['attention_mask']).loss model_engine.backward(loss) model_engine.step() if step % steps_per_print == 0: # 打印训练相关信息 percent=step*args.deepspeed_config['train_micro_batch_size_per_gpu']*num_gpus/len(ds) cost_time=time()-start_time total_cost_time=time()-total_start_time print(f'epoch {epoch} step {step} {percent:.2%} loss:{loss:8.3f} time:{cost_time:.2f} total time:{total_cost_time:.2f}\n') meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle) print(f'GPU State --- Free:{meminfo.free/1024**2} Used:{meminfo.used/1024**2} Total:{meminfo.total/1024**2}\n') #总的显存大小（float） utilization = pynvml.nvmlDeviceGetUtilizationRates(handle) print(f'GPU Util --- {utilization.gpu}%\n') # gpu利用率 # 保存训练相关参数 training_params_save['epoch'].append(epoch) training_params_save['step'].append(step) training_params_save['loss'].append(float(loss)) training_params_save['lr'].append(optimizer.param_groups[0]['lr']) training_params_save['cost_time'].append(round(cost_time,4)) training_params_save['total_cost_time'].append(round(total_cost_time,4)) training_params_save['percent'].append(round(percent,4)) # 每隔steps_per_save step保存一次 if save_flag and step%steps_per_save==0: save_model_tokenizer_config_params(args,model_engine,tokenizer,training_params_save,epoch,step) # 每个epoch保存一次 if save_flag: save_model_tokenizer_config_params(args,model_engine,tokenizer,training_params_save,epoch) if __name__ == '__main__': args=argparse.Namespace() args.seed = 0 model_name='Qwen2p5Coder_1p5B' args.model = f'/home/chenzy/models/{model_name}' args.training_dataset = '/home/chenzy/sysu_datasets/SolRet_training' args.save_dir=f'/home/chenzy/models/save/{model_name}_SolRet' ds=load_training_dataset(args, args.training_dataset) lr=1e-4 train_micro_batch_size_per_gpu=128 gradient_accumulation_steps=1 num_gpus=4 dataset_len=len(ds) steps=dataset_len//(num_gpus*train_micro_batch_size_per_gpu) # steps=5 select_dataset_num=steps*num_gpus*train_micro_batch_size_per_gpu ds=ds.select(range(select_dataset_num)) num_epochs=3 steps_per_print=25 steps_per_save=steps//4 args.gradient_checkpointing=True save_flag=True deterministic=True LoRA=False pynvml.nvmlInit() handle = pynvml.nvmlDeviceGetHandleByIndex(0) # 指定显卡号 model,tokenizer=initialize_model_tokenizer(args) DEEPSPEED_CONFIG = \ { # 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'initial_scale_power': 12, 'hysteresis': 2, 'min_loss_scale': 1}, # 'bf16': {'enabled':True}, 'fp16': {'enabled':True}, # 'optimizer': {'type': 'AdamW', 'params': {'lr': 1e-05, 'betas': [0.9, 0.95], 'eps': 1e-08, 'weight_decay': 0.0}}, # 'scheduler': {'type': 'WarmupLR', 'params': {'warmup_min_lr': 0, 'warmup_max_lr': 1e-05, 'warmup_num_steps': 100}}, #stage 1: optimizer states #stage 2: gradients #stage 3: model parameters # 'zero_optimization': { # 'stage': 2, # #pin_memory improve the throughput # #device: or none # # 'offload_optimizer': {'device': 'cpu', 'pin_memory': True}, # # 'offload _param': {'device': 'cpu', 'pin_memory': True}, # #trade off GPU Ram and latency # 'overlap_comm': True, # # reduce memory fragmentation # 'contiguous_gradients': True, # # default =1e9, when not using NVMe # # 'sub_group_size': 1e9, # # 'reduce_bucket_size': 16777216, # # 'stage3_prefetch_bucket_size': 15099494.4, # # 'stage3_param_persistence_threshold': 40960, # # 'gather_16bit_weights_on_model_save': True, # }, 'zero_optimization': {'stage': 0}, # This is the amount of data samples that leads to one step of model update.train_batch_size must be equal to train_micro_batch_size_per_gpu * gradient_accumulation * number of GPUs #'train_batch_size': train_batch_size, # Batch size to be processed by one GPU in one step (without gradient accumulation). 'train_micro_batch_size_per_gpu': train_micro_batch_size_per_gpu, # Number of training steps to accumulate gradients before averaging and applying them. 'gradient_accumulation_steps': gradient_accumulation_steps, 'gradient_clipping': 1.0, # Print progress report every N training steps. The report includes the number of training steps, number of skipped optimizer updates 'steps_per_print': steps_per_print*5, 'wall_clock_breakdown': False, 'compression_training': {'weight_quantization': {'shared_parameters': {}, 'different_groups': {}}, 'activation_quantization': {'shared_parameters': {}, 'different_groups': {}}, 'sparse_pruning': {'shared_parameters': {}, 'different_groups': {}}, 'row_pruning': {'shared_parameters': {}, 'different_groups': {}}, 'head_pruning': {'shared_parameters': {}, 'different_groups': {}}, 'channel_pruning': {'shared_parameters': {}, 'different_groups': {}}} } args.deepspeed_config = DEEPSPEED_CONFIG deepspeed_train(args=args,ds=ds,model=model,tokenizer=tokenizer) ```