Megatron-LM和Megatron-DeepSpeed的这个分支是什么

这是https://github.com/microsoft/Megatron-DeepSpeed的一个独立分支，而该仓库本身是https://github.com/NVIDIA/Megatron-LM的一个分支。前者将DeepSpeed集成到了原始的Megatron-LM代码中。

这个分支将包含BigScience项目所需的对模型的直接修改。这是我们用于该项目的仓库。

此外，在https://github.com/bigscience-workshop/bigscience中可以找到各种代码片段和大量文档。

请注意，此页面的其余部分已经被精简，只包含与BigScience项目相关的信息，并更新为与集成的Deepspeed一起使用。您可以在这里找到原始页面，其中包含所有表格和关于Bert和T5训练的信息。

快速入门

这里有一个文档，只包含从零开始到快速开始训练的说明。

设置

安装bigscience-workshop/Megatron-DeepSpeed

git clone https://github.com/bigscience-workshop/Megatron-DeepSpeed
cd Megatron-DeepSpeed
pip install -r requirements.txt

现在您可以直接使用这个仓库，直接在其中工作。除非您在其他地方编写自己的脚本并使用此仓库中的模块，否则不需要安装它。如果需要安装，可以执行：

pip install -e .

安装apex

git clone https://github.com/NVIDIA/apex
cd apex
pip install --global-option="--cpp_ext" --global-option="--cuda_ext" --no-cache -v --disable-pip-version-check .  2>&1 | tee build.log

(在JZ上，需要以特殊方式安装，请参见这里。)

安装deepspeed

git clone https://github.com/microsoft/deepspeed
cd deepspeed
rm -rf build
TORCH_CUDA_ARCH_LIST="7.0" DS_BUILD_CPU_ADAM=1 DS_BUILD_AIO=1 DS_BUILD_UTILS=1 pip install -e . --global-option="build_ext" --global-option="-j8" --no-cache -v --disable-pip-version-check

将TORCH_CUDA_ARCH_LIST="7.0"调整为您的NVIDIA GPU的架构（如果不确定如何找到，可以完全删除它）。

(在JZ上，需要以特殊方式安装，请参见这里。)

CUDA内核编译

第一次运行训练脚本时，将编译几个CUDA内核。这意味着您需要在环境中设置cuda环境，并且它应该与pytorch构建时使用的版本匹配。

使用

安装后，有几种可能的工作流程。最全面的是：

数据预处理
预训练
微调（对于零样本任务是可选的）
下游任务评估或文本生成

然而，步骤1和2可以通过使用上述预训练模型之一来替代。

我们在examples目录中提供了几个用于预训练BERT和GPT的脚本，以及用于零样本和微调下游任务的脚本，包括MNLI、RACE、WikiText103和LAMBADA评估。还有一个用于GPT交互式文本生成的脚本。

训练

词汇表

GPT的词汇文件和合并表可以直接下载。

数据预处理

训练数据需要预处理。首先，将您的训练数据放在松散的json格式中，每行包含一个文本样本的json。例如：

{"src": "www.nvidia.com", "text": "The quick brown fox", "type": "Eng", "id": "0", "title": "First Part"}
{"src": "The Internet", "text": "jumps over the lazy dog", "type": "Eng", "id": "42", "title": "Second Part"}

可以通过在preprocess_data.py中使用--json-key标志来更改json的text字段的名称。其他元数据是可选的，不用于训练。

然后将松散的json处理成二进制格式以进行训练。要将json转换为mmap、缓存索引文件或惰性加载器格式，请使用preprocess_data.py。将--dataset-impl标志设置为mmap、cached或lazy（默认为mmap）。

用于准备GPT训练数据的示例脚本如下：

python tools/preprocess_data.py \
    --input my-corpus.json \
    --output-prefix my-gpt2 \
    --vocab-file gpt2-vocab.json \
    --dataset-impl mmap \
    --tokenizer-type GPT2BPETokenizer \
    --merge-file gpt2-merges.txt \
    --append-eod \
    --workers 8

输出将是两个文件，在这种情况下命名为my-gpt2_text_document.bin和my-gpt2_text_document.idx。后续GPT训练中指定的--data-path是完整路径和新文件名，但不包括文件扩展名。

源文件preprocess_data.py中描述了更多命令行参数。

如果有大量可用的CPU核心，您也可以使用tools/preprocess_data_many_cores.py。通常在JZ设置中，CPU节点最多有40个物理CPU核心，您应该使用大约60个工作线程运行此脚本，而不是tools/preprocess_data.py。可用的命令行参数相同。

合并数据集

有时很难一次处理非常大的数据集，因此可以分块预处理它，然后将这些数据集合并为一个组合的索引数据集。以下是一个示例：

python tools/merge_preprocessed_data.py \
    --datasets \
    meg-gpt2-oscar-en-500-p1_text_document \
    meg-gpt2-oscar-en-500-p2_text_document \
    meg-gpt2-oscar-en-500-p3_text_document \
    --output-prefix meg-gpt2_oscar_text_document

快速预处理以开始训练

以下是如何使用1GB的79K记录jsonl数据集快速准备训练：

wget https://huggingface.co/bigscience/misc-test-data/resolve/main/stas/oscar-1GB.jsonl.xz
wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-vocab.json
wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-merges.txt
xz -d oscar-1GB.jsonl.xz
python tools/preprocess_data.py \
    --input oscar-1GB.jsonl \
    --output-prefix my-gpt2 \
    --vocab-file gpt2-vocab.json \
    --dataset-impl mmap \
    --tokenizer-type GPT2BPETokenizer \
    --merge-file gpt2-merges.txt \
    --append-eod \
    --workers 8

GPT预训练

注意：您可能想跳到下一节，因为它描述了我们目前实际使用的内容。 examples/pretrain_gpt.sh脚本运行单GPU 345M参数GPT预训练。单GPU训练主要用于调试，因为代码库和命令行参数针对高度分布式训练进行了优化。大多数参数都是自解释的。默认情况下，学习率在训练迭代过程中从--lr开始线性衰减，在--lr-decay-iters次迭代中达到--min-lr设置的最小值。--lr-warmup-fraction设置用于预热的训练迭代比例。虽然这是单GPU训练，但--micro-batch-size指定的批量大小是单次前向-反向路径的批量大小，代码会进行梯度累积步骤，直到达到global-batch-size，即每次迭代的批量大小。

数据按949:50:1的比例分为训练/验证/测试集（默认为969:30:1）。这种分区是即时进行的，但在相同的随机种子下（默认为1234，或使用--seed手动指定）在不同运行中保持一致。我们使用train-iters作为请求的训练迭代次数。或者，可以提供--train-samples，即要训练的总样本数。如果使用此选项，则需要提供--lr-decay-samples而不是--lr-decay-iters。

指定了日志记录、检查点保存和评估间隔。对激活进行检查点有助于训练更大的模型和/或批次。注意，--data-path现在包含预处理中添加的额外_text_sentence后缀，但不包含文件扩展名。

使用的分词方案是BPE（需要合并表和json词汇文件），模型架构允许更长的序列（请注意，最大位置嵌入必须大于或等于最大序列长度），并且--lr-decay-style设置为余弦衰减。注意，--data-path现在包含预处理中添加的额外_text_document后缀，但不包含文件扩展名。

然而，如下所示，你会发现即使使用单个GPU，DeepSpeed也需要分布式环境。因此，请参考pretrain_gpt_single_node.sh，它可以在此仓库中使用。

CHECKPOINT_PATH=checkpoints/gpt2
VOCAB_FILE=gpt2-vocab.json
MERGE_FILE=gpt2-merges.txt
DATA_PATH=my-gpt2_text_document

GPT_ARGS=" \
    --num-layers 24 \
    --hidden-size 1024 \
    --num-attention-heads 16 \
    --seq-length 1024 \
    --max-position-embeddings 1024 \
    --micro-batch-size 4 \
    --global-batch-size 8 \
    --lr 0.00015 \
    --train-iters 500000 \
    --lr-decay-iters 320000 \
    --lr-decay-style cosine \
    --vocab-file $VOCAB_FILE \
    --merge-file $MERGE_FILE \
    --lr-warmup-fraction .01 \
    --fp16 \
    "

OUTPUT_ARGS=" \
    --log-interval 10 \
    --save-interval 500 \
    --eval-interval 100 \
    --eval-iters 10 \
    --checkpoint-activations \
    "

DATA_ARGS=" \
    --save $CHECKPOINT_PATH \
    --load $CHECKPOINT_PATH \
    --data-path $DATA_PATH \
    "

CMD="pretrain_gpt.py GPTARGSGPT_ARGS OUTPUT_ARGS $DATA_ARGS"

N_GPUS=1

LAUNCHER="deepspeed --num_gpus $N_GPUS"

LAUNCHERLAUNCHER CMD

注意，我们用deepspeed --num_gpus 1替换了python。对于多GPU训练，请将--num_gpus更新为你拥有的GPU数量。

对于多节点训练，你需要创建一个定义所有节点的hostfile，如这里所解释的那样，或者在SLURM环境中可能不起作用，你需要使用：

CMD=<同上>

MASTER_ADDR=`perl -le '$_=$ENV{"SLURM_JOB_NODELIST"}; s/,.*//; s/-.*//; s/\[//; print'`
MASTER_PORT=6000
GPUS_PER_NODE=4
NNODES=16

export LAUNCHER="python -u -m torch.distributed.launch \
    --nproc_per_node $GPUS_PER_NODE \
    --nnodes $NNODES \
    --master_addr $MASTER_ADDR \
    --master_port $MASTER_PORT \
    "

srun --jobid $SLURM_JOBID bash -c '$LAUNCHER --node_rank $SLURM_PROCID $CMD'

对于单GPU，另一种方法是模拟distributed：

MASTER_ADDR=localhost MASTER_PORT=9994 RANK=0 LOCAL_RANK=0 python pretrain_gpt.py ...

更多命令行参数在源文件arguments.py中有描述。

Deepspeed PP和ZeRO-DP

为了提供更大的灵活性，我们使用Deepspeed PP（流水线并行）和ZeRO-DP，同时保留Megatron的正常功能。也就是说，我们用Deepspeed的PP替换Megatron的PP，并使用ZERO-DP进行DP。

这与普通的Megatron-LM启动器类似，但增加了一个deepspeed配置文件和一些参数：

CHECKPOINT_PATH=checkpoints/gpt2
VOCAB_FILE=data/gpt2-vocab.json
MERGE_FILE=data/gpt2-merges.txt
DATA_PATH=data/meg-gpt2_oscar-combined_text_document
TENSORBOARD_PATH=output_dir/tensorboard
CODECARBON_PATH=output_dir/codecarbon

MICRO_BATCH_SIZE=1
GLOBAL_BATCH_SIZE=16
TP_SIZE=1
PP_SIZE=1

N_GPUS=2
SAVE_INTERVAL=100

#    --train-samples 10_000 \
#    --exit-interval $EXIT_INTERVAL \

#    --exit-interval 100 \
GPT_ARGS=" \
    --num-layers 2 \
    --hidden-size 64 \
    --num-attention-heads 2 \
    --seq-length 1024 \
    --max-position-embeddings 1024 \
    --micro-batch-size $MICRO_BATCH_SIZE \
    --rampup-batch-size 2 2 1_000 \
    --global-batch-size $GLOBAL_BATCH_SIZE \
    --train-samples 100 \
    --optimizer adam \
    --adam-beta1 0.9 \
    --adam-beta2 0.95 \
    --adam-eps 1e-8 \
    --lr 1e-4 \
    --lr-warmup-samples 5 \
    --clip-grad 1.0 \
    --weight-decay 1e-1 \
    --vocab-file $VOCAB_FILE \
    --merge-file $MERGE_FILE \
    --fp16 \
    "
#    --train-iters 500 \

OUTPUT_ARGS=" \
    --log-interval 10 \
    --save-interval $SAVE_INTERVAL \
    --eval-interval 100 \
    --eval-iters 10 \
    --checkpoint-activations \
    "

DATA_ARGS="
--save $CHECKPOINT_PATH
--load $CHECKPOINT_PATH
--data-path $DATA_PATH
--tensorboard-dir $TENSORBOARD_PATH
--tensorboard-queue-size 5
--log-timers-to-tensorboard
--log-batch-size-to-tensorboard
--log-validation-ppl-to-tensorboard
"

ZERO_STAGE=1

config_json="./ds_config.json"

Deepspeed通过set_train_batch_size()动态确定GAS

cat <<EOT > $config_json { "train_micro_batch_size_per_gpu": $MICRO_BATCH_SIZE, "train_batch_size": $GLOBAL_BATCH_SIZE, "gradient_clipping": 1.0, "zero_optimization": { "stage": $ZERO_STAGE }, "fp16": { "enabled": true, "loss_scale": 0, "loss_scale_window": 500, "hysteresis": 2, "min_loss_scale": 1, "initial_scale_power": 12 }, "steps_per_print": 2000, "wall_clock_breakdown": false } EOT

DEEPSPEED_ARGS="
--deepspeed
--deepspeed_config ${config_json}
--zero-stage ${ZERO_STAGE}
--deepspeed-activation-checkpointing
"

ALL_ARGS="$GPT_ARGS $OUTPUT_ARGS $DATA_ARGS $DEEPSPEED_ARGS"

如果你无法忍受pt-1.9启动器的噪音

export LOGLEVEL=WARNING

LAUNCHER="deepspeed --num_gpus $N_GPUS" export CMD="
$LAUNCHER pretrain_gpt.py
--tensor-model-parallel-size $TP_SIZE
--pipeline-model-parallel-size $PP_SIZE
--distributed-backend nccl
$ALL_ARGS
"

echo $CMD

$CMD

在JZ上我们使用不同的启动命令，例如参见tr1-13B-round1.slurm的末尾，但这也是一个完全可用的脚本。只是它是为SLURM环境编写的。

使用任何预训练的分词器

感谢@sbmaruf，现在可以使用任何HF预训练的分词器来代替Megatron提供的BERT/GPT/T5分词器。你需要自己运行预处理(tools/preprocess_data.py)，使用tokenizer-type=PretrainedFromHF和tokenizer-name-or-path=<your_tokenizer>。例如，python tools/preprocess_data.py --input ~/c4_en_train.jsonl --output-prefix c4_en_train --dataset-impl mmap --tokenizer-type PretrainedFromHF --tokenizer-name-or-path t5-small --workers 30 --append-eod

分布式预训练

examples/pretrain_{bert,gpt,t5}_distributed.sh脚本使用PyTorch分布式启动器进行分布式训练。因此，通过正确设置环境变量并在启动器中使用init_method='env://'，可以实现多节点训练。有关这些环境变量的进一步描述，请参阅官方PyTorch文档。默认情况下，多节点训练使用nccl分布式后端。只需要一组简单的额外参数和使用Python标志-m torch.distributed.launch的PyTorch分布式模块，就可以采用分布式训练。

我们使用两种类型的并行性：数据并行和模型并行。我们提供了两种分布式数据并行实现：一种是我们自己的简单实现，在反向传播步骤结束时执行梯度全归约；另一种是Torch的分布式数据并行包装器，它将梯度归约与反向传播计算重叠。要在这两个选项之间切换，请分别使用--DDP-impl local或--DDP-impl torch。正如预期的那样，Torch分布式数据并行在较大的模型规模下更加高效。例如，对于在512个GPU上运行的83亿参数模型，当使用Torch的分布式数据并行时，扩展性从60%提高到76%。然而，重叠方法需要更多的内存，对于某些配置（例如，使用2路模型并行的25亿参数和没有模型并行的12亿参数）可能会使整体训练变慢。我们通过实验发现，在这些情况下使用较小的模型可以改善训练时间。

其次，我们开发了一种简单高效的二维模型并行方法。要使用张量模型并行（在多个GPU上分割单个transformer模块的执行），请添加--tensor-model-parallel-size标志来指定要分割模型的GPU数量，以及上面提到的传递给分布式启动器的参数。要使用流水线模型并行（将transformer模块分成具有相等数量transformer模块的阶段，然后通过将批次分解为更小的微批次来流水线执行），请使用--pipeline-model-parallel-size标志来指定要将模型分割成的阶段数（例如，将具有24个transformer层的模型分成4个阶段，意味着每个阶段获得6个transformer层）。

我们在以distributed_with_mp.sh结尾的示例脚本中展示了如何使用这两种不同形式的模型并行，请注意，T5模型当前不支持流水线并行：

除了这些微小的变化外，分布式训练与单GPU训练相同。

分布式训练：

有关如何使用deepspeed启动器进行分布式训练的详细信息，请参见上面几节 XXX: 以下内容需要更新：

WORLD_SIZE=8
TENSOR_MP_SIZE=2
PIPELINE_MP_SIZE=2

DISTRIBUTED_ARGS="--nproc_per_node $WORLD_SIZE \
    --nnodes 1 \
    --node_rank 0 \
    --master_addr localhost \
    --master_port 6000"

CHECKPOINT_PATH=<与上面相同>
VOCAB_FILE=<与上面相同>
DATA_PATH=<与上面相同>
MODEL_ARGS=<与上面相同>
OUTPUT_ARGS=<与上面相同>

python -m torch.distributed.launch $DISTRIBUTED_ARGS ./pretrain_<model>.py \
    $MODEL_ARGS \
    $OUTPUT_ARGS \
    --save $CHECKPOINT_PATH \
    --load $CHECKPOINT_PATH \
    --data-path $DATA_PATH \
    --tensor-model-parallel-size $TENSOR_MP_SIZE \
    --pipeline-model-parallel-size $PIPELINE_MP_SIZE \
    --DDP-impl torch

GPT-3示例

在 examples/pretrain_gpt3_175B.sh 中,我们提供了一个如何配置 Megatron 以在 1024 个 GPU 上运行拥有 1750 亿参数的 GPT-3 的示例。该脚本是为带有 pyxis 插件的 slurm 设计的,但可以轻松适用于任何其他调度程序。它分别使用了 8 路和 16 路的张量并行和管道并行。使用选项 global-batch-size 1536 和 rampup-batch-size 16 16 5859375,训练将从全局批量大小 16 开始,并在 5,859,375 个样本上以 16 的增量步骤线性增加全局批量大小至 1536。训练数据集可以是单个数据集,也可以是多个数据集与一组权重的组合。

在 1024 个 A100 GPU 上使用完整的 1536 全局批量大小时,每次迭代大约需要 32 秒,每个 GPU 的计算速度达到 138 teraFLOPs,这是理论峰值 FLOPs 的 44%。

评估和任务

我们提供了几个命令行参数,详见下面列出的脚本,用于处理各种零样本和微调下游任务。但是,您也可以根据需要从预训练检查点对其他语料库进行微调。为此,只需添加 --finetune 标志,并在原始训练脚本中调整输入文件和训练参数。迭代计数将重置为零,优化器和内部状态将重新初始化。如果微调因任何原因中断,请务必在继续之前删除 --finetune 标志,否则训练将从头开始。

由于评估所需的内存远少于训练,因此合并并行训练的模型以在单个 GPU 上用于下游任务可能更有优势。以下脚本实现了这一点。目前仅支持输入的张量模型并行和输出的管道模型并行。此示例读取具有 2 路张量模型并行的模型,并输出具有 2 路管道模型并行的模型。

TENSOR_MODEL_PARALLEL_SIZE=2
TARGET_PIPELINE_MODEL_PARALLEL_SIZE=2

VOCAB_FILE=bert-vocab.txt
CHECKPOINT_PATH=checkpoints/bert_345m

WORLD_SIZE=$TENSOR_MODEL_PARALLEL_SIZE python tools/merge_mp_partitions.py \
    --model-type BERT \
    --tensor-model-parallel-size $TENSOR_MODEL_PARALLEL_SIZE \
    --pipeline-model-parallel-size 1 \
    --target-pipeline-model-parallel-size $TARGET_PIPELINE_MODEL_PARALLEL_SIZE \
    --tokenizer-type BertWordPieceLowerCase \
    --vocab-file $VOCAB_FILE \
    --num-layers 24 \
    --hidden-size 1024 \
    --num-attention-heads 16 \
    --seq-length 512 \
    --max-position-embeddings 512 \
    --load $CHECKPOINT_PATH
    --save $CHECKPOINT_PATH/merged

下面描述了 GPT 和 BERT 模型的几个下游任务。它们可以使用与训练脚本中相同的更改在分布式和模型并行模式下运行。

GPT 文本生成

bash examples/generate_text.sh

我们使用 GPT 预训练脚本生成文本样本。需要做一些更改,例如我们需要提供预训练检查点的路径、输出样本的长度、是无条件生成文本（使用 --num-samples 表示要生成多少个样本）还是有条件生成（需要传递 --sample-input-file <filename>,文件中的每一行将用作条件文本）。还有一些可选参数可以调整,例如 top-k、top-p 或 greedy（将 top-k 和 top-p 设置为 0）采样。

CHECKPOINT_PATH=checkpoints/gpt2_345m
VOCAB_FILE=gpt2-vocab.json
MERGE_FILE=gpt2-merges.txt
GPT_ARGS=<与上面 GPT 预训练中相同>

MAX_OUTPUT_SEQUENCE_LENGTH=1024
TEMPERATURE=1.0
TOP_P=0.9
NUMBER_OF_SAMPLES=2
OUTPUT_FILE=samples.json

python tools/generate_samples_gpt.py \
    $GPT_ARGS \
    --load $CHECKPOINT_PATH \
    --out-seq-length $MAX_OUTPUT_SEQUENCE_LENGTH \
    --temperature $TEMPERATURE \
    --genfile $OUTPUT_FILE \
    --num-samples $NUMBER_OF_SAMPLES \
    --top_p $TOP_P \
    --recompute

GPT 评估

我们包含了 GPT 评估的示例脚本,用于 WikiText 困惑度评估和 LAMBADA 完形填空准确性。

WikiText 困惑度评估

为了与之前的工作进行公平比较,我们在词级 WikiText-103 测试数据集上评估困惑度,并考虑到使用我们的子词分词器时标记的变化,适当地计算困惑度。

我们使用以下命令在一个 3.45 亿参数的模型上运行 WikiText-103 评估。

TASK="WIKITEXT103"

VALID_DATA=<wikitext 路径>.txt
VOCAB_FILE=gpt2-vocab.json
MERGE_FILE=gpt2-merges.txt
CHECKPOINT_PATH=checkpoints/gpt2_345m

COMMON_TASK_ARGS=" \
    --num-layers 24 \
    --hidden-size 1024 \
    --num-attention-heads 16 \
    --seq-length 1024 \
    --max-position-embeddings 1024 \
    --fp16 \
    --vocab-file $VOCAB_FILE"

python tasks/main.py \
    --task $TASK \
    $COMMON_TASK_ARGS \
    --valid-data $VALID_DATA \
    --tokenizer-type GPT2BPETokenizer \
    --merge-file $MERGE_FILE \
    --load $CHECKPOINT_PATH \
    --micro-batch-size 8 \
    --checkpoint-activations \
    --log-interval 10 \
    --no-load-optim \
    --no-load-rng