代码连连看——mamba

传送门

arxiv

github

模型解读

【汇报】 Mamba模型及其公式推导

Mamba系列日积月累（一）：状态空间模型SSM的离散化过程推导

连连看

首先从用法开始入手:

import torch
from mamba_ssm import Mamba

batch, length, dim = 2, 64, 16
x = torch.randn(batch, length, dim).to("cuda")
model = Mamba(
    # This module uses roughly 3 * expand * d_model^2 parameters
    d_model=dim, # Model dimension d_model
    d_state=16,  # SSM state expansion factor
    d_conv=4,    # Local convolution width
    expand=2,    # Block expansion factor
).to("cuda")
y = model(x)
assert y.shape == x.shape

IDE会告诉我们: 入口在mamba_ssm/modules/mamba_simple.py中. 接下来就来会会它

出发!

先来看看Mamba.forward(), 很容易就能看到Mamba.step()是关键点

def forward(self, hidden_states, inference_params=None):
    """
    hidden_states: (B, L, D)
    Returns: same shape as hidden_states
    """
    batch, seqlen, dim = hidden_states.shape

    conv_state, ssm_state = None, None
    if inference_params is not None:
        conv_state, ssm_state = self._get_states_from_cache(inference_params, batch)
        if inference_params.seqlen_offset > 0:
            # The states are updated inplace
            out, _, _ = self.step(hidden_states, conv_state, ssm_state)
            return out
    ...

Mamba.step()完整实现了Figure3, 为了和源码中的注释区分, 混合了C风格的注释(#//, 下同)

点我查看完整代码

def step(self, hidden_states, conv_state, ssm_state):
    dtype = hidden_states.dtype
    assert hidden_states.shape[1] == 1, "Only support decoding with 1 token at a time for now"
    xz = self.in_proj(hidden_states.squeeze(1))  # (B 2D)
    x, z = xz.chunk(2, dim=-1)  # (B D) #// 两个部分, x走Figure3中左侧通路, z走右侧通路

    # Conv step
    if causal_conv1d_update is None:
        conv_state.copy_(torch.roll(conv_state, shifts=-1, dims=-1))  # Update state (B D W)
        conv_state[:, :, -1] = x
        x = torch.sum(conv_state * rearrange(self.conv1d.weight, "d 1 w -> d w"), dim=-1)  # (B D)
        if self.conv1d.bias is not None:
            x = x + self.conv1d.bias
        x = self.act(x).to(dtype=dtype)
    else:
        x = causal_conv1d_update(
            x,
            conv_state,
            rearrange(self.conv1d.weight, "d 1 w -> d w"),
            self.conv1d.bias,
            self.activation,
        )

    #// Algorithm 2
    x_db = self.x_proj(x)  # (B dt_rank+2*d_state) #// x_proj 把 s_delta(的linear部分), s_b, s_c 合并了, 所以下面要从中分割出dt,B,C
    dt, B, C = torch.split(x_db, [self.dt_rank, self.d_state, self.d_state], dim=-1)
    # Don't add dt_bias here
    dt = F.linear(dt, self.dt_proj.weight)  # (B d_inner) #// s_delta的broadcasr部分
    A = -torch.exp(self.A_log.float())  # (d_inner, d_state) #// A_log在init()中定义为torch.log(A)<TODO>为什么添加负号</TODO>
    #// 至此, A, B, C和softplus之前的delta已经准备完毕
    # SSM step
    if selective_state_update is None:
        # Discretize A and B
        dt = F.softplus(dt + self.dt_proj.bias.to(dtype=dt.dtype)) #// 现在, 我们有真正的delta了
        dA = torch.exp(torch.einsum("bd,dn->bdn", dt, A)) #// A_bar = exp(delta*A)
        dB = torch.einsum("bd,bn->bdn", dt, B) #// dB = delta*B
        ssm_state.copy_(ssm_state * dA + rearrange(x, "b d -> b d 1") * dB) #// h_t = A_bar * h_{t-1} + B_bar * x_{t}
        y = torch.einsum("bdn,bn->bd", ssm_state.to(dtype), C) #// y = C * h_{t}
        y = y + self.D.to(dtype) * x #// <TODO>论文中似乎没有提及SSM内部有残差连接</TODO>
        y = y * self.act(z)  # (B D)
    else:
        y = selective_state_update(
            ssm_state, x, dt, A, B, C, self.D, z=z, dt_bias=self.dt_proj.bias, dt_softplus=True
        )

    out = self.out_proj(y)
    return out.unsqueeze(1), conv_state, ssm_state

默认的selective_state_update实现, 位于mamba_ssm/ops/triton/selective_state_update.py中, 因为使用了一些优化手段, 阅读起来比较复杂, 但是作者给了方便阅读的对应的_ref代码, 这里就不赘述了.

再回到Mamba.forward(), 可以知道位于mamba_ssm/ops/selective_scan_interface.py中的SelectiveScanFn和MambaInnerFn是实际运行时采用的代码, 作者同样提供了对应的_ref版本, 基本逻辑一致, 直接跳过了

总结

1. Mamba.step()中已经有了S6的计算, 为何又在selective_scan_fn中重写了操作?
2. 有些代码细节和论文不一致, 以后读论文和代码要注意
3. A矩阵为什么要取log再存储, 还原的时候又要加负号?