大模型微调实战：使用INT8/FP4/NF4微调大模型

bitsandbytes 是自定义 CUDA 函数的轻量级包装器，特别是 8 比特优化器、矩阵乘法和量化函数。具有混合精度分解的 8 比特矩阵乘法LLM.int8() 推理8 比特优化器：Adam、AdamW、RMSProp、LARS、LAMB、Lion（节省 75% 内存）稳定的嵌入层：通过更好的初始化和标准化提高稳定性8 比特量化：分位数、线性和动态量化快速的分位数估计：比其他算法快 100

爱喝白开水a

1360人浏览 · 2025-01-26 09:30:00

爱喝白开水a · 2025-01-26 09:30:00 发布

随着，ChatGPT 迅速爆火，引发了大模型的时代变革。然而对于普通大众来说，进行大模型的预训练或者全量微调遥不可及。由此，催生了各种参数高效微调技术，让科研人员或者普通开发者有机会尝试微调大模型。

因此，该技术值得我们进行深入分析其背后的机理，之前分享了大模型参数高效微调技术原理综述的文章。下面给大家分享大模型参数高效微调技术实战系列文章，相关代码均放置在GitHub：llm-action。

本文为大模型参数高效微调技术实战的第八篇。本文将结合 bitsandbytes（使用 INT8 量化来加载大模型）和 LoRA 技术来微调Bloom大模型。

量化降低了浮点数据类型的精度，减少了存储模型权重所需的内存。因此，量化会降低推理性能，因为降低精度时会丢失信息。INT8量化只使用四分之一的精度，但它不会降低训练性能，因为它不仅丢弃位或数据。相反，INT8量化从一种数据类型舍入到另一种。

数据集和模型准备

本文使用english_quotes数据集，该数据集可用于多标签文本分类和文本生成。数据集下载地址：huggingface.co/datasets/Ab…

模型下载地址：huggingface.co/bigscience/…

bitsandbytes 简介

bitsandbytes 是自定义 CUDA 函数的轻量级包装器，特别是 8 比特优化器、矩阵乘法和量化函数。主要特征如下：

具有混合精度分解的 8 比特矩阵乘法
LLM.int8() 推理
8 比特优化器：Adam、AdamW、RMSProp、LARS、LAMB、Lion（节省 75% 内存）
稳定的嵌入层：通过更好的初始化和标准化提高稳定性
8 比特量化：分位数、线性和动态量化
快速的分位数估计：比其他算法快 100 倍

目前，transformers 库已经集成并原生支持了 bitsandbytes 这个量化库。而且bitsandbytes 是量化任何模型的最简单方法之一，因为它不需要量化校准数据及校准过程 (即零样本量化)。任何模型只要含有 torch.nn.Linear 模块，就可以对其进行开箱即用的量化。每当在 transformers 库中添加新架构时，只要其可以用 accelerate 库的 device_map="auto" 加载，用户就可以直接受益于开箱即用的 bitsandbytes 量化，同时该方法对性能的影响也是最小的。量化是在模型加载时执行的，无需运行任何后处理或准备步骤。

由于量化模型的唯一条件是包含 torch.nn.Linear 层，因此量化对于任何模态都可以实现开箱即用。用户可以开箱即用地加载诸如 Whisper、ViT、Blip2 之类的 8 比特或 4 比特(FP4/NF4)模型。

如果你在量化基础模型之上使用PEFT库基于Lora进行训练，则可以将训练得到的Apapter合并在基础模型之上进行部署，而不会降低推理性能。你甚至还可以在反量化模型之上合并 Apapter！

下面是使用 NF4 量化加载 4 比特模型的示例：

from transformers import BitsAndBytesConfig

nf4_config = BitsAndBytesConfig(
   load_in_4bit=True,
   bnb_4bit_quant_type="nf4",
   bnb_4bit_use_double_quant=True,
   bnb_4bit_compute_dtype=torch.bfloat16
)

model_nf4 = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=nf4_config)

下面是使用 FP4 量化加载 4 比特模型的示例：

import torch
from transformers import BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(
   load_in_4bit=True,
   bnb_4bit_compute_dtype=torch.bfloat16
)

LoRA 简介

LoRA方法的核心思想就是通过低秩分解来模拟参数的改变量，从而以极小的参数量来实现大模型的间接训练。

模型训练及推理

为了不影响阅读体验，详细的微调代码放置在GitHub：llm-action 项目中 finetune_bloom_bnb_peft.ipynb文件，这里仅列出关键步骤。

第一步，加载模型及Tokenizer。

import os

import torch
import torch.nn as nn
import bitsandbytes as bnb
from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("/workspace/model/bloomz-3b", load_in_8bit=True)
tokenizer = AutoTokenizer.from_pretrained("/workspace/model/bloomz-3b")

注意：这里演示的是INT8，如果希望使用NF4或者FP4，在模型加载时修改quantization_config参数即可，具体可以参考PEFT官方示例：inetune_fp4_opt_bnb_peft.py。

NF4+QLoRA：

model = AutoModelForCausalLM.from_pretrained(
    model_name_or_path='/name/or/path/to/your/model',
    load_in_4bit=True,
    device_map='auto',
    max_memory=max_memory,
    torch_dtype=torch.bfloat16,
    quantization_config=BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type='nf4'
    ),
)

第二步，训练模型做准备。在使用 peft 训练 int8 模型之前，需要完成一些预处理，需导入一个实用函数prepare_model_for_int8_training，它将完成如下操作：

将所有非 int8 模块转换为全精度 (fp32) 以确保稳定性
将forward_hook添加到输入嵌入层以启用输入隐藏状态的梯度计算
启用梯度检查点（gradient checkpointing）以提高训练的内存效率

from peft import prepare_model_for_int8_training
model = prepare_model_for_int8_training(model)

第三步，创建 LoRA 微调方法对应的配置；同时，通过调用 get_peft_model 方法包装基础的 Transformer 模型。

from peft import LoraConfig, get_peft_model

config = LoraConfig(
    r=16, lora_alpha=32, target_modules=["query_key_value"], lora_dropout=0.05, bias="none", task_type="CAUSAL_LM"
)

model = get_peft_model(model, config)
model = model.to(device)
print_trainable_parameters(model)

第四步，进行模型微调。

from transformers import Seq2SeqTrainer, TrainerCallback, TrainingArguments, TrainerState, TrainerControl
from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR

# 回调保存Peft模型
class SavePeftModelCallback(TrainerCallback):
    def on_save(
        self,
        args: TrainingArguments,
        state: TrainerState,
        control: TrainerControl,
        **kwargs,
    ):
        checkpoint_folder = os.path.join(args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}")
        print("checkpoint folder: ",checkpoint_folder)
        peft_model_path = os.path.join(checkpoint_folder, "adapter_model")
        kwargs["model"].save_pretrained(peft_model_path)

        
        files = os.listdir(checkpoint_folder)
        print("checkpoint folder list: ", files)
        adapter_files = os.listdir(peft_model_path)
        print("checkpoint adapter folder list: ", adapter_files)
        
        pytorch_model_path = os.path.join(checkpoint_folder, "pytorch_model.bin")
        if os.path.exists(pytorch_model_path):
            os.remove(pytorch_model_path)
        return control

args = transformers.TrainingArguments(
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        warmup_steps=5,
        max_steps=20,
        learning_rate=2e-4,
        fp16=True,
        logging_steps=1,
        output_dir="outputs",
        save_strategy = 'steps',
        save_steps = 10
    )

trainer = transformers.Trainer(
    model=model,
    train_dataset=data["train"],
    args=args,
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
    callbacks=[SavePeftModelCallback()],
)

model.config.use_cache = False  # silence the warnings. Please re-enable for inference!

trainer.train()

第五步，重新加载训练的Apapter进行推理。

import torch
from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM, AutoTokenizer

peft_model_id = "outputs/checkpoint-20/"
config = PeftConfig.from_pretrained(peft_model_id)
model = AutoModelForCausalLM.from_pretrained(
    config.base_model_name_or_path, return_dict=True, load_in_8bit=True, device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)

# Load the Lora model
model = PeftModel.from_pretrained(model, peft_model_id)

batch = tokenizer("Two things are infinite: ", return_tensors="pt")

# 使用 Pytorch 的 autocast 运行推理
# 以 autocast 选择的特定op数据类型运行ops，以提高性能，同时保持准确性。
with torch.cuda.amp.autocast():
    output_tokens = model.generate(**batch, max_new_tokens=50)

print("output：\n\n", tokenizer.decode(output_tokens[0], skip_special_tokens=True))