Skip to content

GitLab

Unverified Commit bf0d97d7 authored by Michael Goin's avatar Michael Goin Committed by GitHub
Browse files

Update requirements-tpu.txt to support python 3.9 and 3.11 (#11695) 


Signed-off-by: default avatarmgoin <michael@neuralmagic.com>
parent a655eb30
Hide whitespace changes
Inline Side-by-side
Showing with 7 additions and 3 deletions
requirements-tpu.txt
View file @ bf0d97d7
......@@ -18,6 +18,8 @@ ray[default]
--find-links https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
torch==2.6.0.dev20241126+cpu
torchvision==0.20.0.dev20241126+cpu
torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.6.0.dev20241126-cp310-cp310-linux_x86_64.whl
torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.6.0.dev20241126-cp39-cp39-linux_x86_64.whl ; python_version == "3.9"
torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.6.0.dev20241126-cp310-cp310-linux_x86_64.whl ; python_version == "3.10"
torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.6.0.dev20241126-cp311-cp311-linux_x86_64.whl ; python_version == "3.11"
jaxlib==0.4.36.dev20241122
jax==0.4.36.dev20241122
vllm/worker/tpu_model_runner.py
View file @ bf0d97d7
......@@ -126,8 +126,10 @@ class TPUModelRunner(ModelRunnerBase[ModelInputForTPU]):
logger.warning(
"The max_model_len (%d) is too large. This may degrade the "
"performance due to the insufficient smem size. Consider "
"setting --max-model-len to a smaller value.",
self.model_config.max_model_len)
"setting --max-model-len to a smaller value, like %d.",
self.model_config.max_model_len,
self.model_config.max_model_len /
(block_table_size / smem_size))
def load_model(self) -> None:
self.device = self.device_config.device
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment