Repository: instructlab/training
Pull request #521: fix: Use default torch timeout for nccl watchdog unless overridden
Author: login: booxter
Labels: name: testing
description: Relates to testing
Created: 2025-05-02T13:57:14Z | Merged: 2025-05-20T23:43:01Z

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>



Diff:
diff --git a/src/instructlab/training/main_ds.py b/src/instructlab/training/main_ds.py
index 4ea8787a..a7511bcd 100644
--- a/src/instructlab/training/main_ds.py
+++ b/src/instructlab/training/main_ds.py
@@ -5,6 +5,7 @@
 from pathlib import Path
 import argparse
 import datetime
+import functools
 import logging
 import math
 import os
@@ -544,6 +545,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):
     if args.distributed_training_framework == "deepspeed" and not FusedAdam:
         raise ImportError(
@@ -571,15 +594,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)
diff --git a/tests/unit/test_main_ds.py b/tests/unit/test_main_ds.py
new file mode 100644
index 00000000..12b35127
--- /dev/null
+++ b/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()


Reviews:
user: login: RobotSail
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-03T23:00:53Z
user: login: booxter
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-05T16:27:14Z
user: login: RobotSail
body: We just need some comments explaining why we're doing it this way, otherwise I think it's good. 👍 
state: CHANGES_REQUESTED
author_association: MEMBER
submitted_at: 2025-05-05T22:03:55Z
user: login: JamesKunstle
body: I like this solution.
I'm curious whether the torch environment variable `TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC` from the docs overrides this or this overrides that var. 
https://pytorch.org/docs/stable/torch_nccl_environment_variables.html
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-06T04:15:57Z
user: login: RobotSail
body: @booxter Since this is an important debug env, could you also please update the README.md to let people know that we have this env? Otherwise LGTM
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-12T23:11:06Z
user: login: booxter
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-14T01:49:09Z
user: login: cdoern
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-16T18:06:02Z
user: login: RobotSail
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-20T23:42:13Z
user: login: RobotSail
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-20T23:42:27Z

Review comments:
diff_hunk: @@ -566,15 +585,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()
path: src/instructlab/training/main_ds.py
user: login: RobotSail
body: @booxter What is the reason to do it this way? It seems like this would be much simpler: 
```suggestion
# set timeout to 10 mins as a default: https://pytorch.org/docs/stable/distributed.html 
timeout = os.getenv("INSTRUCTLAB_NCCL_TIMEOUT_MS", 600_000) 
torch.distributed.init_process_group("nccl", timeout=timeout)
```
created_at: 2025-05-03T23:00:53Z
author_association: MEMBER
diff_hunk: @@ -566,15 +585,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()
path: src/instructlab/training/main_ds.py
user: login: booxter
body: The timeout from pytorch may be changed by upstream, in which case we will start overriding the new default value.
created_at: 2025-05-05T16:27:14Z
author_association: CONTRIBUTOR
diff_hunk: @@ -566,15 +585,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()
path: src/instructlab/training/main_ds.py
user: login: RobotSail
body: @booxter I see, could you please add some comments here explaining this then? Otherwise people reading this may not understand why we're doing it this way.
created_at: 2025-05-05T22:03:33Z
author_association: MEMBER
diff_hunk: @@ -544,6 +545,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")
path: src/instructlab/training/main_ds.py
user: login: RobotSail
body: It might be worthwhile to also update the `README.md` to let people know that this env can be used to set the timeout
created_at: 2025-05-12T23:10:14Z
author_association: MEMBER
diff_hunk: @@ -544,6 +545,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")
path: src/instructlab/training/main_ds.py
user: login: booxter
body: It's already documented in README:

```
## Environment variables

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.
```
created_at: 2025-05-14T01:49:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -544,6 +545,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")
path: src/instructlab/training/main_ds.py
user: login: RobotSail
body: Ah I see, I must have missed that. Thank you! 
created_at: 2025-05-20T23:42:13Z
author_association: MEMBER

PR comments:
user: login: booxter
created_at: 2025-05-06T15:38:02Z
body: @JamesKunstle below is my understanding of how this works. I may be wrong since I'm new to the topic, but I'll try to link the relevant code for reference. Please double check me: it's important we understand how this works.

---

So, there are three separate entities - a timeout for process group, a NCCL monitoring thread and a NCCL watchdog. They are separate entities serving separate needs.

---

**The process group timeout** is what you configure when passing `timeout=` to `init_process_group`. Torch uses this timeout (or its default torch value if not passed) to monitor collectives making progress (each collective is allocated the specified time).

Each backend will implement it in some way, for example, NCCL will [assign](https://github.com/pytorch/pytorch/blob/f47bf38e302f3849a829f9d9e3faf56fb0e50db9/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp#L3622) the timeout to each work item.

This parameter is not backend specific. If we ever use a different backend (not `nccl`), the timeout should still be honored. (This is why I'm planning to [rename](https://github.com/instructlab/training/pull/534) the variable in a follow-up patch since currently it has a NCCL specific name - a better name would be e.g. `INSTRUCTLAB_PROCESS_GROUP_TIMEOUT_SEC`, not `INSTRUCTLAB_NCCL_TIMEOUT_SEC`.)

---

**Now, to watchdogs.**

These are backend specific. It's up to the backend to run a watchdog or some other mechanism to implement the PG timeout.

Each NCCL rank starts a separate native [thread](https://github.com/pytorch/pytorch/blob/f47bf38e302f3849a829f9d9e3faf56fb0e50db9/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp#L2273) running a peer watchdog. The watchdog thread will periodically [check](https://github.com/pytorch/pytorch/blob/f47bf38e302f3849a829f9d9e3faf56fb0e50db9/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp#L2339-L2362) on each worker to see if it failed or timed out, and report back. It has its own sleep timer between iterations ([just 100ms](https://github.com/pytorch/pytorch/blob/f47bf38e302f3849a829f9d9e3faf56fb0e50db9/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp#L486)).

---

**There's also a NCCL [monitoring thread](https://github.com/pytorch/pytorch/blob/f47bf38e302f3849a829f9d9e3faf56fb0e50db9/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp#L1695).** Also backend specific. This thread is running separate to a watchdog and monitors the watchdog itself. Specifically, it watches for its heartbeats. If a single heartbeat is not detected in `TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC` seconds ([8 mins by default](https://github.com/pytorch/pytorch/blob/f47bf38e302f3849a829f9d9e3faf56fb0e50db9/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp#L1009)), torch will fail the job because something really bad happened to the watchdog that it couldn't even loop through a busy-wait iteration in such a long time; side note: I find the default 8 minutes value very conservative).

In case you wonder, the monitoring thread is [enabled](https://github.com/pytorch/pytorch/blob/b06cbd49f1ef6343f4595bb967d83e77969ca3c8/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp#L1006) by default [since 2.3.0](https://github.com/pytorch/pytorch/commit/0379c11248d022faa19c2aace6c9f86a0cee5de2). It is controlled by `TORCH_NCCL_ENABLE_MONITORING` that you can use to disable the monitoring thread (why would you tho?)

**What's a heartbeat?** Just an [atomic](https://github.com/pytorch/pytorch/blob/f56bcd240856e1cabe44ea56f0404abaa5678e89/torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp#L1172) integer [increment](https://github.com/pytorch/pytorch/blob/f47bf38e302f3849a829f9d9e3faf56fb0e50db9/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp#L2288) on a shared variable that is executed by the watchdog thread and [watched](https://github.com/pytorch/pytorch/blob/f47bf38e302f3849a829f9d9e3faf56fb0e50db9/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp#L1810-L1819) by the monitoring thread.

Important to note: for a watchdog thread heartbeat to happen, currently running collectives don't have to make any progress: as long as the watchdog thread is alive, it will heartbeat. So tuning `TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC` up is not very helpful.

---

The watchdog may even be disabled with `TORCH_NCCL_BLOCKING_WAIT`, [in which case](https://github.com/pytorch/pytorch/blob/f47bf38e302f3849a829f9d9e3faf56fb0e50db9/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp#L1050) no separate watchdog thread is running (and since [the watchdog thread starts its own monitoring thread](https://github.com/pytorch/pytorch/blob/f47bf38e302f3849a829f9d9e3faf56fb0e50db9/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp#L1981-L1982), neither the latter). Even then, the PG timeout set for each work item should apply. I believe the blocking mode is only used for **debugging purposes**, to surface errors from collectives synchronously (for better error traces or to attach a debugger right where it fails).
author_association: CONTRIBUTOR
user: login: mergify[bot]
created_at: 2025-05-12T20:52:42Z
body: This pull request has merge conflicts that must be resolved before it can be
merged. @booxter please rebase it. https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork
author_association: CONTRIBUTOR
user: login: booxter
created_at: 2025-05-20T21:44:54Z
body: I believe comment from @RobotSail is addressed but please double-check.
author_association: CONTRIBUTOR