[BUG] DeepSpeed generation fails with a CPU/CUDA device mismatch when keep-module-on-host is enabled

[Redempt1onzzZZ]

Describe the bug
DeepSpeed accepts keep-module-on-host, completes model initialization, and passes health checks. The first real generation request fails. The same failure family was reproduced on Qwen3-4B-Instruct-2507, DeepSeek-R1-Distill-Llama-8B, and Meta-Llama-3.1-8B-Instruct.

To Reproduce
Steps to reproduce the behavior:

1. Save the following as a shell command and run it in a shell on a machine with a free CUDA device.
2. Run the command in the target Python environment with the listed dependencies installed.
3. Observe the startup or generation result.

import torch
import deepspeed
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "/root/.cache/huggingface/hub/models--Qwen--Qwen3-4B-Instruct-2507/snapshots/cdbee75f17c01a7cc42f958dc650907174af0554"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name, torch_dtype=torch.bfloat16
).to("cuda")
engine = deepspeed.init_inference(
    model, config={"keep_module_on_host": True}
)
device = next(engine.module.parameters()).device
inputs = tokenizer("Reply with exactly OK.", return_tensors="pt")
inputs = {key: value.to(device) for key, value in inputs.items()}
with torch.no_grad():
    output = engine.generate(**inputs, max_new_tokens=8, do_sample=False)
print(tokenizer.decode(output[0], skip_special_tokens=True))

Actual behavior

Startup and health checks succeed, but generation fails with a device-placement error: Expected all tensors to be on the same device, but index is on cuda:0 while other tensors are on cpu during index_select.

Expected behavior

The option should either preserve compatible device placement for all tensors or be rejected before serving with an actionable validation error.

System info (please complete the following information):

• OS: Ubuntu 22.04.3 LTS
• GPU count and types: 1 x NVIDIA H100 80GB HBM3
• DeepSpeed version: 0.19.1
• DeepSpeed-MII: not used
• PyTorch: 2.9.1+cu128
• Transformers: 5.15.0
• CUDA reported by PyTorch: 12.8
• nvcc: 13.3
• Python version: 3.12.12
• Target runtime: /root/anaconda3/envs/deepspeed128

Docker context

No Docker image was used. The reproduction ran directly on Ubuntu 22.04.3 LTS in the target conda environment.