Skip to main content

layer_norm_backward

Function layer_norm_backward 

Source
pub fn layer_norm_backward<T: Triton, D: Float, const BLOCK_N: i32>(
    dy_ptr: T::Pointer<D>,
    x_ptr: T::Pointer<D>,
    dx_ptr: T::Pointer<D>,
    weight_ptr: T::Pointer<D>,
    dweight_ptr: T::Pointer<D>,
    dbias_ptr: T::Pointer<D>,
    mean_ptr: T::Pointer<D>,
    rstd_ptr: T::Pointer<D>,
    _M: i32,
    N: i32,
)
where T::I32Tensor: Tensor<i32, 1> + Comparison<i32, BoolTensor = T::BoolTensor>, T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
Expand description

Backward pass for LayerNorm.

Given saved mean and rstd from the forward pass:

xhat[m,n]    = (x[m,n] - mean[m]) * rstd[m]
dweight[n]   = Σ_m dy[m,n] * xhat[m,n]
dbias[n]     = Σ_m dy[m,n]
dx[m,n]      = rstd[m] * γ[n] * (dy[m,n]
                 - (Σ_n dy[m,n]*γ[n]) / N
                 - xhat[m,n] * (Σ_n dy[m,n]*γ[n]*xhat[m,n]) / N)

Grid: [M] — one CTA per row.