[Natural Language Processing] [Large Model] LoRA and BLOOM-LORA implementation codes for fine-tuning large model methods with very low resources

Very low resource fine-tuning large model method LoRA and BLOOM-LORA implementation code

Related blog
[Natural Language Processing] [Large Model] ChatGLM-6B model structure code analysis (stand-alone version)
[Natural Language Processing] [Large Model] BLOOM model structure source code analysis (stand-alone version)
[Natural Language Processing] [Large Model] extremely Low-resource fine-tuning of large model methods LoRA and BLOOM-LORA implementation code
[Natural Language Processing] [Large Model] DeepMind's large model Gopher
[Natural Language Processing] [Large Model] Chinchilla: Large language model with optimal training and computing utilization
[Natural Language Processing] [Large Model] Large Language Model BLOOM Reasoning Tool Test
[Natural Language Processing] [Large Model] GLM-130B: an open source bilingual pre-trained language model
[Natural Language Processing] [Large Model] 8- for large Transformers Introduction to bit matrix multiplication
[Natural Language Processing] [Large Model] BLOOM: A multilingual model with 176B parameters and open access
[Natural Language Processing] [Large Model] PaLM: A large language model based on Pathways
[Natural Language Processing] [chatGPT Series] Large language models can improve themselves

1. The principle of LoRA

LoRA is a method of fine-tuning large models with very low resources, which comes from the paper LoRA: Low-Rank Adaptation of Large Language Models .

1. The dilemma of fine-tuning large models

​ As the scale of the model continues to expand, the model will "emerge" with various capabilities. Especially for the large language model (LLM), as the scale expands, its capabilities such as zero-shot and common sense reasoning will be greatly improved. The fine-tuning cost and deployment cost of large models are very high compared to smaller models. For example, GPT-3 175B model fine-tuning requires 1.2TB of video memory. In addition, if multiple models are fine-tuned for different downstream tasks, it is necessary to save a model weight for each downstream task, which is very costly. In some scenarios, it may even be necessary to fine-tune different models for different users, then the cost of model fine-tuning and deployment will be unacceptable .

Therefore , how to reduce the cost of large-scale model fine-tuning and deployment will be an important part of large-scale model commercialization .

2. Pre-LoRA approach

Before the LoRA method was proposed, there were also many methods to try to solve the dilemma of large model fine-tuning. There are two main directions: (1) adding an adapter layer; (2) due to some form of input layer activation. But both approaches have limitations:

2.1 The Adapter layer will introduce inference delay

​ Simply put, the adapter is to fix the original parameters and add some additional parameters for fine-tuning. In the above figure, two adapters will be added to the original transformer block, one behind the multi-head attention, and the other behind the FFN.

Obviously, the adapter will add additional layers to the model, which will cause large models to require more GPU communication during inference, and will also constrain the model parallelism. These problems will lead to slow model inference .

2.2 prefix-tuning is difficult to optimize

The prefix-tuning method is inspired by the in-context learning ability of the language model. As long as there is a suitable context, the language model can solve natural language tasks well. However, it takes a long time to find the prefix of discrete tokens for a specific task, prefix-tuning proposes to use continuous virtual token embedding to replace discrete tokens.

​ Specifically, for each layer in the transformer, a trainable virtual token embedding is inserted in front of the sentence representation. For autoregressive models (GPT series), add a continuous prefix before the sentence, ie $z=[\text{PREFIX};x;y]$ . For the Encoder-Decoder model (T5), add a continuous prefix before Ecoder and Decoder$z=[\text{PREFIX};x|{\text{PREFIX}}^{\prime };y]$ . The process of adding a prefix is ​​shown in the figure above.

​Although , prefix-tuning does not add too many extra parameters. However, prefix-tuning is difficult to optimize and will reduce the sequence length of downstream tasks.

3. Formal formulation of the problem

​Terms and conventions . Due to the introduction of the LoRA principle, the Transformer architecture will be used. Therefore, some terminology conventions are given here first. The input and output dimensions of a Transformer layer are $d_{model}$, using $W_q$、$W_k$、$W_v$Japanese $W_o$Represents the query/key/value/output projection matrix in the self-attention module. $W$ or$W_0$Represents the weight matrix of the pre-trained model, $\Delta W\; representsthe$ gradient update of the model during the fitting process. $r$ to represent the rank of the LoRA module. Using Adam as the model optimizer, the dimension of the Transformer MLP feed-forward layer is$d_{ff}=4\times d_{model}$。

​Problem statement . Although LoRA has nothing to do with the training goal, here is an example of language modeling. Suppose a pre-trained autoregressive language model $P_{}(y|x)$ Φ \Phi$\Phi$ is a model parameter. The goal is to adapt the language model to downstream tasks such as summarization and machine reading comprehension. Each downstream task has a training set consisting of context-target sample pairs:$\mathcal{Z}=\{(x_i,y_i){\}}_{i=1,\dots ,N}$, among which $x_i$sum $y_i$They are all token sequences. For example, for a summary task, $x_i$is the article content, $y_i$is the summary.

​ In the process of complete fine-tuning, the model uses the pre-trained weight ${Phi}_0$to initialize the model, and then update the parameters Φ 0 + Δ Φ \Phi_0+\Delta\Phi by maximizing the conditional language model${Phi}_0+\Delta \Phi$：
$$\underset{Phi}{\max }\sum _{(x,y)\in \mathcal{Z}}\sum _{t=1}^{|y|}\log (P_{}(y_t|x,y_{<t}))\text{\hspace{1em}( 1 )}$$
The main disadvantage of full fine-tuning: For each downstream task, a different parameter update$\Delta \Phi$ , where dimension$|\Delta \Phi |=|{\Phi }_0|$ . Therefore, if the pretrained model is large, it is very challenging to store and deploy many independent fine-tuned model instances.

In order to be more parameter efficient, LoRA uses a relatively very small parameter $\Theta$ to represent the task-related parameter increment$D.F=\Delta \Phi \left(\Theta \right)$，使用$|\Theta |\ll |{\Phi }_0|$ . Find$The\; task\; of\; \Delta \Phi$ becomes to$Default$ max
$$\underset{Th}{\max }\sum _{(x,y)\in \mathcal{Z}}\sum _{t=1}^{|y|}\log (p_{{Phi}_0+\Delta \Phi (\Theta }(y_t|x,y_{<t}))\text{\hspace{1em}( 2 )}$$
LoRA will use a low-rank representation to encode$\Delta \Phi$ , achieving both computational efficiency and storage efficiency. When the pre-training model is 175B GPT-3, the trainable parameters$|\Theta |$ can be as small as$|{\Phi }_0|0.01$ % of$0.01\%$。

4. LoRA

Typically, a neural network will contain many dense layers that perform matrix multiplication, and these layers are usually full rank. Adgajanyan et al.The research by et al. shows that pre-trained language models have low "intrinsic dimensionality". Inspired by this work, weight updates should also have low "intrinsic rank" during model adaptation to downstream tasks. For pre-training weight matrix $W_0\in {\mathbb{R}}^{d\times k}$ , its update can be represented by low-rank decomposition$W_0+\Delta W=W_0+BA$，$B\in R^{d\times r},A\in R^{r\times k}$且秩$r\ll \min (d,k)$ . During training,$W_0$is frozen and does not accept gradient updates, $A$ and$B$ is a trainable parameter. Note that$W_0$Sum $\Delta W=Both\; Band\; A$ will be multiplied by the same input. For$h=W_{0}x$ , forward propagation becomes:
$$h=W_{0}x+\Delta Wx=W_{0}x+BAx\text{\hspace{1em}(3)}$$

pair matrix $A$ is initialized with a random Gaussian, and the matrix$B$ is initialized with 0, so$\Delta W=BA$ is 0 at the beginning of training. Use$\frac{}{r}$to scale $\Delta Wx$ , where$\alpha$ is less thanrrconstant for$r\; .$When using Adam optimization, after proper scaling initialization, tuning$\alpha$ is about the same as tuning the learning rate.

​ When deploying, explicitly calculate and store $W=W_0+BA$ , and perform inference normally. $W_0$Japanese $BA$ is${\mathbb{R}}^{d\times k}$ . When it is necessary to switch to another downstream task, it can be done by subtracting$BA$ to restore$W_0$, and then add different $B^{\prime }{A}^{\text{'}}$ . Crucially, this is guaranteed not to introduce any additional inference latency.

5. LoRA applied to Transformer

In theory, LoRA can be applied to the weight matrix of any neural network, thereby reducing the number of trainable parameters. The self-attention module in the Transformer architecture has 4 weight matrices: $W_q,W_k,W_v,W_o$, and the weight matrices of the two MLP models. W $W_q$(or $W_k,W_v$) as a dimension $d_{model}\times d_{model}$of a single matrix. For simplicity and parameter efficiency, this study is limited to adapting attention weights for downstream tasks and freezing the MLP module.

​Advantages . The most notable advantage is the reduction in video memory and storage space. For a large Transformer trained with Adam, if $r\ll d_{model}$, VRAM usage is reduced by 2/3 since there is no need to store the optimizer state for frozen parameters. For GPT-3 175B, the memory consumption during training is reduced from 1.2TB to 350GB. when $r=4$ And when only adjusting the query matrix and value matrix, the checkpoint size is reduced by 10000 times (from 350GB to 35MB). Another advantage is that tasks can be switched at deployment time at a lower cost, only by exchanging LoRA weights. Furthermore, compared to full fine-tuning, GPT-3 175B trains 25% faster because gradients for the vast majority of parameters do not need to be computed.

2. Code: Implement BLOOM-LoRA

This section shows how to use LoRA to fine-tune the large language model bloom.

• Data : Use the 1 million instructions provided by BELLE to fine-tune the data;

• Model : Use bloomz-7b1-mt, this version of bloomz is also a model after instruction fine-tuning. See the principle of BLOOM:

  [Natural Language Processing] [Large Model] BLOOM: A multilingual model with 176B parameters and open access

• Dependent packages : use transformers to provide model loading and training; use peft to provide LoRA implementation; use DeepSpeed ​​to provide training assumptions.

​ Note: The peft package is still in the process of rapid iteration, and there may be major changes in subsequent interfaces, and there may also be some bugs. Key dependency package version:

transformers==4.26.1
torch==1.13.1
deepspeed==0.8.2
peft==0.2.0

1. Training code

For brevity, assume the training code is located in train.py.

1.1 Import dependent packages

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 Setting parameters

# 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 Load data

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 Loading the model

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 related

# 转换模型
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 Training parameters

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 Model training

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 ​​configuration file

​ The DeepSpeed ​​configuration file is named 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. start

deepspeed --include=localhost:0,1,2,3 train.py

4. Reasoning

​ The inference file name is 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]))

References

https://arxiv.org/pdf/2106.09685.pdf

https://zhuanlan.zhihu.com/p/615235322

https://github.com/tloen/alpaca-lora/blob/main/finetune.py

https://github.com/huggingface/peft/blob/main/examples/conditional_generation/peft_lora_seq2seq_accelerate_ds_zero3_offload.py