SFT（Supervised Fine-Tuning） 是指在已有的大语言模型（LLM，如 GPT-3、LLaMA）基础上，利用人工标注的监督数据进行微调，使模型能够更好地完成特定任务或符合特定风格。
通俗理解：
原始大模型就像一个“万能工具箱”，能干很多事情，但不一定专精。
SFT就是拿一些明确示例（输入→输出），教模型“这样回答是对的”，让模型在特定任务上更精准。
关键特点：
监督学习：有明确的输入和对应正确输出（训练标签）。
微调：不是从零训练，而是在大模型已有知识上微调。
输出通常是 生成式文本，也可用于分类、问答等任务。
如果按照数据类型分类#
指令-响应型（Instruction-Response）
数据格式：
输入（指令） → 输出（回答）特点：常用于 ChatGPT 类模型训练
示例：
Input: 翻译“Hello”成中文
Output: 你好问答型（Question-Answer）
类似指令型，但通常是明确问题 → 明确答案
用于 FAQ、知识问答系统
对话型（Conversation / Multi-turn Dialogue）
数据包含多轮上下文
训练模型理解上下文关系
示例：
User: 你好，今天天气如何？
Model: 今天天气晴，温度25℃。
User: 明天呢？
Model: 明天有小雨。分类型（Text Classification）
输入文本 → 标签输出
用于情感分析、主题分类等任务
生成型（Text Generation / Summarization）
- 输入是文本或文章 → 输出是摘要、改写、翻译等
按训练策略分类#
全量微调（Full Fine-Tuning）
更新模型所有参数
优点：表达能力最强
缺点：算力消耗大，容易过拟合
低秩微调（LoRA / Adapter + SFT）
只训练部分参数或附加模块
优点：节省显存，容易管理
缺点：可能略低于全量微调性能
混合策略（Hybrid SFT）
先用全量 SFT，再加低秩/Adapter 微调特定任务
常用于多任务或多领域微调
上面就是简单的讲解了一下具体都有哪些的东西。下面会详细的讲解全量微调（Full Fine-Tuning）与 低秩微调LoRA。
在任务开始之前，我们先来讲解一下我们的微调数据集。
微调数据集#
我们之前已经做过了预训练，那么这个时候到底什么是预训练呢？我们不是已经训练了吗？为什么还有一个微调？这个又是一个什么训练？
预训练是干嘛？#
想象你要培养一个人，让他成为一个“智能助理”。你不可能一开始就教他写公文、做客服，你得先让他理解语言本身。
这就是「预训练」的目的。
在预训练阶段，我们不会告诉模型“你要翻译”“你要总结”“你要写代码”。
它得到的唯一任务是：
——在海量文本中，给定前面的内容，预测下一个词是什么。
例如，看到
“今天天气很好，我想去___”
模型就学习去预测下一个词：“散步”、“公园”、“外面”等。
通过这种预测，它慢慢学会语言的统计规律、语法结构、常见表达方式、词义之间的关系、甚至隐含的世界知识。
如果它看够多的文本，就能在参数中编码出非常丰富的语言理解能力。
比如，当它看到“北京是中国的”，它知道接下来的词概率最大的是“首都”；当它看到“水在零度以下会”，它知道可能接“结冰”。
预训练的核心是“语言建模”，不是任务学习。
它像一个人，从小到大通过阅读几亿本书，掌握语言和常识，但他并不知道你要他“干什么”。
他能理解语言，但还不会像人一样“照指令行动”。
他可以续写文章，但你让他“写一份律师函”时，他可能写不出规范格式的内容。
微调是干嘛的？#
这就到了第二个阶段——微调。
微调，就是在预训练模型的基础上，教它“怎么做事”。
你不再让它盲目预测下一个词，而是用人类准备好的「输入 → 输出」示例，告诉它“当别人这么问时，你应该这么回答”。
比如：
用户：请帮我写一封辞职信。
模型：尊敬的领导，您好……（写出完整辞职信）通过成千上万这样的样本，模型学会了遵循人类指令、完成特定任务、采用适合的语气。
从训练机制上看，这个阶段依然是优化交叉熵损失（预测正确答案的概率），但训练的数据不同了：
预训练时的数据是无标签的自然文本；
微调时的数据是有明确目标输出的人工标注数据。
因此，微调不是“再训练一次语言模型”，而是“在语言理解能力的基础上，矫正行为模式”。
它让模型不再只是复述知识，而是学会根据任务需求输出特定形式的答案。
这就像一个人：
他在预训练阶段读了所有百科、小说、报纸，语言能力极强；
但如果你要他去做医生、律师、客服，他还需要在那些领域接受“专业训练”——这就是微调。
举个真实例子。
OpenAI 训练 GPT 系列模型时，首先用互联网上几万亿个 token 的语料进行预训练，耗时数月，成本巨大。
这个阶段得到的模型（比如 GPT-3）已经能理解语言，能续写文本，但回答问题常常胡说八道，甚至不遵守指令。
于是，他们在 GPT-3 的基础上，用几万条人工编写的「指令-回答」数据进行监督微调（SFT）。
这让模型开始能理解“问题的上下文”和“人类期望的回答方式”。
后来又加上强化学习（RLHF）——让人类对回答进行打分，再进一步调整模型的输出偏好。
经过这一系列步骤，才有了现在你熟悉的 ChatGPT，它不仅懂语言，还能对话、分析、推理、保持礼貌、避免不当内容。
所以，当你问“我们不是已经训练了吗？为什么还要微调？” 答案是：预训练教模型懂语言，而微调教模型做任务。 预训练相当于给模型“认知能力”，微调是给它“行为准则”。
没有预训练，模型不懂语言，根本无法对输入做出有意义的反应；
没有微调，它虽然懂语言，却无法理解人类的意图，也不会照你的要求去做。
预训练是打地基，让模型拥有语言智能；
微调是装修，让它具备实际用途。
这里大家可以联想一下做CV任务，是不是有很多骨干都是在ImageNet数据集上进行了训练？，这样是不是可以理解了。大家有兴趣可以去看看GPT1这篇报道，里面就很详细的写了，GPT任务是怎么来的，也是现在的大模型的鼻祖。
数据集#
我们使用的是BelleGroup/train_3.5M_CN。链接：
BelleGroup/train_3.5M_CN 这个数据集，说白了就是一大堆中英文聊天记录，主要是“人问一句、AI 回一句”这样的指令对话，用来教模型怎么听懂人话、怎么按要求回答。里面一共有大概三百多万条这样的对话，每条都是 JSON 格式的，内容一般是一个用户提问（标记成 human）和一个模型回答（标记成 assistant），有的只有一问一答，有的能聊上好几百轮。
它被广泛用来做 监督微调（SFT），也就是在模型预训练之后，用这些有“问题—答案”配对的数据再教模型“怎么好好说话、怎么听指令”。不过它的数据来源比较杂，没说清楚是人工写的还是机器生成的，质量参差不齐，可能有重复、废话或者不太准确的内容，而且它的授权是 GPL-3.0，这种许可证在商用时要特别注意合规问题。实际使用前，一般要先清洗过滤、分出验证集，再根据任务需要挑选或裁剪数据。简单讲，这个数据集就是一个能让大模型学会中文对话和执行指令的“大型练习册”，但用之前最好先打磨干净再训练。
那么这个数据集到底是怎么样的呢？到底长什么样子。
我们将数据下载下来，大概是5G的数据集。我们仍然是使用之前我们写的读取大文件的预览代码，将文件读取一下，看看内部具体是什么样子的：
代码地址：preview_jsonl.py
import json
file_path = "/root/autodl-tmp/Dataset/BelleGroup/train_3.5M_CN.json"
k = 1000
def stream_json_objects(file_path, max_items=1000):
with open(file_path, "r", encoding="utf-8") as f:
buf = []
depth = 0
in_string = False
escape = False
count = 0
while True:
ch = f.read(1)
if not ch:
break
buf.append(ch)
if escape:
escape = False
continue
if ch == "\\":
escape = True
continue
if ch == '"':
in_string = not in_string
continue
if not in_string:
if ch == '{':
depth += 1
elif ch == '}':
depth -= 1
if depth == 0:
raw = ''.join(buf)
buf = []
try:
obj = json.loads(raw)
yield obj
count += 1
if count >= max_items:
return
except json.JSONDecodeError:
print("解析失败：", raw[:100])
# 使用
for item in stream_json_objects(file_path, k):
print(item)但是，这个数据格式不是我们想要的，我们需要写一个代码转换一下。那么我们想要一个什么格式？
对于原始的数据集来说：
它是长什么样子呢？
{
"id": "...",
"conversations": [
{"from": "human", "value": "问题文本"},
{"from": "assistant", "value": "回答文本"}
]
}那么我们需要什么样子的呢？
[
{"role": "system", "content": "你是一个AI助手"},
{"role": "user", "content": "..."},
{"role": "assistant", "content": "..."},
...
]我们通过代码来处理一下：
import json
from tqdm import tqdm
sft_data = '/root/autodl-tmp/Dataset/BelleGroup/train_3.5M_CN.json'
output_sft_data = '/root/autodl-tmp/Dataset/BelleGroup/BelleGroup_sft.jsonl'
def convert_message(data):
*"""*
* 将原始数据转换为标准格式*
* """*
* *message = [
{"role": "system", "content": "你是一个AI助手"},
]
for item in data:
if item['from'] == 'human':
message.append({'role': 'user', 'content': item['value']})
elif item['from'] == 'assistant':
message.append({'role': 'assistant', 'content': item['value']})
return message
with open(output_sft_data, 'a', encoding='utf-8') as sft:
with open(sft_data, 'r', encoding='utf-8') as f:
data = f.readlines()
for item in tqdm(data, desc="Processing", unit="lines"):
item = json.loads(item)
message = convert_message(item['conversations'])
sft.write(json.dumps(message, ensure_ascii=False) + '\n')我们可以再次的预览一下：
这个就是我们想要的。
数据加载器#
现在我们是不是得写一个数据加载器。和我们预训练一样。
代码地址：SFTMapDataset.py
import json
import torch
from torch.utils.data import Dataset
class SFTMapDataset(Dataset):
*"""*
* 高效 map-style SFT Dataset*
* 支持 role/content 格式数据，X/Y/loss_mask 对齐*
* """*
* *def __init__(self, data_path, tokenizer, max_length=256, padding=0):
super().__init__()
self.data_path = data_path
self.tokenizer = tokenizer
self.max_length = max_length
self.padding = padding
# 构建文件行偏移表
self.offsets = []
cur = 0
with open(data_path, "rb") as f:
for line in f:
self.offsets.append(cur)
cur += len(line)
self._len = len(self.offsets)
self._fp = None
def __len__(self):
return self._len
def _ensure_file_open(self):
if self._fp is None:
self._fp = open(self.data_path, "r", encoding="utf-8")
def _read_line(self, index):
self._ensure_file_open()
self._fp.seek(self.offsets[index])
return self._fp.readline().rstrip("\n")
def generate_loss_mask(self, input_ids):
# 生成 loss mask, 0 表示不计算损失, 1 表示计算损失
mask = [0] * len(input_ids)
a_sequence = self.tokenizer("<|im_start|>assistant\n")['input_ids'] # <|im_start|>assistant\n
a_length = len(a_sequence)
n = len(input_ids)
i = 0
while i <= n - a_length:
# 检查当前位置是否匹配目标子序列
match = True
for k in range(a_length):
if input_ids[i + k] != a_sequence[k]:
match = False
break
if match:
# 从子序列结束的位置开始查找第一个 4 (eos_token_id)
j = None
for idx in range(i + a_length, n):
if input_ids[idx] == self.tokenizer.eos_token_id:
j = idx
break
if j is not None:
start = i + a_length
end = j # 结束位置设为j（包含4）
# 标记区间为1（包括start到end）
if start <= end:
for pos in range(start, end + 1):
if pos < len(mask):
mask[pos] = 1
# 跳过当前子序列，避免重叠匹配
i += a_length
else:
i += 1
return mask
def __getitem__(self, index):
line = self._read_line(index)
sample = json.loads(line)
text = self.tokenizer.apply_chat_template(sample, tokenize=False, add_generation_prompt=False)
input_ids = self.tokenizer(text).data['input_ids'][:self.max_length]
# padding
pad_len = self.max_length - len(input_ids)
if pad_len > 0:
input_ids += [self.padding] * pad_len
loss_mask = self.generate_loss_mask(input_ids)
# X/Y/loss_mask 对齐
X = torch.tensor(input_ids[:-1], dtype=torch.long)
Y = torch.tensor(input_ids[1:], dtype=torch.long)
loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long)
return X, Y, loss_mask
def __del__(self):
if hasattr(self, "_fp") and self._fp is not None:
try:
self._fp.close()
except:
pass
# ===== 测试 =====
if __name__ == "__main__":
from torch.utils.data import DataLoader
from transformers import AutoTokenizer
import os
test_file = "test_sft.jsonl"
sample_data = [[
{'role': 'system', 'content': '你是一个AI助手'},
{'role': 'user', 'content': '根据以下文本，对此事件进行分类：中国队在足球比赛中赢得了冠军。'},
{'role': 'assistant', 'content': '这个事件可以被分类为体育比赛。具体地，中国队在足球比赛中致胜并赢得了冠军。'}
]]
with open(test_file, "w", encoding="utf-8") as f:
for item in sample_data:
f.write(json.dumps(item, ensure_ascii=False) + "\n")
tokenizer = AutoTokenizer.from_pretrained("/root/StudyLLM/NX_LLM/第二章 动手实现/tokenizer_k")
dataset = SFTMapDataset(test_file, tokenizer, max_length=128)
dataloader = DataLoader(dataset, batch_size=2, shuffle=False)
for X, Y, loss_mask in dataloader:
print("input_ids:\n", X)
print("labels:\n", Y)
print("loss_mask:\n", loss_mask)
break
os.remove(test_file)这个和我们之前的预训练有什么不一样吗？核心逻辑是什么
假设输入数据
[
{"role": "user", "content": "我爱南巷的花猫"},
{"role": "assistant", "content": "我也很喜欢这只猫，它非常可爱。"}
]构建 Chat 模板
<|im_start|>user
我爱南巷的花猫
<|im_end|><|im_start|>assistant
我也很喜欢这只猫，它非常可爱。
<|im_end|>Tokenization
假设 tokenizer 把文本分成 token IDs（用数字代替）：
[3, 1001, 1002, 1003, 1004, 4, 3, 2001, 2002, 2003, 2004, 4]
3=<|im_start|>4=<|im_end|>1001… = 用户输入
2001… = AI回答
最大长度 max_length = 12（假设）
Padding 不需要，因为刚好满。
生成 X, Y, loss_mask
- X = 输入序列去掉最后一个 token
X = [3, 1001, 1002, 1003, 1004, 4, 3, 2001, 2002, 2003, 2004]
- Y = 输入序列去掉第一个 token
Y = [1001, 1002, 1003, 1004, 4, 3, 2001, 2002, 2003, 2004, 4]
- loss_mask 通过
generate_loss_mask生成
generate_loss_mask 逻辑：
找
<|im_start|>assistant\n>的起始位置从子序列结束位置到
<|im_end|>(token_id = 4) 之间的 token 标记为 1，其余为 0
假设 <|im_start|>assistant\n> 对应 token [3] 开头的第 7 个 token：
mask = [0,0,0,0,0,0,0,1,1,1,1,1] # 用户部分为0，助手回答部分为1
去掉第一个 token 与 X/Y 对齐：
loss_mask = [0,0,0,0,0,0,1,1,1,1,1] # 长度 = len(X) = 11
最终返回
X = tensor([3, 1001, 1002, 1003, 1004, 4, 3, 2001, 2002, 2003, 2004]) Y = tensor([1001, 1002, 1003, 1004, 4, 3, 2001, 2002, 2003, 2004, 4]) ``loss_mask = tensor([0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1])
解释：
X是模型的输入序列Y是预测目标（模型要预测的下一个 token）loss_mask只对助手输出部分计算损失，用户输入部分不参与损失计算
本质上和我们之前绘制的那个图是一样的，我们可以回顾一下。
这个loss mask 实现起来也是非常非常简单，直接就：loss * loss_mask
下面就开始我们的全参微调#
直接上代码#
这个怎么说呢，没有任何好讲的，大家发现没有，我们数据构建一模一样，我们全参微调训练方法和预训练方法一模一样！！！
没错，是一模一样，没有任何需要更改的地方！！！！
但是我们依然将代码附上！
代码地址：Code/SFTtrain.py
# -*- coding: utf-8 -*-
import os
import argparse
import time
import math
import torch
from torch import optim
from torch.utils.data import DataLoader, DistributedSampler
import torch.distributed as dist
from contextlib import nullcontext
from transformers import AutoTokenizer
from ModelConfig import ModelConfig
from LLaMA2 import Transformer
from SFTMapDataset import SFTMapDataset
import swanlab
# ------------------- 工具函数 -------------------
def Logger(msg: str, rank=0):
*"""*
* 多卡训练中只在主卡打印日志*
* rank=0 表示主进程*
* """*
* *if rank == 0:
print(msg)
def get_lr(it, total_iters, args):
*"""*
* 计算学习率：*
* - 线性预热阶段：从0线性增长到目标学习率*
* - 余弦退火阶段：按余弦衰减到最小学习率*
* - 超过总迭代次数：保持最小学习率*
* """*
* *warmup_iters = args.warmup_iters
min_lr = args.learning_rate / 10
if it < warmup_iters:
return args.learning_rate * it / warmup_iters
elif it > total_iters:
return min_lr
else:
decay_ratio = (it - warmup_iters) / (total_iters - warmup_iters)
coeff = 0.5 * (1 + math.cos(math.pi * decay_ratio))
return min_lr + coeff * (args.learning_rate - min_lr)
def save_model(model, save_dir, step, lm_config, rank=0):
*"""*
* 保存模型：*
* - DDP 多卡训练时，只让主卡(rank=0)保存模型*
* - 处理 DataParallel 或 DDP 包装的模型*
* """*
* *if rank == 0:
os.makedirs(save_dir, exist_ok=True)
state_dict = model.module.state_dict() if hasattr(model, 'module') else model.state_dict()
path = f"{save_dir}/pretrain_{lm_config.dim}_{lm_config.n_layers}_{lm_config.vocab_size}_step{step + 1}.pth"
torch.save(state_dict, path)
Logger(f"模型保存: {path}", rank)
# ------------------- 模型初始化 -------------------
def init_model(args, lm_config, rank, gpu_id):
*"""*
* 初始化模型 + tokenizer + DDP 包装*
* """*
* *def count_parameters(model):
*"""计算可训练参数量"""*
* *return sum(p.numel() for p in model.parameters() if p.requires_grad)
# 加载 tokenizer
tokenizer = AutoTokenizer.from_pretrained('/root/StudyLLM/HappyLLM/Tokenizer/tokenizer_k')
# 创建 Transformer 模型并移动到当前 GPU
model = Transformer(lm_config).to(args.device)
# 加载预训练权重
ckp = '/root/StudyLLM/NX_LLM/第三章_预训练流程/ModelZoom/base-CodeLab-50M/pretrain_512_16_6144_step67200.pth'
state_dict = torch.load(ckp, map_location=args.device)
unwanted_prefix = '_orig_mod.'
for k, v in list(state_dict.items()):
if k.startswith(unwanted_prefix):
state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
model.load_state_dict(state_dict, strict=False)
# DDP 包装模型，每个进程只处理自己对应的 GPU
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[gpu_id],
output_device=gpu_id,
)
Logger(f'LLM总参数量: {count_parameters(model) / 1e6:.3f} M', rank)
return model, tokenizer
# 训练单个epoch，其实对于大模型来说，基本都是只会训练一次，因为这个数据量实在是非常非常庞大
def train_epoch(epoch, model, train_loader, optimizer, scaler, args, lm_config, ctx, iter_per_epoch, rank):
*"""*
* 训练一个 epoch*
* - 支持梯度累积*
* - 支持混合精度*
* - 支持 DDP*
* """*
* *# 记录开始时间
start_time = time.time()
for step, (X, Y, Attention_mask) in enumerate(train_loader):
# 将数据迁移到显卡
X, Y, Attention_mask = X.to(args.device), Y.to(args.device), Attention_mask.to(args.device)
# 一个简易的学习率调度器
lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
# 使用混合精度前向传播
with ctx:
out = model(X, Y) # 模型前向
# 除以梯度累积步数
loss = out.last_loss / args.accumulation_steps
# 展平 mask
loss_mask_flat = Attention_mask.view(-1)
# 忽略 padding
loss = torch.sum(loss * loss_mask_flat) / loss_mask_flat.sum()
# 反向传播获取梯度
scaler.scale(loss).backward()
# 进行梯度更新
if (step + 1) % args.accumulation_steps == 0:
scaler.unscale_(optimizer)
# 梯度裁剪
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
# 优化器更新
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True)
# ------------------- 日志 -------------------
if step % args.log_interval == 0:
elapsed_time = time.time() - start_time
Logger(
f"Rank[{rank}] Epoch[{epoch + 1}/{args.epochs}] "
f"Step[{step}/{iter_per_epoch}] "
f"Loss:{loss.item() * args.accumulation_steps:.4f} "
f"LR:{lr:.7f} Elapsed:{elapsed_time:.1f}s",
rank
)
if args.use_swanlab and rank == 0:
swanlab.log({"loss": loss.item() * args.accumulation_steps, "lr": lr})
# ------------------- 模型保存 -------------------
if (step + 1) % args.save_interval == 0:
save_model(model, args.save_dir, step, lm_config, rank)
# ------------------- 主函数 -------------------
if __name__ == "__main__":
parser = argparse.ArgumentParser()
# ------------------- 基础训练参数 -------------------
parser.add_argument("--out_dir", type=str, default="../ModelZoom/SFT-CodeLab-50M", help="模型输出目录")
parser.add_argument("--epochs", type=int, default=1, help="训练轮数")
parser.add_argument("--batch_size", type=int, default=64, help="模型训练批次大小")
parser.add_argument("--learning_rate", type=float, default=2e-4, help="训练的学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
parser.add_argument("--dtype", type=str, default="bfloat16")
parser.add_argument("--use_swanlab", action="store_true", default=True, help="是否使用SwanLab进行实验跟踪")
parser.add_argument("--num_workers", type=int, default=8, help="数据加载的工作进程数")
parser.add_argument("--data_path", type=str,
default="/root/autodl-tmp/Dataset/BelleGroup/BelleGroup_sft.jsonl",
help="训练数据路径")
parser.add_argument("--accumulation_steps", type=int, default=8, help="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
parser.add_argument("--warmup_iters", type=int, default=0, help="学习率预热迭代次数")
parser.add_argument("--log_interval", type=int, default=100, help="日志记录间隔")
parser.add_argument("--save_interval", type=int, default=1200, help="模型保存间隔")
# DDP 本地 rank
parser.add_argument("--local_rank", type=int, default=0, help="DDP local rank")
args = parser.parse_args()
# ------------------- 初始化 DDP & 自动适配 GPU -------------------
dist.init_process_group(backend='nccl')
local_rank = int(os.environ.get("LOCAL_RANK", 0))
gpu_id = local_rank
torch.cuda.set_device(gpu_id)
args.device = f'cuda:{gpu_id}'
# 获取总 GPU 数量
ngpus = torch.cuda.device_count()
# 检查 local_rank 是否超出可用 GPU 范围
if gpu_id >= ngpus:
raise RuntimeError(f"local_rank {gpu_id} 超出可用 GPU 范围 (0-{ngpus - 1})")
Logger(f"[Rank {args.local_rank}] 使用 GPU {gpu_id} / 总 GPU 数量 {ngpus}", args.local_rank)
# ------------------- SwanLab -------------------
if args.use_swanlab and args.local_rank == 0:
swanlab.login(api_key="你的 API key")
swanlab.init(project="CodeLab-LLM", experiment_name="SFT-50M", config=args)
# ------------------- 模型配置 -------------------
lm_config = ModelConfig(dim=512, n_layers=16)
max_seq_len = lm_config.max_seq_len
args.save_dir = os.path.join(args.out_dir)
os.makedirs(args.out_dir, exist_ok=True)
torch.manual_seed(42)
# ------------------- 混合精度上下文 -------------------
ctx = nullcontext() if "cpu" in args.device else torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16)
# ------------------- 模型 & 数据 -------------------
model, tokenizer = init_model(args, lm_config, args.local_rank, gpu_id)
train_ds = SFTMapDataset(args.data_path, tokenizer, max_length=max_seq_len)
# DDP 必须使用 DistributedSampler
train_sampler = DistributedSampler(train_ds)
train_loader = DataLoader(
train_ds,
batch_size=args.batch_size,
sampler=train_sampler,
num_workers=args.num_workers,
pin_memory=True
)
# ------------------- 优化器 & 混合精度 scaler -------------------
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
# ------------------- 开始训练 -------------------
iter_per_epoch = len(train_loader)
for epoch in range(args.epochs):
train_sampler.set_epoch(epoch) # DDP 每轮必须重置 epoch 保证 shuffle
train_epoch(epoch, model, train_loader, optimizer, scaler, args, lm_config, ctx, iter_per_epoch,
args.local_rank)
if args.local_rank == 0:
Logger("训练完成！")
dist.destroy_process_group()我们可以看见，模型正常收敛，说明是有效果的。
基于LORA的微调#
LORA详解#
当然可以！我们用矩阵方式一步步解释 LoRA 的思想和来源，这样会更直观。下面我把原理完全矩阵化。
原始线性层#
假设我们有一个线性层：
$ y = W_0 x$
$x \in \mathbb{R}^{k \times 1}$ ——输入向量
$W_0 \in \mathbb{R}^{d \times k}$ ——原始权重矩阵
$y \in \mathbb{R}^{d \times 1}$ ——输出向量
例如：
输出：
LoRA 的低秩增量矩阵#
LoRA 假设我们只需要对 $W_0$做 低秩调整：
$ \Delta W = A B$
$A \in \mathbb{R}^{d \times r}$
$B \in \mathbb{R}^{r \times k}$
$r \ll d, k$ （低秩）
矩阵维度示意：
$ \underbrace{W_0}_{d \times k} + \underbrace{A}_{d \times r} \cdot \underbrace{B}_{r \times k} = \underbrace{W}_{d \times k}$
输出：
$ y = W x = W_0 x + (A B) x = W_0 x + A (B x)$
矩阵内的直观解释#
假设：
- $d = 4, k = 4, r = 2$
计算：
- 先计算 (B x)
- **再计算 **$A (B x)$
这样，低秩矩阵完成了原矩阵大小的调整，但只需训练$d \cdot r + r \cdot k = 16$** 个参数，而不是 (16) 个参数**（示例很小，实际模型差异更大）。
矩阵化的训练思路#
训练 LoRA 的矩阵目标是最小化损失$L(y_\text{pred}, y_\text{true})$：
$ \min_{A,B} L\Big( (W_0 + AB)x, y_\text{true} \Big)$
固定 (W_0)，只优化 (A) 和 (B)
训练参数量小，计算量小
输出仍然是原模型维度 $d \times 1$
LoRA 的直观矩阵理解#
原矩阵 (W_0) ：已经学会通用知识
低秩矩阵 (AB) ：学习“任务特定方向”，在一个小低维子空间里调整输出
输出叠加：
$y = W_0 x + A (B x)$
(B) 将输入 (x) 投影到低维空间
(A) 将低维空间映射回输出维度
低秩矩阵就像给模型加了一个“小补丁”
我们使用peft工具
pip install peft这个工具只需要对原先的代码进行很小的改动就行。
1、包装模型#
from peft import LoraConfig, get_peft_model, TaskType
# LoRA 配置
lora_config = LoraConfig(
task_type=TaskType.CAUSAL_LM, # 自回归语言模型
r=8, # LoRA rank
lora_alpha=16, # LoRA alpha
lora_dropout=0.1, # Dropout
target_modules=["q_proj", "v_proj"] # 哪些权重使用 LoRA
)
# 包装模型
model = get_peft_model(model, lora_config)2、优化器修改#
optimizer = optim.AdamW(
filter(lambda p: p.requires_grad, model.parameters()),
lr=args.learning_rate
)3、保存与加载 LoRA 权重#
- 保存 LoRA 权重：
from peft import PeftModel
if rank == 0:
model.save_pretrained(args.save_dir)- 加载 LoRA 权重：
from peft import PeftModel
base_model = Transformer(lm_config).to(args.device)
model = PeftModel.from_pretrained(base_model, lora_path)下面就将我们的训练代码改一下，改为基于我们的方法进行LORA微调
# -*- coding: utf-8 -*-
import os
import argparse
import time
import math
import torch
from torch import optim
from torch.utils.data import DataLoader, DistributedSampler
import torch.distributed as dist
from contextlib import nullcontext
from transformers import AutoTokenizer
from ModelConfig import ModelConfig
from LLaMA2 import Transformer
from SFTMapDataset import SFTMapDataset
from peft import LoraConfig, get_peft_model, TaskType
import swanlab
# ------------------- 工具函数 -------------------
def Logger(msg: str, rank=0):
if rank == 0:
print(msg)
def get_lr(it, total_iters, args):
warmup_iters = args.warmup_iters
min_lr = args.learning_rate / 10
if it < warmup_iters:
return args.learning_rate * it / warmup_iters
elif it > total_iters:
return min_lr
else:
decay_ratio = (it - warmup_iters) / (total_iters - warmup_iters)
coeff = 0.5 * (1 + math.cos(math.pi * decay_ratio))
return min_lr + coeff * (args.learning_rate - min_lr)
def save_model(model, save_dir, step, lm_config, rank=0):
*"""*
* 保存 LoRA 微调模型（兼容 DDP 和多卡）*
* 文件名和原始预训练保存一致：*
* pretrain_{dim}_{n_layers}_{vocab_size}_step{step}.pth*
* Args:*
* model: 当前模型，可能是 DDP 包装的*
* save_dir: 保存目录*
* step: 当前训练步数*
* lm_config: 模型配置（dim, n_layers, vocab_size）*
* rank: 当前进程 rank，只让主卡保存*
* """*
* *if rank != 0:
return # 只有主卡保存
os.makedirs(save_dir, exist_ok=True)
# 如果是 DDP 包装的模型，需要取 module
model_to_save = model.module if hasattr(model, 'module') else model
# 构建完整的保存路径
path = os.path.join(
save_dir,
f"pretrain_{lm_config.dim}_{lm_config.n_layers}_{lm_config.vocab_size}_step{step + 1}"
)
# 使用 HuggingFace 的 save_pretrained 保存 LoRA 权重
model_to_save.save_pretrained(path)
Logger(f"LoRA 模型已保存: {path}", rank)
# ------------------- 模型初始化 -------------------
def init_model(args, lm_config, rank, gpu_id):
def count_parameters(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
tokenizer = AutoTokenizer.from_pretrained('/root/StudyLLM/HappyLLM/Tokenizer/tokenizer_k')
model = Transformer(lm_config).to(args.device)
# 加载原始预训练权重
ckp = '/root/StudyLLM/NX_LLM/第三章_预训练流程/ModelZoom/base-CodeLab-50M/pretrain_512_16_6144_step67200.pth'
state_dict = torch.load(ckp, map_location=args.device)
unwanted_prefix = '_orig_mod.'
for k, v in list(state_dict.items()):
if k.startswith(unwanted_prefix):
state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
model.load_state_dict(state_dict, strict=False)
# ------------------- LoRA 包装 -------------------
lora_config = LoraConfig(
task_type=TaskType.CAUSAL_LM,
r=4,
lora_alpha=16,
lora_dropout=0.1,
target_modules=["wq", "wv"], # 根据 Transformer 实现选择
)
model = get_peft_model(model, lora_config)
# DDP 包装
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[gpu_id],
output_device=gpu_id,
)
Logger(f'LLM 可训练参数量: {count_parameters(model)/1e6:.3f} M', rank)
return model, tokenizer
# ------------------- 训练 -------------------
def train_epoch(epoch, model, train_loader, optimizer, scaler, args, lm_config, ctx, iter_per_epoch, rank):
start_time = time.time()
for step, (X, Y, Attention_mask) in enumerate(train_loader):
X, Y, Attention_mask = X.to(args.device), Y.to(args.device), Attention_mask.to(args.device)
lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
with ctx:
out = model(X, Y)
loss = out.last_loss / args.accumulation_steps
loss_mask_flat = Attention_mask.view(-1)
loss = torch.sum(loss * loss_mask_flat) / loss_mask_flat.sum()
scaler.scale(loss).backward()
if (step + 1) % args.accumulation_steps == 0:
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True)
if step % args.log_interval == 0:
elapsed_time = time.time() - start_time
Logger(
f"Rank[{rank}] Epoch[{epoch+1}/{args.epochs}] "
f"Step[{step}/{iter_per_epoch}] "
f"Loss:{loss.item()*args.accumulation_steps:.4f} "
f"LR:{lr:.7f} Elapsed:{elapsed_time:.1f}s",
rank
)
if args.use_swanlab and rank == 0:
swanlab.log({"loss": loss.item()*args.accumulation_steps, "lr": lr})
if (step + 1) % args.save_interval == 0:
save_model(model, args.save_dir, step, lm_config, rank)
# ------------------- 主函数 -------------------
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--out_dir", type=str, default="../ModelZoom/SFT-CodeLab-50M-LORA")
parser.add_argument("--epochs", type=int, default=2)
parser.add_argument("--batch_size", type=int, default=64)
parser.add_argument("--learning_rate", type=float, default=2e-4)
parser.add_argument("--device", type=str, default="cuda:0")
parser.add_argument("--dtype", type=str, default="bfloat16")
parser.add_argument("--use_swanlab", action="store_true", default=True)
parser.add_argument("--num_workers", type=int, default=8)
parser.add_argument("--data_path", type=str, default="/root/autodl-tmp/Dataset/BelleGroup/BelleGroup_sft.jsonl")
parser.add_argument("--accumulation_steps", type=int, default=8)
parser.add_argument("--grad_clip", type=float, default=1.0)
parser.add_argument("--warmup_iters", type=int, default=0)
parser.add_argument("--log_interval", type=int, default=100)
parser.add_argument("--save_interval", type=int, default=1200)
parser.add_argument("--local_rank", type=int, default=0)
args = parser.parse_args()
# DDP 初始化
dist.init_process_group(backend='nccl')
local_rank = int(os.environ.get("LOCAL_RANK", 0))
gpu_id = local_rank
torch.cuda.set_device(gpu_id)
args.device = f'cuda:{gpu_id}'
ngpus = torch.cuda.device_count()
if gpu_id >= ngpus:
raise RuntimeError(f"local_rank {gpu_id} 超出可用 GPU 范围 (0-{ngpus-1})")
Logger(f"[Rank {args.local_rank}] 使用 GPU {gpu_id} / 总 GPU 数量 {ngpus}", args.local_rank)
if args.use_swanlab and args.local_rank == 0:
swanlab.login(api_key="你的API KEY")
swanlab.init(project="CodeLab-LLM", experiment_name="SFT-50M-LORA", config=args)
lm_config = ModelConfig(dim=512, n_layers=16)
max_seq_len = lm_config.max_seq_len
args.save_dir = os.path.join(args.out_dir)
os.makedirs(args.out_dir, exist_ok=True)
torch.manual_seed(42)
ctx = nullcontext() if "cpu" in args.device else torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16)
model, tokenizer = init_model(args, lm_config, args.local_rank, gpu_id)
train_ds = SFTMapDataset(args.data_path, tokenizer, max_length=max_seq_len)
train_sampler = DistributedSampler(train_ds)
train_loader = DataLoader(
train_ds,
batch_size=args.batch_size,
sampler=train_sampler,
num_workers=args.num_workers,
pin_memory=True
)
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
optimizer = optim.AdamW(
filter(lambda p: p.requires_grad, model.parameters()), # 只训练 LoRA 参数
lr=args.learning_rate
)
iter_per_epoch = len(train_loader)
for epoch in range(args.epochs):
train_sampler.set_epoch(epoch)
train_epoch(epoch, model, train_loader, optimizer, scaler, args, lm_config, ctx, iter_per_epoch, args.local_rank)
if args.local_rank == 0:
Logger("LoRA 微调完成！")
dist.destroy_process_group()在微调完成后，我们可以写一个简单的控制台交互聊天的代码，进行对话。
from contextlib import nullcontext
import torch
from transformers import AutoTokenizer
from peft import PeftModel
from ModelConfig import ModelConfig
from LLaMA2 import Transformer
import torch.nn.functional as F
class TextGenerator:
def __init__(self,
base_model_path='', # 基础模型 checkpoint，可选
lora_checkpoint='', # LoRA checkpoint，可选
tokenizer_model_path='/root/StudyLLM/NX_LLM/第三章_预训练流程/tokenizer_k',
seed=42,
device=None,
dtype="bfloat16"):
self.device = device or ('cuda:0' if torch.cuda.is_available() else 'cpu')
self.device_type = 'cuda' if 'cuda' in self.device else 'cpu'
self.dtype = dtype
self.seed = seed
torch.manual_seed(seed)
if 'cuda' in self.device:
torch.cuda.manual_seed(seed)
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[self.dtype]
self.ctx = nullcontext() if self.device_type == 'cpu' else torch.amp.autocast(
device_type=self.device_type, dtype=ptdtype
)
# -----------------------------
# 加载基础模型
# -----------------------------
print("加载基础模型...")
self.model = Transformer(ModelConfig(dim=512, n_layers=16))
if base_model_path:
checkpoint_dict = torch.load(base_model_path, map_location='cpu')
self.model.load_state_dict(checkpoint_dict, strict=False)
# -----------------------------
# 加载 LoRA 权重
# -----------------------------
if lora_checkpoint:
print("加载 LoRA 权重...")
self.model = PeftModel.from_pretrained(self.model, lora_checkpoint, device_map={'': self.device})
self.model.eval()
self.model.to(self.device)
# -----------------------------
# 分词器
# -----------------------------
self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_model_path)
# -----------------------------
# 打印参数量
# -----------------------------
num_params = sum(p.numel() for n, p in self.model.named_parameters() if p.requires_grad)
num_lora_params = sum(p.numel() for n, p in self.model.named_parameters() if 'lora' in n)
print(f"总可训练参数量: {num_params / 1e6:.6f} M")
print(f"LoRA 参数量: {num_lora_params / 1e6:.6f} M")
# -----------------------------
# 简单 chat 模板
# -----------------------------
def chat_template(self, prompt):
message = [
{"role": "system", "content": "你是一个AI助手，你的名字叫小明。"},
{"role": "user", "content": prompt}
]
return self.tokenizer.apply_chat_template(message, tokenize=False, add_generation_prompt=True)
# -----------------------------
# 逐步生成函数
# -----------------------------
def sft_sample(self, start="Hello!", num_samples=1, max_new_tokens=128,
temperature=0.7, top_k=50, stop_token_id=None):
*"""*
* 基于模型 forward 逐步生成 token*
* """*
* *start_text = self.chat_template(start)
start_ids = self.tokenizer(start_text).data['input_ids']
idx = torch.tensor(start_ids, dtype=torch.long, device=self.device)[None, ...] # [1, seq_len]
stop_token_id = stop_token_id or self.tokenizer.eos_token_id
outputs = []
with torch.no_grad():
with self.ctx:
for _ in range(num_samples):
cur_idx = idx.clone()
generated = []
for _ in range(max_new_tokens):
logits = self.model(cur_idx).logits[:, -1, :] # [batch, vocab]
logits = logits / temperature
if top_k is not None:
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
logits[logits < v[:, [-1]]] = -float('Inf')
probs = F.softmax(logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1) # [batch, 1]
if next_token.item() == stop_token_id:
break
generated.append(next_token.item())
cur_idx = torch.cat([cur_idx, next_token], dim=1)
outputs.append(self.tokenizer.decode(generated))
return outputs
# -----------------------------
# 控制台聊天函数
# -----------------------------
def chat_console(self):
print("输入 'exit' 退出对话。")
while True:
user_input = input("你: ").strip()
if user_input.lower() in ['exit', 'quit']:
break
answers = self.sft_sample(start=user_input, num_samples=1, max_new_tokens=128,
temperature=0.8, top_k=50)
print(f"小明: {answers[0]}")
# -----------------------------
# 测试
# -----------------------------
if __name__ == "__main__":
generator = TextGenerator(
base_model_path='/root/StudyLLM/NX_LLM/第三章_预训练流程/ModelZoom/base-CodeLab-50M/pretrain_512_16_6144_step67200.pth',
lora_checkpoint='/root/StudyLLM/NX_LLM/第四章_SFT/ModelZoom/SFT-CodeLab-50M-LORA/pretrain_512_16_6144_step4800',
tokenizer_model_path='/root/StudyLLM/NX_LLM/第三章_预训练流程/tokenizer_k'
)
generator.chat_console()结果如下：
是不是有点无厘头。
确确实实是模型容量太小了。
等我有很多显卡，我训练一个大的，然后上传模型到仓库里面。
