719 lines
22 KiB
Python
719 lines
22 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Optional, Tuple
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import triton
|
|
import triton.language as tl
|
|
|
|
# Try to load the C++ extension
|
|
try:
|
|
from fastvideo_kernel._C import fastvideo_kernel_ops
|
|
quant_cuda = getattr(fastvideo_kernel_ops, "quant_cuda", None)
|
|
gemm_cuda = getattr(fastvideo_kernel_ops, "gemm_cuda", None)
|
|
rms_norm_cuda = getattr(fastvideo_kernel_ops, "rms_norm_cuda", None)
|
|
layer_norm_cuda = getattr(fastvideo_kernel_ops, "layer_norm_cuda", None)
|
|
except ImportError:
|
|
quant_cuda = None
|
|
gemm_cuda = None
|
|
rms_norm_cuda = None
|
|
layer_norm_cuda = None
|
|
|
|
|
|
def int8_quant(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Quantize a floating-point tensor to int8 using a custom CUDA kernel.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor of type float16/bfloat16.
|
|
|
|
Returns:
|
|
Tuple[torch.Tensor, torch.Tensor]:
|
|
- x_q: Quantized int8 tensor.
|
|
- x_scale: Per-block scale tensor used for quantization.
|
|
"""
|
|
x_q, x_scale = quant_cuda(x, None, None)
|
|
return x_q, x_scale
|
|
|
|
|
|
def int8_linear(
|
|
x: torch.Tensor,
|
|
w_q: torch.Tensor,
|
|
w_s: torch.Tensor,
|
|
**kwargs,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Perform an int8 GEMM (matrix multiplication) using quantized weights and a
|
|
quantized version of the input. The underlying compute is performed by a
|
|
custom CUDA kernel.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input activation of shape (M, K) in float32.
|
|
w_q (torch.Tensor): Quantized int8 weight tensor of shape (N, K).
|
|
w_s (torch.Tensor): Scale tensor associated with w_q.
|
|
**kwargs: Additional options (reserved for future use).
|
|
|
|
Returns:
|
|
torch.Tensor: Output tensor of shape (M, N) in float32.
|
|
"""
|
|
assert w_q.dtype == torch.int8, "Weight tensor must be int8."
|
|
shape = x.shape
|
|
x = x.reshape(-1, shape[-1])
|
|
m = x.shape[0]
|
|
n = w_q.shape[0]
|
|
y = torch.zeros(m, n, dtype=x.dtype, device=x.device)
|
|
|
|
x_q, x_s = int8_quant(x)
|
|
gemm_cuda(x_q, x_s, w_q, w_s, y)
|
|
return y.reshape(*shape[:-1], n)
|
|
|
|
|
|
def flatten_if_batched(*tensors):
|
|
"""
|
|
Flattens all input tensors from (B, N, D_i) to (B * N, D_i) if they are batched (3D).
|
|
|
|
Args:
|
|
*tensors: Any number of input tensors, each must have shape (B, N, D_i) or (N, D_i)
|
|
|
|
Returns:
|
|
flat_tensors: List of flattened tensors
|
|
batched: Boolean flag indicating whether inputs were batched
|
|
batch_size: Batch size if batched, else None
|
|
"""
|
|
if not tensors:
|
|
raise ValueError("At least one tensor must be provided.")
|
|
|
|
first = tensors[0]
|
|
assert len(first.shape) in [
|
|
2,
|
|
3,
|
|
], "Input tensors must be batched (3D) or not batched (2D)"
|
|
|
|
if len(first.shape) == 3: # batched
|
|
batched = True
|
|
batch_size = first.shape[0]
|
|
assert all(t.shape[0] == batch_size for t in tensors), "All input tensors must have the same batch size"
|
|
assert all(t.shape[1] == first.shape[1]
|
|
for t in tensors), "All input tensors must have the same sequence length"
|
|
flat_tensors = [t.reshape(-1, t.shape[-1]) for t in tensors]
|
|
else:
|
|
batched = False
|
|
batch_size = None
|
|
flat_tensors = list(tensors)
|
|
|
|
return flat_tensors, batched, batch_size
|
|
|
|
|
|
@triton.jit
|
|
def _rms_norm_fwd_fused(
|
|
X,
|
|
Y,
|
|
W,
|
|
Rstd,
|
|
x_stride,
|
|
y_stride,
|
|
N: tl.constexpr, # number of columns in X,
|
|
N2: tl.constexpr,
|
|
eps, # epsilon to avoid division by zero
|
|
BLOCK_M: tl.constexpr,
|
|
):
|
|
# Map the program id to the row of X and Y it should compute.
|
|
pid = tl.program_id(0)
|
|
rows = pid * BLOCK_M + tl.arange(0, BLOCK_M)
|
|
cols = tl.arange(0, N2)
|
|
mask = cols < N
|
|
|
|
x_ptr = X + rows[:, None] * x_stride + cols[None, :]
|
|
y_ptr = Y + rows[:, None] * y_stride + cols[None, :]
|
|
|
|
x = tl.load(x_ptr, mask=mask[None, :], other=0.0).to(tl.float32)
|
|
|
|
# Compute variance
|
|
_var = x * x
|
|
var = tl.sum(_var, axis=1) / N
|
|
rstd = 1 / tl.sqrt(var + eps)
|
|
|
|
# Write mean / rstd
|
|
tl.store(Rstd + rows, rstd)
|
|
rstd = tl.reshape(rstd, (BLOCK_M, 1))
|
|
|
|
# Normalize and apply linear transformation
|
|
w = tl.load(W + cols)
|
|
x_hat = x * rstd
|
|
y = x_hat * w
|
|
|
|
# Write output
|
|
y = y.to(Y.type.element_ty)
|
|
tl.store(y_ptr, y, mask=mask[None, :])
|
|
|
|
|
|
def rmsnorm(x, w, eps):
|
|
"""
|
|
Forward pass of the RMSNorm.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor, High precision.
|
|
w (torch.Tensor): RMSNorm weight tensor.
|
|
eps (float): RMSNorm epsilon value.
|
|
|
|
Returns:
|
|
y (torch.Tensor): Output tensor, High precision.
|
|
rstd (torch.Tensor): Inverse standard deviation, needed for backward.
|
|
"""
|
|
assert x.is_contiguous(), "Input must be contiguous"
|
|
# Change batched 3D input to 2D
|
|
[x], batched, BS = flatten_if_batched(x)
|
|
|
|
# allocate output
|
|
M, N = x.shape
|
|
y = torch.empty_like(x, dtype=x.dtype)
|
|
rstd = torch.empty((M, ), dtype=torch.float32, device=x.device)
|
|
|
|
# heuristics for number of warps
|
|
num_warps = 8
|
|
|
|
# Avoid illegal memory access
|
|
N2 = triton.next_power_of_2(N)
|
|
|
|
if N <= 512:
|
|
BLOCK_M = 32
|
|
else:
|
|
BLOCK_M = 1
|
|
|
|
# Call the triton kernel
|
|
_rms_norm_fwd_fused[(triton.cdiv(M, BLOCK_M), )]( #
|
|
x,
|
|
y,
|
|
w,
|
|
rstd, #
|
|
x.stride(0),
|
|
y.stride(0),
|
|
N,
|
|
N2,
|
|
eps,
|
|
num_warps=num_warps,
|
|
BLOCK_M=BLOCK_M,
|
|
)
|
|
|
|
# Recover 2D to 3D
|
|
if batched:
|
|
y = y.reshape(BS, -1, y.shape[-1])
|
|
|
|
return y, rstd
|
|
|
|
|
|
@triton.jit
|
|
def _layer_norm_param_fwd_fused(
|
|
X, # pointer to the input
|
|
Y, # pointer to the output
|
|
W, # pointer to the weights
|
|
B, # pointer to the biases
|
|
Mean, # pointer to the mean
|
|
Rstd, # pointer to the 1/std
|
|
x_stride, # how much to increase the pointer when moving by 1 row
|
|
y_stride, # how much to increase the pointer when moving by 1 row
|
|
N: tl.constexpr, # number of columns in X,
|
|
N2: tl.constexpr, # number of columns in X,
|
|
eps, # epsilon to avoid division by zero
|
|
BLOCK_M: tl.constexpr,
|
|
):
|
|
# Map the program id to the row of X and Y it should compute.
|
|
pid = tl.program_id(0)
|
|
rows = pid * BLOCK_M + tl.arange(0, BLOCK_M)
|
|
cols = tl.arange(0, N2)
|
|
mask = cols < N
|
|
|
|
x_ptr = X + rows[:, None] * x_stride + cols[None, :]
|
|
y_ptr = Y + rows[:, None] * y_stride + cols[None, :]
|
|
|
|
x = tl.load(x_ptr, mask=mask[None, :], other=0.0).to(tl.float32)
|
|
|
|
# Compute mean and Variance
|
|
mean = tl.sum(x, axis=1, keep_dims=True) / N
|
|
# Compute variance
|
|
_var = (x - mean) * (x - mean)
|
|
var = tl.sum(_var, axis=1, keep_dims=True) / N
|
|
rstd = 1 / tl.sqrt(var + eps)
|
|
|
|
# Write mean / rstd
|
|
_mean = tl.reshape(mean, (BLOCK_M))
|
|
_rstd = tl.reshape(rstd, (BLOCK_M))
|
|
tl.store(Mean + rows, _mean)
|
|
tl.store(Rstd + rows, _rstd)
|
|
|
|
# Normalize and apply linear transformation
|
|
x_hat = (x - mean) * rstd
|
|
|
|
w = tl.load(W + cols)
|
|
b = tl.load(B + cols)
|
|
|
|
x_hat = x_hat * w + b
|
|
|
|
# Write output
|
|
x_hat = x_hat.to(Y.type.element_ty)
|
|
tl.store(y_ptr, x_hat, mask=mask[None, :])
|
|
|
|
|
|
def layernorm_param(x, w, b, eps):
|
|
# Change batched 3D input to 2D
|
|
[x], batched, BS = flatten_if_batched(x)
|
|
|
|
# allocate output
|
|
M, N = x.shape
|
|
y = torch.empty_like(x, dtype=torch.float32)
|
|
mean = torch.empty((M, ), dtype=torch.float32, device=x.device)
|
|
rstd = torch.empty((M, ), dtype=torch.float32, device=x.device)
|
|
# heuristics for number of warps
|
|
num_warps = 8
|
|
|
|
N2 = triton.next_power_of_2(N)
|
|
|
|
if N <= 512:
|
|
BLOCK_M = 32
|
|
else:
|
|
BLOCK_M = 1
|
|
|
|
# enqueue kernel
|
|
_layer_norm_param_fwd_fused[(triton.cdiv(M, BLOCK_M), )]( #
|
|
x,
|
|
y,
|
|
w,
|
|
b,
|
|
mean,
|
|
rstd, #
|
|
x.stride(0),
|
|
y.stride(0),
|
|
N,
|
|
N2,
|
|
eps,
|
|
num_warps=num_warps,
|
|
BLOCK_M=BLOCK_M,
|
|
)
|
|
|
|
# Recover 2D to 3D
|
|
if batched:
|
|
y = y.reshape(BS, -1, y.shape[-1])
|
|
|
|
return y, mean, rstd
|
|
|
|
|
|
########################################################
|
|
# Elementwise_affine=False
|
|
########################################################
|
|
|
|
|
|
@triton.jit
|
|
def _layer_norm_noparam_fwd_fused(
|
|
X, # pointer to the input
|
|
Y, # pointer to the output
|
|
Mean, # pointer to the mean
|
|
Rstd, # pointer to the 1/std
|
|
x_stride, # how much to increase the pointer when moving by 1 row
|
|
y_stride, # how much to increase the pointer when moving by 1 row
|
|
N: tl.constexpr, # number of columns in X,
|
|
N2: tl.constexpr, # number of columns in X,
|
|
eps, # epsilon to avoid division by zero
|
|
BLOCK_M: tl.constexpr,
|
|
):
|
|
# Map the program id to the row of X and Y it should compute.
|
|
pid = tl.program_id(0)
|
|
rows = pid * BLOCK_M + tl.arange(0, BLOCK_M)
|
|
cols = tl.arange(0, N2)
|
|
mask = cols < N
|
|
|
|
x_ptr = X + rows[:, None] * x_stride + cols[None, :]
|
|
y_ptr = Y + rows[:, None] * y_stride + cols[None, :]
|
|
|
|
x = tl.load(x_ptr, mask=mask[None, :], other=0.0).to(tl.float32)
|
|
|
|
# Compute mean and Variance
|
|
mean = tl.sum(x, axis=1, keep_dims=True) / N
|
|
# Compute variance
|
|
_var = (x - mean) * (x - mean)
|
|
var = tl.sum(_var, axis=1, keep_dims=True) / N
|
|
rstd = 1 / tl.sqrt(var + eps)
|
|
|
|
# Write mean / rstd
|
|
_mean = tl.reshape(mean, (BLOCK_M))
|
|
_rstd = tl.reshape(rstd, (BLOCK_M))
|
|
tl.store(Mean + rows, _mean)
|
|
tl.store(Rstd + rows, _rstd)
|
|
|
|
# Normalize and apply linear transformation
|
|
x_hat = (x - mean) * rstd
|
|
|
|
# Write output
|
|
x_hat = x_hat.to(Y.type.element_ty)
|
|
tl.store(y_ptr, x_hat, mask=mask[None, :])
|
|
|
|
|
|
def layernorm_noparam(x, eps):
|
|
assert x.is_contiguous(), "Input must be contiguous"
|
|
|
|
# Change batched 3D input to 2D
|
|
[x], batched, BS = flatten_if_batched(x)
|
|
|
|
# allocate output
|
|
M, N = x.shape
|
|
y = torch.empty_like(x, dtype=torch.float32)
|
|
mean = torch.empty((M, ), dtype=torch.float32, device=x.device)
|
|
rstd = torch.empty((M, ), dtype=torch.float32, device=x.device)
|
|
# heuristics for number of warps
|
|
num_warps = 8
|
|
|
|
N2 = triton.next_power_of_2(N)
|
|
|
|
if N <= 512:
|
|
BLOCK_M = 32
|
|
else:
|
|
BLOCK_M = 1
|
|
|
|
# enqueue kernel
|
|
_layer_norm_noparam_fwd_fused[(triton.cdiv(M, BLOCK_M), )]( #
|
|
x,
|
|
y,
|
|
mean,
|
|
rstd, #
|
|
x.stride(0),
|
|
y.stride(0),
|
|
N,
|
|
N2,
|
|
eps,
|
|
num_warps=num_warps,
|
|
BLOCK_M=BLOCK_M,
|
|
)
|
|
|
|
# Recover 2D to 3D
|
|
if batched:
|
|
y = y.reshape(BS, -1, y.shape[-1])
|
|
|
|
return y, mean, rstd
|
|
|
|
|
|
def layernorm(x, w, b, eps, elementwise_affine=True):
|
|
if elementwise_affine:
|
|
assert w is not None and b is not None
|
|
return layernorm_param(x, w, b, eps)
|
|
else:
|
|
assert w is None and b is None
|
|
return layernorm_noparam(x, eps)
|
|
|
|
|
|
def cdiv(a: int, b: int):
|
|
return (a + b - 1) // b
|
|
|
|
|
|
def dequantize_weight(w_q, w_s, dtype):
|
|
# w_q: (N, K) int8
|
|
# w_s: (NB, KB) float32
|
|
# Block size 128
|
|
BLOCK = 128
|
|
N, K = w_q.shape
|
|
# Expand w_s
|
|
# Repeat interleave
|
|
w_s_exp = w_s.repeat_interleave(BLOCK, dim=0).repeat_interleave(BLOCK, dim=1)
|
|
# Crop
|
|
w_s_exp = w_s_exp[:N, :K]
|
|
return w_q.to(dtype) * w_s_exp.to(dtype)
|
|
|
|
|
|
class Int8LinearFunction(torch.autograd.Function):
|
|
|
|
@staticmethod
|
|
def forward(ctx, x, w_q, w_s, bias=None):
|
|
ctx.save_for_backward(x, w_q, w_s, bias)
|
|
ctx.bias_requires_grad = bias.requires_grad if bias is not None else False
|
|
|
|
out = int8_linear(x, w_q, w_s)
|
|
if bias is not None:
|
|
out = out + bias
|
|
return out
|
|
|
|
@staticmethod
|
|
def backward(ctx, grad_output):
|
|
x, w_q, w_s, bias = ctx.saved_tensors
|
|
|
|
grad_input = None
|
|
grad_weight = None
|
|
grad_scale = None
|
|
grad_bias = None
|
|
|
|
if ctx.needs_input_grad[0]:
|
|
# grad_input = grad_output @ W
|
|
w_float = dequantize_weight(w_q, w_s, grad_output.dtype)
|
|
grad_input = torch.matmul(grad_output, w_float)
|
|
|
|
if ctx.bias_requires_grad:
|
|
dim_to_sum = list(range(grad_output.dim() - 1))
|
|
grad_bias = grad_output.sum(dim=dim_to_sum)
|
|
|
|
return grad_input, grad_weight, grad_scale, grad_bias
|
|
|
|
|
|
class RMSNormFunction(torch.autograd.Function):
|
|
|
|
@staticmethod
|
|
def forward(ctx, x, w, eps):
|
|
y, rstd = rmsnorm(x, w, eps)
|
|
ctx.save_for_backward(x, w, rstd)
|
|
ctx.eps = eps
|
|
return y
|
|
|
|
@staticmethod
|
|
def backward(ctx, grad_output):
|
|
# x, w, rstd are saved
|
|
# grad_output is dL/dy
|
|
# We can implement backward using torch ops for correctness
|
|
x, w, rstd = ctx.saved_tensors
|
|
eps = ctx.eps
|
|
N = x.shape[-1]
|
|
|
|
# dL/dy * w
|
|
dx = grad_output * w
|
|
|
|
# Expand rstd
|
|
rstd = rstd.unsqueeze(-1) # (B, 1) or (M, 1)
|
|
if x.dim() == 3:
|
|
# Flatten if necessary or handle dimensions.
|
|
# forward flattens x to (M, N) internally but saved x might be 3D?
|
|
# rmsnorm takes x, flattens it. But does it modify x in place? No.
|
|
# But ctx.save_for_backward saves the original x (3D if input was 3D).
|
|
# rstd is (M,).
|
|
# We should flatten x and grad_output to match logic
|
|
x_flat = x.reshape(-1, N)
|
|
grad_output_flat = grad_output.reshape(-1, N)
|
|
dx = dx.reshape(-1, N)
|
|
else:
|
|
x_flat = x
|
|
grad_output_flat = grad_output
|
|
dx = dx
|
|
|
|
# Standard RMSNorm backward
|
|
# dy = grad_output_flat
|
|
# x_hat = x * rstd
|
|
# w * dy
|
|
# c1 = mean(dy * w * x) * rstd^2
|
|
# dx = (dy * w - c1 * x) * rstd
|
|
|
|
# More precise:
|
|
# y = x * rstd * w
|
|
# dy = grad_output
|
|
# dw = sum(dy * x * rstd, dim=0)
|
|
|
|
grad_w = None
|
|
if ctx.needs_input_grad[1]: # w
|
|
# grad_w = (grad_output_flat * x_flat * rstd).sum(dim=0)
|
|
# More accurate gradient for weight when considering rstd was computed from x
|
|
# Actually, for Affine part (w), it is just sum(dL/dy * x_hat).
|
|
# y = x_hat * w
|
|
# dL/dw = sum(dL/dy * x_hat)
|
|
x_hat = x_flat * rstd
|
|
grad_w = (grad_output_flat * x_hat).sum(dim=0)
|
|
|
|
grad_input = None
|
|
if ctx.needs_input_grad[0]: # x
|
|
# x_hat = x * rstd
|
|
# dx = rstd * (w * dy - mean(w * dy * x_hat) * x_hat)
|
|
# but for RMSNorm:
|
|
# dx = rstd * (w * dy - (x * rstd^2) * mean(w * dy * x)) -> check formula
|
|
|
|
# Using PyTorch autograd for reference logic:
|
|
# y = x / sqrt(mean(x^2) + eps) * w
|
|
# Let sigma = sqrt(...)
|
|
# dL/dx = dL/dy * w * (1/sigma) + dL/dsigma * dsigma/dx
|
|
# dsigma/dx = 1/(2*sigma) * 2x/N = x / (N * sigma)
|
|
# dL/dsigma = sum(dL/dy * w * x * (-1/sigma^2))
|
|
# dL/dx = (dL/dy * w)/sigma - sum(dL/dy * w * x) * x / (N * sigma^3)
|
|
# = (1/sigma) * [ (dL/dy * w) - x * sum(dL/dy * w * x) / (N * sigma^2) ]
|
|
# = rstd * [ (dL/dy * w) - x * rstd^2 * mean(dL/dy * w * x) ] # Wait, sum/N is mean
|
|
|
|
# Let's compute:
|
|
dy_w = grad_output_flat * w
|
|
# term2 = (dy_w * x_flat).sum(dim=-1, keepdim=True) / N # mean(dy*w*x)
|
|
# grad_input = rstd * (dy_w - x_flat * (rstd ** 2) * term2)
|
|
|
|
# Re-derivation:
|
|
# x_hat = x * rstd
|
|
# y = x_hat * w
|
|
# dL/dx = dL/dy * dy/dx
|
|
# dy/dx = w * dx_hat/dx
|
|
# dx_hat/dx = rstd * (I - x * x^T * rstd^2 / N) ?? No.
|
|
# dx_hat_i / dx_j = rstd * (delta_ij - x_i * x_j * rstd^2 / N)
|
|
|
|
# dL/dx_hat = dL/dy * w
|
|
# dL/dx = rstd * (dL/dx_hat - x * rstd^2 * mean(dL/dx_hat * x))
|
|
|
|
dx_hat = grad_output_flat * w
|
|
dx_hat_x_mean = (dx_hat * x_flat).mean(dim=-1, keepdim=True)
|
|
grad_input = rstd * (dx_hat - x_flat * (rstd**2) * dx_hat_x_mean)
|
|
|
|
if x.dim() == 3:
|
|
grad_input = grad_input.reshape(x.shape)
|
|
|
|
return grad_input, grad_w, None
|
|
|
|
|
|
class LayerNormFunction(torch.autograd.Function):
|
|
|
|
@staticmethod
|
|
def forward(ctx, x, w, b, eps, elementwise_affine):
|
|
if elementwise_affine:
|
|
y, mean, rstd = layernorm_param(x, w, b, eps)
|
|
ctx.save_for_backward(x, w, b, mean, rstd)
|
|
else:
|
|
y, mean, rstd = layernorm_noparam(x, eps)
|
|
ctx.save_for_backward(x, None, None, mean, rstd)
|
|
|
|
ctx.eps = eps
|
|
ctx.elementwise_affine = elementwise_affine
|
|
return y
|
|
|
|
@staticmethod
|
|
def backward(ctx, grad_output):
|
|
x, w, b, mean, rstd = ctx.saved_tensors
|
|
N = x.shape[-1]
|
|
|
|
if x.dim() == 3:
|
|
x_flat = x.reshape(-1, N)
|
|
grad_output_flat = grad_output.reshape(-1, N)
|
|
else:
|
|
x_flat = x
|
|
grad_output_flat = grad_output
|
|
|
|
if w is None:
|
|
w_eff = 1.0
|
|
else:
|
|
w_eff = w
|
|
|
|
# dx calculation
|
|
# x_hat = (x - mean) * rstd
|
|
# y = x_hat * w + b
|
|
# dL/dx_hat = dL/dy * w
|
|
|
|
dy = grad_output_flat
|
|
dx_hat = dy * w_eff
|
|
|
|
# dL/dvar = sum(dL/dx_hat * (x-mean) * (-0.5) * (var+eps)^(-1.5))
|
|
# = -0.5 * rstd^3 * sum(dx_hat * (x-mean))
|
|
# dL/dmean = sum(dL/dx_hat * (-rstd)) + dL/dvar * (-2/N) * sum(x-mean)
|
|
# term sum(x-mean) is 0. So second part vanishes.
|
|
# dL/dmean = -rstd * sum(dx_hat)
|
|
|
|
# dL/dx = dL/dx_hat * rstd + dL/dvar * 2(x-mean)/N + dL/dmean * 1/N
|
|
# = dx_hat * rstd + (-0.5 * rstd^3 * sum(dx_hat * (x-mean))) * 2(x-mean)/N + (-rstd * sum(dx_hat))/N
|
|
# = rstd * [ dx_hat - mean(dx_hat) - (x-mean)*rstd^2 * mean(dx_hat * (x-mean)) ]
|
|
|
|
x_centered = x_flat - mean.unsqueeze(-1)
|
|
dx_hat_mean = dx_hat.mean(dim=-1, keepdim=True)
|
|
dx_hat_x_centered_mean = (dx_hat * x_centered).mean(dim=-1, keepdim=True)
|
|
|
|
grad_input = rstd.unsqueeze(-1) * (dx_hat - dx_hat_mean - x_centered *
|
|
(rstd.unsqueeze(-1)**2) * dx_hat_x_centered_mean)
|
|
|
|
if x.dim() == 3:
|
|
grad_input = grad_input.reshape(x.shape)
|
|
|
|
grad_w = None
|
|
grad_b = None
|
|
|
|
if ctx.elementwise_affine:
|
|
if ctx.needs_input_grad[1]:
|
|
x_hat = x_centered * rstd.unsqueeze(-1)
|
|
grad_w = (dy * x_hat).sum(dim=0)
|
|
if ctx.needs_input_grad[2]:
|
|
grad_b = dy.sum(dim=0)
|
|
|
|
return grad_input, grad_w, grad_b, None, None
|
|
|
|
|
|
class Int8Linear(nn.Module):
|
|
|
|
def __init__(self, in_features, out_features, bias=True, dtype=torch.bfloat16):
|
|
super().__init__()
|
|
self.in_features = in_features
|
|
self.out_features = out_features
|
|
|
|
row_blocks = cdiv(out_features, b=128)
|
|
col_blocks = cdiv(in_features, b=128)
|
|
|
|
self.register_buffer("int8_weight", torch.empty((out_features, in_features), dtype=torch.int8))
|
|
self.register_buffer("scale", torch.empty((row_blocks, col_blocks), dtype=torch.float32))
|
|
if bias:
|
|
self.register_buffer("bias", torch.empty(out_features, dtype=dtype))
|
|
else:
|
|
self.bias = None
|
|
|
|
def forward(self, x):
|
|
return Int8LinearFunction.apply(x, self.int8_weight, self.scale, self.bias)
|
|
|
|
@classmethod
|
|
def from_linear(cls, original_linear: nn.Linear, quantize: bool = True):
|
|
|
|
int8_layer = cls(original_linear.in_features,
|
|
original_linear.out_features,
|
|
bias=original_linear.bias is not None,
|
|
dtype=original_linear.weight.dtype)
|
|
if quantize:
|
|
w_data = original_linear.weight.data.cuda()
|
|
int8_w, scale = int8_quant(w_data)
|
|
|
|
int8_layer.int8_weight.copy_(int8_w)
|
|
int8_layer.scale.copy_(scale)
|
|
if original_linear.bias is not None:
|
|
int8_layer.bias.data.copy_(original_linear.bias.data.cuda())
|
|
|
|
return int8_layer
|
|
|
|
|
|
class FastRMSNorm(nn.Module):
|
|
|
|
def __init__(self, dim: int, eps: float = 1e-5):
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.eps = eps
|
|
self.register_buffer("weight", torch.ones(dim))
|
|
|
|
def forward(self, x):
|
|
return RMSNormFunction.apply(x.float(), self.weight, self.eps).to(x.dtype)
|
|
|
|
@classmethod
|
|
def from_rmsnorm(cls, original_rmsnorm):
|
|
rmsnorm_layer = cls(dim=original_rmsnorm.dim, eps=original_rmsnorm.eps)
|
|
if original_rmsnorm.weight.device != torch.device('meta'):
|
|
rmsnorm_layer.weight.data.copy_(original_rmsnorm.weight.float().data)
|
|
return rmsnorm_layer
|
|
|
|
|
|
class FastLayerNorm(nn.Module):
|
|
|
|
def __init__(self, dim: int, eps: float = 1e-5, elementwise_affine: bool = False, bias: bool = True):
|
|
super().__init__()
|
|
self.dim = dim # type: ignore[arg-type]
|
|
self.eps = eps
|
|
self.elementwise_affine = elementwise_affine
|
|
if self.elementwise_affine:
|
|
self.register_buffer("weight", torch.empty(self.dim))
|
|
if bias:
|
|
self.register_buffer("bias", torch.empty(self.dim))
|
|
else:
|
|
self.bias = None
|
|
else:
|
|
self.register_parameter("weight", None)
|
|
self.register_parameter("bias", None)
|
|
|
|
def forward(self, x):
|
|
return LayerNormFunction.apply(x.float(), self.weight, self.bias, self.eps, self.elementwise_affine).to(x.dtype)
|
|
|
|
@classmethod
|
|
def from_layernorm(cls, original_layernorm):
|
|
layernorm_layer = cls(dim=original_layernorm.normalized_shape[0],
|
|
eps=original_layernorm.eps,
|
|
elementwise_affine=False if original_layernorm.weight is None else True,
|
|
bias=original_layernorm.bias is not None)
|
|
if original_layernorm.weight is not None and original_layernorm.weight.device != torch.device('meta'):
|
|
layernorm_layer.weight.data.copy_(original_layernorm.weight.data)
|
|
if original_layernorm.bias is not None and original_layernorm.bias.device != torch.device('meta'):
|
|
layernorm_layer.bias.data.copy_(original_layernorm.bias.data)
|
|
return layernorm_layer
|