Enable NCCL_ASYNC_ERROR_HANDLING in torchelastic (#133)

Summary:
Pull Request resolved: #133

NCCL Async Error Handling is a new mechanism implemented in ProcessGroupNCCL to provide reliability for DDP training runs using NCCL. See here for a more detailed background and implementation details: pytorch/pytorch#46874.

At a high-level, this system was designed to ensure desynchronization, high GPU utilization, and NCCL errors don't cause indefinite hanging in distributed training runs. This system catches these errors without any perf impact and brings down the training process, and torchelastic can detect this and restart training from the previous checkpoint. The time after which stuck collectives are detected can be tuned using the `timeout` argument to `init_process_group`.

Reviewed By: kiukchung, jiayisuse

Differential Revision: D23610237

fbshipit-source-id: 7a2a496c0b781b68d76e138bd66ca0b7c04f17d0

torchelastic/agent/server/local_elastic_agent.py

@@ -144,6 +144,7 @@ def _start_workers(self, worker_group: WorkerGroup) -> Dict[int, Any]:
                 "TORCHELASTIC_RESTART_COUNT": str(restart_count),
                 "TORCHELASTIC_MAX_RESTARTS": str(spec.max_restarts),
                 "TORCHELASTIC_RUN_ID": spec.rdzv_handler.get_run_id(),
+                "NCCL_ASYNC_ERROR_HANDLING": str(1),
             }
             if "OMP_NUM_THREADS" in os.environ:
                 worker_env["OMP_NUM_THREADS"] = os.environ["OMP_NUM_THREADS"]