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

@@ -5,6 +5,7 @@
 from pathlib import Path
 import argparse
 import datetime
+import functools
 import math
 import os
 import re
@@ -533,6 +534,28 @@ def train(
         )
 
 
+# This function makes an effort to stick to a default value from torch library,
+# whatever it may be. That's why we don't just set to the current (as of the
+# time of writing) default: to cover the unlikely event torch decides to tweak
+# the default.
+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 +589,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,39 @@
+# Standard
+from unittest import mock
+import datetime
+
+# Third Party
+import pytest
+
+# First Party
+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()