在16G的GPU上微调Mixtral-8x7B

news2024/9/20 8:50:06

Mixtral-8x7B是最好的开源llm之一。但是消费级硬件上对其进行微调也是非常具有挑战性的。因为模型需要96.8 GB内存。而微调则需要更多的内存来存储状态和训练数据。比如说80gb RAM的H100 GPU是不够的。

这时我们就想到了QLoRA，它将模型大小除以4，同时通过仅调整LoRA适配器进行微调来减小优化器状态的大小。但是即使使用QLoRA，然需要32 GB的GPU内存来微调Mixtral-8x7B。

如果我们可以将Mixtral-8x7B量化到更低的精度呢?

例如我们可以用AQLM将Mixtral-8x7B量化为2位，同时最小化模型性能的下降。

在本文中，我将展示如何仅使用16 GB的GPU RAM对使用AQLM进行量化的Mixtral-8x7B进行微调。我还讨论了如何优化微调超参数，以进一步减少内存消耗，同时保持良好的性能。对2位的Mixtral进行微调是很快的，并且可能产生比QLoRA更好的模型，同时使用的内存减少了两倍。

用AQLM对2位LLM量化进行微调

AQLM的作者已经在Hugging Face发布了量化版本的Mixtral-8x7B:

ISTA-DASLab/Mixtral-8x7b-AQLM-2Bit-1x16-hf

我们使用这个模型进行微调。它只占用12.6 GB的内存。

首先安装AQLM和最新版本的PEFT、Transformers、Accelerate和TRL。我还安装了bitsandbytes来使用分页AdamW 8位，因为这将进一步减少内存消耗。

 pip install git+https://github.com/huggingface/transformers.git
 pip install git+https://github.com/huggingface/peft
 pip install git+https://github.com/huggingface/trl.git
 pip install git+https://github.com/huggingface/accelerate.git
 pip install aqlm[gpu,cpu]

然后引用库

 from datasets import load_dataset
 from peft import LoraConfig, prepare_model_for_kbit_training
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
     TrainingArguments,
 )
 from trl import SFTTrainer

加载模型

 model = AutoModelForCausalLM.from_pretrained(
     "ISTA-DASLab/Mixtral-8x7b-AQLM-2Bit-1x16-hf",
     trust_remote_code=True, torch_dtype="auto", device_map="cuda", low_cpu_mem_usage=True
 )

tokenizer

 tokenizer = AutoTokenizer.from_pretrained("ISTA-DASLab/Mixtral-8x7b-AQLM-2Bit-1x16-hf")
 tokenizer.pad_token = tokenizer.eos_token

这里使用了prepare_model_for_kbit_training为微调准备模型，它与AQLM一起工作得很好。

 model = prepare_model_for_kbit_training(model)

通过微调，我们将把Mixtral-8x7B变成一个指示/聊天模型。我使用了timdettmers/openassistant-guanaco，因为它小而好，适合本教程。由于它有一个名为“text”的列，因此它的格式也可以使用TRL进行微调。

 dataset = load_dataset("timdettmers/openassistant-guanaco")

LoraConfig定义如下:

 peft_config = LoraConfig(
         lora_alpha=16,
         lora_dropout=0.05,
         r=16,
         bias="none",
         task_type="CAUSAL_LM",
         target_modules= ['k_proj', 'q_proj', 'v_proj', 'o_proj', "gate", "w1", "w2", "w3"]
 )

我们的配置会产生一个242M参数的巨大适配器。这是因为我的目标是模块w1、w2和w3。这些模块都在Mixtral的8位专家内部，总共24个模块。其他模块(自注意力和门控模块)由8位专家共享。它们对可训练参数总数的影响可以忽略不计。

对于训练超参数，我选择如下:

 training_arguments = TrainingArguments(
         output_dir="./fine-tuned_MixtralAQLM_2bit/",
         evaluation_strategy="steps",
         do_eval=True,
         optim="paged_adamw_8bit",
         per_device_train_batch_size=4,
         gradient_accumulation_steps=4,
         per_device_eval_batch_size=4,
         log_level="debug",
         logging_steps=25,
         learning_rate=1e-4,
         eval_steps=25,
         save_strategy='steps',
         max_steps=100,
         warmup_steps=25,
         lr_scheduler_type="linear",
 )

这里重要的超参数是:

per_device_train_batch_size和gradient_accumulation_steps:我将它们设置为4和4，这将产生总批大小为16(4*4)。

learning_rate:我把它设置为1e-4，因为它看起来很好用。这绝对不是最好的值。

lr_scheduler_type:我将其设置为“linear”。

optim: paged_adamw_8bit性能良好，同时比原始AdamW实现消耗的内存少得多。缺点是它减慢了微调速度，特别是如果你有一个旧的CPU。如果您有足够的内存，比如24G，可以将其替换为“adamw_torch”或“adamw_8bit”。

我们将max_steps设置为100。来证实模型在学习，但时这还不足以训练出一个好的模型。因为至少训练1到2个epoch，才能获得一个相当好的适配器。

使用TRL的SFTTrainer进行培训:

 trainer = SFTTrainer(
         model=model,
         train_dataset=dataset['train'],
         eval_dataset=dataset['test'],
         peft_config=peft_config,
         dataset_text_field="text",
         max_seq_length=256,
         tokenizer=tokenizer,
         args=training_arguments,
 )
 trainer.train()

根据任务，可以增加“max_seq_length”，我将其设置为256。这是一个较低的值，但它有助于消耗更少的内存。

根据训练和验证损失，微调进展顺利: