vllm-project / vllm-project/aibrix
StreamLoader library support device_map to allocate tensor to different device
- Dominant language
- Go
- Stars
- 5.1k
- Forks
- 694
- Avg merge
- 1d 19h
- Merged PRs (30d)
- 98
Description
### 🚀 Feature Description and Motivation
In the parameter list of the loading tensor in the StreamLoader library, `device` should be refined to `device_map`. And StreamLoader library should support device_map to allocate tensor to different device.
> device_map (Dict[str, Union[int, str, torch.device]], optional) — A map that specifies where each submodule should go. It doesn’t need to be refined to each parameter/buffer name, once a given module name is inside, every submodule of it will be sent to the same device.
Refs: https://huggingface.co/docs/accelerate/v1.0.1/en/package_reference/utilities#accelerate.utils.load_state_dict.device_map
### Use Case
_No response_
### Proposed Solution
[`infer_auto_device_map` functions](https://huggingface.co/docs/accelerate/v1.0.1/en/package_reference/utilities#accelerate.infer_auto_device_map) in accelerate could be helpful.
Contributor guide
Assessment
This issue has not been assessed yet.