Repository navigation
Conversation
The periodic logging branch in run_pretraining computed average_loss as a per-rank sum and logged it directly, without the div_(world_size) and all_reduce that the final-step branch in the same function performs. Every periodic log line reported rank 0's local-shard loss instead of the true cross-rank mean; only the one-time final_loss logged at the end of training was ever actually averaged. Mirrors the final-step branch's existing pattern: divide by log_freq * gradient_accumulation_steps, then divide by world_size and all_reduce before copying into the host-side stat, all before the is_main_process() gate (the collective op must run unconditionally on every rank or it deadlocks). Verified via a standalone 8-rank simulation: old behavior reports rank 0's value (0.12), fixed behavior reports the true cross-rank mean (0.1235). No torch install available on this host to run the real test suite. Signed-off-by: Siddhardha Nanda <99672439+SID-6921@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
In
PyTorch/LanguageModeling/BERT/run_pretraining.py, the periodic logging branch computedaverage_lossas a per-rank sum and logged it directly, without thediv_(world_size)+all_reducethe final-step branch in the same function performs. Every periodic log line reported rank 0's local-shard loss instead of the true cross-rank mean; only the one-timefinal_losslogged at the end of training was ever actually averaged.Fix
Mirror the final-step branch's existing pattern in the periodic branch: divide by
log_freq * gradient_accumulation_steps, then divide byworld_sizeandall_reduce, before copying into the host-side stat — all placed before theis_main_process()gate, since the collective op must run unconditionally on every rank or it deadlocks.Test plan
No
torchinstall available on this host to run the actual test suite. Verified via:git diff/ast.parsefor a clean, syntactically valid change🤖 Generated with Claude Code