[PP] Ensure loss is visible on console for users#946
Merged
Conversation
facebook-github-bot
added
the
CLA Signed
H-Huang
approved these changes
Mar 7, 2025
Member
H-Huang
left a comment
H-Huang
left a comment
There was a problem hiding this comment.
Nice!! Thanks for adding it so quickly
MaxiBoether
pushed a commit
to eth-easl/torchtitan-mixtera
that referenced
this pull request
Apr 17, 2025
This is similar in spirit to [PR_944](pytorch#944) (cc @lkhphuc) but takes a slightly different approach. Problem - users that default turn on PP training will get -1 for their loss. This is b/c by default, rank 0 is the only one logged. However, for *most* PP schedules, the loss is output on the last rank. Thus, users see -1 for loss and it's a bad/confusing experience. This PR adds a check to review both the current PP schedule (b/c for VBlocks, loss is returned on 0) and if it is a last rank loss schedule, then it checks that the first rank of the last stage is visible in the LOG_RANK environment variable. If not, it warns the user, using Red for the warning if color is enabled, and highlights the rank they should add in yellow: <img width="1236" alt="Screenshot 2025-03-07 at 11 51 46 AM" src="https://github.com/user-attachments/assets/02b18870-90bb-4cfb-89c1-3e92d2fb9bfb" /> Note that I attempted to then modify the LOG_RANK to add the missing last rank...but it has no effect. This is b/c the --log_rank_filter passed into torchrun is fixed and thus the env has no effect. We can fix this by moving to our own filtering via python log filtering (thanks to @d4l3k for this idea) and then it would auto-update. The tradeoff is that we have to init distributed first (to understand the ranks) meaning that at launch, there's a bit of delay before the first logging. From there, then NCCL warnings are not suppressed b/c they are emitted from .cpp file vs torchrun filtering controls that...so we get some additional console spam. This PR thus sticks to a simple warning with Red highlight (assuming color is on) and provide the user how to fix.
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
3 participants
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.
This is similar in spirit to PR_944 (cc @lkhphuc) but takes a slightly different approach.
Problem - users that default turn on PP training will get -1 for their loss. This is b/c by default, rank 0 is the only one logged.
However, for most PP schedules, the loss is output on the last rank.
Thus, users see -1 for loss and it's a bad/confusing experience.
This PR adds a check to review both the current PP schedule (b/c for VBlocks, loss is returned on 0) and if it is a last rank loss schedule, then it checks that the first rank of the last stage is visible in the LOG_RANK environment variable.
If not, it warns the user, using Red for the warning if color is enabled, and highlights the rank they should add in yellow:
Note that I attempted to then modify the LOG_RANK to add the missing last rank...but it has no effect. This is b/c the --log_rank_filter passed into torchrun is fixed and thus the env has no effect.
We can fix this by moving to our own filtering via python log filtering (thanks to @d4l3k for this idea) and then it would auto-update.
The tradeoff is that we have to init distributed first (to understand the ranks) meaning that at launch, there's a bit of delay before the first logging. From there, then NCCL warnings are not suppressed b/c they are emitted from .cpp file vs torchrun filtering controls that...so we get some additional console spam.
This PR thus sticks to a simple warning with Red highlight (assuming color is on) and provide the user how to fix.