From dbd43aa18f1319d443009ed3307efac22dc5ad90 Mon Sep 17 00:00:00 2001
From: Sunny Liu <sunny19981005@outlook.com>
Date: Tue, 18 Feb 2025 17:04:13 -0500
Subject: [PATCH] set tp_plan

---
 src/axolotl/utils/models.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py
index c4c07dd33..b0d81dd56 100644
--- a/src/axolotl/utils/models.py
+++ b/src/axolotl/utils/models.py
@@ -762,6 +762,9 @@ class ModelLoader:
             return hf_ds_cfg
 
         skip_move_to_device = False
+        if self.cfg.tp_size is not None:
+            self.model_kwargs["tp_plan"] = "auto"
+
         if (  # pylint: disable=condition-evals-to-constant)
             (self.cfg.fsdp and self.cfg.fsdp_config.fsdp_cpu_ram_efficient_loading)
             and not qlora_fsdp