NVIDIA / NVIDIA/Megatron-LM

[BUG] Loss logging double-counts the first iteration after checkpoint resume

Open
#6,237 0 comments 0 reactions 0 assignees View on GitHub
community-request waiting-on-maintainers
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 3h
Merged PRs (30d)
272

Description

## Description

When training resumes from a checkpoint, the first iteration of the resumed process is treated as a fresh-run first iteration by `training_log`. This causes the loss accumulator to retain the first resumed loss and add it to the next iteration.

The issue is visible with `--log-interval 1`. It affects reported loss/timing statistics, but not the forward computation, gradients, optimizer update, or checkpoint contents.

## Environment

- Repository: `NVIDIA/Megatron-LM`
- Branch: `core_r0.17.0`
- Commit: `963bf39218e8bb83a1203b40293358498322be50`
- BF16 distributed training
- Observed on an Ascend/NPU run; the affected code path appears device-independent
- `--log-interval 1`
- Checkpoint saved at iteration 50, then resumed at iteration 51

## Reproduction

1. Run a deterministic BF16 training job through iteration 53 with `--save-interval 50` and `--log-interval 1`.
2. Run the same job from the beginning and separately resume from the iteration-50 checkpoint.
3. Compare the loss logs for iterations 51--53.

The relevant control flow is:

```python
# train()
is_first_iteration = True

# training_log()
should_reset = not is_first_iteration

if iteration % args.log_interval == 0 or is_first_iteration:
...
if should_reset:
total_loss_dict[key] = torch.tensor([0.0], ...)
```

On resume, `is_first_iteration` is reset to `True`, while `args.iteration > 0` and `total_loss_dict` starts empty.

## Observed behavior

At iteration 52:

- uninterrupted run: `lm loss = 11.156137466430664`, `advanced iterations = 1`
- resumed run: accumulated loss `= 22.273590087890625`, `advanced iterations = 2`, average `= 11.136795043945312`

Iteration 53 becomes normal again because iteration 52 resets the accumulator. Additional boundary diagnostics show that the batch, local/global loss, and pre-optimizer gradient hashes match between uninterrupted and resumed runs; only the `loss_log` accumulator differs.

## Expected behavior

After resuming, `--log-interval 1` should report the loss for the current iteration only, matching an uninterrupted run. The first resumed iteration should not be retained and double-counted in the next log interval.

## Possible fix

A minimal fix is to preserve the current first-iteration accumulation behavior only for a fresh run:

```python
should_reset = not is_first_iteration or args.iteration > 0
```

Alternatively, the first-iteration logging state could be initialized from the restored global iteration rather than process lifetime.

Contributor guide

Open the contributing guide

Research direction

Start in the train() and training_log() control flow described in the issue, then reproduce the checkpoint-resume scenario with --log-interval 1 and compare iterations 51–53. Done means the first resumed iteration is not retained in the next accumulator interval and the resumed loss logs match an uninterrupted run without changing forward computation, gradients, optimizer updates, or checkpoint contents.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
70/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.