Use spawn for multiprocessing start method (meta-pytorch#3284)

Summary:

CUDA context initialization is not fork-safe. If a CUDA context is created in a parent process, and then the process is forked (using `os.fork()`), the child process may encounter errors or undefined behavior when using CUDA. This is because the CUDA driver and runtime are not designed to be safely duplicated via `fork()`. It's recommended to use `spawn` or `forkserver`.

Among the two, `forkserver` needs to be use carefully and specifically, it's recommended to call `multiprocessing.set_start_method('forkserver')` at the very start of the program, and the parent process also needs to avoid initializing the CUDA context. When upgrading APS to CUDA 12.8, we encountered a test failure, and the test is apparently initializing the CUDA context before starting up two children processes, and I suspect that caused the test to hang - [post](https://fb.workplace.com/groups/319878845696681/posts/1494595861558301).

It's hard to avoid initializing the CUDA context early in this test, because it checks the GPU count in the test method's decorator - [code](https://fburl.com/code/27naz2eg).  Among the `spawn` and `forkserver` start methods, `spawn` is less efficient but it's the most robust. Let's switch to that instead to avoid any potential undefined behaviors with CUDA 12.8 and multiprocessing.

Reviewed By: adamomainz, weifengpy

Differential Revision: D80305233

Author: zhuhan0

torchrec/distributed/test_utils/multi_process.py

@@ -107,9 +107,14 @@ def __init__(
     ) -> None:
         super().__init__(methodName)
 
+        # In CUDA 12.8 we're seeing hangs from using forkserver, so we're
+        # switching to spawn.
         # AMD's HIP runtime doesn't seem to work with forkserver; hipMalloc will fail
         # Therefore we use spawn for HIP runtime until AMD fixes the issue
-        self._mp_init_mode: str = mp_init_mode if torch.version.hip is None else "spawn"
+        if torch.version.cuda >= "12.8" or torch.version.hip is not None:
+            self._mp_init_mode: str = "spawn"
+        else:
+            self._mp_init_mode: str = mp_init_mode
         logging.info(f"Using {self._mp_init_mode} for multiprocessing")
 
     @seed_and_log