del device_map for tp

This commit is contained in:
Sunny Liu
2025-02-18 17:08:32 -05:00
committed by Sung Ching Liu
parent dbd43aa18f
commit d663652216

View File

@@ -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)