[megatron程式碼閱讀] 1. 初始化和組網

SunStriKE發表於2025-01-13
原文網址 : https://www.cnblogs.com/sunstrikes/p/18668809

以pretrain_gpt.py為例, 看megatron的整體邏輯. 本章主要包括megatron初始化相關邏輯, 核心函式為initialize_megatron, setup_model_and_optimizer兩個

initialize_megatron

parse_args

從argparse中直接讀取超引數配置. 如學習率, 正則化等. 從環境變數中獲取rank等

load_args_from_checkpoint

  • 優先從未被持久化的ckpt載入, 並且只載入rank0的args

  • _load_non_persistent_base_checkpoint

    • find_checkpoint_rank_0

      在不知道是否使用pp/ep策略的情況下, 嘗試拼裝出rank0 ckpt的名稱, 如果存在就能定位到實際的存放目錄

    • verify_checkpoint_and_load_strategy

      根據是zarr還是 torch_dist選擇不同的載入策略

    • TorchCommonLoadStrategy->torch.load()

  • 如果沒有非持久化的, 載入遠端ckpt

  • 從ckpt裡的args替換掉之前解析的部分args, 比如tp/pp/vp等超引數

校驗yaml/args, 全域性變數設定

_initialize_distributed

pytorch裡的get_world_size 返回的是gpu總卡數

初始化torch.distributed

mpu.initialize_model_parallel (並行設定,核心函式)

RankGenerator:

  1. 在每塊GPU上啟動一個程序（process），每個程序獨立執行自己所維護的那部分模型的計算，實現並行訓練
  2. 儲存tp/pp/dp/ep/cp 各種並行度配置大小. 並且能夠從 tp-dp str格式的並行配置裡獲取 tp/dp對應的mask和並行度大小設定.
  3. get_ranks: 根據parallel_size和mask, 計算各種並行策略拆分後的rank group.

[!NOTE]

舉例: 假定有2個8卡機器，node1: rank 0-7，node2: rank 8-15 tp-pp-dp: [2,4,2]

  • _TENSOR_MODEL_PARALLEL_GROUP ：[g0, g1], [g2, g3], [g4, g5], [g6, g7], [g8, g9], [g10, g11], [g12, g13], [g14, g15]。
  • _PIPELINE_MODEL_PARALLEL_GROUP ： [g0, g4, g8, g12], [g1, g5, g9, g13], [g2, g6, g10, g14], [g3, g7, g11, g15]。
  • _MODEL_PARALLEL_GROUP ：tp-pp = 2 * 4 = 8 [0, 1, 4, 5, 8, 9, 12, 13]，[2, 3, 6, 7, 10, 11, 14, 15]
  • _DATA_PARALLEL_GROUP ：[g0, g2], [g1, g3], [g4, g6], [g5, g7], [g8, g10], [g9, g11], [g12, g14], [g13, g15]。
分隔樣例

注意在PP內輸入層和輸出層共享一個word_embedding，PP組中的第一個和最後一個rank需要通訊，保證word_embedding完全一致

group全域性變數賦值: 每個並行模式有一個分組全域性變數.透過 generator_wrapper生成, 自己的程序rank如果在group內, 初始化對應的nccl/gloo torch.distributed.new_group

GlobalMemoryBuffer: 儲存每個已經分配出的tensor, 避免視訊記憶體重分配.

setup_model_and_optimizer

主要邏輯是配置模型組網和最佳化器.

model_provider: torch gpt組網

megatron/core/transformer, transformer組網核心邏輯, 基於torch.nn.Module, 將涉及到的子模型結構進行了抽象. 透過subModule的方式嵌入自定義module, 便於程式碼複用

例如

self_attention=ModuleSpec(
    module=SelfAttention,
    params={"attn_mask_type": attn_mask_type},
    submodules=SelfAttentionSubmodules(
        linear_qkv=ColumnParallelLinear,
        core_attention=DotProductAttention,
        linear_proj=RowParallelLinear,
        q_layernorm=IdentityOp,
        k_layernorm=IdentityOp,
    ),
)

在attention.py裡讀到之前moduleSpec中的對應linear_qkv的實現, 即TP列並行的Linear實現. 加上TransformerConfig, 就能定義出最終的網路邏輯. TP相關邏輯在後續專門看的時候再細寫.

self.linear_qkv = build_module(
    submodules.linear_qkv,
    self.config.hidden_size,
    self.query_projection_size + 2 * self.kv_projection_size,
    config=self.config,
    init_method=self.config.init_method,
    gather_output=False,
    bias=self.config.add_bias_linear or self.config.add_qkv_bias,
    skip_bias_add=False,
    is_expert=False,
    tp_comm_buffer_name='qkv',
)

torch裡實現module時, 主要關注__init__()和forward(), bp透過自動微分生成.

配置

配置類 ModelParallelConfig, TransformerConfig

ModelParallelConfig: 主要包括 模型並行/PP/通訊overlap相關最佳化開關/cpuOffload 等相關配置

TransformerConfig: 主要包括 模型結構/MOE/運算元fusion加速/啟用重計算/Context並行 等配置

models/gpt/gpt_model.py

preprocess

分為word_emb和pos_emb兩部分. 輸出為 word_emb(b,s,h) + pos_emb(s,h) + tokentype_emb(b,s,h)(需要轉置適配)

注意在embedding最後要進行dropout處理, 應該是為了減少模型過擬合的風險

WordEmbeddings

tensor_parallel.VocabParallelEmbedding

vocab_size表示詞表維度, 例如分詞預處理後保留能查到的幾千個常用單詞. 將vocab_size個embed均分儲存到global_world_size張卡上, embedding lookup時從對應的儲存卡上拉取. 這裡把非自身rank的emb透過[start_idx, end_idx)的mask操作置0, 然後透過reduce就能獲取完整的詞表.

如果配置開了序列並行, reduce操作會變為reduceScatter操作, lookup之後直接分配好sp的輸入.

RoPE(旋轉位置編碼)

位置編碼需要滿足幾個性質: 1. 不能滿足交換律, 第m個token與第n個token的位置關係，和第n個token與第m個token的位置關係一定要有區分度。 2.需要有遠端衰減性

image-20250108114351771

為了便於加速計算, 可以等價最佳化為下面這種向量乘法的形式:

image-20250108114806801 image-20250108114336830
tokentype_embedding

型別嵌入層，用於區分輸入中不同型別的token, 例如，在BERT中用於區分兩個句子，而在某些GPT變種或特定任務中可能用於區分不同型別的輸入資料，如對話中的提問和回答.

transformer

self.decoder就是上面透過ModuleSpec獲得的module, 可以根據配置選擇普通的selfAttention, 還是MLA.

  1. MLA原理: 在模型能力不變基礎上，透過KV低秩壓縮, 使得推理的KVcache視訊記憶體佔用和計算效率上對比MHA效能有明顯提升.
image-20250107171551928
postprocess
1.output_layer & loss

訓練時output可以並行, 這裡是個TP列並行的方式, 訓練方式如下例子:

<s>
<s> i
<s> i love 
<s> i love maching
<s> i love maching learning <eos/>

訓練階段將這個矩陣直接輸入到decoder，分別得到 5個輸出 \(O_i, i\in [1,2,3,4,5]\)， 理想的輸出應該是[i, love, maching, learning, ] ,然後 比較\(O_i\)和理想輸出的交叉熵，得到loss. 而且這五個序列可以放在一個batch內平行計算.

optimizer

_get_param_groups_and_buffers

從多個model_chunks中遍歷所有的param向量, 對其中某些param進行特殊的處理

  • decoupled_lr是為input/output layer單獨設定的lr
  • no_weight_decay_cond: 配置引數是否應該執行權重衰減。
  • scale_lr_cond: 對某些指定層的引數進行學習率縮放, 匹配到對應的param_map後執行.
_get_megatron_optimizer_based_on_param_groups

主要邏輯是混合精度optimizer的設定(MixedPrecisionOptimizer), TODO: 細看Apex.FusedAdam, 和torch.adamW的區別在哪裡

梯度縮放: DynamicGradScaler

混合精度訓練的時候, 用於動態調整梯度縮放比例，以處理梯度爆炸或消失問題.

主要邏輯是有一個初始化scale值, 當連續hysteresis次迭代中出現NaN，torch.max(scale * backoff_factor, min_scale) 用來減小scale\(backoff\_factor \in (0, 1)\).

當連續growth_interval次沒出現NaN, 按照_scale * growth_factor_, 放大scale, \(growth\_factor > 1\)

DistributedOptimizer

介面繼承自torch.optimizer, 核心邏輯在step(self), 有3個類: FP32Optimizer, ChainedOptimizer, MixedPrecisionOptimizer

FP32Optimizer: fp32訓練使用到的, 主要功能是配置了clip_grad後進行normalization, norm分兩種, 一種是取max_grad, 一種是l2範數, 透過all_reduce拿到total_norm, 最後用這個值分別對每個param tensor進行scale. 在scale之後就呼叫的是torch.optimizer.step進行正常的Adam更新.

MixedPrecisionOptimizer: 混合精度訓練使用

  • prepare_grads: 先從param.grad copy到 param.main_grad, 這一步同時做了fp16->fp32的轉換, 然後檢查所有的grad, 先unscale, 再看是否存在NaN. 注意只有fp16需要, bf16不需要.
  • clip_grad_norm: 與FP32Optimizer一樣的方法scale grad.
  • step_with_ready_grads: optimizer.step後, 再把fp32的main_param copy回用於下一輪bp的fp16 param裡面.

ChainedOptimizer: 用於moe場景, 每個分塊子模型配置不同的optimizer時使用. 多個optimizer之間序列執行.

下一節看megatron的模型儲存&載入, 並行訓練相關程式碼.

參考連結

ROPE位置編碼部落格, 論文

MLA原理部落格

相關文章