mirrors/labmlai--annotated_deep_learning_paper_implementations
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:19:01 +08:00
commit 3b90d1192f
2172 changed files with 594509 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
labml_nn/rwkv/__init__.py
+302
View File
@@ -0,0 +1,302 @@
"""
---
title: Receptance Weighted Key Value (RWKV)
summary: >
This implements the RWKV model
using PyTorch with explanations.
---
# Receptance Weighted Key Value (RWKV)
This is a tutorial/implementation of RWKV
from paper [RWKV: Reinventing RNNs for the Transformer Era](https://arxiv.org/pdf/2305.13048.pdf)
in [PyTorch](https://pytorch.org/).
Full definition of a RWKV Language Model, all of it in this single file.
References:
1) [the official RWKV PyTorch implementation released by Bo Peng](https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v4neo/src/model.py)
2) [huggingface/transformers PyTorch implementation](https://github.com/huggingface/transformers/blob/main/src/transformers/models/rwkv/modeling_rwkv.py)
"""
import torch
import torch.nn as nn
from torch.nn import functional as F
PREV_X_TIME = 0
NUM_STATE = 1
DEN_STATE = 2
MAX_STATE = 3
PREV_X_CHANNEL = 4
class LayerNorm(nn.Module):
"""
### Layer normalization with bias
"""
def __init__(self, ndim, bias):
super().__init__()
self.weight = nn.Parameter(torch.ones(ndim))
self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None
def forward(self, input):
return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)
class L2Wrap(torch.autograd.Function):
"""
### L2 loss wrapper
[ref](https://github.com/BlinkDL/RWKV-LM/blob/cca1b5e8e597cf40675882bb10b46287c844e35c/RWKV-v4/src/model.py#L21)
"""
@staticmethod
def forward(ctx, loss, y):
ctx.save_for_backward(y)
return loss
@staticmethod
def backward(ctx, grad_output):
y = ctx.saved_tensors[0]
# to encourage the logits to be close to 0
factor = 1e-4 / (y.shape[0] * y.shape[1])
maxx, ids = torch.max(y, -1, keepdim=True)
gy = torch.zeros_like(y)
gy.scatter_(-1, ids, maxx * factor)
return grad_output, gy
class ChannelMixing(nn.Module):
"""
### Channel Mixing
"""
def __init__(self, config, layer_id):
super().__init__()
self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
# token shifting
self.layer_id = layer_id
n_embd = config.n_embd
intermediate_size = (
config.intermediate_size if config.intermediate_size is not None else 4 * n_embd
)
# Learnable Matrix
self.key_proj = nn.Linear(n_embd, intermediate_size, bias=False)
self.value_proj = nn.Linear(intermediate_size, n_embd, bias=False)
self.receptance_proj = nn.Linear(n_embd, n_embd, bias=False)
# Learnable Vector
self.time_mix_key = nn.Parameter(torch.empty(1, 1, n_embd))
self.time_mix_receptance = nn.Parameter(torch.empty(1, 1, n_embd))
def forward(self, x, state=None):
"""
# x = (Batch,Time,Channel)
"""
if state is not None:
prev_x = state[self.layer_id, :, [PREV_X_CHANNEL], :]
state[self.layer_id, :, [PREV_X_CHANNEL], :] = x
else:
prev_x = self.time_shift(x)
# $r_t=W_r \cdot (\mu_r x_t + (1-\mu_r)x_{t-1})$
receptance = x * self.time_mix_receptance + prev_x * (1 - self.time_mix_receptance)
receptance = self.receptance_proj(receptance)
# $k_t=W_k \cdot (\mu_k x_t + (1-\mu_k)x_{t-1})$
key = x * self.time_mix_key + prev_x * (1 - self.time_mix_key)
key = self.key_proj(key)
# $V_t=W_v \cdot max(k_t,0)^2$
value = self.value_proj(torch.square(torch.relu(key)))
# $o_t=\sigma(r_t) \odot v_t$
out = F.sigmoid(receptance) * value
return out, state
class TimeMixing(nn.Module):
"""
### Time Mixing
"""
def __init__(self, config, layer_id):
super().__init__()
self.config = config
self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
self.layer_id = layer_id
n_embd = config.n_embd
attn_sz = n_embd
# learnable matrix
self.key_proj = nn.Linear(n_embd, attn_sz, bias=False)
self.value_proj = nn.Linear(n_embd, attn_sz, bias=False)
self.receptance_proj = nn.Linear(n_embd, attn_sz, bias=False)
self.output_proj = nn.Linear(attn_sz, n_embd, bias=False)
# learnable vector
self.time_decay = nn.Parameter(torch.empty(attn_sz))
self.time_first = nn.Parameter(torch.empty(attn_sz))
self.time_mix_key = nn.Parameter(torch.empty(1, 1, n_embd))
self.time_mix_value = nn.Parameter(torch.empty(1, 1, n_embd))
self.time_mix_receptance = nn.Parameter(torch.empty(1, 1, n_embd))
def forward(self, x, state=None):
"""
x = (Batch,Time,Channel)
"""
if state is not None:
prev_x = state[self.layer_id, :, [PREV_X_TIME], :]
state[self.layer_id, :, [PREV_X_TIME], :] = x
else:
prev_x = self.time_shift(x)
# $r_t=W_r \cdot (\mu_r x_t + (1-\mu_r)x_{t-1})$
receptance = x * self.time_mix_receptance + prev_x * (1 - self.time_mix_receptance)
receptance = self.receptance_proj(receptance)
# $k_t=W_k \cdot (\mu_k x_t + (1-\mu_k)x_{t-1})$
key = x * self.time_mix_key + prev_x * (1 - self.time_mix_key)
key = self.key_proj(key)
# $v_t=W_v \cdot (\mu_v x_t + (1-\mu_v)x_{t-1})$
value = x * self.time_mix_value + prev_x * (1 - self.time_mix_value)
value = self.value_proj(value)
# WKV calculation
_, seq_length, _ = key.size()
output = torch.zeros_like(key)
if state is None:
num_state = torch.zeros_like(key[:, 0], dtype=torch.float32)
den_state = torch.zeros_like(key[:, 0], dtype=torch.float32)
max_state = torch.zeros_like(key[:, 0], dtype=torch.float32) - 1e38
else:
num_state = state[self.layer_id, :, NUM_STATE, :]
den_state = state[self.layer_id, :, DEN_STATE, :]
max_state = state[self.layer_id, :, MAX_STATE, :]
time_decay = -torch.exp(self.time_decay)
for current_index in range(seq_length):
current_key = key[:, current_index].float()
current_value = value[:, current_index]
# $wkv_t=\frac{\sum^{t-1}_{i=1}d^{-(t-1-i)w+k_i}v_i+e^{u+k_t}v_t}{\sum^{t-1}_{i=1}e^{-(t-1-i)w+k_i}+e^{u+k_t}}$
max_for_output = torch.maximum(max_state, current_key + self.time_first)
e1 = torch.exp(max_state - max_for_output)
e2 = torch.exp(current_key + self.time_first - max_for_output)
numerator = e1 * num_state + e2 * current_value
denominator = e1 * den_state + e2
output[:, current_index] = (numerator / denominator).to(output.dtype)
# Update state for next iteration
max_for_state = torch.maximum(max_state + time_decay, current_key)
e1 = torch.exp(max_state + time_decay - max_for_state)
e2 = torch.exp(current_key - max_for_state)
num_state = e1 * num_state + e2 * current_value
den_state = e1 * den_state + e2
max_state = max_for_state
# update states
state[self.layer_id, :, NUM_STATE, :] = num_state
state[self.layer_id, :, DEN_STATE, :] = den_state
state[self.layer_id, :, MAX_STATE, :] = max_state
wkv, state = self.wkv_function(key, value, use_customized_cuda_kernel=self.config.use_customized_cuda_kernel,
state=state)
# $o_t=W_o \cdot (\sigma(r_t) \odot wkv_t)$
rwkv = F.sigmoid(receptance) * wkv
rwkv = self.output_proj(rwkv)
return rwkv, state
class Block(nn.Module):
"""
## RWKV block element
"""
def __init__(self, config, layer_id):
super().__init__()
self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
self.attn = TimeMixing(config, layer_id)
self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)
self.ffn = ChannelMixing(config, layer_id)
def forward(self, x, state=None):
# state: [batch_size, 5 , n_embd]
# time mixing
residual = x
x, state = self.attn(self.ln_1(x), state=state)
x = x + residual
# channel mixing
residual = x
x, state = self.ffn(self.ln_2(x), state=state)
x = x + residual
return x, state
class RWKV(nn.Module):
"""
## RWKV
"""
def __init__(self, config, lr_init=0.0008):
super().__init__()
assert config.vocab_size is not None
assert config.block_size is not None
self.config = config
self.lr_init = lr_init ## used to initialize embedding parameters
self.n_layer = config.n_layer
self.n_embd = config.n_embd
# Initiate model layers
self.rwkv = nn.ModuleDict(dict(
wte=nn.Embedding(config.vocab_size, config.n_embd),
ln_p=LayerNorm(config.n_embd, bias=config.bias),
h=nn.ModuleList([Block(config, layer_id) for layer_id in range(config.n_layer)]),
ln_f=LayerNorm(config.n_embd, bias=config.bias),
))
# Output linear layer
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
def forward(self, idx, targets=None, state=None, return_state=False):
b, t = idx.size()
assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
# Embedding Layer
x = self.rwkv.wte(idx)
# Layer Norm
x = self.rwkv.ln_p(x)
# RWKV Blocks
for block_idx, block in enumerate(self.rwkv.h):
x, state = block(x, state)
x = self.rwkv.ln_f(x)
# Logit Layer and loss Function (for training)
if targets is not None:
# if we are given some desired targets also calculate the loss
logits = self.lm_head(x)
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
if self.training:
loss = L2Wrap.apply(loss, logits)
else:
# inference-time mini-optimization: only forward the lm_head on the very last position
logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
loss = None
# Return Logits and loss
if return_state:
return logits, loss, state
else:
return logits, loss
labml_nn/rwkv/configs.py
+24
View File
@@ -0,0 +1,24 @@
from labml.configs import BaseConfigs
class RWKVConfigs(BaseConfigs):
"""
## Transformer Configurations
This defines configurations for a transformer.
The configurations are calculate using option functions.
These are lazy loaded and therefore only the necessary modules
are calculated.
"""
# Number of attention heads
n_heads: int = 8
# Transformer embedding size
d_model: int = 512
# Number of layers
n_layers: int = 6
# Dropout probability
dropout: float = 0.1
# Number of tokens in the source vocabulary (for token embeddings)
n_src_vocab: int
# Number of tokens in the target vocabulary (to generate logits for prediction)
n_tgt_vocab: int
labml_nn/rwkv/experiment.py
+169
View File
@@ -0,0 +1,169 @@
import inspect
import math
import torch
import torch.nn as nn
from labml_nn.rwkv.configs import RWKVConfigs
from labml_nn.rwkv import RWKV
from labml_nn.rwkv import TimeMixing
from labml import experiment
from labml.configs import option
from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
class Configs(NLPAutoRegressionConfigs):
"""
## Configurations
This inherits from
[`NLPAutoRegressionConfigs`](../../experiments/nlp_autoregression.html#NLPAutoRegressionConfigs)
"""
# RWKV model
model: RWKV
rwkv: RWKVConfigs
# number of warmup iterations
warmup_iters: int = 2000
# total number of training iterations
max_iters: int = 600000
# weight decay
weight_decay: float = 1e-1
# Custom optimizer
beta1: float = 0.9
beta2: float = 0.95
optimizer = 'rwkv_optimizer'
@option(Configs.rwkv, 'RWKV')
def _rwkv_configs(c: Configs):
"""
### RWKV configurations
"""
# We use our
# [configurable RWKV implementation](../configs.html#RWKVConfigs)
conf = RWKVConfigs()
# Set the vocabulary sizes for embeddings and generating logits
conf.n_src_vocab = c.n_tokens
conf.n_tgt_vocab = c.n_tokens
return conf
def _init_weights(module, rwkv: RWKVConfigs):
# initialize Vector Parameters in TimeMixing
if isinstance(module, TimeMixing):
layer_id = module.layer_id
n_layer = module.n_layer
n_embd = module.n_embd
attn_sz = n_embd
with torch.no_grad():
ratio_0_to_1 = layer_id / (n_layer - 1) # 0 to 1
ratio_1_to_almost0 = 1.0 - (layer_id / n_layer) # 1 to ~0
ddd = torch.ones(1, 1, n_embd)
for i in range(n_embd):
ddd[0, 0, i] = i / n_embd
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)
module.time_decay = nn.Parameter(decay_speed)
zigzag = torch.tensor([(i + 1) % 3 - 1 for i in range(attn_sz)]) * 0.5
module.time_first = nn.Parameter(torch.ones(attn_sz) * math.log(0.3) + zigzag)
module.time_mix_key = nn.Parameter(torch.pow(ddd, ratio_1_to_almost0))
module.time_mix_value = nn.Parameter(torch.pow(ddd, ratio_1_to_almost0) + 0.3 * ratio_0_to_1)
module.time_mix_receptance = nn.Parameter(torch.pow(ddd, 0.5 * ratio_1_to_almost0))
@option(Configs.model)
def _model(c: Configs):
"""
Create RWKV model and initialize weights
"""
m = RWKV(c.rwkv).to(c.device)
# Apply custom weight initialization
m.apply(_init_weights, c.rwkv)
return m
@option(NLPAutoRegressionConfigs.optimizer)
def _configure_optimizers(c: NLPAutoRegressionConfigs):
# start with all of the candidate parameters
param_dict = {pn: p for pn, p in c.model.named_parameters()}
# filter out those that do not require grad
param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
# create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
# i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.
decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
optim_groups = [
{'params': decay_params, 'weight_decay': c.weight_decay},
{'params': nodecay_params, 'weight_decay': 0.0}
]
num_decay_params = sum(p.numel() for p in decay_params)
num_nodecay_params = sum(p.numel() for p in nodecay_params)
print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")
print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")
# Create AdamW optimizer and use the fused version if it is available
fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
use_fused = fused_available and c.device_type == 'cuda'
extra_args = dict(fused=True) if use_fused else dict()
optimizer = torch.optim.AdamW(optim_groups, lr=c.learning_rate, betas=c.betas, **extra_args)
print(f"using fused AdamW: {use_fused}")
return optimizer
def main():
# Create experiment
experiment.create(name="RWKV")
# Create configs
conf = Configs()
print(conf.model)
# Override configurations
experiment.configs(conf, {
# Use character level tokenizer
'tokenizer': 'character',
# Prompt separator is blank
'prompt_separator': '',
# Starting prompt for sampling
'prompt': 'It is ',
# Use Tiny Shakespeare dataset
'text': 'tiny_shakespeare',
# Use a context size of $128$
'seq_len': 128,
# Train for $32$ epochs
'epochs': 32,
# Batch size $128$
'batch_size': 128,
# Switch between training and validation for $10$ times
# per epoch
'inner_iterations': 10,
'rwkv.block_size': 1024,
# model
'rwkv.n_layer': 12,
'rwkv.n_heads': 12,
'rwkv.n_embd': 768
})
print(conf.model)
# Set models for saving and loading
experiment.add_pytorch_models({'model': conf.model})
# Start the experiment
with experiment.start():
# Run training
conf.run()
#
if __name__ == '__main__':
main()