From 5b4ba2e1e13f7a1df1316052bbfcf76641f7d4f3 Mon Sep 17 00:00:00 2001 From: Chengji Yao Date: Tue, 23 Sep 2025 22:18:08 -0700 Subject: [PATCH] [TPU][Bugfix] fix the missing apply_model in tpu worker (#25526) Signed-off-by: Chengji Yao Signed-off-by: yewentao256 --- tests/v1/tpu/test_tpu_int8.py | 6 +----- vllm/v1/worker/tpu_worker.py | 8 +++++++- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/tests/v1/tpu/test_tpu_int8.py b/tests/v1/tpu/test_tpu_int8.py index 991070dc9239d..f39a8021a29ef 100644 --- a/tests/v1/tpu/test_tpu_int8.py +++ b/tests/v1/tpu/test_tpu_int8.py @@ -48,13 +48,9 @@ def test_model_tpu_int8(vllm_runner, model: str, dtype: str, max_tokens: int, prompts = [ "A robot may not injure a human being", - "It is only with the heart that one can see rightly;", - "The greatest glory in living lies not in never falling,", ] answers = [ - "or, being injured, not kill, except in", - "without the heart, one can only see wrongly.", - "but in rising every time we fall. - Nelson" + "or kill a human being", ] with vllm_runner(model, dtype=dtype, hf_overrides=hf_overrides) as vllm: diff --git a/vllm/v1/worker/tpu_worker.py b/vllm/v1/worker/tpu_worker.py index fc72b954df9cf..d4f0a65f2a164 100644 --- a/vllm/v1/worker/tpu_worker.py +++ b/vllm/v1/worker/tpu_worker.py @@ -3,7 +3,7 @@ """A TPU worker class.""" import os -from typing import Any, Optional +from typing import Any, Callable, Optional, TypeVar import torch import torch.distributed @@ -31,6 +31,8 @@ from vllm.v1.worker.utils import bind_kv_cache logger = init_logger(__name__) +_R = TypeVar("_R") + if not USE_TPU_COMMONS: logger.info("tpu_commons not found, using vLLM's TPUWorker.") import torch_xla.core.xla_model as xm @@ -333,6 +335,10 @@ class TPUWorker: def shutdown(self) -> None: self.model_runner.ensure_kv_transfer_shutdown() + def apply_model(self, fn: Callable[[nn.Module], _R]) -> _R: + """Apply a function on the model inside this worker.""" + return fn(self.get_model()) + if USE_TPU_COMMONS: from tpu_commons.worker import TPUWorker as TPUCommonsWorker