MOE Yuan2.0 预训练 Yuan2.0-M32/examples/pretrain_yuan2.0_moe_2x32B.sh

Yuan2.0 预训练

https://github.com/IEIT-Yuan/Yuan2.0-M32/blob/main/docs/pretrain.md

介绍

本文档提供了Yuan2.0-M32预训练模型的说明。

提供三种型号，主要参数如下：

	层数	隐藏尺寸	注意头	专家编号
2*32B	24	2048	16	三十二

用法

这些脚本描述了Yuan2.0中的三种模型：

• 2B：pretrain_yuan2.0_moe_2x32B.sh

例子

运行 Yuan-2.1B-M32 预训练的示例脚本是：

bash examples/pretrain_yuan2.0_moe_2x32B.sh

参数设置

在运行脚本之前，应正确设置相关参数。

CHECKPOINT_PATH首先，进行任何所需的修改，包括设置、DATA_PATH和 TOKENIZER_MODEL_PATH 的环境变量TENSORBOARD_PATH。

如果数据集路径是：

/path/dataset.bin

可以DATA_PATH设置：

#DATA_PATH='weight dataset_path'
DATA_PATH='1 /path/dataset'

数据集预处理可以参见此处的文档。

--tensor-model-parallel-size和--pipeline-model-parallel-size 标志可以控制简单而高效的三维模型并行方法。如果--pipeline-model-parallel-method将标志设置为，则每个管道阶段的block变压器层数应由指定。--pipeline-model-parallel-blocks

可以通过 ' 标志激活基于局部过滤的注意力 (LFA) --use-lf-gate。并且该--lf-conv2d-num-pad标志应设置1为用于训练和0推理。

并且--use-distributed-optimizer可以--recompute-method控制训练期间内存的使用。

源文件arguments.py和Megatron-LM中描述了进一步的命令行参数

Yuan2.0-M32/examples/pretrain_yuan2.0_moe_2x32B.sh

#!/bin/bash

# Runs the "Yuan-2.1B" parameter model

export CUDA_DEVICE_MAX_CONNECTIONS=1

GPUS_PER_NODE=8
# Change for multinode config
MASTER_ADDR=localhost
MASTER_PORT=6000
NNODES=1
NODE_RANK=0
WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))

CHECKPOINT_PATH=<Specify path>
DATA_PATH=<Specify path and file prefix>_text_document
TOKENIZER_MODEL_PATH=<Specify path to file>
TENSORBOARD_PATH=<Specify path to file>

DISTRIBUTED_ARGS="
    --nproc_per_node $GPUS_PER_NODE \
    --nnodes $NNODES \
    --node_rank $NODE_RANK \
    --master_addr $MASTER_ADDR \
    --master_port $MASTER_PORT
"

GPT_ARGS="
    --tensor-model-parallel-size 1 \
    --pipeline-model-parallel-size 8 \
    --timing-log-level 2 \
    --num-workers 2 \
    --num-layers 24 \
    --hidden-size 2048 \
    --num-attention-heads 16 \
    --kv-channels 256 \
    --use-lf-gate \
    --lf-conv2d-group 1 \
    --lf-conv2d-num-pad 1 \
    --position-embedding-type rope \
    --no-embedding-dropout \
    --flash-attn-drop 0.1 \
    --fim-rate 0.5 \
    --fim-spm-rate 0.5 \
    --norm-dtype RMSNorm \
    --attention-dropout 0 \
    --hidden-dropout 0 \
    --disable-bias-linear \
    --reset-position-ids \
    --use-flash-attn \
    --swiglu \
    --adam-beta1 0.9 \
    --adam-beta2 0.95 \
    --seq-length 4096 \
    --max-position-embeddings 4096 \
    --micro-batch-size 2 \
    --global-batch-size 1536 \
    --lr 0.0001 \
    --train-iters 318000 \
    --lr-decay-iters 318000 \
    --lr-decay-style cosine \
    --min-lr 1.0e-5 \
    --weight-decay 1e-1 \
    --lr-warmup-iters 6400 \
    --clip-grad 1.0 \
    --recompute-method block \
    --recompute-granularity full \
    --recompute-num-layers 1 \
    --bf16 \
    --rotary-percent 0.5 \
    --use-attention-router \
    --no-masked-softmax-fusion \
    --use-fp32-router \
    --num-experts 32 \
    --moe-router-load-balancing-type none \
    --moe-router-topk 2 \
    --moe-grouped-gemm \

    
"


DATA_ARGS="
    --data-path $DATA_PATH \
    --tokenizer-type YuanTokenizer \
    --tokenizer-model-path $TOKENIZER_MODEL_PATH \
    --data-impl mmap \
    --split 10,0,0
"

OUTPUT_ARGS="
    --log-interval 1 \
    --save-interval 10000 \
    --eval-interval 1000000 \
    --eval-iters 10
"

LOG_ARGS="
    --tensorboard-dir $TENSORBOARD_PATH \
    --tensorboard-log-interval 1 \
    --tensorboard-queue-size 1000 \
    --log-timers-to-tensorboard \
    --log-batch-size-to-tensorboard \
    --log-memory-to-tensorboard \
    --log-world-size-to-tensorboard
"

torchrun $DISTRIBUTED_ARGS pretrain_yuan.py \
    $GPT_ARGS \
    $DATA_ARGS \
    $OUTPUT_ARGS \
    $LOG_ARGS \
    --distributed-backend nccl \
    --save $CHECKPOINT_PATH \
    --load $CHECKPOINT_PATH

下面是逐行解释 "Yuan-2.1B" 参数模型预训练脚本的详细说明：

bash

复制

# 设置 CUDA 设备的最大连接数为 1
export CUDA_DEVICE_MAX_CONNECTIONS=1

# 每台节点上的 GPU 数量
GPUS_PER_NODE=8
# 多节点配置相关参数
MASTER_ADDR=localhost  # 主节点地址
MASTER_PORT=6000       # 主节点端口
NNODES=1               # 节点数量
NODE_RANK=0            # 当前节点的排名
WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))  # 全局 GPU 数量（所有节点的总和）

# 检查点保存路径
CHECKPOINT_PATH=<Specify path>
# 数据路径和文件前缀
DATA_PATH=<Specify path and file prefix>_text_document
# Tokenizer 模型路径
TOKENIZER_MODEL_PATH=<Specify path to file>
# TensorBoard 日志保存路径
TENSORBOARD_PATH=<Specify path to file>

# 分布式训练参数
DISTRIBUTED_ARGS="
    --nproc_per_node $GPUS_PER_NODE \  # 每个节点上的进程数
    --nnodes $NNODES \                 # 节点总数
    --node_rank $NODE_RANK \           # 当前节点的排名
    --master_addr $MASTER_ADDR \       # 主节点地址
    --master_port $MASTER_PORT         # 主节点端口
"

# GPT 模型参数
GPT_ARGS="
    --tensor-model-parallel-size 1 \            # 张量模型并行大小
    --pipeline-model-parallel-size 8 \          # 管道模型并行大小
    --timing-log-level 2 \                      # 时间记录级别
    --num-workers 2 \                           # 工作线程数
    --num-layers 24 \                           # 模型层数
    --hidden-size 2048 \                        # 隐藏层大小
    --num-attention-heads 16 \                  # 注意力头数量
    --kv-channels 256 \                         # 键/值通道数量
    --use-lf-gate \                             # 使用 LF 门
    --lf-conv2d-group 1 \                       # LF 卷积组数
    --lf-conv2d-num-pad 1 \                     # LF 卷积填充数
    --position-embedding-type rope \            # 位置嵌入类型
    --no-embedding-dropout \                    # 不使用嵌入丢弃
    --flash-attn-drop 0.1 \                     # Flash 注意力丢弃率
    --fim-rate 0.5 \                            # FIM 率
    --fim-spm-rate 0.5 \                        # FIM SPM 率
    --norm-dtype RMSNorm \                      # 归一化类型
    --attention-dropout 0 \                     # 注意力丢弃率
    --hidden-dropout 0 \                        # 隐藏层丢弃率
    --disable-bias-linear \                     # 禁用偏置线性
    --reset-position-ids \                      # 重置位置 ID
    --use-flash-attn \                          # 使用 Flash 注意力
    --swiglu \                                  # 使用 SwiGLU 激活函数
    --adam-beta1 0.9 \                          # Adam 优化器 beta1 参数
    --adam-beta2 0.95 \                         # Adam 优化器 beta2 参数
    --seq-length 4096 \                         # 序列长度
    --max-position-embeddings 4096 \            # 最大位置嵌入
    --micro-batch-size 2 \                      # 微批次大小
    --global-batch-size 1536 \                  # 全局批次大小
    --lr 0.0001 \                               # 学习率
    --train-iters 318000 \                      # 训练迭代次数
    --lr-decay-iters 318000 \                   # 学习率衰减迭代次数
    --lr-decay-style cosine \                   # 学习率衰减方式（余弦）
    --min-lr 1.0e-5 \                           # 最小学习率
    --weight-decay 1e-1 \                       # 权重衰减
    --lr-warmup-iters 6400 \                    # 学习率预热迭代次数
    --clip-grad 1.0 \                           # 梯度裁剪阈值
    --recompute-method block \                  # 重计算方法（块）
    --recompute-granularity full \              # 重计算粒度（完整）
    --recompute-num-layers 1 \                  # 重计算层数
    --bf16 \                                    # 使用 bfloat16 精度
    --rotary-percent 0.5 \                      # 旋转百分比
    --use-attention-router \                    # 使用注意力路由
    --no-masked-softmax-fusion \                # 不使用掩码 Softmax 融合
    --use-fp32-router \                         # 使用 FP32 路由
    --num-experts 32 \                          # 专家数量
    --moe-router-load-balancing-type none \     # MoE 路由负载均衡类型
    --moe-router-topk 2 \                       # MoE 路由 Top-K
    --moe-grouped-gemm \                        # 使用分组 GEMM
"

# 数据相关参数
DATA_ARGS="
    --data-path $DATA_PATH \                    # 数据路径
    --tokenizer-type YuanTokenizer \            # Tokenizer 类型
    --tokenizer-model-path $TOKENIZER_MODEL_PATH \ # Tokenizer 模型路径
    --data-impl mmap \                          # 数据实现方式（内存映射）
    --split 10,0,0                              # 数据划分比例
"

# 输出相关参数
OUTPUT_ARGS="
    --log-interval 1 \                          # 日志记录间隔
    --save-interval 10000 \                     # 模型保存间隔
    --eval-interval 1000000 \                   # 评估间隔
    --eval-iters 10                             # 评估迭代次数
"

# 日志相关参数
LOG_ARGS="
    --tensorboard-dir $TENSORBOARD_PATH \       # TensorBoard 日志目录
    --tensorboard-log-interval 1 \              # TensorBoard 日志记录间隔
    --tensorboard-queue-size 1000 \             # TensorBoard 队列大小
    --log-timers-to-tensorboard \               # 将计时器记录到 TensorBoard
    --log-batch-size-to-tensorboard \           # 将批次大小记录到 TensorBoard
    --log-memory-to-tensorboard \               # 将内存使用记录到 TensorBoard
    --log-world-size-to-tensorboard             # 将世界大小记录到 TensorBoard
"

# 运行分布式训练
torchrun $DISTRIBUTED_ARGS pretrain_yuan.py \
    $GPT_ARGS \
    $DATA_ARGS \
    $OUTPUT_ARGS \
    $LOG_ARGS \
    --distributed-backend nccl \                # 分布式后端使用 NCCL
    --save $CHECKPOINT_PATH \                   # 保存检查点路径
    --load $CHECKPOINT_PATH                     # 加载检查点路径

这个脚本配置并运行了一个名为 "Yuan-2.1B" 的大规模预训练模型，涉及多个 GPU 和可能的多节点分布式训练。主要步骤包括：

1. 设置环境变量：如 CUDA_DEVICE_MAX_CONNECTIONS。
2. 定义节点和 GPU 配置：包括每台节点上的 GPU 数量、节点数量、主节点地址和端口等。
3. 指定检查点和数据路径：设置模型检查点保存路径、数据路径、Tokenizer 模型路径和 TensorBoard 日志路径。
4. 配置分布式参数：如每个节点的进程数、节点总数、当前节点排名等。
5. 配置模型参数：详细设置 GPT 模型的各项参数，包括层数、隐藏层大小、注意力头数量、学习率、批次大小等。
6. 配置数据参数：如数据路径、Tokenizer 类型和模型路径、数据实现方式和数据划分比例等。
7. 配置输出参数：日志记录间隔、模型保存间隔、评估间隔等。
8. 配置日志参数：TensorBoard 日志目录、记录间隔、队列大小等。
9. 运行分布式训练：最终通过 torchrun 命令执行分布式训练脚本 pretrain_yuan.py，并传入所有的配置参数。

思维链（扩展阅读）：

分布式训练工具torchrun、accelerate、deepspeed、Megatron-CSDN博客

->

Megatron-LM github主页-CSDN博客