feat(CRITICAL): 从 GitHub 扫描搬运 ixformer SDK + xllm 完整 GDN/MoE 代码
来源:
1. Chranos/ixformer (GitHub) → ixformer_sdk/ (230 files, 70K lines)
- inference/functions/vllm.py: vllm_moe_topk_softmax 完整实现 (2033 lines)
- inference/functions/moe.py: MoE ops 完整实现 (1380 lines)
- contrib/vllm_flash_attn/: FA2 Python 接口 (1018 lines)
- contrib/tgi/fused_moe.py: TGI fused MoE (429 lines)
- csrc/include/ixformer/: C++ kernel headers + cmake
2. Deep-Spark/xllm (GitHub) → upstream_ref/xllm_latest/ (+15 files)
- npu_torch/qwen3_5_decoder_layer_impl.cpp/.h
- npu_torch/qwen3_5_gated_delta_net.cpp/.h
- npu_torch/qwen3_next_*.cpp/.h (6 files)
- npu_torch/attention.cpp/.h + fused_moe.cpp/.h + CMakeLists.txt
- models/llm/qwen3_5.h + qwen3_5_mtp.h + qwen3_next.h
- models/vlm/qwen3_5.h
调用链完整性:
ixformer_sdk/inference/functions/vllm.py
→ ops.infer.moe_topk_softmax() (C++ 层)
→ 这就是 base 镜像 libixformer.so 里的实现
upstream_ref/xllm_latest/core/layers/ilu/fused_moe.cpp
→ ixformer::infer::topk_softmax() (直接 C++ 调用)
→ ixformer::infer::group_gemm() → 完整 7-step MoE pipeline
This commit is contained in:
project6-dev
1
ixformer_sdk/distributed/__init__.py
Normal file
1
ixformer_sdk/distributed/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
from ._distributed import *
|
||||
481
ixformer_sdk/distributed/_distributed.py
Normal file
481
ixformer_sdk/distributed/_distributed.py
Normal file
@@ -0,0 +1,481 @@
|
||||
import warnings
|
||||
from collections import defaultdict
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.distributed.distributed_c10d as c10d
|
||||
from ixformer._C import _distributed as cdist
|
||||
from ixformer._C._distributed import comm
|
||||
from ixformer._C._distributed.comm import (
|
||||
AllGatherAlgo,
|
||||
AllReduceAlgo,
|
||||
BroadcastAlgo,
|
||||
ReduceAlgo,
|
||||
ReduceOp,
|
||||
ReduceScatterAlgo,
|
||||
SendAlgo,
|
||||
)
|
||||
from ixformer.core.multi_level_cache import MultiLevelCache
|
||||
from torch import Tensor
|
||||
from torch.distributed import ProcessGroup
|
||||
|
||||
from ixformer.core import config
|
||||
|
||||
IxformerCommType = int
|
||||
RecvAlgo = SendAlgo
|
||||
|
||||
_GROUP_TO_IXFC_COMM_CACHE = MultiLevelCache()
|
||||
_IXFC_COMM_TO_GROUP_CACHE = MultiLevelCache()
|
||||
|
||||
|
||||
def get_store(group: dist.ProcessGroup = None) -> dist.Store:
|
||||
if group is None:
|
||||
group = c10d._get_default_group()
|
||||
|
||||
return c10d._pg_map[group][1]
|
||||
|
||||
|
||||
class StoreWrapper(cdist.comm.C10dStoreWrapper):
|
||||
_GROUP_COUNT = defaultdict(dict)
|
||||
|
||||
def __init__(self, group: ProcessGroup):
|
||||
super().__init__()
|
||||
self.store = get_store()
|
||||
|
||||
ranks = dist.get_process_group_ranks(group)
|
||||
|
||||
group_key = "_".join([str(r) for r in ranks])
|
||||
if group not in self._GROUP_COUNT[group_key]:
|
||||
self._GROUP_COUNT[group_key][group] = len(self._GROUP_COUNT[group_key])
|
||||
group_count = self._GROUP_COUNT[group_key][group]
|
||||
|
||||
self.prefix = f"gid_{group_count}_" + group_key
|
||||
|
||||
def _gen_unique_key(self, key):
|
||||
return f"{self.prefix}_{key}"
|
||||
|
||||
def set(self, key: str, value: str):
|
||||
key = self._gen_unique_key(key)
|
||||
self.store.set(key, value)
|
||||
|
||||
def get(self, key: str) -> str:
|
||||
key = self._gen_unique_key(key)
|
||||
self.store.wait([key])
|
||||
return self.store.get(key).decode("utf8")
|
||||
|
||||
|
||||
def init_comm_with_store(group=None, shmsize: int = None):
|
||||
if group is None:
|
||||
group = c10d._get_default_group()
|
||||
|
||||
world_size = dist.get_world_size(group=group)
|
||||
rank = dist.get_group_rank(group=group, global_rank=dist.get_rank())
|
||||
|
||||
if shmsize is None:
|
||||
shmsize = config.IXFORMER_COMM_SHM_SIZE
|
||||
|
||||
store_wrapper = StoreWrapper(group=group)
|
||||
ixfc_comm = cdist.comm.init_communicator_by_store(
|
||||
store=store_wrapper, world_size=world_size, rank=rank, max_shm_mem_size=shmsize
|
||||
)
|
||||
|
||||
_GROUP_TO_IXFC_COMM_CACHE.set(group, ixfc_comm)
|
||||
_IXFC_COMM_TO_GROUP_CACHE.set(ixfc_comm, group)
|
||||
return ixfc_comm
|
||||
|
||||
|
||||
_sub_store = None
|
||||
|
||||
|
||||
def create_nccl_unique_id(addr: str, port: str, world_size: int, rank: int):
|
||||
global _sub_store
|
||||
_sub_store = dist.TCPStore(
|
||||
host_name=addr, port=int(port), world_size=world_size, is_master=rank == 0
|
||||
)
|
||||
store_key = "ncclUniqueId"
|
||||
if rank == 0:
|
||||
commid = cdist.comm.create_nccl_unique_id()
|
||||
_sub_store.set(store_key, commid)
|
||||
else:
|
||||
_sub_store.wait([store_key])
|
||||
commid = _sub_store.get(store_key).decode("utf8")
|
||||
|
||||
return commid
|
||||
|
||||
|
||||
def init_comm_with_eth(
|
||||
addr: str, port: str, world_size: int, rank: int, shmsize: int = None
|
||||
):
|
||||
commid = create_nccl_unique_id(addr, port, world_size=world_size, rank=rank)
|
||||
return cdist.comm.init_communicator_by_nccl_id(commid, world_size, rank, shmsize)
|
||||
|
||||
|
||||
def _check_group(group: Optional[ProcessGroup] = None):
|
||||
if group is None:
|
||||
group = c10d._get_default_group()
|
||||
|
||||
if isinstance(group, ProcessGroup):
|
||||
ixfc_comm = _GROUP_TO_IXFC_COMM_CACHE.get(group, None)
|
||||
if ixfc_comm is None:
|
||||
return init_comm_with_store(group)
|
||||
return ixfc_comm
|
||||
|
||||
return group
|
||||
|
||||
|
||||
def get_comm_group_stream(group: Optional[ProcessGroup] = None):
|
||||
group = _check_group(group)
|
||||
return comm.get_comm_group_stream(group)
|
||||
|
||||
|
||||
def set_comm_group_stream(stream: int, group: Optional[ProcessGroup] = None):
|
||||
group = _check_group(group)
|
||||
return comm.set_comm_group_stream(group, stream)
|
||||
|
||||
|
||||
def get_group_rank(group: Optional[ProcessGroup], global_rank) -> int:
|
||||
"""将 global rank 映射到 group 中的相对 rank"""
|
||||
if isinstance(group, IxformerCommType):
|
||||
_pg = _IXFC_COMM_TO_GROUP_CACHE.get(group, None)
|
||||
if _pg is None:
|
||||
return global_rank
|
||||
else:
|
||||
group = _IXFC_COMM_TO_GROUP_CACHE.get(group)
|
||||
|
||||
if group is None:
|
||||
group = c10d._get_default_group()
|
||||
|
||||
return dist.get_group_rank(group, global_rank)
|
||||
|
||||
|
||||
def get_global_rank(group: Optional[ProcessGroup], group_rank: int) -> int:
|
||||
"""将一个 group rank 映射到 global rank"""
|
||||
if group is None:
|
||||
group = c10d._get_default_group()
|
||||
return c10d.get_global_rank(group, group_rank)
|
||||
|
||||
|
||||
def get_process_group_ranks(group: Optional[ProcessGroup] = None) -> List[int]:
|
||||
"""获取 Group 的 global ranks"""
|
||||
if group is None:
|
||||
group = c10d._get_default_group()
|
||||
return c10d.get_process_group_ranks(group)
|
||||
|
||||
|
||||
def new_group(ranks: List[int] = None, shmsize=None, *args, **kwargs):
|
||||
"""通过 global ranks 去创建一个通讯组"""
|
||||
group = c10d.new_group(ranks, *args, **kwargs)
|
||||
|
||||
if ranks is None:
|
||||
ranks = dist.get_process_group_ranks(group)
|
||||
|
||||
if get_rank() in ranks:
|
||||
init_comm_with_store(group=group, shmsize=shmsize)
|
||||
return group
|
||||
|
||||
|
||||
def new_subgroups_by_enumeration(
|
||||
ranks_per_subgroup_list, shmsize=None, *args, **kwargs
|
||||
) -> Tuple[ProcessGroup, List[ProcessGroup]]:
|
||||
"""
|
||||
通过一组 global ranks 去创建通讯组
|
||||
|
||||
:param ranks_per_subgroup_list: global ranks
|
||||
:return: 返回当前 rank 所在的通讯组 和 新的 subgroups
|
||||
"""
|
||||
self_group, other_group = c10d.new_subgroups_by_enumeration(
|
||||
ranks_per_subgroup_list, *args, **kwargs
|
||||
)
|
||||
init_comm_with_store(self_group, shmsize=shmsize)
|
||||
return self_group, other_group
|
||||
|
||||
|
||||
def destroy_process_group(group: Optional[ProcessGroup] = None):
|
||||
"""销毁 Group"""
|
||||
if group is None:
|
||||
group = c10d._get_default_group()
|
||||
ixfc_comm = _GROUP_TO_IXFC_COMM_CACHE.get(group, None)
|
||||
|
||||
if ixfc_comm is None:
|
||||
dist.destroy_process_group(group)
|
||||
else:
|
||||
comm.destroy(ixfc_comm)
|
||||
dist.destroy_process_group(group)
|
||||
|
||||
|
||||
def get_rank(group: Optional[ProcessGroup] = None) -> int:
|
||||
"""获取当前进程的 Rank，如果 group 是 null，那么返回的是 Global Rank, 否则返回的相对的 Rank，即在当前组中的 rank"""
|
||||
return c10d.get_rank(group)
|
||||
|
||||
|
||||
def get_world_size(group: Optional[ProcessGroup] = None) -> int:
|
||||
"""获取 Group 中的成员大小"""
|
||||
return c10d.get_world_size(group)
|
||||
|
||||
|
||||
def barrier(group: Optional[ProcessGroup] = None, use_comm_stream: bool = False):
|
||||
"""同步 Group 中的 rank"""
|
||||
group = _check_group(group)
|
||||
comm.barrier(group, use_comm_stream)
|
||||
|
||||
|
||||
def isend(
|
||||
tensor: Tensor,
|
||||
dst: int,
|
||||
group: Optional[ProcessGroup] = None,
|
||||
use_comm_stream: bool = False,
|
||||
):
|
||||
dst = get_group_rank(group, dst)
|
||||
group = _check_group(group)
|
||||
return comm.send(group, tensor, dst, use_comm_stream, SendAlgo.kNone)
|
||||
|
||||
|
||||
def send(*args, **kwargs):
|
||||
warnings.warn("not support sync mode, as async to call.")
|
||||
return isend(*args, **kwargs)
|
||||
|
||||
|
||||
def irecv(
|
||||
tensor: torch.Tensor,
|
||||
src: int,
|
||||
group: Optional[ProcessGroup] = None,
|
||||
use_comm_stream: bool = False,
|
||||
):
|
||||
src = get_group_rank(group, src)
|
||||
group = _check_group(group)
|
||||
return comm.recv(group, tensor, src, use_comm_stream, SendAlgo.kNone)
|
||||
|
||||
|
||||
def recv(*args, **kwargs):
|
||||
warnings.warn("not support sync mode, as async to call.")
|
||||
return irecv(*args, **kwargs)
|
||||
|
||||
|
||||
def point_to_point(
|
||||
tensor: Tensor,
|
||||
src: int,
|
||||
dst: int,
|
||||
group: Optional[ProcessGroup] = None,
|
||||
use_comm_stream: bool = False,
|
||||
):
|
||||
"""在 src rank 发送 tensor，在 dst_rank 上接收数据到 tensor 中"""
|
||||
src = get_group_rank(group, src)
|
||||
dst = get_group_rank(group, dst)
|
||||
group = _check_group(group)
|
||||
return comm.p2p(group, tensor, src, dst, use_comm_stream)
|
||||
|
||||
|
||||
def reduce(
|
||||
tensor,
|
||||
root: int,
|
||||
op=ReduceOp.SUM,
|
||||
group: Optional[ProcessGroup] = None,
|
||||
async_op=False,
|
||||
out: Tensor = None,
|
||||
use_comm_stream: bool = False,
|
||||
):
|
||||
"""
|
||||
Example:
|
||||
ixf_tensor = torch.tensor([1], device="cuda")
|
||||
ixfd.reduce(ixf_tensor, 1, async_op=True)
|
||||
print("rank {rank}:", ixf_tensor)
|
||||
|
||||
# output
|
||||
rank 0: tensor([1], device='cuda:0')
|
||||
rank 1: tensor([4], device='cuda:1')
|
||||
rank 2: tensor([1], device='cuda:2')
|
||||
rank 3: tensor([1], device='cuda:3')
|
||||
"""
|
||||
|
||||
if not async_op:
|
||||
raise RuntimeError("Not support sync operation now.")
|
||||
|
||||
if out is None:
|
||||
out = tensor
|
||||
|
||||
root = get_group_rank(group, root)
|
||||
group = _check_group(group)
|
||||
return comm.reduce(group, tensor, out, op, root, use_comm_stream, ReduceAlgo.kNone)
|
||||
|
||||
|
||||
def broadcast(
|
||||
tensor: Tensor,
|
||||
src: int,
|
||||
group: Optional[ProcessGroup] = None,
|
||||
async_op=False,
|
||||
out: Tensor = None,
|
||||
use_comm_stream: bool = False,
|
||||
):
|
||||
"""
|
||||
Example:
|
||||
ixf_tensor = torch.tensor([rank], device="cuda")
|
||||
ixfd.broadcast(ixf_tensor, 1, async_op=True)
|
||||
print("rank {rank}: ", ixf_tensor)
|
||||
|
||||
# output
|
||||
rank 0: tensor([1], device='cuda:0')
|
||||
rank 1: tensor([1], device='cuda:1')
|
||||
rank 2: tensor([1], device='cuda:2')
|
||||
rank 3: tensor([1], device='cuda:3')
|
||||
"""
|
||||
if not async_op:
|
||||
raise RuntimeError("Not support sync operation now.")
|
||||
|
||||
if out is None:
|
||||
out = tensor
|
||||
|
||||
src = get_group_rank(group, src)
|
||||
group = _check_group(group)
|
||||
return comm.broadcast(group, tensor, out, src, use_comm_stream, BroadcastAlgo.kNone)
|
||||
|
||||
|
||||
def reduce_scatter_tensor(
|
||||
output: Tensor,
|
||||
input: Tensor,
|
||||
op=ReduceOp.SUM,
|
||||
group: Optional[ProcessGroup] = None,
|
||||
async_op=False,
|
||||
use_comm_stream: bool = False,
|
||||
):
|
||||
"""
|
||||
Example:
|
||||
ixf_tensor_out = torch.zeros(2, dtype=torch.int64, device="cuda")
|
||||
tensor_in = torch.arange(world_size * 2, dtype=torch.int64, device="cuda")
|
||||
# tensor_in: tensor([0, 1, 2, 3, 4, 5, 6, 7], device='cuda:0')
|
||||
|
||||
ixfd.reduce_scatter_tensor(ixf_tensor_out, tensor_in, async_op=True)
|
||||
print("rank {rank}:", ixf_tensor_out)
|
||||
|
||||
# output
|
||||
rank 0: tensor([0, 4], device='cuda:0')
|
||||
rank 1: tensor([ 8, 12], device='cuda:1')
|
||||
rank 2: tensor([16, 20], device='cuda:2')
|
||||
rank 3: tensor([24, 28], device='cuda:3')
|
||||
"""
|
||||
if not async_op:
|
||||
raise RuntimeError("Not support sync operation now.")
|
||||
|
||||
group = _check_group(group)
|
||||
return comm.reduce_scatter(
|
||||
group, input, output, op, use_comm_stream, ReduceScatterAlgo.kNone
|
||||
)
|
||||
|
||||
|
||||
def all_reduce(
|
||||
tensor: Tensor,
|
||||
op=ReduceOp.SUM,
|
||||
group: Optional[ProcessGroup] = None,
|
||||
async_op=False,
|
||||
out: Tensor = None,
|
||||
algo: AllReduceAlgo = AllReduceAlgo.kNone,
|
||||
use_comm_stream: bool = False,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
tensor: inpute tensor
|
||||
op: ReduceOp: SUM, MIN or MAX
|
||||
group: communicator group
|
||||
async_op: ixformer support async mode
|
||||
out: output tensor
|
||||
algo: AllReduce Algo: Auto, Quant, QuantL1, QuantL2, NCCL, Ring, AllGatherSum, BroadcastSum
|
||||
use_comm_stream: ixformer support set communication stream by ixformer.distributed.set_comm_group_stream,
|
||||
if true, submit the kernels of communication to communication stream,
|
||||
if false, use current stream by torch.cuda.current_stream
|
||||
Returns: out
|
||||
|
||||
Example:
|
||||
>>> # All tensors below are of torch.int64 type.
|
||||
>>> # We have 2 process groups, 2 ranks.
|
||||
>>> tensor = torch.arange(2, dtype=torch.int64) + 1 + 2 * rank
|
||||
>>> tensor
|
||||
tensor([1, 2]) # Rank 0
|
||||
tensor([3, 4]) # Rank 1
|
||||
>>> ixfd.all_reduce(tensor, op=ReduceOp.SUM, async_op=True)
|
||||
>>> tensor
|
||||
tensor([4, 6]) # Rank 0
|
||||
tensor([4, 6]) # Rank 1
|
||||
"""
|
||||
if not async_op:
|
||||
raise RuntimeError("Not support sync operation now.")
|
||||
|
||||
group = _check_group(group)
|
||||
|
||||
if out is None:
|
||||
out = tensor
|
||||
|
||||
comm.all_reduce(
|
||||
group,
|
||||
tensor,
|
||||
out,
|
||||
op,
|
||||
use_comm_stream=use_comm_stream,
|
||||
algo=algo,
|
||||
)
|
||||
|
||||
|
||||
def all_gather_into_tensor(
|
||||
output: Tensor,
|
||||
input: Tensor,
|
||||
group: Optional[ProcessGroup] = None,
|
||||
async_op=False,
|
||||
use_comm_stream: bool = False,
|
||||
):
|
||||
"""
|
||||
Example:
|
||||
tensor_in = torch.arange(2, dtype=torch.int64, device="cuda") + 1 + 2 * rank
|
||||
rank 0: tensor in: tensor([1, 2], device='cuda:0')
|
||||
rank 1: tensor in: tensor([3, 4], device='cuda:1')
|
||||
rank 2: tensor in: tensor([5, 6], device='cuda:2')
|
||||
rank 3: tensor in: tensor([7, 8], device='cuda:3')
|
||||
|
||||
ixf_tensor_out = torch.zeros(world_size * 2, dtype=torch.int64, device="cuda")
|
||||
ixfd.all_gather_into_tensor(ixf_tensor_out, tensor_in, async_op=True)
|
||||
print("rank {rank}:", ixf_tensor_out)
|
||||
|
||||
# output:
|
||||
rank 0: tensor([1, 2, 3, 4, 5, 6, 7, 8], device='cuda:0')
|
||||
rank 1: tensor([1, 2, 3, 4, 5, 6, 7, 8], device='cuda:1')
|
||||
rank 2: tensor([1, 2, 3, 4, 5, 6, 7, 8], device='cuda:2')
|
||||
rank 3: tensor([1, 2, 3, 4, 5, 6, 7, 8], device='cuda:3')
|
||||
"""
|
||||
if not async_op:
|
||||
raise RuntimeError("Not support sync operation now.")
|
||||
|
||||
group = _check_group(group)
|
||||
return comm.all_gather(
|
||||
group, input, output, use_comm_stream, algo=AllGatherAlgo.kNone
|
||||
)
|
||||
|
||||
|
||||
def gather(
|
||||
tensor,
|
||||
gather_list=None,
|
||||
dst=0,
|
||||
group: Optional[ProcessGroup] = None,
|
||||
async_op=False,
|
||||
use_comm_stream: bool = False,
|
||||
):
|
||||
"""
|
||||
Example:
|
||||
>>> # We have 2 process groups, 2 ranks.
|
||||
>>> tensor = torch.tensor(rank+1,dtype=torch.float32).cuda()
|
||||
>>> tensor
|
||||
tensor(1.) # Rank 0
|
||||
tensor(2.) # Rank 1
|
||||
>>> gather_list = [torch.zeros(1).cuda() for _ in range(rank)] if rank == dst else None
|
||||
>>> gather_list
|
||||
[tensor([0,]),tensor([1,])] # Rank 0
|
||||
None # Rank 1
|
||||
ixfd.gather(tensor,gather_list,0,async_op=True)
|
||||
>>> gather_list
|
||||
[tensor([1.]),tensor([2.])] # Rank 0
|
||||
None # Rank 1
|
||||
"""
|
||||
gather_list = gather_list if gather_list is not None else []
|
||||
if not async_op:
|
||||
raise RuntimeError("Not support sync operation now.")
|
||||
|
||||
dst = get_group_rank(group, dst)
|
||||
group = _check_group(group)
|
||||
return comm.gather(group, tensor, gather_list, dst, use_comm_stream)
|
||||
412
ixformer_sdk/distributed/overlap_comm.py
Normal file
412
ixformer_sdk/distributed/overlap_comm.py
Normal file
@@ -0,0 +1,412 @@
|
||||
import abc
|
||||
import enum
|
||||
import os
|
||||
from contextlib import contextmanager, nullcontext
|
||||
from typing import List, Optional
|
||||
|
||||
import torch.cuda
|
||||
from ixformer.core.dispatcher import Dispatcher
|
||||
|
||||
from ixformer.core import config
|
||||
|
||||
from . import _distributed as ixfd
|
||||
|
||||
|
||||
class SplitOverlapComm(Dispatcher):
|
||||
def __init__(self, num_chunks, num_compute_streams=None, comm_group=None):
|
||||
"""
|
||||
Args:
|
||||
num_chunks: the number of chunks
|
||||
num_compute_streams: the number of compute streams, default: 1
|
||||
comm_group: communicator group
|
||||
"""
|
||||
|
||||
self._num_chunks = num_chunks
|
||||
self._num_compute_streams = num_compute_streams or 1
|
||||
self._comm_group = comm_group
|
||||
|
||||
self._compute_streams: List[torch.cuda.Stream] = self.create_compute_streams()
|
||||
self._comm_stream: torch.cuda.Stream = torch.cuda.Stream(priority=-1)
|
||||
|
||||
self._start_compute_event: torch.cuda.Event = torch.cuda.Event()
|
||||
self._stop_compute_event: torch.cuda.Event = torch.cuda.Event()
|
||||
|
||||
self._start_comm_event: torch.cuda.Event = torch.cuda.Event()
|
||||
self._stop_comm_event: torch.cuda.Event = torch.cuda.Event()
|
||||
|
||||
# keep origin state
|
||||
self._main_stream: Optional[torch.cuda.Stream] = None
|
||||
self._origin_ixf_comm_stream = None
|
||||
self._ixformer_streams = dict()
|
||||
|
||||
@classmethod
|
||||
def dispatcher_key(
|
||||
cls, num_chunks, num_compute_streams=None, comm_group=None, *args, **kwargs
|
||||
):
|
||||
"""
|
||||
the key of SplitOverlapComm
|
||||
Args:
|
||||
num_chunks: the number of chunks
|
||||
num_compute_streams: the number of compute streams, default: 1
|
||||
comm_group: communicator group
|
||||
Returns: unique key
|
||||
"""
|
||||
# warn: keey same function parameters with init
|
||||
return (cls.__name__, num_chunks, num_compute_streams, comm_group)
|
||||
|
||||
@classmethod
|
||||
def enable(cls):
|
||||
return config.IXFORMER_ENABLE_OVERLAP_COMM
|
||||
|
||||
@property
|
||||
def num_chunks(self):
|
||||
return self._num_chunks
|
||||
|
||||
@property
|
||||
def num_compute_streams(self):
|
||||
return self._num_compute_streams
|
||||
|
||||
@property
|
||||
def comm_group(self):
|
||||
return self._comm_group
|
||||
|
||||
def create_compute_streams(self):
|
||||
streams = []
|
||||
for _ in range(self.num_compute_streams):
|
||||
streams.append(torch.cuda.Stream())
|
||||
return streams
|
||||
|
||||
def start_overlap(self):
|
||||
self._main_stream = torch.cuda.current_stream()
|
||||
|
||||
self._start_compute_event.record(torch.cuda.current_stream())
|
||||
for compute_stream in self._compute_streams:
|
||||
compute_stream.wait_event(self._start_compute_event)
|
||||
|
||||
self._origin_ixf_comm_stream = ixfd.get_comm_group_stream(self._comm_group)
|
||||
ixfd.set_comm_group_stream(self._comm_stream.cuda_stream, self._comm_group)
|
||||
|
||||
def stop_overlap(self):
|
||||
last_compute_stream_id = (
|
||||
self.num_chunks + self.num_compute_streams - 1
|
||||
) % self.num_compute_streams
|
||||
self._stop_compute_event.record(self._compute_streams[last_compute_stream_id])
|
||||
self._stop_comm_event.record(self._comm_stream)
|
||||
torch.cuda.current_stream().wait_event(self._stop_compute_event)
|
||||
torch.cuda.current_stream().wait_event(self._stop_comm_event)
|
||||
|
||||
ixfd.set_comm_group_stream(self._origin_ixf_comm_stream, self._comm_group)
|
||||
|
||||
def start_comm(self, chunk_idx):
|
||||
"""
|
||||
prepare communication stream and wait event.
|
||||
Args:
|
||||
chunk_idx: the index of chunk
|
||||
"""
|
||||
|
||||
self._start_comm_event.record(
|
||||
self._compute_streams[chunk_idx % self.num_compute_streams]
|
||||
)
|
||||
self._comm_stream.wait_event(self._start_comm_event)
|
||||
|
||||
@contextmanager
|
||||
def compute_stream_context(self, chunk_idx):
|
||||
"""
|
||||
open python context and switch to compute stream in torch context
|
||||
Args:
|
||||
chunk_idx: the index of chunk
|
||||
"""
|
||||
|
||||
stream = self._compute_streams[chunk_idx % self.num_compute_streams]
|
||||
|
||||
# print("before stream:", torch.cuda.current_stream())
|
||||
torch.cuda.set_stream(stream)
|
||||
|
||||
# print("after stream:", torch.cuda.current_stream(), ixformer.cuda.current_stream())
|
||||
yield stream
|
||||
|
||||
torch.cuda.set_stream(self._main_stream)
|
||||
|
||||
@contextmanager
|
||||
def stream_context(self, stream):
|
||||
# print("before stream:", torch.cuda.current_stream())
|
||||
torch.cuda.set_stream(stream)
|
||||
|
||||
# print("after stream:", torch.cuda.current_stream(), ixformer.cuda.current_stream())
|
||||
yield stream
|
||||
|
||||
torch.cuda.set_stream(self._main_stream)
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
self.start_overlap()
|
||||
out = self.compute(*args, **kwargs)
|
||||
self.stop_overlap()
|
||||
return out
|
||||
|
||||
@abc.abstractmethod
|
||||
def compute(self, *args, **kwargs):
|
||||
"""
|
||||
it is abstract method to execute compute and communication.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class GemmMethod(enum.IntEnum):
|
||||
kCUINFER = 0
|
||||
kCUBLAS = 1
|
||||
kLIMITED_GEMM = 2
|
||||
|
||||
|
||||
class GemmWithLimitedBlock:
|
||||
def __init__(self, limit_algo=0) -> None:
|
||||
self.limit_algo = limit_algo
|
||||
self.env_key = "PYTORCH_GEMM_BLOCK_LIMITATION"
|
||||
|
||||
def __enter__(self) -> None:
|
||||
os.environ[self.env_key] = str(self.limit_algo)
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback) -> None:
|
||||
del os.environ[self.env_key]
|
||||
|
||||
|
||||
class IxFormerLimitedGemmContext:
|
||||
def __init__(self) -> None:
|
||||
self.env_key = "IXFORMER_ENABLE_PERSISTENT_GEMM"
|
||||
|
||||
def __enter__(self) -> None:
|
||||
os.environ[self.env_key] = "1"
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback) -> None:
|
||||
os.environ[self.env_key] = "0"
|
||||
|
||||
|
||||
class GemmAllReduceSplitOverlapComm(SplitOverlapComm):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
self.gemm_method_env = config.IXFORMER_OVERLAP_GEMM_METHOD
|
||||
|
||||
if self.gemm_method_env is None:
|
||||
if ixfd.get_world_size(self.comm_group) == 2:
|
||||
self.gemm_method_env = 0
|
||||
else:
|
||||
self.gemm_method_env = 2
|
||||
|
||||
self.gemm_method = GemmMethod(int(self.gemm_method_env))
|
||||
self.limited_gemm_ctx = GemmWithLimitedBlock()
|
||||
self.ixf_limited_gemm_ctx = IxFormerLimitedGemmContext()
|
||||
self.split_ratio = config.IXFORMER_OVERLAP_SPLIT_RATIO
|
||||
|
||||
@classmethod
|
||||
def compute_row_parallel_dims(cls, input):
|
||||
batch = 1
|
||||
if input.ndim == 2:
|
||||
seqlen = input.shape[0]
|
||||
else:
|
||||
batch = input.shape[0]
|
||||
seqlen = input.shape[1]
|
||||
|
||||
parallel_dims = batch * seqlen
|
||||
return parallel_dims
|
||||
|
||||
def compute(self, input, weight, bias=None, out=None, *args, **kwargs):
|
||||
"""
|
||||
:param input: [Batch, SeqLen, Hidden]
|
||||
:param weight: [OutChannel, InChannel]
|
||||
:param bias: [OutChannel]
|
||||
"""
|
||||
|
||||
is_update_shape = input.ndim > 2
|
||||
batch = 1
|
||||
if input.ndim == 2:
|
||||
seqlen = input.shape[0]
|
||||
else:
|
||||
batch = input.shape[0]
|
||||
seqlen = input.shape[1]
|
||||
|
||||
parallel_dims = batch * seqlen
|
||||
|
||||
if is_update_shape:
|
||||
input = input.reshape(parallel_dims, -1)
|
||||
|
||||
if out is None:
|
||||
out_shape = [parallel_dims, weight.shape[0]]
|
||||
out_dtype = kwargs["out_dtype"] if "out_dtype" in kwargs else input.dtype
|
||||
out = torch.empty(out_shape, dtype=out_dtype, device=input.device)
|
||||
|
||||
if self.split_ratio is not None:
|
||||
round_multiples = 256 if parallel_dims >= 256 else parallel_dims
|
||||
first_chunk_size = (
|
||||
round((parallel_dims * float(self.split_ratio)) / round_multiples)
|
||||
* round_multiples
|
||||
)
|
||||
middle_chunk_size = (parallel_dims - first_chunk_size) // (
|
||||
self.num_chunks - 1
|
||||
)
|
||||
middle_chunk_size = (middle_chunk_size // round_multiples) * round_multiples
|
||||
last_chunk_size = (
|
||||
parallel_dims
|
||||
- first_chunk_size
|
||||
- middle_chunk_size * (self.num_chunks - 2)
|
||||
)
|
||||
|
||||
chunk_sizes = (
|
||||
[first_chunk_size]
|
||||
+ [middle_chunk_size] * (self.num_chunks - 2)
|
||||
+ [last_chunk_size]
|
||||
)
|
||||
input_chunks = torch.split_with_sizes(input, chunk_sizes, dim=0)
|
||||
out_chunks = torch.split_with_sizes(out, chunk_sizes, dim=0)
|
||||
|
||||
# print(first_chunk_size, middle_chunk_size, last_chunk_size, chunk_sizes)
|
||||
else:
|
||||
input_chunks = torch.chunk(input, self.num_chunks, dim=0)
|
||||
out_chunks = torch.chunk(out, self.num_chunks, dim=0)
|
||||
|
||||
for chunk_idx in range(len(input_chunks)):
|
||||
with self.compute_stream_context(chunk_idx):
|
||||
chunk_out = self.gemm_dispatcher(
|
||||
chunk_idx,
|
||||
input_chunks[chunk_idx],
|
||||
weight,
|
||||
out_chunks[chunk_idx],
|
||||
*args,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
self.start_comm(chunk_idx)
|
||||
|
||||
ixfd.all_reduce(
|
||||
chunk_out, async_op=True, group=self.comm_group, use_comm_stream=True
|
||||
)
|
||||
|
||||
if is_update_shape:
|
||||
out = out.reshape(batch, seqlen, -1)
|
||||
|
||||
if bias is not None:
|
||||
out = out + bias
|
||||
|
||||
return out
|
||||
|
||||
def gemm_dispatcher(
|
||||
self,
|
||||
chunk_idx,
|
||||
chunk_input,
|
||||
weight,
|
||||
chunk_out=None,
|
||||
user_gemm_method=None,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
if user_gemm_method is not None and callable(user_gemm_method):
|
||||
ctx = nullcontext() if chunk_idx == 0 else self.ixf_limited_gemm_ctx
|
||||
with ctx:
|
||||
return user_gemm_method(
|
||||
chunk_input, weight, out=chunk_out, *args, **kwargs
|
||||
)
|
||||
|
||||
if user_gemm_method is None:
|
||||
user_gemm_method = self.gemm_method
|
||||
|
||||
if user_gemm_method == GemmMethod.kCUINFER:
|
||||
import ixformer.functions as ixff
|
||||
|
||||
return ixff.linear(chunk_input, weight, output=chunk_out)
|
||||
elif user_gemm_method == GemmMethod.kCUBLAS:
|
||||
return torch.matmul(chunk_input, weight.T, out=chunk_out)
|
||||
elif user_gemm_method == GemmMethod.kLIMITED_GEMM:
|
||||
ctx = self.limited_gemm_ctx
|
||||
with ctx:
|
||||
return torch.matmul(chunk_input, weight.T, out=chunk_out)
|
||||
elif user_gemm_method == GemmMethod.kCUBLAS:
|
||||
return torch.matmul(chunk_input, weight.T, out=chunk_out)
|
||||
else:
|
||||
raise RuntimeError(f"Invalid gemm method, got {self.gemm_method}.")
|
||||
|
||||
@classmethod
|
||||
def native_forward(
|
||||
cls,
|
||||
input,
|
||||
weight,
|
||||
bias=None,
|
||||
out=None,
|
||||
group=None,
|
||||
user_gemm_method=None,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
if user_gemm_method is not None and callable(user_gemm_method):
|
||||
gemm_out = user_gemm_method(
|
||||
input, weight, bias=bias, out=out, *args, **kwargs
|
||||
)
|
||||
out = out if gemm_out is None else gemm_out
|
||||
else:
|
||||
import ixformer.functions as ixff
|
||||
|
||||
# warning: 下面的两种 gemm 可能存在精度不一致
|
||||
# out = torch.matmul(input, weight.T, out=out)
|
||||
out = ixff.linear(input=input, weight=weight, bias=bias, output=out)
|
||||
ixfd.all_reduce(out, async_op=True, group=group)
|
||||
return out
|
||||
|
||||
@classmethod
|
||||
def is_supported(cls, input, num_chunks, comm_group):
|
||||
if not cls.enable():
|
||||
return False
|
||||
|
||||
ndim = input.ndim
|
||||
shape = input.shape
|
||||
|
||||
if ndim == 1:
|
||||
m, k = 1, shape[0]
|
||||
elif ndim == 2:
|
||||
m, k = shape
|
||||
else:
|
||||
m, k = sum(shape[:-1]), shape[-1]
|
||||
|
||||
return m >= 512
|
||||
|
||||
|
||||
_DEFAULT_OVERLAP_GROUP = None
|
||||
_DEFAULT_OVERLAP_COMM_N2 = None
|
||||
_DEFAULT_OVERLAP_COMM_N4 = None
|
||||
_DEFAULT_OVERLAP_CHUNKS = config.IXFORMER_OVERLAP_CHUNKS
|
||||
|
||||
|
||||
def linear_allreduce_overlap(
|
||||
input, weight, bias=None, out=None, group=None, num_chunks=None, *args, **kwargs
|
||||
):
|
||||
num_chunks = num_chunks or _DEFAULT_OVERLAP_CHUNKS
|
||||
|
||||
# print("call overlap:", GemmAllReduceSplitOverlapComm.is_supported(input, num_chunks=num_chunks, comm_group=group), input.shape, weight.shape if torch.is_tensor(weight) else None, "WorldSize:", ixfd.get_group_world_size(group), ", NumChunks:", num_chunks)
|
||||
if not GemmAllReduceSplitOverlapComm.is_supported(
|
||||
input, num_chunks=num_chunks, comm_group=group
|
||||
):
|
||||
return GemmAllReduceSplitOverlapComm.native_forward(
|
||||
input, weight, bias=bias, out=out, group=group, *args, **kwargs
|
||||
)
|
||||
|
||||
global _DEFAULT_OVERLAP_GROUP
|
||||
global _DEFAULT_OVERLAP_COMM_N2
|
||||
global _DEFAULT_OVERLAP_COMM_N4
|
||||
|
||||
if _DEFAULT_OVERLAP_GROUP is None:
|
||||
_DEFAULT_OVERLAP_GROUP = group
|
||||
|
||||
if num_chunks == 2 and group == _DEFAULT_OVERLAP_GROUP:
|
||||
if _DEFAULT_OVERLAP_COMM_N2 is None:
|
||||
_DEFAULT_OVERLAP_COMM_N2 = GemmAllReduceSplitOverlapComm.dispatcher(
|
||||
num_chunks=num_chunks, comm_group=group
|
||||
)
|
||||
overlap_comm = _DEFAULT_OVERLAP_COMM_N2
|
||||
elif num_chunks == 4 and group == _DEFAULT_OVERLAP_GROUP:
|
||||
if _DEFAULT_OVERLAP_COMM_N4 is None:
|
||||
_DEFAULT_OVERLAP_COMM_N4 = GemmAllReduceSplitOverlapComm.dispatcher(
|
||||
num_chunks=num_chunks, comm_group=group
|
||||
)
|
||||
overlap_comm = _DEFAULT_OVERLAP_COMM_N4
|
||||
else:
|
||||
overlap_comm = GemmAllReduceSplitOverlapComm.dispatcher(
|
||||
num_chunks=num_chunks, comm_group=group
|
||||
)
|
||||
return overlap_comm.forward(input, weight, bias=bias, out=out, *args, **kwargs)
|
||||
Reference in New Issue
Block a user