project_6/ixformer_sdk/train/functions/rms_norm.py

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__)