|
当前已经全面进入从模型框架、SFT训练数据以及模型开源井喷的时期，一日恍若一月。 此外，更小，更快的训练和部署模型也逐步成为目前开源的一个趋势。 因此，本文主要介绍极低资源微调大模型方法LoRA以及BLOOM-LORA实现代码，供大家一起参考。 一、LoRA的原理 LoRA是一种以极低资源微调大模型的方法，其来自于论文LoRA: Low-Rank Adaptation of Large Language Models。 1. 大模型微调的困境 随着模型规模的不断扩大，模型会'涌现'出各种能力。特别是对大语言模型(LLM)来说，随着规模的扩大其在zero-shot、常识推理等能力上会有大幅度的提高。相比于规模较小的模型，大模型的微调成本和部署成本都非常高。例如，GPT-3 175B模型微调需要1.2TB的显存。此外，若针对不同下游任务微调多个模型，那么就需要为每个下游任务保存一份模型权重，成本非常高。在某些场景下，甚至可能需要针对不同的用户微调不同的模型，那么模型微调和部署的成本将不可接受。 因此，如何降低大模型微调和部署成本，将是大模型商用的重要一环。 2. LoRA之前的方法 在LoRA方法提出之前，也有很多方法尝试解决大模型微调困境的方法。其中有两个主要的方向：(1) 添加adapter层；(2) 由于某种形式的输入层激活。但是这两种方法都有局限性： 2.1 Adapter层会引入推理时延 简单来说，adapter就是固定原有的参数，并添加一些额外参数用于微调。上图中会在原始的transformer block中添加2个adapter，一个在多头注意力后面，另一个这是FFN后面。  显然，adapter会在模型中添加额外的层，这些层会导致大模型在推理时需要更多的GPU通信，而且也会约束模型并行。这些问题都将导致模型推理变慢。 2.2 prefix-tuning难以优化 prefix-tuning方法是受语言模型in-context learning能力的启发，只要有合适的上下文则语言模型可以很好的解决自然语言任务。但是，针对特定的任务找到离散token的前缀需要花费很长时间，prefix-tuning提出使用连续的virtual token embedding来替换离散token。  具体来说，对于transformer中的每一层，都在句子表征前面插入可训练的virtual token embedding。对于自回归模型(GPT系列)，在句子前添加连续前缀，即z = [PREFIX;x;y]。对于Encoder-Decoder模型(T5)，则在Ecoder和Decoder前都添加连续前缀 z = [PREFIX;x|PREFIX';y]。添加前缀的过程如上图所示。 虽然，prefix-tuning并没有添加太多的额外参数。但是，prefix-tuning难以优化，且会减少下游任务的序列长度。 3. 问题的正式表述 术语与约定。由于LoRA原理的介绍，会使用Transformer架构。因此，这里先给出一些术语约定。一个Transformer层的输入和输出维度尺寸为d_model，使用Wq、Wk、Wv和Wo表示自注意力模块中的query/key/value/output投影矩阵。W或w0表示预训练模型的权重矩阵， delata_W表示模型在适配过程中的梯度更新。r来表示LoRA模块的秩。使用Adam作为模型优化器，Transformer MLP前馈层的维度为 d_ffn=4xd_model。 问题表述。LoRA虽然与训练目标无关，这里还是以语言建模为例。假设给定一个预训练的自回归语言模型  目标是使该语言模型适应下游的摘要、机器阅读理解等任务。每个下游任务都有context-target样本对组成的训练集：  其中xi和yi都是token序列。例如，对于摘要任务，xi是文章内容， yi是摘要。 在完整微调的过程中，模型使用预训练好的权重来初始化模型，然后通过最大化条件语言模型来更新参数:  完整微调的主要缺点：对于每个下游任务，都需要学习不同的参数更新，因此，如果预训练模型很大，存储和部署许多独立的微调模型实例非常有挑战。 LoRA为了更加的参数高效，使用相对非常小的参数来表示任务相关的参数增量。 LoRA将会使用低秩表示来编码，同时实现计算高效和存储高效。 4. LoRA 通常，神经网络中会包含许多进行矩阵乘法的稠密层，这些层通常是满秩的。Adgajanyan et al.等人的研究表示预训练语言模型具有低的'内在维度'。  受该工作的启发，在模型适配下游任务的过程中，权重更新也应该具有低的“内在秩”。对于预训练权重矩阵：  可以通过低秩分解来表示其更新：  并且秩  在训练过程中，W0被冻结且不接受梯度更新，A和B则是可训练参数。注意，W0和delata-W=BA都会乘以相同的输入。对于  ，前向传播变为：  对矩阵A 使用随机高斯初始化，对矩阵B使用0进行初始化，因此delata-W=BA 在训练的开始为0。使用alpha/r来缩放delata-Wx。当使用Adam优化时，经过适当的缩放初始化，调优alpha与调优学习率大致相同。 当进行部署时，以显式的计算和存储W=W0+BA ，并正常执行推理。W0和BA都是 。当需要转换至另一个下游任务，可以通过减去BA来恢复W0，然后添加不同的B'A'。至关重要的是，这保证不会引人任何额外的推理时延。 5. LoRA应用于Transformer 理论上，LoRA可以应用于任何神经网络的权重矩阵，从而减少可训练参数的数量。Transformer架构中的自注意力模块有4个权重矩阵： ，以及两个MLP模型的权重矩阵。将Wq(或者Wk,Wv )作为一个维度为 的单个矩阵。为了简单和参数高效，本研究仅限于适配下游任务的注意力权重，并冻结MLP模块。 优点。最显著的优点是显存和存储空间的减少。对于使用Adam训练的大型Transformer，若 ，由于不需要存储被冻结参数的优化器状态，VRAM使用量减少2/3。 对于GPT-3 175B，训练中的显存消耗从1.2TB减少自350GB。当r=4并且仅调整query矩阵和value矩阵时，checkpoint大小减少10000倍(从350GB减少自35MB)。 另一个优点是，可以在部署时以更低的成本切换任务，仅需要交换LoRA权重即可。此外，与完全微调相比，GPT-3 175B训练速度提高了25%，因为不需要计算绝大多数参数的梯度。 二、代码：实现BLOOM-LoRA 本小节展示如何使用LoRA微调大语言模型bloom。 数据：使用BELLE提供的100万指令微调数据； 模型：使用bloomz-7b1-mt，该版本的bloomz也是经过指令微调后的模型。 依赖包：使用transformers提供模型加载和训练；使用peft提供LoRA实现；使用DeepSpeed提供训练加速。 注意：peft包目前还处于快速迭代当中，后续接口可能会有大的变动，也可能存在一些bug。关键依赖包版本： transformers==4.26.1 torch==1.13.1 deepspeed==0.8.2 peft==0.2. 1. 训练代码 为了简洁，假设训练代码位于train.py。 1.1 导入依赖包import os import torch import random import datasets import numpy as np from tqdm import tqdm from typing import Dict from torch.utils.data import DataLoader from transformers import ( AutoModelForCausalLM, AutoTokenizer, DataCollatorForSeq2Seq, TrainingArguments, Trainer ) from peft import ( LoraConfig, TaskType, get_peft_model, get_peft_model_state_dict, set_peft_model_state_dict ) def set_random_seed(seed): if seed is not None and seed > 0: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.random.manual_seed(seed) torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True set_random_seed(1234) 1.2 设置参数 # LoRA参数 LORA_R = 8 LORA_ALPHA = 32 LORA_DROPOUT = 0.1 # 训练参数 EPOCHS=3 LEARNING_RATE=5e-5 OUTPUT_DIR='./checkpoints' BATCH_SIZE=4 # 2 GRADIENT_ACCUMULATION_STEPS=3 # 其他参数 MODEL_PATH = 'bigscience/bloomz-7b1-mt' DATA_PATH = './data/belle_open_source_1M.train.json' MAX_LENGTH = 512 PATTERN = '{}\n{}' DS_CONFIG = 'ds_zero2_config.json' tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH) # 加载tokenizer 1.3 加载数据dataset = datasets.load_dataset('json', data_files=DATA_PATH) # print(dataset['train'][0]) 1.4 tokenize def tokenize(text: str, add_eos_token=True): result = tokenizer( text, truncation=True, max_length=MAX_LENGTH, padding=False, return_tensors=None) # 判断是否要添加eos_token if (result['input_ids'][-1] != tokenizer.eos_token_id and len(result['input_ids']) < MAX_LENGTH and add_eos_token): result['input_ids'].append(tokenizer.eos_token_id) result['attention_mask'].append(1) result['labels'] = result['input_ids'].copy() return result def preprocess(example: Dict, train_on_inputs: bool = False): prompt = example['input'] response = example['target'] text = PATTERN.format(prompt, response) tokenized_inp = tokenize(text) # 若train_on_inputs为False，则将label中与input相关的token替换为-100 if not train_on_inputs: tokenized_prompt = tokenize(prompt,add_eos_token=False) prompt_tokens_len = len(tokenized_prompt['input_ids']) tokenized_inp['labels'] = [-100]*prompt_tokens_len + tokenized_inp['labels'][prompt_tokens_len:] return tokenized_inp train_data = dataset['train'].shuffle().map(preprocess, remove_columns=['id', 'input', 'target']) print(train_data[0]) 1.5 collate_fn # pad_to_multiple_of=8表示padding的长度是8的倍数 collate_fn = DataCollatorForSeq2Seq(tokenizer, pad_to_multiple_of=8, return_tensors='pt', padding=True) 1.6 加载模型 device_map = {'': int(os.environ.get('LOCAL_RANK') or 0)} # device_map指定模型加载的GPU;troch_dtype=torch.float16表示半精度加载模型 model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, torch_dtype=torch.float16, device_map=device_map) 1.7 LoRA相关# 转换模型 model = get_peft_model(model, lora_config) model.config.use_cache = False old_state_dict = model.state_dict model.state_dict = ( lambda self, *_, **__: get_peft_model_state_dict(self, old_state_dict()) ).__get__(model, type(model)) # 打印模型中的可训练参数 model.print_trainable_parameters() 1.8 训练参数 args = TrainingArguments( output_dir=OUTPUT_DIR, # checkpoint的存储目录 per_device_train_batch_size=BATCH_SIZE, # 单设备上的batch size gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS, # 梯度累加的step数 warmup_steps=100, num_train_epochs=EPOCHS, learning_rate=LEARNING_RATE, fp16=True, # 使用混合精度训练 logging_steps=50, evaluation_strategy='no', # 不进行评估 save_strategy='steps', save_steps=2000, # 保存checkpoint的step数 save_total_limit=5, # 最多保存5个checkpoint deepspeed=DS_CONFIG ) 1.9 模型训练trainer = Trainer( model=model, train_dataset=train_data, eval_dataset=None, args=args, data_collator=collate_fn ) trainer.train() model.save_pretrained('best_model') 2. DeepSpeed配置文件 DeepSpeed配置文件名为ds_zero2_config.json。 { 'train_micro_batch_size_per_gpu': 'auto', 'gradient_accumulation_steps': 'auto', 'steps_per_print': 50, 'gradient_clipping': 1.0, 'zero_optimization': { 'stage': 2, 'offload_optimizer': { 'device': 'cpu' }, 'contiguous_gradients': true, 'overlap_comm': true }, 'zero_allow_untested_optimizer': true, 'fp16': { 'enabled': true, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1 }, 'optimizer': { 'type': 'Adam', 'params': { 'lr': 'auto', 'betas': 'auto', 'eps': 'auto', 'weight_decay': 'auto' } }, 'activation_checkpointing': { 'partition_activations': true, 'contiguous_memory_optimization': true }, 'wall_clock_breakdown': false } 3. 启动deepspeed --include=localhost:0,1,2,3 train.py 4. 推理 推理文件名为inference.py import torch from peft import PeftModel from transformers import AutoModelForCausalLM, AutoTokenizer BASE_MODEL = 'bigscience/bloomz-7b1-mt' LORA_WEIGHTS = 'best_model' tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL) model = AutoModelForCausalLM.from_pretrained( BASE_MODEL, torch_dtype=torch.float16, # 加载半精度 device_map={'':0}, # 指定GPU 0 ) model.eval() # 加载LoRA权重 model = PeftModel.from_pretrained(model, LORA_WEIGHTS, torch_dtype=torch.float16) model.half() prompt = '' inp = tokenizer(prompt, max_length=512, return_tensors='pt').to('cuda') outputs = model.generate(input_ids=inp['input_ids'], max_new_tokens=256) print(tokenizer.decode(outputs[0])) 参考文献 1、https://zhuanlan.zhihu.com/p/618073170 2、arxiv.org/pdf/2106.09685.pdf 3、https://zhuanlan.zhihu.com/p/615235322 4、https://github.com/tloen/alpaca-lora/blob/main/finetune.py 5、https://github.com/huggingface/peft/blob/main/examples/conditional_generation/peft_lora_seq2seq_accelerate_ds_zero3_offload.py |
|
|
