Test the outputs of a simple geglu mlp with liger and kernels

compare-simple-liger-kernels.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from liger_kernel.transformers import LigerTiledGEGLUMLP
from kernels import get_kernel

# create the inputs and the conigurations
device = "cuda"
batch = 128
seq = 16
dim = 128
hidden = 512
x = torch.randn(batch, seq, dim, device=device, dtype=torch.bfloat16)

class Config:
    hidden_size = dim
    intermediate_size = hidden
    hidden_act = "gelu_pytorch_tanh"

# Build the simple geglu mlp using torch.nn.Module
class SimpleGeGLUMLP(nn.Module):
    def __init__(self, dim, hidden):
        super().__init__()
        self.gate_proj = nn.Linear(dim, hidden, bias=False)
        self.up_proj = nn.Linear(dim, hidden, bias=False)
        self.down_proj = nn.Linear(hidden, dim, bias=False)

    def forward(self, x):
        return self.down_proj(F.gelu(self.gate_proj(x), approximate="tanh") * self.up_proj(x))


simple_geglu_mlp = SimpleGeGLUMLP(dim, hidden).to(device, dtype=torch.bfloat16).eval()
liger_geglu_mlp = LigerTiledGEGLUMLP(config=Config()).to(device=device, dtype=torch.bfloat16).eval()

# make the weights equivalent (liger)
liger_geglu_mlp.gate_proj.weight.data.copy_(simple_geglu_mlp.gate_proj.weight.data)
liger_geglu_mlp.up_proj.weight.data.copy_(simple_geglu_mlp.up_proj.weight.data)
liger_geglu_mlp.down_proj.weight.data.copy_(simple_geglu_mlp.down_proj.weight.data)

kernels_layers = get_kernel("kernels-community/liger-kernels", version=1).layers
kernels_geglu_mlp = kernels_layers.LigerGEGLUMLP
kernels_geglu_mlp = kernels_geglu_mlp(Config()).to(device=device, dtype=torch.bfloat16).eval()

# make the weights equivalent (kernels)
kernels_geglu_mlp.gate_proj.weight.data.copy_(simple_geglu_mlp.gate_proj.weight.data)
kernels_geglu_mlp.up_proj.weight.data.copy_(simple_geglu_mlp.up_proj.weight.data)
kernels_geglu_mlp.down_proj.weight.data.copy_(simple_geglu_mlp.down_proj.weight.data)

simple_outputs = simple_geglu_mlp(x)
liger_outputs = liger_geglu_mlp(x)
kernels_outputs = kernels_geglu_mlp(x)

torch.testing.assert_close(
    simple_outputs,
    liger_outputs,
    atol=1e-4,
    rtol=1e-4,
)

torch.testing.assert_close(
    simple_outputs,
    kernels_outputs,
    atol=1e-4,
    rtol=1e-4,
)