[MISC] Support multi node inference with Neuron #8692

sssrijan-amazon · 2024-09-21T10:38:01Z

FILL IN THE PR DESCRIPTION HERE

This change adds support for multi node tensor parallelism inference on neuron device using AWS Neuron Transformers Neuronx library. Refer multi-node inference guide here

# Neuron specific env required for multi node inference support 
NEURON_RT_ROOT_COMM_ID=The master node’s <IP address>:<port>
NEURON_RANK_ID=Rank of the node, 0 means master node
NEURON_LOCAL_TP=The local tensor parallel degree on each node. Eg: 32 for for trn1 instance type.

# Vllm specific env required for inter node communication for broadcasting data.
VLLM_HOST_IP=The master node’s <IP address>
VLLM_PORT=The master node’s port

Example to run meta llama 3.1 70B model on 2 node trn1n cluster with TP = 64

Pre-req to setup the cluster - https://awsdocs-neuron.readthedocs-hosted.com/en/latest/frameworks/torch/torch-neuronx/setup-trn1-multi-node-execution.html

offline_multi_node_inference_neuron.py

import os

from vllm import LLM, SamplingParams

# creates XLA hlo graphs for all the context length buckets.
os.environ['NEURON_CONTEXT_LENGTH_BUCKETS'] = "128"
# creates XLA hlo graphs for all the token gen buckets.
os.environ['NEURON_TOKEN_GEN_BUCKETS'] = "128"

# Sample prompts.
prompts = [
     "The capital of France is",
]
# Create a sampling params object.
sampling_params = SamplingParams(max_tokens=100, top_k=1)
# Create an LLM.
llm = LLM(
    model="meta-llama-3-1/Meta-Llama-3-1-70B",
    max_num_seqs=2,
    max_model_len=128,
    block_size=128,
    device="neuron",
    tensor_parallel_size=64,
    world_size=2 # Number of nodes in the cluster
)

outputs = llm.generate(prompts, sampling_params)

# Print the outputs.
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

Node 1 command line:

NEURON_RT_ROOT_COMM_ID=10.1.201.64:63423 NEURON_RANK_ID=0 NEURON_LOCAL_TP=32 VLLM_HOST_IP=10.1.201.64 VLLM_PORT=8989 python3 offline_multi_node_inference_neuron.py

Node 2 command line:

NEURON_RT_ROOT_COMM_ID=10.1.201.64:63423 NEURON_RANK_ID=1 NEURON_LOCAL_TP=32 VLLM_HOST_IP=10.1.201.64 VLLM_PORT=8989 python3 offline_multi_node_inference_neuron.py

BEFORE SUBMITTING, PLEASE READ THE CHECKLIST BELOW AND FILL IN THE DESCRIPTION ABOVE

PR Checklist (Click to Expand)

Thank you for your contribution to vLLM! Before submitting the pull request, please ensure the PR meets the following criteria. This helps vLLM maintain the code quality and improve the efficiency of the review process.

PR Title and Classification

Only specific types of PRs will be reviewed. The PR title is prefixed appropriately to indicate the type of change. Please use one of the following:

[Bugfix] for bug fixes.
[CI/Build] for build or continuous integration improvements.
[Doc] for documentation fixes and improvements.
[Model] for adding a new model or improving an existing model. Model name should appear in the title.
[Frontend] For changes on the vLLM frontend (e.g., OpenAI API server, LLM class, etc.)
[Kernel] for changes affecting CUDA kernels or other compute kernels.
[Core] for changes in the core vLLM logic (e.g., LLMEngine, AsyncLLMEngine, Scheduler, etc.)
[Hardware][Vendor] for hardware-specific changes. Vendor name should appear in the prefix (e.g., [Hardware][AMD]).
[Misc] for PRs that do not fit the above categories. Please use this sparingly.

Note: If the PR spans more than one category, please include all relevant prefixes.

Code Quality

The PR need to meet the following code quality standards:

We adhere to Google Python style guide and Google C++ style guide.
Pass all linter checks. Please use format.sh to format your code.
The code need to be well-documented to ensure future contributors can easily understand the code.
Include sufficient tests to ensure the project to stay correct and robust. This includes both unit tests and integration tests.
Please add documentation to docs/source/ if the PR modifies the user-facing behaviors of vLLM. It helps vLLM user understand and utilize the new features or changes.

Adding or changing kernels

Each custom kernel needs a schema and one or more implementations to be registered with PyTorch.

Make sure custom ops are registered following PyTorch guidelines: Custom C++ and CUDA Operators and The Custom Operators Manual
Custom operations that return Tensors require meta-functions. Meta-functions should be implemented and registered in python so that dynamic dims can be handled automatically. See above documents for a description of meta-functions.
Use torch.libary.opcheck() to test the function registration and meta-function for any registered ops. See tests/kernels for examples.
When changing the C++ signature of an existing op, the schema must be updated to reflect the changes.
If a new custom type is needed, see the following document: Custom Class Support in PT2.

Notes for Large Changes

Please keep the changes as concise as possible. For major architectural changes (>500 LOC excluding kernel/data/config/test), we would expect a GitHub issue (RFC) discussing the technical design and justification. Otherwise, we will tag it with rfc-required and might not go through the PR.

What to Expect for the Reviews

The goal of the vLLM team is to be a transparent reviewing machine. We would like to make the review process transparent and efficient and make sure no contributor feel confused or frustrated. However, the vLLM team is small, so we need to prioritize some PRs over others. Here is what you can expect from the review process:

After the PR is submitted, the PR will be assigned to a reviewer. Every reviewer will pick up the PRs based on their expertise and availability.
After the PR is assigned, the reviewer will provide status update every 2-3 days. If the PR is not reviewed within 7 days, please feel free to ping the reviewer or the vLLM team.
After the review, the reviewer will put an action-required label on the PR if there are changes required. The contributor should address the comments and ping the reviewer to re-review the PR.
Please respond to all comments within a reasonable time frame. If a comment isn't clear or you disagree with a suggestion, feel free to ask for clarification or discuss the suggestion.

Thank You

Finally, thank you for taking the time to read these guidelines and for your interest in contributing to vLLM. Your contributions make vLLM a great tool for everyone!

github-actions · 2024-09-21T10:38:14Z

👋 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 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 do one of these:

Add ready label to the PR
Enable auto-merge.

🚀

liangfu

Thanks @sssrijan-amazon for contributing to multi-node support. Please try to reuse existing world size prediction logic, instead of introducing a new world_size argument explicitly.

liangfu · 2024-09-21T17:03:14Z

examples/offline_inference_neuron.py

I'm confuse with world_size=1, while TP size=2. Isn't world_size predictable from number of nodes times local TP size?

Compared to single node tensor parallel, multi-node tensor parallel shards the model weights in the same way but having mores cores across nodes. In the meantime, it requires each node’s model.forward() receives the exact same input, otherwise there would be unexpected behaviors (runtime failure, wrong output).

For Neuron multi node inference, we need to broadcast the API input params across all the nodes. So the world size for Neuron within the scope of Vllm is the number of nodes on which we perform inference instead of TP*PP as model parallelism is handled within the context of neuron (outside the scope of vllm)

we can still reuse existing vLLM config argument to calculate world size, right? I'm trying to understand what blocker did you face in reusing the existing config arguments.

If we use the existing flow of using world_size = TP*PP then it would come out to be 64 for a cluster of 2 trn1/trn1n instances. This causes the init_process_group to hang as we only initialize it with 2 ranks (0 and 1 from each node of the cluster) instead of all the ranks.

This was currently working for single node inference with TP=32 because we had hard coded the world size to 1 while initializing the process groups here instead of using world_size = TP*PP which would be 32.

we can send fake input to init_process_group so that we won't get hang there, and send real inputs to where it actually matters.

liangfu · 2024-09-24T19:45:39Z

examples/offline_inference_neuron.py

i think the consistency of user interface matters. it's hard to understand what to expect when we set default as TP=2, PP=1, world_size=1

Done, refactored to keep interface consistent.

liangfu · 2024-09-25T01:37:52Z

examples/neuron_multi_node/run_llama_3_70b_2_neuron_nodes.sh

i guess we don't need these local setup in the upstream.

export NEURON_RT_ROOT_COMM_ID export NEURON_RANK_ID export NEURON_LOCAL_TP export VLLM_HOST_IP export VLLM_PORT

These are required for multi node inference using transformers-neuronx library. I can get rid of the RT install part.

omrishiv · 2024-09-25T16:21:43Z

vllm/entrypoints/neuron_multi_node/api_server.py

Does it make sense to refactor this + the base api_server.py so we have something we can deduplicate this and extend in the future?

omrishiv · 2024-09-25T16:44:29Z

vllm/entrypoints/neuron_multi_node/api_server.py

let's rename to run_driver

omrishiv

Can we add some documentation on the neuron side on how to run multi node inference?

omrishiv · 2024-09-25T16:45:59Z

vllm/executor/neuron_executor.py

Let's change to

rank=int(os.getenv("NEURON_RANK_ID", 0)), local_rank=rank, # Make sure we explain this

Local rank is not utilized with neuron multi instance flow. Using local rank as 0 given we have 1 process per node for neuron.

sssrijan-amazon · 2024-09-25T22:38:01Z

Can we add some documentation on the neuron side on how to run multi node inference?

Added instructions for neuron cluster setup and details in the run script.

omrishiv

Thanks @sssrijan-amazon for working through this! LGTM from a code perspective, but I'll defer to @liangfu for neuron compatibility.

@simon-mo we've talked internally about making the api_server more extensible. The change has helped dedupe code and should make it more reusable in the future. Hopefully that's ok?

liangfu

Thanks for the updates.

There seems to be some unnecessary complications here in the current design. My impression is that we should be able to follow existing practise in https://docs.vllm.ai/en/latest/serving/distributed_serving.html , and simply run multi-node with

vllm serve /path/to/the/model/in/the/container \
--tensor-parallel-size 128

liangfu · 2024-09-30T05:14:31Z

examples/neuron_multi_node/run_llama_3_70b_2_neuron_nodes.sh

it would be better if we detect this with ray support for neuron, in order to convert this into a scalable design.

liangfu · 2024-09-30T05:18:45Z

examples/neuron_multi_node/run_llama_3_70b_2_neuron_nodes.sh

i think there is room to simply this design, so that we don't need to introduce the dependency to either slurm or kubernetes.

liangfu · 2024-09-30T05:22:58Z

examples/neuron_multi_node/run_llama_3_70b_2_neuron_nodes.sh

i would prefer to write these in the document, if this is the required setup for every time we run the follow python script.
if this is required purely for setup purpose, it will be better if we mention in the document and explain the reason.

liangfu · 2024-09-30T05:24:23Z

examples/neuron_multi_node/run_llama_3_70b_2_neuron_nodes.sh

i think we can read these into the python script, to avoid the complication here. or is that even feasible ?

liangfu · 2024-09-30T05:27:19Z

docs/source/getting_started/neuron-installation.rst

can we follow the practise in https://docs.vllm.ai/en/latest/serving/distributed_serving.html , and simply run multi-node with

vllm serve /path/to/the/model/in/the/container \ --tensor-parallel-size 128

simon-mo · 2024-09-30T18:25:26Z

Sorry dumb question, why do we need a separate api server?

liangfu suggested changes Sep 21, 2024

View reviewed changes

liangfu reviewed Sep 24, 2024

View reviewed changes

liangfu reviewed Sep 25, 2024

View reviewed changes

omrishiv reviewed Sep 25, 2024

View reviewed changes

omrishiv approved these changes Sep 26, 2024

View reviewed changes

liangfu suggested changes Sep 30, 2024

View reviewed changes

WoosukKwon added the aws-neuron Related to AWS Inferentia & Trainium label Oct 17, 2024

sssrijan-amazon closed this Oct 31, 2024

sssrijan-amazon force-pushed the main branch from 2c9f1ce to 3ea2dc2 Compare October 31, 2024 07:33

liangfu mentioned this pull request Apr 30, 2025

Add NeuronxDistributedInference support, Speculative Decoding, Dynamic on-device sampling #16357

Merged

Uh oh!

[MISC] Support multi node inference with Neuron #8692

[MISC] Support multi node inference with Neuron #8692

Uh oh!

Conversation

sssrijan-amazon commented Sep 21, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Title and Classification

Code Quality

Adding or changing kernels

Notes for Large Changes

What to Expect for the Reviews

Thank You

Uh oh!

github-actions bot commented Sep 21, 2024

Uh oh!

liangfu left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

liangfu Sep 21, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

liangfu Sep 23, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

liangfu Sep 24, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

omrishiv left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sssrijan-amazon commented Sep 25, 2024

Uh oh!

omrishiv left a comment

Choose a reason for hiding this comment

Uh oh!

liangfu left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

liangfu Sep 30, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

liangfu Sep 30, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sssrijan-amazon commented Sep 21, 2024 •

edited

Loading

liangfu left a comment •

edited

Loading

liangfu Sep 21, 2024 •

edited

Loading

liangfu Sep 23, 2024 •

edited

Loading

liangfu Sep 24, 2024 •

edited

Loading

liangfu left a comment •

edited

Loading

liangfu Sep 30, 2024 •

edited

Loading

liangfu Sep 30, 2024 •

edited

Loading