vllm-project / vllm-project/aibrix

StreamLoader library support device_map to allocate tensor to different device

Open
#404 2 comments 0 reactions 1 assignee Claimed by @brosoul View on GitHub
area/model-loader kind/enhancement priority/important-soon
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.