loader.py — HF 权重加载到 mcore 模型

文件路径

verl/models/mcore/loader.py

文件概述

将 HuggingFace 格式的合并权重加载到分片的 Megatron-Core 模型中。核心挑战是将完整权重按照 TP (Tensor Parallel)、PP (Pipeline Parallel) 的切分方式正确分发到各个 GPU。

关键代码讲解

核心流程

def load_state_dict_to_megatron_gptmodel(state_dict, wrapped_models, config, params_dtype):
    """加载合并的 state_dict 到分片的 Megatron 模型"""
    layer_map = _megatron_calc_layer_map(config)

    for name, param in state_dict.items():
        # 1. 确定该参数属于哪个 PP stage
        pp_rank = layer_map[layer_idx][0]

        # 2. 从 rank 0 广播到对应的 PP stage
        param = _broadcast_tensor(param, src_rank=0, target_pp_rank=pp_rank)

        # 3. 按 TP 切分
        if "q_proj" in name or "k_proj" in name or "v_proj" in name:
            param = _broadcast_tp_shard_tensor_qkv(param, ...)
        elif "gate_proj" in name or "up_proj" in name:
            param = _broadcast_tp_shard_tensor_gate_up(param, ...)
        else:
            param = _broadcast_tp_shard_tensor(param, ...)

        # 4. 写入模型参数
        model_param.data.copy_(param)

TP 切分辅助函数

def _broadcast_tp_shard_tensor(param, dim=0):
    """按 TP 切分权重"""
    tp_rank = mpu.get_tensor_model_parallel_rank()
    tp_size = mpu.get_tensor_model_parallel_world_size()
    chunk_size = param.shape[dim] // tp_size
    return param.narrow(dim, tp_rank * chunk_size, chunk_size)

def _broadcast_tp_shard_tensor_qkv(param):
    """QKV 权重的 TP 切分（需要考虑 GQA 的头数分布）"""
    # Q: 按头数切分
    # K/V: 按 GQA 组切分（可能比 Q 少）

def _broadcast_tp_shard_tensor_gate_up(param):
    """gate_proj + up_proj 合并后按 TP 切分"""
    # 先合并两个权重，然后按列切分

PP 层映射

def _megatron_calc_layer_map(config):
    """计算全局层号到 (pp_rank, vpp_rank, local_layer_idx) 的映射"""
    for pp_rank_idx in range(pp_size):
        for virtual_pp_rank_idx in range(virtual_pp_size):
            layer_offset = virtual_pp_rank_idx * (num_layers // virtual_pp_size) + ...
            for layer_idx in range(num_layers_per_model):
                layer_map[layer_offset + layer_idx] = (pp_rank_idx, virtual_pp_rank_idx, layer_idx)

核心函数列表

函数名	作用
load_state_dict_to_megatron_gptmodel()	主入口：加载 HF 权重到 mcore
_megatron_calc_layer_map()	计算层到 PP stage 的映射
_broadcast_tensor()	跨 PP 广播权重
_broadcast_tp_shard_tensor()	按 TP 切分权重
_broadcast_tp_shard_tensor_qkv()	QKV 权重的 TP 切分
_broadcast_tp_shard_tensor_gate_up()	gate/up 权重的合并 + 切分

与其他模块的关系

• 被 weight_loader_registry.py 注册为权重加载器
• 使用 saver.py 的 _megatron_calc_global_rank() 函数
• 被训练流程中的权重同步代码调用

小结

权重加载器处理了从完整 HF 权重到分片 mcore 模型的分发。关键挑战是正确处理 TP 切分（特别是 QKV 和 gate/up 的合并权重）和 PP 层映射。