diff --git a/modelopt/torch/distill/loss_balancers.py b/modelopt/torch/distill/loss_balancers.py index 0ceaca3e3c2..e7640209f85 100644 --- a/modelopt/torch/distill/loss_balancers.py +++ b/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: diff --git a/tests/unit/torch/distill/test_loss_balancers.py b/tests/unit/torch/distill/test_loss_balancers.py new file mode 100644 index 00000000000..94274494ad0 --- /dev/null +++ b/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])