Keep the tensors a model names for fp32 in fp32 in init_inference

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()

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