使用GaLore在本地GPU進行高效的LLM調優
訓練大型語言模型(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最接近于經典的全參數微調。
微調效果對比
給定超參數的訓練損失與全參數調優的軌跡非常相似，表明GaLore分層方法確實是等效的。
用GaLore訓練的模型得分與全參數微調非常相似。
GaLore可以節省大約15 GB的VRAM，但由于定期投影更新，它需要更長的訓練時間。
上圖為2個3090的內存占用對比
訓練事件對比，微調:~58分鐘。GaLore:約130分鐘
最后我們再看看GaLore和LoRA的對比
上圖為LoRA微調所有線性層，rank64,alpha 16的損失圖
從數值上可以看到GaLore是一種近似全參數訓練的新方法，性能與微調相當，比LoRA要好得多。
總結
GaLore可以節省VRAM，允許在消費級GPU上訓練7B模型，但是速度較慢，比微調和LoRA的時間要長差不多兩倍的時間。
