From 5de58d595547f57dd4df2f3572158eb9cbdd24c9 Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Tue, 19 Nov 2024 19:44:44 +0800 Subject: [PATCH] fix cpu bnb path (#34647) * fix cpu bnb path * Update src/transformers/generation/utils.py Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com> * fix awq quantizer env check * fix awq quantizer device check Signed-off-by: jiqing-feng --------- Signed-off-by: jiqing-feng Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com> --- src/transformers/generation/utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index 53cd2df3a49c84..97a294fd427987 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -1636,7 +1636,10 @@ def get_layer_device_map(execution_device_map: Optional[dict] = None): # This is needed here if we don't want to make changes in accelerate in order to save execution_device # For offloaded case, we need to get the execution device, not just the device where it is offloaded if hasattr(self, "hf_device_map"): - main_device = [d for d in self.hf_device_map.values() if d not in ["cpu", "disk"]][0] + if set(self.hf_device_map.values()) == {"cpu"} or set(self.hf_device_map.values()) == {"cpu", "disk"}: + main_device = "cpu" + else: + main_device = [d for d in self.hf_device_map.values() if d not in ["cpu", "disk"]][0] execution_device_map = { name: main_device if device in ["cpu", "disk"] else device for name, device in self.hf_device_map.items()