mirror/OpenLLM
1
0
Fork 0
mirror of https://github.com/bentoml/OpenLLM.git synced 2026-04-22 16:07:24 -04:00
Code Issues Packages Projects Releases Wiki Activity

fix(dolly): make sure to use GPU when available

Browse Source 
map device_map to auto when GPU is available

Signed-off-by: aarnphm-ec2-dev <29749331+aarnphm@users.noreply.github.com>
This commit is contained in:
aarnphm-ec2-dev
2023-06-15 05:52:25 +00:00
parent dfe71d7867
commit b3d924e6d6
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/openllm/models/dolly_v2/modeling_dolly_v2.py
View File

@@ -38,7 +38,7 @@ class DollyV2(openllm.LLM["transformers.Pipeline", "transformers.PreTrainedToken
@property
def import_kwargs(self):
model_kwds = {
"device_map": "auto" if torch.cuda.is_available() and torch.cuda.device_count() > 1 else None,
"device_map": "auto" if torch.cuda.is_available() else None,
"torch_dtype": torch.bfloat16,
}
tokenizer_kwds = {"padding_side": "left"}