Fix training gradient underflow in quantization tests (#13539)

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

* fix autocast dtype check

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>

---------

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

Author: jiqing-feng

tests/models/testing_utils/quantization.py

@@ -407,7 +407,9 @@ 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 on XPU to avoid gradient underflow with quantized layers
+        autocast_dtype = torch.bfloat16 if torch_device == "xpu" else torch.float16
+        with torch.amp.autocast(torch_device, dtype=autocast_dtype):
             out = model(**inputs, return_dict=False)[0]
             out.norm().backward()