DeepSpeed支持

TencentPretrain支持使用DeepSpeed进行超大模型预训练、微调和推理。

预训练

使用DeepSpeed进行预训练时需要指定 --deepspeed 选项，并使用 --deepspeed_config 指定DeepSpeed的配置文件。这里以Megatron-LM中的超大模型为例介绍如何在TencentPretrain中使用DeepSpeed。注意到Megatron BERT和Megatron GPT-2都使用了前置layernorm。

Megatron BERT：

使用DeepSpeed训练Megatron BERT模型示例： 单机8GPU预训练示例：

python3 preprocess.py --corpus_path corpora/CLUECorpusSmall_5000_lines.txt --vocab_path models/google_zh_vocab.txt \
                      --dataset_path dataset.pt --processes_num 8 --dynamic_masking \
                      --data_processor mlm

deepspeed pretrain.py --deepspeed --deepspeed_config models/deepspeed_config.json \
                      --dataset_path dataset.pt --vocab_path models/google_zh_vocab.txt \
                      --config_path models/megatron/bert_3.9B_config.json \
                      --output_model_path models/output_model \
                      --world_size 8 --batch_size 16 \ 
                      --total_steps 10000 --save_checkpoint_steps 5000 --report_steps 100 --deep_init

加载PyTorch格式预训练模型进行增量预训练示例：

deepspeed pretrain.py --deepspeed --deepspeed_config models/deepspeed_config.json \
                      --pretrained_model_path models/input_model.bin \
                      --dataset_path dataset.pt --vocab_path models/google_zh_vocab.txt \
                      --config_path models/megatron/bert_3.9B_config.json \
                      --output_model_path models/output_model \
                      --world_size 8 --batch_size 16 \ 
                      --total_steps 10000 --save_checkpoint_steps 5000 --report_steps 100 --deep_init

2机每机8GPU预训练示例： 使用多机训练时需要先配置hostfile.txt文件，文件格式为：ip slots=gpu数 例如：

1.1.1.1 slots=8
2.2.2.2 slots=8

在进行多机训练时，只需要在主节点启动训练脚本。

python3 preprocess.py --corpus_path corpora/CLUECorpusSmall_5000_lines.txt --vocab_path models/google_zh_vocab.txt \
                      --dataset_path dataset.pt --processes_num 8 --dynamic_masking \
                      --data_processor mlm

deepspeed --hostfile=hostfile.txt pretrain.py --deepspeed --deepspeed_config models/deepspeed_config.json \
                                              --dataset_path dataset.pt --vocab_path models/google_zh_vocab.txt \
                                              --config_path models/megatron/bert_3.9B_config.json \
                                              --output_model_path models/output_model \
                                              --world_size 16 --batch_size 16 \
                                              --total_steps 10000 --save_checkpoint_steps 5000 --report_steps 100 --deep_init

Megatron GPT-2：

使用DeepSpeed训练Megatron GPT-2模型示例： 单机8GPU预训练示例：

python3 preprocess.py --corpus_path corpora/CLUECorpusSmall_5000_lines.txt --vocab_path models/google_zh_vocab.txt \
                      --dataset_path dataset.pt --processes_num 8 --data_processor lm

deepspeed pretrain.py --deepspeed --deepspeed_config models/deepspeed_config.json \
                      --dataset_path dataset.pt --vocab_path models/google_zh_vocab.txt \
                      --config_path models/megatron/gpt2_8.3B_config.json \
                      --output_model_path models/output_model \
                      --world_size 8 --batch_size 4 \ 
                      --total_steps 10000 --save_checkpoint_steps 5000 --report_steps 100 --deep_init

2机每机8GPU预训练示例：

python3 preprocess.py --corpus_path corpora/CLUECorpusSmall_5000_lines.txt --vocab_path models/google_zh_vocab.txt \
                      --dataset_path dataset.pt --processes_num 8 --data_processor lm

deepspeed --hostfile=hostfile.txt pretrain.py --deepspeed --deepspeed_config models/deepspeed_config.json \
                                              --dataset_path dataset.pt --vocab_path models/google_zh_vocab.txt \
                                              --config_path models/megatron/gpt2_8.3B_config.json \
                                              --output_model_path models/output_model \
                                              --world_size 16 --batch_size 4 \ 
                                              --total_steps 10000 --save_checkpoint_steps 5000 --report_steps 100 --deep_init

DeepSpeed模型转换

预训练完成后，会根据设置的保存步数在 models/output_model 文件中生成相应的模型文件，以及模型转换脚本 zero_to_fp32.py ，通过该脚本可以将DeepSpeed格式的模型转换为PyTorch格式的模型。

usage: zero_to_fp32.py [-h] checkpoint_dir output_file

positional arguments:
  checkpoint_dir path to the deepspeed checkpoint folder, e.g., path/checkpoint-1/global_step1
  output_file path to the pytorch fp32 state_dict output file (e.g.,path/checkpoint-1/pytorch_model.bin)

optional arguments:
  -h, --help how this help message and exit

以转换Megatron BERT为例：

python3 models/output_model/zero_to_fp32.py models/output_model/10000 models/output_model/megatron_bert.bin-10000

下游任务微调

通过 finetune/run_classifier_deepspeed.py 脚本可以使用DeepSpeed对超大模型在分类任务上进行微调， run_classifier_deepspeed.py 脚本与普通分类脚本 run_classifier.py 有以下区别：

• 使用 run_classifier_deepspeed.py 脚本时需要通过 --world_size 指定使用的GPU数量。
• run_classifier_deepspeed.py 中的 --batch_size 为单卡的batch size。模型实际的batch size大小为 --batch_size 乘以 --world_size
• run_classifier_deepspeed.py 脚本会将每一轮微调后的模型都保存在 --output_model_path 指定的目录下。

以Megatron BERT为例 单机8GPU微调示例：

deepspeed finetune/run_classifier_deepspeed.py --pretrained_model_path models/output_model/megatron_bert.bin-10000 \
                                               --deepspeed_config models/deepspeed_config.json \
                                               --vocab_path models/google_zh_vocab.txt \
                                               --config_path models/megatron/bert_3.9B_config.json \
                                               --output_model_path models/classifier_model \
                                               --train_path  datasets/chnsenticorp/train.tsv \
                                               --dev_path datasets/chnsenticorp/dev.tsv \
                                               --test_path datasets/chnsenticorp/test.tsv \
                                               --epochs_num 3 --batch_size 8  --world_size 8

2机每机8GPU微调示例：

deepspeed --hostfile=hostfile.txt finetune/run_classifier_deepspeed.py --pretrained_model_path models/output_model/megatron_bert.bin-10000 \
                                                                       --deepspeed_config models/deepspeed_config.json \
                                                                       --vocab_path models/google_zh_vocab.txt \
                                                                       --config_path models/megatron/bert_3.9B_config.json \
                                                                       --output_model_path models/classifier_model \
                                                                       --train_path  datasets/chnsenticorp/train.tsv \
                                                                       --dev_path datasets/chnsenticorp/dev.tsv \
                                                                       --test_path datasets/chnsenticorp/test.tsv \
                                                                       --epochs_num 3 --batch_size 4  --world_size 16

将保存的DeepSpeed格式模型转换为PyTorch格式模型：

python3 models/classifier_model/zero_to_fp32.py models/classifier_model/3 models/classifier_model/megatron_bert_classifier.bin

推理

通过 run_classifier_deepspeed_infer.py 脚本可以使用DeepSpeed对超大模型在分类任务上进行推理。 以Megatron BERT为例 单机推理示例：

deepspeed inference/run_classifier_deepspeed_infer.py --load_model_path models/classifier_model/megatron_bert_classifier.bin \
                                                     --vocab_path models/google_zh_vocab.txt \
                                                     --config_path models/megatron/bert_3.9B_config.json \
                                                     --test_path datasets/chnsenticorp/test_nolabel.tsv \
                                                     --prediction_path prediction.txt --labels_num 2

2机推理示例：

deepspeed --hostfile=hostfile.txt inference/run_classifier_deepspeed_infer.py --load_model_path models/classifier_model/megatron_bert_classifier.bin \
                                                                             --vocab_path models/google_zh_vocab.txt \
                                                                             --config_path models/megatron/bert_3.9B_config.json \
                                                                             --test_path datasets/chnsenticorp/test_nolabel.tsv \
                                                                             --prediction_path prediction.txt  --labels_num 2

文本生成

通过 generate_lm_deepspeed.py 脚本可以使用DeepSpeed在超大语言模型上进行文本生成，给定文本开头，模型根据开头续写。 以Megatron GPT-2为例： 单机推理示例：

deepspeed scripts/generate_lm_deepspeed.py --load_model_path models/megatron_gpt2.bin \
                                           --vocab_path models/google_zh_vocab.txt \
                                           --config_path models/megatron/gpt2_8.3B_config.json \
                                           --test_path beginning.txt  --prediction_path generated_sentence.txt

2机推理示例：

deepspeed --hostfile=hostfile.txt scripts/generate_lm_deepspeed.py --load_model_path models/megatron_gpt2.bin \
                                                                   --vocab_path models/google_zh_vocab.txt \
                                                                   --config_path models/megatron/gpt2_8.3B_config.json \
                                                                   --test_path beginning.txt  --prediction_path generated_sentence.txt

通过 generate_seq2seq_deepspeed.py 脚本可以使用DeepSpeed在超大Seq2seq模型上进行文本生成。 单机推理示例：

deepspeed scripts/generate_seq2seq_deepspeed.py --load_model_path models/input_model.bin \
                                                --vocab_path models/google_zh_vocab.txt \
                                                --config_path models/encoder_decoder_config.json \
                                                --test_path input.txt  --prediction_path output.txt

2机推理示例：

deepspeed --hostfile=hostfile.txt scripts/generate_seq2seq_deepspeed.py --load_model_path models/input_model.bin \
                                                                        --vocab_path models/google_zh_vocab.txt \
                                                                        --config_path models/encoder_decoder_config.json \
                                                                        --test_path input.txt  --prediction_path output.txt