以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.需要有远程衰减性 。

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

tokentype_embedding

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

transformer

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

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

最后此篇关于[megatron代码阅读]1.初始化和组网的文章就讲到这里了,如果你想了解更多关于[megatron代码阅读]1.初始化和组网的内容请搜索CFSDN的文章或继续浏览相关文章，希望大家以后支持我的博客！ 。