AI全量微调折腾手记

我们的数据主要有三类:

  1. 代码-注释对:从内部仓库提取的代码文件和对应注释
  2. 问题-解答对:团队内部的 Q&A 文档
  3. 代码评审记录:PR review 里的问题和建议

如果用 7B 或者更小的模型,可能连基础编程能力都不稳;要是上 70B,训练成本就吃不消了。

背景和需求

先说清楚场景:我们要训练一个代码生成模型,要求:

  1. 能熟练使用公司内部的工具库和框架
  2. 生成代码符合团队的编码规范
  3. 对常见 Bug 有预判,能主动规避
  4. 训练完成后,推理速度不能太慢

模型基础选了 Llama-3.1-8B,原因很简单:参数量够大,能力够强,而且社区资源丰富。如果用 7B 或者更小的模型,可能连基础编程能力都不稳;要是上 70B,训练成本就吃不消了。

硬件环境:

  • 4× A100 80GB(总共 320GB 显存)
  • Ubuntu 22.04
  • CUDA 12.1
  • PyTorch 2.3.1

这个配置不是充钱充的,而是算出来的。全量微调 Llama-3.1-8B,用 AdamW 优化器、混合精度训练(BF16),每个参数的优化器状态大概需要 8 字节。8B 参数的模型,光是优化器状态就要 64GB;再加上模型权重、梯度、激活值,单张 80GB 显存卡勉强能装,但训练起来会频繁 OOM(Out of Memory)。

所以至少需要 4 卡并行,才能有足够显存同时装下训练状态和维持合理的 batch size。

实现过程

数据准备

数据是微调的核心,但也是最容易踩坑的地方。

我们的数据主要有三类:

  1. 代码-注释对:从内部仓库提取的代码文件和对应注释
  2. 问题-解答对:团队内部的 Q&A 文档
  3. 代码评审记录:PR review 里的问题和建议

原始数据清洗是个大头活儿:

  • 去掉大段重复代码(避免模型背书)
  • 过滤掉低质量注释(比如 “// TODO: fix this”)
  • 样本均衡(不能全是某个项目的代码)
  • 重新格式化成对话格式

清洗后的数据格式类似:

{
  "messages": [
    {"role": "user", "content": "写一个函数,从 API 获取用户信息并缓存"},
    {"role": "assistant", "content": "```python\ndef get_user_info(user_id: int, cache_ttl: int = 300) -> dict:\n    '''获取用户信息并缓存'''\n    # ... 具体实现 ...\n```"}
  ]
}

最终我们准备了约 15万条高质量样本,训练集 12万条,验证集 3万条。

训练脚本和环境

训练框架选了 Hugging Face Transformers + Accelerate + DeepSpeed,这套组合对多卡并行支持比较成熟。

环境安装:

pip install torch==2.3.1 transformers==4.43.0 accelerate==0.33.0 deepspeed==0.14.0
pip install datasets==2.20.0 wandb==0.17.5
pip install flash-attn==2.6.3  # 加速 attention 计算

训练脚本的核心配置:

from transformers import (
    AutoTokenizer, AutoModelForCausalLM,
    TrainingArguments, Trainer,
    DataCollatorForLanguageModeling
)
import torch
from datasets import load_dataset
import os

# 模型和分词器
model_name = "meta-llama/Meta-Llama-3.1-8B"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token  # Llama 没有 pad_token,用 eos_token 代替

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True
)

# 数据加载
train_dataset = load_dataset("json", data_files="train_data.jsonl", split="train")
eval_dataset = load_dataset("json", data_files="eval_data.jsonl", split="train")

def tokenize_function(examples):
    # 将对话格式转换为训练文本
    texts = []
    for messages in examples["messages"]:
        formatted = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=False
        )
        texts.append(formatted)

    tokenized = tokenizer(
        texts,
        truncation=True,
        max_length=2048,
        padding=False,
        return_tensors=None
    )

    tokenized["labels"] = tokenized["input_ids"].copy()
    return tokenized

tokenized_train = train_dataset.map(
    tokenize_function,
    batched=True,
    remove_columns=train_dataset.column_names,
    num_proc=16
)

tokenized_eval = eval_dataset.map(
    tokenize_function,
    batched=True,
    remove_columns=eval_dataset.column_names,
    num_proc=16
)

# 训练参数
training_args = TrainingArguments(
    output_dir="./llama-finetune",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    gradient_accumulation_steps=8,
    learning_rate=2e-5,
    warmup_steps=500,
    logging_steps=10,
    save_steps=500,
    eval_steps=500,
    fp16=False,
    bf16=True,
    gradient_checkpointing=True,
    deepspeed="ds_config.json",  # DeepSpeed 配置
    logging_dir="./logs",
    report_to="wandb",
    run_name="llama-8b-finetune-v1"
)

# 数据整理器(因果 LM 使用自回归)
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train,
    eval_dataset=tokenized_eval,
    data_collator=data_collator,
)

# 开始训练
trainer.train()

DeepSpeed 配置文件 ds_config.json

{
  "train_batch_size": 128,
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 8,
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 2e-5,
      "betas": [0.9, 0.95],
      "eps": 1e-8,
      "weight_decay": 0.01
    }
  },
  "scheduler": {
    "type": "WarmupLR",
    "params": {
      "warmup_min_lr": 0,
      "warmup_max_lr": 2e-5,
      "warmup_num_steps": 500
    }
  },
  "fp16": {
    "enabled": false,
    "loss_scale": 0,
    "initial_scale_power": 16,
    "loss_scale_window": 1000,
    "hysteresis": 2,
    "min_loss_scale": 1
  },
  "bf16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 5e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 5e8,
    "contiguous_gradients": true
  },
  "gradient_clipping": 1.0,
  "activation_checkpointing": {
    "partition_activations": false,
    "cpu_checkpointing": false,
    "contiguous_memory_optimization": false,
    "number_checkpoints": null,
    "synchronize_checkpoint_boundary": false,
    "profile": false
  }
}

训练流程

整个训练流程可以用下面这个图来概括:

graph TD A[原始数据<br/>代码/QA/评审记录] --> B[数据清洗<br/>去重/过滤/格式化] B --> C[数据分词<br/>chat template + truncation] C --> D[模型加载<br/>Llama-3.1-8B + BF16] D --> E[多卡并行<br/>DeepSpeed Zero-2] E --> F[训练循环<br/>3 epochs + warmup] F --> G[评估检查<br/>每500步] G --> H{Loss 是否<br/>稳定下降?} H -->|是| I[继续训练] H -->|否| J[调整超参数<br/>或检查数据] I --> F J --> F G --> K{训练完成?} K -->|否| F K -->|是| L[模型保存<br/>+ 推理测试]

训练过程中的关键监控指标:

  • Training Loss:应该平稳下降,如果突然飙升,可能数据有问题
  • Validation Loss:用来检测过拟合
  • Learning Rate:warmup 后应该保持在设定值
  • Gradient Norm:如果持续超过 1.0,可能需要更强的梯度裁剪
  • GPU 显存使用率:理想是 80-90%,太高会 OOM,太低说明资源浪费

我们用 Weights & Biases(wandb)做可视化监控,每 10 步记录一次指标。

踩坑实录

显存爆炸

第一次训练,batch size 设得太激进,几乎马上就 OOM 了。

RuntimeError: CUDA out of memory. Tried to allocate 2.34 GiB

排查下来,发现三个问题:

  1. per_device_train_batch_size 设置为 8,每张卡吃不消
  2. 没开 gradient_checkpointing,激活值占用太多显存
  3. 数据长度没控制好,有些样本超过 2048 tokens

调整策略:

  • 降低 per_device_train_batch_size 到 4
  • 开启 gradient_checkpointing=True(会慢一点,但省显存)
  • 数据预处理时严格截断到 2048 tokens
  • gradient_accumulation_steps 提高到 8,模拟更大的 batch size

梯度爆炸

训练到 500 步左右,突然 loss 爆炸,从 2.3 飙到 100+。

查了半天,发现是学习率设置不当。我们的数据质量参差不齐,有些样本噪声较大,初始学习率 5e-5 太高了。

解决方案:

  • 降低学习率到 2e-5
  • 增加梯度裁剪强度,从 0.5 改到 1.0
  • 调整 warmup 步数从 100 到 500,让模型更平稳地进入训练
  • 检查数据,去掉明显的异常样本

过拟合

第二个 epoch 开始,validation loss 不降反升,但 training loss 还在下降——典型的过拟合。

分析原因:

  • 数据量相对模型参数量太小(15万 vs 8B 参数)
  • 训练步数过多(3 epochs 对于这个任务来说可能太多了)
  • 数据多样性不足,大量相似的代码模式

解决方法:

  • 增加 dropout,但 Llama 架构本身 dropout 很少,效果有限
  • 引入正则化,调高 weight_decay 到 0.01
  • 早停策略:validation loss 连续 3 次不下降就停止
  • 数据增强:对同一段代码生成不同的注释/解释

生成质量退化

训练完的模型,代码风格是对了,但能力反而变差了——比如不会写复杂的算法,或者胡编不存在的库。

这实际上是"灾难性遗忘":模型过度拟合训练数据,忘掉了预训练时学到的通用能力。

缓解策略:

  • 混合训练:在训练数据中加入一部分通用编程数据(比如 CodeAlpaca、HumanEval)
  • 保留部分原始预训练数据的权重,不要让模型跑偏太多
  • 控制训练强度:降低学习率、减少训练步数
  • 定期在通用任务上评估,及时发现问题

推理速度变慢

全量微调后,模型文件从 15GB 变到了 15GB(没变),但推理速度反而变慢了。

这主要是因为:

  • 训练时的一些配置影响了推理效率
  • 可能需要重新量化和优化

优化方法:

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 加载训练好的模型
model = AutoModelForCausalLM.from_pretrained(
    "./llama-finetune/final",
    torch_dtype=torch.bfloat16,
    device_map="auto",
    low_cpu_mem_usage=True
)

# 启用缓存加速
model.config.use_cache = True

# 量化到 8-bit(可选)
from bitsandbytes import BitsAndBytesConfig
bnb_config = BitsAndBytesConfig(
    load_in_8bit=True,
    llm_int8_threshold=6.0,
    llm_int8_has_fp16_weight=False,
)

model_8bit = AutoModelForCausalLM.from_pretrained(
    "./llama-finetune/final",
    quantization_config=bnb_config,
    device_map="auto"
)

结果和评估

训练完成后,我们在多个维度上评估了模型:

定量评估

指标基础模型微调后模型
Validation Loss1.890.72
HumanEval Pass@145.2%62.8%
代码规范符合率67%94%
推理速度(tokens/s)128124

HumanEval 和规范符合率的大幅提升,是全量微调最值得看的收益;推理速度几乎没变,说明领域定制没有明显拖慢吞吐。

Llama-3.1-8B 全量微调后 HumanEval、代码规范符合率与推理速度对比

Pass@1 和规范符合率都显著改善,而 tokens/s 只小幅下降——说明这次微调主要改的是「写对、写规范」,而不是牺牲推理效率。

定性评估

最重要的还是实际使用体验:

  1. 风格一致性:模型生成的代码完全符合团队规范,变量命名、注释风格都很统一
  2. 领域适配:能正确使用内部工具库,不会胡编 API
  3. 错误预判:对常见 Bug 有意识,比如会主动处理空指针、边界检查
  4. 通用能力保留:依然能解决复杂算法问题,没有完全遗忘

但也有不足:

  1. 偶尔会"背书",生成和训练数据几乎一模一样的代码
  2. 对训练数据里没见过的技术栈,理解能力下降
  3. 推理成本略高(相比 LoRA 微调)

成本分析

算一笔账,方便后来人参考:

  • 硬件成本:4× A100 80GB,云上租约 8-10 美元/小时,训练约 48 小时,总计约 400-500 美元
  • 人力成本:数据准备 2 周,调试训练 1 周,评估优化 1 周,算下来约 4 人周
  • 后续维护:每季度更新一次数据重新训练,每次约 2-3 人天

对比 LoRA 方案:

  • LoRA 训练成本低得多(单卡 40GB 就够),但效果确实差一截
  • 全量微调虽然贵,但在我们的场景下,多花的成本是值得的

结语

回过头看,这次全量微调的折腾算是"贵但值"。

LoRA 虽然便宜,但对于真正需要深度定制、彻底改变模型行为的场景,还是全量微调更稳。当然,前提是你有足够的显存、时间和耐心。

技术选型没有银弹。模型也好、框架也好,终究要回到"解决什么问题"这个原点。如果你的任务只需要轻量调整,别上全量微调;但如果你发现 LoRA 的效果总是差临门一脚,别犹豫,全量训练可能是唯一解法。

最后说一句:微调这件事,数据比算法重要。花时间准备高质量数据,比疯狂调参更有用。模型只能学好你教它的东西,喂进去什么,它就吐出来什么——这个道理,在这个领域尤其真实。

版权声明: 本文首发于 指尖魔法屋-AI全量微调折腾手记https://blog.thinkmoon.cn/post/393-ai-full-finetuning-complete-practice-guide/) 转载或引用必须申明原指尖魔法屋来源及源地址!