under test, not sure no errors

2026-09-02 07:01:29 +00:00
commit 9f4fcd5039
4224 changed files with 1015180 additions and 0 deletions
--- a/ixformer_sdk/train/functions/residual_bias.py
+++ b/ixformer_sdk/train/functions/residual_bias.py
@@ -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)