Update nccl error handling flag

daviswer · web-flow · commit 514ffd5c90f5 · 2024-10-10T11:26:21.000-04:00
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):
