[TPU] optimize the all-reduce performance

[yaochengji]

Before the PR, the all-reduce performance is not optimal due to two reasons:

• on v6e-8, XLA compiler accidentally apply 2D-ring strategy, while 1D-ring is expected
• The ring-order cannot be automatically adjusted

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[yaochengji]

I have two questions:

• Currently I wrap the all-reduce into a pytorch python custom op to make it compatible with dynamo, should I wrap it in vllm/distributed/parallel_state.py instead?
• I need to another environment variable LIBTPU_INIT_ARGS="--xla_tpu_force_1d_allreduce_at_chunk_count=1" for the performance optimization, should I add it in the code?

[alexm-redhat]

@yaochengji thanks for this important PR!

1. About dynamo, I don't have a strong opinion there.
2. About libtpu flag, I think you can detect inside init_device() of tpu_worker that it is a v6 and simply add the env var. Similar to the code below (that adds PJRT_DEVICE flag):

def init_device(self):
        os.environ["PJRT_DEVICE"] = "TPU"

[yaochengji]

Thanks @alexm-redhat for your suggestion.

Hi @youkaichao , do you have any suggestion on the dynamo custom op?

[yaochengji]

My local experiments show that the throughput can improve from ~4.2 reqs/s to ~4.9reqs/s for Llama 70B on 8 v6e.

[youkaichao]

is it possible to use custom op call in vllm directly? e.g. extending this line

vllm/vllm/distributed/parallel_state.py

Line 222 in 44f9905

self.use_custom_op_call = current_platform.is_cuda_alike()

to include tpu.

i didn't do it initially, because i remember tpu has some custom dynamo-related logic.

[yaochengji]

Thanks for you suggestion!

I removed the custom op in tpu_communicator.py and made use of the custom op for TPU in parallel_state.py

[alexm-redhat]

LGTM!

[robertgshaw2-redhat]

Magical incantations!

[robertgshaw2-redhat]

NOTE: V1 test failing fixed by #15969