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/inference/overlap/__init__.py
Normal file
1
ixformer_sdk/inference/overlap/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
from .llama_decoder_layer_overlap import LlamaDecoderLayerOverlapProtocol, LlamaDecoderLayerOverlapDefault, create_vllm_llama_decoder_layer
|
||||
@@ -0,0 +1,333 @@
|
||||
import dataclasses
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import ixformer.distributed as ixfd
|
||||
import ixformer.functions as F
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from ixformer.contrib.vllm_flash_attn import flash_attn_varlen_func
|
||||
from ixformer.distributed.overlap_comm import SplitOverlapComm
|
||||
|
||||
from ixformer.core import config as ixff_config
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class FmhaOProjAllReduceLnGatingParams:
|
||||
# ==============================
|
||||
# attention
|
||||
# ==============================
|
||||
|
||||
# shape: [Batch * SeqLen, NumHeads / TP, HeadDim]
|
||||
q: torch.Tensor
|
||||
|
||||
# shape: [Batch * SeqLen, NumHeads / TP, HeadDim]
|
||||
k: torch.Tensor
|
||||
|
||||
# shape: [Batch * SeqLen, NumHeads / TP, HeadDim]
|
||||
v: torch.Tensor
|
||||
|
||||
# shape [Batch + 1], dtype torch.int32. The cumulative sequence lengths
|
||||
# of the sequences in the batch, used to index into q.
|
||||
cu_seqlens_q: torch.Tensor
|
||||
|
||||
# shape: [Batch + 1], dtype torch.int32. The cumulative sequence lengths
|
||||
# of the sequences in the batch, used to index into kv.
|
||||
cu_seqlens_k: torch.Tensor
|
||||
|
||||
# Maximum query sequence length in the batch.
|
||||
max_seqlen_q: int
|
||||
|
||||
# Maximum key sequence length in the batch.
|
||||
max_seqlen_k: int
|
||||
|
||||
# ==============================
|
||||
# o_proj
|
||||
# ==============================
|
||||
|
||||
# shape: [HiddenSize, NumHeads * HeadDim / TP], dtype: int8
|
||||
o_proj_weight: torch.Tensor
|
||||
|
||||
# shape: [HiddenSize], dtype: float32
|
||||
o_proj_weight_scale: torch.Tensor
|
||||
|
||||
# shape: [HiddenSize]
|
||||
o_proj_bias: torch.Tensor
|
||||
|
||||
# shape: [NumHeads * HeadDim / TP], dtype: float16 or bfloat16
|
||||
o_proj_smooth_scale: torch.Tensor
|
||||
|
||||
# ==============================
|
||||
# ln
|
||||
# ==============================
|
||||
|
||||
# shape: [Batch * SeqLen, HiddenSize], dtype: float16 or bfloat16
|
||||
residual: torch.Tensor
|
||||
|
||||
# shape: [HiddenSize], dtype: float16 or bfloat16
|
||||
ln_weight: torch.Tensor
|
||||
|
||||
# shape: [HiddenSize], dtype: float16 or bfloat16
|
||||
ln_bias: torch.Tensor
|
||||
|
||||
# ==============================
|
||||
# gating linear
|
||||
# ==============================
|
||||
|
||||
# shape: [TopK, HiddenSize], dtype: float16 or bfloat16
|
||||
gating_weight: torch.Tensor
|
||||
|
||||
# shape: [SeqLen, TopK], dtype: float16 or bfloat16
|
||||
out: Optional[torch.Tensor] = None
|
||||
|
||||
# ==============================
|
||||
# default parameters
|
||||
# ==============================
|
||||
|
||||
# the seqlens of q for per chunk when using overlap,
|
||||
# the parameter can be initiated by params.prepare_overlap_params(),
|
||||
# and only need to initialize once during the model's forward.
|
||||
cu_seqlens_q_chunks = None
|
||||
cu_seqlens_k_chunks = None
|
||||
|
||||
softmax_scale: Optional[float] = None
|
||||
ln_eps: float = 1e-5
|
||||
|
||||
@property
|
||||
def batch(self):
|
||||
return len(self.cu_seqlens_q) - 1
|
||||
|
||||
@property
|
||||
def seqlen(self):
|
||||
return self.q.shape[0]
|
||||
|
||||
@property
|
||||
def topk(self):
|
||||
return self.gating_weight.shape[0]
|
||||
|
||||
def prepare_overlap_params(self):
|
||||
"""compute the cu_seqlens qk of chunk when using overlap"""
|
||||
first_chunk_size = int(self.q.shape[0] // 2)
|
||||
|
||||
if not hasattr(self.cu_seqlens_q, "q_chunks"):
|
||||
first_q_chunks_cu_seqlens = self.cu_seqlens_q.clone()
|
||||
first_q_chunks_cu_seqlens[-1] = first_chunk_size
|
||||
|
||||
self.cu_seqlens_q.q_chunks = [
|
||||
first_q_chunks_cu_seqlens,
|
||||
first_q_chunks_cu_seqlens,
|
||||
]
|
||||
|
||||
self.cu_seqlens_q_chunks = self.cu_seqlens_q.q_chunks
|
||||
|
||||
if not hasattr(self.cu_seqlens_k, "k_chunks"):
|
||||
first_chunk = self.cu_seqlens_k.clone()
|
||||
last_chunk = self.cu_seqlens_k
|
||||
first_chunk[-1] = first_chunk_size
|
||||
self.cu_seqlens_k.k_chunks = [first_chunk, last_chunk]
|
||||
|
||||
self.cu_seqlens_k_chunks = self.cu_seqlens_k.k_chunks
|
||||
|
||||
return self
|
||||
|
||||
|
||||
class FmhaOProjAllreduceLnGatingOverlap(SplitOverlapComm):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
if self.num_chunks != 2:
|
||||
raise RuntimeError(
|
||||
f"Overlap only support num_chunks == 2, but got {self.num_chunks}."
|
||||
)
|
||||
|
||||
self.allreduce_end_events = [torch.cuda.Event() for _ in range(self.num_chunks)]
|
||||
|
||||
def start_ln_gating(self, chunk_idx):
|
||||
compute_stream = self._compute_streams[chunk_idx % self.num_compute_streams]
|
||||
compute_stream.wait_event(self.allreduce_end_events[chunk_idx])
|
||||
|
||||
def compute(self, params: FmhaOProjAllReduceLnGatingParams):
|
||||
if params.out is None:
|
||||
params.out = torch.empty(
|
||||
[params.seqlen, params.topk], device="cuda", dtype=torch.float
|
||||
)
|
||||
|
||||
seqlen_chunks = [int(params.q.shape[0] // 2)]
|
||||
seqlen_chunks.append(params.q.shape[0] - seqlen_chunks[0])
|
||||
|
||||
q_chunks = torch.split_with_sizes(params.q, seqlen_chunks, dim=0)
|
||||
|
||||
ar_out_chunks = []
|
||||
for chunk_idx in range(len(seqlen_chunks)):
|
||||
with self.compute_stream_context(chunk_idx):
|
||||
hidden_states = flash_attn_varlen_func(
|
||||
q=q_chunks[chunk_idx],
|
||||
k=params.k,
|
||||
v=params.v,
|
||||
cu_seqlens_q=params.cu_seqlens_q_chunks[chunk_idx],
|
||||
cu_seqlens_k=params.cu_seqlens_k_chunks[chunk_idx],
|
||||
max_seqlen_q=seqlen_chunks[chunk_idx],
|
||||
max_seqlen_k=seqlen_chunks[0]
|
||||
if chunk_idx == 0
|
||||
else params.max_seqlen_k,
|
||||
softmax_scale=params.softmax_scale,
|
||||
causal=True,
|
||||
window_size=(-1, -1),
|
||||
alibi_slopes=None,
|
||||
softcap=0,
|
||||
)
|
||||
|
||||
hidden_states = hidden_states.view(hidden_states.shape[0], -1)
|
||||
hidden_states, i_scales = F.dynamic_scaled_quant_dynamic_int8(
|
||||
hidden_states, params.o_proj_smooth_scale
|
||||
)
|
||||
|
||||
out_chunk = F.w8a8(
|
||||
hidden_states,
|
||||
params.o_proj_weight,
|
||||
i_scales,
|
||||
params.o_proj_weight_scale,
|
||||
bias=params.o_proj_bias,
|
||||
out_dtype=params.residual.dtype,
|
||||
output=None,
|
||||
persistent=True,
|
||||
)
|
||||
|
||||
self.start_comm(chunk_idx)
|
||||
|
||||
ixfd.all_reduce(
|
||||
out_chunk, async_op=True, group=self.comm_group, use_comm_stream=True
|
||||
)
|
||||
ar_out_chunks.append(out_chunk)
|
||||
|
||||
self.allreduce_end_events[chunk_idx].record(self._comm_stream)
|
||||
|
||||
ln_out = torch.empty_like(params.residual)
|
||||
ln_out_chunks = ln_out.chunk(2, dim=0)
|
||||
residual_chunks = torch.split_with_sizes(params.residual, seqlen_chunks, dim=0)
|
||||
|
||||
if params.out is None:
|
||||
params.out = torch.empty(
|
||||
[params.seqlen, params.topk], dtype=params.residual.dtype, device="cuda"
|
||||
)
|
||||
|
||||
out_chunks = list(torch.split_with_sizes(params.out, seqlen_chunks, dim=0))
|
||||
|
||||
for chunk_idx in range(len(seqlen_chunks)):
|
||||
self.start_ln_gating(chunk_idx)
|
||||
with self.compute_stream_context(chunk_idx):
|
||||
ln_out_chunk, residual_chunk = F.residual_layer_norm(
|
||||
input=ar_out_chunks[chunk_idx],
|
||||
weight=params.ln_weight,
|
||||
bias=params.ln_bias,
|
||||
residual=residual_chunks[chunk_idx].reshape(
|
||||
ar_out_chunks[chunk_idx].shape
|
||||
),
|
||||
eps=params.ln_eps,
|
||||
output=ln_out_chunks[chunk_idx],
|
||||
)
|
||||
if ln_out_chunk.dtype == params.gating_weight.dtype:
|
||||
F.linear(
|
||||
ln_out_chunk, params.gating_weight, output=out_chunks[chunk_idx]
|
||||
)
|
||||
else:
|
||||
F.mixed_type_linear(
|
||||
ln_out_chunk, params.gating_weight, output=out_chunks[chunk_idx]
|
||||
)
|
||||
|
||||
return (
|
||||
params.residual.reshape(params.batch, params.seqlen, -1),
|
||||
ln_out.reshape(params.batch, params.seqlen, -1),
|
||||
params.out,
|
||||
)
|
||||
|
||||
|
||||
_fa_o_proj_allreduce_ln_gating_overlap = None
|
||||
|
||||
|
||||
def fmha_oproj_allreduce_ln_gating(
|
||||
params: FmhaOProjAllReduceLnGatingParams,
|
||||
enable_overlap: bool = False,
|
||||
comm_group=None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
FMHA + OProjLinear + AllReduce + LayerNorm + GatingLinear
|
||||
|
||||
Args:
|
||||
params: fused operator params
|
||||
enable_overlap: whether enable overlap
|
||||
comm_group: communication group
|
||||
Returns:
|
||||
Residual: shape: [Batch * SeqLen, HiddenSize], dtype: float16 or bfloat16
|
||||
HiddenStates: shape: [Batch * SeqLen, HiddenSize], dtype: float16 or bfloat16
|
||||
GatingLinearOutput: shape: [SeqLen, TopK], dtype: float16 or bfloat16
|
||||
"""
|
||||
|
||||
global _fa_o_proj_allreduce_ln_gating_overlap
|
||||
if _fa_o_proj_allreduce_ln_gating_overlap is None:
|
||||
_fa_o_proj_allreduce_ln_gating_overlap = (
|
||||
FmhaOProjAllreduceLnGatingOverlap.dispatcher(
|
||||
num_chunks=2, comm_group=comm_group
|
||||
).forward
|
||||
)
|
||||
|
||||
if (
|
||||
enable_overlap
|
||||
and ixff_config.IXFORMER_ENABLE_OVERLAP_COMM
|
||||
and dist.is_initialized()
|
||||
and dist.get_world_size(comm_group) > 1
|
||||
and params.batch == 1
|
||||
and params.seqlen > 1
|
||||
):
|
||||
if params.cu_seqlens_q_chunks is None:
|
||||
params = params.prepare_overlap_params()
|
||||
return _fa_o_proj_allreduce_ln_gating_overlap(params)
|
||||
|
||||
hidden_states = flash_attn_varlen_func(
|
||||
q=params.q,
|
||||
k=params.k,
|
||||
v=params.v,
|
||||
cu_seqlens_q=params.cu_seqlens_q,
|
||||
cu_seqlens_k=params.cu_seqlens_k,
|
||||
max_seqlen_q=params.max_seqlen_q,
|
||||
max_seqlen_k=params.max_seqlen_k,
|
||||
softmax_scale=params.softmax_scale,
|
||||
causal=True,
|
||||
window_size=(-1, -1),
|
||||
alibi_slopes=None,
|
||||
softcap=0,
|
||||
)
|
||||
|
||||
input = hidden_states.view(hidden_states.shape[0], -1)
|
||||
input, i_scales = F.dynamic_scaled_quant_smoothquant(
|
||||
input, params.o_proj_smooth_scale
|
||||
)
|
||||
|
||||
hidden_states = F.w8a8(
|
||||
input,
|
||||
params.o_proj_weight,
|
||||
i_scales,
|
||||
params.o_proj_weight_scale,
|
||||
bias=params.o_proj_bias,
|
||||
out_dtype=params.residual.dtype,
|
||||
output=None,
|
||||
)
|
||||
ixfd.all_reduce(hidden_states, async_op=True, group=comm_group)
|
||||
|
||||
hidden_states, residual = F.residual_layer_norm(
|
||||
input=hidden_states,
|
||||
weight=params.ln_weight,
|
||||
bias=params.ln_bias,
|
||||
residual=params.residual.reshape(hidden_states.shape),
|
||||
eps=params.ln_eps,
|
||||
)
|
||||
if hidden_states.dtype == params.gating_weight.dtype:
|
||||
out = F.linear(hidden_states, params.gating_weight, output=params.out)
|
||||
else:
|
||||
out = F.mixed_type_linear(
|
||||
hidden_states, params.gating_weight, output=params.out
|
||||
)
|
||||
return (
|
||||
residual.reshape(params.batch, params.seqlen, -1),
|
||||
hidden_states.reshape(params.batch, params.seqlen, -1),
|
||||
out,
|
||||
)
|
||||
@@ -0,0 +1,219 @@
|
||||
import dataclasses
|
||||
import math
|
||||
from contextlib import contextmanager
|
||||
from typing import List, Optional
|
||||
|
||||
import ixformer.distributed as ixfd
|
||||
import ixformer.functions as F
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from ixformer.distributed.overlap_comm import SplitOverlapComm
|
||||
|
||||
from ixformer.core import config as ixff_config
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class GroupGemmMoeReduceSumAllReduceParams:
|
||||
# M: NumTokens * TopK
|
||||
# K: InnerSize // TP
|
||||
# N: HiddenSize
|
||||
# NumTokens: M // TopK
|
||||
|
||||
# =================================
|
||||
# group gemm
|
||||
# =================================
|
||||
|
||||
# the top k of experts
|
||||
topk: int
|
||||
|
||||
# shape: [M, K] if format[1]=="N" else [K, M], dtype: int8
|
||||
input: torch.Tensor
|
||||
|
||||
# shape: [NumExperts, N, K] if format[0]=="T" else [NumExperts, K, N], dtype: int8
|
||||
weight: torch.Tensor
|
||||
|
||||
# shape: [M], dtype: float32
|
||||
i_scales: torch.Tensor
|
||||
|
||||
# shape: [NumExperts, N], dtype: float32
|
||||
w_scales: torch.Tensor
|
||||
|
||||
# shape: [NumExperts], dtype: int32
|
||||
tokens_per_experts: torch.Tensor
|
||||
|
||||
# the dtype of output, support float16 and bfloat16
|
||||
out_dtype: torch.dtype = None
|
||||
|
||||
# index of dst to src, shape: [M], dtype: int32
|
||||
dst_to_src: torch.Tensor = None
|
||||
|
||||
# only support TN now
|
||||
format: str = "TN"
|
||||
|
||||
# =================================
|
||||
# moe reduce sum
|
||||
# =================================
|
||||
|
||||
# shape: [M // TopK, TopK], dtype: torch.float16 or torch.bfloat16
|
||||
topk_weight: torch.Tensor = None
|
||||
|
||||
# shape: [M // TopK, N], dtype: torch.float16 or torch.bfloat16
|
||||
output: torch.Tensor = None
|
||||
|
||||
# overlap
|
||||
output_chunks: Optional[List[torch.Tensor]] = None
|
||||
|
||||
@property
|
||||
def M(self):
|
||||
return self.input.shape[0]
|
||||
|
||||
@property
|
||||
def N(self):
|
||||
if torch.is_tensor(self.weight):
|
||||
return self.weight.shape[1]
|
||||
return sum(t.shape[1] for t in self.weight)
|
||||
|
||||
@property
|
||||
def K(self):
|
||||
return self.input.shape[-1]
|
||||
|
||||
def prepare_overlap_params(
|
||||
self, num_chunks: int, split_ratio: Optional[float] = None
|
||||
):
|
||||
if num_chunks == 2 and split_ratio not in [0, None]:
|
||||
return self.prepare_overla_params_with_ratio(split_ratio)
|
||||
return self.prepare_overlap_params_with_chunks(num_chunks)
|
||||
|
||||
def prepare_overlap_params_with_chunks(self, num_chunks: int):
|
||||
if torch.is_tensor(self.weight):
|
||||
weight_chunks = torch.chunk(self.weight, num_chunks, dim=1)
|
||||
self.weight = list(weight_chunks)
|
||||
|
||||
if torch.is_tensor(self.w_scales):
|
||||
weight_scale_chunks = torch.chunk(self.w_scales, num_chunks, dim=1)
|
||||
self.w_scales = list(weight_scale_chunks)
|
||||
|
||||
if self.output is None:
|
||||
self.output = torch.empty(
|
||||
self.M // self.topk, self.N, dtype=self.out_dtype, device="cuda"
|
||||
)
|
||||
|
||||
if torch.is_tensor(self.output):
|
||||
output_chunks = torch.chunk(self.output, num_chunks, dim=1)
|
||||
self.output_chunks = list(output_chunks)
|
||||
|
||||
def prepare_overla_params_with_ratio(self, split_ratio: float):
|
||||
N = self.N
|
||||
n_chunks = [int(math.ceil(N * split_ratio))]
|
||||
n_chunks.append(N - n_chunks[0])
|
||||
|
||||
if torch.is_tensor(self.weight):
|
||||
weight_chunks = torch.split(self.weight, n_chunks, dim=1)
|
||||
self.weight = list(weight_chunks)
|
||||
|
||||
if torch.is_tensor(self.w_scales):
|
||||
weight_scale_chunks = torch.split(self.w_scales, n_chunks, dim=1)
|
||||
self.w_scales = list(weight_scale_chunks)
|
||||
|
||||
if self.output is None:
|
||||
self.output = torch.empty(
|
||||
self.M // self.topk, self.N, dtype=self.out_dtype, device="cuda"
|
||||
)
|
||||
|
||||
if torch.is_tensor(self.output):
|
||||
output_chunks = torch.split(self.output, n_chunks, dim=1)
|
||||
self.output_chunks = list(output_chunks)
|
||||
|
||||
|
||||
class GroupGemmMoeReduceSumAllReduceSplitNOverlap(SplitOverlapComm):
|
||||
def compute(self, params: GroupGemmMoeReduceSumAllReduceParams):
|
||||
for chunk_idx, (weight, weight_scale) in enumerate(
|
||||
zip(params.weight, params.w_scales)
|
||||
):
|
||||
with self.compute_stream_context(chunk_idx):
|
||||
out = F.moe_w8a8_group_gemm(
|
||||
input=params.input,
|
||||
weight=weight,
|
||||
i_scales=params.i_scales,
|
||||
w_scales=weight_scale,
|
||||
output_dtype=params.out_dtype,
|
||||
tokens_per_experts=params.tokens_per_experts,
|
||||
dst_to_src=params.dst_to_src,
|
||||
format=params.format,
|
||||
)
|
||||
out = out.reshape(-1, params.topk, out.shape[-1])
|
||||
out = F.moe_output_reduce_sum(
|
||||
input=out,
|
||||
topk_weight=params.topk_weight,
|
||||
output=params.output_chunks[chunk_idx],
|
||||
)
|
||||
|
||||
self.start_comm(chunk_idx)
|
||||
ixfd.all_reduce(
|
||||
out,
|
||||
async_op=True,
|
||||
group=self.comm_group,
|
||||
use_comm_stream=True,
|
||||
algo=ixfd.AllReduceAlgo.Stride,
|
||||
)
|
||||
|
||||
# if chunk_idx == 0: torch.cuda.synchronize()
|
||||
|
||||
return params.output
|
||||
|
||||
|
||||
_group_gemm_moe_reduce_sum_all_reduce_overlap = None
|
||||
|
||||
|
||||
def group_gemm_moe_reduce_sum_allreduce(
|
||||
params: GroupGemmMoeReduceSumAllReduceParams,
|
||||
enable_overlap: bool = False,
|
||||
comm_group=None,
|
||||
num_chunks=2,
|
||||
split_ratio: Optional[float] = None,
|
||||
):
|
||||
if params.output is None and params.out_dtype is None:
|
||||
raise RuntimeError(
|
||||
"group_gemm_moe_reduce_sum_all_reduce need out_dtype argument when output is none."
|
||||
)
|
||||
|
||||
if params.out_dtype is None:
|
||||
params.out_dtype = params.output.dtype
|
||||
|
||||
if (
|
||||
enable_overlap
|
||||
and ixff_config.IXFORMER_ENABLE_OVERLAP_COMM
|
||||
and dist.is_initialized()
|
||||
and dist.get_world_size(comm_group) > 1
|
||||
):
|
||||
global _group_gemm_moe_reduce_sum_all_reduce_overlap
|
||||
if _group_gemm_moe_reduce_sum_all_reduce_overlap is None:
|
||||
_group_gemm_moe_reduce_sum_all_reduce_overlap = (
|
||||
GroupGemmMoeReduceSumAllReduceSplitNOverlap.dispatcher(
|
||||
num_chunks=num_chunks, comm_group=comm_group
|
||||
).forward
|
||||
)
|
||||
|
||||
params.prepare_overlap_params(num_chunks=num_chunks, split_ratio=split_ratio)
|
||||
return _group_gemm_moe_reduce_sum_all_reduce_overlap(params)
|
||||
|
||||
out = F.moe_w8a8_group_gemm(
|
||||
input=params.input,
|
||||
weight=params.weight,
|
||||
i_scales=params.i_scales,
|
||||
w_scales=params.w_scales,
|
||||
output_dtype=params.out_dtype,
|
||||
tokens_per_experts=params.tokens_per_experts,
|
||||
dst_to_src=params.dst_to_src,
|
||||
format=params.format,
|
||||
)
|
||||
|
||||
out = out.reshape(-1, params.topk, out.shape[-1])
|
||||
out = F.moe_output_reduce_sum(
|
||||
input=out, topk_weight=params.topk_weight, output=params.output
|
||||
)
|
||||
|
||||
if dist.is_initialized() and dist.get_world_size(comm_group) > 1:
|
||||
ixfd.all_reduce(out, group=comm_group, async_op=True)
|
||||
|
||||
return out
|
||||
305
ixformer_sdk/inference/overlap/linear_mlp_overlap_comm.py
Normal file
305
ixformer_sdk/inference/overlap/linear_mlp_overlap_comm.py
Normal file
@@ -0,0 +1,305 @@
|
||||
from contextlib import nullcontext
|
||||
from typing import List
|
||||
|
||||
import torch.cuda
|
||||
|
||||
from ...distributed import _distributed as ixfd
|
||||
from ...distributed import overlap_comm as base_overlap_comm
|
||||
from ...distributed.overlap_comm import GemmAllReduceSplitOverlapComm
|
||||
from .. import overlap as overlap_base
|
||||
|
||||
|
||||
class LinearMLPOverlapCommHook:
|
||||
def on_mlp_linear2_finished(
|
||||
self,
|
||||
overlap_comm: "LinearMLPOverlapComm",
|
||||
num_chunks,
|
||||
chunk_idx,
|
||||
hidden_states_chunk,
|
||||
residual_chunk,
|
||||
):
|
||||
pass
|
||||
|
||||
def on_mlp_finished(
|
||||
self,
|
||||
overlap_comm: "LinearMLPOverlapComm",
|
||||
hidden_states_chunks,
|
||||
residual_chunks,
|
||||
):
|
||||
pass
|
||||
|
||||
|
||||
class LinearMLPOverlapComm(GemmAllReduceSplitOverlapComm):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(num_compute_streams=1, *args, **kwargs)
|
||||
|
||||
self._mlp_linear1_start_events: List[torch.cuda.Event] = [
|
||||
torch.cuda.Event() for _ in range(self.num_chunks)
|
||||
]
|
||||
self._mlp_linear1_end_events: List[torch.cuda.Event] = [
|
||||
torch.cuda.Event() for _ in range(self.num_chunks)
|
||||
]
|
||||
|
||||
self._mlp_linear1_stream: torch.cuda.Stream = torch.cuda.Stream()
|
||||
|
||||
def stop_linear_comm(self, chunk_idx):
|
||||
event = self._mlp_linear1_start_events[chunk_idx]
|
||||
event.record(self._comm_stream)
|
||||
|
||||
def start_mlp_linear1(self, chunk_idx):
|
||||
self._mlp_linear1_stream.wait_event(self._mlp_linear1_start_events[chunk_idx])
|
||||
|
||||
def stop_mlp_linear1(self, chunk_idx):
|
||||
event = self._mlp_linear1_end_events[chunk_idx]
|
||||
event.record(self._mlp_linear1_stream)
|
||||
|
||||
def start_mlp_linear2(self, chunk_idx):
|
||||
compute_stream = self._compute_streams[chunk_idx % self.num_compute_streams]
|
||||
compute_stream.wait_event(self._mlp_linear1_end_events[chunk_idx])
|
||||
|
||||
def compute(
|
||||
self,
|
||||
protocol: "overlap_base.LlamaDecoderLayerOverlapDefault",
|
||||
attn_output,
|
||||
residual,
|
||||
*,
|
||||
mlp_linear2_finished_callback=None,
|
||||
mlp_finished_callback=None,
|
||||
):
|
||||
""" """
|
||||
attn_output_shape = attn_output.shape
|
||||
residual_shape = None if residual is None else residual.shape
|
||||
|
||||
is_update_shape = attn_output.ndim > 2
|
||||
batch = 1
|
||||
if attn_output.ndim == 2:
|
||||
seqlen = attn_output_shape[0]
|
||||
else:
|
||||
batch = attn_output_shape[0]
|
||||
seqlen = attn_output_shape[1]
|
||||
|
||||
parallel_dims = batch * seqlen
|
||||
|
||||
if is_update_shape:
|
||||
attn_output = attn_output.reshape(parallel_dims, -1)
|
||||
if residual is not None:
|
||||
residual = residual.reshape(-1, residual_shape[-1])
|
||||
|
||||
attn_output_chunks, residual_chunks = protocol.split_mlp_inputs(
|
||||
attn_output, residual, self.num_chunks
|
||||
)
|
||||
|
||||
out = protocol.create_mlp_output()
|
||||
out_chunks = protocol.split_mlp_output(out, self.num_chunks)
|
||||
|
||||
res_chunks = []
|
||||
|
||||
# 1. output project linear
|
||||
for chunk_idx, (attn_output_chunk, residual_chunk) in enumerate(
|
||||
zip(attn_output_chunks, residual_chunks)
|
||||
):
|
||||
with self.compute_stream_context(chunk_idx):
|
||||
hidden_states = protocol.attn_output_proj_linear(
|
||||
self.num_chunks,
|
||||
chunk_idx,
|
||||
attn_output_chunk,
|
||||
use_limited_gemm=chunk_idx != 0,
|
||||
)
|
||||
|
||||
self.start_comm(chunk_idx)
|
||||
ixfd.all_reduce(
|
||||
hidden_states,
|
||||
async_op=True,
|
||||
group=self.comm_group,
|
||||
use_comm_stream=True,
|
||||
)
|
||||
self.stop_linear_comm(chunk_idx)
|
||||
attn_output_chunks[chunk_idx] = hidden_states
|
||||
|
||||
# 2. ln, mlp_linear1 and act
|
||||
for chunk_idx, (hidden_states, residual_chunk) in enumerate(
|
||||
zip(attn_output_chunks, residual_chunks)
|
||||
):
|
||||
self.start_mlp_linear1(chunk_idx)
|
||||
with self.stream_context(self._mlp_linear1_stream):
|
||||
(
|
||||
hidden_states,
|
||||
residual_chunk,
|
||||
) = protocol.attn_output_proj_linear_layer_norm(
|
||||
self.num_chunks, chunk_idx, hidden_states, residual_chunk
|
||||
)
|
||||
|
||||
hidden_states = protocol.mlp_linear1(
|
||||
self.num_chunks, chunk_idx, hidden_states, use_limited_gemm=True
|
||||
)
|
||||
hidden_states = protocol.mlp_activation(hidden_states)
|
||||
|
||||
attn_output_chunks[chunk_idx] = hidden_states
|
||||
res_chunks.append(residual_chunk)
|
||||
|
||||
self.stop_mlp_linear1(chunk_idx)
|
||||
|
||||
# 3. mlp_linear2
|
||||
for chunk_idx, hidden_states in enumerate(attn_output_chunks):
|
||||
self.start_mlp_linear2(chunk_idx)
|
||||
with self.compute_stream_context(chunk_idx):
|
||||
hidden_states = protocol.mlp_linear2(
|
||||
self.num_chunks,
|
||||
chunk_idx,
|
||||
hidden_states,
|
||||
out=out_chunks[chunk_idx],
|
||||
use_limited_gemm=True,
|
||||
)
|
||||
|
||||
self.start_comm(chunk_idx)
|
||||
ixfd.all_reduce(
|
||||
hidden_states,
|
||||
async_op=True,
|
||||
group=self.comm_group,
|
||||
use_comm_stream=True,
|
||||
)
|
||||
|
||||
if mlp_linear2_finished_callback is not None:
|
||||
mlp_linear2_finished_callback(
|
||||
self,
|
||||
self.num_chunks,
|
||||
chunk_idx,
|
||||
hidden_states,
|
||||
res_chunks[chunk_idx],
|
||||
)
|
||||
|
||||
if mlp_finished_callback is not None:
|
||||
mlp_finished_callback(self, out_chunks, res_chunks)
|
||||
|
||||
if is_update_shape:
|
||||
out = out.reshape(attn_output_shape)
|
||||
if residual is not None:
|
||||
residual = residual.reshape(residual_shape)
|
||||
|
||||
return out, residual
|
||||
|
||||
def gemm_dispatcher(
|
||||
self,
|
||||
chunk_idx,
|
||||
chunk_input,
|
||||
weight,
|
||||
chunk_out,
|
||||
use_limited_gemm=False,
|
||||
user_gemm_method=None,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
if user_gemm_method is not None and callable(user_gemm_method):
|
||||
ctx = self.ixf_limited_gemm_ctx if use_limited_gemm else nullcontext()
|
||||
with ctx:
|
||||
return user_gemm_method(
|
||||
chunk_input, weight, out=chunk_out, *args, **kwargs
|
||||
)
|
||||
|
||||
ctx = self.limited_gemm_ctx if use_limited_gemm else nullcontext()
|
||||
with ctx:
|
||||
return torch.matmul(chunk_input, weight.T, out=chunk_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
|
||||
|
||||
@classmethod
|
||||
def native_forward(
|
||||
cls,
|
||||
attn_output,
|
||||
residual,
|
||||
linear_weight,
|
||||
ln_layer,
|
||||
mlp_weight1,
|
||||
mlp_weight2,
|
||||
mlp_activation,
|
||||
linear_method=None,
|
||||
mlp_linear1_method=None,
|
||||
mlp_linear2_method=None,
|
||||
group=None,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
import ixformer.functions as ixff
|
||||
|
||||
linear_method = linear_method or ixff.linear
|
||||
mlp_linear1_method = mlp_linear1_method or ixff.linear
|
||||
mlp_linear2_method = mlp_linear2_method or ixff.linear
|
||||
|
||||
hidden_states = linear_method(attn_output, linear_weight)
|
||||
ixfd.all_reduce(hidden_states, async_op=True, group=group)
|
||||
|
||||
if ln_layer is not None:
|
||||
hidden_states, residual = ln_layer(hidden_states, residual)
|
||||
|
||||
hidden_states = mlp_linear1_method(hidden_states, mlp_weight1)
|
||||
hidden_states = mlp_activation(hidden_states)
|
||||
hidden_states = mlp_linear2_method(hidden_states, mlp_weight2)
|
||||
ixfd.all_reduce(hidden_states, async_op=True, group=group)
|
||||
|
||||
return hidden_states, residual
|
||||
|
||||
|
||||
_DEFAULT_OVERLAP_GROUP = None
|
||||
_DEFAULT_OVERLAP_COMM_N2 = None
|
||||
_DEFAULT_OVERLAP_COMM_N4 = None
|
||||
_DEFAULT_OVERLAP_CHUNKS = base_overlap_comm._DEFAULT_OVERLAP_CHUNKS
|
||||
|
||||
|
||||
def linear_mlp_overlap(
|
||||
protocol: "overlap_base.LlamaDecoderLayerOverlapProtocol",
|
||||
attn_output,
|
||||
residual,
|
||||
num_chunks=None,
|
||||
group=None,
|
||||
*,
|
||||
mlp_linear2_finished_callback=None,
|
||||
mlp_finished_callback=None,
|
||||
):
|
||||
num_chunks = num_chunks or _DEFAULT_OVERLAP_CHUNKS
|
||||
|
||||
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 = LinearMLPOverlapComm.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 = LinearMLPOverlapComm.dispatcher(
|
||||
num_chunks=num_chunks, comm_group=group
|
||||
)
|
||||
overlap_comm = _DEFAULT_OVERLAP_COMM_N4
|
||||
else:
|
||||
overlap_comm = LinearMLPOverlapComm.dispatcher(
|
||||
num_chunks=num_chunks, comm_group=group
|
||||
)
|
||||
|
||||
return overlap_comm.forward(
|
||||
protocol,
|
||||
attn_output,
|
||||
residual,
|
||||
mlp_linear2_finished_callback=mlp_linear2_finished_callback,
|
||||
mlp_finished_callback=mlp_finished_callback,
|
||||
)
|
||||
1396
ixformer_sdk/inference/overlap/llama_decoder_layer_overlap.py
Normal file
1396
ixformer_sdk/inference/overlap/llama_decoder_layer_overlap.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,190 @@
|
||||
import dataclasses
|
||||
import math
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import ixformer.distributed as ixfd
|
||||
import ixformer.functions as F
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from ixformer.distributed.overlap_comm import SplitOverlapComm
|
||||
|
||||
from ixformer.core import config as ixff_config
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class MoeReduceAllReduceLnQkvLinearParams:
|
||||
# ==============================
|
||||
# MOE Reduce Sum
|
||||
# ==============================
|
||||
|
||||
# shape: [Batch * SeqLen, TopK, HiddenSize], dtype: float16 or bfloat16
|
||||
input: torch.Tensor
|
||||
|
||||
# shape: [Batch * SeqLen, TopK], dtype: float32
|
||||
topk_weight: Optional[torch.Tensor]
|
||||
|
||||
# ==============================
|
||||
# Ln
|
||||
# ==============================
|
||||
|
||||
# shape: [Batch * SeqLen, HiddenSize], dtype: float16 or bfloat16
|
||||
residual: torch.Tensor
|
||||
|
||||
# shape: [HiddenSize], dtype: float16 or bfloat16
|
||||
ln_weight: torch.Tensor
|
||||
|
||||
# shape: [HiddenSize], dtype: float16 or bfloat16
|
||||
ln_bias: torch.Tensor
|
||||
ln_eps: float
|
||||
|
||||
# ==============================
|
||||
# QkvLinear
|
||||
# ==============================
|
||||
|
||||
# shape: [(NumHeads + 2 * NumKvHeads) * HeadDim / TP, HiddenSize], dtype: float16 or bfloat16
|
||||
qkv_weight: torch.Tensor
|
||||
|
||||
# shape: [(NumHeads + 2 * NumKvHeads) * HeadDim / TP]
|
||||
qkv_weight_scale: torch.Tensor
|
||||
|
||||
# shape: [Batch * SeqLen, (NumHeads + 2 * NumKvHeads) * HeadDim / TP], dtype: float16 or bfloat16
|
||||
qkv_out: torch.Tensor
|
||||
|
||||
|
||||
class MoeReduceSumAllReduceLnQkvLinearOverlap(SplitOverlapComm):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
self.allreduce_end_events = [torch.cuda.Event() for _ in range(self.num_chunks)]
|
||||
|
||||
def start_qkv_linear(self, chunk_idx):
|
||||
compute_stream = self._compute_streams[chunk_idx % self.num_compute_streams]
|
||||
compute_stream.wait_event(self.allreduce_end_events[chunk_idx])
|
||||
|
||||
def compute(self, params: MoeReduceAllReduceLnQkvLinearParams, split_ratio=0.5):
|
||||
input_chunk_sizes = [int(math.ceil(params.input.shape[0] * split_ratio))]
|
||||
input_chunk_sizes.append(params.input.shape[0] - input_chunk_sizes[0])
|
||||
|
||||
input_chunks = list(
|
||||
torch.split_with_sizes(params.input, input_chunk_sizes, dim=0)
|
||||
)
|
||||
topk_weight_chunks = list(
|
||||
torch.split_with_sizes(params.topk_weight, input_chunk_sizes, dim=0)
|
||||
)
|
||||
out_chunks = []
|
||||
|
||||
for chunk_idx in range(len(input_chunks)):
|
||||
with self.compute_stream_context(chunk_idx):
|
||||
out_chunks.append(
|
||||
F.moe_output_reduce_sum(
|
||||
input_chunks[chunk_idx],
|
||||
topk_weight=topk_weight_chunks[chunk_idx],
|
||||
)
|
||||
)
|
||||
|
||||
self.start_comm(chunk_idx)
|
||||
|
||||
ixfd.all_reduce(
|
||||
out_chunks[chunk_idx],
|
||||
async_op=True,
|
||||
group=self.comm_group,
|
||||
use_comm_stream=True,
|
||||
)
|
||||
|
||||
self.allreduce_end_events[chunk_idx].record(self._comm_stream)
|
||||
|
||||
residual_chunk_sizes = [int(params.residual.shape[0] * split_ratio)]
|
||||
residual_chunk_sizes.append(params.residual.shape[0] - residual_chunk_sizes[0])
|
||||
residual_chunks = torch.split_with_sizes(
|
||||
params.residual, residual_chunk_sizes, dim=0
|
||||
)
|
||||
|
||||
qkv_out_chunk_sizes = [int(params.qkv_out.shape[0] * split_ratio)]
|
||||
qkv_out_chunk_sizes.append(params.qkv_out.shape[0] - qkv_out_chunk_sizes[0])
|
||||
qkv_out_chunks = torch.split_with_sizes(
|
||||
params.qkv_out, qkv_out_chunk_sizes, dim=0
|
||||
)
|
||||
|
||||
for chunk_idx in range(len(input_chunks)):
|
||||
self.start_qkv_linear(chunk_idx)
|
||||
with self.compute_stream_context(chunk_idx):
|
||||
(
|
||||
i8_hidden_states,
|
||||
residual,
|
||||
i_scales,
|
||||
) = F.residual_layer_norm_dynamic_int8(
|
||||
input=out_chunks[chunk_idx],
|
||||
residual=residual_chunks[chunk_idx],
|
||||
weight=params.ln_weight,
|
||||
bias=params.ln_bias,
|
||||
eps=params.ln_eps,
|
||||
)
|
||||
|
||||
F.w8a8(
|
||||
i8_hidden_states,
|
||||
params.qkv_weight,
|
||||
i_scales,
|
||||
params.qkv_weight_scale,
|
||||
output=qkv_out_chunks[chunk_idx],
|
||||
)
|
||||
|
||||
return params.qkv_out, params.residual
|
||||
|
||||
|
||||
_moe_reduce_with_allreduce_overlap = None
|
||||
|
||||
|
||||
def moe_reduce_sum_allreduce_ln_qkv_linear(
|
||||
params: MoeReduceAllReduceLnQkvLinearParams,
|
||||
enable_overlap=False,
|
||||
comm_group=None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
MOE Reduce Sum + AllReduce + LayerNorm + QkvLinear
|
||||
|
||||
Args:
|
||||
params: fused operator params
|
||||
enable_overlap: whether enable overlap
|
||||
comm_group: communication group
|
||||
Returns:
|
||||
QkvLinearOutput: [Batch * SeqLen, (NumHeads + 2 * NumKvHeads) * HeadDim / TP], dtype: float16 or bfloat16
|
||||
Residual: [Batch * SeqLen, HiddenSize], dtype: float16 or bfloat16
|
||||
"""
|
||||
|
||||
global _moe_reduce_with_allreduce_overlap
|
||||
if _moe_reduce_with_allreduce_overlap is None:
|
||||
_moe_reduce_with_allreduce_overlap = (
|
||||
MoeReduceSumAllReduceLnQkvLinearOverlap.dispatcher(
|
||||
num_chunks=2, comm_group=comm_group
|
||||
).forward
|
||||
)
|
||||
|
||||
if (
|
||||
enable_overlap
|
||||
and ixff_config.IXFORMER_ENABLE_OVERLAP_COMM
|
||||
and dist.is_initialized()
|
||||
and dist.get_world_size(comm_group) > 1
|
||||
and params.input.shape[0] > 1
|
||||
):
|
||||
return _moe_reduce_with_allreduce_overlap(params, split_ratio=0.5)
|
||||
|
||||
out = F.moe_output_reduce_sum(params.input, topk_weight=params.topk_weight)
|
||||
ixfd.all_reduce(out, async_op=True, group=comm_group)
|
||||
|
||||
i8_hidden_states, residual, i_scales = F.residual_layer_norm_dynamic_int8(
|
||||
input=out,
|
||||
residual=params.residual,
|
||||
weight=params.ln_weight,
|
||||
bias=params.ln_bias,
|
||||
eps=params.ln_eps,
|
||||
)
|
||||
|
||||
out = F.w8a8(
|
||||
i8_hidden_states,
|
||||
params.qkv_weight,
|
||||
i_scales,
|
||||
params.qkv_weight_scale,
|
||||
output=params.qkv_out,
|
||||
)
|
||||
|
||||
return out, residual
|
||||
250
ixformer_sdk/inference/overlap/overlap_comm.py
Normal file
250
ixformer_sdk/inference/overlap/overlap_comm.py
Normal file
@@ -0,0 +1,250 @@
|
||||
import enum
|
||||
from typing import Optional
|
||||
|
||||
import ixformer.functions as ixff
|
||||
import torch
|
||||
from ixformer.inference.overlap.linear_mlp_overlap_comm import (
|
||||
LinearMLPOverlapComm,
|
||||
LinearMLPOverlapCommHook,
|
||||
)
|
||||
|
||||
|
||||
def get_overlap_linear_method(layer):
|
||||
if hasattr(layer, "_overlap_comm_gemm_fn"):
|
||||
return layer._overlap_comm_gemm_fn
|
||||
|
||||
if layer.linear_weights["weight"].itemsize == 2:
|
||||
layer._overlap_comm_gemm_fn = None
|
||||
return None
|
||||
|
||||
def overlap_linear_fn(input, weight, bias=None, out: torch.Tensor = None, **kwargs):
|
||||
return layer.linear_method.apply_weights(
|
||||
layer.linear_weights, input, output=out
|
||||
)
|
||||
|
||||
layer._overlap_comm_gemm_fn = overlap_linear_fn
|
||||
return overlap_linear_fn
|
||||
|
||||
|
||||
class DecoderLayerOverlapComm(LinearMLPOverlapCommHook):
|
||||
class HookStage(enum.IntEnum):
|
||||
kExited = 0
|
||||
kTracing = 1
|
||||
|
||||
class HookState:
|
||||
def __init__(self, max_num_chunks):
|
||||
self.max_num_chunks = max_num_chunks
|
||||
self.stage = DecoderLayerOverlapComm.HookStage.kExited
|
||||
|
||||
self.mlp_linaer2_end_events = [
|
||||
torch.cuda.Event() for _ in range(max_num_chunks)
|
||||
]
|
||||
self.ln_attn_end_event = torch.cuda.Event()
|
||||
|
||||
self.overlap_comm: Optional[LinearMLPOverlapComm] = None
|
||||
|
||||
def is_tracing_stage(self):
|
||||
return self.stage == DecoderLayerOverlapComm.HookStage.kTracing
|
||||
|
||||
def enter(self, overlap_comm, chunk_idx):
|
||||
self.stage = DecoderLayerOverlapComm.HookStage.kTracing
|
||||
|
||||
self.overlap_comm = overlap_comm
|
||||
self.mlp_linaer2_end_events[chunk_idx].record(overlap_comm._comm_stream)
|
||||
|
||||
def exit(self):
|
||||
self.overlap_comm = None
|
||||
self.stage = DecoderLayerOverlapComm.HookStage.kExited
|
||||
|
||||
def __str__(self):
|
||||
return f"HookState(overlap_comm={self.overlap_comm}, stage={self.stage})"
|
||||
|
||||
def __repr__(self):
|
||||
return self.__str__()
|
||||
|
||||
_overlap_comm_hook_state = dict()
|
||||
|
||||
def __init__(self, model_id, layer_idx, max_num_chunks: int = 4):
|
||||
"""
|
||||
DecoderLayer 的流程：
|
||||
ln_qkv: InputLayerNorm(hidden_states, [residual]) -> qkv_proj(hidden_states) -> q, k, v = split(hidden_states) -> Attention(q, k, v)
|
||||
linear_mlp: AttentionOutputProj(hidden_states) -> PostLayerNorm(hidden_states) -> MLPLinear1 -> MLPActivation -> MLPLinear2
|
||||
|
||||
其中：AttentionOutputProj 和 MLPLinear2 之后如果使用 TP，那么需要进行 AllReduce
|
||||
|
||||
通过上述流程，该类的目的是将 MLPLinear2 后的 AllReduce 和 DecoderLayer 最开始的 ln_qkv 进行 Overlap。
|
||||
其中，第一层 DecoderLayer 不进行 ln_qkv 的 Overlap，因为在第一层之前没有通讯。
|
||||
我们需要将第 i 层 MLPLinear2 后的通讯 和 第 i + 1 层的 ln_qkv 进行 Overlap。
|
||||
|
||||
为了管理当前的状态和获取前一层的状态，从而设计了 DecoderLayerOverlapComm 类。
|
||||
该类需要 model_id 来推断当前正在运行的模型，用 layer_idx 来标记每一层的开始和结束，
|
||||
以及通过 layer_idx 去获取前一层的状态。
|
||||
|
||||
注：
|
||||
- 在 call_ln_qkv_overlap 中对 Tensor 进行切分时，
|
||||
需要保持和 linear_mlp 切分的大小是一致的，否则会出现 Tensor 的数据不对应；
|
||||
- 如果需要使用 ln_qkv 进行 Overlap，那么必须使用该类的 linear_mlp 去替换 linear_mlp_overlap
|
||||
|
||||
:param model_id: 模型的 id，可以使用 id(model) 去设置
|
||||
:param layer_idx: layer 的索引，注意，需要从 0 到 NumLayers 的顺序去完成构造
|
||||
:param max_num_chunks: 最大能进行切分的次数
|
||||
"""
|
||||
self._model_id = model_id
|
||||
self._layer_idx = layer_idx
|
||||
self._max_num_chunks = max_num_chunks
|
||||
|
||||
self._state = self.HookState(max_num_chunks)
|
||||
self._overlap_comm_hook_state[(model_id, layer_idx)] = self._state
|
||||
self._prev_layer_state = (
|
||||
None
|
||||
if layer_idx == 0
|
||||
else self._overlap_comm_hook_state[(model_id, layer_idx - 1)]
|
||||
)
|
||||
|
||||
@property
|
||||
def model_id(self):
|
||||
return self._model_id
|
||||
|
||||
@property
|
||||
def layer_idx(self):
|
||||
return self._layer_idx
|
||||
|
||||
@property
|
||||
def max_num_chunks(self):
|
||||
return self._max_num_chunks
|
||||
|
||||
@property
|
||||
def state(self) -> "DecoderLayerOverlapComm.HookState":
|
||||
return self._state
|
||||
|
||||
@property
|
||||
def prev_layer_state(self) -> "DecoderLayerOverlapComm.HookState":
|
||||
return self._prev_layer_state
|
||||
|
||||
def is_ln_qkv_overlap(self):
|
||||
return not (
|
||||
self.layer_idx == 0
|
||||
or not self.prev_layer_state.is_tracing_stage()
|
||||
or self.prev_layer_state.overlap_comm is None
|
||||
)
|
||||
|
||||
def ln_qkv(self, hidden_states, residual, ln_layer, qkv_layer, out_last_dim):
|
||||
"""
|
||||
:param hidden_state: shape[Batch * SeqLen, HiddenSize]
|
||||
:param residual: shape[Batch * SeqLen, HiddenSize]
|
||||
:param ln_layer: torch.nn.Module or Function(hidden_state, residual=None)
|
||||
:param qkv_layer: vllm.QKVParallelLinear
|
||||
:param out_last_dim: qkv_layer 输出 Tensor 的最后一个维度
|
||||
:return: qkv, residual
|
||||
"""
|
||||
if self.is_ln_qkv_overlap():
|
||||
qkv, residual = self.call_ln_qkv_overlap(
|
||||
hidden_states, residual, ln_layer, qkv_layer, out_last_dim
|
||||
)
|
||||
else:
|
||||
qkv, residual = self.call_ln_qkv(
|
||||
hidden_states, residual, ln_layer, qkv_layer
|
||||
)
|
||||
|
||||
return qkv, residual
|
||||
|
||||
def call_ln_qkv_overlap(
|
||||
self, hidden_states, residual, ln_layer, qkv_layer, out_last_dim
|
||||
):
|
||||
if hidden_states.ndim != 2:
|
||||
raise RuntimeError(
|
||||
f"Expected 2-dim for hidden state, but got {hidden_states.ndim}."
|
||||
)
|
||||
|
||||
num_chunks = self.prev_layer_state.overlap_comm.num_chunks
|
||||
overlap_comm: LinearMLPOverlapComm = self.prev_layer_state.overlap_comm
|
||||
|
||||
if num_chunks > self.max_num_chunks:
|
||||
raise RuntimeError(
|
||||
f"The layer is not support more than {self.max_num_chunks}, got {num_chunks}."
|
||||
)
|
||||
|
||||
hidden_state_chunks = list(torch.chunk(hidden_states, num_chunks, dim=0))
|
||||
if residual is None:
|
||||
residual = hidden_states
|
||||
residual_chunks = [None] * num_chunks
|
||||
else:
|
||||
residual_chunks = torch.chunk(residual, num_chunks, dim=0)
|
||||
|
||||
out = torch.empty(
|
||||
(hidden_states.shape[0], out_last_dim),
|
||||
device=hidden_states.device,
|
||||
dtype=hidden_states.dtype,
|
||||
)
|
||||
out_chunks = list(torch.chunk(out, num_chunks, dim=0))
|
||||
|
||||
for chunk_idx, (hidden_state_chunk, residual_chunk, out_chunk) in enumerate(
|
||||
zip(hidden_state_chunks, residual_chunks, out_chunks)
|
||||
):
|
||||
overlap_comm._compute_streams[
|
||||
chunk_idx % overlap_comm.num_compute_streams
|
||||
].wait_event(self.prev_layer_state.mlp_linaer2_end_events[chunk_idx])
|
||||
with overlap_comm.compute_stream_context(chunk_idx):
|
||||
self.call_ln_qkv(
|
||||
hidden_state_chunk,
|
||||
residual_chunk,
|
||||
ln_layer,
|
||||
qkv_layer,
|
||||
chunk_idx,
|
||||
use_limited_gemm=chunk_idx != (num_chunks - 1),
|
||||
out=out_chunk,
|
||||
overlap_comm=overlap_comm,
|
||||
)
|
||||
|
||||
self.prev_layer_state.exit()
|
||||
overlap_comm.stop_overlap()
|
||||
return out, residual
|
||||
|
||||
def call_ln_qkv(
|
||||
self,
|
||||
hidden_state,
|
||||
residual,
|
||||
ln_layer,
|
||||
qkv_layer,
|
||||
chunk_idx=0,
|
||||
use_limited_gemm=False,
|
||||
out=None,
|
||||
overlap_comm: LinearMLPOverlapComm = None,
|
||||
):
|
||||
if residual is None:
|
||||
residual = hidden_state
|
||||
if ln_layer is not None:
|
||||
hidden_state = ln_layer(hidden_state)
|
||||
else:
|
||||
hidden_state, residual = ln_layer(hidden_state, residual)
|
||||
|
||||
if out is None:
|
||||
qkv, _ = qkv_layer(hidden_state)
|
||||
else:
|
||||
gemm_method = get_overlap_linear_method(qkv_layer)
|
||||
qkv = overlap_comm.gemm_dispatcher(
|
||||
chunk_idx=chunk_idx,
|
||||
chunk_input=hidden_state,
|
||||
weight=qkv_layer.linear_weights["weight"],
|
||||
chunk_out=out,
|
||||
user_gemm_method=gemm_method,
|
||||
use_limited_gemm=use_limited_gemm,
|
||||
)
|
||||
|
||||
return qkv, residual
|
||||
|
||||
def linear_mlp(self, *args, **kwargs):
|
||||
"""ref: linear_mlp_overlap"""
|
||||
return ixff.linear_mlp_overlap(
|
||||
*args, **kwargs, mlp_linear2_finished_callback=self.on_mlp_linear2_finished
|
||||
)
|
||||
|
||||
def on_mlp_linear2_finished(
|
||||
self,
|
||||
overlap_comm: LinearMLPOverlapComm,
|
||||
num_chunks,
|
||||
chunk_idx,
|
||||
hidden_states_chunk,
|
||||
residual_chunk,
|
||||
):
|
||||
self.state.enter(overlap_comm, chunk_idx)
|
||||
154
ixformer_sdk/inference/overlap/w8a8_allreduce.py
Normal file
154
ixformer_sdk/inference/overlap/w8a8_allreduce.py
Normal file
@@ -0,0 +1,154 @@
|
||||
import math
|
||||
from typing import Optional
|
||||
|
||||
import ixformer.distributed as ixfd
|
||||
import ixformer.functions as F
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from ixformer.distributed.overlap_comm import SplitOverlapComm
|
||||
|
||||
from ixformer.core import config as ixff_config
|
||||
|
||||
__all__ = ["w8a8_allreduce"]
|
||||
|
||||
|
||||
class W8A8AllReduceOverlap(SplitOverlapComm):
|
||||
def compute(
|
||||
self,
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
input_scale: torch.Tensor,
|
||||
weight_scale: torch.Tensor,
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
output: Optional[torch.Tensor] = None,
|
||||
format: str = "TN",
|
||||
out_dtype: torch.dtype = None,
|
||||
comm_group=None,
|
||||
split_ratio=0.5,
|
||||
):
|
||||
# compute the chunk size of input
|
||||
input_chunk_sizes = [int(math.ceil(input.shape[0] * split_ratio))]
|
||||
input_chunk_sizes.append(input.shape[0] - input_chunk_sizes[0])
|
||||
|
||||
# split input and input_scale
|
||||
input_chunks = list(torch.split_with_sizes(input, input_chunk_sizes, dim=0))
|
||||
input_scale_chunks = torch.split(
|
||||
input_scale,
|
||||
input_chunk_sizes,
|
||||
)
|
||||
|
||||
# create output and split it
|
||||
if output is None:
|
||||
if out_dtype is None:
|
||||
raise RuntimeError(
|
||||
"w8a8 gemm need out_dtype argument when output is none."
|
||||
)
|
||||
output = torch.empty(
|
||||
(input.shape[:-1] + (weight.shape[0],)),
|
||||
dtype=out_dtype,
|
||||
device=input.device,
|
||||
)
|
||||
out_chunks = torch.split(output, input_chunk_sizes)
|
||||
|
||||
# overlap gemm and allreduce
|
||||
for chunk_idx in range(len(input_chunks)):
|
||||
# submit gemm kernel into compute stream
|
||||
with self.compute_stream_context(chunk_idx):
|
||||
F.w8a8(
|
||||
input=input_chunks[chunk_idx],
|
||||
weight=weight,
|
||||
i_scales=input_scale_chunks[chunk_idx],
|
||||
w_scales=weight_scale,
|
||||
bias=bias,
|
||||
output=out_chunks[chunk_idx],
|
||||
format=format,
|
||||
persistent=chunk_idx != 0,
|
||||
)
|
||||
|
||||
# recode compute stream and wait gemm
|
||||
self.start_comm(chunk_idx)
|
||||
|
||||
# submit allreduce kernel into communication stream by set use_comm_stream to true
|
||||
ixfd.all_reduce(
|
||||
out_chunks[chunk_idx],
|
||||
async_op=True,
|
||||
group=self.comm_group,
|
||||
use_comm_stream=True,
|
||||
)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
_w8a8_allreduce_overlap = None
|
||||
|
||||
|
||||
def w8a8_allreduce(
|
||||
enable_overlap: bool,
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
input_scale: torch.Tensor,
|
||||
weight_scale: torch.Tensor,
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
output: Optional[torch.Tensor] = None,
|
||||
format: str = "TN",
|
||||
out_dtype: torch.dtype = None,
|
||||
comm_group=None,
|
||||
split_ratio=0.5,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Gemm(w8a8) + AllReduce
|
||||
|
||||
Args:
|
||||
enable_overlap: whether enable gemm and allreduce overlap
|
||||
input: shape: [M, K], dtype: int8, linear input
|
||||
weight: shape: [N, K], dtype: int8, linear weight
|
||||
input_scale: shape: [M], dtype: float32, quantized scale of input
|
||||
weight_scale: shape: [N], dtype: float32, quantized scale of weight
|
||||
bias: shape: [N], dtype: float16 or bfloat16, linear bias
|
||||
output: shape: [M, N], dtype: float16 or bfloat16, allreduce output
|
||||
format: options include TN, NN and NT
|
||||
out_dtype: use the argument to decide to the dtype of output when output is None
|
||||
comm_group: communication group
|
||||
split_ratio: split the ratio of input.shape[0] when using overlap, range: (0, 1),
|
||||
it will affect area of the overlap for gemm and allreduce.
|
||||
Returns: output
|
||||
"""
|
||||
if (
|
||||
enable_overlap
|
||||
and ixff_config.IXFORMER_ENABLE_OVERLAP_COMM
|
||||
and dist.is_initialized()
|
||||
and dist.get_world_size(comm_group) > 1
|
||||
and input.shape[0] > 1
|
||||
):
|
||||
global _w8a8_allreduce_overlap
|
||||
if _w8a8_allreduce_overlap is None:
|
||||
_w8a8_allreduce_overlap = W8A8AllReduceOverlap.dispatcher(
|
||||
num_chunks=2, comm_group=comm_group
|
||||
).forward
|
||||
return _w8a8_allreduce_overlap(
|
||||
input=input,
|
||||
weight=weight,
|
||||
input_scale=input_scale,
|
||||
weight_scale=weight_scale,
|
||||
bias=bias,
|
||||
output=output,
|
||||
format=format,
|
||||
out_dtype=out_dtype,
|
||||
split_ratio=split_ratio,
|
||||
)
|
||||
|
||||
out = F.w8a8(
|
||||
input=input,
|
||||
weight=weight,
|
||||
i_scales=input_scale,
|
||||
w_scales=weight_scale,
|
||||
bias=bias,
|
||||
output=output,
|
||||
format=format,
|
||||
out_dtype=out_dtype,
|
||||
)
|
||||
|
||||
if dist.get_world_size() > 1:
|
||||
ixfd.all_reduce(out, op=ixfd.ReduceOp.SUM, async_op=True, group=comm_group)
|
||||
|
||||
return out
|
||||
Reference in New Issue
Block a user