tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

set os environ RANK

Browse Source
This commit is contained in:
bursteratom
2024-12-11 11:40:20 -05:00
parent b17b1aada7
commit b5f9dd44f2
1 changed files with 1 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
src/axolotl/utils/models.py
View File

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