Fix mixed precision call in group norm sharded. #1380
Fix mixed precision call in group norm sharded. #1380coreyjadams wants to merge 4 commits intoNVIDIA:mainfrom
Conversation
…ow variances are combined across GPUs.
| if weight is not None: | ||
| weight = weight.to(input.dtype) | ||
| if bias is not None: | ||
| bias = bias.to(input.dtype) | ||
|
|
There was a problem hiding this comment.
This fixes a mixed precision crash.
torch tweaking the python dispatch behavior and DTensor. - adding more layers to handle select - add more reliable handling of casting DTensor to ShardTensor. In particular, the focus is on making sure we maintain proper autograd graphs. - switch to a first-principles implemetation of group norm. It's more stable, simpler, and while it might be a little slower the upcoming torch.compile work can address that. - add a dedicated view handler at functional and dispatch level. It's necessary at this point to wrap our own view implementation due to the differences with DTensor.
Greptile OverviewGreptile SummaryThis PR fixes two critical bugs in the sharded group normalization implementation and adds comprehensive AMP testing. Key Changes:
Additional Changes:
Important Files Changed
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| ) | ||
|
|
||
| if check_version_spec("torch", "2.10.0"): | ||
| if check_version_spec("torch", "2.10.0a"): |
There was a problem hiding this comment.
This is to get pre-release versions. Fixes #1394
|
This is getting broken up in to smaller PRs for easier review. |
1 participant
Also fix a math error in how variances are combined across GPUs.
PhysicsNeMo Pull Request
Description
Checklist
Dependencies
Review Process
All PRs are reviewed by the PhysicsNeMo team before merging.
Depending on which files are changed, GitHub may automatically assign a maintainer for review.
We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.
AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.