using spawn instead of fork for XPU device (#3884)

* using `spawn` instead of `fork` for XPU device

Signed-off-by: Liu, Kaixuan <kaixuan.liu@intel.com>

* update code

Signed-off-by: Liu, Kaixuan <kaixuan.liu@intel.com>

* remove comment

Signed-off-by: Liu, Kaixuan <kaixuan.liu@intel.com>

---------

Signed-off-by: Liu, Kaixuan <kaixuan.liu@intel.com>

Author: kaixuanliu

src/accelerate/launchers.py

@@ -206,10 +206,12 @@ def train(*args):
                 # First dummy launch
                 # Determine device type without initializing any device (which would break fork)
                 device_type, distributed_type = get_current_device_type()
+                # XPU requires spawn instead of fork
+                start_method = "spawn" if device_type == "xpu" else "fork"
                 if os.environ.get("ACCELERATE_DEBUG_MODE", "false").lower() == "true":
                     launcher = PrepareForLaunch(test_launch, distributed_type=distributed_type)
                     try:
-                        start_processes(launcher, args=(), nprocs=num_processes, start_method="fork")
+                        start_processes(launcher, args=(), nprocs=num_processes, start_method=start_method)
                     except ProcessRaisedException as e:
                         err = "An issue was found when verifying a stable environment for the notebook launcher."
                         if f"Cannot re-initialize {device_type.upper()} in forked subprocess" in e.args[0]:
@@ -241,7 +243,7 @@ def train(*args):
                         rdzv_configs=rdzv_conf,
                         max_restarts=max_restarts,
                         monitor_interval=monitor_interval,
-                        start_method="fork",
+                        start_method=start_method,
                     )
                     if is_torch_version(">=", ELASTIC_LOG_LINE_PREFIX_TEMPLATE_PYTORCH_VERSION):
                         launch_config_kwargs["log_line_prefix_template"] = log_line_prefix_template

src/accelerate/test_utils/scripts/test_notebook.py

@@ -17,22 +17,21 @@
 
 import os
 import time
-from multiprocessing import Queue
 
 from pytest import mark, raises
 from torch.distributed.elastic.multiprocessing.errors import ChildFailedError
 
 from accelerate import PartialState, notebook_launcher
 from accelerate.test_utils import require_bnb
-from accelerate.utils import is_bnb_available
+from accelerate.utils import is_bnb_available, is_xpu_available
 
 
 def basic_function():
     # Just prints the PartialState
     print(f"PartialState:\n{PartialState()}")
 
 
-def tough_nut_function(queue: Queue):
+def tough_nut_function(queue):
     if queue.empty():
         return
     trial = queue.get()
@@ -70,7 +69,15 @@ def test_c10d_rdzv_backend():
 
 @mark.skipif(NUM_PROCESSES < 2, reason="Need at least 2 processes to test fault tolerance")
 def test_fault_tolerant(max_restarts: int = 3):
-    queue = Queue()
+    # Use torch.multiprocessing to get the right context for the current device
+    import torch.multiprocessing as mp
+
+    # Get appropriate context - 'spawn' for XPU, 'fork' for others
+    if is_xpu_available():
+        ctx = mp.get_context("spawn")
+    else:
+        ctx = mp.get_context("fork")
+    queue = ctx.Queue()
     queue.put(max_restarts)
     notebook_launcher(tough_nut_function, (queue,), num_processes=NUM_PROCESSES, max_restarts=max_restarts)