From b5f9dd44f2d7342ec018b92c1d6ffbcc977e5fe2 Mon Sep 17 00:00:00 2001 From: bursteratom Date: Wed, 11 Dec 2024 11:40:20 -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 fbe23430d..ca695846f 100644 --- a/src/axolotl/utils/models.py +++ b/src/axolotl/utils/models.py @@ -828,6 +828,7 @@ class ModelLoader: if self.cfg.tensor_parallel == "auto": rank = int(os.environ.get("LOCAL_RANK", 0)) + os.environ["RANK"] = str(rank) device = torch.device(f"cuda:{rank}") torch.distributed.init_process_group("nccl", device_id=device)