NVIDIA / NVIDIA/Megatron-LM

[ENHANCEMENT]When load_ckpt is called and the obtained iteration count equals args.train_iters, the train_step process will be directly skipped. If, at this point, the save_checkpoint function may encounter an error.

Open
#1,310 4 comments 1 reaction 0 assignees View on GitHub
bug community-request
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 3h
Merged PRs (30d)
272

Description

**Is your feature request related to a problem? Please describe.**
When load_ckpt is called and the obtained iteration count equals args.train_iters, the train_step process will be directly skipped. If, at this point, the condition if args.save and iteration != 0 and iteration % args.save_interval != 0: is entered, the save_checkpoint function may encounter an error due to the absence of optimizer-related parameters, leading to the failure of the training task. In the case of Torch Elastic mode with asynchronous checkpointing (async ckpt) enabled, this can result in infinite resumption of training.

**Describe alternatives you've considered**
Exit process when adding iteration==args.train_iters

if iteration == args.train_iters:
sys.exit(0)
while iteration < args.train_iters:
if args.profile and torch.distributed.get_rank() in args.profile_ranks:
if args.use_pytorch_profiler:
prof.step()
elif iteration == args.profile_step_start:
torch.cuda.cudart().cudaProfilerStart()

Contributor guide

Open the contributing guide

Research direction

Start at the load_ckpt and train_step entry points, then trace the save_checkpoint call under the shown save condition when iteration equals args.train_iters. Reproduce with Torch Elastic and asynchronous checkpointing enabled, and verify that a checkpoint loaded at the target iteration exits cleanly without attempting a save that lacks optimizer parameters or repeatedly resuming training.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.