使用GaLore在本地GPU進行高效的LLM調優

作者：佚名 2024-03-25 14:22:07

訓練大型語言模型(llm)，即使是那些“只有”70億個參數的模型，也是一項計算密集型的任務。這種水平的訓練需要的資源超出了大多數個人愛好者的能力范圍。為了彌補這一差距，出現了低秩適應(LoRA)等參數高效方法，可以在消費級gpu上對大量模型進行微調。

GaLore是一種新的方法，它不是通過直接減少參數的數量，而是通過優化這些參數的訓練方式來降低VRAM需求，也就是說GaLore是一種新的模型訓練策略，可讓模型使用全部參數進行學習，并且比LoRA更省內存。

GaLore將這些梯度投影到低秩空間上，顯著減少了計算負荷，同時保留了訓練所需的基本信息。與傳統的優化器在反向傳播后同時更新所有層的方法不同，GaLore在反向傳播期間實現逐層更新。這種方法進一步減少了整個訓練過程中的內存占用。

就像LoRA一樣，GaLore可以讓我們在具有24 GB VRAM的消費級GPU上微調7B模型。結果模型的性能與全參數微調相當，并且似乎優于LoRA。

優于目前Hugging Face還沒有官方代碼，我們就來手動使用論文的代碼進行訓練，并與LoRA進行對比

安裝依賴

首先就要安裝GaLore

pip install galore-torch

然后我們還要一下這些庫，并且請注意版本

datasets==2.18.0
 transformers==4.39.1
 trl==0.8.1
 accelerate==0.28.0
 torch==2.2.1

調度器和優化器的類

Galore分層優化器是通過模型權重掛鉤激活的。由于我們使用Hugging Face Trainer，還需要自己實現一個優化器和調度器的抽象類。這些類的結構不執行任何操作。

from typing import Optional
 import torch
 
 # Approach taken from Hugging Face transformers https://github.com/huggingface/transformers/blob/main/src/transformers/optimization.py
 class LayerWiseDummyOptimizer(torch.optim.Optimizer):
    def __init__(self, optimizer_dict=None, *args, **kwargs):
        dummy_tensor = torch.randn(1, 1)
        self.optimizer_dict = optimizer_dict
        super().__init__([dummy_tensor], {"lr": 1e-03})
 
    def zero_grad(self, set_to_none: bool = True) -> None: 
      pass
 
    def step(self, closure=None) -> Optional[float]: 
      pass
 
 class LayerWiseDummyScheduler(torch.optim.lr_scheduler.LRScheduler):
    def __init__(self, *args, **kwargs):
        optimizer = LayerWiseDummyOptimizer()
        last_epoch = -1
        verbose = False
        super().__init__(optimizer, last_epoch, verbose)
 
    def get_lr(self): 
      return [group["lr"] for group in self.optimizer.param_groups]
 
    def _get_closed_form_lr(self): 
      return self.base_lrs

加載GaLore優化器

GaLore優化器的目標是特定的參數，主要是那些在線性層中以attn或mlp命名的參數。通過系統地將函數與這些目標參數掛鉤，GaLore 8位優化器就會開始工作。

from transformers import get_constant_schedule
 from functools import partial
 import torch.nn
 import bitsandbytes as bnb
 
 from galore_torch import GaLoreAdamW8bit
         
 def load_galore_optimizer(model, lr, galore_config):    
    # function to hook optimizer and scheduler to a given parameter 
    def optimizer_hook(p, optimizer, scheduler):
        if p.grad is not None: 
            optimizer.step()
            optimizer.zero_grad()
            scheduler.step()
 
    # Parameters to optimize with Galore
    galore_params = [
        (module.weight, module_name) for module_name, module in model.named_modules() 
        if isinstance(module, nn.Linear) and any(target_key in module_name for target_key in galore_config["target_modules_list"])
    ] 
    id_galore_params = {id(p) for p, _ in galore_params}
     
    # Hook Galore optim to all target params, Adam8bit to all others
    for p in model.parameters():
        if p.requires_grad:
            if id(p) in id_galore_params:
                optimizer = GaLoreAdamW8bit([dict(params=[p], **galore_config)], lr=lr)
            else:
                optimizer = bnb.optim.Adam8bit([p], lr = lr)
            scheduler = get_constant_schedule(optimizer)
             
            p.register_post_accumulate_grad_hook(partial(optimizer_hook, optimizer=optimizer, scheduler=scheduler))
             
    # return dummies, stepping is done with hooks 
    return LayerWiseDummyOptimizer(), LayerWiseDummyScheduler()

HF Trainer

準備好優化器后，我們開始使用Trainer進行訓練。下面是一個簡單的例子，使用TRL的SFTTrainer (Trainer的子類)在Open Assistant數據集上微調llama2-7b，并在RTX 3090/4090等24 GB VRAM GPU上運行。

from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, set_seed, get_constant_schedule
 from trl import SFTTrainer, setup_chat_format, DataCollatorForCompletionOnlyLM
 from datasets import load_dataset
 import torch, torch.nn as nn, uuid, wandb
 
 lr = 1e-5
 
 # GaLore optimizer hyperparameters
 galore_config = dict(
    target_modules_list = ["attn", "mlp"], 
    rank = 1024, 
    update_proj_gap = 200, 
    scale = 2, 
    proj_type="std"
 )
 
 modelpath = "meta-llama/Llama-2-7b"
 model = AutoModelForCausalLM.from_pretrained(
    modelpath,    
    torch_dtype=torch.bfloat16,
    attn_implementation = "flash_attention_2",  
    device_map = "auto",
    use_cache = False,
 )
 tokenizer = AutoTokenizer.from_pretrained(modelpath, use_fast = False)
 
 # Setup for ChatML
 model, tokenizer = setup_chat_format(model, tokenizer)
 if tokenizer.pad_token in [None, tokenizer.eos_token]: 
    tokenizer.pad_token = tokenizer.unk_token
 
 # subset of the Open Assistant 2 dataset, 4000 of the top ranking conversations
 dataset = load_dataset("g-ronimo/oasst2_top4k_en")
 
 training_arguments = TrainingArguments(
    output_dir = f"out_{run_id}",
    evaluation_strategy = "steps",
    label_names = ["labels"],
    per_device_train_batch_size = 16,
    gradient_accumulation_steps = 1,
    save_steps = 250,
    eval_steps = 250,
    logging_steps = 1, 
    learning_rate = lr,
    num_train_epochs = 3,
    lr_scheduler_type = "constant",
    gradient_checkpointing = True,
    group_by_length = False,
 )
 
 optimizers = load_galore_optimizer(model, lr, galore_config)
 
 trainer = SFTTrainer(
    model = model,
    tokenizer = tokenizer,
    train_dataset = dataset["train"],
    eval_dataset = dataset['test'],
    data_collator = DataCollatorForCompletionOnlyLM(
        instruction_template = "<|im_start|>user", 
        response_template = "<|im_start|>assistant", 
        tokenizer = tokenizer, 
        mlm = False),
    max_seq_length = 256,
    dataset_kwargs = dict(add_special_tokens = False),
    optimizers = optimizers,
    args = training_arguments,
 )
 
 trainer.train()

GaLore優化器帶有一些需要設置的超參數如下:

target_modules_list:指定GaLore針對的層

rank:投影矩陣的秩。與LoRA類似，秩越高，微調就越接近全參數微調。GaLore的作者建議7B使用1024

update_proj_gap:更新投影的步驟數。這是一個昂貴的步驟，對于7B來說大約需要15分鐘。定義更新投影的間隔，建議范圍在50到1000步之間。

scale:類似于LoRA的alpha的比例因子，用于調整更新強度。在嘗試了幾個值之后，我發現scale=2最接近于經典的全參數微調。