fix: device could be in meta, transformers#33154 (#2089)

Signed-off-by: Yu Chin Fabian Lim <[email protected]>
huggingface · Sep 21, 2024 · 9b80f3d · 9b80f3d
1 parent 2038e52
commit 9b80f3d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/trl/trainer/sft_trainer.py b/trl/trainer/sft_trainer.py
@@ -232,7 +232,7 @@ def __init__(
                 if getattr(model, "is_loaded_in_4bit", False):
                     for _, param in model.named_parameters():
                         if param.__class__.__name__ == "Params4bit":
-                            is_sharded_qlora = param.data.device.type == "cpu"
+                            is_sharded_qlora = param.data.device.type in {"cpu", "meta"}
                             break
                 if getattr(model, "is_loaded_in_8bit", False) or (
                     getattr(model, "is_loaded_in_4bit", False) and not is_sharded_qlora