Set torch_dtype in TransformersModel (#13088)

hmellor · web-flow · commit ad9776353e6b · 2025-02-11T23:51:19.000+08:00
Signed-off-by: Harry Mellor &lt;19981378+hmellor@users.noreply.github.com&gt;
diff --git a/vllm/model_executor/models/transformers.py b/vllm/model_executor/models/transformers.py
@@ -143,6 +143,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         self.model: PreTrainedModel = AutoModel.from_config(
             self.config,
             attn_implementation="vllm",
+            torch_dtype=vllm_config.model_config.dtype,
             trust_remote_code=vllm_config.model_config.trust_remote_code,
         )
         prefix = self.model.base_model_prefix

Original file line number	Diff line number	Diff line change
`@@ -143,6 +143,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:`
`143`	`143`	`self.model: PreTrainedModel = AutoModel.from_config(`
`144`	`144`	`self.config,`
`145`	`145`	`attn_implementation="vllm",`
	`146`	`+ torch_dtype=vllm_config.model_config.dtype,`
`146`	`147`	`trust_remote_code=vllm_config.model_config.trust_remote_code,`
`147`	`148`	`)`
`148`	`149`	`prefix = self.model.base_model_prefix`