fix: validate kd_loss_weight type and reject negative individual weights

modelopt/torch/distill/loss_balancers.py

@@ -92,8 +92,13 @@ def __init__(self, kd_loss_weight: float | list[float] = 0.5):
             ValueError if kd_loss_weight is out of bounds.
         """
         super().__init__()
-        if isinstance(kd_loss_weight, float):
-            kd_loss_weight = [kd_loss_weight]
+        if isinstance(kd_loss_weight, (int, float)):
+            kd_loss_weight = [float(kd_loss_weight)]
+
+        if any(w < 0.0 for w in kd_loss_weight):
+            raise ValueError(
+                f"Individual kd_loss_weight values must be non-negative, got {kd_loss_weight}"
+            )
 
         sum_kd_loss_weight = sum(kd_loss_weight)
         if sum_kd_loss_weight < 0.0 or sum_kd_loss_weight > 1.0:

tests/unit/torch/distill/test_loss_balancers.py

@@ -0,0 +1,18 @@
+"""Test distillation loss balancers."""
+# pyrefly: ignore [missing-import]
+import pytest
+from modelopt.torch.distill.loss_balancers import StaticLossBalancer
+
+@pytest.mark.parametrize(
+    ("weight", "expected"),
+    [(1, [1.0]), (0.5, [0.5])],
+)
+def test_static_loss_balancer_weight_validation(weight, expected):
+    """Test that StaticLossBalancer correctly validates scalar and negative weights."""
+    # 1. Verify scalar weights are accepted (and cast to list of float)
+    balancer = StaticLossBalancer(weight)
+    assert balancer._kd_loss_weight == expected
+
+    # 2. Verify negative individual weights are rejected even if sum is valid
+    with pytest.raises(ValueError, match="non-negative"):
+        StaticLossBalancer([0.5, -0.3])