biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[hardware] unify usage of is_tpu to current_platform.is_tpu() (#7102)

Browse Source
This commit is contained in:
youkaichao
2024-08-13 00:16:42 -07:00
committed by GitHub
parent 7025b11d94
commit 4d2dc5072b
8 changed files with 29 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
vllm/model_executor/model_loader/loader.py
View File

@@ -41,7 +41,7 @@ from vllm.model_executor.models.interfaces import (has_inner_state,
supports_vision)
from vllm.model_executor.utils import set_weight_attrs
from vllm.platforms import current_platform
from vllm.utils import is_pin_memory_available, is_tpu
from vllm.utils import is_pin_memory_available
@contextmanager
@@ -94,7 +94,7 @@ def _get_quantization_config(
"""Get the quantization config."""
if model_config.quantization is not None:
quant_config = get_quant_config(model_config, load_config)
if not is_tpu():
if not current_platform.is_tpu():
capability = current_platform.get_device_capability()
capability = capability[0] * 10 + capability[1]
if capability < quant_config.get_min_capability():
@@ -320,7 +320,7 @@ class DefaultModelLoader(BaseModelLoader):
else:
weights_iterator = pt_weights_iterator(hf_weights_files)
if is_tpu():
if current_platform.is_tpu():
# In PyTorch XLA, we should call `xm.mark_step` frequently so that
# not too many ops are accumulated in the XLA program.
import torch_xla.core.xla_model as xm