Accept torch.int8 in init_inference's dtype validation gate

deepspeed/inference/engine.py

@@ -72,8 +72,26 @@ def __init__(self, model, config):
             TransformerPolicy.hf_model_config = self.module.config
 
         if config.dtype not in get_accelerator().supported_dtypes():
-            raise ValueError(
-                f"Data type {config.dtype} is not supported by {get_accelerator().device_name()} accelerator")
+            if config.dtype != torch.int8:
+                raise ValueError(
+                    f"Data type {config.dtype} is not supported by {get_accelerator().device_name()} accelerator")
+            # torch.int8 is only ever consumed by replace_transformer_layer's quantizer
+            # construction, reached solely through kernel injection, a user injection policy, or
+            # automatic tensor parallelism. Without one of those, _convert_to_dtype has no int8
+            # branch (see its `if False:` guard), so weights would silently stay in their
+            # original dtype while config.dtype claimed int8.
+            #
+            # AutoTP is reached two ways: an explicit tensor_parallel.tp_size > 1, or a
+            # tensor_parallel.mpu. An mpu caller does not set tp_size here; it is derived from
+            # dist.get_world_size(self.mpu.get_model_parallel_group()) further down in __init__,
+            # after this gate. So the mpu case must be recognized directly, or int8 + mpu is
+            # wrongly refused even though the AutoTP branch would run and quantize.
+            if not (config.injection_policy or config.replace_with_kernel_inject
+                    or config.tensor_parallel.tp_size > 1 or config.tensor_parallel.mpu):
+                raise ValueError("Data type torch.int8 requires kernel injection or a replacement policy "
+                                 "(replace_with_kernel_inject=True, injection_policy=..., "
+                                 "tensor_parallel.tp_size > 1, or tensor_parallel.mpu) to actually "
+                                 "quantize weights.")
 
         # todo: keep this self.injection_dict because we don't use to change config.injection_policy API
         # todo: this will get changed when Molly's PR on auto injection dict is merged

tests/unit/inference/test_inference_config.py

@@ -92,3 +92,54 @@ def test_moe_backward_compat_bool(self):
             config = DeepSpeedInferenceConfig(moe=value)
             assert isinstance(config.moe, DeepSpeedMoEConfig)
             assert config.moe.enabled == value
+
+    def test_int8_dtype_accepted_with_injection_policy(self):
+        # Regression test: the dtype gate added in #6528 ("add bfloat16 to inference support
+        # dtypes") replaced a narrow fp16-only check with `get_accelerator().supported_dtypes()`,
+        # whose per-accelerator lists (cpu/cuda/hpu/mlu/xpu) have never included torch.int8. That
+        # unintentionally made every `init_inference(..., dtype=torch.int8)` call raise ValueError
+        # on every accelerator, even when an injection_policy is given and
+        # replace_module.py's replace_transformer_layer would actually consume config.dtype ==
+        # torch.int8 to build a quantizer.
+        class _WithLinear(torch.nn.Module):
+
+            def __init__(self):
+                super().__init__()
+                self.lin = torch.nn.Linear(4, 4)
+
+        engine = deepspeed.init_inference(_WithLinear(),
+                                          dtype=torch.int8,
+                                          injection_policy={torch.nn.Linear: ("lin", )})
+        assert engine._config.dtype == torch.int8
+
+    def test_int8_dtype_without_consumer_raises(self):
+        # torch.int8 is only consumed by replace_transformer_layer's quantizer construction,
+        # reached solely via kernel injection / an injection_policy / AutoTP
+        # (tensor_parallel.tp_size > 1 or tensor_parallel.mpu). Plain dtype conversion
+        # (InferenceEngine._convert_to_dtype) never acts on int8, so without one of those,
+        # accepting the dtype would silently leave weights in their original dtype while
+        # `_config.dtype` claimed int8. Assert this now fails loudly instead.
+        with pytest.raises(ValueError, match="requires kernel injection or a replacement policy"):
+            deepspeed.init_inference(torch.nn.Linear(4, 4), dtype=torch.int8)
+
+    def test_int8_dtype_accepted_with_mpu(self):
+        # tensor_parallel.mpu is a fourth route into AutoTP. An mpu caller does not set
+        # tensor_parallel.tp_size at gate time; tp_size is derived from
+        # dist.get_world_size(mpu.get_model_parallel_group()) later in InferenceEngine.__init__,
+        # after the dtype gate. So the gate must key off tensor_parallel.mpu directly, or it
+        # wrongly refuses int8 for a configuration whose AutoTP branch would run and quantize.
+        # This pins that the gate no longer raises its int8 refusal for the mpu path (the gate
+        # is what regressed; a single-rank test cannot exercise the real tp_size > 1 quantize).
+        class _StubMpu:
+
+            def get_model_parallel_group(self):
+                return None
+
+            def get_model_parallel_rank(self):
+                return 0
+
+        try:
+            deepspeed.init_inference(torch.nn.Linear(4, 4), dtype=torch.int8, tensor_parallel={"mpu": _StubMpu()})
+        except Exception as e:  # may fail further in on this single-rank stub; only the gate is under test
+            assert "requires kernel injection or a replacement policy" not in str(e), \
+                "int8 dtype gate wrongly refused the tensor_parallel.mpu AutoTP path"