Add fork start_method support for xmp.spawn (#4236)

* Add fork start_method support for xmp.spawn

* Fix spacing format

* Make start_method keyword only for run_multiprocess

Author: darisoy

torch_xla/distributed/xla_multiprocessing.py

@@ -386,7 +386,7 @@ def spawn(fn,
     return None.
   """
   if pjrt.using_pjrt():
-    return pjrt.spawn(fn, args)
+    return pjrt.spawn(fn, start_method, args)
 
   if not _is_xla_config():
     # If this is not an XLA setup, jump to normal multi-processing.

torch_xla/experimental/pjrt.py

@@ -192,14 +192,19 @@ def _thread_fn(device: torch.device):
 
 
 @requires_pjrt
-def _run_multiprocess(fn: Callable[..., R], *args, **kwargs) -> Dict[int, R]:
+def _run_multiprocess(fn: Callable[..., R],
+                      *args,
+                      start_method: str = 'spawn',
+                      **kwargs) -> Dict[int, R]:
   """Runs `fn` on all devices available to PjRt.
 
   Spawns one process per physical device (e.g. TPU chip).
 
   Args:
     fn: Function to run on all devices
     args: args to pass to `fn`
+    start_method: The Python `multiprocessing` process creation method.
+      Default: `spawn`
     kwargs: kwargs to pass to `fn`
 
   Returns:
@@ -213,7 +218,7 @@ def _run_multiprocess(fn: Callable[..., R], *args, **kwargs) -> Dict[int, R]:
 
   with concurrent.futures.ProcessPoolExecutor(
       max_workers=num_processes,
-      mp_context=torch.multiprocessing.get_context('spawn')) as executor:
+      mp_context=torch.multiprocessing.get_context(start_method)) as executor:
 
     mp_fn = functools.partial(
         _run_thread_per_device,
@@ -239,15 +244,17 @@ def __call__(self) -> None:
     self.fn(global_ordinal(), *self.args, **self.kwargs)
 
 
-def spawn(fn: Callable, args: Tuple = ()) -> None:
+def spawn(fn: Callable, start_method: str = 'spawn', args: Tuple = ()) -> None:
   """Run functions compatible with xmp.spawn.
 
   Args:
     fn: Callable that takes the process index as the first argument.
     args: args to pass to `fn`
+    start_method: The Python `multiprocessing` process creation method.
+      Default: `spawn`
   """
   spawn_fn = _SpawnFn(fn, *args)
-  _run_multiprocess(spawn_fn)
+  _run_multiprocess(spawn_fn, start_method=start_method)
 
 
 def broadcast_master_param(model: nn.Module) -> None: