From 2014f581812a1914ed6bcca2afeca14b8696599f Mon Sep 17 00:00:00 2001 From: bursteratom Date: Wed, 11 Dec 2024 11:45:07 -0500 Subject: [PATCH] set os environ RANK --- 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 ca695846f..a3f90bca5 100644 --- a/src/axolotl/utils/models.py +++ b/src/axolotl/utils/models.py @@ -829,6 +829,7 @@ class ModelLoader: if self.cfg.tensor_parallel == "auto": rank = int(os.environ.get("LOCAL_RANK", 0)) os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = os.getenv("WORLD_SIZE", "1") device = torch.device(f"cuda:{rank}") torch.distributed.init_process_group("nccl", device_id=device)