A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[OOT] Support sync_model_loading for OOT (#25126)
Signed-off-by: Chendi Xue <Chendi.Xue@intel.com>
Chendi.Xue committed
Sep 19, 2025 at 00:41 UTC
a6149aa587d6582545b7878a2dffed3a2419605d
4 files changed
+33
-17
vllm/model_executor/parameter.py
+3
-3
@@ -12,7 +12,6 @@ from torch.nn import Parameter
12
from vllm.distributed import (get_tensor_model_parallel_rank,
13
get_tensor_model_parallel_world_size)
14
from vllm.logger import init_logger
15
-from vllm.model_executor.utils import _make_synced_weight_loader
15
16
__all__ = [
17
"BasevLLMParameter", "PackedvLLMParameter", "PerTensorScaleParameter",
@@ -53,8 +52,9 @@ class BasevLLMParameter(Parameter):
52
# This sometimes causes OOM errors during model loading. To avoid this,
53
# we sync the param tensor after its weight loader is called.
54
from vllm.platforms import current_platform
56
- if current_platform.is_tpu():
57
- weight_loader = _make_synced_weight_loader(weight_loader)
55
+ if current_platform.use_sync_weight_loader():
56
+ weight_loader = current_platform.make_synced_weight_loader(
57
+ weight_loader)
58
59
self._weight_loader = weight_loader
60
self.tp_rank = get_tensor_model_parallel_rank()
vllm/model_executor/utils.py
+3
-14
@@ -44,23 +44,12 @@ def set_weight_attrs(
44
# TODO(woosuk): Remove this hack once we have a better solution.
45
from vllm.platforms import current_platform
46
47
- if current_platform.is_tpu() and key == "weight_loader":
48
- value = _make_synced_weight_loader(value)
47
+ if current_platform.use_sync_weight_loader(
48
+ ) and key == "weight_loader":
49
+ value = current_platform.make_synced_weight_loader(value)
50
setattr(weight, key, value)
51
52
52
-def _make_synced_weight_loader(original_weight_loader):
53
-
54
- def _synced_weight_loader(param, *args, **kwargs):
55
- out = original_weight_loader(param, *args, **kwargs)
56
- # torch._sync doesn't support, is not needed for CPU tensors.
57
- if param.device != torch.device("cpu"):
58
- torch._sync(param)
59
- return out
60
-
61
- return _synced_weight_loader
62
-
63
-
53
def get_packed_modules_mapping(model: torch.nn.Module) -> dict[str, list[str]]:
54
parent_map = getattr(model, "packed_modules_mapping", None)
55
parent_map = copy.deepcopy(parent_map) if parent_map is not None else {}
vllm/platforms/interface.py
+23
@@ -594,6 +594,29 @@ class Platform:
594
"""
595
return False
596
597
+ @classmethod
598
+ def use_sync_weight_loader(cls) -> bool:
599
+ """
600
+ Returns if the current platform needs to sync weight loader.
601
+ """
602
+ return False
603
+
604
+ @classmethod
605
+ def make_synced_weight_loader(cls, original_weight_loader):
606
+ """
607
+ Wrap the original weight loader to make it synced.
608
+ """
609
+ if not cls.use_sync_weight_loader():
610
+ return original_weight_loader
611
+
612
+ def _synced_weight_loader(param, *args, **kwargs):
613
+ out = original_weight_loader(param, *args, **kwargs)
614
+ if param.device != torch.device("cpu"):
615
+ torch._sync(param)
616
+ return out
617
+
618
+ return _synced_weight_loader
619
+
620
621
class UnspecifiedPlatform(Platform):
622
_enum = PlatformEnum.UNSPECIFIED
vllm/platforms/tpu.py
+4
@@ -226,6 +226,10 @@ class TpuPlatform(Platform):
226
torch.ops.xla.dynamo_set_buffer_donor_(src_cache, True)
227
dst_cache[dst_block_indices] = src_cache[src_block_indices].cpu()
228
229
+ @classmethod
230
+ def use_sync_weight_loader(cls) -> bool:
231
+ return True
232
+
233
234
try:
235
from tpu_commons.platforms import TpuPlatform as TpuCommonsPlatform