Merge pull request #516 from booxter/pull-req

Implicit distributed backend selection

Author: mergify[bot]

requirements.txt

@@ -2,9 +2,7 @@ packaging>=20.9
 wheel>=0.43
 pyyaml
 py-cpuinfo
-# we set this to be above 0a0 so that it doesn't
-# replace custom pytorch images with the 2.3.0
-torch>=2.3.0a0
+torch>=2.6.0
 transformers>=4.45.2
 
 datasets>=2.15.0

src/instructlab/training/main_ds.py

@@ -4,7 +4,6 @@
 from copy import deepcopy
 import argparse
 import datetime
-import functools
 import logging
 import os
 import subprocess
@@ -328,11 +327,10 @@ def main(args):
     args.local_rank = int(os.environ["LOCAL_RANK"])
 
     timeout = _get_collective_timeout()
-    init = functools.partial(torch.distributed.init_process_group, "nccl")
     if timeout is not None:
-        init(timeout=timeout)
+        torch.distributed.init_process_group(timeout=timeout)
     else:
-        init()
+        torch.distributed.init_process_group()
 
     args.global_rank = torch.distributed.get_rank()
     tensor = torch.ByteTensor([False]).cuda()