Auto3DSeg: NCCL timeout error

[pwrightkcl]

For the auto3dseg experts:

I'm getting an error about NCCL timeouts, which some searching suggests is due to loading a large dataset. I've found potential solutions online but I'm hoping for some advice on which is correct how to implement in autoseg3d.

Here's the error:

[E ProcessGroupNCCL.cpp:457] [Rank 3] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=825, OpType=ALLREDUCE, NumelIn=1, NumelOut=1, Timeout(ms)=1800000) ran for 2209706 milliseconds before timing out.
[E ProcessGroupNCCL.cpp:457] [Rank 2] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=825, OpType=ALLREDUCE, NumelIn=1, NumelOut=1, Timeout(ms)=1800000) ran for 2209707 milliseconds before timing out.
[E ProcessGroupNCCL.cpp:457] [Rank 1] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=825, OpType=ALLREDUCE, NumelIn=1, NumelOut=1, Timeout(ms)=1800000) ran for 2209706 milliseconds before timing out.
[E ProcessGroupNCCL.cpp:457] [Rank 0] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=825, OpType=ALLREDUCE, NumelIn=1, NumelOut=1, Timeout(ms)=1800000) ran for 2209706 milliseconds before timing out.
[E ProcessGroupNCCL.cpp:471] Some NCCL operations have failed or timed out. Due to the asynchronous nature of CUDA kernels, subsequent GPU operations might run on corrupted/incomplete data.
[E ProcessGroupNCCL.cpp:477] To avoid data inconsistency, we are taking the entire process down.
[E ProcessGroupNCCL.cpp:842] [Rank 0] NCCL watchdog thread terminated with exception: [Rank 0] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=825, OpType=ALLREDUCE, NumelIn=1, NumelOut=1, Timeout(ms)=1800000) ran for 2209706 milliseconds before timing out.

In this SO thread the main answer suggests setting the SHM or export NCCL_P2P_LEVEL=NVL, which I can test. The second answer talks about setting the timeout in a Python line initialising a torch process group (which seems likely as I have large datasets). I have no idea how to drill down to that in autoseg3d, so could use some advice on how to do that or that I should leave it alone.

Many thanks!