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

initialise process group for tp

Browse Source
This commit is contained in:
bursteratom
2024-12-11 11:37:21 -05:00
parent 85381b6b15
commit b17b1aada7
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -827,7 +827,7 @@ class ModelLoader:
_ = _configure_zero3_memory_efficient_loading() _ = _configure_zero3_memory_efficient_loading()
if self.cfg.tensor_parallel == "auto": 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}") device = torch.device(f"cuda:{rank}")
torch.distributed.init_process_group("nccl", device_id=device) torch.distributed.init_process_group("nccl", device_id=device)