Fix casting in SongUNetPosEmbd and shape in CorrDiff generation #982
Conversation
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
juliusberner
changed the title
juliusberner
changed the title
juliusberner
changed the title
juliusberner
force-pushed
the
jberner/fix_corrdiff
branch
from
92a58a3 to
b7e0382
Compare
juliusberner
force-pushed
the
jberner/fix_corrdiff
branch
from
b7e0382 to
f124ba9
Compare
juliusberner
commented
Jun 18, 2025
CharlelieLrt
added
bug
CharlelieLrt
requested changes
Jun 25, 2025
| @@ -196,7 +196,7 @@ def generate_fn(): | |||
| net=net_reg, | |||
| img_lr=img_lr, | |||
| latents_shape=( | |||
| cfg.generation.seed_batch_size, | |||
| sum(map(len, rank_batches)), | |||
There was a problem hiding this comment.
@juliusberner could you explain the reason for this change? AFAIK the batch dimension of latents_shape is never really used, right?
There was a problem hiding this comment.
This is the batch-size that the output of regression_step is expanded to. Since we compute image_out = image_reg + image_res later, this needs to match the batch-size of the output of diffusion_step.
Comment on lines
-883
to
+889
| if self.prob_channels and out.dtype != self.scalar.dtype: | ||
| self.scalar.data = self.scalar.data.to(out.dtype) | ||
| if self.prob_channels and (not self.training): | ||
| out[:, self.prob_channels] = ( | ||
| out[:, self.prob_channels] * self.scalar | ||
| ).softmax(dim=1) | ||
| elif self.prob_channels and self.training: | ||
| scalar = self.scalar | ||
| if out.dtype != scalar.dtype: | ||
| scalar = scalar.to(out.dtype) | ||
| if self.training: | ||
| out[:, self.prob_channels] = out[:, self.prob_channels] * scalar | ||
| else: | ||
| out[:, self.prob_channels] = ( | ||
| out[:, self.prob_channels] * self.scalar | ||
| (out[:, self.prob_channels] * scalar) | ||
| .softmax(dim=1) | ||
| .to(out.dtype) |
There was a problem hiding this comment.
LGTM, but could you just post below an MRE of the bug you encountered with the former casting logic?
There was a problem hiding this comment.
In amp-bf16 training, the output of softmax is float32, while out.dtype = bfloat16, which gives a RuntimeError: Index put requires the source and destination dtypes match, got BFloat16 for the destination and Float for the source..
Comment on lines
+948
to
+950
| pos_embd = self.pos_embd | ||
| if x.dtype != pos_embd.dtype: | ||
| pos_embd = pos_embd.to(x.dtype) |
There was a problem hiding this comment.
Two remarks:
- Same as above: could you post below a MRE of the bug that you would get with the former casting logic (the MRE could be grouped with the one above)
- Is there some logic problem here? We are accessing
pos_embd.dtypeand right below we are checkingif pos_embd is not None? I think ifself.pos_embdisNonethen we doreturn Noneright away?
There was a problem hiding this comment.
- The assignment
self.pos_embd = self.pos_embd.to(dtype)only works ifpos_embdis a buffer but not if it is a parameter (which is the case ifself.gridtype == "learnable"). Thus, we define a new local variable which works in both cases. positional_embedding_indexingis only called in theforwardifself.pos_embd is not None. If it is called from outside, it would return an empty list[]. How should we handle it?
Comment on lines
+1088
to
+1090
| embeddings = self.pos_embd | ||
| if x.dtype != embeddings.dtype: | ||
| embeddings = embeddings.to(x.dtype) |
There was a problem hiding this comment.
Two remarks:
- Same as above, it would be great if you could post an MRE below (can be grouped with other RMEs for these castings bugs).
- Is there a specific reason to call it
embeddingshere, whereas it was calledpos_embdin thepositional_embedding_indexingmethod? If not, let's try to remain consistent in the names
There was a problem hiding this comment.
- Copying from above: The assignment
self.pos_embd = self.pos_embd.to(dtype)only works ifpos_embdis a buffer but not if it is a parameter (which is the case ifself.gridtype == "learnable"). Thus, we define a new local variable which works in both cases. - I took it from the existing code, but it makes sense to rename it to
pos_embd.
…ding, lead time aware, with compile, apex_gn, etc... Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
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.
PhysicsNeMo Pull Request
Description
fused_actisTrueinApexGroupNormself.pos_embdcan be buffer or parameter) and fix dtype of softmax output (which isfp32) forSongUNetPosEmbd.dataattribute ofself.scalarto enable.compileforSongUNetPosEmbd.Checklist