04-Mamba与状态空间模型¶

import torch
import torch.nn as nn
import math

def discretize_zoh(A, B, delta):
    """
    零阶保持离散化
    A: (N, N) 状态转移矩阵
    B: (N, 1) 输入矩阵
    delta: (batch, seq_len, 1) 或标量，时间步长
    """
    # 简化版本：一阶近似
    A_bar = torch.eye(A.shape[0], device=A.device) + delta * A
    B_bar = delta * B
    return A_bar, B_bar

def ssm_conv_kernel(A_bar, B_bar, C, L):
    """
    计算SSM的卷积核
    A_bar: (N, N) 离散化状态矩阵
    B_bar: (N, 1) 离散化输入矩阵  
    C: (1, N) 输出矩阵
    L: 序列长度
    返回: (L,) 卷积核
    """
    N = A_bar.shape[0]
    K = torch.zeros(L)

    A_power = torch.eye(N)  # A^0
    for i in range(L):
        # K[i] = C @ A^i @ B
        K[i] = (C @ A_power @ B_bar).squeeze()
        A_power = A_power @ A_bar  # A^(i+1)

    return K

def hippo_matrix(N):
    """
    构造HiPPO矩阵（LegS变体）
    N: 隐状态维度
    """
    A = torch.zeros(N, N)
    for n in range(N):
        for k in range(N):
            if n > k:
                A[n, k] = -((2*n + 1) ** 0.5) * ((2*k + 1) ** 0.5)
            elif n == k:
                A[n, k] = -(n + 1)
            else:
                A[n, k] = 0
    return A

def hippo_init(N):
    """
    HiPPO初始化的完整版本（S4/Mamba使用）
    """
    # HiPPO-LegS矩阵
    A = torch.zeros(N, N)
    for i in range(N):
        for j in range(N):
            if i > j:
                A[i, j] = (2*i + 1) ** 0.5 * (2*j + 1) ** 0.5
            elif i == j:
                A[i, j] = i + 1

    A = -A  # 加负号

    # B矩阵
    B = torch.zeros(N, 1)
    for i in range(N):
        B[i, 0] = (2*i + 1) ** 0.5

    return A, B

class SelectiveSSM(nn.Module):
    """
    选择性状态空间模块（Mamba核心）
    """
    def __init__(self, d_model, d_state=16, d_conv=4, expand=2):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state  # N, 隐状态维度
        self.d_conv = d_conv    # 局部卷积核大小
        self.d_inner = d_model * expand  # 扩展维度

        # 输入投影
        self.in_proj = nn.Linear(d_model, self.d_inner * 2, bias=False)

        # 局部卷积
        self.conv1d = nn.Conv1d(
            self.d_inner, self.d_inner, 
            kernel_size=d_conv, 
            padding=d_conv - 1,
            groups=self.d_inner
        )

        # SSM参数投影
        self.x_proj = nn.Linear(self.d_inner, d_state * 2 + 1, bias=False)

        # 可学习的A参数（对角线）
        self.A_log = nn.Parameter(torch.randn(self.d_inner, d_state))
        self.D = nn.Parameter(torch.ones(self.d_inner))  # Skip connection

        # 输出投影
        self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)

    def forward(self, x):
        """
        x: (batch, seq_len, d_model)
        """
        batch, seq_len, _ = x.shape

        # 1. 输入投影并分割
        xz = self.in_proj(x)  # (batch, seq_len, d_inner * 2)
        x, z = xz.chunk(2, dim=-1)  # 各 (batch, seq_len, d_inner)

        # 2. 局部卷积 + 激活
        x = x.transpose(1, 2)  # (batch, d_inner, seq_len)
        x = self.conv1d(x)[:, :, :seq_len]
        x = x.transpose(1, 2)  # (batch, seq_len, d_inner)
        x = F.silu(x)

        # 3. 计算时变参数
        x_dbl = self.x_proj(x)  # (batch, seq_len, d_state * 2 + 1)
        delta, B, C = x_dbl.split([1, self.d_state, self.d_state], dim=-1)
        delta = F.softplus(delta)  # 确保正值

        # 4. SSM计算
        A = -torch.exp(self.A_log)  # (d_inner, d_state)
        y = self.selective_scan(x, delta, A, B, C, self.D)

        # 5. 门控 + 输出投影
        y = y * F.silu(z)
        return self.out_proj(y)

    def selective_scan(self, u, delta, A, B, C, D):
        """
        选择性扫描算法（核心计算）
        u: (batch, seq_len, d_inner) 输入
        delta: (batch, seq_len, 1) 时间步长
        A: (d_inner, d_state) 状态矩阵（对角）
        B: (batch, seq_len, d_state) 输入矩阵
        C: (batch, seq_len, d_state) 输出矩阵
        D: (d_inner,) skip connection
        """
        batch, seq_len, d_inner = u.shape
        d_state = A.shape[1]

        # 离散化 A_bar = exp(delta * A), B_bar = delta * B
        deltaA = torch.exp(delta.unsqueeze(-1) * A)  # (batch, seq_len, d_inner, d_state)
        deltaB_u = delta * B * u.unsqueeze(-1)  # (batch, seq_len, d_inner, d_state)

        # 递归计算
        h = torch.zeros(batch, d_inner, d_state, device=u.device)
        ys = []

        for i in range(seq_len):
            h = deltaA[:, i] * h + deltaB_u[:, i]  # 状态更新
            y = torch.sum(h * C[:, i].unsqueeze(1), dim=-1)  # 输出
            ys.append(y)

        y = torch.stack(ys, dim=1)  # (batch, seq_len, d_inner)
        y = y + u * D  # skip connection

        return y

def parallel_scan(A, B, h0):
    """
    并行扫描算法示意（简化版）
    实际实现需要CUDA kernel
    h[t] = A[t] * h[t-1] + B[t]
    """
    L = A.shape[0]

    # 向上扫描（reduce）
    # 实际实现使用平衡二叉树

    # 向下扫描（downsweep）
    # 实际实现使用平衡二叉树

    pass  # 实际实现需要自定义CUDA kernel

class LinearAttention(nn.Module):
    """
    线性注意力实现
    """
    def __init__(self, dim, heads=8, dim_head=64):
        super().__init__()
        self.heads = heads
        self.dim_head = dim_head
        inner_dim = heads * dim_head

        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
        self.to_out = nn.Linear(inner_dim, dim)

    def forward(self, x):
        """
        x: (batch, seq_len, dim)
        """
        batch, seq_len, _ = x.shape

        # QKV投影
        qkv = self.to_qkv(x).reshape(batch, seq_len, 3, self.heads, self.dim_head)
        q, k, v = qkv.unbind(2)  # 各 (batch, seq_len, heads, dim_head)

        # 应用核函数（这里用ELU+1）
        q = F.elu(q) + 1
        k = F.elu(k) + 1

        # 重排维度
        q = q.transpose(1, 2)  # (batch, heads, seq_len, dim_head)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        # 线性注意力: Q(K^T V) 而不是 (Q K^T)V
        # K^T V: (batch, heads, dim_head, dim_head)
        kv = torch.einsum('bhnd,bhne->bhde', k, v)

        # Q (K^T V): (batch, heads, seq_len, dim_head)
        qkv = torch.einsum('bhnd,bhde->bhne', q, kv)

        # 归一化
        k_sum = k.sum(dim=2, keepdim=True)  # (batch, heads, 1, dim_head)
        normalizer = torch.einsum('bhnd,bhmd->bhnm', q, k_sum).squeeze(-1)  # (batch, heads, seq_len)
        normalizer = normalizer.unsqueeze(-1) + 1e-6

        out = qkv / normalizer

        # 输出
        out = out.transpose(1, 2).reshape(batch, seq_len, -1)
        return self.to_out(out)

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from einops import rearrange

class MambaBlock(nn.Module):
    """
    Mamba块的完整实现

    参考: "Mamba: Linear-Time Sequence Modeling with Selective State Spaces"
    作者: Gu & Dao (2023)
    """

    def __init__(
        self,
        d_model: int,
        d_state: int = 16,
        d_conv: int = 4,
        expand: int = 2,
        dt_rank: str | int = "auto",
        dt_min: float = 0.001,
        dt_max: float = 0.1,
        dt_init: str = "random",
        dt_scale: float = 1.0,
        dt_init_floor: float = 1e-4,
        conv_bias: bool = True,
        bias: bool = False,
    ):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.d_conv = d_conv
        self.expand = expand
        self.d_inner = int(self.expand * self.d_model)
        self.dt_rank = math.ceil(self.d_model / 16) if dt_rank == "auto" else dt_rank

        # 输入投影: d_model -> d_inner * 2
        self.in_proj = nn.Linear(d_model, self.d_inner * 2, bias=bias)

        # 局部卷积
        self.conv1d = nn.Conv1d(
            in_channels=self.d_inner,
            out_channels=self.d_inner,
            bias=conv_bias,
            kernel_size=d_conv,
            groups=self.d_inner,
            padding=d_conv - 1,
        )

        # x投影: d_inner -> dt_rank + d_state * 2
        self.x_proj = nn.Linear(self.d_inner, self.dt_rank + self.d_state * 2, bias=False)

        # dt投影: dt_rank -> d_inner
        self.dt_proj = nn.Linear(self.dt_rank, self.d_inner, bias=True)

        # 初始化dt_proj特殊bias
        dt_init_std = self.dt_rank ** -0.5 * dt_scale
        if dt_init == "constant":
            nn.init.constant_(self.dt_proj.weight, dt_init_std)
        elif dt_init == "random":
            nn.init.uniform_(self.dt_proj.weight, -dt_init_std, dt_init_std)
        else:
            raise NotImplementedError

        # 初始化dt bias
        dt = torch.exp(
            torch.rand(self.d_inner) * (math.log(dt_max) - math.log(dt_min)) + math.log(dt_min)
        ).clamp(min=dt_init_floor)
        inv_dt = dt + torch.log(-torch.expm1(-dt))
        with torch.no_grad():
            self.dt_proj.bias.copy_(inv_dt)

        # A参数（对数空间）
        A = torch.arange(1, self.d_state + 1, dtype=torch.float32).repeat(self.d_inner, 1)
        self.A_log = nn.Parameter(torch.log(A))
        self.A_log._no_weight_decay = True

        # D参数（skip connection）
        self.D = nn.Parameter(torch.ones(self.d_inner))
        self.D._no_weight_decay = True

        # 输出投影
        self.out_proj = nn.Linear(self.d_inner, d_model, bias=bias)

    def forward(self, x, cache=None):
        """
        x: (batch, seq_len, d_model)
        返回: (batch, seq_len, d_model)
        """
        batch, seq_len, d_model = x.shape

        # 1. 输入投影
        xz = self.in_proj(x)  # (batch, seq_len, d_inner * 2)
        x, z = xz.chunk(2, dim=-1)  # 各 (batch, seq_len, d_inner)

        # 2. 局部卷积
        x = rearrange(x, 'b l d -> b d l')
        x = self.conv1d(x)[:, :, :seq_len]
        x = rearrange(x, 'b d l -> b l d')
        x = F.silu(x)

        # 3. 计算SSM参数
        x_dbl = self.x_proj(x)  # (batch, seq_len, dt_rank + d_state * 2)
        dt, B, C = torch.split(x_dbl, [self.dt_rank, self.d_state, self.d_state], dim=-1)
        dt = self.dt_proj(dt)  # (batch, seq_len, d_inner)

        # 4. 选择性扫描
        A = -torch.exp(self.A_log.float())  # (d_inner, d_state)
        y = self.selective_scan(x, dt, A, B, C)

        # 5. 门控和输出
        y = y * F.silu(z)
        output = self.out_proj(y)

        return output

    def selective_scan(self, u, delta, A, B, C):
        """
        选择性扫描（简化版，实际应用需要CUDA优化）
        u: (batch, seq_len, d_inner) 输入
        delta: (batch, seq_len, d_inner) 时间步长
        A: (d_inner, d_state) 状态矩阵
        B: (batch, seq_len, d_state) 输入矩阵
        C: (batch, seq_len, d_state) 输出矩阵
        """
        batch, seq_len, d_inner = u.shape
        d_state = A.shape[1]

        # 确保delta为正
        delta = F.softplus(delta)

        # 离散化
        deltaA = torch.exp(torch.einsum('bld,dn->bldn', delta, A))
        deltaB_u = torch.einsum('bld,bln,bld->bldn', delta, B, u)

        # 递归计算（这里用Python循环示意，实际需要并行扫描）
        h = torch.zeros(batch, d_inner, d_state, device=u.device, dtype=u.dtype)
        ys = []

        for i in range(seq_len):
            h = deltaA[:, i] * h + deltaB_u[:, i]
            y = torch.einsum('bdn,bn->bd', h, C[:, i])
            ys.append(y)

        y = torch.stack(ys, dim=1)  # (batch, seq_len, d_inner)
        y = y + u * self.D  # skip connection

        return y


class MambaLayer(nn.Module):
    """
    Mamba层 = Mamba块 + 归一化
    """
    def __init__(self, d_model, **kwargs):
        super().__init__()
        self.mamba = MambaBlock(d_model, **kwargs)
        self.norm = nn.RMSNorm(d_model)

    def forward(self, x):
        return self.norm(x + self.mamba(x))


class MambaModel(nn.Module):
    """
    完整的Mamba语言模型
    """
    def __init__(
        self,
        vocab_size: int,
        d_model: int = 768,
        n_layer: int = 24,
        d_state: int = 16,
        d_conv: int = 4,
        expand: int = 2,
    ):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.layers = nn.ModuleList([
            MambaLayer(d_model, d_state=d_state, d_conv=d_conv, expand=expand)
            for _ in range(n_layer)
        ])
        self.norm_f = nn.RMSNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, input_ids):
        """
        input_ids: (batch, seq_len)
        """
        x = self.embedding(input_ids)

        for layer in self.layers:
            x = layer(x)

        x = self.norm_f(x)
        logits = self.lm_head(x)

        return logits

class SSMLayer(nn.Module):
    """
    基础SSM层（非选择性，用于理解原理）
    """
    def __init__(self, d_model, d_state=64, dropout=0.0):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state

        # HiPPO初始化
        A, B = hippo_init(d_state)
        self.A = nn.Parameter(A)  # (d_state, d_state)
        self.B = nn.Parameter(B)  # (d_state, 1)
        self.C = nn.Parameter(torch.randn(1, d_state) * 0.01)  # (1, d_state)
        self.D = nn.Parameter(torch.zeros(1))  # skip connection

        # 时间步长（可学习）
        self.delta = nn.Parameter(torch.ones(1) * 0.01)

        # 投影层
        self.in_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)
        self.norm = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        """
        x: (batch, seq_len, d_model)
        """
        residual = x
        x = self.norm(x)
        x = self.in_proj(x)

        # 离散化
        delta = F.softplus(self.delta)
        A_bar = torch.matrix_exp(delta * self.A)
        B_bar = torch.linalg.solve(self.A, (A_bar - torch.eye(self.d_state, device=x.device))) @ self.B
        B_bar = delta * self.B  # 简化近似

        # SSM递归
        batch, seq_len, _ = x.shape
        h = torch.zeros(batch, self.d_state, device=x.device)
        outputs = []

        for t in range(seq_len):
            # h_t = A_bar @ h_{t-1} + B_bar @ x_t
            h = A_bar @ h + B_bar.squeeze(-1) * x[:, t]
            # y_t = C @ h_t + D @ x_t
            y = self.C @ h + self.D * x[:, t]
            outputs.append(y)

        y = torch.stack(outputs, dim=1)  # (batch, seq_len, d_model)
        y = self.out_proj(y)
        y = self.dropout(y)

        return residual + y

class TransformerBlock(nn.Module):
    """
    标准Transformer块（用于对比）
    """
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super().__init__()
        self.attention = nn.MultiheadAttention(d_model, n_heads, dropout=dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        # 自注意力
        residual = x
        x = self.norm1(x)
        x, _ = self.attention(x, x, x)
        x = residual + x

        # 前馈网络
        residual = x
        x = self.norm2(x)
        x = self.ffn(x)
        x = residual + x

        return x


def compare_complexity():
    """
    复杂度对比实验
    """
    import time

    d_model = 768
    batch_size = 1
    seq_lengths = [128, 512, 1024, 2048, 4096]

    transformer = TransformerBlock(d_model, n_heads=12, d_ff=3072)
    mamba = MambaBlock(d_model)

    print(f"{'Seq Len':<10} {'Transformer (ms)':<20} {'Mamba (ms)':<15} {'Speedup':<10}")
    print("-" * 60)

    for seq_len in seq_lengths:
        x = torch.randn(batch_size, seq_len, d_model)

        # Transformer
        start = time.time()
        with torch.no_grad():
            _ = transformer(x)
        trans_time = (time.time() - start) * 1000

        # Mamba
        start = time.time()
        with torch.no_grad():
            _ = mamba(x)
        mamba_time = (time.time() - start) * 1000

        speedup = trans_time / mamba_time
        print(f"{seq_len:<10} {trans_time:<20.2f} {mamba_time:<15.2f} {speedup:<10.2f}x")


if __name__ == "__main__":
    # 测试Mamba块
    batch, seq_len, d_model = 2, 128, 256
    x = torch.randn(batch, seq_len, d_model)

    mamba = MambaBlock(d_model)
    out = mamba(x)
    print(f"Input shape: {x.shape}")
    print(f"Output shape: {out.shape}")

    # 复杂度对比
    compare_complexity()

# 内存对比
seq_len = 100000  # 100K tokens
d_model = 4096
n_layers = 32

# Transformer KV Cache (假设8头)
kv_cache_size = 2 * n_layers * seq_len * d_model * 4  # bytes
print(f"Transformer KV Cache: {kv_cache_size / 1e9:.2f} GB")

# Mamba 状态
d_state = 16
mamba_state_size = n_layers * d_model * d_state * 4
print(f"Mamba State: {mamba_state_size / 1e6:.2f} MB")

class JambaBlock(nn.Module):
    """
    Jamba混合块：Transformer + Mamba
    """
    def __init__(self, d_model, n_heads, d_ff, d_state=16, layer_type="mamba"):
        super().__init__()
        self.layer_type = layer_type

        if layer_type == "transformer":
            self.layer = TransformerBlock(d_model, n_heads, d_ff)
        else:
            self.layer = MambaLayer(d_model, d_state=d_state)

    def forward(self, x):
        return self.layer(x)


class JambaModel(nn.Module):
    """
    Jamba风格混合模型
    """
    def __init__(self, vocab_size, d_model, n_layers, n_heads, d_ff):
        super().__init__()

        self.embedding = nn.Embedding(vocab_size, d_model)

        # 混合层配置：每4层中1个Transformer，3个Mamba
        layers = []
        for i in range(n_layers):
            if i % 4 == 0:
                layers.append(JambaBlock(d_model, n_heads, d_ff, layer_type="transformer"))
            else:
                layers.append(JambaBlock(d_model, n_heads, d_ff, layer_type="mamba"))

        self.layers = nn.ModuleList(layers)
        self.norm = nn.RMSNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size)

    def forward(self, input_ids):
        x = self.embedding(input_ids)
        for layer in self.layers:
            x = layer(x)
        x = self.norm(x)
        return self.lm_head(x)

class Mamba2Block(nn.Module):
    """
    Mamba-2 块（简化示意）
    主要改进：使用块并行算法
    """
    def __init__(self, d_model, d_state=128, block_size=64):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.block_size = block_size

        # 参数
        self.A = nn.Parameter(torch.randn(d_model, d_state))
        self.B = nn.Parameter(torch.randn(d_model, d_state))
        self.C = nn.Parameter(torch.randn(d_model, d_state))

        self.in_proj = nn.Linear(d_model, d_model * 2)
        self.out_proj = nn.Linear(d_model, d_model)

    def forward(self, x):
        """
        使用块并行算法
        """
        batch, seq_len, _ = x.shape

        # 投影
        xz = self.in_proj(x)
        x, z = xz.chunk(2, dim=-1)

        # 块并行SSM（示意）
        n_blocks = (seq_len + self.block_size - 1) // self.block_size
        outputs = []

        for block_idx in range(n_blocks):
            start = block_idx * self.block_size
            end = min(start + self.block_size, seq_len)
            x_block = x[:, start:end, :]

            # 块内并行计算
            # 实际实现使用优化的CUDA kernel
            y_block = self._block_ssm(x_block)
            outputs.append(y_block)

        y = torch.cat(outputs, dim=1)
        y = y * F.silu(z)

        return self.out_proj(y)

    def _block_ssm(self, x_block):
        """块内SSM计算"""
        # 简化实现
        return x_block

# 1. 学习率设置
# Mamba的A_log参数需要较小学习率
optimizer = torch.optim.AdamW([
    {'params': model.embedding.parameters(), 'lr': 3e-4},
    {'params': model.layers.parameters(), 'lr': 3e-4},
    {'params': [p for n, p in model.named_parameters() if 'A_log' in n], 'lr': 1e-4}
])

# 2. 梯度裁剪
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# 3. 混合精度训练
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()

with autocast():
    output = model(input_ids)
    loss = criterion(output, targets)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

# 练习1：实现简化版SSM
class SimpleSSM(nn.Module):
    """
    TODO: 实现一个简化版SSM
    要求：
    1. 使用HiPPO初始化
    2. 实现递归形式的前向传播
    3. 实现卷积形式的前向传播
    4. 验证两种形式输出一致
    """
    pass

# 练习2：实现线性Attention
class LinearAttention(nn.Module):
    """
    TODO: 实现线性Attention
    要求：
    1. 使用ELU+1核函数
    2. 实现O(L*d²)复杂度的前向传播
    3. 与标准Attention对比输出差异
    """
    pass

# 练习3：复杂度实验
def complexity_experiment():
    """
    TODO: 设计实验对比以下模型在不同序列长度下的性能
    1. 标准Transformer
    2. 线性Attention
    3. Mamba

    测量指标：
    - 前向时间
    - 反向时间
    - 显存占用
    """
    pass