Skip to content

Accept torch.int8 in init_inference's dtype validation gate #8578

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
udsy19 wants to merge 3 commits into
base: master
Choose a base branch
from
Open

Accept torch.int8 in init_inference's dtype validation gate #8578

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e57bded
Accept torch.int8 in init_inference's dtype validation gate
udsy19 Sep 17, 2026
ce500b2
Require an injection path for torch.int8 in init_inference
udsy19 Sep 20, 2026
6a116ac
Recognize tensor_parallel.mpu in the int8 dtype gate
udsy19 Oct 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion deepspeed/inference/engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ def __init__(self, model, config):
if hasattr(self.module, "config"):
TransformerPolicy.hf_model_config = self.module.config

if config.dtype not in get_accelerator().supported_dtypes():
if config.dtype not in get_accelerator().supported_dtypes() and config.dtype != torch.int8:

@ebarkhordar ebarkhordar Sep 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 ebarkhordar Sep 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

raise ValueError(
f"Data type {config.dtype} is not supported by {get_accelerator().device_name()} accelerator")

Expand Down
10 changes: 10 additions & 0 deletions tests/unit/inference/test_inference_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,3 +92,13 @@ 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(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 though DeepSpeedSelfAttention/DeepSpeedMLP still special-case
# torch.int8 weights and the inference tutorial documents int8 as a supported dtype.
engine = deepspeed.init_inference(torch.nn.Linear(4, 4), dtype=torch.int8)
assert engine._config.dtype == torch.int8
Loading