dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

feat(CRITICAL): 从 GitHub 扫描搬运 ixformer SDK + xllm 完整 GDN/MoE 代码

Browse Source 
来源:
  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
2026-08-11 02:31:56 +00:00
parent a8b16da5da
commit 87a19d2d00
250 changed files with 76690 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
ixformer_sdk/train/functions/__init__.py Normal file
View File

@@ -0,0 +1,12 @@
from .cross_entropy_loss import *
from .fused_rope import *
from .geglu import *
from .gelu import *
from .layernorm import *
from .linear import *
from .matmul import *
from .residual_bias import *
from .residual_bias_ln import *
from .rms_norm import *
from .swiglu import *
from .group_norm import *

239
ixformer_sdk/train/functions/cross_entropy_loss.py Normal file
View File

@@ -0,0 +1,239 @@
from typing import Union
import ixformer._C as ops
import torch
from torch.autograd.function import Function
from torch.nn import init
from torch.nn.parameter import Parameter
__all__ = ["vocab_parallel_cross_entropy"]
class _VocabParallelCrossEntropyCustom(Function):
@staticmethod
def forward(ctx, vocab_parallel_logits, target, label_smoothing=0.0):
device = vocab_parallel_logits.device
xnumel = vocab_parallel_logits.shape[0]
rnumel = vocab_parallel_logits.shape[-1]
exp_logits = torch.empty(
(xnumel, 1, rnumel), device=device, dtype=torch.float32
)
masked_target_1d = torch.empty((xnumel,), device=device, dtype=torch.int32)
loss = torch.empty((xnumel, 1), device=device, dtype=torch.float32)
ops.train.cross_entropy_loss_forward(
vocab_parallel_logits, target.int(), exp_logits, masked_target_1d, loss
)
vocab_size = exp_logits.size(-1)
if label_smoothing > 0:
"""
We'd like to assign 1 / (K - 1) probability mass to every index that is not the ground truth.
= (1 - alpha) * y_gt + alpha * mean(y_{i for i != gt})
= (1 - alpha) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i
= ((K - 1) * (1 - alpha) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i
= (K * (1 - alpha) - 1) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i} y_i
= (1 - (alpha * K) / (K - 1)) * y_gt + ( (alpha * K) / (K - 1) ) * \sum_{i} y_i / K
From: https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/common/losses/smoothed_cross_entropy.py
"""
assert 1.0 > label_smoothing > 0.0
smoothing = label_smoothing * vocab_size / (vocab_size - 1)
# Exp logits at this point are normalized probabilities. So we can just take the log to get log-probs.
log_probs = torch.log(exp_logits)
mean_log_probs = log_probs.mean(dim=-1)
loss = (1.0 - smoothing) * loss - smoothing * mean_log_probs
ctx.label_smoothing, ctx.vocab_size = label_smoothing, vocab_size
# Store softmax, target-mask and masked-target for backward pass.
ctx.save_for_backward(exp_logits, masked_target_1d)
return loss
@staticmethod
def backward(ctx, grad_output):
# Retreive tensors from the forward path.
softmax, masked_target_1d = ctx.saved_tensors
label_smoothing, vocab_size = ctx.label_smoothing, ctx.vocab_size
# All the inputs have softmax as thier gradient.
grad_input = softmax
# For simplicity, work with the 2D gradient.
partition_vocab_size = softmax.size()[-1]
grad_2d = grad_input.view(-1, partition_vocab_size)
# Add the gradient from matching classes.
arange_1d = torch.arange(start=0, end=grad_2d.size()[0], device=grad_2d.device)
softmax_update = 1.0
if label_smoothing > 0:
smoothing = label_smoothing * vocab_size / (vocab_size - 1)
grad_2d[arange_1d, masked_target_1d] -= (1.0 - smoothing) * softmax_update
average_grad = 1 / vocab_size
grad_2d[arange_1d, :] -= smoothing * average_grad
else:
grad_2d[arange_1d, masked_target_1d] -= softmax_update
# Finally elementwise multiplication with the output gradients.
grad_input = torch.mul(grad_input, grad_output.unsqueeze(dim=-1))
return grad_input, None, None
class _VocabParallelCrossEntropy(Function):
@staticmethod
def forward(
ctx,
vocab_parallel_logits,
target,
label_smoothing=0.0,
vocab_start_index=0,
vocab_end_index=320000,
group=None,
):
# Maximum value along vocab dimension across all GPUs.
logits_max = torch.max(vocab_parallel_logits, dim=-1)[0]
torch.distributed.all_reduce(
logits_max, op=torch.distributed.ReduceOp.MAX, group=group
)
# Subtract the maximum value.
vocab_parallel_logits = vocab_parallel_logits - logits_max.unsqueeze(dim=-1)
# Get the partition's vocab indecies
partition_vocab_size = vocab_parallel_logits.size()[-1]
# Create a mask of valid vocab ids (1 means it needs to be masked).
target_mask = (target < vocab_start_index) | (target >= vocab_end_index)
masked_target = target.clone() - vocab_start_index
masked_target[target_mask] = 0
# Get predicted-logits = logits[target].
# For Simplicity, we convert logits to a 2-D tensor with size
# [*, partition-vocab-size] and target to a 1-D tensor of size [*].
logits_2d = vocab_parallel_logits.view(-1, partition_vocab_size)
masked_target_1d = masked_target.view(-1)
arange_1d = torch.arange(
start=0, end=logits_2d.size()[0], device=logits_2d.device
)
predicted_logits_1d = logits_2d[arange_1d, masked_target_1d]
predicted_logits_1d = predicted_logits_1d.clone().contiguous()
predicted_logits = predicted_logits_1d.view_as(target)
predicted_logits[target_mask] = 0.0
# All reduce is needed to get the chunks from other GPUs.
torch.distributed.all_reduce(
predicted_logits,
op=torch.distributed.ReduceOp.SUM,
group=group,
)
# Sum of exponential of logits along vocab dimension across all GPUs.
exp_logits = vocab_parallel_logits
torch.exp(vocab_parallel_logits, out=exp_logits)
sum_exp_logits = exp_logits.sum(dim=-1)
torch.distributed.all_reduce(
sum_exp_logits,
op=torch.distributed.ReduceOp.SUM,
group=group,
)
# Loss = log(sum(exp(logits))) - predicted-logit.
loss = torch.log(sum_exp_logits) - predicted_logits
# Normalize and optionally smooth logits
exp_logits.div_(sum_exp_logits.unsqueeze(dim=-1))
vocab_size = exp_logits.size(-1)
if label_smoothing > 0:
"""
We'd like to assign 1 / (K - 1) probability mass to every index that is not the ground truth.
= (1 - alpha) * y_gt + alpha * mean(y_{i for i != gt})
= (1 - alpha) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i
= ((K - 1) * (1 - alpha) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i
= (K * (1 - alpha) - 1) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i} y_i
= (1 - (alpha * K) / (K - 1)) * y_gt + ( (alpha * K) / (K - 1) ) * \sum_{i} y_i / K
From: https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/common/losses/smoothed_cross_entropy.py
"""
assert 1.0 > label_smoothing > 0.0
smoothing = label_smoothing * vocab_size / (vocab_size - 1)
# Exp logits at this point are normalized probabilities. So we can just take the log to get log-probs.
log_probs = torch.log(exp_logits)
mean_log_probs = log_probs.mean(dim=-1)
loss = (1.0 - smoothing) * loss - smoothing * mean_log_probs
ctx.label_smoothing, ctx.vocab_size = label_smoothing, vocab_size
# Store softmax, target-mask and masked-target for backward pass.
ctx.save_for_backward(exp_logits, target_mask, masked_target_1d)
return loss
@staticmethod
def backward(ctx, grad_output):
# Retreive tensors from the forward path.
softmax, target_mask, masked_target_1d = ctx.saved_tensors
label_smoothing, vocab_size = ctx.label_smoothing, ctx.vocab_size
# All the inputs have softmax as thier gradient.
grad_input = softmax
# For simplicity, work with the 2D gradient.
partition_vocab_size = softmax.size()[-1]
grad_2d = grad_input.view(-1, partition_vocab_size)
# Add the gradient from matching classes.
arange_1d = torch.arange(start=0, end=grad_2d.size()[0], device=grad_2d.device)
softmax_update = 1.0 - target_mask.view(-1).float()
if label_smoothing > 0:
smoothing = label_smoothing * vocab_size / (vocab_size - 1)
grad_2d[arange_1d, masked_target_1d] -= (1.0 - smoothing) * softmax_update
average_grad = 1 / vocab_size
grad_2d[arange_1d, :] -= smoothing * average_grad
else:
grad_2d[arange_1d, masked_target_1d] -= softmax_update
# Finally elementwise multiplication with the output gradients.
grad_input.mul_(grad_output.unsqueeze(dim=-1))
return grad_input, None, None, None, None, None
def vocab_parallel_cross_entropy(
vocab_parallel_logits: torch.Tensor,
target: torch.Tensor,
label_smoothing: float = 0.0,
world_size: int = 1,
vocab_start_index: int = 0,
vocab_end_index: int = 320000,
group=None,
):
"""
参数说明:
目前只支持batch_size = 1 的情况,当batch_size >1时,计算不正确
Args:
vocab_parallel_logits: shape : [seq_len,1,vocal_size] dtype : torch.bfloat16,torch.float,torch.half
target: shape : [seq_len,1] dtype : torch.int64
group: TP 并行组
return:
loss: shape : [seq_len,1] dtype : torch.float32
"""
if world_size == 1:
return _VocabParallelCrossEntropyCustom.apply(
vocab_parallel_logits, target, label_smoothing
)
else:
return _VocabParallelCrossEntropy.apply(
vocab_parallel_logits,
target,
label_smoothing,
vocab_start_index,
vocab_end_index,
group,
)

214
ixformer_sdk/train/functions/fused_rope.py Normal file
View File

@@ -0,0 +1,214 @@
from typing import List, Tuple, Union
import ixformer._C as ops
import torch
from torch.autograd.function import Function
# adding by xuelu.peng 20240417
# from https://github.com/NVIDIA/apex/blob/master/apex/transformer/functional/fused_rope.py#L59
__all__ = ["fused_apply_rotary_pos_emb", "fused_apply_split_rotary_pos_emb", "fused_apply_rotary_pos_emb_cache"]
class FusedRoPEFunc(Function):
"""
Fused RoPE function
This implementation assumes the input tensor to be in `sbhd` format and the RoPE tensor to be
of shape (s, 1, 1, d). It accepts arbitrary memory layouts to avoid the expensive
`.contiguous()` calls, thus it may not achieve the best memory access pattern.
"""
@staticmethod
def forward(
ctx,
t: torch.Tensor,
freqs: torch.Tensor,
transpose_output_memory: bool = False,
) -> torch.Tensor:
# assert transpose_output_memory == False
output = ops.train.fused_rope_forward(t, freqs, transpose_output_memory)
ctx.save_for_backward(freqs)
ctx.transpose_output_memory = transpose_output_memory
return output
@staticmethod
def backward(
ctx, grad_output: torch.Tensor
) -> Tuple[Union[torch.Tensor, None], ...]:
(freqs,) = ctx.saved_tensors
grad_input = ops.train.fused_rope_backward(
grad_output, freqs, ctx.transpose_output_memory
)
return grad_input, None, None
class FusedFluxRoPEFunc(Function):
"""
Fused FluxRoPE function
This implementation assumes the input tensor to be in `bshd` format and the RoPE tensor to be
of shape (s, d), and output shape is the same as input shape.
"""
@staticmethod
def forward(
ctx,
t: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
imp_mode : int = 1
) -> torch.Tensor:
# assert transpose_output_memory == False
output = ops.train.fused_rope_forward_cached(t, cos, sin, imp_mode)
ctx.save_for_backward(cos, sin)
return output
@staticmethod
def backward(
ctx, grad_output: torch.Tensor
) -> Tuple[Union[torch.Tensor, None], ...]:
(cos, sin) = ctx.saved_tensors
grad_input = ops.train.fused_rope_backward_cached(
grad_output, cos, sin
)
return grad_input, None, None, None
def fused_apply_rotary_pos_emb(
t: torch.Tensor,
freqs: torch.Tensor,
transpose_output_memory: bool = False,
) -> torch.Tensor:
"""Apply rotary positional embedding to input tensor T in `sbhd` format, where
s: sequence length
b: batch size
h: head num
d: dim of each head
Args:
t (Tensor): Input tensor T is of shape [s, b, h, d], dtype : torch.float32, torch.half
freqs (Tensor): Rotary Positional embedding tensor freq is of shape [s, 1, 1, d] and
`float` dtype
transpose_output_memory (bool): Default to False. Whether to transpose the 's' and 'b'
dimension of the output's underlying memory format. This is very helpful when you want to
get a contiguous tensor after calling `output.transpose(0, 1)`.
Returns:
Tensor: The input tensor after applying RoPE
"""
return FusedRoPEFunc.apply(t, freqs, transpose_output_memory)
def fused_apply_rotary_pos_emb_cache(
t: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
imp_mode: int = 1,
) -> torch.Tensor:
"""Apply rotary positional embedding to input tensor T in `bshd` format, where
s: sequence length
b: batch size
h: head num
d: dim of each head
Args:
t (Tensor): Input tensor T is of shape [b, s, h, d], dtype : torch.float32, torch.half, torch.bfloat16
cos/sin (Tensor): Rotary Positional embedding tensor freq is of shape [s, d] and
`float` dtype
imp_mode (bool): Default to 1. 1 for flux/cogvideox/hunyuan-dit, img_mode=0 for Stable Audio. For now, only img_mode = 1 is supported.
Returns:
Tensor: The input tensor after applying RoPE
"""
return FusedFluxRoPEFunc.apply(t, cos, sin, imp_mode)
class FusedSplitRoPEFunc(torch.autograd.Function):
"""
Fused Split and RoPE function
This implementation assumes the input tensor to be in `sbh3d` format and the RoPE tensor to be
of shape (s, 1, 1, d). It accepts arbitrary memory layouts to avoid the expensive
`.contiguous()` calls, thus it may not achieve the best memory access pattern.
input: mix_q_k_v [s,b,hn_kv,h/hn_kv+2,d]
output: output_q, output_k, output_v [s,b,h,d]
"""
@staticmethod
def forward(
ctx,
mixed_q_k_v: torch.Tensor,
freqs: torch.Tensor,
transpose_output_memory: bool = False,
) -> torch.Tensor:
assert transpose_output_memory == False, "do not support transpose_output now"
assert mixed_q_k_v.is_contiguous() == True, "mixed_q_k_v should be contiguous in FusedSplitRoPEFunc."
s, b, hn_kv, repplus2, d = mixed_q_k_v.size()
num_key_value_groups = repplus2-2
q, k, v = torch.split(mixed_q_k_v, (num_key_value_groups,1,1), dim=3)
ctx.hn_kv = hn_kv
output_q, output_k, output_v = torch.empty_like(q).view(s,b,-1,d),torch.empty_like(q).view(s,b,-1,d),torch.empty_like(q).view(s,b,-1,d)
ops.train.fused_split_rope_forward(
q, k, v, freqs, output_q, output_k, output_v, transpose_output_memory, hn_kv, num_key_value_groups
)
ctx.save_for_backward(freqs)
ctx.transpose_output_memory = transpose_output_memory
return output_q, output_k, output_v
@staticmethod
def backward(
ctx, grad_o_q: torch.Tensor, grad_o_k: torch.Tensor, grad_o_v: torch.Tensor
) -> Tuple[Union[torch.Tensor, None], ...]:
# grad_o_q: [s,b,h,d]
s,b,h,d = grad_o_q.size()
hn_kv = ctx.hn_kv
mixed_shape = (s, b, hn_kv,(h//hn_kv+2), d)
if hn_kv == h:
grad_mixed_q_k_v = torch.empty(mixed_shape, dtype=grad_o_q.dtype, device=grad_o_q.device,memory_format=torch.contiguous_format) # torch.empty效率比torch.zeros高
else:
grad_mixed_q_k_v = torch.zeros(mixed_shape, dtype=grad_o_q.dtype, device=grad_o_q.device) # 支持 gqa 的情况，kernel内需要进行累加，需要把qkv的梯度置零
grad_q, grad_k, grad_v = torch.split(grad_mixed_q_k_v.view(s,b,hn_kv,-1,d), (h//hn_kv,1,1), dim=3)
(freqs,) = ctx.saved_tensors
ops.train.fused_split_rope_backward(
grad_o_q, grad_o_k, grad_o_v, freqs, grad_q, grad_k, grad_v, ctx.transpose_output_memory
)
return grad_mixed_q_k_v, None, None
def fused_apply_split_rotary_pos_emb(
mixed_q_k_v: torch.Tensor,
freqs: torch.Tensor,
transpose_output_memory: bool = False,
) -> torch.Tensor:
""" Split mixed_q_k_v and apply rotary positional embedding to q and k in `sbhd` format, where
s: sequence length
b: batch size
h: head num
d: dim of each head
hn_kv: num head of key and value
Args:
mixed_q_k_v (Tensor): Input tensor T is of shape [s,b,hn_kv,h/hn_kv+2,d]
freqs (Tensor): Rotary Positional embedding tensor freq is of shape [s, 1, 1, d] and
`float` dtype
transpose_output_memory (bool): Default to False. Whether to transpose the 's' and 'b'
dimension of the output's underlying memory format. This is very helpful when you want to
get a contiguous tensor after calling `output.transpose(0, 1)`.
Returns:
Tensors: The input tensors after split and applying RoPE
"""
return FusedSplitRoPEFunc.apply(mixed_q_k_v, freqs, transpose_output_memory)

43
ixformer_sdk/train/functions/geglu.py Normal file
View File

@@ -0,0 +1,43 @@
from typing import Union
import ixformer._C as ops
import torch
from torch.autograd.function import Function, FunctionCtx
__all__ = ["geglu"]
class GegluFunction(torch.autograd.Function):
@staticmethod
def forward(ctx, input):
output_shape = list(input.shape)
output_shape[-1] = output_shape[-1] // 2
output = input.new_empty(output_shape)
ops.train.geglu_training_forward(input, output)
ctx.save_for_backward(input)
return output
@staticmethod
def backward(ctx: FunctionCtx, grad_output):
input = ctx.saved_tensors[0]
grad_input = torch.empty_like(input)
ops.train.geglu_training_backward(input, grad_output, grad_input)
return grad_input
def geglu(input: "torch.Tensor"):
"""
等价实现：
def ref_gelu_and_mul(x: torch.Tensor) -> torch.Tensor:
dtype = x.dtype
x = x.float()
x1, x2 = x.chunk(chunks=2, dim=-1)
res = NNF.gelu(x2) * x1
return res.to(dtype)
Args:
input: dtype:[torch.float, torch.half, torch.bfloat16]
Returns:
output: (....,input.shape[-1] //2), dtype:[torch.float, torch.half, torch.bfloat16]
"""
return GegluFunction.apply(input)

47
ixformer_sdk/train/functions/gelu.py Normal file
View File

@@ -0,0 +1,47 @@
from typing import List, Union
import ixformer._C as ops
import torch
from torch.autograd.function import Function, FunctionCtx
__all__ = [
"gelu",
]
class GeluFunction(Function):
@staticmethod
def forward(
ctx, input: torch.Tensor, in_place: bool = False, training: bool = False
):
if training:
if in_place:
ctx.save_for_backward(input.clone())
else:
ctx.save_for_backward(input)
if in_place:
return ops.train.gelu_forward(input, input)
else:
return ops.train.gelu_forward(input)
@staticmethod
def backward(ctx: FunctionCtx, grad_outputs):
input = ctx.saved_tensors[0]
grad_input = ops.train.gelu_backward(input, grad_outputs)
return grad_input, None, None
def gelu(
input: torch.Tensor, in_place: bool = False, training: bool = False
) -> torch.Tensor:
"""
等价实现：
torch.nn.functional.gelu
Args:
input: dtype:[torch.float, torch.half, torch.bfloat16]
in place: bool. Whether to operate directly on the original input data.
Returns:
output: dtype:[torch.float, torch.half, torch.bfloat16]
"""
return GeluFunction.apply(input, in_place, training)

61
ixformer_sdk/train/functions/group_norm.py Normal file
View File

@@ -0,0 +1,61 @@
import ixformer._C as ops
import torch
from torch.nn import init
from torch.nn.parameter import Parameter
class GN_NHWC_Func(torch.autograd.Function):
@staticmethod
def forward(ctx, X: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, G: int, eps: float, activation: str):
X_out, means, rstds = ops.train.gn_nhwc_fwd(X, weight, bias, G, eps, activation)
ctx.save_for_backward(X, weight, bias, means, rstds)
ctx.G = G
ctx.activation = activation
return X_out
@staticmethod
def backward(ctx, dy: torch.Tensor):
dy = dy.contiguous(memory_format=torch.channels_last)
X, weight, bias, means, rstds = ctx.saved_tensors
dx, dgamma, dbeta = ops.train.gn_nhwc_bwd(dy, X, weight, bias, means, rstds, ctx.G, ctx.activation)
return dx, dgamma, dbeta, None, None, None
class GroupNorm_nhwc(torch.nn.GroupNorm):
def __init__(self, num_groups: int, nc: int, activation='identity', **kwargs):
super().__init__(num_groups, nc, **kwargs)
assert activation in {'identity', 'silu', 'relu', 'gelu', 'gelu_tanh'}
if activation == 'identity':
self.activation = 0
if activation == 'relu':
self.activation = 1
if activation == 'silu':
self.activation = 2
if activation == 'gelu':
self.activation = 3
if activation == 'gelu_tanh':
self.activation = 4
@torch._dynamo.disable
def forward(self, x):
#print(x.shape, self.num_channels)
if len(x.size()) == 3:
N, C, L = x.shape
elif len(x.size()) == 4:
N, C, H, W = x.shape
else:
raise ValueError
G = self.num_groups
#if C // G > 512:
# raise ValueError(f'Error in fwd for X.shape={x.shape}, G={G}: C // G = {C // G} which is greater than 512. This input is not supported.')
#if H * W % 8 != 0:
# raise ValueError(f'Error in fwd for X.shape={x.shape}, G={G}: H * W is not a multiple of 8. This input is not supported.')
if self.affine:
return GN_NHWC_Func.apply(x, self.weight, self.bias, self.num_groups, self.eps, self.activation)
else:
w = torch.ones((self.num_channels,), device=x.device, dtype=x.dtype)
b = torch.zeros((self.num_channels,), device=x.device, dtype=x.dtype)
return GN_NHWC_Func.apply(x, w, b, self.num_groups, self.eps, self.activation)

98
ixformer_sdk/train/functions/layernorm.py Normal file
View File

@@ -0,0 +1,98 @@
from typing import List, Tuple, Union
import ixformer._C as ops
import torch
from torch.autograd.function import Function, FunctionCtx
__all__ = ["layernorm"]
class LayerNormFunction(Function):
@staticmethod
def forward(
ctx,
input: torch.Tensor,
ln_weight: torch.Tensor,
ln_bias: torch.Tensor,
output: torch.Tensor,
normalized_shape=None,
training: bool = False,
):
if ln_weight is None or ln_bias is None:
raise NotImplementedError()
# normalized_shape 需要是list或者tuple，并且不能为空
if normalized_shape == None:
norm_size = ln_weight.size(-1)
else:
norm_size = 1
if isinstance(normalized_shape, int):
norm_size = normalized_shape
normalized_shape = [normalized_shape]
elif (
isinstance(normalized_shape, list)
or isinstance(normalized_shape, tuple)
) and len(normalized_shape) >= 1:
for i in normalized_shape:
norm_size = i * norm_size
else:
raise f"layer_norm(): argument 'normalized_shape' (position 2) must be tuple of ints, not {type(normalized_shape)}"
if norm_size != ln_weight.size(-1):
raise f"layer_norm(): argument 'norm_size' must == ln_weight.size(-1)"
if output is None:
output = torch.empty_like(input)
if training:
mean_size = input.numel() // norm_size
input_hat = torch.empty_like(input)
rstd = torch.empty([mean_size], dtype=input.dtype, device=input.device)
ops.train.layernorm_training_forward(
input, ln_weight, ln_bias, output, input_hat, rstd
)
ctx.norm_size = norm_size
ctx.save_for_backward(input_hat, rstd, ln_weight)
else:
ops.train.layernorm_forward(input, ln_weight, ln_bias, output)
return output
@staticmethod
# def backward(ctx: FunctionCtx, grad_output, dh, dr):
def backward(ctx: FunctionCtx, grad_output):
input_hat, rstd, ln_weight = ctx.saved_tensors
grad_input = torch.empty_like(input_hat)
grad_weight = torch.empty_like(ln_weight)
grad_bias = torch.empty_like(ln_weight)
ops.train.layernorm_weightbias_backward(
input_hat, grad_output, grad_weight, grad_bias
)
ops.train.layernorm_input_backward(
input_hat, rstd, grad_output, ln_weight, grad_input
)
return grad_input, grad_weight, grad_bias, None, None, None
def layernorm(
input: torch.Tensor,
ln_weight: torch.Tensor,
ln_bias: torch.Tensor,
normalized_shape=None,
output: torch.Tensor = None,
training: bool = False,
):
"""
等价实现：
torch.nn.functional.layer_norm( input, normalized_shape, ln_weight, ln_bias, eps=0.000001)
Arguments:
input: (batch_count * seq_len, hidden_size), dtype:[torch.half]
ln_weight: (hidden_size), dtype:[torch.half]
ln_bias:(hidden_size),dtype:[torch.half]
normalized_shape: list[int], [hidden_size]
Return:
output: (batch_count * seq_len, hidden_size), dtype:[torch.half]
"""
return LayerNormFunction.apply(
input, ln_weight, ln_bias, output, normalized_shape, training
)

89
ixformer_sdk/train/functions/linear.py Normal file
View File

@@ -0,0 +1,89 @@
import os
from typing import Union
import ixformer._C as ops
import torch
from torch.autograd.function import Function, FunctionCtx
__all__ = ["linear"]
class LinearFunction(Function):
@staticmethod
def forward(
ctx,
input: torch.Tensor,
weight: torch.Tensor,
bias: torch.Tensor = None,
output: torch.Tensor = None,
):
if bias is not None:
if output is None:
output = ops.train.linear_forward(input, weight, bias)
else:
ops.train.linear_forward_(input, weight, bias, output)
else:
if output is None:
output = ops.train.linear_forward(input, weight)
else:
ops.train.linear_forward_(input, weight, output)
ctx.has_bias = bias is not None
ctx.save_for_backward(input, weight)
return output
@staticmethod
def backward(ctx: FunctionCtx, dy: torch.Tensor):
x, w = ctx.saved_tensors
dx = ops.train.linear_backward_dx(w, dy, x.shape)
dw = ops.train.linear_backward_dw(x, dy, w.shape)
if ctx.has_bias:
reduce_dims = list(range(dy.ndim - 1))
db = torch.sum(dy, reduce_dims)
return dx, dw, db, None
else:
return dx, dw, None, None
def gemv_conditions(input, weight, bias, gemv_max_batch):
# gemv 使用的条件 input:[m,k] weight:[n,k]
# 1. m<=gemv_max_batch
# 2. k%2==0 n%2==0
# 3. bias is None
input = input.view(-1, input.shape[-1])
weight = weight.view(-1, weight.shape[-1])
m = input.shape[0]
k = input.shape[1]
n = weight.shape[0]
if bias is None and m <= gemv_max_batch and k % 2 == 0 and n % 2 == 0:
return True
return False
def linear(
input: torch.Tensor,
weight: torch.Tensor,
bias: torch.Tensor = None,
output: torch.Tensor = None,
use_gemv: bool = True,
gemv_max_batch=1,
):
"""
Arguments:
input : [...,k] dtype: [torch.half, torch.bfloat16]
weights : [n,k] dtype: [torch.half, torch.bfloat16]
use_gemv: bool 是否使用gemv
gemv 使用的条件 input:[m,k] weight:[n,k]
1. m<=gemv_max_batch
2. k%2==0 n%2==0
3. bias is None
gemv_max_batch: int 用于是否满足gemv使用条件的判断
Return:
output : [...,n] dtype: [torch.half, torch.bfloat16]
"""
return LinearFunction.apply(input, weight, bias, output)

108
ixformer_sdk/train/functions/matmul.py Normal file
View File

@@ -0,0 +1,108 @@
import ixformer._C as ops
import torch
from torch.autograd.function import Function, FunctionCtx
__all__ = ["matmul", "MatmulFunction"]
class MatmulFunction(Function):
@staticmethod
def forward(
ctx: FunctionCtx,
input: torch.Tensor,
other: torch.Tensor,
out: torch.Tensor = None,
transa: bool = False,
transb: bool = False,
alpha: float = 1.0,
beta: float = 0.0,
):
ctx.save_for_backward(input, other)
ctx.params = (transa, transb, alpha, beta)
if out is None:
return ops.train.matmul(
input, other, transa=transa, transb=transb, alpha=alpha, beta=beta
)
else:
return ops.train.matmul(
input,
other,
out=out,
transa=transa,
transb=transb,
alpha=alpha,
beta=beta,
)
@staticmethod
def backward(ctx: FunctionCtx, dy):
input, other = ctx.saved_tensors
transa, transb, alpha, beta = ctx.params
if beta in [1, None]:
raise RuntimeError("Backward don't support beta == 1.0f")
if not transa and not transb:
dx = matmul(dy, other, transb=True, alpha=alpha)
do = matmul(input, dy, transa=True, alpha=alpha)
return dx, do, None, None, None, None, None
if transa and not transb:
dx = matmul(other, dy, transb=True, alpha=alpha)
do = matmul(input, dy, alpha=alpha)
return dx, do, None, None, None, None, None
if not transa and transb:
dx = matmul(dy, other, alpha=alpha)
do = matmul(dy, input, transa=True, alpha=alpha)
return dx, do, None, None, None, None, None
if transa and transb:
dx = matmul(other, dy, transa=True, transb=True, alpha=alpha)
do = matmul(dy, input, transa=True, transb=True, alpha=alpha)
return dx, do, None, None, None, None, None
def matmul(
input: torch.Tensor,
other: torch.Tensor,
*,
out: torch.Tensor = None,
transa: bool = False,
transb: bool = False,
alpha: float = 1.0,
beta: float = 0.0
) -> torch.Tensor:
"""
等价实现:
def pt_matmul(a, b, transa, transb, alpha):
if transa:
dims = list(range(a.ndim))
dims[-1], dims[-2] = dims[-2], dims[-1]
a = a.permute(*dims).contiguous()
if transb:
dims = list(range(b.ndim))
dims[-1], dims[-2] = dims[-2], dims[-1]
b = b.permute(*dims).contiguous()
return alpha * torch.matmul(a, b)
Arguments:
input:
当transa为False shape : [...,m,k] dtype: torch.half
当transa为True shape : [...,k,m] dtype: torch.half
other:
当transb为False shape : [...,k,n] dtype: torch.half
当transb为True shape : [...,n,k] dtype: torch.half
Return:
output: [...m,n] dtype: [torch.half]
"""
if not input.is_contiguous():
input = input.contiguous()
if not other.is_contiguous():
if not other.transpose(-2, -1).is_contiguous():
other = other.contiguous()
return MatmulFunction.apply(input, other, out, transa, transb, alpha, beta)

82
ixformer_sdk/train/functions/residual_bias.py Normal file
View File

@@ -0,0 +1,82 @@
from typing import Union
import ixformer._C as ops
import torch
from torch.autograd.function import Function, FunctionCtx
import ixformer
__all__ = ["residual_bias"]
class ResidualBiasFunction(Function):
@staticmethod
def forward(
ctx,
input: torch.Tensor,
residual: torch.Tensor,
bias: torch.Tensor,
output: torch.Tensor,
alpha=1,
):
if output is None:
output = torch.empty_like(input)
if alpha is None:
alpha = 1
if bias is not None:
ops.train.add_residual_bias_forward(input, residual, bias, alpha, output)
else:
ops.train.add_residual_bias_forward(input, residual, alpha, output)
ctx.has_bias = bias is not None
ctx.alpha = alpha
return output
@staticmethod
def backward(ctx: FunctionCtx, grad_output):
grad_input = torch.empty_like(grad_output)
grad_residual = torch.empty_like(grad_output)
if ctx.has_bias:
grad_bias = torch.empty(
[grad_output.size(-1)],
dtype=grad_output.dtype,
device=grad_output.device,
)
ops.train.add_residual_bias_backward(
grad_output,
grad_input,
grad_residual,
grad_bias,
ctx.alpha,
)
return (grad_input, grad_residual, grad_bias, None, None)
else:
ops.train.add_residual_bias_backward(
grad_output,
grad_input,
grad_residual,
ctx.alpha,
)
return (grad_input, grad_residual, None, None, None)
def residual_bias(
input: torch.Tensor,
residual: torch.Tensor,
bias: torch.Tensor = None,
output: torch.Tensor = None,
alpha=1,
):
"""
等价实现：
input = residual.float() * alpha + input.float() + bias.float()
参数说明:
Args:
input: shape:[batch_count, seq_len, hidden_size],dtype:[torch.half]
residual: shape:[batch_count, seq_len, hidden_size],dtype:[torch.half]
bias: shape:[hidden_size],dtype:[torch.half]
alpha: float
return:
output: shape:[batch_count, seq_len, hidden_size],dtype:[torch.half]
"""
return ResidualBiasFunction.apply(input, residual, bias, output, alpha)

170
ixformer_sdk/train/functions/residual_bias_ln.py Normal file
View File

@@ -0,0 +1,170 @@
from typing import Union
import ixformer._C as ops
import torch
from torch.autograd.function import Function, FunctionCtx
__all__ = ["residual_bias_ln"]
class ResidualBiasLnFunction(Function):
@staticmethod
def forward(
ctx,
input: torch.Tensor,
residual: torch.Tensor,
bias: torch.Tensor,
ln_weight: torch.Tensor,
ln_bias: torch.Tensor,
output: torch.Tensor,
alpha=1,
is_post_ln=True,
):
norm_size = ln_weight.size(-1)
mean_size = input.numel() // norm_size
input_hat = torch.empty_like(input)
rstd = torch.empty([mean_size], dtype=input.dtype, device=input.device)
if bias is not None:
ops.train.add_residual_bias_ln_training_forward(
input,
residual,
bias,
ln_weight,
ln_bias,
alpha,
is_post_ln,
output,
input_hat,
rstd,
)
else:
ops.train.add_residual_bias_ln_training_forward(
input,
residual,
ln_weight,
ln_bias,
alpha,
is_post_ln,
output,
input_hat,
rstd,
)
ctx.norm_size = norm_size
ctx.has_bias = bias is not None
ctx.alpha = alpha
ctx.save_for_backward(input_hat, rstd, ln_weight)
return output
@staticmethod
def backward(ctx: FunctionCtx, grad_output):
input_hat, rstd_data, ln_weight = ctx.saved_tensors
grad_input = torch.empty_like(input_hat)
grad_residual = torch.empty_like(input_hat)
grad_ln_weight = torch.empty_like(ln_weight)
grad_ln_bias = torch.empty_like(ln_weight)
if ctx.has_bias:
grad_bias = torch.empty_like(ln_weight)
ops.train.add_residual_bias_ln_backward(
input_hat,
rstd_data,
ln_weight,
grad_output,
grad_ln_weight,
grad_ln_bias,
grad_input,
grad_residual,
grad_bias,
ctx.alpha,
)
return (
grad_input,
grad_residual,
grad_bias,
grad_ln_weight,
grad_ln_bias,
None,
None,
None,
None,
)
else:
ops.train.add_residual_bias_ln_backward(
input_hat,
rstd_data,
ln_weight,
grad_output,
grad_ln_weight,
grad_ln_bias,
grad_input,
grad_residual,
ctx.alpha,
)
return (
grad_input,
grad_residual,
None,
grad_ln_weight,
grad_ln_bias,
None,
None,
None,
None,
)
def residual_bias_ln(
input: torch.Tensor,
residual: torch.Tensor,
bias: torch.Tensor,
ln_weight: torch.Tensor,
ln_bias: torch.Tensor,
alpha=1,
is_post_ln=True,
output: torch.Tensor = None,
training: bool = False,
):
"""
等价实现：
input = residual.float() * alpha + input.float() + bias.float()
output = torch.nn.functional.layer_norm(
input, [input.shape[-1]], ln_weight.float(), ln_bias.float(), eps=1e-5)
参数说明:
Args:
input: shape:[batch_count * seq_len, hidden_size],dtype:[torch.half]
residual: shape:[batch_count * seq_len, hidden_size],dtype:[torch.half]
bias: shape:[hidden_size],dtype:[torch.half]
ln_weight:shape:[hidden_size],,dtype:[torch.half]
ln_bias:shape:[hidden_size],,dtype:[torch.half]
alpha: float
is_post_ln: bool, 是否应用layernorm 后处理
return:
output: shape:[batch_count * seq_len, hidden_size],dtype:[torch.half]
"""
if alpha is None:
alpha = 1
if not is_post_ln:
raise NotImplementedError()
if ln_weight is None or ln_bias is None:
raise NotImplementedError()
if output is None:
output = torch.empty_like(input)
if not training:
if bias is not None:
ops.infer.add_residual_bias_ln_forward(
input, residual, bias, ln_weight, ln_bias, alpha, is_post_ln, output
)
else:
ops.infer.add_residual_bias_ln_forward(
input, residual, ln_weight, ln_bias, alpha, is_post_ln, output
)
return output
else:
return ResidualBiasLnFunction.apply(
input, residual, bias, ln_weight, ln_bias, output, alpha, is_post_ln
)

324
ixformer_sdk/train/functions/rms_norm.py Normal file
View File

@@ -0,0 +1,324 @@
import numbers
from typing import Union
import ixformer._C as ops
import torch
from torch.nn import init
from torch.nn.parameter import Parameter
# apex interface for trainning add by xuelu.peng 2024/04/07
class FusedRMSNormAffineFunction(torch.autograd.Function):
@staticmethod
def forward(ctx, input, weight, normalized_shape, eps, memory_efficient=False, gradient_accumulation_fusion=False):
ctx.normalized_shape = normalized_shape
ctx.eps = eps
ctx.memory_efficient = memory_efficient
ctx.gradient_accumulation_fusion = gradient_accumulation_fusion
input_ = input.contiguous()
weight_ = weight.contiguous()
output = torch.empty_like(input_)
normalized_shape_size = len(normalized_shape)
assert normalized_shape_size == 1 # 目前只支持normalized_shape_size=1
invvar = torch.empty(
input_.shape[:-normalized_shape_size],
dtype=torch.float,
device=input_.device,
)
ops.train.rms_norm_forward_training(input_, weight_, output, invvar, ctx.eps)
ctx.save_for_backward(input_, weight_, invvar)
return output
@staticmethod
def backward(ctx, grad_output):
input_, weight_, invvar = ctx.saved_tensors
if ctx.gradient_accumulation_fusion:
if weight_.grad == None:
weight_.grad = torch.zeros_like(weight_)
grad_weight = weight_.grad
else:
grad_weight = torch.zeros_like(weight_) # 支持权重梯度累积融合，使用zeros_like，而不是emtpy_like 。
grad_input = torch.empty_like(input_)
if input_.numel() < 4096 * 8192:
ops.train.rms_norm_backward_training(
input_, invvar, weight_, grad_output, grad_weight, grad_input
)
else: ##llama 34b
ops.train.rms_norm_backward_training_opt(
input_, invvar, weight_, grad_output, grad_weight, grad_input
)
if ctx.gradient_accumulation_fusion:
grad_weight = None
return grad_input, grad_weight, None, None, None, None
def fused_rms_norm_affine(
input, weight, normalized_shape, eps=1e-6, memory_efficient=False, gradient_accumulation_fusion = False
):
return FusedRMSNormAffineFunction.apply(
input, weight, normalized_shape, eps, memory_efficient, gradient_accumulation_fusion
)
class FusedRMSNorm(torch.nn.Module):
r"""Applies RMS Normalization over a mini-batch of inputs
Currently only runs on cuda() tensors.
.. math::
y = \frac{x}{\mathrm{RMS}[x]} * \gamma
The root-mean-square is calculated separately over the last
certain number dimensions which have to be of the shape specified by
:attr:`normalized_shape`.
:math:`\gamma` is a learnable affine transform parameter of
:attr:`normalized_shape` if :attr:`elementwise_affine` is ``True``.
`epsilon` is added to the mean-square, then the root of the sum is taken.
.. note::
Unlike Batch Normalization and Instance Normalization, which applies
scalar scale and bias for each entire channel/plane with the
:attr:`affine` option, RMS Normalization applies per-element scale
with :attr:`elementwise_affine`.
This layer uses statistics computed from input data in both training and
evaluation modes.
Args:
normalized_shape (int or list or torch.Size): input shape from an expected input
of size
.. math::
[* \times \text{normalized}\_\text{shape}[0] \times \text{normalized}\_\text{shape}[1]
\times \ldots \times \text{normalized}\_\text{shape}[-1]]
If a single integer is used, it is treated as a singleton list, and this module will
normalize over the last dimension which is expected to be of that specific size.
eps: a value added to the denominator for numerical stability. Default: 1e-5
elementwise_affine: a boolean value that when set to ``True``, this module
has learnable per-element affine parameters initialized to ones (for weights)
and zeros (for biases). Default: ``True``.
Shape:
- Input: :math:`(N, *)`
- Output: :math:`(N, *)` (same shape as input)
Examples::
>>> input = torch.randn(20, 5, 10, 10)
>>> # With Learnable Parameters
>>> m = ixformer.FusedRMSNorm(10)
>>> # Without Learnable Parameters
>>> m = ixformer.FusedRMSNorm(input.size()[1:], elementwise_affine=False)
>>> # Normalize over last dimension of size 10 #目前只支持在最后一维norm
>>> m = ixformer.FusedRMSNorm(10)
>>> # Activating the module
>>> output = m(input)
.. _`Root Mean Square Layer Normalization`: https://arxiv.org/pdf/1910.07467.pdf
"""
def __init__(
self,
normalized_shape,
eps=1e-5,
elementwise_affine=True,
memory_efficient=False,
gradient_accumulation_fusion=False
):
super().__init__()
if isinstance(normalized_shape, numbers.Integral):
normalized_shape = (normalized_shape,)
self.normalized_shape = torch.Size(normalized_shape)
self.eps = eps
self.elementwise_affine = elementwise_affine
self.memory_efficient = memory_efficient
self.gradient_accumulation_fusion = gradient_accumulation_fusion
if self.elementwise_affine:
self.weight = Parameter(torch.empty(*normalized_shape))
else:
self.register_parameter("weight", None)
self.reset_parameters()
def reset_parameters(self):
if self.elementwise_affine:
init.ones_(self.weight)
def forward(self, input):
if torch.jit.is_tracing() or torch.jit.is_scripting() or not input.is_cuda:
raise NotImplementedError()
if self.elementwise_affine:
return fused_rms_norm_affine(
input,
self.weight,
self.normalized_shape,
self.eps,
self.memory_efficient,
self.gradient_accumulation_fusion
)
else:
raise NotImplementedError()
def extra_repr(self):
return "{normalized_shape}, eps={eps}, " "elementwise_affine={elementwise_affine}".format(**self.__dict__)
class FusedRMSNormResFunction(torch.autograd.Function):
@staticmethod
def forward(ctx, input, weight, residual, normalized_shape, eps, gradient_accumulation_fusion=False, memory_efficient=False):
ctx.normalized_shape = normalized_shape
ctx.eps = eps
ctx.memory_efficient = memory_efficient
ctx.gradient_accumulation_fusion = gradient_accumulation_fusion
input_ = input.contiguous()
weight_ = weight.contiguous()
output = torch.empty_like(input_)
normalized_shape_size=len(normalized_shape)
assert normalized_shape_size == 1 #目前只支持normalized_shape_size=1
invvar = torch.empty(input_.shape[:-normalized_shape_size], dtype=torch.float, device=input_.device)
if residual is not None:
ctx.input_res = True
out_res = torch.empty_like(input_)
ops.train.rms_norm_res_forward_training(input_, weight_, output, invvar, ctx.eps, residual, out_res)
else:
ctx.input_res = False
ops.train.rms_norm_forward_training(input_, weight_, output, invvar, ctx.eps)
out_res = input_
# input_res 为 True 时 LN 的 input 为 input+redidual
ctx.save_for_backward(out_res, weight_, invvar)
return output, out_res
@staticmethod
def backward(ctx, grad_output, grad_out_res):
input_, weight_, invvar = ctx.saved_tensors
if ctx.gradient_accumulation_fusion:
if weight_.grad == None:
weight_.grad = torch.zeros_like(weight_)
grad_weight = weight_.grad
else:
grad_weight = torch.zeros_like(weight_) # 算子kernel 支持权重梯度累积融合，使用zeros_like，而不是emtpy_like 。
grad_input = torch.empty_like(input_)
# rms_norm_res_backward_training 本身支持权重梯度累积融合，当不进行融合时，其输入 grad_weight 必须为 zero_like 。
if input_.numel()< 4096*8192:
ops.train.rms_norm_res_backward_training(input_, invvar, weight_,
grad_output, grad_weight, grad_input, grad_out_res)
else:##llama 34b
ops.train.rms_norm_res_backward_training_opt(input_,invvar, weight_,
grad_output,grad_weight,grad_input,grad_out_res)
if ctx.input_res:
grad_res = grad_input
else:
grad_res = None
if ctx.gradient_accumulation_fusion:
grad_weight = None
return grad_input, grad_weight, grad_res, None, None, None, None
class FusedRMSNormRes(torch.nn.Module):
r"""Applies RMS Normalization and resdiual over a mini-batch of inputs, RMS Normalization part comes from FusedRMSNorm.
Currently only runs on cuda() tensors.
.. math::
y = \frac{x}{\mathrm{RMS}[x]} * \gamma
if residual None, x is input and output is equal to x, otherwise, x is input+residual and out_res is equal to x.
The root-mean-square is calculated separately over the last
certain number dimensions which have to be of the shape specified by
:attr:`normalized_shape`.
:math:`\gamma` is a learnable affine transform parameter of
:attr:`normalized_shape` if :attr:`elementwise_affine` is ``True``.
`epsilon` is added to the mean-square, then the root of the sum is taken.
.. note::
Unlike Batch Normalization and Instance Normalization, which applies
scalar scale and bias for each entire channel/plane with the
:attr:`affine` option, RMS Normalization applies per-element scale
with :attr:`elementwise_affine`.
This layer uses statistics computed from input data in both training and
evaluation modes.
Args:
normalized_shape (int or list or torch.Size): input shape from an expected input
of size
.. math::
[* \times \text{normalized}\_\text{shape}[0] \times \text{normalized}\_\text{shape}[1]
\times \ldots \times \text{normalized}\_\text{shape}[-1]]
If a single integer is used, it is treated as a singleton list, and this module will
normalize over the last dimension which is expected to be of that specific size.
eps: a value added to the denominator for numerical stability. Default: 1e-5
elementwise_affine: a boolean value that when set to ``True``, this module
has learnable per-element affine parameters initialized to ones (for weights)
and zeros (for biases). Default: ``True``.
Shape:
- Input: :math:`(N, *)`
- residual: :math:`(N, *)` (if not None)
- Output: :math:`(N, *)` (same shape as input)
- out_res: :math:`(N, *)`
Examples::
>>> input = torch.randn(20, 5, 10, 10)
>>> res = torch.randn(20, 5, 10, 10)
>>> # With Learnable Parameters
>>> m = ixformer.FusedRMSNorm(10)
>>> # Without Learnable Parameters
>>> m = ixformer.FusedRMSNorm(input.size()[1:], elementwise_affine=False)
>>> # Normalize over last dimension of size 10 #目前只支持在最后一维norm
>>> m = ixformer.FusedRMSNorm(10)
>>> # Activating the module
>>> output, output_res = m(input, res)
.. _`Root Mean Square Layer Normalization`: https://arxiv.org/pdf/1910.07467.pdf
"""
def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True, memory_efficient=False, gradient_accumulation_fusion=False):
super().__init__()
if isinstance(normalized_shape, numbers.Integral):
normalized_shape = (normalized_shape,)
self.normalized_shape = torch.Size(normalized_shape)
self.eps = eps
self.elementwise_affine = elementwise_affine
self.gradient_accumulation_fusion = gradient_accumulation_fusion
self.memory_efficient = memory_efficient
if self.elementwise_affine:
self.weight = Parameter(torch.empty(*normalized_shape))
else:
self.register_parameter("weight", None)
self.reset_parameters()
def reset_parameters(self):
if self.elementwise_affine:
init.ones_(self.weight)
def forward(self, input, residual=None):
if torch.jit.is_tracing() or torch.jit.is_scripting() or not input.is_cuda:
raise NotImplementedError()
if self.elementwise_affine:
return FusedRMSNormResFunction.apply(input, self.weight, residual, self.normalized_shape, self.eps, self.gradient_accumulation_fusion, self.memory_efficient)
else:
raise NotImplementedError()
def extra_repr(self):
return "{normalized_shape}, eps={eps}, " "elementwise_affine={elementwise_affine}".format(**self.__dict__)

45
ixformer_sdk/train/functions/swiglu.py Normal file
View File

@@ -0,0 +1,45 @@
from typing import Union
import ixformer._C as ops
import torch
from torch.autograd.function import Function, FunctionCtx
__all__ = ["swiglu"]
class SwigluFunction(Function):
@staticmethod
def forward(ctx, input):
output_shape = list(input.shape)
output_shape[-1] = output_shape[-1] // 2
output = input.new_empty(output_shape)
ops.train.swiglu_training_forward(input, output)
ctx.save_for_backward(input)
return output
@staticmethod
def backward(ctx: FunctionCtx, grad_output):
input = ctx.saved_tensors[0]
grad_input = torch.empty_like(input)
ops.train.swiglu_training_backward(input, grad_output, grad_input)
return grad_input
def swiglu(input):
"""
等价实现：
def ref_silu_and_mul(x: torch.Tensor) -> torch.Tensor:
dtype = x.dtype
x = x.float()
x1, x2 = x.chunk(chunks=2, dim=-1)
res = torch.nn.functional.silu(x1) * x2
return res.to(dtype)
参数说明:
Args:
input: dtype:torch.float, torch.half, torch.bfloat16
return:
output: dtype:torch.float, torch.half, torch.bfloat16
"""
return SwigluFunction.apply(input)