megatron_utils.py — Megatron 模型转换工具

文件路径: verl/utils/megatron_utils.py

文件概述

这是 verl 中最大的工具文件（约 68.5KB），提供 HuggingFace 模型格式与 Megatron 模型格式之间的双向转换。Megatron 使用特定的参数命名和分片策略，与 HuggingFace 的格式不同。

背景知识

Megatron-LM 是 NVIDIA 的大模型训练框架，使用张量并行和流水线并行。它的模型参数命名（如 layers.0.self_attention.query_key_value.weight）与 HuggingFace 的命名（如 model.layers.0.self_attn.q_proj.weight）完全不同。

核心功能

1. HF → Megatron 转换

def hf_to_megatron(hf_state_dict, config):
    """将 HuggingFace 模型权重转为 Megatron 格式"""
    # 重命名 key
    # 合并 Q/K/V 为 QKV
    # 按 TP 并行度分片

2. Megatron → HF 转换

def megatron_to_hf(megatron_state_dict, config):
    """将 Megatron 模型权重转为 HuggingFace 格式"""
    # 反向重命名
    # 拆分 QKV 为 Q/K/V
    # 合并 TP 分片

支持的模型架构

• Llama / Llama2 / Llama3
• Qwen / Qwen2
• DeepSeek
• Mistral
• 以及其他 HuggingFace 标准架构

与其他模块的关系

• 被 Megatron worker 在模型初始化时调用
• 被 checkpoint 模块在保存/加载时使用
• 依赖 megatron/tensor_parallel.py 的分片信息

小结

megatron_utils.py 是连接 HuggingFace 生态和 Megatron 训练框架的桥梁。它处理了复杂的参数名映射和张量分片逻辑，让用户可以用 HuggingFace 预训练权重开始 Megatron 训练。