Use math.fsum for StaticLossBalancer weight checks

[aniruddh-krovvidi]

What does this PR do?

Type of change: Bug fix

StaticLossBalancer.__init__ checks sum(kd_loss_weight) against 1.0. On Python 3.10 and 3.11 (both in the CI matrix) plain sum() is not compensated, so weights that sum to 1 in decimal can come out at 0.9999999999999999 and trigger the spurious "weights do not sum to 1.0" warning (or, past 1.0, the ValueError). Python 3.12+ compensates inside sum(), which is why the check looks fine there. math.fsum gives the correctly rounded sum on every supported version. The same sum() in forward() derived the residual student_loss weight, so normalized weights left a 1e-16 coefficient on it; it now uses the same fsum.

Fixes #2488

Usage

mtd.StaticLossBalancer([0.7, 0.2, 0.1])  # no warning on 3.10/3.11 either

Testing

tests/unit/torch/distill/test_distill.py::test_static_loss_balancer_weights_summing_to_one_do_not_warn: fails on Python 3.11 before the change, passes after. Whole test_distill.py module run on 3.11 CPU.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: N/A
• Did you get Claude approval on this PR?: N/A

Additional Information

The issue's snippet passes a dict, which sum() cannot take at all; the warning it describes is the list form on a pre-3.12 interpreter.

Summary by CodeRabbit

• Bug Fixes

  • Improved numerical accuracy when calculating knowledge-distillation loss weights.
  • Prevented an incorrect warning when valid loss weights sum to 1 due to floating-point rounding.

• Tests

  • Added regression coverage for valid loss weights affected by floating-point precision.

[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.

[coderabbitai]

Understand this PR’s impact

Explore downstream dependencies and potential security impact with Blast Radius.

View blast radius →

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository: NVIDIA/Model-Optimizer/.coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 7654c7eb-6d00-445e-8478-f6074fa5ac3e

📥 Commits

Reviewing files that changed from the base of the PR and between fc4c40f and 209cbde.

📒 Files selected for processing (2)

• modelopt/torch/distill/loss_balancers.py
• tests/unit/torch/distill/test_distill.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

📝 Walkthrough

Walkthrough

StaticLossBalancer now uses math.fsum for weight calculations. A regression test verifies that weights [0.7, 0.2, 0.1] do not trigger a warning.

Changes

Loss balancer precision

Layer / File(s)	Summary
Summation precision and regression coverage modelopt/torch/distill/loss_balancers.py, tests/unit/torch/distill/test_distill.py	The balancer uses math.fsum for KD-loss totals and the student-loss remainder. A warning-as-error test covers weights that sum to 1.0 despite floating-point representation differences.

Priority: ➖ Normal

Estimated code review effort: 1 (Trivial) | ~5 minutes

Change: Bug fix · Severity of issue fixed: Medium

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 40.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 5 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the primary change: using math.fsum for StaticLossBalancer weight checks.
Linked Issues check	✅ Passed	For issue #2488, StaticLossBalancer now uses math.fsum for weight validation and for the residual student_loss coefficient. This prevents the documented 0.9999999999999999 result from causing …
Out of Scope Changes check	✅ Passed	The pull request changes only weight summation in modelopt/torch/distill/loss_balancers.py and adds a focused regression test in tests/unit/torch/distill/test_distill.py. Both changes directly add…
Security Anti-Patterns	✅ Passed	PASS. The authoritative PR range changes only modelopt/torch/distill/loss_balancers.py and a unit test. The modelopt change adds math and replaces numeric sum() calls with math.fsum; it does n…

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

---

_{Comment @coderabbitai help to get the list of available commands.}

[AAnoosheh]

Thanks, but this is likely overkill and not an issue in practice.