From 750440e83eadceff3969267603d25c96f60f45a4 Mon Sep 17 00:00:00 2001 From: Abhishek Kulkarni Date: Fri, 4 Sep 2026 18:37:45 +0000 Subject: [PATCH 1/3] Skip narrow dtypes in the mixed-precision cast _cast_module_mixed_precision casts every parameter torch reports as floating point, which includes FP8, MX scales and NVFP4. Casting those discards the quantized encoding and doubles memory, and NVFP4 has no copy_ so it raises outright. Cast only the standard floating-point dtypes. Fixes #8414 Signed-off-by: Abhishek Kulkarni --- deepspeed/runtime/engine.py | 8 +++++-- .../test_mixed_precision_dtype.py | 24 +++++++++++++++++++ 2 files changed, 30 insertions(+), 2 deletions(-) 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..30fe9e419dad 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,21 @@ def _module_with_fp32_buffer(hidden_dim=8): return module +# FP8/MX/NVFP4 storage dtypes, whichever this torch exposes. +NARROW_DTYPES = [ + getattr(torch, name) for name in ("float8_e4m3fn", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz", + "float8_e8m0fnu", "float4_e2m1fn_x2") if hasattr(torch, name) +] + + +def _module_with_narrow_param(dtype, hidden_dim=8): + """Linear layer plus a frozen quantized param and its scale buffer.""" + module = torch.nn.Sequential(torch.nn.Linear(hidden_dim, hidden_dim)) + module.quantized = torch.nn.Parameter(torch.zeros(hidden_dim, hidden_dim, dtype=dtype), requires_grad=False) + module.register_buffer("scale", torch.zeros(hidden_dim, dtype=dtype)) + return module + + class TestMixedPrecisionDtypeResolution: def _engine(self, param_dtype=None, buffer_dtype=None, fp16=False, bf16=False): @@ -94,6 +109,15 @@ 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", NARROW_DTYPES, ids=lambda d: str(d).rsplit(".", 1)[-1]) + def test_narrow_dtypes_preserved(self, dtype): + # NVFP4 has no copy_, so casting it used to raise rather than silently degrade. + module = _module_with_narrow_param(dtype) + DeepSpeedEngine._cast_module_mixed_precision(self._engine(module), torch.bfloat16, torch.bfloat16, False) + assert module.quantized.dtype == dtype + assert module.scale.dtype == dtype + assert module[0].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]) From 50bc86355b5d8e57bc3411a2a90851ee8a4a48fa Mon Sep 17 00:00:00 2001 From: Abhishek Kulkarni Date: Fri, 4 Sep 2026 19:36:46 +0000 Subject: [PATCH 2/3] Cover narrow dtypes through deepspeed.initialize The private-helper test alone did not verify that initialization preserves frozen narrow parameters, and it tied the contract to that helper. Add the same assertions to the end-to-end class for the fp8 pairs, at ZeRO 0 and 3. e8m0 and float4 stay on the helper test: NCCL rejects e8m0 in the parameter broadcast and float4 has no fill_, both before any casting happens. Move the quantized tensors into a submodule so a standard-dtype parameter comes first, since ZeRO-3 reads the model dtype from list(module.parameters())[0]. Signed-off-by: Abhishek Kulkarni --- .../test_mixed_precision_dtype.py | 41 +++++++++++++++---- 1 file changed, 34 insertions(+), 7 deletions(-) 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 30fe9e419dad..ec00b18a6843 100644 --- a/tests/unit/v1/half_precision/test_mixed_precision_dtype.py +++ b/tests/unit/v1/half_precision/test_mixed_precision_dtype.py @@ -27,12 +27,28 @@ def _module_with_fp32_buffer(hidden_dim=8): "float8_e8m0fnu", "float4_e2m1fn_x2") if hasattr(torch, name) ] +# The subset that survives a full deepspeed.initialize. NCCL rejects e8m0 in the +# parameter broadcast and float4 has no fill_, both before any casting happens. +E2E_NARROW_DTYPES = [ + getattr(torch, name) for name in ("float8_e4m3fn", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz") + if hasattr(torch, name) +] + def _module_with_narrow_param(dtype, hidden_dim=8): - """Linear layer plus a frozen quantized param and its scale buffer.""" - module = torch.nn.Sequential(torch.nn.Linear(hidden_dim, hidden_dim)) - module.quantized = torch.nn.Parameter(torch.zeros(hidden_dim, hidden_dim, dtype=dtype), requires_grad=False) - module.register_buffer("scale", torch.zeros(hidden_dim, dtype=dtype)) + """Linear layer plus a frozen quantized weight and its scale buffer. + + The quantized tensors sit in a submodule so that a standard-dtype parameter + comes first, as in a real model: ZeRO-3 takes the model dtype from + ``list(module.parameters())[0]``. + """ + 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 @@ -114,9 +130,9 @@ def test_narrow_dtypes_preserved(self, dtype): # NVFP4 has no copy_, so casting it used to raise rather than silently degrade. module = _module_with_narrow_param(dtype) DeepSpeedEngine._cast_module_mixed_precision(self._engine(module), torch.bfloat16, torch.bfloat16, False) - assert module.quantized.dtype == dtype - assert module.scale.dtype == dtype - assert module[0].weight.dtype == torch.bfloat16 + 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") @@ -168,3 +184,14 @@ 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", E2E_NARROW_DTYPES, ids=lambda d: str(d).rsplit(".", 1)[-1]) + def test_narrow_dtypes_preserved(self, zero_stage, dtype): + """The same contract as above, through the public entry point.""" + 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 From a081850b88d27956e4d0252b133f7f7e4061537a Mon Sep 17 00:00:00 2001 From: Abhishek Kulkarni Date: Fri, 4 Sep 2026 20:08:40 +0000 Subject: [PATCH 3/3] Drop the helper-test cases the end-to-end test subsumes The fp8 pairs are now asserted through deepspeed.initialize, so covering them twice added nothing. The helper test keeps only e8m0 and float4, which cannot reach the cast through initialize. Signed-off-by: Abhishek Kulkarni --- .../test_mixed_precision_dtype.py | 23 ++++--------------- 1 file changed, 5 insertions(+), 18 deletions(-) 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 ec00b18a6843..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,27 +21,16 @@ def _module_with_fp32_buffer(hidden_dim=8): return module -# FP8/MX/NVFP4 storage dtypes, whichever this torch exposes. NARROW_DTYPES = [ - getattr(torch, name) for name in ("float8_e4m3fn", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz", - "float8_e8m0fnu", "float4_e2m1fn_x2") if hasattr(torch, name) -] - -# The subset that survives a full deepspeed.initialize. NCCL rejects e8m0 in the -# parameter broadcast and float4 has no fill_, both before any casting happens. -E2E_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 frozen quantized weight and its scale buffer. - The quantized tensors sit in a submodule so that a standard-dtype parameter - comes first, as in a real model: ZeRO-3 takes the model dtype from - ``list(module.parameters())[0]``. - """ +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)) @@ -125,9 +114,8 @@ 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", NARROW_DTYPES, ids=lambda d: str(d).rsplit(".", 1)[-1]) + @pytest.mark.parametrize("dtype", HELPER_ONLY_DTYPES, ids=lambda d: str(d).rsplit(".", 1)[-1]) def test_narrow_dtypes_preserved(self, dtype): - # NVFP4 has no copy_, so casting it used to raise rather than silently degrade. 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 @@ -185,9 +173,8 @@ def test_buffer_dtype_opt_in(self, zero_stage): model_parameters=model.parameters()) assert engine.module.inv_freq.dtype == torch.bfloat16 - @pytest.mark.parametrize("dtype", E2E_NARROW_DTYPES, ids=lambda d: str(d).rsplit(".", 1)[-1]) + @pytest.mark.parametrize("dtype", NARROW_DTYPES, ids=lambda d: str(d).rsplit(".", 1)[-1]) def test_narrow_dtypes_preserved(self, zero_stage, dtype): - """The same contract as above, through the public entry point.""" model = _module_with_narrow_param(dtype, 1024) engine, _, _, _ = deepspeed.initialize(config=self._config(zero_stage), model=model,