LLMs之minimind：minimind源码解读(lora_sft.py)——基于LoRA技术的因果语言模型微调与训练实现—初始化配置参数(最大序列长度/训练周期数/批处理大小/学习率,设置训练设备/创建输出目录/wandb监控等)→初始化模型和分词器(加载模型与分词器+识别并标记需要进行LoRA的线性层+配置LoRA参数+集成LoRA模型+迁移GPU)→初始化Adam优化器和学习率缩放器→模型编译（可选）优化推理速度→模型训练(每个周期调用+保存模型)
目录
minimind源码解读(lora_sft.py)——基于LoRA技术的因果语言模型微调与训练实现—初始化配置参数(最大序列长度/训练周期数/批处理大小/学习率,设置训练设备/创建输出目录/wandb监控等)→初始化模型和分词器(加载模型与分词器+识别并标记需要进行LoRA的线性层+配置LoRA参数+集成LoRA模型+迁移GPU)→初始化Adam优化器和学习率缩放器→模型编译（可选）优化推理速度→模型训练(每个周期调用+保存模型)
# LLMs之minimind：minimind源码解读(lora_sft.py)——基于LoRA技术的因果语言模型微调与训练实现
# 初始化配置参数(最大序列长度/训练周期数/批处理大小/学习率,设置训练设备/创建输出目录/wandb监控等)→初始化模型和分词器(加载模型与分词器+识别并标记需要进行LoRA的线性层+配置LoRA参数+集成LoRA模型+迁移GPU)→初始化Adam优化器和学习率缩放器→模型编译（可选）优化推理速度→模型训练(每个周期调用+保存模型)
'''
源代码地址:https://github.com/jingyaogong/minimind/blob/master/4-lora_sft.py
该代码实现了一个基于 LoRA（Low-Rank Adaptation）技术的语言模型微调流程。整个过程包括模型初始化、数据集加载、学习率调度、混合精度训练以及模型的保存。
代码中通过 PEFT 提供的 LoRA 技术，能够有效减少模型参数的更新，提高训练效率。其训练过程支持自动混合精度优化，并且可以通过可选的 wandb 进行训练日志记录。
模型训练过程中，还使用了动态的学习率调度策略来提升训练效果。
'''
# 导入PyTorch 和 Hugging Face 的 transformer 模型相关库，以及PEFT（Parameter-Efficient Fine-Tuning）的 LoraConfig 和 get_peft_model 进行模型的 LoRA（Low-Rank Adaptation）微调。
# 核心技术点：主要依赖 PyTorch 和 transformers 库，PEFT 用于高效的参数微调。
import os
import platform
import time
import math
import warnings
import torch
import pandas as pd
import torch.nn.functional as F
from contextlib import nullcontext
from torch import optim
from transformers import AutoTokenizer
from transformers import AutoModelForCausalLM
from peft import get_peft_model, LoraConfig, TaskType
from torch.utils.data import DataLoader
from model.LMConfig import LMConfig
from model.dataset import SFTDataset
warnings.filterwarnings('ignore', category=UserWarning)
# 学习率调度函数 get_lr:定义一个随训练步数变化的学习率调度函数，结合 warmup 和 cosine decay 进行动态学习率调整。
# 核心技术点：学习率策略使用了 warmup 和 cosine decay 来优化训练过程。
def get_lr(it):
warmup_iters = 1000
lr_decay_iters = 80000
min_lr = 1e-5
if it < warmup_iters:
return learning_rate * it / warmup_iters
if it > lr_decay_iters:
return min_lr
decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
assert 0 <= decay_ratio <= 1
coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
return min_lr + coeff * (learning_rate - min_lr)
# 训练单个 epoch 的 train_epoch 函数：进行模型训练，计算损失并更新梯度。每 100 步打印一次训练日志，包含损失、学习率等信息。支持通过 wandb（可选）来记录训练过程。
# 核心技术点：使用自动混合精度训练 (torch.cuda.amp)，结合梯度裁剪防止梯度爆炸，以及基于 cross-entropy 损失函数进行优化。
# ------------------------------------------------------------------------------
def train_epoch(epoch, wandb):
start_time = time.time()
for step, (X, Y, loss_mask) in enumerate(train_loader):
X = X.to(device)
Y = Y.to(device)
loss_mask = loss_mask.to(device)
lr = get_lr(epoch * iter_per_epoch + step)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
with ctx:
logits = model(X, Y).logits
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), Y.view(-1), ignore_index=0, reduction='none')
loss_mask = loss_mask.view(-1)
loss = torch.sum(loss * loss_mask) / loss_mask.sum()
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True)
if step % 100 == 0:
spend_time = time.time() - start_time
print(
'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.7f} epoch_Time:{}min:'.format(
epoch,
epochs,
step,
iter_per_epoch,
loss.item(),
optimizer.param_groups[-1]['lr'],
spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))
if wandb is not None:
wandb.log({"loss": loss.item(), "lr": optimizer.param_groups[-1]['lr'],
"epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60})
# 线性层名查找 find_all_linear_names 函数:找出模型中所有线性层的名称，用于后续的 LoRA 模型微调
# 核心技术点：遍历模型，定位需要应用 LoRA 的目标模块（线性层）。
def find_all_linear_names(model):
cls = torch.nn.Linear
lora_module_names = set()
for name, module in model.named_modules():
if isinstance(module, cls):
names = name.split('.')
lora_module_names.add(names[0] if len(names) == 1 else names[-1])
if 'lm_head' in lora_module_names:
lora_module_names.remove('lm_head')
return list(lora_module_names)
# 利用init_model函数初始化模型：加载模型与分词器+识别并标记需要进行LoRA的线性层+配置LoRA参数+集成LoRA模型+迁移GPU
def init_model():
model_name_or_path = "./minimind-v1-small"
tokenizer_name_or_path = "./minimind-v1-small"
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path, trust_remote_code=True, use_fast=False)
model = AutoModelForCausalLM.from_pretrained(model_name_or_path, trust_remote_code=True).to(device)
target_modules = find_all_linear_names(model)
peft_config = LoraConfig(
task_type=TaskType.CAUSAL_LM,
r=8,
lora_alpha=16,
lora_dropout=0.1,
inference_mode=False,
target_modules=target_modules
)
model = get_peft_model(model, peft_config)
model.print_trainable_parameters()
model = model.to(device)
return model, tokenizer
# I/O
if __name__ == "__main__":
# 1、初始化配置参数，如最大序列长度、训练周期数、批处理大小、学习率,设置训练设备、创建输出目录、是否wandb监控等
# -----------------------------------------------------------------------------
lm_config = LMConfig()
max_seq_len = lm_config.max_seq_len
out_dir = 'out'
epochs = 20
gradient_accumulation_steps = 1
batch_size = 16
learning_rate = 1e-4
weight_decay = 1e-1
# 设置训练设备，创建输出目录
device = 'cuda:0'
dtype = 'bfloat16'
save_dir = os.path.join(out_dir)
os.makedirs(save_dir, exist_ok=True)
tokens_per_iter = gradient_accumulation_steps * batch_size * max_seq_len
os.makedirs(out_dir, exist_ok=True)
torch.manual_seed(1337)
device_type = device if "cuda" in device else "cpu"
use_wandb = False # 是否使用wandb
wandb_project = "MiniMind-LoRA-SFT"
wandb_run_name = f"MiniMind-LoRA-SFT-Epoch-{epochs}-BatchSize-{batch_size}-LearningRate-{learning_rate}"
if use_wandb:
import wandb
wandb.init(project=wandb_project, name=wandb_run_name)
else:
wandb = None
ctx = (
nullcontext()
if device_type == "cpu"
else torch.cuda.amp.autocast()
)
# -----------------------------------------------------------------------------
# 2、利用init_model函数初始化模型和分词器：加载模型与分词器+识别并标记需要进行LoRA的线性层+配置LoRA参数+集成LoRA模型+迁移GPU
model, tokenizer = init_model()
# 3、加载数据集SFTDataset并创建数据加载器实现批处理
# -----init dataloader------
df = pd.read_csv('./dataset/sft_data_single.csv')
df = df.sample(frac=1.0)
train_ds = SFTDataset(df, tokenizer, max_length=max_seq_len)
train_loader = DataLoader(
train_ds,
batch_size=batch_size,
pin_memory=False,
drop_last=False,
shuffle=False,
num_workers=0,
)
# 4、初始化Adam优化器和学习率缩放器
scaler = torch.cuda.amp.GradScaler(enabled=(dtype == 'float16'))
# optimizer
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
iter_per_epoch = len(train_loader)
# 5、模型编译（可选）优化推理速度：如果使用的是 Linux 且 PyTorch 版本 2.0 及以上，支持编译模型以加快推理速度。
# 核心技术点：在特定平台上支持 PyTorch 2.0 的模型编译功能（torch.compile），优化推理性能。
# compile the model
if False and platform.system() != 'Windows' and float(torch.__version__.split('.')[0]) >= 2:
print("compiling the model... (takes a ~minute)")
unoptimized_model = model
model = torch.compile(model)
# 6、模型训练，每个周期调用train_epoch函数进行训练，并在训练结束后保存模型。
raw_model = model
# training loop
for epoch in range(epochs):
train_epoch(epoch, wandb)
model.save_pretrained('minimind')
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
