Skip to content

Safely access state dict args in load ckpt #1957

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
maanug-nv merged 4 commits into from
Oct 29, 2025
Merged

Safely access state dict args in load ckpt #1957

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
8103049
safely access state dict args in load ckpt
maanug-nv Oct 24, 2025
681ce76
rename var
maanug-nv Oct 26, 2025
e42f8de
update tp pp default
maanug-nv Oct 27, 2025
8d96d7b
Merge branch 'main' into safe-statedict-args-access
maanug-nv Oct 28, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 17 additions & 7 deletions megatron/training/checkpointing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import shutil
import sys
import threading
import types
from argparse import Namespace
from enum import Enum, auto
from logging import getLogger
Expand Down Expand Up @@ -1412,26 +1413,35 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, load_arg='load',
ignore_rng_state = False
ignore_rerun_state = True
if ckpt_format == "torch_dist":
ckpt_args = types.SimpleNamespace()
if state_dict is not None and "args" in state_dict:
ckpt_args = state_dict.get("args")

if not hasattr(ckpt_args, "tensor_model_parallel_size"):
print_rank_0("WARNING: TP size not found in checkpoint args, using 0 as default.")
if not hasattr(ckpt_args, "pipeline_model_parallel_size"):
print_rank_0("WARNING: PP size not found in checkpoint args, using 0 as default.")

ckpt_tp_pp = (
state_dict['args'].tensor_model_parallel_size,
state_dict['args'].pipeline_model_parallel_size,
getattr(ckpt_args, "tensor_model_parallel_size", 0),
getattr(ckpt_args, "pipeline_model_parallel_size", 0),
)
Comment on lines +1420 to 1428

@ananthsub ananthsub Oct 27, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what about instead of adding fallback values, the check for RNG states in L1443 is modified to directly check if the state dict contains args. we could safely assume that if using a converted checkpoint, we do not need to resume the rng states. the same goes for the re-run state machine check below. the point would be that we should avoid collisions with checkpoints saved with TP1/PP1

@maanug-nv maanug-nv Oct 27, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

good point. i think using 0,0 as the fallback should accomplish this. let me know if you see any issue with that.

run_tp_pp = (
args.tensor_model_parallel_size,
args.pipeline_model_parallel_size,
)

ckpt_world_size = getattr(state_dict['args'], 'world_size', 0)
ckpt_world_size = getattr(ckpt_args, 'world_size', 0)
run_world_size = getattr(args, 'world_size', 0)
ckpt_dp = getattr(state_dict['args'], 'data_parallel_size', 0)
ckpt_dp = getattr(ckpt_args, 'data_parallel_size', 0)
run_dp = getattr(args, 'data_parallel_size', 0)
mismatch_msg = "(TP, PP) mismatch after resume ({} vs {} from checkpoint)".format(
run_tp_pp, ckpt_tp_pp
)

# Determine if RNG state will be loaded
if (ckpt_tp_pp == run_tp_pp and not release and not args.finetune and not args.no_load_rng
and not getattr(state_dict['args'], 'no_save_rng', False)):
and not getattr(ckpt_args, 'no_save_rng', False)):
gen_sd_rng_state = get_rng_state(args.ckpt_format) # we can load the rng state
else:
ignore_rng_state = True
Expand All @@ -1446,7 +1456,7 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, load_arg='load',
print_rank_0(f'sharded_state_dict metadata loaded from the checkpoint: {sharded_sd_metadata}')
# Determine if optimizer state will be loaded
if (not release and not args.finetune and not args.no_load_optim
and not getattr(state_dict['args'], 'no_save_optim', False)):
and not getattr(ckpt_args, 'no_save_optim', False)):
gen_sd_optim = optimizer
gen_sd_opt_param_scheduler = opt_param_scheduler

Expand All @@ -1457,7 +1467,7 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, load_arg='load',
# (for MCore v0.13+ checkpoints `sharded_sd_metadata is not None`)
sharded_sd_metadata = {
'distrib_optim_sharding_type': ('fully_sharded_model_space'
if getattr(state_dict['args'], 'ckpt_fully_parallel_save', False)
if getattr(ckpt_args, 'ckpt_fully_parallel_save', False)
else 'dp_zero_gather_scatter'),
}
if (
Expand Down
Loading