From 31525856c59f75f3389e0f6ac4532c4491a36fa2 Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Wed, 22 Apr 2026 09:45:00 +0800 Subject: [PATCH] Fix training gradient underflow in quantization tests Change autocast dtype from float16 to bfloat16 in _test_quantization_training. Float16's limited dynamic range causes gradients to underflow to zero when passing through quantized tensor subclass operations. --- tests/models/testing_utils/quantization.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py index 4403cacc6966..7474632dc5bb 100644 --- a/tests/models/testing_utils/quantization.py +++ b/tests/models/testing_utils/quantization.py @@ -415,7 +415,8 @@ def _test_quantization_training(self, config_kwargs): # Step 3: run forward and backward pass inputs = self.get_dummy_inputs() - with torch.amp.autocast(torch_device, dtype=torch.float16): + # Use bfloat16 instead of float16 to avoid gradient underflow with quantized layers + with torch.amp.autocast(torch_device, dtype=torch.bfloat16): out = model(**inputs, return_dict=False)[0] out.norm().backward()