initialise process group for tp

2024-12-11 11:37:21 -05:00
parent 85381b6b15
commit b17b1aada7
1 changed files with 1 additions and 1 deletions
--- a/src/axolotl/utils/models.py
+++ b/src/axolotl/utils/models.py
@@ -827,7 +827,7 @@ class ModelLoader:
            _ = _configure_zero3_memory_efficient_loading()

            if self.cfg.tensor_parallel == "auto":
-                rank = int(os.environ["RANK"])
+                rank = int(os.environ.get("LOCAL_RANK", 0))
                device = torch.device(f"cuda:{rank}")
                torch.distributed.init_process_group("nccl", device_id=device)