Repository: vllm-project/vllm
Pull request: [Bugfix] neuron: enable tensor parallelism
FILL IN THE PR DESCRIPTION HERE

FIX #6269

Following #7175, this enables tensor parallelism on neuron with vLLM > 0.5.0. The `block-size` choices are also reverted back to the previous values and vLLM will take the `max-model-len` as the `block-size` on neuron devices. 


Problem statement:
f[Bug]: TypeError: Can't instantiate abstract class NeuronWorker with abstract method execute_worker
### Your current environment

```text
The output of `python collect_env.py`
```
Collecting environment information...
WARNING 07-09 19:46:51 _custom_ops.py:14] Failed to import from vllm._C with ModuleNotFoundError("No module named 'vllm._C'")
PyTorch version: 2.1.2+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.4 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.30.0
Libc version: glibc-2.35

Python version: 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-1031-aws-x86_64-with-glibc2.35
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
Address sizes:                   48 bits physical, 48 bits virtual
Byte Order:                      Little Endian
CPU(s):                          32
On-line CPU(s) list:             0-31
Vendor ID:                       AuthenticAMD
Model name:                      AMD EPYC 7R13 Processor
CPU family:                      25
Model:                           1
Thread(s) per core:              2
Core(s) per socket:              16
Socket(s):                       1
Stepping:                        1
BogoMIPS:                        5299.99
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy cr8_legacy abm sse4a misalignsse 3dnowprefetch topoext invpcid_single ssbd ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 clzero xsaveerptr rdpru wbnoinvd arat npt nrip_save vaes vpclmulqdq rdpid
Hypervisor vendor:               KVM
Virtualization type:             full
L1d cache:                       512 KiB (16 instances)
L1i cache:                       512 KiB (16 instances)
L2 cache:                        8 MiB (16 instances)
L3 cache:                        64 MiB (2 instances)
NUMA node(s):                    2
NUMA node0 CPU(s):               0-7,16-23
NUMA node1 CPU(s):               8-15,24-31
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 and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP always-on, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected

Versions of relevant libraries:
[pip3] numpy==1.25.2
[pip3] nvidia-nccl-cu12==2.18.1
[pip3] torch==2.1.2
[pip3] torch-neuronx==2.1.2.2.2.0
[pip3] torch-xla==2.1.3
[pip3] torchvision==0.16.2
[pip3] transformers==4.42.3
[pip3] transformers-neuronx==0.11.351
[pip3] triton==2.1.0
[conda] Could not collect
ROCM Version: Could not collect
Neuron SDK Version: (0, 'instance-type: inf2.8xlarge\ninstance-id: i-0d1255be16fb43afc\n+--------+--------+--------+---------+\n| NEURON | NEURON | NEURON |   PCI   |\n| DEVICE | CORES  | MEMORY |   BDF   |\n+--------+--------+--------+---------+\n| 0      | 2      | 32 GB  | 00:1f.0 |\n+--------+--------+--------+---------+', '')
vLLM Version: 0.5.1
vLLM Build Flags:
CUDA Archs: Not Set; ROCm: Disabled; Neuron: Disabled
GPU Topology:
Could not collect


### 🐛 Describe the bug

Offline Batched Inference Example from your Quickstart page, running on Neuron inf2 (AWS inf2.8xlarge)  returns: "TypeError: Can't instantiate abstract class NeuronWorker with abstract method execute_worker". The only change from your example is the example code is adding device="neuron" argument to LLM. 

This code should reproduce the error:
------------------------------------------
from vllm import LLM, SamplingParams

prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

llm = LLM(model="facebook/opt-125m",device="neuron")
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}")

Hints from issue discussion:
I still get this error with the latest https://github.com/vllm-project/vllm/pull/6313 merge.

```
Traceback (most recent call last):
  File "/home/ubuntu/vllm/simple_test.py", line 11, in <module>
    llm = LLM(model="facebook/opt-125m",device="neuron")
  File "/home/ubuntu/vllm/vllm/entrypoints/llm.py", line 150, in __init__
    self.llm_engine = LLMEngine.from_engine_args(
  File "/home/ubuntu/vllm/vllm/engine/llm_engine.py", line 421, in from_engine_args
    engine = cls(
  File "/home/ubuntu/vllm/vllm/engine/llm_engine.py", line 249, in __init__
    self.model_executor = executor_class(
  File "/home/ubuntu/vllm/vllm/executor/executor_base.py", line 46, in __init__
    self._init_executor()
  File "/home/ubuntu/vllm/vllm/executor/neuron_executor.py", line 21, in _init_executor
    self._init_worker()
  File "/home/ubuntu/vllm/vllm/executor/neuron_executor.py", line 26, in _init_worker
    self.driver_worker = NeuronWorker(
TypeError: Can't instantiate abstract class NeuronWorker with abstract method execute_worker
```

@liangfu Could you please take a look at @jianyinglangaws's comment?
Saw that the patch to #6313 was merged to main, so I did a git pull origin main to update to the latest. I can confirm that the behavior seen in #6269 is still present, namely:

TypeError: Can't instantiate abstract class NeuronWorker with abstract method execute_worker

I posted this info on #6313 as well because it was speculated that #6313 would fix #6269. I am beginning to wonder about the status of support for the Neuron architecture in vLLM. Is it vLLM's intention to drop support for Neuron? I only ask this because the vLLM Documentation suggests that at least in V 0.3 vLLM did support neuron. Now, of course, those quickstart examples do not work. Any information regarding the direction vLLM is headed relative to Neuron would be appreciated.
One extra question: can you recommend a previous known-working version of vLLM that we can pull that would work on Neuron and can you provide a tag for that?
I tried to clone and install vLLM v0.3.0 but the requirements for that include pkg_resources which has been deprecated.
The vllm 0.5.0 works for me with the latest Neuron SDK 2.19.0 and transformers-neuronx.
> The vllm 0.5.0 works for me with the latest Neuron SDK 2.19.0 and transformers-neuronx.

That's very interesting to know. So AFAIK, the DLAMI instance one can get by default from AWS in 2.18.2. How are you accessing 2.19.0? How do I install/upgrade to it?
The default Neuron DLAMI (Deep Learning AMI Neuron (Ubuntu 22.04) 20240703) was updated to 2.19.0 last week.
I see, thank you - I will try to use a new instance.
I installed vllm on DLAMI 2.19.0 following these procedures (steps 2 and 3) VERBATIM : https://docs.vllm.ai/en/latest/getting_started/neuron-installation.html. The example does not work with the same error with the Quickstart example. So I guess we need to understand how your working instance on 2.19.0 is differing from my broken? Can you provide your collect_env.py output? Or do you want me to provide mine?

> The vllm 0.5.0 works for me with the latest Neuron SDK 2.19.0 and transformers-neuronx.

Commit message: block_size should equal max_model_len on neuron, not max_num_seqs

Signed-off-by: omrishiv <327609+omrishiv@users.noreply.github.com>

Patch to review:
diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py
index 48d01fcfd8f..5a490863631 100644
--- a/vllm/engine/arg_utils.py
+++ b/vllm/engine/arg_utils.py
@@ -316,9 +316,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         parser.add_argument('--block-size',
                             type=int,
                             default=EngineArgs.block_size,
-                            choices=[8, 16, 32, 128, 256, 512, 1024, 2048],
+                            choices=[8, 16, 32],
                             help='Token block size for contiguous chunks of '
-                            'tokens.')
+                            'tokens. This is ignored on neuron devices and '
+                            'set to max-model-len')
 
         parser.add_argument('--enable-prefix-caching',
                             action='store_true',
@@ -780,7 +781,8 @@ def create_engine_config(self, ) -> EngineConfig:
             served_model_name=self.served_model_name,
             multimodal_config=multimodal_config)
         cache_config = CacheConfig(
-            block_size=self.block_size,
+            block_size=self.block_size if self.device != "neuron" else
+            self.max_model_len,  # neuron needs block_size = max_model_len
             gpu_memory_utilization=self.gpu_memory_utilization,
             swap_space=self.swap_space,
             cache_dtype=self.kv_cache_dtype,
diff --git a/vllm/executor/neuron_executor.py b/vllm/executor/neuron_executor.py
index b45d5d86b54..02627de3e0b 100644
--- a/vllm/executor/neuron_executor.py
+++ b/vllm/executor/neuron_executor.py
@@ -4,7 +4,8 @@
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
 from vllm.sequence import ExecuteModelRequest, SamplerOutput
-from vllm.utils import make_async
+from vllm.utils import (get_distributed_init_method, get_ip, get_open_port,
+                        make_async)
 
 logger = init_logger(__name__)
 
@@ -24,14 +25,17 @@ def _init_executor(self) -> None:
 
     def _init_worker(self):
         from vllm.worker.neuron_worker import NeuronWorker
-
+        distributed_init_method = get_distributed_init_method(
+            get_ip(), get_open_port())
         self.driver_worker = NeuronWorker(
-            self.model_config,
-            self.parallel_config,
-            self.scheduler_config,
-            self.device_config,
-            self.cache_config,
-        )
+            model_config=self.model_config,
+            parallel_config=self.parallel_config,
+            scheduler_config=self.scheduler_config,
+            device_config=self.device_config,
+            cache_config=self.cache_config,
+            local_rank=0,
+            rank=0,
+            distributed_init_method=distributed_init_method)
         self.driver_worker.init_device()
         self.driver_worker.load_model()
 
diff --git a/vllm/worker/neuron_worker.py b/vllm/worker/neuron_worker.py
index 3b0ded36ca1..9b4367d761b 100644
--- a/vllm/worker/neuron_worker.py
+++ b/vllm/worker/neuron_worker.py
@@ -1,4 +1,5 @@
 """A Neuron worker class."""
+import os
 from typing import List, Optional, Tuple
 
 import torch
@@ -6,6 +7,8 @@
 
 from vllm.config import (CacheConfig, DeviceConfig, ModelConfig,
                          ParallelConfig, SchedulerConfig)
+from vllm.distributed import (ensure_model_parallel_initialized,
+                              init_distributed_environment)
 from vllm.model_executor import set_random_seed
 from vllm.sequence import ExecuteModelRequest
 from vllm.worker.neuron_model_runner import NeuronModelRunner
@@ -24,12 +27,18 @@ def __init__(
         scheduler_config: SchedulerConfig,
         device_config: DeviceConfig,
         cache_config: CacheConfig,
+        local_rank: int,
+        rank: int,
+        distributed_init_method: str,
     ) -> None:
         self.model_config = model_config
         self.parallel_config = parallel_config
         self.scheduler_config = scheduler_config
         self.device_config = device_config
         self.cache_config = cache_config
+        self.local_rank = local_rank
+        self.rank = rank
+        self.distributed_init_method = distributed_init_method
         if self.model_config.trust_remote_code:
             # note: lazy import to avoid importing torch before initializing
             from vllm.utils import init_cached_hf_modules
@@ -40,6 +49,9 @@ def __init__(
         self.is_driver_worker = True
 
     def init_device(self) -> None:
+        os.environ["PJRT_DEVICE"] = "NEURON"
+        self.init_distributed_environment()
+
         # Set random seed.
         set_random_seed(self.model_config.seed)
 
@@ -98,3 +110,20 @@ def get_cache_block_size_bytes(self) -> int:
         This is required for speculative decoding; it is not yet implemented.
         """
         raise NotImplementedError
+
+    def init_distributed_environment(self):
+        """Neuron uses transformers-neuronx for tensor parallelism.
+
+        vLLM still needs the environment inited when TP/PP > 1
+        """
+        init_distributed_environment(
+            world_size=1,
+            rank=self.rank,
+            local_rank=self.local_rank,
+            distributed_init_method=self.distributed_init_method,
+            backend="gloo",
+        )
+        ensure_model_parallel_initialized(
+            1,
+            1,
+        )


Code review comments:
diff_hunk: @@ -40,6 +49,9 @@ def __init__(
         self.is_driver_worker = True
 
     def init_device(self) -> None:
+        os.environ["PJRT_DEVICE"] = "NEURON"
path: vllm/worker/neuron_worker.py
text: @user1:
Help take this out?