初始化项目，由ModelHub XC社区提供模型
Model: pathcosmos/frankenstallm Source: Original Platform
This commit is contained in:
ModelHub XC
280
source/model/mamba_block.py
Normal file
280
source/model/mamba_block.py
Normal file
@@ -0,0 +1,280 @@
|
||||
"""
|
||||
Mamba-2 block based on the Structured State Space Duality (SSD) formulation.
|
||||
|
||||
Reference: "Transformers are SSMs: Generalized Models and Efficient Algorithms
|
||||
Through Structured State Space Duality" (Dao & Gu, 2024).
|
||||
|
||||
This implements a pure-PyTorch sequential scan for correctness and generality.
|
||||
A chunked SSD kernel can be swapped in later for speed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .layers import RMSNorm
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Selective Scan (sequential, numerically stable in float32)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def selective_scan(
|
||||
x: torch.Tensor,
|
||||
dt: torch.Tensor,
|
||||
A_log: torch.Tensor,
|
||||
B: torch.Tensor,
|
||||
C: torch.Tensor,
|
||||
D: torch.Tensor,
|
||||
n_groups: int,
|
||||
) -> torch.Tensor:
|
||||
"""Run the SSM recurrence sequentially over the time axis.
|
||||
|
||||
Args:
|
||||
x: (B, L, n_heads, head_dim) — input after conv + activation.
|
||||
dt: (B, L, n_heads) — discretisation time-steps (after softplus).
|
||||
A_log: (n_heads,) — log(-A), learnable diagonal decay.
|
||||
B: (B, L, n_groups, d_state) — input-to-state projection per step.
|
||||
C: (B, L, n_groups, d_state) — state-to-output projection per step.
|
||||
D: (n_heads,) — skip/residual connection per head.
|
||||
n_groups: int — number of B/C groups (heads per group share B/C).
|
||||
|
||||
Returns:
|
||||
y: (B, L, n_heads, head_dim) — SSM output.
|
||||
"""
|
||||
batch, seq_len, n_heads, head_dim = x.shape
|
||||
d_state = B.shape[-1]
|
||||
heads_per_group = n_heads // n_groups
|
||||
|
||||
# Compute decay: dA = exp(-exp(A_log) * dt) — shape (B, L, n_heads)
|
||||
neg_A = A_log.exp() # (n_heads,)
|
||||
dA = torch.exp(-neg_A.unsqueeze(0).unsqueeze(0) * dt) # (B, L, n_heads)
|
||||
|
||||
# Scale input by dt: dBx will be accumulated into state
|
||||
# dt: (B, L, n_heads) -> (B, L, n_heads, 1)
|
||||
dt_x = dt.unsqueeze(-1) * x # (B, L, n_heads, head_dim)
|
||||
|
||||
# Allocate output
|
||||
y = torch.zeros_like(x)
|
||||
|
||||
# State: (B, n_heads, head_dim, d_state) — accumulated in float32
|
||||
h = torch.zeros(
|
||||
batch, n_heads, head_dim, d_state,
|
||||
dtype=torch.float32, device=x.device,
|
||||
)
|
||||
|
||||
# Expand B/C from groups to heads: (B, L, n_groups, d_state) -> indexing
|
||||
# For efficiency we index into the group dimension during the loop.
|
||||
# group_idx[head] -> which group this head belongs to
|
||||
group_idx = torch.arange(n_heads, device=x.device) // heads_per_group # (n_heads,)
|
||||
|
||||
for t in range(seq_len):
|
||||
# --- Decay state ---
|
||||
# dA_t: (B, n_heads) -> (B, n_heads, 1, 1)
|
||||
dA_t = dA[:, t, :].float().unsqueeze(-1).unsqueeze(-1)
|
||||
h = h * dA_t # (B, n_heads, head_dim, d_state)
|
||||
|
||||
# --- Input contribution ---
|
||||
# B_t: (B, n_groups, d_state) -> (B, n_heads, d_state) via group expansion
|
||||
B_t = B[:, t, :, :][:, group_idx, :] # (B, n_heads, d_state)
|
||||
# dt_x_t: (B, n_heads, head_dim)
|
||||
dt_x_t = dt_x[:, t, :, :].float() # (B, n_heads, head_dim)
|
||||
# Outer product: (B, n_heads, head_dim, 1) * (B, n_heads, 1, d_state)
|
||||
h = h + dt_x_t.unsqueeze(-1) * B_t.float().unsqueeze(-2)
|
||||
|
||||
# --- Output ---
|
||||
# C_t: (B, n_groups, d_state) -> (B, n_heads, d_state)
|
||||
C_t = C[:, t, :, :][:, group_idx, :] # (B, n_heads, d_state)
|
||||
# y_t = sum_over_d_state( h * C_t ) -> (B, n_heads, head_dim)
|
||||
y_t = torch.einsum("bnhd,bnd->bnh", h, C_t.float())
|
||||
y[:, t, :, :] = y_t.to(x.dtype)
|
||||
|
||||
# Skip connection: D * x
|
||||
y = y + D.view(1, 1, n_heads, 1) * x
|
||||
|
||||
return y
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Mamba-2 Block
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class Mamba2Block(nn.Module):
|
||||
"""Mamba-2 block with pre-norm residual connection.
|
||||
|
||||
Implements:
|
||||
1. RMSNorm (pre-norm)
|
||||
2. Input projection -> (z, x, B, C, dt)
|
||||
3. Causal depth-wise Conv1d on x
|
||||
4. SiLU activation on x
|
||||
5. Selective scan (SSM recurrence)
|
||||
6. Gated output: y * SiLU(z)
|
||||
7. Output projection + residual
|
||||
|
||||
Args:
|
||||
d_model: Model hidden dimension.
|
||||
d_state: SSM state dimension N (default 128).
|
||||
head_dim: Per-head dimension for SSD (default 64).
|
||||
expand: Expansion factor for inner dimension (default 2).
|
||||
conv_kernel: Causal 1D convolution kernel size (default 4).
|
||||
n_groups: Number of groups for B/C projections (default 1).
|
||||
chunk_size: Chunk size for SSD algorithm — reserved for future use (default 256).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
d_model: int,
|
||||
d_state: int = 128,
|
||||
head_dim: int = 64,
|
||||
expand: int = 2,
|
||||
conv_kernel: int = 4,
|
||||
n_groups: int = 1,
|
||||
chunk_size: int = 256,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.d_model = d_model
|
||||
self.d_state = d_state
|
||||
self.head_dim = head_dim
|
||||
self.expand = expand
|
||||
self.n_groups = n_groups
|
||||
self.chunk_size = chunk_size
|
||||
|
||||
# Derived dimensions
|
||||
self.d_inner = expand * d_model
|
||||
self.n_heads = self.d_inner // head_dim
|
||||
assert self.d_inner % head_dim == 0, (
|
||||
f"d_inner ({self.d_inner}) must be divisible by head_dim ({head_dim})"
|
||||
)
|
||||
assert self.n_heads % n_groups == 0, (
|
||||
f"n_heads ({self.n_heads}) must be divisible by n_groups ({n_groups})"
|
||||
)
|
||||
|
||||
# Pre-norm
|
||||
self.norm = RMSNorm(d_model)
|
||||
|
||||
# Input projection: d_model -> z + x + B + C + dt
|
||||
self.d_proj = (
|
||||
self.d_inner # z (gate)
|
||||
+ self.d_inner # x (input to conv + SSM)
|
||||
+ n_groups * d_state # B
|
||||
+ n_groups * d_state # C
|
||||
+ self.n_heads # dt (one per head)
|
||||
)
|
||||
self.in_proj = nn.Linear(d_model, self.d_proj, bias=False)
|
||||
|
||||
# Causal depth-wise conv1d over x
|
||||
self.conv1d = nn.Conv1d(
|
||||
in_channels=self.d_inner,
|
||||
out_channels=self.d_inner,
|
||||
kernel_size=conv_kernel,
|
||||
groups=self.d_inner,
|
||||
padding=conv_kernel - 1, # causal: trim trailing values
|
||||
)
|
||||
|
||||
# SSM parameters
|
||||
# A_log: log(-A) where A is the diagonal decay — init from log(uniform(1, 16))
|
||||
A_init = torch.log(torch.rand(self.n_heads) * 15.0 + 1.0) # log(U(1,16))
|
||||
self.A_log = nn.Parameter(A_init)
|
||||
|
||||
# D: skip connection per head — init to ones
|
||||
self.D = nn.Parameter(torch.ones(self.n_heads))
|
||||
|
||||
# dt_bias: added before softplus — init from log(uniform(0.001, 0.1))
|
||||
dt_bias_init = torch.log(torch.rand(self.n_heads) * 0.099 + 0.001)
|
||||
self.dt_bias = nn.Parameter(dt_bias_init)
|
||||
|
||||
# Output projection
|
||||
self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _split_projection(
|
||||
self, proj: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Split the fused input projection into (z, x, B, C, dt).
|
||||
|
||||
Args:
|
||||
proj: (B, L, d_proj)
|
||||
|
||||
Returns:
|
||||
z: (B, L, d_inner)
|
||||
x: (B, L, d_inner)
|
||||
B: (B, L, n_groups, d_state)
|
||||
C: (B, L, n_groups, d_state)
|
||||
dt: (B, L, n_heads)
|
||||
"""
|
||||
batch, seq_len, _ = proj.shape
|
||||
i = 0
|
||||
|
||||
z = proj[:, :, i : i + self.d_inner]
|
||||
i += self.d_inner
|
||||
|
||||
x = proj[:, :, i : i + self.d_inner]
|
||||
i += self.d_inner
|
||||
|
||||
bc_dim = self.n_groups * self.d_state
|
||||
B = proj[:, :, i : i + bc_dim].reshape(batch, seq_len, self.n_groups, self.d_state)
|
||||
i += bc_dim
|
||||
|
||||
C = proj[:, :, i : i + bc_dim].reshape(batch, seq_len, self.n_groups, self.d_state)
|
||||
i += bc_dim
|
||||
|
||||
dt = proj[:, :, i : i + self.n_heads]
|
||||
return z, x, B, C, dt
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Forward
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
x: (B, L, d_model) — input hidden states.
|
||||
|
||||
Returns:
|
||||
(B, L, d_model) — output with residual connection applied.
|
||||
"""
|
||||
residual = x
|
||||
x = self.norm(x)
|
||||
|
||||
# --- Input projection ---
|
||||
proj = self.in_proj(x) # (B, L, d_proj)
|
||||
z, x_ssm, B, C, dt_raw = self._split_projection(proj)
|
||||
|
||||
# --- Causal conv1d on x ---
|
||||
# Conv1d expects (B, C, L)
|
||||
x_conv = x_ssm.transpose(1, 2) # (B, d_inner, L)
|
||||
x_conv = self.conv1d(x_conv)
|
||||
# Trim to causal: remove the (kernel-1) trailing padding
|
||||
x_conv = x_conv[:, :, :x_ssm.shape[1]] # (B, d_inner, L)
|
||||
x_conv = x_conv.transpose(1, 2) # (B, L, d_inner)
|
||||
x_conv = F.silu(x_conv)
|
||||
|
||||
# --- Discretise dt ---
|
||||
dt = F.softplus(dt_raw + self.dt_bias) # (B, L, n_heads)
|
||||
|
||||
# --- Reshape x for multi-head scan ---
|
||||
batch, seq_len, _ = x_conv.shape
|
||||
x_heads = x_conv.reshape(batch, seq_len, self.n_heads, self.head_dim)
|
||||
|
||||
# --- Selective scan (SSM recurrence) ---
|
||||
y = selective_scan(
|
||||
x_heads, dt, self.A_log, B, C, self.D,
|
||||
n_groups=self.n_groups,
|
||||
) # (B, L, n_heads, head_dim)
|
||||
|
||||
# --- Flatten heads back ---
|
||||
y = y.reshape(batch, seq_len, self.d_inner) # (B, L, d_inner)
|
||||
|
||||
# --- Gated output ---
|
||||
y = y * F.silu(z)
|
||||
|
||||
# --- Output projection + residual ---
|
||||
return residual + self.out_proj(y)
|
||||
Reference in New Issue
Block a user