Skip to content

Commit 321556f

Browse files
committed
fix dist not exist case in checkpoint conversion
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
1 parent 982ec64 commit 321556f

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

src/megatron/bridge/training/checkpointing.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -638,7 +638,8 @@ def save_checkpoint(
638638
content_metadata=sharded_sd_metadata,
639639
)
640640
# [ModelOpt]: save sharded modelopt_state
641-
if not cfg.dist.use_decentralized_pg:
641+
# cfg.dist can be None during checkpoint conversion (save_megatron_model)
642+
if not (cfg.dist and cfg.dist.use_decentralized_pg):
642643
save_sharded_modelopt_state(model, checkpoint_name, (ckpt_cfg.ckpt_format, 1))
643644
else:
644645
# [ModelOpt]: Inject modelopt_state into state_dict

0 commit comments

Comments
 (0)