阿里雲PAI-靈駿大模型訓練工具 Pai-Megatron-Patch 正式開源！

Pai-Megatron-Patch是什麼

Pai-Megatron-Patch工具是阿里雲機器學習平臺PAI算法團隊研發，基於阿里雲智算服務PAI-靈駿平臺的大模型最佳實踐解決方案配套工具，旨在幫助大模型開發者快速上手靈駿產品，完成大語言模型（LLM）的高效分佈式訓練，有監督指令微調，模型離線推理驗證等完整大模型開發鏈路。該項目提供了業界主流開源大模型基於Megatron-LM的訓練&離線推理驗證流程，方便用戶快速上手大模型訓練。

主要特性

多款熱門大模型支持：llama，llama-2，codellama, 百川，通義千問，Falcon，GLM，Starcoder，Bloom，chatglm等
支持模型權重互轉轉換：在Huggingface，Megatron和Transformer Engine之間進行算子命名空間映射
支持Flash Attention 2.0和Transformer Engine模式下的FP8訓練加速且確保收斂
豐富且簡單易用的使用示例，支持大模型預訓練，微調，評估和推理，強化學習全流程最佳實踐

開源地址

https://github.com/alibaba/Pai-Megatron-Patch

技術架構

Pai-Megatron-Patch的設計理念是不對Megatron-LM的源碼進行侵入式修改，即不在Megatron-LM裏面添加新的功能特性，將需要擴充完善的部分以patch補丁的方式呈現。在patch中構建LLM訓練鏈路通過依賴Megatron-LM核心庫的方法實現和Megatron-LM的解耦合。這樣解耦合的好處就是Megatron-LM的升級不會影響用戶的LLM最佳實踐體驗。

Pai-Megatron-Patch中包含模型庫，分詞器，模型轉換，強化學習，離線文本生成以及使用示例和工具集等用於構建LLM訓練的關鍵要素。在模型庫中包含熱門大模型的Megatron版本實現，例如baichuan，bloom，chatglm，falcon，galactica，glm，llama，qwen和starcoder，後續還會根據需要及時添加新的Megatron版大模型實現。同時patch還提供了huggingface模型權重和Megatron模型權重之間的雙向轉換。一方面是方便用戶加載huggingface的權重在Megatron中繼續預訓練或者微調，另一方面是方便用戶對訓練好的Megatron模型使用huggingface的評估/推理流程對模型質量進行客觀評估。在強化學習部分，patch提供了PPO訓練流程等，方便用戶使用SFT模型和RM模型進行強化學習。最後patch提供了大量的使用示例幫助用戶快速開始大模型訓練&離線推理。具體請參考阿里雲靈駿產品的使用流程: 智算服務PAI靈駿大模型分佈式訓練方案

關鍵技術

模型權重轉換

研發Megatron-Patch的初衷之一就是能將世界各地研發機構在Huggingface上放出的熱門大模型使用Megatron引擎進行繼續預訓練或者繼續微調。這就需要首先將Huggingface模型格式的ckpt轉換成Megatron模型格式，才能正確加載進來，否則會出現pytorch加載模型失敗。Megatron-Patch的一個核心可靠性保障特徵就是在採用算子拆分，流水並行，序列並行，Zero顯存優化，BF16混合精度，梯度檢查點等訓練加速技術確保模型訓練吞吐速度平均提升1.5倍以上的同時，在評估任務模式下的單一樣本前向loss值，預訓練/微調任務模式下的loss曲線，離線文本生成任務模式下的生成效果這三個方面和Huggingface是對齊的，從而確保Megatron版模型的可靠性。

另一方面，Megatron版的transformer實現方式提供了一種讓用戶僅僅通過設置開關就能實現不同種類GPT模式的能力。比如llama模型打開如下開關即可

--swiglu \
  --use-rotary-position-embeddings \
  --no-position-embedding \
  --untie-embeddings-and-output-weights \
  --disable-bias-linear

如果想將llama模式變成baichuan模型，那麼僅僅需要添加採用--use-alibi-mask開關，同時關閉Rotary Embeeding開關即可，具體配置如下所示：

--swiglu \
  --use-alibi-mask \
  --position-embedding-type none \
  --untie-embeddings-and-output-weights \
  --disable-bias-linear

下面我們以llama-2爲例，詳解從huggingface到megatron的模型權重轉換技術。下表總結了兩者在不同module上的命名對應關係。在patch實現過程中，我們首先將HF格式的ckpt轉換到一種內部格式，然後再把這種內部格式轉換成對應的外部格式。這樣做可以最大程度複用已有的轉換邏輯來處理新模型。在轉換爲內部格式的過程中，

q_proj, k_proj, v_proj需要沿着第0維拼接在一起後賦值給內部變量query_key_value。

當用戶在資源受限情況下需要按照TP>1來拆分權重的時候，這裏需要注意的是針對MLP層的gate_proj和up_proj的操作。不能像qkv那樣在轉換成內部格式的時候進行merge再執行算子拆分。需要在拆分前加入如下針對MLP層的權重合並的代碼邏輯才能確保正確收斂。

for i in range(tp_size):
    params_dict = get_element_from_dict_by_path(output_state_dict[i],
                                                "model.language_model.encoder")
    dense_h_to_4h_1_name = 'mlp.dense_h_to_4h_1.weight'
    dense_h_to_4h_1_layer_name = f"layers.{layer}.{dense_h_to_4h_1_name}"
    dense_h_to_4h_1_weight = params_dict[dense_h_to_4h_1_layer_name]
    dense_h_to_4h_2_name = 'mlp.dense_h_to_4h_2.weight'
    dense_h_to_4h_2_layer_name = f"layers.{layer}.{dense_h_to_4h_2_name}"
    dense_h_to_4h_2_weight = params_dict[dense_h_to_4h_2_layer_name]
    dense_h_to_4h_name = 'mlp.dense_h_to_4h.weight'
    dense_h_to_4h_layer_name = f"layers.{layer}.{dense_h_to_4h_name}"
    params_dict[dense_h_to_4h_layer_name] = torch.cat(
    [dense_h_to_4h_1_weight, dense_h_to_4h_2_weight], dim=0)

基於TE的FP8訓練收斂

Transformer Engine(TE)是一個在英偉達GPUS上運行的針對Transformer模型的加速庫，其中包括針對Hopper GPU的FP8混合精度，該精度可以在較低的顯存利用率下提供更好的訓練&推理速度。在TE內部封裝了Flash Attention實現，同時TE還提供了一組高度優化後的算子用來構建Transformer模型。比如LayerNormLinear就是將LayerNorm和QKV-Proojection進行算子融合，LayerNormMLP就是將layernorm和mlp進行算子融合。如下圖所示：

從Huggingface到TE模型的權重轉換技術和之前是類似的，也需要事先找到兩者之間的映射關係。從下表可以看出，TE中多了_extra_state是用來存fp8訓練的scale和history的，這些在加載的時候會出現衝突，這時只要將load_state_dict函數的strict設置成False就可以了，比如load_state_dict(state_dict_, strict=False)。

在Megatron-Patch中使用示例中打開FP8混合精度訓練開關也很容易，如下所示：

if [ $PR = fp16 ]; then
    pr_options=" \
        --fp16"
elif [ $PR = bf16 ]; then
    pr_options=" \
        --bf16"
elif [ $PR = fp8 ]; then
    pr_options=" \
        --bf16
        --fp8-hybrid \
        --fp8-amax-compute-algo max \
        --fp8-amax-history-len 1024 \
        --transformer-impl transformer_engine"
fi

我們可以使用如下訓練腳本run_pretrain_megatron_llama_enwiki.sh來測試打開FP8開關後的預訓練收斂性。下圖展示了llama-7B和llama-2-70B模型在打開和關閉FP8時的loss曲線對比，可以看出基本是重合的。

LLama-7B

LLama2-70B

大模型訓練&推理

從github上獲取Megatron模型訓練工具PAI-Megatron-Patch（https://github.com/alibaba/Pai-Megatron-Patch）源代碼並拷貝到工作目錄/mnt/workspace/下。

模型格式轉換

使用我們提供的模型轉換腳本，將huggingface格式的模型文件轉換爲megatron格式：

cd /mnt/workspace/
mkdir llama2-ckpts
cd llama2-ckpts
wget https://atp-modelzoo-wlcb-pai.oss-cn-wulanchabu.aliyuncs.com/release/models/pai-megatron-patch/llama2-ckpts/Llama-2-7b-hf.tgz
tar -zxf Llama-2-7b-hf.tgz
mv Llama-2-7b-hf llama2-7b-hf
cd /mnt/workspace/PAI-Megatron-Patch/toolkits/model_checkpoints_convertor/llama
sh model_convertor.sh \
/root/Megatron-LM-23.04        \
/mnt/workspace/llama2-ckpts/llama2-7b-hf         \
/mnt/workspace/llama2-ckpts/llama2-7b-hf-to-megatron-tp1-pp1  \
1  \
1  \
llama-7b \
0 \
false

繼續預訓練

中文繼續預訓練漢化指引

Step1: 獲取需要擴充詞表的模型（如llama-13b-hf）

Step2: 獲取需要擴充的詞表

使用sentence-piece代碼庫從自有文本語料中學習詞表，得到randeng-sp.model文件

Step3: 詞表擴充

擴充模型tokenizer：將randeng-sp.model中的詞表添加到llama-13b-hf文件夾下tokenizer.model中
擴充模型詞表對應的參數矩陣
- word_embedding、lm_head
- 新詞向量可以使用原詞向量均值作爲初始化，比如“天氣”=mean([“天”，“氣”])
修改與詞表大小相關的文件並保存，如config.json

運行繼續預訓練腳本 run_pretrain_megatron_llama.sh，需要傳入的參數列表如下

ENV=$1                          # 運行環境: dlc, dsw
MEGATRON_PATH=$2                # 設置開源Megatron的代碼路徑
MEGATRON_PATCH_PATH=$3          # 設置Megatron Patch的代碼路徑
MODEL_SIZE=$4                   # 模型結構參數量級：7B, 13B
BATCH_SIZE=$5                   # 每卡訓練一次迭代樣本數: 4, 8
GLOBAL_BATCH_SIZE=$6            # 全局batch size
LR=$7                           # 學習率: 1e-5, 5e-5
MIN_LR=$8                       # 最小學習率: 1e-6, 5e-6
SEQ_LEN=$9                      # 序列長度
PAD_LEN=${10}                   # Padding長度：100
EXTRA_VOCAB_SIZE=${11}          # 詞表擴充大小
PR=${12}                        # 訓練精度: fp16, bf16
TP=${13}                        # 模型並行度
PP=${14}                        # 流水並行度
AC=${15}                        # 激活檢查點模式: sel, full
DO=${16}                        # 是否使用Megatron版Zero-1降顯存優化器: true, false
FL=${17}                        # 是否使用Flash Attention: true, false
SP=${18}                        # 是否使用序列並行: true, false
SAVE_INTERVAL=${19}             # 保存ckpt的間隔
DATASET_PATH=${20}              # 訓練數據集路徑
PRETRAIN_CHECKPOINT_PATH=${21}  # 預訓練模型路徑
TRAIN_TOKENS=${22}              # 訓練token數
WARMUP_TOKENS=${23}             # 預熱token數
OUTPUT_BASEPATH=${24}           # 訓練輸出文件路徑

注意設置正確的數據集掛載路徑WORK_DIR以及運行環境ENV，運行示例如下所示：

export WORK_DIR=/mnt/workspace
cd ${WORK_DIR}/PAI-Megatron-Patch/examples/llama2
bash run_pretrain_megatron_llama.sh \
dlc \
/root/Megatron-LM-23.04   \
${WORK_DIR}/PAI-Megatron-Patch  \
7B   \
1    \
16 \
1e-5   \
1e-6   \
2048  \
80  \
0   \
fp16  \
1   \
1  \
sel  \
true   \
false  \
false   \
100000  \
${WORK_DIR}/llama2-datasets/wudao/wudao_llamabpe_text_document   \
${WORK_DIR}/llama2-ckpts/llama2-7b-hf-to-megatron-tp1-pp1   \
100000000   \
10000   \
${WORK_DIR}/output_megatron_llama2/

有監督微調

在微調開始之前，請先進入https://github.com/alibaba/Pai-Megatron-Patch/blob/main/toolkits/pretrain_data_preprocessing/README.md 獲取json文件。運行run_finetune_megatron_llama.sh腳本，需要傳入的參數列表如下

ENV=$1                          # 運行環境: dlc, dsw
MEGATRON_PATH=$2                # 設置開源Megatron的代碼路徑
MEGATRON_PATCH_PATH=$3          # 設置Megatron Patch的代碼路徑
MODEL_SIZE=$4                   # 模型結構參數量級: 7B, 13B
BATCH_SIZE=$5                   # 每卡訓練一次迭代樣本數: 4, 8
LR=$6                           # 學習率: 1e-5, 5e-5
MIN_LR=$7                       # 最小學習率: 1e-6, 5e-6
SEQ_LEN=$8                      # 序列長度
PAD_LEN=$9                      # Padding長度：100
EXTRA_VOCAB_SIZE=${10}          # 詞表擴充大小
PR=${11}                        # 訓練精度: fp16, bf16
TP=${12}                        # 模型並行度
PP=${13}                        # 流水並行度
AC=${14}                        # 激活檢查點模式: sel, full
DO=${15}                        # 是否使用Megatron版Zero-1降顯存優化器: true, false
FL=${16}                        # 是否使用Flash Attention: true, false
SP=${17}                        # 是否使用序列並行: true, false
TRAIN_DATASET_PATH=${18}        # 訓練數據集路徑
VALID_DATASET_PATH=${19}        # 驗證數據集路徑
PRETRAIN_CHECKPOINT_PATH=${20}  # 預訓練模型路徑
EPOCH=${21}                     # 訓練迭代輪次
OUTPUT_BASEPATH=${22}           # 訓練輸出文件路徑

多節點運行示例如下所示：

export WORK_DIR=/mnt/workspace
cd ${WORK_DIR}/PAI-Megatron-Patch/examples/llama2
sh run_finetune_megatron_llama.sh  \
dlc    \
/root/Megatron-LM-23.04   \
${WORK_DIR}/PAI-Megatron-Patch  \
7B     \
1      \
1e-5   \
1e-6   \
2048   \
80     \
0      \
fp16   \
1      \
1      \
sel    \
true   \
false  \
false  \
${WORK_DIR}/llama2-datasets/wudao_train.json   \
${WORK_DIR}/llama2-datasets/wudao_valid.json   \
${WORK_DIR}/llama2-ckpts/llama2-7b-hf-to-megatron-tp1-pp1   \
2   \
${WORK_DIR}/output_megatron_llama2/

離線推理

模型訓練完成後，可以進行離線推理，評估模型效果。根據上面的訓練流程不同，我們提供了Megatron格式的推理鏈路。對於Megatron訓練的模型，可以直接用Megatron框架進行推理。

ENV=$1                          # 運行環境: dlc, dsw
MEGATRON_PATH=$2                # 設置開源Megatron的代碼路徑
MEGATRON_PATCH_PATH=$3          # 設置Megatron Patch的代碼路徑
CHECKPOINT_PATH=$4              # 模型微調階段的模型保存路徑
MODEL_SIZE=$5                   # 模型結構參數量級: 1.1B, 1.7B, 7.1B
TP=$6                           # 模型並行度
BS=$7                           # 每卡推理一次迭代樣本數: 1, 4, 8
SEQ_LEN=$8                      # 序列長度: 256, 512, 1024
PAD_LEN=$9                      # PAD長度：需要將文本拼接到的長度
EXTRA_VOCAB_SIZE=${10}          # 模型轉換時增加的token數量
PR=${11}                        # 推理採用的精度: fp16, bf16
TOP_K=${12}                     # 採樣策略中選擇排在前面的候選詞數量(0-n): 0, 5, 10, 20
INPUT_SEQ_LEN=${13}             # 輸入序列長度: 512
OUTPUT_SEQ_LEN=${14}            # 輸出序列長度: 256
INPUT_FILE=${15}                # 需要推理的文本文件: input.txt, 每行爲一個樣本
OUTPUT_FILE=${16}               # 推理輸出的文件: output.txt
# TOP_K和TOP_P必須有一個爲0
TOP_P=${17}                     # 採樣策略中選擇排在前面的候選詞百分比(0-1): 0, 0.85, 0.95
TEMPERATURE=${18}               # 採樣策略中溫度懲罰: 1-n
REPETITION_PENALTY=${19}        # 避免生成是產生大量重複，可以設置爲(1-2)默認爲1.2

此處提供一個離線推理輸出的文件，推理的數據組織形式需要與微調時的保持一致。
- 測試樣本：https://atp-modelzoo-wlcb-pai.oss-cn-wulanchabu.aliyuncs.com/release/models/pai-megatron-patch/llama2-datasets/pred_input.jsonl
注意：
- 模型保存的路徑下缺少tokenizer依賴的文件，需要將微調前模型路徑下所有json文件及tokenizer.model拷貝至保存模型的路徑下（位於{OUTPUT_BASEPATH }/checkpoint），與latest_checkpointed_iteration.txt同級。

以下有監督微調過程保存模型的推理代碼，需要將run_text_generation_megatron_llama.sh腳本中CUDA_VISIBLE_DEVICES參數設置爲0；GPUS_PER_NODE參數設置爲1；同時使用下列代碼進行推理。此時使用單卡進行推理。注意：此處模型tp爲1，可使用單卡推理；如果tp>1，則需使用相應卡數進行推理。

export WORK_DIR=/mnt/workspace
cd ${WORK_DIR}/PAI-Megatron-Patch/examples/llama2
bash run_text_generation_megatron_llama.sh \
dsw \
/root/Megatron-LM-23.04 \
${WORK_DIR}/PAI-Megatron-Patch \
../../../llama2-train \
7B \
1 \
1 \
1024 \
1024 \
0 \
fp16 \
10 \
512 \
512 \
${WORK_DIR}/pred_input.jsonl \
${WORK_DIR}/llama2_pred.txt \
0 \
1.0 \
1.2

大模型強化學習

一般來說，SFT微調過的模型在對話場景已經會有不錯的表現了。如果想進一步提升模型效果，可以再加上RLHF訓練。包括獎勵模型（Reward Model）的訓練和強化學習（PPO）的訓練。這裏展示瞭如何使用當前最常用的RLHF開源代碼框架，DeepSpeed-Chat和trlx，來進行獎勵函數訓練（RM），以及強化學習優化（PPO）。

模型格式轉換

如果基於huggingface格式的模型直接進行獎勵模型訓練（RM）和強化學習優化（PPO），可以跳過此步驟。

如果基於Megatron格式的模型，如PAI-Megatron-Patch訓練好的SFT模型，進行RM和PPO訓練，需要使用我們提供的模型轉換腳本，先將Megatron格式的模型文件轉換爲huggingface格式。

LLaMA2模型轉換：

cd PAI-Megatron-Patch/toolkits/model_checkpoints_convertor/gpt3_llama
bash model_convertor.sh \
/path/to/Megatron-LM \
/path/to/megatron_llama2_ckpt \
/path/to/hf_llama2_ckpt \
1 \
1 \
llama-7b \
0 \
true

BLOOM模型轉換：

cd PAI-Megatron-Patch/toolkits/model_checkpoints_convertor/bloom
bash model_convertor_huggingface_megatron.sh \
/path/to/Megatron-LM \
/path/to/megatron_bloom_ckpt \
/path/to/hf_bloom_ckpt \
1 \
1 \
true

DeepSpeed-Chat

下載安裝開源社區DeepSpeed-Chat源代碼：

cd PAI-Megatron-Patch/rlhf/deepspeed-chat
git clone https://github.com/microsoft/DeepSpeedExamples.git
cp -f rm_main.py DeepSpeedExamples/applications/DeepSpeed-Chat/training/step2_reward_model_finetuning/main.py
cp -f utils.py DeepSpeedExamples/applications/DeepSpeed-Chat/training/utils/utils.py
cd DeepSpeedExamples/applications/DeepSpeed-Chat/
pip install -r requirements.txt

基於LLaMA2模型訓練獎勵模型（RM）：

cd training/step2_reward_model_finetuning/ && bash training_scripts/llama2/run_llama2_7b.sh

基於LLaMA2進行強化學習優化訓練（PPO）：

cd training/step3_rlhf_finetuning/ && bash training_scripts/llama2/run_llama2_7b_lora.sh

trlx

下載安裝開源社區trlx源代碼：

cd PAI-Megatron-Patch/rlhf/trlx
git clone https://github.com/CarperAI/trlx.git
cp trlx_bloom_rlhf.py trlx_bloom_rlhf_test.py trlx/examples/summarize_rlhf/
cp train_reward_model_bloom.py reward_model_bloom.py ds_config_bloom.json trlx/examples/summarize_rlhf/reward_model/
cp -f ds_config_trlx_gptj_summarize.json trlx/examples/summarize_rlhf/configs/
cd trlx
pip install -e .

基於BLOOM模型訓練獎勵模型（RM）：

cd examples/summarize_rlhf/reward_model/ && deepspeed train_reward_model_bloom.py

基於GPT-J模型訓練獎勵模型（RM）：

cd examples/summarize_rlhf/reward_model/ && deepspeed train_reward_model_gptj.py

基於BLOOM模型進行強化學習優化訓練（PPO）：

cd examples/summarize_rlhf/ && accelerate launch --config_file configs/default_accelerate_config.yaml trlx_bloom_rlhf.py

基於GPT-J模型進行強化學習優化訓練（PPO）：

cd examples/summarize_rlhf/ && accelerate launch --config_file configs/default_accelerate_config.yaml trlx_gptj_text_summarization.py

PPO單測

如果您想跳過有監督微調（SFT）與獎勵模型訓練（RM）兩個步驟，只單獨測試PPO模塊的性能，可以運行如下指令單測PPO：

cd examples/summarize_rlhf/ && accelerate launch --config_file configs/default_accelerate_config.yaml trlx_bloom_rlhf_test.py

開源生態——構想和未來

在PAI-Megatron-Patch的開發過程中，我們圍繞中文大模型訓練加速落地沉澱了以下幾個方面的內容：

Huggingface的模型權重無損轉換成Megatron或者Transformer Engine可讀的模型權重。
H800集羣開啓FP8混合精度訓練確保收斂。
LLM大模型在PAI靈駿智算平臺上的最佳實踐。
強化學習技術在PAI靈駿智算平臺上的最佳實踐。

後續在PAI-Megatron-Patch中還會陸續放出更多高質量的大模型和最佳實踐。

參考文獻

[1]. Attention Is All You Need

[2]. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

[3]. Reducing Activation Recomputation in Large Transformer Models

[4]. FP8 Formats for Deep Learning

[5]. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

[6]. LLaMA: Open and Efficient Foundation Language Models

[7]. Llama 2: Open Foundation and Fine-Tuned Chat Models

[8]. Benchmarking Large Language Models on NVIDIA H100 GPUs with CoreWeave

點擊立即免費試用雲產品開啓雲上實踐之旅！

原文鏈接

本文爲阿里雲原創內容，未經允許不得轉載。

阿里雲PAI-靈駿大模型訓練工具 Pai-Megatron-Patch 正式開源！

Pai-Megatron-Patch是什麼

主要特性

開源地址

技術架構

關鍵技術

模型權重轉換

基於TE的FP8訓練收斂

大模型訓練&推理

大模型強化學習

開源生態——構想和未來

參考文獻

【面試準備】又一次失敗的面試經歷，題目離譜～資深軟件測試工程師

哈哈哈哈或

【自用】關於微信小程序的合法域名列表，

使用 @NoRepositoryBean 簡化數據庫訪問

MySQL查出時間比實際晚8小時的解決方案

什麼是IPD項目管理模式？聊聊IPD下的產品研發流程

https://yachay.unat.edu.pe/blog/index.php?comment_area=format_blog&comment_component=blog&comment_co

linux以太網驅動總結