代码连连看——RWKV

传送门

arxiv

github

模型解读

Huggingface rwkv代码解读

RWKV 模型解析

连连看

先用v4版本做连连看, v5和v6对应arxiv:2404.05892, 有时间再做连连看

老规矩, 找入口(RWKV-v4/run.py):

print(f'Loading {MODEL_NAME}...')
from src.model_run import RWKV_RNN
model = RWKV_RNN(MODEL_NAME, os.environ['RWKV_RUN_DEVICE'], model_type, n_layer, n_embd, ctx_len)
tokenizer = TOKENIZER(WORD_NAME, UNKNOWN_CHAR=UNKNOWN_CHAR)

后面调用的是 model.run(x), 很明显, 我们要进到 RWKV-v4/src/model_run.py 里面去

出发!

注释和mamba一样, 都是用 #// 表示

首先直奔 RWKV_RNN.run(), 看到下面这句话就知道还要继续找下去

w = torch.load(MODEL_NAME + '.pth', map_location=torch.device(RUN_DEVICE))

看看同文件夹的有什么文件就知道该去哪里找了: RWKV-v4/src/model.py

先找最顶层的模块:

点我查看完整代码

class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.step = 0
        self.config = config

        self.emb = nn.Embedding(config.vocab_size, config.n_embd)

        self.blocks = nn.Sequential(*[Block(config, i)
                                    for i in range(config.n_layer)])

        self.ln_out = nn.LayerNorm(config.n_embd)
        self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

        if RWKV_HEAD_QK_DIM > 0:
            self.head_q = nn.Linear(config.n_embd, RWKV_HEAD_QK_DIM, bias=False)
            self.head_q.scale_init = 0
            self.head_k = nn.Linear(config.n_embd, RWKV_HEAD_QK_DIM, bias=False)
            self.head_k.scale_init = 0.1
            self.register_buffer("copy_mask", torch.tril(
                torch.ones(config.ctx_len, config.ctx_len)))

        self.ctx_len = config.ctx_len

        try:
            if os.environ['RWKV_LOAD_MODEL'] == str(False):
                RWKV_Init(self, config) 
        except:
            pass

        logger.info("number of parameters: %e", sum(p.numel()
                    for p in self.parameters()))

    def forward(self, idx, targets=None):
        idx = idx.to(self.emb.weight.device)

        self.step += 1
        B, T = idx.size()
        assert T <= self.ctx_len, "Cannot forward, because len(input) > model ctx_len."

        x = self.emb(idx) #// Input Embeding
        x = self.blocks(x) #// blocks
        x = self.ln_out(x) #// LayerNorm in RWKV-LM Head

        if RWKV_HEAD_QK_DIM > 0:
            q = self.head_q(x)[:, :T, :]
            k = self.head_k(x)[:, :T, :]
            c = (q @ k.transpose(-2, -1)) * (1.0 / RWKV_HEAD_QK_DIM)
            c = c.masked_fill(self.copy_mask[:T, :T] == 0, 0)
            
            if '32' in os.environ['RWKV_FLOAT_MODE']:
                c = c @ F.one_hot(idx, num_classes=self.config.vocab_size)
            elif os.environ['RWKV_FLOAT_MODE'] == 'fp16':
                c = c @ F.one_hot(idx, num_classes=self.config.vocab_size).half()
            elif os.environ['RWKV_FLOAT_MODE'] == 'bf16':
                c = c @ F.one_hot(idx, num_classes=self.config.vocab_size).bfloat16()

            x = self.head(x) + c #// other part of RWKV-LM Head
        else:
            x = self.head(x) #// other part of RWKV-LM Head

        loss = None
        if targets is not None:
            loss = F.cross_entropy(x.view(-1, x.size(-1)), targets.to(x.device).view(-1))

        return L2Wrap.apply(loss, x)

接着来看看Block是怎么实现的

点我查看完整代码

class Block(nn.Module):
    def __init__(self, config, layer_id):
        super().__init__()
        self.config = config
        self.layer_id = layer_id

        self.ln1 = nn.LayerNorm(config.n_embd)
        self.ln2 = nn.LayerNorm(config.n_embd)

        if self.layer_id == 0:
            self.ln0 = nn.LayerNorm(config.n_embd)

        if self.layer_id == 0 and self.config.model_type == 'RWKV-ffnPre':
            self.ffnPre = RWKV_ChannelMix(config, 0)
        else:
            self.att = RWKV_TimeMix(config, layer_id)

        self.ffn = RWKV_ChannelMix(config, layer_id)

    def forward(self, x):
        if self.layer_id == 0:
            x = self.ln0(x)  #// LayerNorm between Input Embedding and Blocks
        if self.layer_id == 0 and self.config.model_type == 'RWKV-ffnPre':
            x = x + self.ffnPre(self.ln1(x))  # better in some cases
        else:
            x = x + self.att(self.ln1(x)) #// LayerNorm and TimeMixing and ResidualConnection
        x = x + self.ffn(self.ln2(x)) #// LayerNorm and ChannelMixing and ResidualConnection
        return x

先看看 RWKV_TimeMix

点我查看完整代码

class RWKV_TimeMix(torch.jit.ScriptModule):
    def __init__(self, config, layer_id):
        super().__init__()
        self.layer_id = layer_id
        self.ctx_len = config.ctx_len
        self.n_embd = config.n_embd

        attn_sz = config.n_embd

        with torch.no_grad(): # fancy init
            ratio_0_to_1 = (layer_id / (config.n_layer - 1)) # 0 to 1
            ratio_1_to_almost0 = (1.0 - (layer_id / config.n_layer)) # 1 to ~0
            
            # fancy time_decay
            decay_speed = torch.ones(attn_sz)
            for h in range(attn_sz):
                decay_speed[h] = -5 + 8 * (h / (attn_sz-1)) ** (0.7 + 1.3 * ratio_0_to_1)
            self.time_decay = nn.Parameter(decay_speed)
            # print(layer_id, self.time_decay.flatten()[:3].cpu().numpy(), '...', self.time_decay.flatten()[-3:].cpu().numpy())

            # fancy time_first
            zigzag = (torch.tensor([(i+1)%3 - 1 for i in range(attn_sz)]) * 0.5)
            self.time_first = nn.Parameter(torch.ones(attn_sz) * math.log(0.3) + zigzag)
            
            # fancy time_mix
            x = torch.ones(1, 1, config.n_embd)
            for i in range(config.n_embd):
                x[0, 0, i] = i / config.n_embd
            self.time_mix_k = nn.Parameter(torch.pow(x, ratio_1_to_almost0))
            self.time_mix_v = nn.Parameter(torch.pow(x, ratio_1_to_almost0) + 0.3 * ratio_0_to_1)
            self.time_mix_r = nn.Parameter(torch.pow(x, 0.5 * ratio_1_to_almost0))
            
        self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))

        self.key = nn.Linear(config.n_embd, attn_sz, bias=False)
        self.value = nn.Linear(config.n_embd, attn_sz, bias=False)
        self.receptance = nn.Linear(config.n_embd, attn_sz, bias=False)

        self.output = nn.Linear(attn_sz, config.n_embd, bias=False)

        self.key.scale_init = 0
        self.receptance.scale_init = 0
        self.output.scale_init = 0

    @torch.jit.script_method
    def jit_func(self, x):

        # Mix x with the previous timestep to produce xk, xv, xr
        xx = self.time_shift(x) #// x_{t-1}
        xk = x * self.time_mix_k + xx * (1 - self.time_mix_k) #// k_{t} before multiple w_{k}
        xv = x * self.time_mix_v + xx * (1 - self.time_mix_v) #// v_{t} before multiple w_{v}
        xr = x * self.time_mix_r + xx * (1 - self.time_mix_r) #// r_{t} before multiple w_{r}

        # Use xk, xv, xr to produce k, v, r
        k = self.key(xk) #// k_{t}
        v = self.value(xv) #// v_{t}
        r = self.receptance(xr) #// r_{t}
        sr = torch.sigmoid(r) #// sigma

        return sr, k, v

    def forward(self, x):
        B, T, C = x.size() # x = (Batch,Time,Channel)

        sr, k, v = self.jit_func(x)

        rwkv = sr * RUN_CUDA(B, T, C, self.time_decay, self.time_first, k, v) #// sigma * WKV
        rwkv = self.output(rwkv) #// Output Gating
        return rwkv

里面的 RUN_CUDA() 实际上是 class WKV 套娃, 而 class WKV 也是套娃, 真正的实现是用CUDA写的, 在 RWKV-v4/cuda/wkv_cuda.cu 和 RWKV-v4/cuda/wkv_op.cpp 里

接下来看看 ChannelMix

点我查看完整代码

class RWKV_ChannelMix(torch.jit.ScriptModule):
    def __init__(self, config, layer_id):
        super().__init__()
        self.layer_id = layer_id

        self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))

        with torch.no_grad(): # fancy init of time_mix
            ratio_1_to_almost0 = (1.0 - (layer_id / config.n_layer)) # 1 to ~0

            x = torch.ones(1, 1, config.n_embd)
            for i in range(config.n_embd):
                x[0, 0, i] = i / config.n_embd

            self.time_mix_k = nn.Parameter(torch.pow(x, ratio_1_to_almost0))
            self.time_mix_r = nn.Parameter(torch.pow(x, ratio_1_to_almost0))

        hidden_sz = 4 * config.n_embd
        self.key = nn.Linear(config.n_embd, hidden_sz, bias=False)
        self.receptance = nn.Linear(config.n_embd, config.n_embd, bias=False)
        self.value = nn.Linear(hidden_sz, config.n_embd, bias=False)

        self.value.scale_init = 0
        self.receptance.scale_init = 0

    @torch.jit.script_method
    def forward(self, x):
        xx = self.time_shift(x)
        xk = x * self.time_mix_k + xx * (1 - self.time_mix_k) #// k'_{t} before multiple w'_{k}
        xr = x * self.time_mix_r + xx * (1 - self.time_mix_r) #// r'_{t} before multiple w'_{r}

        k = self.key(xk) #// k'_{t}
        k = torch.square(torch.relu(k)) #// (max(k'_{t}, 0))^{2}
        kv = self.value(k) #// v' = w'_{v} * (max(k'_{t}, 0))^{2}

        rkv = torch.sigmoid(self.receptance(xr)) * kv #// sigma' * v'
        return rkv

是不是忘了什么? 对, 还有一步很重要的 States 的传递没有体现, 因为我们看的不是 RNN 形式, 让我们回到 RWKV_RNN, States 的传递是通过 save() 和 load() 实现的, 其他代码和上面的基本一致, 就不赘述了

总结