Unable to distribute using vLLM

[greg-8090]

To reproduce:

1. trace Llama-3.2-1B using NXD as shown in the generation_demo.py example.

2. try to deploy in a Docker container using this command vLLM command:
   vllm serve /mnt/models
   --device neuron
   --tensor-parallel-size 2
   --max-model-len 64
   --max-num-seqs 32

3. you will see the following error:

Traceback (most recent call last):
File "/opt/conda/lib/python3.10/multiprocessing/process.py", line 314, in _bootstrap
self.run()
File "/opt/conda/lib/python3.10/multiprocessing/process.py", line 108, in run
self._target(*self._args, **self._kwargs)
File "/app/vllm/vllm/vllm/engine/multiprocessing/engine.py", line 388, in run_mp_engine
engine = MQLLMEngine.from_engine_args(engine_args=engine_args,
File "/app/vllm/vllm/vllm/engine/multiprocessing/engine.py", line 138, in from_engine_args
return cls(
File "/app/vllm/vllm/vllm/engine/multiprocessing/engine.py", line 78, in init
self.engine = LLMEngine(*args,
File "/app/vllm/vllm/vllm/engine/llm_engine.py", line 325, in init
self.model_executor = executor_class(
File "/app/vllm/vllm/vllm/executor/executor_base.py", line 47, in init
self._init_executor()
File "/app/vllm/vllm/vllm/executor/neuron_executor.py", line 23, in _init_executor
self._init_worker()
File "/app/vllm/vllm/vllm/executor/neuron_executor.py", line 39, in _init_worker
self.driver_worker.load_model()
File "/app/vllm/vllm/vllm/worker/neuron_worker.py", line 134, in load_model
self.model_runner.load_model()
File "/app/vllm/vllm/vllm/worker/neuron_model_runner.py", line 120, in load_model
self.model = get_neuron_model(
File "/app/vllm/vllm/vllm/model_executor/model_loader/neuron.py", line 301, in get_neuron_model
model.load_weights(model_config.model,
File "/app/vllm/vllm/vllm/model_executor/model_loader/neuron.py", line 115, in load_weights
self.model.to_neuron()
File "/opt/conda/lib/python3.10/site-packages/transformers_neuronx/base.py", line 84, in to_neuron
self.load_weights()
File "/opt/conda/lib/python3.10/site-packages/transformers_neuronx/llama/model.py", line 118, in load_weights
self.materialize_embeddings()
File "/opt/conda/lib/python3.10/site-packages/transformers_neuronx/llama/model.py", line 234, in materialize_embeddings
self.chkpt_model.model.embed_tokens.materialize()
File "/opt/conda/lib/python3.10/site-packages/transformers_neuronx/module.py", line 80, in materialize
raise FileNotFoundError(f'Could not find a weight for {param._global_key} in {param._file_path}')
FileNotFoundError: Could not find a weight for model.embed_tokens.weight in /mnt/models/model-0002-of-0002.safetensors

[sssrijan-amazon]

Thanks for reaching out. Looking at the stack trace, vllm is attempting to serve the model using transformers-neuronx (which suggests you have both transformers-neuronx and neuronx-distributed-inference installed). Given that you mentioned in step 1 that you traced the model using generation demo example script in neuronx-distributed-inference framework, can you try setting VLLM_NEURON_FRAMEWORK environment variable to neuronx-distributed-inference as part of step 2 (vllm serve) to force vLLM to serve using neuronx-distributed-inference framework. Reference - https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-inference/developer_guides/vllm-user-guide.html#neuron-framework-selection
We are able to successfully trace and run the model using the neuronx-distributed-inference generation demo script using release 2.21.1 dependencies with the model checkpoint downloaded from huggingface llama3.2 1B.
Additionally trying out tracing and compilation along with vllm serving also works with the following command:

VLLM_NEURON_FRAMEWORK=neuronx-distributed-inference vllm serve /mnt/models --device neuron --tensor-parallel-size 2 --max-model-len 64 --max-num-seqs 32

In case you are trying to use the traced model from generation demo example script to serve vLLM requests, its recommended to use is_continuous_batching=True and ctx_batch_size=1 while initializing the NeuronConfig object here. After tracing the model using the script, use the following command to serve the compiled model using vllm with the traced model path VLLM_NEURON_FRAMEWORK=neuronx-distributed-inference NEURON_COMPILED_ARTIFACTS=/mnt/traced-models/Llama-3.2-1B vllm serve /mnt/models --device neuron --tensor-parallel-size 32 --max-model-len 64 --max-num-seqs 2 --override-neuron-config "{}"
More details on using traced model path to serve with vllm can be found here - https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-inference/developer_guides/vllm-user-guide.html#loading-pre-compiled-models-serialization-support

Please let us know if this resolves the issue you're seeing.

[aarondou]

Please let us know if the issue resolved or you are still having issues. Thanks

[aarondou]

Close the issue assuming it is resolved based on the recommendation. Feel free to reopen if you still have issues.