From b3d924e6d62ce19e8d7754650ce4de13d4e8a8b2 Mon Sep 17 00:00:00 2001
From: aarnphm-ec2-dev <29749331+aarnphm@users.noreply.github.com>
Date: Thu, 15 Jun 2023 05:52:25 +0000
Subject: [PATCH] fix(dolly): make sure to use GPU when available

map device_map to auto when GPU is available

Signed-off-by: aarnphm-ec2-dev <29749331+aarnphm@users.noreply.github.com>
---
 src/openllm/models/dolly_v2/modeling_dolly_v2.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/openllm/models/dolly_v2/modeling_dolly_v2.py b/src/openllm/models/dolly_v2/modeling_dolly_v2.py
index cf416012..3175d692 100644
--- a/src/openllm/models/dolly_v2/modeling_dolly_v2.py
+++ b/src/openllm/models/dolly_v2/modeling_dolly_v2.py
@@ -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"}