LLMs之minimind：minimind源码解读(pretrain.py)——实现基于Transformer架 ...

LLMs之minimind：minimind源码解读(pretrain.py)——实现基于Transformer架构的大规模语言模型预练习及wandb监控—支持余弦退火学习率调度/分布式预练习/自动肴杂精度优化/梯度累积/梯度裁剪/定期生存模型

目录
minimind源码解读(pretrain.py)——实现基于Transformer架构的大规模语言模型预练习及wandb监控—支持余弦退火学习率调度/分布式预练习/自动肴杂精度优化/梯度累积/梯度裁剪/定期生存模型

---

minimind源码解读(pretrain.py)——实现基于Transformer架构的大规模语言模型预练习及wandb监控—支持余弦退火学习率调度/分布式预练习/自动肴杂精度优化/梯度累积/梯度裁剪/定期生存模型

2. # LLMs之minimind：minimind源码解读(pretrain.py)——实现基于Transformer架构的大规模语言模型预训练及wandb监控—支持余弦退火学习率调度/分布式预训练/自动混合精度优化/梯度累积/梯度裁剪/定期保存模型
4. '''
5. 源代码地址:https://github.com/jingyaogong/minimind/blob/master/1-pretrain.py
7. 这段代码主要实现了基于 Transformer 架构的语言模型的预训练流程。它支持单机多卡以及跨机器的分布式训练。通过精心设计的数据加载机制、
8. 灵活的学习率调整策略、高效的优化算法以及混合精度训练等技术手段，该脚本能够高效稳定地完成大规模语言模型的训练任务。
9. 同时，它还提供了对训练过程的监控和支持实验管理的功能。
11. 该Py文件主要实现了一个基于Transformer模型的分布式预训练任务。文件的实现流程包括：
12. 模型与分布式环境的初始化。
13. 使用余弦退火学习率调度和**自动混合精度训练（AMP）**提升训练效率。
14. 通过torch.nn.parallel.DistributedDataParallel实现模型的分布式训练，加速训练过程。
15. 实现了包括梯度累积、梯度裁剪和定期保存模型等细节，确保模型训练的稳定性和可扩展性。
16. 支持可选的wandb（Weights and Biases）日志记录，用于实验管理与可视化。
17. '''
20. # torch的深度学习库和相关分布式训练模块（torch.distributed、DistributedDataParallel等
21. # 核心技术点：PyTorch分布式训练与自动混合精度训练（AMP）
22. import os
23. import platform
24. import time
25. import math
26. import warnings
27. import torch
28. import torch.distributed as dist
29. from torch import optim
30. from torch.nn.parallel import DistributedDataParallel
31. from torch.optim.lr_scheduler import CosineAnnealingLR
32. from torch.utils.data import DataLoader, DistributedSampler
33. from contextlib import nullcontext
34. from model.model import Transformer
35. from model.LMConfig import LMConfig
36. from model.dataset import PretrainDataset
39. # 过滤不必要的警告信息（warnings.filterwarnings('ignore')）
40. warnings.filterwarnings('ignore')
43. # 定义Logger日志记录函数，负责记录日志输出。在分布式训练（DDP）中，只有主进程会输出日志信息。
44. def Logger(content):
45.     if not ddp or dist.get_rank() == 0:
46.         print(content)
49. # 实现了get_lr学习率调度器函数，根据当前的迭代次数调整学习率。主要使用余弦退火学习率调度策略，允许动态调整学习率
50. def get_lr(it, all):
51.     warmup_iters = 0
52.     lr_decay_iters = all
53.     min_lr = learning_rate / 10
55.     if it < warmup_iters:
56.         return learning_rate * it / warmup_iters
57.     if it > lr_decay_iters:
58.         return min_lr
59.     decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
60.     assert 0 <= decay_ratio <= 1
61.     coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
62.     return min_lr + coeff * (learning_rate - min_lr)
65. # 训练单个Epoch的实现
66. # 负责一个完整的训练周期的操作：数据加载和模型训练→计算当前的学习率，更新优化器的参数→使用torch.cuda.amp进行自动混合精度训练，进行前向传播与梯度回传→每经过设定的步数后，记录日志并保存模型的状态
67. # 核心技术点：自动混合精度（AMP）、梯度累积、梯度裁剪、分布式数据并行、模型保存
68. def train_epoch(epoch, wandb, accumulation_steps=8):
69.     start_time = time.time()
70.     for step, (X, Y) in enumerate(train_loader):
71.         X = X.to(device)
72.         Y = Y.to(device)
74.         lr = get_lr(epoch * iter_per_epoch + step, epochs * iter_per_epoch)
75.         for param_group in optimizer.param_groups:
76.             param_group['lr'] = lr
78.         with ctx:
79.             out = model(X, Y)
80.             loss = out.last_loss / accumulation_steps
82.         scaler.scale(loss).backward()
84.         if (step + 1) % accumulation_steps == 0:
85.             scaler.unscale_(optimizer)
86.             torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
88.             scaler.step(optimizer)
89.             scaler.update()
91.             optimizer.zero_grad(set_to_none=True)
93.         if step % 100 == 0:
94.             spend_time = time.time() - start_time
95.             Logger(
96.                 'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.7f} epoch_Time:{}min:'.format(
97.                     epoch,
98.                     epochs,
99.                     step,
100.                     iter_per_epoch,
101.                     loss.item() * accumulation_steps,
102.                     optimizer.param_groups[-1]['lr'],
103.                     spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))
105.             if (wandb is not None) and (not ddp or dist.get_rank() == 0):
106.                 wandb.log({"loss": loss.item() * accumulation_steps,
107.                            "lr": optimizer.param_groups[-1]['lr'],
108.                            "epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60})
110.         if (step + 1) % 1000 == 0 and (not ddp or dist.get_rank() == 0):
111.             model.eval()
112.             # torch.save(model.state_dict(), '{}/iter_{}.pth'.format(save_dir, int(step + epoch * iter_per_epoch)))
113.             moe_path = '_moe' if lm_config.use_moe else ''
114.             ckp = f'{save_dir}/pretrain_{lm_config.dim}{moe_path}.pth'
116.             if isinstance(model, torch.nn.parallel.DistributedDataParallel):
117.                 state_dict = model.module.state_dict()
118.             else:
119.                 state_dict = model.state_dict()
121.             torch.save(state_dict, ckp)
122.             model.train()
124. # 模型初始化: init_model函数用于初始化Transformer模型，计算并输出模型的参数数量
125. def init_model():
126.     def count_parameters(model):
127.         return sum(p.numel() for p in model.parameters() if p.requires_grad)
129.     # model init
130.     model = Transformer(lm_config).to(device)
131.     moe_path = '_moe' if lm_config.use_moe else ''
132.     # ckp = f'{save_dir}/pretrain_{lm_config.dim}{moe_path}.pth'
133.     #
134.     # state_dict = torch.load(ckp, map_location=device)
135.     # unwanted_prefix = '_orig_mod.'
136.     # for k, v in list(state_dict.items()):
137.     #     if k.startswith(unwanted_prefix):
138.     #         state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
139.     # model.load_state_dict(state_dict, strict=False)
141.     Logger(f'LLM总参数量：{count_parameters(model) / 1e6:.3f} 百万')
142.     return model
145. # 分布式模式初始化: init_distributed_mode函数用于初始化分布式训练模式，包括设置设备编号、进程组初始化和NCCL后端的配置。
146. # 核心技术点：分布式训练模式的设置（NCCL、DDP）。
147. def init_distributed_mode():
148.     if not ddp: return
149.     global ddp_local_rank, DEVICE
151.     dist.init_process_group(backend="nccl")
152.     ddp_rank = int(os.environ["RANK"])
153.     ddp_local_rank = int(os.environ["LOCAL_RANK"])
154.     ddp_world_size = int(os.environ["WORLD_SIZE"])
155.     DEVICE = f"cuda:{ddp_local_rank}"
156.     torch.cuda.set_device(DEVICE)
159. # torchrun --nproc_per_node 2 1-pretrain.py
160. # I/O
161. if __name__ == "__main__":
162.     # -----------------------------------------------------------------------------
163.     # 1、配置环境（GPU设备、随机种子、数据集路径等）
164.     lm_config = LMConfig()
165.     max_seq_len = lm_config.max_seq_len
166.     out_dir = 'out'
167.     epochs = 20
168.     batch_size = 64
169.     learning_rate = 2e-4
170.     device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
171.     dtype = 'bfloat16'
172.     save_dir = os.path.join(out_dir)
173.     os.makedirs(save_dir, exist_ok=True)
174.     os.makedirs(out_dir, exist_ok=True)
175.     tokens_per_iter = batch_size * max_seq_len
176.     torch.manual_seed(1337)
177.     device_type = device if "cuda" in device else "cpu"
179.     use_wandb = False  # 是否使用wandb
180.     wandb_project = "MiniMind-Pretrain"
181.     wandb_run_name = f"MiniMind-Pretrain-Epoch-{epochs}-BatchSize-{batch_size}-LearningRate-{learning_rate}"
182.     if use_wandb:
183.         import wandb
184.         wandb.init(project=wandb_project, name=wandb_run_name)
185.     else:
186.         wandb = None
188.     ctx = (
189.         nullcontext()
190.         if device_type == "cpu"
191.         else torch.cuda.amp.autocast()
192.     )
193.     ddp = int(os.environ.get("RANK", -1)) != -1  # is this a ddp run?
194.     ddp_local_rank, DEVICE = 0, "cuda:0"
195.     if ddp:
196.         init_distributed_mode()
197.         device = torch.device(DEVICE)
198.     # -----------------------------------------------------------------------------
199.     # 2、训练数据加载器的配置：初始化训练数据集与数据加载器，使用PretrainDataset处理数据，并根据是否为分布式训练设置数据采样器DistributedSampler。
200.     # 核心技术点：数据加载器与分布式数据采样器的配置。
201.     # -----init dataloader------
202.     data_path_list = ['./dataset/pretrain_data.bin']
203.     train_ds = PretrainDataset(data_path_list, max_length=max_seq_len, memmap=True)
204.     train_sampler = DistributedSampler(train_ds) if ddp else None
205.     num_workers = 8  # 可以根据系统的 CPU 核心数来调整
206.     train_loader = DataLoader(
207.         train_ds,
208.         batch_size=batch_size,
209.         pin_memory=True,
210.         drop_last=False,
211.         shuffle=False,
212.         num_workers=num_workers,
213.         sampler=train_sampler
214.     )
216.     # 3、初始化模型与优化器，选择自动混合精度训练和分布式训练模式
217.     # init model
218.     model = init_model()
220.     scaler = torch.cuda.amp.GradScaler(enabled=(dtype == dtype))
221.     # optimizer
222.     optimizer = optim.Adam(model.parameters(), lr=learning_rate)
223.     # compile the model
224.     if False and platform.system() != 'Windows' and float(torch.__version__.split('.')[0]) >= 2:
225.         Logger("compiling the model... (takes a ~minute)")
226.         unoptimized_model = model
227.         model = torch.compile(model)
229.     if ddp:
230.         # Ignore the freqs_cis buffer so that DDP does not broadcast it at
231.         # construction time since NCCL does not support ComplexFloat
232.         model._ddp_params_and_buffers_to_ignore = {"pos_cis"}
233.         model = DistributedDataParallel(model, device_ids=[ddp_local_rank])
235.     # 4、通过for循环控制训练多个Epoch，每个Epoch调用train_epoch函数进行训练
236.     # training loop
237.     iter_per_epoch = len(train_loader)
238.     for epoch in range(epochs):
239.         train_epoch(epoch, wandb)

复制代码

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。