megatron_utils.py — 提供 Megatron-Core 模型的 CPU 参数备份和恢复功能

文件路径: verl/experimental/fully_async_policy/megatron_utils.py

文件概述

提供 Megatron-Core 模型的 CPU 参数备份和恢复功能。在全异步训练中，当 trigger_parameter_sync_step > 1 时，Trainer 需要在多个训练步之间保存/恢复模型参数，以支持 MIS（Multiple Importance Sampling）修正。

关键代码讲解

1. copy_megatron_model_to_cpu - 复制模型到 CPU

@torch.no_grad()
def copy_megatron_model_to_cpu(models):
    """
    将 Megatron 模型参数复制到 CPU（非破坏性复制，GPU 数据保持不变）。
    """
    cpu_state = {}

    for model_idx, model_chunk in enumerate(models):
        if isinstance(model_chunk, DDP):
            # 处理 DDP 包装的模型
            model_chunk_all_buffers = [model_chunk.buffers, model_chunk.expert_parallel_buffers]
            buffer_states = []

            for buffers in model_chunk_all_buffers:
                buffer_list = []
                for buffer in buffers:
                    buffer_state = {}
                    if buffer.param_data.storage().size() > 0:
                        # 复制到 CPU 并 pin_memory（加速后续 GPU 传输）
                        buffer_state["param_data"] = buffer.param_data.data.cpu().clone().pin_memory()
                    buffer_list.append(buffer_state)
                buffer_states.append(buffer_list)

            cpu_state[f"model_chunk_{model_idx}"] = {"buffer_states": buffer_states, "is_ddp": True}
        else:
            # 非 DDP 模型（如 ref 模型）
            model_state = {}
            for name, param in model_chunk.named_parameters():
                model_state[name] = {"data": param.data.cpu().clone().pin_memory()}
            cpu_state[f"model_chunk_{model_idx}"] = {"model_state": model_state, "is_ddp": False}

    return cpu_state

2. restore_megatron_model_from_cpu - 从 CPU 恢复模型

@torch.no_grad()
def restore_megatron_model_from_cpu(models, cpu_state):
    for model_idx, model_chunk in enumerate(models):
        chunk_state = cpu_state[f"model_chunk_{model_idx}"]

        if chunk_state["is_ddp"]:
            # 恢复 DDP 缓冲区
            for buffers, buffer_list in zip(model_chunk_all_buffers, buffer_states):
                for buffer, buffer_state in zip(buffers, buffer_list):
                    if "param_data" in buffer_state:
                        buffer.param_data.data.copy_(
                            buffer_state["param_data"].to(buffer.param_data.device)
                        )
        else:
            # 恢复非 DDP 模型
            for name, param in model_chunk.named_parameters():
                if name in model_state:
                    param.data.copy_(model_state[name]["data"].to(param.device))

使用场景

# 训练第 1 步：保存参数到 CPU
cpu_state_v1 = copy_megatron_model_to_cpu(models)

# 训练第 2 步：需要用 v1 参数计算 old_log_prob
cpu_state_v2 = copy_megatron_model_to_cpu(models)  # 先保存 v2
restore_megatron_model_from_cpu(models, cpu_state_v1)  # 恢复 v1
old_log_prob = compute_log_prob(models, batch)  # 用 v1 计算
restore_megatron_model_from_cpu(models, cpu_state_v2)  # 恢复 v2 继续训练

核心类/函数列表

名称	类型	说明
copy_megatron_model_to_cpu()	函数	将模型参数复制到 CPU
restore_megatron_model_from_cpu()	函数	从 CPU 恢复模型参数到 GPU

与其他模块的关系

• 被 FullyAsyncTrainer._compute_old_log_prob() 间接使用（通过 actor_wg.save_model_to_cpu / restore_model_from_cpu）
• 专用于 Megatron-Core 分布式训练框架

小结

megatron_utils.py 解决了 Megatron-Core 模型在异步训练中的参数快照问题。通过 pin_memory() 加速 CPU-GPU 传输，通过区分 DDP/非 DDP 模型确保兼容性。这是实现 MIS 修正的基础设施。