[Installation]: VLLM does not support TPU v5p-16 (Multi-Host) with Ray Cluster

[Bihan]

Your current environment

The output of `python collect_env.py`

Collecting environment information...
WARNING:root:libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU.
INFO 11-04 16:11:44 importing.py:15] Triton not installed or not compatible; certain GPU-related functions will not be available.
PyTorch version: 2.6.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Debian GNU/Linux 11 (bullseye) (x86_64)
GCC version: (Debian 10.2.1-6) 10.2.1 20210110
Clang version: Could not collect
CMake version: version 3.30.5
Libc version: glibc-2.31

Python version: 3.10.15 (main, Oct 17 2024, 02:58:23) [GCC 10.2.1 20210110] (64-bit runtime)
Python platform: Linux-5.19.0-1022-gcp-x86_64-with-glibc2.31
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Byte Order:                      Little Endian
Address sizes:                   52 bits physical, 57 bits virtual
CPU(s):                          208
On-line CPU(s) list:             0-207
Thread(s) per core:              2
Core(s) per socket:              52
Socket(s):                       2
NUMA node(s):                    2
Vendor ID:                       GenuineIntel
CPU family:                      6
Model:                           143
Model name:                      Intel(R) Xeon(R) Platinum 8481C CPU @ 2.70GHz
Stepping:                        8
CPU MHz:                         2699.998
BogoMIPS:                        5399.99
Hypervisor vendor:               KVM
Virtualization type:             full
L1d cache:                       4.9 MiB
L1i cache:                       3.3 MiB
L2 cache:                        208 MiB
L3 cache:                        210 MiB
NUMA node0 CPU(s):               0-51,104-155
NUMA node1 CPU(s):               52-103,156-207
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Not affected
Vulnerability Retbleed:          Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid rtm avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves avx_vnni avx512_bf16 arat avx512vbmi umip avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq la57 rdpid cldemote movdiri movdir64b fsrm md_clear serialize tsxldtrk amx_bf16 avx512_fp16 amx_tile amx_int8 arch_capabilities

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] pyzmq==26.2.0
[pip3] torch==2.6.0
[pip3] torch-xla==2.6.0+gita0f81e5
[pip3] torchvision==0.19.0a0+d23a6e1
[pip3] transformers==4.46.1
[conda] Could not collect
ROCM Version: Could not collect
Neuron SDK Version: N/A
vLLM Version: 0.6.3.post2.dev217+gccb5376a
vLLM Build Flags:
CUDA Archs: Not Set; ROCm: Disabled; Neuron: Disabled
GPU Topology:
Could not collect

How you are installing vllm

Create a TPU VM

gcloud compute tpus tpu-vm create tpu-v5p-benchmark \
  --zone=europe-west4-b \
  --accelerator-type=v5p-16 \
  --version=tpu-ubuntu2204-base

On head node

ray start --block --head --port=6379

On other node (Note: TPU v5p-16 has 2 nodes)

ray start --block --address=<head-node-address>:6379

#Below is the ray-status

Ray status shows both the nodes active

======== Autoscaler status: 2024-11-05 15:48:55.473751 ========
Node status
---------------------------------------------------------------
Active:
 1 node_49fc62d654acc1939448a1668ee1770feef20f763ab4bedface3ccf7
 1 node_88790deb8765d60a25e104192e33414fedcdc89d40e339316235547c
Pending:
 (no pending nodes)
Recent failures:
 (no failures)

Resources
---------------------------------------------------------------
Usage:
 0.0/416.0 CPU
 0B/839.91GiB memory
 0B/30.40GiB object_store_memory

Demands:
 (no resource demands)

However, when I run vllm serve from master node it issues error "The number of required TPUs exceeds the total number of available TPUs in the placement group.", even when it is connected to cluster.
I have tried with --tensor-parallel-size 2, 4, 8, 16 and the output is same.

$ vllm serve /root/.llama/checkpoints/Llama3.1-70B --tensor-parallel-size 8
WARNING:root:libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU.
INFO 11-05 15:51:16 importing.py:15] Triton not installed or not compatible; certain GPU-related functions will not be available.
INFO 11-05 15:51:18 api_server.py:551] vLLM API server version 0.6.3.post2.dev217+gccb5376a
INFO 11-05 15:51:18 api_server.py:552] args: Namespace(subparser='serve', model_tag='/root/.llama/checkpoints/Llama3.1-70B', config='', host=None, port=8000, uvicorn_log_level='info', allow_credentials=False, allowed_origins=['*'], allowed_methods=['*'], allowed_headers=['*'], api_key=None, lora_modules=None, prompt_adapters=None, chat_template=None, response_role='assistant', ssl_keyfile=None, ssl_certfile=None, ssl_ca_certs=None, ssl_cert_reqs=0, root_path=None, middleware=[], return_tokens_as_token_ids=False, disable_frontend_multiprocessing=False, enable_auto_tool_choice=False, tool_call_parser=None, tool_parser_plugin='', model='/root/.llama/checkpoints/Llama3.1-70B', task='auto', tokenizer=None, skip_tokenizer_init=False, revision=None, code_revision=None, tokenizer_revision=None, tokenizer_mode='auto', chat_template_text_format='string', trust_remote_code=False, download_dir=None, load_format='auto', config_format=<ConfigFormat.AUTO: 'auto'>, dtype='auto', kv_cache_dtype='auto', quantization_param_path=None, max_model_len=None, guided_decoding_backend='outlines', distributed_executor_backend=None, worker_use_ray=False, pipeline_parallel_size=1, tensor_parallel_size=8, max_parallel_loading_workers=None, ray_workers_use_nsight=False, block_size=16, enable_prefix_caching=False, disable_sliding_window=False, use_v2_block_manager=False, num_lookahead_slots=0, seed=0, swap_space=4, cpu_offload_gb=0, gpu_memory_utilization=0.9, num_gpu_blocks_override=None, max_num_batched_tokens=None, max_num_seqs=256, max_logprobs=20, disable_log_stats=False, quantization=None, rope_scaling=None, rope_theta=None, enforce_eager=False, max_seq_len_to_capture=8192, disable_custom_all_reduce=False, tokenizer_pool_size=0, tokenizer_pool_type='ray', tokenizer_pool_extra_config=None, limit_mm_per_prompt=None, mm_processor_kwargs=None, enable_lora=False, max_loras=1, max_lora_rank=16, lora_extra_vocab_size=256, lora_dtype='auto', long_lora_scaling_factors=None, max_cpu_loras=None, fully_sharded_loras=False, enable_prompt_adapter=False, max_prompt_adapters=1, max_prompt_adapter_token=0, device='auto', num_scheduler_steps=1, multi_step_stream_outputs=True, scheduler_delay_factor=0.0, enable_chunked_prefill=None, speculative_model=None, speculative_model_quantization=None, num_speculative_tokens=None, speculative_disable_mqa_scorer=False, speculative_draft_tensor_parallel_size=None, speculative_max_model_len=None, speculative_disable_by_batch_size=None, ngram_prompt_lookup_max=None, ngram_prompt_lookup_min=None, spec_decoding_acceptance_method='rejection_sampler', typical_acceptance_sampler_posterior_threshold=None, typical_acceptance_sampler_posterior_alpha=None, disable_logprobs_during_spec_decoding=None, model_loader_extra_config=None, ignore_patterns=[], preemption_mode=None, served_model_name=None, qlora_adapter_name_or_path=None, otlp_traces_endpoint=None, collect_detailed_traces=None, disable_async_output_proc=False, override_neuron_config=None, scheduling_policy='fcfs', pooling_type=None, pooling_norm=None, pooling_softmax=None, pooling_step_tag_id=None, pooling_returned_token_ids=None, disable_log_requests=False, max_log_len=None, disable_fastapi_docs=False, dispatch_function=<function serve at 0x7ff7515e1750>)
INFO 11-05 15:51:18 api_server.py:166] Multiprocessing frontend to use ipc:///tmp/5be6afd9-60d4-4329-8202-6ecbec5a18be for IPC Path.
INFO 11-05 15:51:18 api_server.py:181] Started engine process with PID 1132
INFO 11-05 15:51:18 config.py:1752] Downcasting torch.float32 to torch.float16.
INFO 11-05 15:51:21 importing.py:15] Triton not installed or not compatible; certain GPU-related functions will not be available.
INFO 11-05 15:51:22 config.py:1752] Downcasting torch.float32 to torch.float16.
INFO 11-05 15:51:23 config.py:323] This model supports multiple tasks: {'generate', 'embedding'}. Defaulting to 'generate'.
WARNING 11-05 15:51:23 arg_utils.py:1051] The model has a long context length (128000). This may cause OOM errors during the initial memory profiling phase, or result in low performance due to small KV cache space. Consider setting --max-model-len to a smaller value.
WARNING 11-05 15:51:23 arg_utils.py:1103] [DEPRECATED] Block manager v1 has been removed, and setting --use-v2-block-manager to True or False has no effect on vLLM behavior. Please remove --use-v2-block-manager in your engine argument. If your use case is not supported by SelfAttnBlockSpaceManager (i.e. block manager v2), please file an issue with detailed information.
You are using the default legacy behaviour of the <class 'transformers.models.llama.tokenization_llama.LlamaTokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565 - if you loaded a llama tokenizer from a GGUF file you can ignore this message
You are using the default legacy behaviour of the <class 'transformers.models.llama.tokenization_llama_fast.LlamaTokenizerFast'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565 - if you loaded a llama tokenizer from a GGUF file you can ignore this message.
INFO 11-05 15:51:27 config.py:323] This model supports multiple tasks: {'embedding', 'generate'}. Defaulting to 'generate'.
WARNING 11-05 15:51:27 arg_utils.py:1051] The model has a long context length (128000). This may cause OOM errors during the initial memory profiling phase, or result in low performance due to small KV cache space. Consider setting --max-model-len to a smaller value.
WARNING 11-05 15:51:27 arg_utils.py:1103] [DEPRECATED] Block manager v1 has been removed, and setting --use-v2-block-manager to True or False has no effect on vLLM behavior. Please remove --use-v2-block-manager in your engine argument. If your use case is not supported by SelfAttnBlockSpaceManager (i.e. block manager v2), please file an issue with detailed information.
2024-11-05 15:51:27,430 INFO worker.py:1631 -- Connecting to existing Ray cluster at address: 10.164.15.221:6379...
2024-11-05 15:51:27,438 INFO worker.py:1807 -- Connected to Ray cluster. View the dashboard at 127.0.0.1:8265 
Process SpawnProcess-1:
Traceback (most recent call last):
  File "/usr/local/lib/python3.10/multiprocessing/process.py", line 314, in _bootstrap
    self.run()
  File "/usr/local/lib/python3.10/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/workspace/vllm/vllm/engine/multiprocessing/engine.py", line 361, in run_mp_engine
    engine = MQLLMEngine.from_engine_args(engine_args=engine_args,
  File "/workspace/vllm/vllm/engine/multiprocessing/engine.py", line 122, in from_engine_args
    executor_class = LLMEngine._get_executor_cls(engine_config)
  File "/workspace/vllm/vllm/engine/llm_engine.py", line 521, in _get_executor_cls
    initialize_ray_cluster(engine_config.parallel_config)
  File "/workspace/vllm/vllm/executor/ray_utils.py", line 277, in initialize_ray_cluster
    raise ValueError(
ValueError: The number of required TPUs exceeds the total number of available TPUs in the placement group.
Traceback (most recent call last):
  File "/usr/local/bin/vllm", line 33, in <module>
    sys.exit(load_entry_point('vllm', 'console_scripts', 'vllm')())
  File "/workspace/vllm/vllm/scripts.py", line 195, in main
    args.dispatch_function(args)
  File "/workspace/vllm/vllm/scripts.py", line 41, in serve
    uvloop.run(run_server(args))
  File "/usr/local/lib/python3.10/site-packages/uvloop/__init__.py", line 82, in run
    return loop.run_until_complete(wrapper())
  File "uvloop/loop.pyx", line 1518, in uvloop.loop.Loop.run_until_complete
  File "/usr/local/lib/python3.10/site-packages/uvloop/__init__.py", line 61, in wrapper
    return await main
  File "/workspace/vllm/vllm/entrypoints/openai/api_server.py", line 575, in run_server
    async with build_async_engine_client(args) as engine_client:
  File "/usr/local/lib/python3.10/contextlib.py", line 199, in __aenter__
    return await anext(self.gen)
  File "/workspace/vllm/vllm/entrypoints/openai/api_server.py", line 107, in build_async_engine_client
    async with build_async_engine_client_from_engine_args(
  File "/usr/local/lib/python3.10/contextlib.py", line 199, in __aenter__
    return await anext(self.gen)
  File "/workspace/vllm/vllm/entrypoints/openai/api_server.py", line 197, in build_async_engine_client_from_engine_args
    raise RuntimeError(
RuntimeError: Engine process failed to start

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[CortexEdgeUser]

Ray doesn't detect TPU.. I have the same issue

[totorochina]

I see you are using tpu-ubuntu2204-base, I encountered a similar problem recently, and I solved it after changing the image to v2-alpha-tpuv5 according to the documentation.
https://cloud.google.com/tpu/docs/runtimes#pytorch_and_jax
You can try python3 -c 'import jax; print("Total TPU chips:", jax.device_count())' to check if that is the problem.
And I remembered tpu-ubuntu2204-base used to work, but it just does not work recently.

[ruisearch42]

This will probably be fixed by #11257

[youkaichao]

@ruisearch42 why #11257 can fix it?

[youkaichao]

is it because the change from ray.get_gpu_ids() to ray.get_runtime_context().get_accelerator_ids() ?

[ruisearch42]

@youkaichao actually #11257 won't fix it.

@Bihan I think for some reason TPU resource was not detected by Ray. Here is how the detection works:
https://github.com/ray-project/ray/blob/916f534e571278b26733812b24a7b3dee08f24e4/python/ray/_private/resource_spec.py#L199
https://github.com/ray-project/ray/blob/916f534e571278b26733812b24a7b3dee08f24e4/python/ray/_private/accelerators/tpu.py#L97

You can add some debug code into your ray installation, and see what these would print when you run ray start. It may be a configuration issue or easy fix in ray.

[Bihan]

@youkaichao actually #11257 won't fix it.

@Bihan I think for some reason TPU resource was not detected by Ray. Here is how the detection works: https://github.com/ray-project/ray/blob/916f534e571278b26733812b24a7b3dee08f24e4/python/ray/_private/resource_spec.py#L199 https://github.com/ray-project/ray/blob/916f534e571278b26733812b24a7b3dee08f24e4/python/ray/_private/accelerators/tpu.py#L97

You can add some debug code into your ray installation, and see what these would print when you run ray start. It may be a configuration issue or easy fix in ray.

@ruisearch42 Thank you. Also should I use v2-alpha-tpuv5 instead of tpu-ubuntu2204-base?

[ruisearch42]

@Bihan you can try that, since it was reported to be working.

[richardsliu]

Hi, I believe #10155 (comment) should fix this issue. Can this be closed?

[ruisearch42]

@richardsliu looks like your link is not correct?

[richardsliu]

Not sure why. I meant this fix:

use v2-alpha-tpuv5 instead of tpu-ubuntu2204-base

[ruisearch42]

@richardsliu Sounds good, thanks. I'm closing the issue. @Bihan feel free to reopen if it is not fixed.