[BUG] Loss logging double-counts the first iteration after checkpoint resume
- 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
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