[Common][PyTorch] Fuse scaling and unscaling of bf16 momentums into kernels #2632
+217
−111
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
Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: Xin Yao <xiny@nvidia.com>
Member
Author
Contributor
Greptile OverviewGreptile SummaryThis PR optimizes BF16 momentum handling in FusedAdam by fusing scaling/unscaling operations directly into the CUDA kernels, eliminating explicit FP32 copies and reducing peak memory footprint. It also enables CUDA Graphs (capturable mode) for BF16 momentums. Key changes:
Implementation approach: Confidence Score: 4/5
Important Files Changed
Sequence DiagramsequenceDiagram
participant User
participant FusedAdam as FusedAdam (Python)
participant AdamKernel as Adam CUDA Kernel
participant Moments as BF16 Moments (exp_avg/exp_avg_sq)
Note over User,Moments: BF16 Moments with Fused Scaling Flow
User->>FusedAdam: step() with BF16 moments
FusedAdam->>FusedAdam: Check fuse_unscale flag (BF16 moments)
FusedAdam->>FusedAdam: get_unscaled_state(skip_unscale=True)
Note over FusedAdam: Skip explicit BF16→FP32 conversion
FusedAdam->>AdamKernel: Call kernel with BF16 moments directly
AdamKernel->>AdamKernel: Load BF16 moments, cast to FP32 for math
AdamKernel->>AdamKernel: Perform Adam update in FP32
AdamKernel->>Moments: Store updated moments as BF16
Note over AdamKernel,Moments: Fused scaling: FP32→BF16 cast in kernel
AdamKernel-->>FusedAdam: Return
FusedAdam->>FusedAdam: Skip explicit scaling (fuse_unscale=True)
Note over FusedAdam: Avoid explicit FP32 copy & scaling
FusedAdam-->>User: Complete
|
Signed-off-by: Xin Yao <xiny@nvidia.com>
for more information, see https://pre-commit.ci
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Comment on lines
+410
to
+422
| @pytest.mark.skipif(not is_bf16_available(), reason="bf16 if not supported") | ||
| def test_bf16_exp_avg_and_exp_avg_sq(self): | ||
| self.gen_precision_aware_test( | ||
| use_fp8_params=False, | ||
| param_dtype=torch.bfloat16, | ||
| use_master_weights=True, | ||
| master_weight_dtype=torch.float32, | ||
| grad_dtype=torch.float32, | ||
| exp_avg_dtype=torch.bfloat16, | ||
| exp_avg_sq_dtype=torch.bfloat16, | ||
| master_rtol=2e-3, | ||
| master_atol=2e-3, | ||
| ) |
Contributor
There was a problem hiding this comment.
Consider adding a test for capturable mode (CUDA Graphs) with BF16 momentums, since the PR description mentions "Enable CUDA Graphs for BF16 momentums" as a key feature. The current test only covers non-capturable mode.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
Signed-off-by: Xin Yao <xiny@nvidia.com>
Comment on lines
+410
to
+422
| @pytest.mark.skipif(not is_bf16_available(), reason="bf16 if not supported") | ||
| def test_bf16_exp_avg_and_exp_avg_sq(self): | ||
| self.gen_precision_aware_test( | ||
| use_fp8_params=False, | ||
| param_dtype=torch.bfloat16, | ||
| use_master_weights=True, | ||
| master_weight_dtype=torch.float32, | ||
| grad_dtype=torch.float32, | ||
| exp_avg_dtype=torch.bfloat16, | ||
| exp_avg_sq_dtype=torch.bfloat16, | ||
| master_rtol=2e-3, | ||
| master_atol=2e-3, | ||
| ) |
Contributor
There was a problem hiding this comment.
Test only covers non-capturable mode. Add test for capturable mode with BF16 momentums since PR enables CUDA Graphs for this case.
Note: Check that gen_precision_aware_test supports a capturable parameter, or create a separate test method.
kunlunl
approved these changes
Feb 2, 2026
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 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.
Description
This PR only enables this feature for BF16 momentums, because BF16 doesn't require real scaling and unscaling, but just type converting.
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: