Update nccl error handling flag

daviswer · daviswer · commit c54c7e59c3c9 · 2024-10-10T14:52:58.000-04:00
Signed-off-by: Davis Wertheimer &lt;davis.wertheimer@ibm.com&gt;
diff --git a/fms_fsdp/utils/train_utils.py b/fms_fsdp/utils/train_utils.py
@@ -184,7 +184,7 @@ def setup():
 
 def setup_environ_flags():
     os.environ["TORCH_SHOW_CPP_STACKTRACES"] = str(1)
-    os.environ["NCCL_ASYNC_ERROR_HANDLING"] = str(1)
+    os.environ["TORCH_NCCL_ASYNC_ERROR_HANDLING"] = str(1)
 
 
 def get_mixed_precision_policy(cfg, rank):
