当前位置:AIGC资讯 > AIGC > 正文

使用deepspeed继续训练LLAMA

目录

1. 数据训练配置

2. 模型载入

3. 优化器设置

4. DeepSpeed 设置

5. DeepSpeed 初始化

6. 模型训练

LLAMA 模型子结构:

1. 数据训练配置

利用 PyTorch 和 Transformers 库创建数据加载器,它支持单机或多机分布式训练环境下的数据加载与采样。涉及的模块包括:

DataLoader: 由 PyTorch 提供,用于数据集到模型的数据加载。 RandomSampler 和 SequentialSampler: PyTorch 提供的随机和顺序数据采样器。 DistributedSampler: 专为分布式训练设计的采样器。 default_data_collator: Transformers 库的默认数据整合器,用于批量数据处理。 create_pretrain_dataset: 创建预训练数据集的自定义函数。

根据 args.local_rank 的值,选择单机采样器或分布式采样器。DistributedSampler 确保每个训练节点获得唯一数据子集,而单机环境下则使用随机或顺序采样器。

2. 模型载入

通过 Transformers 库,加载并配置 LLaMA 模型及其分词器。使用 from_pretrained 方法加载预训练模型、分词器和配置。设置分词器以处理不同文本长度,并设定填充符号为 [PAD],确保填充发生在句子右侧。模型配置中也设置了句子结束和填充符号的 ID,并优化了词汇表嵌入大小以提升硬件性能。

3. 优化器设置

DeepSpeed 库提供了优化的优化器算法,如 DeepSpeedCPUAdam 和 FusedAdam,提高了大规模数据和模型训练速度。优化器设置涉及:

参数分组: 通过 get_optimizer_grouped_parameters 函数实现参数分组,一组应用权重衰减,另一组不应用。 优化器选择: 根据训练环境选择 DeepSpeedCPUAdam 或 FusedAdam。 学习率调度: 动态调整学习率,考虑预热步骤和总训练步数。
4. DeepSpeed 设置

定义全局批次大小 (GLOBAL_BATCH_SIZE) 和每 GPU 微批次大小 (MICRO_BATCH_SIZE)。get_train_ds_config 训练配置函数包括:

ZeRO 优化: 减少冗余并加速训练。 混合精度训练: 通过设置 fp16 字段使用 16 位浮点数。 梯度裁剪: 防止梯度爆炸。 混合引擎配置: 优化输出分词数量和张量大小。 TensorBoard 集成: 方便跟踪训练过程。

get_eval_ds_config 函数提供简洁的验证集配置,专注于模型推理。

5. DeepSpeed 初始化

初始化包括:

设备确定: 检查本地 GPU 或使用 CUDA。 分布式后端初始化: 使用 deepspeed.init_distributed() 同步进程。 设置 DeepSpeed 配置: 根据用户参数构建训练设置。 同步工作进程: 使用 torch.distributed.barrier() 确保进程同步。 初始化: 通过 deepspeed.initialize 优化模型和优化器。 梯度检查点: 启用时,使用 model.gradient_checkpointing_enable()。
6. 模型训练

DeepSpeed 框架下的训练步骤:

训练前准备: 使用 print_rank_0 函数输出训练状态,避免多进程重复输出。 训练循环: 打印周期信息,进行前向传播、梯度计算和参数更新。 模型保存: 保存模型状态和配置,支持 Hugging Face 和 DeepSpeed Zero Stage 3 格式。

更新时间 2024-02-15