在单张 RTX 5090 上用 LoRA 微调 Llama 3.1 8B,一个下午即可完成

数据集格式化、能装进 32 GB 显存的 QLoRA 设置、一个 epoch 实际要多久,以及在上线 adapter 之前如何评估。

CGCheapServ GPU 团队作者 9 分钟阅读
一块 GPU 显卡,上方是一个形似大脑的小型神经网络
本页目录8
  1. 32 GB 显存能装下什么
  2. 环境
  3. 数据:决定成败的环节
  4. 训练脚本
  5. 需要多长时间
  6. 合并之前先评估
  7. 合并与部署
  8. 成本

过去做微调,需要一个插着八块卡的节点。如今,对于 8B 模型的大多数实际适配任务,一块拥有 32 GB 显存的 RTX 5090 就足够了,整个任务一个下午就能做完。本文的实操流程以 Llama 3.1 8B Instruct、一份包含 10,000 条样本的指令数据集,以及一套每个 epoch 约需 40 分钟的 QLoRA 配置为例,然后介绍如何在合并并部署 adapter 之前,先对它做评估。

32 GB 显存能装下什么

用 BF16 对 8B 参数的模型做全量微调,需要存放权重、梯度和优化器状态:大致是 16 + 16 + 64 GB,根本放不下。LoRA 训练的是小型的 adapter 矩阵,而不是权重本身;QLoRA 则进一步把冻结的基座模型以 4-bit 存放。这样一来,5090 上的显存占用情况如下:

项目QLoRA 4-bit,8B,4k 上下文,batch 4
基座权重(NF4)~5.5 GB
LoRA 参数 + 优化器(rank 32,所有线性层)~0.8 GB
启用梯度检查点(gradient checkpointing)后的激活值~9 GB
CUDA 上下文、缓存、显存碎片~3 GB
合计~19 GB

这样还有余量,可以调大 batch size,或把上下文提高到 8k。用 BF16 的普通 LoRA(不量化)也放得下,约占 26 GB,训练速度大约快 30%;如果您的每条样本都不超过 4k token,就选它。

环境

可以从 LLaMA-Factory 或 Unsloth 模板开始,也可以在 Ubuntu 24.04 CUDA 12.8 基础镜像上自行搭建:

uv venv /opt/ft && source /opt/ft/bin/activate
uv pip install torch==2.5.1 --index-url https://download.pytorch.org/whl/cu124
uv pip install "transformers==4.46.*" "peft==0.13.*" "trl==0.12.*" bitsandbytes datasets accelerate
huggingface-cli download meta-llama/Llama-3.1-8B-Instruct --local-dir /data/models/llama-8b

数据:决定成败的环节

大多数失败的微调,都是数据出了问题。请把每条样本都格式化成模型已经熟悉的对话模板;如果您在生产环境中使用 system 消息,这里也要带上;并保持回答风格一致:模型学到的格式,和它学到的内容一样多。

{"messages": [
  {"role": "system", "content": "You are a support assistant for Acme Cloud."},
  {"role": "user", "content": "How do I rotate my API key?"},
  {"role": "assistant", "content": "Open Account → API keys, click Rotate next to the key…"}
]}

一万条这种形态的样本,足以完成风格和领域的适配。训练前先留出 500 条用于评估,调超参数时绝不要去看它们。对近乎相同的提示词做去重;重复样本会让 loss 曲线变得很漂亮,模型却变得更差。

训练脚本

from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import LoraConfig
from trl import SFTTrainer, SFTConfig
import torch

base = "/data/models/llama-8b"
tok = AutoTokenizer.from_pretrained(base)
model = AutoModelForCausalLM.from_pretrained(
    base, torch_dtype=torch.bfloat16,
    quantization_config=BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                                           bnb_4bit_compute_dtype=torch.bfloat16))
ds = load_dataset("json", data_files={"train": "train.jsonl", "eval": "eval.jsonl"})
cfg = SFTConfig(output_dir="/data/runs/support-v1", num_train_epochs=2,
    per_device_train_batch_size=4, gradient_accumulation_steps=4, learning_rate=2e-4,
    lr_scheduler_type="cosine", warmup_ratio=0.03, bf16=True, gradient_checkpointing=True,
    max_seq_length=4096, logging_steps=10, eval_strategy="steps", eval_steps=100, save_steps=100)
lora = LoraConfig(r=32, lora_alpha=64, lora_dropout=0.05, task_type="CAUSAL_LM",
    target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"])
SFTTrainer(model=model, args=cfg, peft_config=lora, train_dataset=ds["train"],
           eval_dataset=ds["eval"], tokenizer=tok).train()

值得解释的几个选择:在我们的测试中,对所有线性层使用 rank 32,是 8B 模型的最佳平衡点;学习率 2e-4 对 LoRA 来说是标准值,对全量微调则会高得离谱;而有效 batch 为 16、配合两个 epoch 的余弦衰减,通常正是这种规模的数据集上 eval loss 触底的位置。

需要多长时间

配置Tokens/s10k 条样本 × 2 个 epoch(平均 600 tokens)
RTX 5090,QLoRA 4-bit,梯度检查点~5,000每个 epoch 约 40 分钟,共 80 分钟
RTX 5090,LoRA BF16,梯度检查点~6,500每个 epoch 约 31 分钟
RTX 4090,QLoRA 4-bit~3,100每个 epoch 约 65 分钟
H100 SXM,LoRA BF16,不使用梯度检查点~14,000每个 epoch 约 14 分钟

在最初的几分钟里盯着 nvidia-smi:SM 利用率应当保持在 90% 以上。如果不是,那瓶颈就在 dataloader;请预先对数据集做 tokenize,或者增加 worker 的数量。

合并之前先评估

eval loss 下降,只能证明模型记住了您的格式,并不能证明它变得更好。请把留出的那些提示词分别交给基座模型和加载了 adapter 的模型,对比两边的回答,最好有一套评分标准(rubric),并让另一个模型当裁判,或者由人工抽查 50 条样本。有三种失败模式只会在这里暴露,其他地方都看不出来:

  • 对措辞过拟合:每个回答的开头都一模一样。把 epoch 降到 1,或者降低 rank。
  • 遗忘:模型在原本能处理的通用问题上变差了。在训练集中混入 10% 到 20% 的通用指令数据。
  • 拒答能力丢失:如果您的数据里从未出现过模型拒绝回答的情形,它就不再拒绝。请加入体现您希望保留的那种行为的样本。

合并与部署

from peft import PeftModel
m = AutoModelForCausalLM.from_pretrained(base, torch_dtype=torch.bfloat16)
m = PeftModel.from_pretrained(m, "/data/runs/support-v1/checkpoint-1250").merge_and_unload()
m.save_pretrained("/data/models/support-v1"); tok.save_pretrained("/data/models/support-v1")

合并后的模型用 vLLM 部署,方式与基座模型完全相同;vLLM 指南原样适用,同一块 5090 上的 BF16 8B 模型,单个请求的生成速度约为每秒 90 个 token。也请保留 adapter 的 checkpoint:它只有 300 MB,日后可以让您把它重新合并到更新的基座模型上。

成本

在 RTX 5090 套餐上花一整个下午,包括下载,以及调参过程中的三次训练,折合下来只占月费中的几美元。同样的任务放在按小时计费的云 GPU 上,跑一次的成本相当,但真正的微调离不开反复迭代,这部分费用要高得多;一块按月租用、可以把模型一直留在上面的显卡,才让第二次、第三次尝试变得便宜。

CG
CheapServ GPU 团队

GPU 服务器集群的基准测试、镜像和模板。

一分钟内部署您的第一台服务器。

以 BTC、ETH、XMR 或 USDT 充值,$25 起。余额永不过期,未使用的部分可退款。

立即注册