From d663652216e23b38138d041c0d5c0fdb12ded91d Mon Sep 17 00:00:00 2001 From: Sunny Liu Date: Tue, 18 Feb 2025 17:08:32 -0500 Subject: [PATCH] del device_map for tp --- src/axolotl/utils/models.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py index b0d81dd56..6d1bb12c3 100644 --- a/src/axolotl/utils/models.py +++ b/src/axolotl/utils/models.py @@ -764,6 +764,7 @@ class ModelLoader: skip_move_to_device = False if self.cfg.tp_size is not None: self.model_kwargs["tp_plan"] = "auto" + del self.model_kwargs["device_map"] if ( # pylint: disable=condition-evals-to-constant) (self.cfg.fsdp and self.cfg.fsdp_config.fsdp_cpu_ram_efficient_loading)