qwen2_saver.py — Qwen2 权重合并保存

文件路径

verl/models/qwen2/megatron/checkpoint_utils/qwen2_saver.py

文件概述

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

与 LLaMA saver 的区别

1. QKV bias 的收集与拆分

# 收集 QKV weight
_broadcast_tp_shard_tensor_qkv(
    sync_layer.self_attn.qkv_proj.weight,
    f"{layer_name}.self_attn.q_proj.weight",
    f"{layer_name}.self_attn.k_proj.weight",
    f"{layer_name}.self_attn.v_proj.weight",
    src_pp_rank=src_pp_rank,
)
# 收集 QKV bias（Qwen2 特有）
_broadcast_tp_shard_tensor_qkv(
    sync_layer.self_attn.qkv_proj.bias,
    f"{layer_name}.self_attn.q_proj.bias",
    f"{layer_name}.self_attn.k_proj.bias",
    f"{layer_name}.self_attn.v_proj.bias",
    src_pp_rank=src_pp_rank,
)

2. tie_word_embeddings 处理

if tie_word_embeddings:
    print_rank_0("tie word embedding skip load lm_head...")
else:
    # 正常收集 lm_head

合并流程

与 LLaMA saver 对称： 1. 收集 embedding（TP 分片拼接） 2. 逐层收集权重（包括 QKV bias） 3. 收集 final norm + lm_head 4. 只有 rank 0 持有完整 state_dict

小结

Qwen2 权重保存器与 LLaMA 版本结构一致，增加了 QKV bias 收集和 tie_word_embeddings 跳过逻辑。