From d8fb74feb5bec69aeb6b5598d482c383cc3e9ec8 Mon Sep 17 00:00:00 2001 From: Robert joseph Date: Thu, 6 Aug 2026 03:37:02 +0000 Subject: [PATCH] Fix clip backward with broadcasted constant bounds Signed-off-by: Robert joseph --- modelopt/torch/quantization/nn/functional.py | 4 +- .../torch/quantization/test_functional.py | 38 +++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) create mode 100644 tests/unit/torch/quantization/test_functional.py diff --git a/modelopt/torch/quantization/nn/functional.py b/modelopt/torch/quantization/nn/functional.py index 533e7ebf68d..692687031b3 100644 --- a/modelopt/torch/quantization/nn/functional.py +++ b/modelopt/torch/quantization/nn/functional.py @@ -49,7 +49,9 @@ def backward(ctx, grad_output): if clip_value_min.requires_grad or clip_value_max.requires_grad: warnings.warn("Learning enabled for clip min/max. This is an experimental feature.") - if clip_value_min.numel() != 1 or clip_value_max.numel() != 1: + if (clip_value_min.requires_grad and clip_value_min.numel() != 1) or ( + clip_value_max.requires_grad and clip_value_max.numel() != 1 + ): raise ValueError( f"Learnable min/max can only be scalar, got size {clip_value_min.size()} and {clip_value_max.size()}." ) diff --git a/tests/unit/torch/quantization/test_functional.py b/tests/unit/torch/quantization/test_functional.py new file mode 100644 index 00000000000..458b42eb8be --- /dev/null +++ b/tests/unit/torch/quantization/test_functional.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU unit tests for quantization functions.""" + +import torch + +from modelopt.torch.quantization.nn.functional import clip + + +def test_clip_backward_with_broadcast_constant_bounds(): + inputs = torch.tensor([[-2.0, 0.5, 4.0], [-0.5, 2.0, 2.0]], requires_grad=True) + clip_value_min = torch.tensor([-1.0, 0.0, 1.0]) + clip_value_max = torch.tensor([0.0, 1.0, 3.0]) + + outputs = clip(inputs, clip_value_min, clip_value_max) + outputs.sum().backward() + + reference_inputs = inputs.detach().clone().requires_grad_() + reference_outputs = torch.maximum( + torch.minimum(reference_inputs, clip_value_max), clip_value_min + ) + reference_outputs.sum().backward() + + torch.testing.assert_close(outputs, reference_outputs) + torch.testing.assert_close(inputs.grad, reference_inputs.grad)