refactor(parallel): 重构parallel模块

khaosz/parallel/__init__.py

@@ -1,12 +1,10 @@
-from khaosz.parallel.utils import (
+from khaosz.parallel.setup import (
     get_world_size, 
-    get_rank, 
-    get_device_count, 
-    get_current_device, 
-    get_available_backend, 
-    setup_parallel, 
+    get_rank,
+    get_current_device,
+    
     only_on_rank,
-    run_on_rank,
+    setup_parallel, 
     spawn_parallel_fn
 )
 
@@ -18,12 +16,10 @@ from khaosz.parallel.module import (
 __all__ = [
     "get_world_size",
     "get_rank",
-    "get_device_count",
     "get_current_device",
-    "get_available_backend",
-    "setup_parallel",
+    
     "only_on_rank",
-    "run_on_rank",
+    "setup_parallel",
     "spawn_parallel_fn",
     
     "RowParallelLinear",

khaosz/parallel/setup.py

@@ -3,38 +3,11 @@ import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
 
+from typing import Callable
 from functools import wraps
 from contextlib import contextmanager
 
 
-def get_device_count() -> int:
-    if torch.cuda.is_available():
-        return torch.cuda.device_count()
-    elif hasattr(torch, 'xpu') and torch.xpu.is_available():
-        return torch.xpu.device_count()
-    elif hasattr(torch, 'mps') and torch.mps.is_available():
-        return 1
-    else:
-        return 1
-
-def get_current_device() -> torch.device:
-    if torch.cuda.is_available():
-        return torch.device(f"cuda:{torch.cuda.current_device()}")
-    elif hasattr(torch, 'xpu') and torch.xpu.is_available():
-        return torch.device(f"xpu:{torch.xpu.current_device()}")
-    elif hasattr(torch, 'mps') and torch.mps.is_available():
-        return torch.device("mps")
-    else:
-        return torch.device("cpu")
-
-def get_available_backend():
-    if torch.cuda.is_available():
-        return "nccl"
-    elif hasattr(torch, 'xpu') and torch.xpu.is_available():
-        return "ccl"  # Intel XPU use ccl
-    else:
-        return "gloo"
-
 def get_world_size() -> int:
     if dist.is_available() and dist.is_initialized():
         return dist.get_world_size()
@@ -47,10 +20,21 @@ def get_rank() -> int:
     else:
         return 0
 
+def get_current_device():
+    if torch.cuda.is_available():
+        return torch.device(f"cuda:{torch.cuda.current_device()}")
+    elif hasattr(torch, 'xpu') and torch.xpu.is_available():
+        return torch.device(f"xpu:{torch.xpu.current_device()}")
+    elif hasattr(torch, 'mps') and torch.mps.is_available():
+        return torch.device("mps")
+    else:
+        return torch.device("cpu")
+
 @contextmanager
 def setup_parallel(
-    rank: int = 0, 
-    world_size: int = 1,
+    rank: int, 
+    world_size: int,
+    backend: str = "nccl",
     master_addr: str = "localhost",
     master_port: str = "29500"
 ):
@@ -69,11 +53,9 @@ def setup_parallel(
     os.environ['WORLD_SIZE'] = str(world_size)
     os.environ['LOCAL_RANK'] = str(rank)
     
-    backend = get_available_backend()
-    
     dist.init_process_group(
         backend=backend,
-        init_method="env://",
+        init_method=f"tcp://{master_addr}:{master_port}",
         rank=rank,
         world_size=world_size
     )
@@ -89,24 +71,7 @@ def setup_parallel(
         if dist.is_initialized():
             dist.destroy_process_group()
 
-@contextmanager
-def run_on_rank(rank=0, sync_before=True, sync_after=True):
-    """
-    context manager to run a function only on a specific rank.
-    """
-    is_main_proc = (get_rank() == rank)
-    
-    if dist.is_initialized() and sync_before:
-        dist.barrier()
-    
-    try:
-        yield is_main_proc
-        
-    finally:
-        if dist.is_initialized() and sync_after:
-            dist.barrier()
-
-def only_on_rank(rank=0):
+def only_on_rank(rank, sync=False):
     """
     decorator to run a function only on a specific rank.
     """
@@ -116,36 +81,31 @@ def only_on_rank(rank=0):
         def wrapper(*args, **kwargs):
             if get_rank() == rank:
                 return func(*args, **kwargs)
-            else: 
-                return None
+            if sync:
+                dist.barrier()
         
         return wrapper
     
     return decorator
 
-def wrapper_spawn_func(rank, world_size, func, kwargs_dict):
-    with setup_parallel(rank, world_size):
-        func(**kwargs_dict)
+def wrapper_spawn_func(rank, world_size, backend, func, kwargs):
+    with setup_parallel(rank, world_size, backend):
+        func(**kwargs)
 
-def spawn_parallel_fn(func, world_size=None, **kwargs):
- 
-    if world_size is None:
-        world_size = get_device_count()
-    
-    if world_size < 1:
-        raise ValueError("world_size must be greater than 0")
-    
-    device_count = get_device_count()
-    if world_size > device_count:
-        raise ValueError(f"world_size ({world_size}) exceeds available devices ({device_count})")
+def spawn_parallel_fn(func: Callable, world_size: int, backend: str, **kwargs):
     
     if world_size == 1:
         func(**kwargs)
         return
+
+    # clear environment variables
+    for key in ['MASTER_ADDR', 'MASTER_PORT', 'RANK', 'WORLD_SIZE', 'LOCAL_RANK']:
+        if key in os.environ:
+            del os.environ[key]
     
     mp.spawn(
         wrapper_spawn_func,
         nprocs=world_size,
-        args=(world_size, func, kwargs),  
+        args=(world_size, backend, func, kwargs),  
         join=True
     )