fix: Rename variable to control PG timeout to not refer to NCCL

The variable is not backend specific. In the future, if/when we support
other backends, this will become more evidently a problem.

Signed-off-by: Ihar Hrachyshka <[email protected]>

Author: booxter

README.md

@@ -375,4 +375,4 @@ run_training(
 
 Below is a list of custom environment variables users can set in the training library.
 
-1. `INSTRUCTLAB_NCCL_TIMEOUT_MS`, this environment variable controls the NCCL timeout in milliseconds. Consider increasing if seeing FSDP related NCCL errors.
+1. `INSTRUCTLAB_PROCESS_GROUP_TIMEOUT_MS`, this environment variable controls the process group timeout in milliseconds. Consider increasing if seeing FSDP related errors.

src/instructlab/training/const.py

@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: Apache-2.0
+
+INSTRUCTLAB_PROCESS_GROUP_TIMEOUT_MS = "INSTRUCTLAB_PROCESS_GROUP_TIMEOUT_MS"

src/instructlab/training/main_ds.py

@@ -61,6 +61,7 @@
 
 # pylint: disable=no-name-in-module
 from instructlab.training.config import DistributedBackend, TorchrunArgs, TrainingArgs
+from instructlab.training.const import INSTRUCTLAB_PROCESS_GROUP_TIMEOUT_MS
 from instructlab.training.logger import setup_metric_logger, setup_root_logger
 from instructlab.training.multipack_sampler import (
     find_packing_max_batch_len_and_grad_accum,
@@ -550,7 +551,7 @@ def train(
 # 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")
+    timeout_var = os.getenv(INSTRUCTLAB_PROCESS_GROUP_TIMEOUT_MS)
     if timeout_var is None:
         return None
 
@@ -561,7 +562,7 @@ def _get_collective_timeout() -> datetime.timedelta | None:
 
     if timeout <= 0:
         raise ValueError(
-            f"Invalid value for INSTRUCTLAB_NCCL_TIMEOUT_MS: {timeout_var}. Must be a positive integer."
+            f"Invalid value for {INSTRUCTLAB_PROCESS_GROUP_TIMEOUT_MS}: {timeout_var}. Must be a positive integer."
         )
 
     return datetime.timedelta(milliseconds=timeout)

tests/unit/test_main_ds.py

@@ -7,6 +7,7 @@
 
 # First Party
 from instructlab.training import main_ds
+from instructlab.training.const import INSTRUCTLAB_PROCESS_GROUP_TIMEOUT_MS
 
 
 def test__get_collective_timeout():
@@ -16,7 +17,7 @@ def test__get_collective_timeout():
     # Test with custom timeout
     timeout = 1234
     with mock.patch.dict(
-        main_ds.os.environ, {"INSTRUCTLAB_NCCL_TIMEOUT_MS": str(timeout)}
+        main_ds.os.environ, {INSTRUCTLAB_PROCESS_GROUP_TIMEOUT_MS: str(timeout)}
     ):
         assert main_ds._get_collective_timeout() == datetime.timedelta(
             milliseconds=timeout
@@ -25,15 +26,15 @@ def test__get_collective_timeout():
     # Test with invalid timeout (negative)
     invalid_timeout = "-100"
     with mock.patch.dict(
-        main_ds.os.environ, {"INSTRUCTLAB_NCCL_TIMEOUT_MS": invalid_timeout}
+        main_ds.os.environ, {INSTRUCTLAB_PROCESS_GROUP_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}
+        main_ds.os.environ, {INSTRUCTLAB_PROCESS_GROUP_TIMEOUT_MS: invalid_timeout}
     ):
         with pytest.raises(ValueError):
             main_ds._get_collective_timeout()

tox.ini

@@ -17,8 +17,8 @@ basepython = python3.11
 [testenv:py3-unit]
 description = run unit tests with pytest
 passenv =
-	HF_HOME
-        INSTRUCTLAB_NCCL_TIMEOUT_MS
+  HF_HOME
+  INSTRUCTLAB_PROCESS_GROUP_TIMEOUT_MS
 deps = 
     pytest
     wandb
@@ -33,8 +33,8 @@ commands = {envpython} -m pytest tests/unit {posargs}
 [testenv:py3-smoke]
 description = run accelerated smoke tests with pytest
 passenv =
-	HF_HOME
-        INSTRUCTLAB_NCCL_TIMEOUT_MS
+    HF_HOME
+    INSTRUCTLAB_PROCESS_GROUP_TIMEOUT_MS
 deps = 
     pytest
     -r requirements-dev.txt