[TPU] Support multi-host inference

[WoosukKwon]

Changing global rank and world size into local rank and local world size.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which consists a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of default ones by unblocking the steps in your fast-check build on Buildkite UI.

Once the PR is approved and ready to go, please make sure to run full CI as it is required to merge (or just use auto-merge).

To run full CI, you can do one of these:

• Comment /ready on the PR
• Add ready label to the PR
• Enable auto-merge.

🚀

[youkaichao]

In ray gpu executor, there are these lines:

vllm/vllm/executor/ray_gpu_executor.py

Lines 175 to 191 in 7025b11

	def sort_by_driver_then_worker_ip(worker):
	"""
	Sort the workers based on 3 properties:
	1. If the worker is on the same node as the driver (vllm engine),
	it should be placed first.
	2. Then, if the worker is on a node with fewer workers, it should
	be placed first.
	3. Finally, if the work is on a node with smaller IP address, it
	should be placed first.
	"""
	ip = ray.get(worker.get_node_ip.remote())
	return (ip != driver_ip, ip_counts[ip], ip)
	# After sorting, the workers on the same node will be
	# close to each other, and the workers on the driver
	# node will be placed first.
	self.workers = sorted(self.workers, key=sort_by_driver_then_worker_ip)

to make sure the worker index aligns with machine boundary. you might need it in TPU, too. Otherwise local ranks can be wrong. for example, rank 0, 1, 2, 4 in one node, and 3, 5, 6, 7 in another node.

[WoosukKwon]

@youkaichao Can you explain more?

[youkaichao]

say you have 2 nodes, 8 TPUs in total.

ray actors are launched one by one. when you launch the first actor, it might live in node 0; when you launch the second actor, it might live in node 1. If you use the index of worker as global rank, then it will cause a problem.

[WoosukKwon]

@youkaichao Thanks for the explanation. Let me merge this PR first as there are users waiting for this and the scope of this PR is isolated to the TPU backend. I will address your comment in a followup PR.

[jaisong123]

Hi @WoosukKwon - thanks for the office hours today!

And thanks for your hard work on this! I’m eager to start using the multi-host inference support on TPUs. Do you know when this feature will be available for general use?

Thanks again!

[sparsh35]

And thanks for your hard work on this! I’m eager to start using the multi-host inference support on TPUs. Do you know when this feature will be available for general use?

I think it should work now , though I haven't tried it yet especially for bigger model like mistral large v2 , any doc or tutorial would help, @WoosukKwon