Repository: vllm-project/vllm
Pull request: [Bugfix] Specify device when loading LoRA and embedding tensors
[Bugfix] Assign device when loading LoRA modules from file

Fixes [#3374](https://github.com/vllm-project/vllm/issues/3374). Note — existing PR to address this issue got stale; so bumping with some light updates.

This PR addresses the issue of CUDA device mismatch when loading LoRA modules from files. It fixes the `Attempting to deserialize object on CUDA device X but torch.cuda.device_count() is Y` error by explicitly specifying the device during tensor loading.

Changes:
- Add `map_location="device"` when loading LoRA tensors from .bin files
- Add `map_location="device"` when loading new embeddings from .bin files


Problem statement:
lora load failed
```
Exception in callback functools.partial(<function _raise_exception_on_finish at 0x7f9123d01870>, request_tracker=<vllm.engine.async_llm_engine.RequestTracker object at 0x7f91197ea980>)
handle: <Handle functools.partial(<function _raise_exception_on_finish at 0x7f9123d01870>, request_tracker=<vllm.engine.async_llm_engine.RequestTracker object at 0x7f91197ea980>)>
Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/vllm/engine/async_llm_engine.py", line 29, in _raise_exception_on_finish
    task.result()
  File "/usr/local/lib/python3.10/dist-packages/vllm/engine/async_llm_engine.py", line 414, in run_engine_loop
    has_requests_in_progress = await self.engine_step()
  File "/usr/local/lib/python3.10/dist-packages/vllm/engine/async_llm_engine.py", line 393, in engine_step
    request_outputs = await self.engine.step_async()
  File "/usr/local/lib/python3.10/dist-packages/vllm/engine/async_llm_engine.py", line 189, in step_async
    all_outputs = await self._run_workers_async(
  File "/usr/local/lib/python3.10/dist-packages/vllm/engine/async_llm_engine.py", line 276, in _run_workers_async
    all_outputs = await asyncio.gather(*coros)
  File "/usr/lib/python3.10/asyncio/tasks.py", line 650, in _wrap_awaitable
    return (yield from awaitable.__await__())
ray.exceptions.RayTaskError(RuntimeError): ray::RayWorkerVllm.execute_method() (pid=59396, ip=10.192.82.3, actor_id=81f1bd9d1e1f68a1ddb5a26301000000, repr=<vllm.engine.ray_utils.RayWorkerVll
m object at 0x7f3fe6e189a0>)
  File "/usr/local/lib/python3.10/dist-packages/vllm/lora/models.py", line 212, in from_local_checkpoint
    tensors = torch.load(lora_bin_file_path)
  File "/usr/local/lib/python3.10/dist-packages/torch/serialization.py", line 1014, in load
    return _load(opened_zipfile,
  File "/usr/local/lib/python3.10/dist-packages/torch/serialization.py", line 1422, in _load
    result = unpickler.load()
  File "/usr/local/lib/python3.10/dist-packages/torch/serialization.py", line 1392, in persistent_load
    typed_storage = load_tensor(dtype, nbytes, key, _maybe_decode_ascii(location))
  File "/usr/local/lib/python3.10/dist-packages/torch/serialization.py", line 1366, in load_tensor
    wrap_storage=restore_location(storage, location),
  File "/usr/local/lib/python3.10/dist-packages/torch/serialization.py", line 381, in default_restore_location
    result = fn(storage, location)
  File "/usr/local/lib/python3.10/dist-packages/torch/serialization.py", line 274, in _cuda_deserialize
    device = validate_cuda_device(location)
  File "/usr/local/lib/python3.10/dist-packages/torch/serialization.py", line 265, in validate_cuda_device
    raise RuntimeError('Attempting to deserialize object on CUDA device '
RuntimeError: Attempting to deserialize object on CUDA device 5 but torch.cuda.device_count() is 2. Please use torch.load with map_location to map your storages to an existing device.
The above exception was the direct cause of the following exception:

ray::RayWorkerVllm.execute_method() (pid=59396, ip=10.192.82.3, actor_id=81f1bd9d1e1f68a1ddb5a26301000000, repr=<vllm.engine.ray_utils.RayWorkerVllm object at 0x7f3fe6e189a0>)
  File "/usr/local/lib/python3.10/dist-packages/vllm/engine/ray_utils.py", line 37, in execute_method
    return executor(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/vllm/worker/worker.py", line 223, in execute_model
    output = self.model_runner.execute_model(seq_group_metadata_list,
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/vllm/worker/model_runner.py", line 574, in execute_model
    self.set_active_loras(lora_requests, lora_mapping)
  File "/usr/local/lib/python3.10/dist-packages/vllm/worker/model_runner.py", line 660, in set_active_loras
    self.lora_manager.set_active_loras(lora_requests, lora_mapping)
  File "/usr/local/lib/python3.10/dist-packages/vllm/lora/worker_manager.py", line 112, in set_active_loras
    self._apply_loras(lora_requests)
  File "/usr/local/lib/python3.10/dist-packages/vllm/lora/worker_manager.py", line 226, in _apply_loras
    self.add_lora(lora)
  File "/usr/local/lib/python3.10/dist-packages/vllm/lora/worker_manager.py", line 233, in add_lora
    lora = self._load_lora(lora_request)
  File "/usr/local/lib/python3.10/dist-packages/vllm/lora/worker_manager.py", line 151, in _load_lora
    raise RuntimeError(
RuntimeError: Loading lora /root/.cache/huggingface/hub/models--shibing624--ziya-llama-13b-medical-lora/snapshots/eadf08120f701bbdf0abaa4b1f2bca6cc1b23c1d/ failed
```

Hints from issue discussion:
Similar error (Jul 29) here: 
```shell
File [/usr/local/lib/python3.10/dist-packages/vllm/lora/worker_manager.py:195](https://566mynctadei05-8888.proxy.runpod.net/lab/tree/workspace/graph-llm/usr/local/lib/python3.10/dist-packages/vllm/lora/worker_manager.py#line=194), in LRUCacheWorkerLoRAManager._apply_adapters(self, lora_requests)
    190     raise RuntimeError(
    191         f"Number of requested LoRAs ({len(loras_map)}) is greater "
    192         "than the number of GPU LoRA slots "
    193         f"({self._adapter_manager.lora_slots}).")
    194 for lora in loras_map.values():
--> 195     self.add_adapter(lora)

File [/usr/local/lib/python3.10/dist-packages/vllm/lora/worker_manager.py:204](https://566mynctadei05-8888.proxy.runpod.net/lab/tree/workspace/graph-llm/usr/local/lib/python3.10/dist-packages/vllm/lora/worker_manager.py#line=203), in LRUCacheWorkerLoRAManager.add_adapter(self, lora_request)
    201         assert isinstance(self._adapter_manager,
    202                           LRUCacheLoRAModelManager)
    203         self._adapter_manager.remove_oldest_adapter()
--> 204     lora = self._load_adapter(lora_request)
    205     loaded = self._adapter_manager.add_adapter(lora)
    206 else:
    207     # If the lora is already loaded, just touch it to
    208     # update its position in the caches

File [/usr/local/lib/python3.10/dist-packages/vllm/lora/worker_manager.py:107](https://566mynctadei05-8888.proxy.runpod.net/lab/tree/workspace/graph-llm/usr/local/lib/python3.10/dist-packages/vllm/lora/worker_manager.py#line=106), in WorkerLoRAManager._load_adapter(self, lora_request)
     94     lora = self._lora_model_cls.from_local_checkpoint(
     95         lora_path,
     96         expected_lora_modules,
   (...)
    104         embedding_padding_modules=self.embedding_padding_modules,
    105     )
    106 except Exception as e:
--> 107     raise RuntimeError(f"Loading lora {lora_path} failed") from e
    108 if lora.rank > self.lora_config.max_lora_rank:
    109     raise ValueError(
    110         f"LoRA rank {lora.rank} is greater than max_lora_rank "
    111         f"{self.lora_config.max_lora_rank}.")

RuntimeError: Loading lora /root/.cache/huggingface/hub/models--Ksgk-fy--IGR-Adaptor-Meta-Llama-3-8B-Instruct-1/snapshots/37a69f049c3d0233bc77e9c8ac481d909cd6699b failed
```
> Similar error (Jul 29) here:
> 
> ```shell
> File [/usr/local/lib/python3.10/dist-packages/vllm/lora/worker_manager.py:195](https://566mynctadei05-8888.proxy.runpod.net/lab/tree/workspace/graph-llm/usr/local/lib/python3.10/dist-packages/vllm/lora/worker_manager.py#line=194), in LRUCacheWorkerLoRAManager._apply_adapters(self, lora_requests)
>     190     raise RuntimeError(
>     191         f"Number of requested LoRAs ({len(loras_map)}) is greater "
>     192         "than the number of GPU LoRA slots "
>     193         f"({self._adapter_manager.lora_slots}).")
>     194 for lora in loras_map.values():
> --> 195     self.add_adapter(lora)
> 
> File [/usr/local/lib/python3.10/dist-packages/vllm/lora/worker_manager.py:204](https://566mynctadei05-8888.proxy.runpod.net/lab/tree/workspace/graph-llm/usr/local/lib/python3.10/dist-packages/vllm/lora/worker_manager.py#line=203), in LRUCacheWorkerLoRAManager.add_adapter(self, lora_request)
>     201         assert isinstance(self._adapter_manager,
>     202                           LRUCacheLoRAModelManager)
>     203         self._adapter_manager.remove_oldest_adapter()
> --> 204     lora = self._load_adapter(lora_request)
>     205     loaded = self._adapter_manager.add_adapter(lora)
>     206 else:
>     207     # If the lora is already loaded, just touch it to
>     208     # update its position in the caches
> 
> File [/usr/local/lib/python3.10/dist-packages/vllm/lora/worker_manager.py:107](https://566mynctadei05-8888.proxy.runpod.net/lab/tree/workspace/graph-llm/usr/local/lib/python3.10/dist-packages/vllm/lora/worker_manager.py#line=106), in WorkerLoRAManager._load_adapter(self, lora_request)
>      94     lora = self._lora_model_cls.from_local_checkpoint(
>      95         lora_path,
>      96         expected_lora_modules,
>    (...)
>     104         embedding_padding_modules=self.embedding_padding_modules,
>     105     )
>     106 except Exception as e:
> --> 107     raise RuntimeError(f"Loading lora {lora_path} failed") from e
>     108 if lora.rank > self.lora_config.max_lora_rank:
>     109     raise ValueError(
>     110         f"LoRA rank {lora.rank} is greater than max_lora_rank "
>     111         f"{self.lora_config.max_lora_rank}.")
> 
> RuntimeError: Loading lora /root/.cache/huggingface/hub/models--Ksgk-fy--IGR-Adaptor-Meta-Llama-3-8B-Instruct-1/snapshots/37a69f049c3d0233bc77e9c8ac481d909cd6699b failed
> ```

it's fixed. see the related pr upstairs.

Commit message: fix: Specify device when loading LoRA and embedding tensors

- Add map_location="device" when loading LoRA tensors from .bin files
- Add map_location="device" when loading new embeddings from .bin files

Patch to review:
diff --git a/vllm/lora/models.py b/vllm/lora/models.py
index 017a1002bb9..2ce46204ab9 100644
--- a/vllm/lora/models.py
+++ b/vllm/lora/models.py
@@ -248,7 +248,7 @@ def from_local_checkpoint(
                     f" target modules in {expected_lora_modules}"
                     f" but received {unexpected_modules}."
                     f" Please verify that the loaded LoRA module is correct")
-            tensors = torch.load(lora_bin_file_path)
+            tensors = torch.load(lora_bin_file_path, map_location="device")
         else:
             raise ValueError(f"{lora_dir} doesn't contain tensors")
 
@@ -257,7 +257,8 @@ def from_local_checkpoint(
             embeddings = safetensors.torch.load_file(
                 new_embeddings_tensor_path)
         elif os.path.isfile(new_embeddings_bin_file_path):
-            embeddings = torch.load(new_embeddings_bin_file_path)
+            embeddings = torch.load(new_embeddings_bin_file_path,
+                                    map_location="device")
 
         rank = config["r"]
         lora_alpha = config["lora_alpha"]


Code review comments:
diff_hunk: @@ -248,7 +248,7 @@ def from_local_checkpoint(
                     f" target modules in {expected_lora_modules}"
                     f" but received {unexpected_modules}."
                     f" Please verify that the loaded LoRA module is correct")
-            tensors = torch.load(lora_bin_file_path)
+            tensors = torch.load(lora_bin_file_path, map_location="device")
path: vllm/lora/models.py
text: @user1:
did you test it? I don't think `"device"` works. this is a string.