[skyw]

When fp32 matmul precision is set to medium, we explicitly cast tensor to bfloat16 because PyTorch lack of bfloat16 I/O kernels and will fallback to TF32. However, that has introduced unintentional behavior that all-reduce in TP happens in bf16.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Confidence Score: 4/5

The PR appears safe to merge, with a non-blocking gap in regression coverage for its precision fix.

Findings

1.  Precision fix lacks regression coverage ▶

Summary

The PR performs tensor-parallel Newton–Schulz all-reduces in FP32, then restores the Gram matrix to the iteration dtype, across all four step variants.

• The changed BF16 reduction behavior needs a numerical regression test.

_{Reviews (1) · Last reviewed commit: "all reduce in fp32 in tp newton schulz"}

[greptile-apps]

Precision fix lacks regression coverage

The new FP32 all-reduce path has no numerical test for BF16 tensor-parallel shards. The distributed step tests use FP32 inputs, where these casts are no-ops, while the medium-precision test checks only argument forwarding. A test with values sensitive to reduction rounding would help catch regressions in the behavior this change fixes.

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!

[skyw]

/ok to test 5cf479e

[github-actions]

Test Results

   97 files  ±0    191 suites  ±0   2m 8s ⏱️ -5s
1 464 tests ±0  1 464 ✅ +1  0 💤 ±0  0 ❌  - 1 
3 344 runs  ±0  3 344 ✅ +1  0 💤 ±0  0 ❌  - 1 

Results for commit 5cf479e. ± Comparison against base commit 0c5cdcd.

[codecov]

Codecov Report

❌ Patch coverage is 50.00000% with 4 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...optimizers/orthogonalized_optimizers/muon_utils.py	50.00%	4 Missing ⚠️

📢 Thoughts on this report? Let us know!