fix: Use default torch timeout for nccl watchdog unless overridden

The default value is recommended, and we should not change it in
production. The knob may still be useful for debugging or testing
purposes though.

Signed-off-by: Ihar Hrachyshka <ihar.hrachyshka@gmail.com>

Author: booxter

src/instructlab/training/main_ds.py

@@ -2,6 +2,7 @@
 
 # Standard
 from copy import deepcopy
+import functools
 from pathlib import Path
 import argparse
 import datetime
@@ -533,6 +534,22 @@ def train(
         )
 
 
+def _get_collective_timeout() -> datetime.timedelta | None:
+    timeout_var = os.getenv("INSTRUCTLAB_NCCL_TIMEOUT_MS")
+    if timeout_var is None:
+        return None
+
+    try:
+        timeout = int(timeout_var)
+    except ValueError:
+        timeout = -1
+
+    if timeout <= 0:
+        raise ValueError(f"Invalid value for INSTRUCTLAB_NCCL_TIMEOUT_MS: {timeout_var}. Must be a positive integer.")
+
+    return datetime.timedelta(milliseconds=timeout)
+
+
 def main(args):
     # Third Party
     import yaml
@@ -566,15 +583,17 @@ def main(args):
     model_conf = AutoConfig.from_pretrained(args.model_name_or_path)
     args.model_type = model_conf.model_type
 
-    # solution discovered from torchtune https://github.com/pytorch/torchtune/issues/2093
-    # gets converted to a timedelta of 1:40:00 if the default is kept
-    nccl_timeout = int(os.getenv("INSTRUCTLAB_NCCL_TIMEOUT_MS", "6000000"))
     #### distributed init #####
     torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
     args.local_rank = int(os.environ["LOCAL_RANK"])
-    torch.distributed.init_process_group(
-        "nccl", timeout=datetime.timedelta(milliseconds=nccl_timeout)
-    )
+
+    timeout = _get_collective_timeout()
+    init = functools.partial(torch.distributed.init_process_group, "nccl")
+    if timeout is not None:
+        init(timeout=timeout)
+    else:
+        init()
+
     args.global_rank = torch.distributed.get_rank()
     tensor = torch.ByteTensor([False]).cuda()
     torch.distributed.all_reduce(tensor)

tests/unit/test_main_ds.py

@@ -0,0 +1,28 @@
+import datetime
+from unittest import mock
+
+import pytest
+
+from instructlab.training import main_ds
+
+
+def test__get_collective_timeout():
+    # Test with default timeout
+    assert main_ds._get_collective_timeout() is None
+
+    # Test with custom timeout
+    timeout = 1234
+    with mock.patch.dict(main_ds.os.environ, {"INSTRUCTLAB_NCCL_TIMEOUT_MS": str(timeout)}):
+        assert main_ds._get_collective_timeout() == datetime.timedelta(milliseconds=timeout)
+
+    # Test with invalid timeout (negative)
+    invalid_timeout = "-100"
+    with mock.patch.dict(main_ds.os.environ, {"INSTRUCTLAB_NCCL_TIMEOUT_MS": invalid_timeout}):
+        with pytest.raises(ValueError):
+            main_ds._get_collective_timeout()
+
+    # Test with invalid timeout (string)
+    invalid_timeout = "invalid"
+    with mock.patch.dict(main_ds.os.environ, {"INSTRUCTLAB_NCCL_TIMEOUT_MS": invalid_timeout}):
+        with pytest.raises(ValueError):
+            main_ds._get_collective_timeout()