fix cpu bnb path (#34647)

* fix cpu bnb path * Update src/transformers/generation/utils.py Co-authored-by: Marc Sun <[email protected]> * fix awq quantizer env check * fix awq quantizer device check Signed-off-by: jiqing-feng <[email protected]> --------- Signed-off-by: jiqing-feng <[email protected]> Co-authored-by: Marc Sun <[email protected]>
huggingface · Nov 19, 2024 · 5de58d5 · 5de58d5
1 parent 3cd78be
commit 5de58d5
Showing 1 changed file with 4 additions and 1 deletion.
diff --git 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()