Loss disagreement between TP=1 and TP=2 #631
Draft
+24
−15
Conversation
This file contains 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
sichu2023
requested review from
pstjohn,
jstjohn,
skothenhill-nv,
farhadrgh,
dorotat-nv and
malcolmgreaves
as code owners
|
Loss curve when initializing |
❌ 31 Tests Failed:
View the top 3 failed tests by shortest run time
To view more test analytics, go to the Test Analytics Dashboard |
jstjohn
reviewed
Jan 22, 2025
| # embedding_activation_buffer=self.embedding_activation_buffer, | ||
| # grad_output_buffer=self.grad_output_buffer, | ||
| # ) | ||
| self.output_layer = torch.nn.Linear(config.hidden_size, self.vocab_size, bias=True) |
There was a problem hiding this comment.
This probably has bugs with bias getting out of sync between TP ranks during training FYI. See https://jirasw.nvidia.com/browse/BIONEMO-668 and https://jirasw.nvidia.com/browse/BIONEMO-666 and https://nvidia.slack.com/archives/C074Z808N05/p1737508003987919 and https://nvidia.slack.com/archives/C0434FDLPQV/p1733963545314469
Also if your concern is when you do TP=2 that the logit dim is 1/2 that may be because columnparallellinear splits along the logit vocab dimension, and ideally vocab parallel cross entropy knows how to reduce across this.
There was a problem hiding this comment.
Yup. Only temporarily placed torch.nn.Linear to debug. Will revert back to ColumnParallelLinear after so.
There was a problem hiding this comment.
I am less concerned about "1/2 logits dim" (128 dim) but more concerned about torch.nn.Linear giving 256 dim on TP=2. 128 dim should be the correct dim (33 + padding).
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.
Description
Loss values from
unreduced_token_loss_fndisagree even though thelogitsare the same when TP=1 and TP=2. Interestingly, this introduces a nearly-constant offset in training curve and does not impactgrad_norm. The offset slowly shrinks over 100 training steps.Suspected root cause:
vocab_parallel_cross_entropytakes into account oflogitsat the padded vocab dimensions while it should have ignored them. As training progresses, more probabilities are then assigned to "real vocabs" so the padding vacabs slowly becomes less relevant, leading to the closing gap between TP=1 and TP=2 loss curves.logitsdimensions also expanded from 128 to 256 when replacingColumnParallelLinearwithtorch.nn.Linearalthough the reason is still unclear. Ifvocab_parallel_cross_entropytakes into account of the padding vocab dimensions, this can potentially explain the constant offset observed.On a related issue,
unreduced_token_loss_fnintroduces inplace operation on the token logits.https://github.com/NVIDIA/bionemo-framework/blob/sichu/loss-curve-tp-shift/sub-packages/bionemo-llm/tests/bionemo/llm/model/test_loss.py#L158
See Slack channel here.
Type of changes
Usage
Pre-submit Checklist