以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,
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原理博客
![题解：AT_abc353_f [ABC353F] Tile Distance](https://cdn.luogu.com.cn/upload/image_hosting/58f4e31e.png)
