[v1] add accelerator (#9607)

src/llamafactory/v1/accelerator/helper.py

@@ -1,4 +1,7 @@
-# Copyright 2025 the LlamaFactory team.
+# Copyright 2025 Bytedance Ltd. and the LlamaFactory team.
+#
+# This code is inspired by the Bytedance's VeOmni library.
+# https://github.com/ByteDance-Seed/VeOmni/blob/v0.1.4/veomni/utils/dist_utils.py
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,12 +15,68 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+from contextlib import contextmanager
+from enum import Enum, unique
 from functools import lru_cache
+from typing import TYPE_CHECKING, Optional
 
+import numpy as np
 import torch
+import torch.distributed as dist
+
+from ..utils.types import Tensor, TensorLike
 
 
-def get_current_accelerator(check_available: bool = True):
+if TYPE_CHECKING:
+    from torch.distributed import ProcessGroup
+
+
+@unique
+class DeviceType(str, Enum):
+    CPU = "cpu"
+    CUDA = "cuda"
+    META = "meta"
+    MPS = "mps"
+    NPU = "npu"
+    XPU = "xpu"
+
+
+@unique
+class ReduceOp(str, Enum):
+    SUM = "sum"
+    MEAN = "mean"
+    MAX = "max"
+    MIN = "min"
+
+
+def is_distributed() -> bool:
+    """Check if distributed environment is available."""
+    return os.getenv("RANK") is not None
+
+
+def get_rank() -> int:
+    """Get rank."""
+    return int(os.getenv("RANK", "0"))
+
+
+def get_local_rank() -> int:
+    """Get local rank."""
+    return int(os.getenv("LOCAL_RANK", "0"))
+
+
+def get_world_size() -> int:
+    """Get world size."""
+    return int(os.getenv("WORLD_SIZE", "1"))
+
+
+def get_local_world_size() -> int:
+    """Get local world size."""
+    return int(os.getenv("LOCAL_WORLD_SIZE", "1"))
+
+
+@lru_cache
+def get_current_accelerator(check_available: bool = True) -> torch.device:
     """Get current accelerator.
 
     Note: this api requires torch>=2.7.0, 2.6 or lower will get an AttributeError or RuntimeError
@@ -27,26 +86,78 @@ def get_current_accelerator(check_available: bool = True):
 
     accelerator = torch.accelerator.current_accelerator(check_available=check_available)
     if accelerator is None:
-        return torch.device("cpu")
+        return torch.device(DeviceType.CPU.value)
 
     return accelerator
 
 
-@lru_cache
-def is_torch_npu_available():
-    return get_current_accelerator().type == "npu"
-
-
-@lru_cache
 def is_torch_cuda_available():
-    return get_current_accelerator().type == "cuda"
+    return get_current_accelerator().type == DeviceType.CUDA
 
 
-@lru_cache
-def is_torch_xpu_available():
-    return get_current_accelerator().type == "xpu"
-
-
-@lru_cache
 def is_torch_mps_available():
-    return get_current_accelerator().type == "mps"
+    return get_current_accelerator().type == DeviceType.MPS
+
+
+def is_torch_npu_available():
+    return get_current_accelerator().type == DeviceType.NPU
+
+
+def is_torch_xpu_available():
+    return get_current_accelerator().type == DeviceType.XPU
+
+
+def all_gather(tensor: Tensor, group: Optional["ProcessGroup"] = None) -> Tensor:
+    """Gathers the tensor from all ranks and concats them along the first dim."""
+    world_size = get_world_size()
+    device = get_current_accelerator()
+    output_tensor = torch.empty(world_size * tensor.numel(), dtype=tensor.dtype, device=device)
+    dist.all_gather_into_tensor(output_tensor, tensor, group=group)
+    return output_tensor.view(-1, *tensor.size()[1:])
+
+
+def all_reduce(data: TensorLike, op: ReduceOp = ReduceOp.MEAN, group: Optional["ProcessGroup"] = None) -> TensorLike:
+    """Performs all reduce in the given process group."""
+    device = get_current_accelerator()
+    is_ndarray = isinstance(data, np.ndarray)
+    is_tensor = isinstance(data, torch.Tensor)
+
+    if is_ndarray:
+        data = torch.from_numpy(data)
+    elif not is_tensor:
+        data = torch.tensor(data, dtype=torch.float, device=device)
+
+    reduce_ops = {
+        ReduceOp.MEAN: dist.ReduceOp.SUM,
+        ReduceOp.SUM: dist.ReduceOp.SUM,
+        ReduceOp.MAX: dist.ReduceOp.MAX,
+        ReduceOp.MIN: dist.ReduceOp.MIN,
+    }
+    dist.all_reduce(data, op=reduce_ops[op], group=group)
+    if op == ReduceOp.MEAN:  # ReduceOp.AVG is not supported by the NPU backend
+        data /= dist.get_world_size(group=group)
+
+    if is_tensor:
+        return data
+    elif is_ndarray:
+        return data.numpy()
+    elif data.numel() == 1:
+        return data.item()
+    else:
+        return data.tolist()
+
+
+@contextmanager
+def main_process_first(local_only: bool = True) -> None:
+    """A context manager for torch distributed environment to do something on the main process firstly."""
+    if get_world_size() > 1:
+        is_main_process = get_local_rank() == 0 if local_only else get_rank() == 0
+        try:
+            if not is_main_process:
+                dist.barrier()
+            yield
+        finally:
+            if is_main_process:
+                dist.barrier()
+    else:
+        yield