From 003707e00f7614d9129cbbedf512ddcc3fcd40b0 Mon Sep 17 00:00:00 2001 From: Udaya Tejas Date: Tue, 20 Oct 2026 05:22:05 +0000 Subject: [PATCH] Keep the tensors a model names for fp32 in fp32 in init_inference transformers loads the tensors a model lists in `_keep_in_fp32_modules` (under fp16) and `_keep_in_fp32_modules_strict` (under fp16 and bf16) in fp32, because they overflow or lose precision in half precision: T5's `wo` projection (huggingface/transformers#20683, added for FlanT5-XL), the MoE routing bias `e_score_correction_bias` of DeepSeek-V3/GLM-4.5 and others. `InferenceEngine._convert_to_dtype` then calls `self.module.half()` / `.bfloat16()` on the whole model, which casts those tensors along with everything else. For FlanT5 in fp16 that is the divergence from Hugging Face reported in #3177: `wo` overflows and generation goes wrong with no error. Record the fp32 tensors the model's lists name (reusing `keep_in_fp32_pattern` from the ZeRO-3 fix in #8699), cast the module, and restore them, so `init_inference` leaves them as transformers loaded them. A listed tensor that arrives in a lower precision is upcast, as transformers' loader and #8699 do. Measured on CPU with google/flan-t5-base in fp16: before this change `init_inference(model, dtype=torch.float16)` generates different text than the same model in Hugging Face; after it the outputs are identical. Fixes #3177 Signed-off-by: Udaya Tejas --- deepspeed/inference/engine.py | 31 ++++++++ .../inference/test_inference_keep_in_fp32.py | 78 +++++++++++++++++++ 2 files changed, 109 insertions(+) create mode 100644 tests/unit/inference/test_inference_keep_in_fp32.py diff --git a/deepspeed/inference/engine.py b/deepspeed/inference/engine.py index b7c38f63747c..33285e0cea52 100755 --- a/deepspeed/inference/engine.py +++ b/deepspeed/inference/engine.py @@ -16,6 +16,7 @@ from deepspeed.runtime.checkpoint_engine.torch_checkpoint_engine import TorchCheckpointEngine from deepspeed.utils.timer import SynchronizedWallClockTimer from deepspeed.runtime.compiler import is_compile_supported +from deepspeed.runtime.keep_in_fp32 import KEEP_IN_FP32_AUTO, keep_in_fp32_pattern from ..runtime.state_dict_factory import SDLoaderFactory from ..runtime.weight_quantizer import WeightQuantization from ..module_inject import replace_transformer_layer, generic_injection @@ -519,12 +520,42 @@ def _convert_to_dtype(self, config): model, self.quantization_scales = quantizer.model_quantize(self.module, self.injection_dict, self.quantize_bits, self.quantize_groups) elif config.dtype == torch.half: + keep = self._tensors_to_keep_in_fp32(config.dtype) self.module.half() + self._restore_fp32_tensors(keep) elif config.dtype == torch.bfloat16: + keep = self._tensors_to_keep_in_fp32(config.dtype) self.module.bfloat16() + self._restore_fp32_tensors(keep) elif config.dtype == torch.float: self.module.float() + def _tensors_to_keep_in_fp32(self, dtype): + """The fp32 copies of the tensors the model names for fp32, keyed by name. + + transformers keeps the tensors a model lists in ``_keep_in_fp32_modules`` (under fp16) and + ``_keep_in_fp32_modules_strict`` (under fp16 and bf16) in fp32 when it loads the model, because + they overflow or lose precision in half precision (T5's ``wo`` projection, MoE routing biases). + Casting the whole module to ``dtype`` would undo that, so they are restored after the cast. + """ + pattern = keep_in_fp32_pattern(self.module, KEEP_IN_FP32_AUTO, dtype) + if pattern is None: + return {} + keep = {} + for name, tensor in list(self.module.named_parameters()) + list(self.module.named_buffers()): + if tensor.is_floating_point() and pattern.search(name): + keep[name] = tensor.data if tensor.dtype == torch.float32 else tensor.data.float() + if keep: + log_dist(f"Keeping {len(keep)} tensors the model names for fp32 in fp32 under {dtype}", [0]) + return keep + + def _restore_fp32_tensors(self, keep): + if not keep: + return + for name, tensor in list(self.module.named_parameters()) + list(self.module.named_buffers()): + if name in keep: + tensor.data = keep[name] + def _create_cuda_graph(self, *inputs, **kwargs): # warmup to create the workspace and cublas handle cuda_stream = get_accelerator().Stream() diff --git a/tests/unit/inference/test_inference_keep_in_fp32.py b/tests/unit/inference/test_inference_keep_in_fp32.py new file mode 100644 index 000000000000..e7e4a1c6742a --- /dev/null +++ b/tests/unit/inference/test_inference_keep_in_fp32.py @@ -0,0 +1,78 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""init_inference(dtype=fp16/bf16) keeps the tensors a model names for fp32 in fp32. + +transformers loads the tensors a model lists in `_keep_in_fp32_modules` (under fp16) and +`_keep_in_fp32_modules_strict` (under fp16 and bf16) in fp32, because they overflow or lose precision in +half precision. Regression test for https://github.com/deepspeedai/DeepSpeed/issues/3177, where +init_inference cast T5's `wo` projection to fp16 along with the rest of the model. +""" + +import pytest +import torch + +import deepspeed +from deepspeed.accelerator import get_accelerator + +HIDDEN = 8 +# Distinct in fp32, but fp16 and bf16 round all of them to 8.0. +EXACT_VALUES = 8.0 + torch.arange(HIDDEN, dtype=torch.float32) * 1e-4 + + +class ToyModel(torch.nn.Module): + _keep_in_fp32_modules = ["wo"] # fp16 only, as T5 lists its `wo` projection + _keep_in_fp32_modules_strict = ["router_bias"] # fp16 and bf16, as the MoE models list their routing bias + + def __init__(self): + super().__init__() + self.wi = torch.nn.Linear(HIDDEN, HIDDEN, bias=False) + self.wo = torch.nn.Linear(HIDDEN, HIDDEN, bias=False) + self.register_buffer("router_bias", EXACT_VALUES.clone()) + self.register_buffer("scale", torch.ones(HIDDEN)) + + +def _dtypes(model): + return {name: t.dtype for name, t in list(model.named_parameters()) + list(model.named_buffers())} + + +@pytest.mark.inference +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_init_inference_keeps_listed_tensors_in_fp32(dtype): + if dtype not in get_accelerator().supported_dtypes(): + pytest.skip(f"{dtype} not supported by {get_accelerator().device_name()}") + model = ToyModel().to(get_accelerator().device_name()) + + engine = deepspeed.init_inference(model, dtype=dtype) + + dtypes = _dtypes(engine.module) + assert dtypes["wi.weight"] == dtype + assert dtypes["scale"] == dtype + assert dtypes["router_bias"] == torch.float32 + # Kept as loaded, not rounded through half precision and cast back. + assert torch.equal(engine.module.router_bias.cpu(), EXACT_VALUES) + # `_keep_in_fp32_modules` applies under fp16 only, the way transformers applies it. + assert dtypes["wo.weight"] == (torch.float32 if dtype == torch.float16 else dtype) + + +@pytest.mark.inference +@pytest.mark.skipif(torch.float16 not in get_accelerator().supported_dtypes(), reason="fp16 not supported") +def test_init_inference_fp16_keeps_t5_wo_in_fp32(): + from transformers import T5Config, T5ForConditionalGeneration + + config = T5Config(vocab_size=64, + d_model=32, + d_kv=8, + d_ff=64, + num_layers=1, + num_decoder_layers=1, + num_heads=2, + decoder_start_token_id=0) + model = T5ForConditionalGeneration(config).to(get_accelerator().device_name()) + + engine = deepspeed.init_inference(model, dtype=torch.float16) + + for name, param in engine.module.named_parameters(): + expected = torch.float32 if name.endswith("wo.weight") else torch.float16 + assert param.dtype == expected, name