LLMs之Colossal-LLaMA-2：源码解读(train.py文件)基于给定数据集实现持续预训练LLaMA-2—解析命令行参数→初始化配置(分布式训练环境colossalai+训练日志+加速插

LLMs之Colossal-LLaMA-2：源码解读(train.py文件)基于给定数据集实现持续预训练LLaMA-2—解析命令行参数→初始化配置(分布式训练环境colossalai+训练日志+加速插件)→数据预处理(初始化分词器+数据处理器+数据加载器)→模型训练(初始化模型/优化器/学习率调度器/梯度检查点/Flash-Attention/设置数据类型/是否加载预训练模型/从上一次训练点继续训练+开启训练循环【分布式训练汇总全局均值/模型定期保存/等】)

目录

源码解读(train.py文件)基于给定数据集实现持续预训练LLaMA-2—解析命令行参数→初始化配置(分布式训练环境colossalai+训练日志+加速插件)→数据预处理(初始化分词器+数据处理器+数据加载器)→模型训练(初始化模型/优化器/学习率调度器/梯度检查点/Flash-Attention/设置数据类型/是否加载预训练模型/从上一次训练点继续训练+开启训练循环【分布式训练汇总全局均值/模型定期保存/等】)

0、定义辅助函数：

get_model_numel(model)：计算模型中参数的总数

format_numel_str(numel)：将参数数量格式化为易读的字符串形式

all_reduce_mean(tensor)：使用分布式计算平均值，这是一个用于分布式训练的函数

# 1、解析命令行参数

# 1.1、包括预训练模型、数据集、加速插件，检查点相关选项(载入/保存模型的间隔/检查点保存的目录)，指定Tensorboard日志保存的目录、指定配置文件的路径，训练的轮数、每个进程的批量大小、学习率、最大长度，选择混合精度的模式、梯度裁剪的阈值、权重衰减系数、预热步数，是否使用梯度检查点、是否使用 flash-attention、是否冻结非嵌入参数；

# 1.2、创建配置文件并保存，记录了所有的命令行参数

# 2、初始化配置

# 2.1、初始化分布式训练环境，启动分布式训练任务，创建DistCoordinator实例协调和管理分布式训练中的各个任务(如同步/通信等)

# (1)、启动分布式深度学习任务：colossalai 是一个与 PyTorch 配合使用的库，用于简化分布式训练的管理和设置。

# (2)、协调和管理分布式训练中的各个任务。它可以用来进行各种与分布式训练相关的操作，如同步、通信等。

# 2.2、初始化 Tensorboard训练日志，用于记录训练过程中的指标

# 2.3、初始化 Booster，选择训练加速插件（例如 GeminiPlugin、LowLevelZeroPlugin）。

# 3、数据预处理

# 3.1、初始化分词器

# (1)、从预训练模型中加载了一个分词器，用于将文本转换为模型可以理解的输入

# (2)、将分词器的填充（pad）标记设置为unk标记，这通常是为了在模型输入中使用相同的标记来表示填充。

# (3)、分别将分词器的开始BOS标记和结束EOS标记的添加设置为False，意味着不会在模型输入中添加这些特殊标记。

# (4)、打印一些配置信息到主节点（master node）上，配置文件的保存路径、Tensorboard日志的保存路、模型检查点将保存的路径、待加载的数据集

# 3.2、初始化数据集以及数据处理器和数据加载器,构建DataLoader供训练使用

# 数据处理器，用于将数据批次（batch）中的样本整理成模型所需的输入格式。它接受一个分词器 tokenizer 和一个最大长度参数

# 数据加载器，用于在训练过程中迭代数据批次，包每个批次的样本数、是否在每个epoch之前打乱数据、是否在最后一个批次中舍弃不足一个批次的样本

# 3.3、打消息到主节点上，消息包括 CUDA 显存的最大使用量（以 MB 为单位）。这可以用于监控训练过程中 GPU 内存的使用情况

# 4、模型训练与评估

# 4.1、初始化模型、优化器和学习率调度器，可以选择是否启用梯度检查点和 Flash-Attention

# (1)、初始化init_ctx上下文管理器，用于控制初始化模型的上下文。如果使用的并行处理插件是GeminiPlugin,那么使用LazyInitContext来延迟初始化；否则直接使用nullcontext。

# (2)、if判断是否启用梯度检查点技术，用于减少显存占用的技术，通常在模型很大的情况下使用。

# (3)、if判断是否启用了 Flash-Attention技术，这是一种特殊类型的注意力机制。

# (4)、统计模型总参数量并输出

# (5)、初始化HybridAdam优化器,并根据是否冻结非嵌入层参数来选择优化器的参数

# (6)、初始化CosineAnnealingWarmup学习率调度器：学习率衰减策略

# 4.2、设置数据类型

# (1)、设置默认数据类型为混合精度类型(半精度浮点数，如fp16或bf16)

# (2)、通过 booster 对象的 boost方法来进行模型、优化器、学习率调度器和数据加载器的初始化和优化。

# (3)、恢复默认数据类型为标准的单精度浮点数 torch.float

# 4.3、加载预训练模型（如果有的话），包括前向、反向传播、优化器更新、检查点保存等步骤

# (1)、if判断是否从预训练模型加载初始权重

# (2)、打印 Booster 初始化后 GPU 和 CPU 内存占用

# 4.4、确保加载最新的训练状态,从上一次训练点继续训练。

# (1)、if判断如果是预训练模型地址,直接加载模型参数；否则,调用函数加载最新的训练状态,包括 epoch、step 数、样本索引等

# (2)、打印加载的 checkpoint 信息

# (3)、检查加载 checkpoint 对内存占用的影响

# (4)、计算每个训练周期中的步骤数量，这个值等于数据加载器 dataloader 的长度，即数据集的总批次数。

# (5)、确保数据加载器的采样器是 StatefulDistributedSampler 的实例。这是为了保证后续可以设置采样器的状态。

# (6)、设置数据加载器的采样器的起始索引为 sampler_start_idx。这是为了从正确的数据位置开始训练，用于从检查点中恢复或继续训练。

# 4.5、模型训练，开始训练循环，包括多个训练周期。

# (1)、在每个新的训练周期开始时，设置数据加载器的采样器的当前周期为 epoch。这是为了确保数据加载器按照周期重新排列数据，以增加训练的多样性。

# (2)、创建一个进度条，迭代数据加载器中的批次数据，并显示训练进度。

# (3)、内部循环，处理每个批次数据

# 3.1)、将批次数据中的张量移动到当前设备上，一般是加载到GPU上以进行加速

# 3.2)、使用模型前向传播计算批次的输出，其中**batch表示将批次数据解包为关键字参数传递给模型

# 3.3)、从模型的输出中提取损失值

# 3.4)、根据计算的损失值执行反向传播，计算梯度

# 3.5)、执行优化器的一步参数更新

# 3.6)、执行学习率调度器的一步，更新学习率

# 3.7)、清零模型参数的梯度，为下一个批次的训练做准备

# 3.8)、基于分布式训练时，将损失值进行全局均值汇总，以便在多个设备上计算的损失值得到同步

# 3.9)、更新进度条的后缀信息，显示当前批次的损失值

# 3.10)、检查当前进程是否为主进程，如果是主进程，则执行，计算全局步数 global_step记录训练的总步数、使用Tensorboard的writer记录损失值和学习率等训练信息。

# 3.11)、if判断定期保存模型检查点，记录损失和学习率等信息到 Tensorboard。

# 3.12)、打印一条日志，指示开始保存模型检查点，并记录当前训练状态。

# 3.13)、调用 save_checkpoint() 函数保存模型的检查点，包括模型、优化器、学习率调度器等信息

# 3.14)、使用 coordinator.print_on_master() 打印保存检查点的信息，包括保存的周期和步骤。

# 3.15)、删除CUDA缓存

# (4)、在每个训练周期结束后，需要进行一些重置操作

# 4.6、训练结束后，最后保存最终的模型检查点

# (1)、打印一条日志，指示开始保存最终的模型检查点。

# (2)、保存模型的最终检查点，将模型保存到指定的文件夹中，文件夹名称为 "modeling"。

# (3)、打印保存最终检查点的信息，包括保存的周期信息

# 4.7、打印一条日志，记录最大的 CUDA 内存使用量，以MB为单位。这个值表示模型训练期间 CUDA 设备的最大内存使用量

实现代码

---