Zhangzhe's Blog

The projection of my life.

0%

上手训练大模型(1)——用Alpaca-cleaned指令微调Llama-3.2-3B

TL;DR

  • 本文是上手训练大模型系列的第一篇,主要介绍了如何用指令微调数据集去微调一个大模型。
  • 为了快速走通流程,本文选择了较小的指令微调数据集和较小的模型。

详细介绍

SFT

  • SFT 全称是 Supervised Fine-Tuning,监督微调,是预训练大模型在变成实际可用模型必须经过的一个步骤。
  • 训练的基础配置(例如:optimizerlearning reteweight decay)等一般都是从预训练模型继承过来的。
  • 但实际上训练范式已经完全变化了,从无监督的预训练变成了有监督的微调,具体来说:
    • 预训练过程中,模型只需要预测下一个 token,直到输出 EOS 或者达到最大长度。
    • Instruct-SFT 阶段,模型的输入是一条指令(通常是 Instruction + Input,即指令和输入),输出是指令对应的输出。指令部分的预测不会被用于 loss 的计算,只有输入部分的预测会被用于 loss 的计算。
  • 用一个例子来说明,假如有一条数据:“请计算:1 + 1 = 2”
    • 在预训练阶段
      • 模型 input 是 [“请”,“计算”,“:”,“1”,“+”,“1”,“=”,“2”]
      • 模型 label 是 [“计算”,“:”,“1”,“+”,“1”,“=”,“2”, “pad”],其中 pad 是占位符,表示这个位置不需要预测。
      • pad 之外,其他位置的预测都会被用于 loss 的计算。
    • Instruct-SFT 阶段
      • 模型 input 是 [“请”,“计算”,“:”,“1”,“+”,“1”,“=”]
      • 模型 label 是 [“pad”, “pad”, “pad”, “pad”, “pad”, “pad”, “2”]
      • 只有最后一个位置的预测会被用于 loss 的计算。
  • 从实际推理效果看:
    • SFT 前,模型的输出重复性较高,指令遵循能力很差。
    • SFT 后,模型的输出重复性降低,指令遵循能力明显提升,但可能产生格式依赖问题。

Alpaca-cleaned 数据集

  • Alpaca-cleaned 是一个指令微调数据集,数据集地址:Alpaca-cleaned
  • 数据集包含 51,760 条数据,每条数据包含 instructioninputoutput 三个字段。
  • 其中 instruction 是指令,input 是输入,output 是输出,input 字段可能为空。
  • 例如:
    • instruction 是 “请计算如下数学题:”
    • input 是 “1 + 1 =”
    • output 是 “2”

Llama-3.2-3B

  • Llama-3.2-3Bmeta 开源的预训练大模型(未经过 SFTRLHF),模型地址:Llama-3.2-3B
  • 词表大小 128,256,支持英语、德语、法语、意大利语、葡萄牙语、印地语、西班牙语和泰语等多种语言(就是不支持中文,好烦)。
  • 模型有 32Transformer,每个 Transformer32Head,使用了 GQA 降低了 Transformer 的计算复杂度,hidden size4096

指令微调

  • 指令微调的流程如下:
    1. 加载预训练模型和 tokenizer
    2. 加载数据集并进行预处理
    3. 分布式训练
  • 具体实现代码如下:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
import os
import torch
import torch.distributed as dist
from datasets import load_dataset
from datasets import Split
from transformers import AutoTokenizer, AutoModelForCausalLM
from torch.utils.data import IterableDataset, DataLoader
import deepspeed
from torch.utils.tensorboard import SummaryWriter
# 初始化 DeepSpeed
deepspeed.init_distributed()
# 设置环境变量以避免 tokenizers 并行化问题
os.environ["TOKENIZERS_PARALLELISM"] = "false"
# 设置分布式训练所需的环境变量
local_rank = int(os.getenv("LOCAL_RANK", "0"))
world_size = int(os.getenv("WORLD_SIZE", "1"))
os.environ["RANK"] = str(local_rank)
os.environ["WORLD_SIZE"] = str(world_size)
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "23333"
# NCCL 超时参数配置
os.environ["NCCL_BLOCKING_WAIT"] = "1"
os.environ["NCCL_ASYNC_ERROR_HANDLING"] = "1"
os.environ["NCCL_DEBUG"] = "INFO"
os.environ["NCCL_TIMEOUT"] = "600"
# 初始化进程组
if not dist.is_initialized():
dist.init_process_group(backend="nccl")
# config
model_path = "./models/meta-llama/Llama-3.2-3B"
dataset_path = "./dataset/LLM/yahma___alpaca-cleaned"
ds_config = "./ds_config.json"
log_dir = "./logs"
max_length = 1024
batch_size = 8
train_epochs = 3
# Load the local alpaca dataset
dataset_stream = load_dataset(dataset_path, split=Split.TRAIN, streaming=True)
# tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_path, token=True)
tokenizer.pad_token = tokenizer.eos_token # llama-3.2 tokenizer padding token not set
# model
model = AutoModelForCausalLM.from_pretrained(model_path, token=True)
model.to(f"cuda:{local_rank}") # 将模型移动到对应的 GPU
# dataset
class StreamDataset(IterableDataset):
def __init__(self, data_stream, tokenizer, max_length=max_length):
self.data_stream = data_stream
self.tokenizer = tokenizer
self.max_length = max_length
def __iter__(self):
for data in self.data_stream:
instruction = data["instruction"]
input_text = data["input"]
output = data["output"]
# 构造 Prompt
if input_text.strip():
prompt = f"### Instruction:\n{instruction}\n\n### Input:\n{input_text}\n\n### Response:\n"
else:
prompt = f"### Instruction:\n{instruction}\n\n### Response:\n"
# 拼接完整输入序列:prompt + output
full_text = prompt + output
# Tokenize
input_ids = tokenizer(
full_text,
truncation=True,
max_length=self.max_length,
padding="max_length",
)["input_ids"]
# 构造 labels: 仅监督 output 部分
# 找到 output 起始位置
prompt_length = len(
tokenizer(prompt, padding="do_not_pad", truncation=False)["input_ids"]
)
labels = [-100] * prompt_length + input_ids[prompt_length:] # 忽略 prompt 的部分
labels = labels[: self.max_length] # 截断
labels = labels + [-100] * (max_length - len(labels)) # Padding 保持一致长度
yield {
"input_ids": torch.tensor(input_ids, dtype=torch.long),
"labels": torch.tensor(labels, dtype=torch.long),
}
def __len__(self):
return 51760 # cleaned alpaca dataset rows
def train():
# DataLoader
train_dataset = StreamDataset(dataset_stream, tokenizer)
train_dataloader = DataLoader(train_dataset, batch_size=batch_size)
# TensorBoard SummaryWriter
writer = SummaryWriter(log_dir=log_dir) if dist.get_rank() == 0 else None
# DeepSpeed 初始化
model_engine, optimizer, _, _ = deepspeed.initialize(
model=model,
model_parameters=model.parameters(),
config=ds_config,
training_data=train_dataset,
)
# 训练循环
for epoch in range(train_epochs): # num_train_epochs
for step, batch in enumerate(train_dataloader):
inputs = {
key: value.to(f"cuda:{local_rank}") for key, value in batch.items()
}
outputs = model_engine(**inputs) # 添加 labels 参数
loss = outputs.loss
model_engine.backward(loss)
model_engine.step()
# 日志打印和 TensorBoard 记录
if dist.get_rank() == 0:
if step % 10 == 0:
print(
f"Epoch: {epoch}, Step: {step}/{len(train_dataloader)}, Loss: {loss.item()}"
)
writer.add_scalar(
"Loss/train", loss.item(), epoch * len(train_dataloader) + step
)
checkpoint_path = f"./checkpoints"
model_engine.save_checkpoint(checkpoint_path, tag=f"epoch_{epoch}")
print(f"Checkpoint saved at {checkpoint_path}")
# 关闭 TensorBoard SummaryWriter
if writer:
writer.close()
if __name__ == "__main__":
train()
  • 总体上使用了:
    • transformers 加载模型和 tokenizer
    • datasets 加载开源数据集。
    • deepspeed 进行分布式训练,包含大规模分布式训练和零冗余优化以及混合精度训练等。

模型格式转换

  • 使用 deepspeed 训练和保存的模型格式是 deepspeed 特有的,无法直接加载到 transformers 中,因此需要进行模型格式转换。
  • deepspeed 保存的模型结构如下:
1
2
3
4
5
6
7
8
checkpoints/
├── epoch_xxx
│ ├── mp_rank_00_model_states.pt
│ ├── zero_pp_rank_0_mp_rank_00_optim_states.pt
│ ├── zero_pp_rank_1_mp_rank_00_optim_states.pt
│ ├── zero_pp_rank_2_mp_rank_00_optim_states.pt
│ └── zero_pp_rank_3_mp_rank_00_optim_states.pt
└── latest
  • 转换过程分成两步:
    1. 使用 deepspeed 在保存模型同时提供的 zero_to_fp32.py 脚本将模型转换为 fp32 格式(主要是将不同的 partfp16 模型参数用 ring reduce 聚合成完整的并转成 fp32),转换后的模型格式是 pytorch 适用的(zero_to_fp32.py 只会转换 latest 指向的 epoch)。
    2. pytorch 格式的模型转成 transformers 格式。
  • 第一步代码 deepseed 提供了,第二步代码如下:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
import os
from transformers import AutoModelForCausalLM, AutoTokenizer
final_model_save_path = "./checkpoints/final_transformers_model" # 保存最终模型的路径
original_model_path = "./models/meta-llama/Llama-3.2-3B" # 原始模型的路径
sharded_model_dir = "./checkpoints/fp32_checkpoints" # SFT fp32 分片模型的路径
# 1. 将原始模型配置复制到最终模型目录,需要让 transformers 自动识别模型
os.system(
f"cp {os.path.join(original_model_path, 'config.json')} {final_model_save_path}"
)
# 2. 将分片模型转换为 transformers 模型
model = AutoModelForCausalLM.from_pretrained(sharded_model_dir)
model.save_pretrained(final_model_save_path)
# 3. 补全 tokenizer 配置
tokenizer = AutoTokenizer.from_pretrained(original_model_path, token=True)
tokenizer.save_pretrained(final_model_save_path)
  • 转换后的模型变成了 transformers 可以识别的格式,可以直接加载到 transformers 中进行推理。

推理

  • 将转换后的模型加载到 transformers 中,即可使用 pipeline 进行推理。
  • 推理代码如下:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
import torch
from transformers import pipeline
model_id = "./checkpoints/final_transformers_model"
pipe = pipeline(
"text-generation", model=model_id, torch_dtype=torch.bfloat16, device_map="auto"
)
instruction = "Use Python code to implement."
input_text = "Find prime numbers within 100."
total_input = f"### Instruction:\n{instruction}\n\n### Input:\n{input_text}\n\n### Response:\n"
print(
pipe(total_input, max_length=1024)[0][
"generated_text"
]
)

Thoughts

  • 大模型是个非常依赖工具的领域,transformersdatasetsdeepspeed 等工具的使用对于大模型的最终效果非常重要。
  • 大模型即使在微调阶段,能做的也不多,什么都不能改,数据集的质量对于模型的影响非常大。
  • 指令微调前后,模型差距不是一星半点。