basic.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint
#%%
def create_xielu_params(alpha_p_init=0.8, alpha_n_init=0.8, beta=0.5, eps=-1e-6, device=None, dtype=None):
    dev_kwargs = {"device": device, "dtype": dtype}
    alpha_p = torch.nn.Parameter(
        torch.log(torch.exp(torch.tensor(alpha_p_init, **dev_kwargs)) - 1).unsqueeze(0))
    alpha_n = torch.nn.Parameter(
        torch.log(torch.exp(torch.tensor(alpha_n_init - beta, **dev_kwargs)) - 1).unsqueeze(0))
    return alpha_p, alpha_n, beta, eps

class XIELUPy(torch.nn.Module):
    def __init__(self, alpha_p_init=0.8, alpha_n_init=0.8, beta=0.5, eps=-1e-6, device=None, dtype=None):
        super().__init__()
        self.alpha_p, self.alpha_n, self.beta, self.eps = create_xielu_params(
            alpha_p_init, alpha_n_init, beta, eps, device, dtype)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        alpha_p = F.softplus(self.alpha_p)
        alpha_n = self.beta + F.softplus(self.alpha_n)
        return torch.where(x > 0,
                           alpha_p * x * x + self.beta * x,
                           alpha_n * torch.expm1(torch.clamp_max(x, self.eps)) - alpha_n * x + self.beta * x)

#%%
# xielu = XIELUPy(device="cuda", dtype=torch.bfloat16)
# #test the xielu function
# x = torch.randn(1,1,1,16).to(torch.bfloat16).to("cuda")
# x = xielu(x)
# print(x.shape)
#%%


def norm(x):
    return F.rms_norm(x, (x.size(-1),),eps=1e-6)

# ---------- FeedForward ----------
class FeedForward(nn.Module):
    def __init__(self, d_model, mult=4):
        super().__init__()
        d_hidden = int(mult * d_model)
        self.w1 = nn.Linear(d_model, d_hidden, bias=False)
        self.wo = nn.Linear(d_hidden, d_model, bias=True)
        self.nonlin_act = XIELUPy()

    # @torch.compile(mode="max-autotune")  # Disabled for testing
    def forward(self, x):
        x = self.w1(x)
        x = self.nonlin_act(x)
        x = self.wo(x)
        return x


class TrainableRoPE(nn.Module):
    def __init__(self, dim: int, max_len: int = 64, base: float = 64.0, dtype=torch.float32,trainable=True):
        super().__init__()
        assert dim % 2 == 0
        self.trainable = trainable
        h = dim // 2
        pos = torch.arange(max_len, dtype=dtype).unsqueeze(1)          # (L,1)
        inv = base ** (-torch.arange(0, h, dtype=dtype) / dim)         # (h,)
        th  = pos * inv                                                # (L,h)
        self.cos = nn.Parameter(torch.cos(th), requires_grad=trainable)  # trainable (L,h)
        self.sin = nn.Parameter(torch.sin(th), requires_grad=trainable)  # trainable (L,h)

    # @torch.compile(mode="max-autotune")  # Disabled for testing
    def forward(self, x: torch.Tensor, offset: int = 0):
        # x: (..., L, D) - stable at the end (last position = 0, counting backwards)
        # offset: how many positions from the end these tokens are
        h = x.size(-1) // 2
        seq_len = x.size(-2)
        cos = self.cos[offset:offset+seq_len, :h].flip(0).reshape((1,)*(x.ndim-2) + (seq_len, h)).to(x.dtype)
        sin = self.sin[offset:offset+seq_len, :h].flip(0).reshape((1,)*(x.ndim-2) + (seq_len, h)).to(x.dtype)
        x1, x2 = x[..., :h], x[..., h:]
        y1 = x1 * cos - x2 * sin
        y2 = x1 * sin + x2 * cos
        return torch.cat([y1, y2], dim=-1)