Accept torch.int8 in init_inference's dtype validation gate

[udsy19]

Fixes #8577

Related to #1454 (the original compute_attention()/input_mask TypeError reported there is
already fixed upstream — deepspeed/ops/transformer/inference/ds_attention.py has required
input_mask on compute_attention() for years — so this does not use a Fixes #1454 keyword;
it closes a different, still-live gap in the same int8/quantized-inference workflow the issue is
about, filed separately as #8577).

Problem

InferenceEngine.__init__ (deepspeed/inference/engine.py:74) rejects any dtype not in
get_accelerator().supported_dtypes(). That per-accelerator list — checked across every
implementation that defines it (cpu, cuda, hpu, mlu, xpu) — has never included torch.int8. So
deepspeed.init_inference(model, dtype=torch.int8) raises ValueError: Data type torch.int8 is not supported by <device> accelerator unconditionally, on every accelerator DeepSpeed ships,
regardless of what hardware is actually running.

Impact: crash-on-valid-input. Who reaches it / triggered by: anyone following
docs/_tutorials/inference-tutorial.md's own "Datatypes and Quantized Models" section (last
touched yesterday by #8535, which still asserts "DeepSpeed inference supports fp32, fp16 and int8
parameters" and "INT8 inference quantization remains supported"), or the dtype=torch.int8
fixture already present in tests/unit/inference/test_checkpoint_sharding.py:54-58. What is
observed: a hard ValueError at engine construction instead of a working (or even attempted)
int8 inference engine.

This is a regression, not a deliberate removal of int8 support:

• Before add bfloat16 to inference support dtypes #6528 ("add bfloat16 to inference support dtypes", 2024-09-27), this same check only
  rejected an unsupported torch.half; it never gated torch.int8 at all. add bfloat16 to inference support dtypes #6528's own stated
  intent was narrowly "to allow running inference tasks using bfloat16" and its diff never
  mentions int8.
• A separate, since-deleted _validate_args() explicitly listed
  supported_dtypes = [None, torch.half, torch.int8, torch.float, torch.bfloat16] right up until
  it was removed as unused dead code by inference: remove unused _validate_args function #5505 (2025-01-07) — independent evidence int8 was still
  considered valid after add bfloat16 to inference support dtypes #6528 landed, just checked by a function nobody called.
• DeepSpeedSelfAttention/DeepSpeedMLP (deepspeed/ops/transformer/inference/ds_attention.py)
  still special-case config.dtype == torch.int8 for parameter dtype construction today — dead
  code for a dtype the engine can no longer be constructed with.
• tests/unit/inference/test_checkpoint_sharding.py's dtype fixture already defensively
  pytest.skips when dtype not in get_accelerator().supported_dtypes(), which, given the
  enumeration above, means its int8 parametrization has been silently skipped on every
  accelerator since add bfloat16 to inference support dtypes #6528, quietly losing that coverage.

Fix

Exempt torch.int8 from the get_accelerator().supported_dtypes() gate, restoring the pre-#6528
scope of this one check. This does not touch the accelerator capability lists themselves (they
describe genuine hardware tensor-dtype support and are used elsewhere too) — it only stops that
list from also gating DeepSpeed's own int8 weight-quantization dtype, which is not what it was
written to describe.

Testing

Added test_int8_dtype_accepted to tests/unit/inference/test_inference_config.py::TestInferenceConfig.

Negative control (revert only deepspeed/inference/engine.py to upstream/master, keep the new
test):

FAILED ...test_int8_dtype_accepted - ValueError: Data type torch.int8 is not supported by cpu accelerator
1 failed

Restored: 1 passed. Full test_inference_config.py: unchanged pre-existing results before/after
this diff (10 passed, 1 pre-existing failure — TestInferenceCudaGraphConfig::test_cuda_graph_with_kernel_inject_raises
fails identically on unmodified upstream/master because this CPU lacks torch.float16 support,
unrelated to torch.int8/this change).

yapf/flake8/codespell clean on both changed files.

Signed-off-by: Udaya Tejas udayatejas2004@gmail.com

[ebarkhordar]

I ran this against e57bdeda67e6430658c2f8cd542331544cd0f6cd in a clean python:3.11-slim container, CPU accelerator, torch 2.14.0+cpu.

The gate change does let int8 through, but the path it opens does not quantize. _convert_to_dtype (engine.py:517) still keeps the int8 branch behind if False:, and no remaining branch matches int8, so the module keeps the dtype it arrived with:

dtype=torch.int8      config.dtype=torch.int8       weight: torch.float32 -> torch.float32
dtype=torch.bfloat16  config.dtype=torch.bfloat16   weight: torch.float32 -> torch.bfloat16

bfloat16 is a control, so the weight read is not just always printing float32. On base fe8c4b10 the same script raises ValueError: Data type torch.int8 is not supported by cpu accelerator. So with no checkpoint and replace_with_kernel_inject at its default False, this swaps a loud error for an engine reporting _config.dtype == torch.int8 over fp32 weights.

Both places that do act on int8 are unreached there: engine.py:480 inside _load_checkpoint, and replace_module.py:201 inside replace_transformer_layer, which __init__ calls only for one of its three injection modes (lines 135-169).

One correction to the body: 279bf743d does delete the if False: block, but it is not an ancestor of master, so it never landed. git blame still puts that line at b5d18a6ab.

test_int8_dtype_accepted asserts _config.dtype only, which passes whether or not anything quantized. Asserting the weight dtype the engine produces would pin the real behaviour, and if that is fp32 then narrowing the exemption to the cases that consume int8 seems better than opening it for all.

CPU only, no checkpoint, no kernel injection.

[ebarkhordar]

Thanks for picking this up. I ran ce500b2 in a clean container (python:3.11-slim, torch 2.14.0+cpu, pip install . from the PR head, reported version 0.19.8+ce500b21e). The narrowed gate refuses one configuration that does reach the quantizer.

tensor_parallel.mpu is a fourth way into AutoTP. The gate at engine.py:84 reads config.tensor_parallel.tp_size, but an mpu caller does not set that field: it is written at engine.py:125 from dist.get_world_size(group=self.mpu.get_model_parallel_group()), after the gate has already raised. So init_inference(model, dtype=torch.int8, tensor_parallel={"mpu": mpu}) is rejected even when that group would have given tp_size > 1 and mode 3 would have run.

Measured on 2 ranks (gloo, world_size 2, an mpu stub returning the default group). Only dtype differs between the two runs:

world_size = 2
[bf16 + mpu] AssertionError: Not able to determine model policy automatically.
             Please provide policy.          <- auto_tp.py:260, inside AutoTP.tp_parser
[int8 + mpu] ValueError: Data type torch.int8 requires kernel injection or a
             replacement policy ...

The bf16 run gets past the gate and into the AutoTP branch. The int8 run never gets there. That assertion is my toy model having no AutoTP policy, which is what makes it a usable marker that the branch ran at all.

Your new test docstring already names the case: "AutoTP (tensor_parallel.tp_size > 1 or tensor_parallel.mpu)". The guard checks only the first half. One more term covers it:

if not (config.injection_policy or config.replace_with_kernel_inject
        or config.tensor_parallel.tp_size > 1 or config.tensor_parallel.mpu):

I did not run a real model-parallel mpu, so what tp_size becomes for a genuine TP group is read from engine.py:125 rather than executed.