在16G的GPU上微调Mixtral-8x7B（附代码）

最新推荐文章于 2024-08-26 08:43:50 发布

数据派THU

最新推荐文章于 2024-08-26 08:43:50 发布

阅读量227

点赞数

原文链接：https://mp.weixin.qq.com/s?__biz=MzI1MjQ2OTQ3Ng==&mid=2247633402&idx=2&sn=1d2952609efe41ae8fa74396de5ef4cf&chksm=e8f7bdaeea224fc4b035ce5f83e7958fad46ea7597b5a186df44a7918c21128826a2ad3bc4f7&scene=126&sessionid=0

版权

来源：DeepHub IMBA
本文约1700字，建议阅读5分钟
本文展示了如何仅使用16 GB的GPU RAM对使用AQLM进行量化的Mixtral-8x7B进行微调。

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。这是一个较低的值，但它有助于消耗更少的内存。

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