diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 86918bd71c5a..cb4c3df07fc3 100644 --- a/deepspeed/runtime/engine.py +++ b/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): diff --git a/tests/unit/v1/half_precision/test_mixed_precision_dtype.py b/tests/unit/v1/half_precision/test_mixed_precision_dtype.py index 2ca883d35e9d..b13e24af695f 100644 --- a/tests/unit/v1/half_precision/test_mixed_precision_dtype.py +++ b/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