megatron_checkpoint_manager.py — Megatron 检查点管理

文件路径: verl/utils/checkpoint/megatron_checkpoint_manager.py

文件概述

Megatron 训练策略下的检查点管理器。利用 Megatron Core 的分布式检查点 API 进行高效的并行保存和加载。

核心功能

保存

def save(self, step):
    # 使用 megatron.core.dist_checkpointing 保存
    # 支持异步保存（训练不停等保存完成）
    # 支持完全并行保存（所有 rank 同时写入）

加载

def load(self, step=None):
    # 使用 megatron.core.dist_checkpointing 加载
    # 支持完全并行加载
    # 处理 TP/PP 并行度变化时的重新分片

Megatron 检查点的特殊性

• 使用 FullyParallelSaveStrategyWrapper 实现数据并行组内的完全并行保存
• 使用 FullyParallelLoadStrategyWrapper 实现并行加载
• 支持版本兼容（mcore >= 0.14.0 支持 content_metadata）

与其他模块的关系

• 继承自 checkpoint_manager.py
• 依赖 megatron/dist_checkpointing.py 进行底层操作
• 被 Megatron worker 在训练中调用

小结

利用 Megatron Core 的分布式检查点能力，实现高效的大模型检查点管理。