执行摘要
- 一句话:修复 Megatron 非 vanilla 路径保存 ckpt 时忽略 strict 配置
- 推荐动作:值得快速阅读即可:单行改动澄清了 Megatron-Bridge 保存路径的 strict 透传设计。关注两点:一是
checkpoint_config.strict 的默认值及其文档语义,二是后续是否应为 strict=False 的导出场景补充测试。总体属于低风险修复。
功能与动机
PR body 仅有模板内容,未提供额外背景;commit 信息为“fix megatron save checkpoints with strict false”。结合代码推断:当用户配置 checkpoint_config.strict=False 时(典型场景是 Megatron 训练的自定义模型结构与 HF 结构存在差异,导出时允许缺失或多余的权重),非 vanilla_bridge 路径的 save_hf_weights 调用未传 strict,bridge 默认按严格模式校验并抛错,导致保存失败。该修复让 strict 配置对保存路径生效,与 vanilla 路径保持一致。
实现拆解
实现拆解:
- 问题定位:
verl/utils/checkpoint/megatron_checkpoint_manager.py 的 _save_model_as_hf_via_bridge 方法,负责通过 Megatron-Bridge 将模型权重导出为 HF 格式。该方法分为 vanilla_bridge、PEFT、普通导出三条路径。
- 核心改动:将普通导出分支的调用由
self.bridge.save_hf_weights(self.model, hf_ckpt_path) 改为 self.bridge.save_hf_weights(self.model, hf_ckpt_path, strict=self.checkpoint_config.strict),把 checkpoint 配置中的 strict 语义透传给 bridge,使 strict=False 时允许权重不完全匹配。
- 影响范围:改动仅影响非 vanilla_bridge、非 PEFT 的保存路径;vanilla 路径与 PEFT 路径不受影响。该行为依赖
checkpoint_config.strict 的默认值,若默认 True 则原有行为完全不变。
- 配套改动:本次无测试、配置或文档变更,属于单行修复。
关键文件:
verl/utils/checkpoint/megatron_checkpoint_manager.py(模块 检查点管理;类别 source;类型 core-logic;符号 _save_model_as_hf_via_bridge): 唯一改动文件,位于 Megatron checkpoint 保存主路径,修复非 vanilla bridge 导出 HF 权重时 strict 配置失效的问题。
关键符号:_save_model_as_hf_via_bridge
关键源码片段
verl/utils/checkpoint/megatron_checkpoint_manager.py
唯一改动文件,位于 Megatron checkpoint 保存主路径,修复非 vanilla bridge 导出 HF 权重时 strict 配置失效的问题。
def _save_model_as_hf_via_bridge(self, hf_ckpt_path: str):
"""Save model weights through megatron-bridge."""
if self.vanilla_bridge:
# vanilla 路径:直接调用 bridge 保存完整权重(带扩展参数)
self.bridge.save_weights(self.model, hf_ckpt_path, **self._get_bridge_extended_args())
else:
if self.peft_cls is not None:
# PEFT 场景:只导出 adapter,与基础权重解耦
hf_adapter_ckpt_path = os.path.join(hf_ckpt_path, "adapter")
self.bridge.save_hf_adapter(self.model, hf_adapter_ckpt_path, self.peft_cls)
log_with_rank(
f"Saved HF PEFT adapter checkpoint to {hf_adapter_ckpt_path}",
rank=self.rank,
logger=logger,
log_only_rank_0=True,
)
else:
# 普通 HF 权重导出:透传 strict 配置,使 strict=False 时允许
# 权重结构与 HF 不完全匹配(修复前未传 strict,强制严格校验,
# 自定义 Megatron 结构模型导出会失败)
self.bridge.save_hf_weights(self.model, hf_ckpt_path, strict=self.checkpoint_config.strict)
评论区精华
该 PR 没有任何 review 评论或 issue 讨论。唯一的 review 记录来自合并者 wuxibin89,状态为 DISMISSED(无正文),说明该改动未经实质性评审讨论即被合并,可能作为低风险单行修复直接通过。
风险与影响
关联脉络
- PR #7223 [megatron] feat: PP/VPP support for the steady delta export: 同属 Megatron checkpoint 导出链路,改动 delta_export.py 与 transformer_impl.py,均在加固导出正确性。
- PR #7181 [megatron] feat: delta_sharded on Megatron-Bridge param mappings (TP+EP, hybrid-Mamba): 同一 Megatron-Bridge 权重导出/同步功能线,涉及 delta_export 与 checkpoint engine。
- PR #7199 [megatron] fix: preserve generation config during merge: 同为 Megatron 保存/合并链路 bugfix,关注导出产物完整性。
参与讨论