[megatron代码阅读] 1. 初始化和组网

以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:

在每块GPU上启动一个进程（process），每个进程独立执行自己所维护的那部分模型的计算，实现并行训练
存储tp/pp/dp/ep/cp 各种并行度配置大小. 并且能够从 tp-dp str格式的并行配置里获取 tp/dp对应的mask和并行度大小设置.
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.需要有远程衰减性

为了便于加速计算, 可以等价优化为下面这种向量乘法的形式:

tokentype_embedding

类型嵌入层，用于区分输入中不同类型的token, 例如，在BERT中用于区分两个句子，而在某些GPT变种或特定任务中可能用于区分不同类型的输入数据，如对话中的提问和回答.

transformer

self.decoder就是上面通过ModuleSpec获得的module, 可以根据配置选择普通的selfAttention, 还是MLA.

MLA原理: 在模型能力不变基础上，通过KV低秩压缩, 使得推理的KVcache显存占用和计算效率上对比MHA性能有明显提升.

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原理博客