Skip narrow dtypes in the mixed-precision cast

deepspeed/runtime/engine.py

@@ -162,6 +162,10 @@
 
 MEMORY_OPT_ALLREDUCE_SIZE = 500000000
 
+# torch reports FP8/MX/NVFP4 as floating point, but casting them discards the quantized
+# encoding, and NVFP4 has no copy_ at all.
+CASTABLE_DTYPES = (torch.float16, torch.bfloat16, torch.float32, torch.float64)
+
 DeepSpeedOptimizerCallable = \
     Callable[[Union[Iterable[Parameter], Dict[str, Iterable]]], Optimizer]
 DeepSpeedSchedulerCallable = Callable[[Optimizer], _LRScheduler]
@@ -1611,13 +1615,13 @@ def _cast_module_mixed_precision(self, param_dtype, buffer_dtype, is_zero_init_m
         # the per-parameter cast applies only in the non-zero-init path.
         if param_dtype is not None and not is_zero_init_model:
             for p in self.module.parameters(recurse=True):
-                if p.is_floating_point() and p.dtype != param_dtype:
+                if p.dtype in CASTABLE_DTYPES and p.dtype != param_dtype:
                     p.data = p.data.to(param_dtype)
 
         # Buffers are never ZeRO-partitioned.
         if buffer_dtype is not None:
             for b in self.module.buffers(recurse=True):
-                if b.is_floating_point() and b.dtype != buffer_dtype:
+                if b.dtype in CASTABLE_DTYPES and b.dtype != buffer_dtype:
                     b.data = b.data.to(buffer_dtype)
 
     def _optimizer_has_ckpt_event_prologue(self):

tests/unit/v1/half_precision/test_mixed_precision_dtype.py

@@ -21,6 +21,26 @@ def _module_with_fp32_buffer(hidden_dim=8):
     return module
 
 
+NARROW_DTYPES = [
+    getattr(torch, name) for name in ("float8_e4m3fn", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz")
+    if hasattr(torch, name)
+]
+
+HELPER_ONLY_DTYPES = [getattr(torch, name) for name in ("float8_e8m0fnu", "float4_e2m1fn_x2") if hasattr(torch, name)]
+
+
+def _module_with_narrow_param(dtype, hidden_dim=8):
+    """Linear layer plus a submodule holding a frozen quantized weight and its scale."""
+    quantized = torch.nn.Module()
+    quantized.weight = torch.nn.Parameter(torch.zeros(hidden_dim, hidden_dim, dtype=dtype), requires_grad=False)
+    quantized.register_buffer("scale", torch.zeros(hidden_dim, dtype=dtype))
+
+    module = torch.nn.Module()
+    module.linear = torch.nn.Linear(hidden_dim, hidden_dim)
+    module.quantized = quantized
+    return module
+
+
 class TestMixedPrecisionDtypeResolution:
 
     def _engine(self, param_dtype=None, buffer_dtype=None, fp16=False, bf16=False):
@@ -94,6 +114,14 @@ def test_param_dtype_none_leaves_params(self):
         assert all(p.dtype == torch.float32 for p in module.parameters())
         assert module.inv_freq.dtype == torch.bfloat16
 
+    @pytest.mark.parametrize("dtype", HELPER_ONLY_DTYPES, ids=lambda d: str(d).rsplit(".", 1)[-1])
+    def test_narrow_dtypes_preserved(self, dtype):
+        module = _module_with_narrow_param(dtype)
+        DeepSpeedEngine._cast_module_mixed_precision(self._engine(module), torch.bfloat16, torch.bfloat16, False)
+        assert module.quantized.weight.dtype == dtype
+        assert module.quantized.scale.dtype == dtype
+        assert module.linear.weight.dtype == torch.bfloat16
+
 
 @pytest.mark.skipif(torch.bfloat16 not in get_accelerator().supported_dtypes(), reason="bf16 not supported")
 @pytest.mark.parametrize("zero_stage", [0, 3])
@@ -144,3 +172,13 @@ def test_buffer_dtype_opt_in(self, zero_stage):
                                                model=model,
                                                model_parameters=model.parameters())
         assert engine.module.inv_freq.dtype == torch.bfloat16
+
+    @pytest.mark.parametrize("dtype", NARROW_DTYPES, ids=lambda d: str(d).rsplit(".", 1)[-1])
+    def test_narrow_dtypes_preserved(self, zero_stage, dtype):
+        model = _module_with_narrow_param(dtype, 1024)
+        engine, _, _, _ = deepspeed.initialize(config=self._config(zero_stage),
+                                               model=model,
+                                               model_parameters=[p for p in model.parameters() if p.requires_grad])
+        assert engine.module.quantized.weight.dtype == dtype
+        assert engine.module.quantized.scale.dtype == dtype
+        assert engine.module.linear.weight.dtype == torch.bfloat16