saver.py — mcore 权重保存/合并

文件路径

verl/models/mcore/saver.py

文件概述

将分片的 Megatron-Core 模型权重收集并合并回 HuggingFace 格式。与 loader.py 互为逆操作。

关键代码讲解

核心流程

def merge_megatron_ckpt_gptmodel(models, config, params_dtype):
    """收集所有 TP/PP 分片的权重，合并为完整的 HF state_dict"""
    state_dict = {}
    for name, param in model.state_dict().items():
        # 1. 从所有 TP rank 收集分片
        gathered_params = all_gather_across_tp(param)

        # 2. 按 TP 维度拼接
        if "linear_qkv" in name:
            # QKV: 拼接后拆分为 q/k/v
            full_param = torch.cat(gathered_params, dim=0)
            q, k, v = split_qkv(full_param)
        elif "linear_fc1" in name:
            # gate+up: 拼接后拆分
            full_param = torch.cat(gathered_params, dim=0)
            gate, up = split_gate_up(full_param)

        # 3. 转换名称并保存
        state_dict[hf_name] = full_param
    return state_dict

全局 rank 计算

def _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=0, cp_rank=0, ep_rank=0):
    """计算给定并行维度坐标的全局 rank"""
    return ((pp_rank * dp_size + dp_rank) * cp_size + cp_rank) * tp_size + tp_rank

核心函数列表

函数名	作用
merge_megatron_ckpt_gptmodel()	Dense 模型权重合并
_megatron_calc_global_rank()	并行坐标到全局 rank 的映射
_megatron_calc_layer_map()	层号到 PP stage 的映射

与其他模块的关系

• 被 weight_loader_registry.py 注册为权重保存器
• 与 loader.py 互为逆操作
• 被 checkpoint 保存流程调用

小结

权重合并是加载的逆过程：收集所有 TP 分片 -> 拼接 -> 拆分合并的 QKV/gate_up -> 转换名称。这个过程在保存 checkpoint 时自动执行。