[Reland] Use torch.multiprocessing instead of multiprocessing module (#976)

Reland #970 due to ghstack
issues.

Author: fegin

torchtitan/components/checkpoint.py

@@ -6,19 +6,18 @@
 
 import enum
 import functools
-import multiprocessing as mp
 import os
 import queue
 import re
 import shutil
 import threading
 import time
-from multiprocessing import get_context
 from typing import Any, Dict, List, Optional, Union
 
 import torch
 import torch.distributed as dist
 import torch.distributed.checkpoint as dcp
+import torch.multiprocessing as mp
 import torch.nn as nn
 from torch.distributed._state_dict_utils import _copy_state_dict, _create_cpu_state_dict
 from torch.distributed.checkpoint.state_dict import (
@@ -294,7 +293,7 @@ def load_state_dict(state_dict):
             self.async_future = None
         elif async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
             self.async_mode = AsyncMode.ASYNC_WITH_PINNED_MEM
-            ctx = get_context("spawn")
+            ctx = mp.get_context("spawn")
             self.mp_queue_send = ctx.Queue()
             self.mp_queue_recv = ctx.Queue()
             self.mp = ctx.Process(