Support inference with transformers-neuronx #2569

liangfu · 2024-01-24T00:13:49Z

This PR enables llama model inference on Inferentia with transformers-neuronx backend.

To demonstrate offline inference with transformers-neuronx, run

python3 examples/offline_inference_neuron.py

WoosukKwon · 2024-01-29T19:22:21Z

Just for a note: We may want to merge #2503 before this PR.

liangfu · 2024-01-30T05:49:45Z

Just for a note: We may want to merge #2503 before this PR.

Sure, I can work on a rebase once that's merged.

WoosukKwon · 2024-01-30T05:53:54Z

Hey @liangfu, thanks for understanding. We are prioritizing some PRs for the upcoming v0.3.0 release and the meetup event. Apologies for the delay.

liangfu · 2024-01-30T06:11:22Z

No worries. Removing hardcoded device is necessary.

WoosukKwon · 2024-02-02T00:01:58Z

Hi @liangfu, we just merged #2503. Could you rebase the PR? Thanks!

liangfu · 2024-02-03T00:05:07Z

Hi @WoosukKwon , i rebased upon latest main branch. Feel free to take a look.

WoosukKwon

@liangfu, thanks for submitting the PR! I've checked that the PR works on inf2 instance and the outputs of offline_inference_neuron match those of offline_inference.py on NVIDIA GPUs (when greedy sampling is used). Happy to see that we have a working implementation!

Also, I liked that the PR minimized the modifications to the codebase, except that it duplicated worker in neuron_worker for some reason I don't understand. Could you clarify this?

Also, I'd like to know more about the details of the implementations made in transformers-neurons. For now, things like max_model_len=128, block_size=128, and os.environ['MASTER_PORT'] = '12355' are quite mysterious to me.

vllm/model_executor/__init__.py

WoosukKwon · 2024-02-06T22:03:26Z

examples/offline_inference_neuron.py

+          max_num_seqs=8,
+          max_model_len=128,
+          block_size=128,
+          device="cpu")


Why does the device need to be "cpu" here?

What are the constraints on block_size and max_model_len and max_num_seqs? If there's any, please specify it.

removed the device='cpu' here, and the device can be automatically detected.

Also, specified the constraints on supporting neuron device explicitly

examples/offline_inference_neuron.py

vllm/model_executor/neuron_model_loader.py

vllm/worker/neuron_worker.py

WoosukKwon · 2024-02-06T22:19:15Z

examples/offline_inference_neuron.py

It seems the online server APIs of vLLM such as API server and OpenAI-compatible servers do not work at the moment. Could you support them as well?

It's a bit tricky for now. I extended the --block-size argument with 128 being an option.

We will be able to run API server with following setup:

python3 -m vllm.entrypoints.api_server --model openlm-research/open_llama_3b --tensor-parallel-size 2 --max-num-seqs 2 --max-model-len 128 --block-size 128 &

On the client side, we need to assume --n=1

python3 api_client.py --stream --n=1 --prompt="one two three four "

WoosukKwon

LGTM! Thanks for submitting the PR and addressing my reviews. Happy to have the initial support for Inferentia! Let's keep optimizing the performance and making this integration more mature 🚀

vllm/worker/model_runner.py

vllm/model_executor/models/neuron/llama.py

vllm/lora/layers.py

vllm/model_executor/__init__.py

vllm/engine/llm_engine.py

vllm/engine/arg_utils.py

zhuohan123

Thanks liangfu for the work! Left more comments on styling.

zhuohan123 · 2024-02-17T20:19:17Z

vllm/config.py

@@ -374,13 +374,17 @@ def __init__(
        disable_custom_all_reduce: bool = False,
    ) -> None:
        self.pipeline_parallel_size = pipeline_parallel_size
-        self.tensor_parallel_size = tensor_parallel_size
+        if is_neuron():


Please comment on why the is_neuron here is needed.

comment added above.

zhuohan123 · 2024-02-17T20:19:32Z

vllm/config.py

-        self.world_size = pipeline_parallel_size * tensor_parallel_size
-        if self.world_size > 1:
+        self.world_size = pipeline_parallel_size * self.tensor_parallel_size
+        if self.world_size > 1 and not is_neuron():


Please comment on why the is_neuron here is needed.

comment added above.

zhuohan123 · 2024-02-17T20:20:45Z

vllm/engine/llm_engine.py

+        if self.is_neuron:
+            from vllm.worker.neuron_worker import Worker
+        else:
+            from vllm.worker.worker import Worker


Can we change this to a dictionary from hardware to workers?

changed to a dict to load the modules dynamically based on device config.

zhuohan123 · 2024-02-17T20:21:25Z

vllm/engine/llm_engine.py

+        if self.is_neuron:
+            from vllm.worker.neuron_worker import Worker
+        else:
+            from vllm.worker.worker import Worker


Again, I personally believe a dictionary will look more clear

reusing the same dispatch_worker function, which implements the dict for load the device-specific worker.

zhuohan123 · 2024-02-17T20:24:09Z

vllm/engine/llm_engine.py

+        if not self.is_neuron:
+            self._run_workers("warm_up_model")


Why don't we warm up for neuron?

I feel like the warmup here is a no-op. We should minimize the number of if is_neurons

Nice catch. The warm-up function becomes a no-op for the neuron backend.

zhuohan123 · 2024-02-17T20:27:04Z

vllm/model_executor/models/__init__.py

+        if is_neuron():
+            module_name = _NEURON_SUPPORTED_MODELS[model_arch]


Why do we need these lines?

The module_name would become neuron.llama instead of llama. We use separate module for neuron support.

vllm/model_executor/models/llama.py

becomes

vllm/model_executor/models/neuron/llama.py

zhuohan123 · 2024-02-17T20:28:15Z

vllm/worker/cache_engine.py

@@ -39,6 +38,9 @@ def __init__(
        self.num_gpu_blocks = cache_config.num_gpu_blocks
        self.num_cpu_blocks = cache_config.num_cpu_blocks

+        if is_neuron():


Please comment on why the is_neuron here is needed and the logic here.

added self.logits_as_hidden_states flag

zhuohan123 · 2024-02-17T20:29:04Z

vllm/worker/model_runner.py

+        if self.is_neuron:
+            self.model = get_model(self.model_config, self.parallel_config,
+                                   self.scheduler_config)
+        else:
+            self.model = get_model(self.model_config, self.device_config,
+                                   self.lora_config)


Why do we have different function arguments for neuron get_model?

The function arguments are different because:
1/ transformer-neuronx doesn't support lora inference yet.
2/ transformer-neuronx would require parallel_config and scheduler_config to compile the kernels.

This has been changed to utils.get_model that fetches the device-specific get_model dynamically.
The call site of the get_model is exactly the same among different devices.

zhuohan123 · 2024-02-17T20:29:44Z

vllm/worker/model_runner.py

@@ -309,7 +314,8 @@ def _prepare_decode(
        use_captured_graph = (
            not self.model_config.enforce_eager
            and batch_size <= _BATCH_SIZES_TO_CAPTURE[-1]
-            and max_context_len <= self.max_context_len_to_capture)
+            and max_context_len <= self.max_context_len_to_capture
+            and not self.is_neuron)


In config, can we just set enforce_eager for neuron? This can help reduce the number of is_neuron in the code.

changed to set enforce_eager for neuron support.

vllm/worker/model_runner.py

Co-authored-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

vllm/model_executor/layers/sampler.py

WoosukKwon · 2024-02-28T17:36:44Z

@liangfu Thanks for submitting the PR and sorry for the delays in the review! Let's merge the PR and do refactoring to isolate the neuron backend from others.

zhuohan123 · 2024-03-05T06:39:52Z

vllm/model_executor/model_loader.py

Why is the change in this file needed?

Care is taken in the code to avoid initializing CUDA prior to CUDA_VISIBLE_DEVICES being set in the worker, but an instance of this was inadvertently introduced in vllm-project#2569.

WoosukKwon added the Inferentia/Trainium label Jan 24, 2024

liangfu mentioned this pull request Jan 24, 2024

[RFC] Initial Support for AWS Inferentia #1866

Closed

5 tasks

liangfu force-pushed the neuron-2 branch 2 times, most recently from a9dbe4b to 242d6d1 Compare January 25, 2024 00:08

WoosukKwon self-requested a review January 25, 2024 18:07

zhuohan123 mentioned this pull request Jan 31, 2024

[Roadmap] vLLM Roadmap Q1 2024 #2681

Closed

30 tasks

liangfu force-pushed the neuron-2 branch from 8629083 to 5c0e2cf Compare February 2, 2024 18:31

liangfu force-pushed the neuron-2 branch from 6b1305c to 5394a21 Compare February 5, 2024 18:22

liangfu mentioned this pull request Feb 6, 2024

[DOC] add setup document to support neuron backend #2777

Merged

WoosukKwon reviewed Feb 6, 2024

View reviewed changes

liangfu force-pushed the neuron-2 branch from 1bbfafb to 1e7dd38 Compare February 9, 2024 06:36

liangfu requested a review from WoosukKwon February 13, 2024 17:40

WoosukKwon approved these changes Feb 14, 2024

View reviewed changes

liangfu force-pushed the neuron-2 branch from de1077b to f02465c Compare February 14, 2024 06:58

zhuohan123 reviewed Feb 17, 2024

View reviewed changes

liangfu force-pushed the neuron-2 branch 2 times, most recently from 614a4f0 to 9645f08 Compare February 21, 2024 17:51

liangfu added 7 commits February 23, 2024 00:49

support inference with transformers-neuronx

ee2f2ba

add the missing file

9ef2a92

fix format

eb63ed1

support external split model directory

d904235

address review comments

5b5624b

fix format

3ccac8d

fix format

d757a5a

liangfu and others added 11 commits February 23, 2024 00:49

Update vllm/model_executor/__init__.py

fab6267

Co-authored-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

specify containts of neuron backend support

a9f094b

specify device argument explicitly

2841688

fix format

8f20019

Apply suggestions from code review

c507d83

Co-authored-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

address review comments

b6c9da7

more device agnostic changes

2359ace

fix typo

9fc4b0b

fix lora sampler

470e737

fix lora test

a7ffe34

fix sample indices

1fb8b8d

liangfu force-pushed the neuron-2 branch from 9645f08 to 1fb8b8d Compare February 23, 2024 01:04

fix format

288b2d7

liangfu requested a review from zhuohan123 February 26, 2024 18:46

WoosukKwon reviewed Feb 26, 2024

View reviewed changes

vllm/model_executor/layers/sampler.py Show resolved Hide resolved

WoosukKwon merged commit 3b7178c into vllm-project:main Feb 28, 2024
22 checks passed

xjpang pushed a commit to xjpang/vllm that referenced this pull request Mar 4, 2024

[Neuron] Support inference with transformers-neuronx (vllm-project#2569)

2df1c2b

zhuohan123 reviewed Mar 5, 2024

View reviewed changes

vllm/model_executor/model_loader.py

Copy link

Member

zhuohan123 Mar 5, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is the change in this file needed?

njhill mentioned this pull request Mar 5, 2024

fix cuda device not found error when LLM is initialized in ray actor #3198

Closed

njhill mentioned this pull request Mar 19, 2024

[BugFix] Avoid initializing CUDA too early #3487

Merged

sangstar mentioned this pull request Mar 21, 2024

[Bug]: Using pytest --forked when CUDA is not generally fork-safe #3557

Closed

Temirulan pushed a commit to Temirulan/vllm-whisper that referenced this pull request Sep 6, 2024

[Neuron] Support inference with transformers-neuronx (vllm-project#2569)

4ceddaf

		if is_neuron():
		module_name = _NEURON_SUPPORTED_MODELS[model_arch]

Support inference with transformers-neuronx #2569

Support inference with transformers-neuronx #2569

Conversation

liangfu commented Jan 24, 2024 • edited Loading

WoosukKwon commented Jan 29, 2024

liangfu commented Jan 30, 2024

WoosukKwon commented Jan 30, 2024

liangfu commented Jan 30, 2024

WoosukKwon commented Feb 2, 2024

liangfu commented Feb 3, 2024

WoosukKwon left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

WoosukKwon left a comment

Choose a reason for hiding this comment

zhuohan123 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

liangfu Feb 21, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

WoosukKwon commented Feb 28, 2024

Choose a reason for hiding this comment

liangfu commented Jan 24, 2024 •

edited

Loading

liangfu Feb 21, 2024 •

edited

Loading