code/flash-linear-attention/utils/convert_from_rwkv6.py

# scripts for converting pretrained hf model weights to fla style
# calling the code to make conversions for RWKV/rwkv-6-world-7b would achieve the following results:
# |    Tasks     |Version|Filter|n-shot|    Metric     | Value |   |Stderr|
# |--------------|------:|------|-----:|---------------|------:|---|------|
# |arc_challenge |      1|none  |     0|acc            | 0.4130|±  |0.0144|
# |              |       |none  |     0|acc_norm       | 0.4403|±  |0.0145|
# |arc_easy      |      1|none  |     0|acc            | 0.7382|±  |0.0090|
# |              |       |none  |     0|acc_norm       | 0.7079|±  |0.0093|
# |boolq         |      2|none  |     0|acc            | 0.6823|±  |0.0081|
# |copa          |      1|none  |     0|acc            | 0.8700|±  |0.0338|
# |hellaswag     |      1|none  |     0|acc            | 0.5508|±  |0.0050|
# |              |       |none  |     0|acc_norm       | 0.7171|±  |0.0045|
# |lambada_openai|      1|none  |     0|perplexity     | 3.2989|±  |0.0634|
# |              |       |none  |     0|acc            | 0.7493|±  |0.0060|
# |openbookqa    |      1|none  |     0|acc            | 0.3200|±  |0.0209|
# |              |       |none  |     0|acc_norm       | 0.4440|±  |0.0222|
# |piqa          |      1|none  |     0|acc            | 0.7753|±  |0.0097|
# |              |       |none  |     0|acc_norm       | 0.7894|±  |0.0095|
# |sciq          |      1|none  |     0|acc            | 0.9370|±  |0.0077|
# |              |       |none  |     0|acc_norm       | 0.8860|±  |0.0101|
# |winogrande    |      1|none  |     0|acc            | 0.6867|±  |0.0130|

import argparse

import torch
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

import fla  # noqa
from fla.utils import device


def sizeof_fmt(num, suffix='B'):
    for unit in ('', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi'):
        if abs(num) < 1024.0:
            return f'{num:.2f}{unit}{suffix}'
        num /= 1024.0
    return f'{num:.2f}Yi{suffix}'


def convert(
    rwkv6: str,
    config: str,
    output: str,
):
    torch.manual_seed(1)
    AutoTokenizer.from_pretrained(rwkv6, trust_remote_code=True).save_pretrained(output)
    rwkv6 = AutoModelForCausalLM.from_pretrained(rwkv6, trust_remote_code=True).to(device)
    print(f"Loading rwkv6 ...\n{rwkv6}")

    config = AutoConfig.from_pretrained(config)
    model = AutoModelForCausalLM.from_config(config).to(device)
    num_parameters = model.num_parameters()
    print(f"Initializing the model from the config:\n{config}\n{model}")
    print(f"Number of parameters in total: {num_parameters} ({sizeof_fmt(num_parameters)})")

    print("Copying the weights from rwkv6 to the model ...")
    print("rwkv6.rwkv.embeddings                        -> model.model.embeddings")
    model.model.embeddings.weight.data.copy_(rwkv6.rwkv.embeddings.weight)
    torch.testing.assert_close(model.model.embeddings.weight, rwkv6.rwkv.embeddings.weight)
    for i in range(config.num_hidden_layers):
        if hasattr(model.model.layers[i], 'pre_norm'):
            if model.model.layers[i].pre_norm.weight is not None:
                print(f"rwkv6.rwkv.blocks{i}.pre_ln.weight -> model.model.layers{i}.pre_norm.weight")
                model.model.layers[i].pre_norm.weight.data.copy_(rwkv6.rwkv.blocks[i].pre_ln.weight)
                torch.testing.assert_close(model.model.layers[i].pre_norm.weight, rwkv6.rwkv.blocks[i].pre_ln.weight)
            if model.model.layers[i].pre_norm.bias is not None:
                print(f"rwkv6.rwkv.blocks{i}.pre_ln.bias -> model.model.layers{i}.pre_norm.bias")
                model.model.layers[i].pre_norm.bias.data.copy_(rwkv6.rwkv.blocks[i].pre_ln.bias)
                torch.testing.assert_close(model.model.layers[i].pre_norm.bias, rwkv6.rwkv.blocks[i].pre_ln.bias)
            model.model.layers[i].pre_norm.eps = rwkv6.rwkv.blocks[i].pre_ln.eps
        if model.model.layers[i].attn_norm.weight is not None:
            print(f"rwkv6.rwkv.blocks{i}.ln1.weight -> model.model.layers{i}.attn_norm.weight")
            model.model.layers[i].attn_norm.weight.data.copy_(rwkv6.rwkv.blocks[i].ln1.weight)
            torch.testing.assert_close(model.model.layers[i].attn_norm.weight, rwkv6.rwkv.blocks[i].ln1.weight)
        if model.model.layers[i].attn_norm.bias is not None:
            print(f"rwkv6.rwkv.blocks{i}.ln1.bias -> model.model.layers{i}.attn_norm.bias")
            model.model.layers[i].attn_norm.bias.data.copy_(rwkv6.rwkv.blocks[i].ln1.bias)
            torch.testing.assert_close(model.model.layers[i].attn_norm.bias, rwkv6.rwkv.blocks[i].ln1.bias)
        model.model.layers[i].attn_norm.eps = rwkv6.rwkv.blocks[i].ln1.eps

        print(f"rwkv6.rwkv.blocks{i}.attention.time_maa_x -> model.model.layers.{i}.attn.x_proj0.mu")
        model.model.layers[i].attn.x_proj[0].mu.data.copy_(rwkv6.rwkv.blocks[i].attention.time_maa_x.view(-1))
        torch.testing.assert_close(model.model.layers[i].attn.x_proj[0].mu,
                                   rwkv6.rwkv.blocks[i].attention.time_maa_x.view(-1))
        print(f"rwkv6.rwkv.blocks{i}.attention.time_maa_w1.weight  -> model.model.layers{i}.attn.x_proj0.linear.weight")
        ww, wk, wv, wr, wg = rwkv6.rwkv.blocks[i].attention.time_maa_w1.view(config.hidden_size, 5, -1).unbind(-2)
        w = torch.cat((wr, ww, wk, wv, wg), -1).t()
        model.model.layers[i].attn.x_proj[0].linear.weight.data.copy_(w)
        torch.testing.assert_close(model.model.layers[i].attn.x_proj[0].linear.weight, w)

        print(f"rwkv6.rwkv.blocks{i}.attention.time_maa_w2.weight  -> model.model.layers{i}.attn.x_proj2.weight")
        ww, wk, wv, wr, wg = rwkv6.rwkv.blocks[i].attention.time_maa_w2.unbind(0)
        w = torch.cat((wr, ww, wk, wv, wg), 0).t()
        model.model.layers[i].attn.x_proj[2].weight.data.copy_(w)
        torch.testing.assert_close(model.model.layers[i].attn.x_proj[2].weight, w)

        print(f"rwkv6.rwkv.blocks{i}.attention.time_maa_wkvrg  -> model.model.layers{i}.attn.x_bias")
        bias = torch.stack((rwkv6.rwkv.blocks[i].attention.time_maa_r.view(-1),
                            rwkv6.rwkv.blocks[i].attention.time_maa_w.view(-1),
                            rwkv6.rwkv.blocks[i].attention.time_maa_k.view(-1),
                            rwkv6.rwkv.blocks[i].attention.time_maa_v.view(-1),
                            rwkv6.rwkv.blocks[i].attention.time_maa_g.view(-1)))
        model.model.layers[i].attn.x_bias.data.copy_(bias)
        torch.testing.assert_close(model.model.layers[i].attn.x_bias, bias)

        print(f"rwkv6.rwkv.blocks{i}.attention.receptance.weight  -> model.model.layers{i}.attn.r_proj.linear.weight")
        model.model.layers[i].attn.r_proj.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.receptance.weight)
        torch.testing.assert_close(model.model.layers[i].attn.r_proj.linear.weight,
                                   rwkv6.rwkv.blocks[i].attention.receptance.weight)
        print(f"rwkv6.rwkv.blocks{i}.attention.time_decay_w1  -> model.model.layers{i}.attn.w_proj.linear.lora0.weight")
        model.model.layers[i].attn.w_proj.linear.lora[0].weight.data.copy_(rwkv6.rwkv.blocks[i].attention.time_decay_w1.t())
        torch.testing.assert_close(model.model.layers[i].attn.w_proj.linear.lora[0].weight,
                                   rwkv6.rwkv.blocks[i].attention.time_decay_w1.t())
        print(f"rwkv6.rwkv.blocks{i}.attention.time_decay_w2  -> model.model.layers{i}.attn.w_proj.linear.lora2.weight")
        model.model.layers[i].attn.w_proj.linear.lora[2].weight.data.copy_(rwkv6.rwkv.blocks[i].attention.time_decay_w2.t())
        torch.testing.assert_close(model.model.layers[i].attn.w_proj.linear.lora[2].weight,
                                   rwkv6.rwkv.blocks[i].attention.time_decay_w2.t())
        print(f"rwkv6.rwkv.blocks{i}.attention.time_decay  -> model.model.layers{i}.attn.w_proj.linear.lora2.bias")
        model.model.layers[i].attn.w_proj.linear.lora[2].bias.data.copy_(rwkv6.rwkv.blocks[i].attention.time_decay.view(-1))
        torch.testing.assert_close(model.model.layers[i].attn.w_proj.linear.lora[2].bias,
                                   rwkv6.rwkv.blocks[i].attention.time_decay.view(-1))

        print(f"rwkv6.rwkv.blocks{i}.attention.key.weight -> model.model.layers.{i}.attn.k_proj.linear.weight")
        model.model.layers[i].attn.k_proj.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.key.weight)
        torch.testing.assert_close(model.model.layers[i].attn.k_proj.linear.weight,
                                   rwkv6.rwkv.blocks[i].attention.key.weight)
        print(f"rwkv6.rwkv.blocks{i}.attention.value.weight -> model.model.layers.{i}.attn.v_proj.linear.weight")
        model.model.layers[i].attn.v_proj.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.value.weight)
        torch.testing.assert_close(model.model.layers[i].attn.v_proj.linear.weight,
                                   rwkv6.rwkv.blocks[i].attention.value.weight)
        print(f"rwkv6.rwkv.blocks{i}.attention.gate.weight -> model.model.layers.{i}.attn.g_proj.linear.weight")
        model.model.layers[i].attn.g_proj.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.gate.weight)
        torch.testing.assert_close(model.model.layers[i].attn.g_proj.linear.weight,
                                   rwkv6.rwkv.blocks[i].attention.gate.weight)
        print(f"rwkv6.rwkv.blocks{i}.attention.time_faaaa -> model.model.layers.{i}.attn.bonus")
        bonus = rwkv6.rwkv.blocks[i].attention.time_faaaa.view(config.num_heads, -1)
        model.model.layers[i].attn.bonus.data.copy_(bonus)
        torch.testing.assert_close(model.model.layers[i].attn.bonus, bonus)

        if model.model.layers[i].attn.g_norm.weight is not None:
            print(f"rwkv6.rwkv.blocks{i}.attention.ln_x.weight -> model.model.layers[i].attn.g_norm.weight")
            model.model.layers[i].attn.g_norm.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.ln_x.weight)
            torch.testing.assert_close(model.model.layers[i].attn.g_norm.weight, rwkv6.rwkv.blocks[i].attention.ln_x.weight)
        if model.model.layers[i].attn.g_norm.bias is not None:
            print(f"rwkv6.rwkv.blocks{i}.attention.ln_x.bias -> model.model.layers[i].attn.g_norm.bias")
            model.model.layers[i].attn.g_norm.bias.data.copy_(rwkv6.rwkv.blocks[i].attention.ln_x.bias)
            torch.testing.assert_close(model.model.layers[i].attn.g_norm.bias, rwkv6.rwkv.blocks[i].attention.ln_x.bias)
        model.model.layers[i].attn.g_norm.eps = rwkv6.rwkv.blocks[i].attention.ln_x.eps

        print(f"rwkv6.rwkv.blocks{i}.attention.output.weight -> model.model.layers.{i}.attn.o_proj.weight")
        model.model.layers[i].attn.o_proj.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.output.weight)
        torch.testing.assert_close(model.model.layers[i].attn.o_proj.weight, rwkv6.rwkv.blocks[i].attention.output.weight)

        if model.model.layers[i].ffn_norm.weight is not None:
            print(f"rwkv6.rwkv.blocks{i}.ln2.weight -> model.model.layers{i}.ffn_norm.weight")
            model.model.layers[i].ffn_norm.weight.data.copy_(rwkv6.rwkv.blocks[i].ln2.weight)
            torch.testing.assert_close(model.model.layers[i].ffn_norm.weight, rwkv6.rwkv.blocks[i].ln2.weight)
        if model.model.layers[i].ffn_norm.bias is not None:
            print(f"rwkv6.rwkv.blocks{i}.ln2.bias -> model.model.layers{i}.ffn_norm.bias")
            model.model.layers[i].ffn_norm.bias.data.copy_(rwkv6.rwkv.blocks[i].ln2.bias)
            torch.testing.assert_close(model.model.layers[i].ffn_norm.bias, rwkv6.rwkv.blocks[i].ln2.bias)
        model.model.layers[i].ffn_norm.eps = rwkv6.rwkv.blocks[i].ln2.eps

        print(f"rwkv6.rwkv.blocks{i}.feed_forward.key.weight -> model.model.layers.{i}.ffn.key.linear.weight")
        model.model.layers[i].ffn.key.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].feed_forward.key.weight)
        torch.testing.assert_close(model.model.layers[i].ffn.key.linear.weight,
                                   rwkv6.rwkv.blocks[i].feed_forward.key.weight)
        print(f"rwkv6.rwkv.blocks{i}.feed_forward.time_maa_k -> model.model.layers.{i}.ffn.key.mu")
        model.model.layers[i].ffn.key.mu.data.copy_(rwkv6.rwkv.blocks[i].feed_forward.time_maa_k.view(-1))
        torch.testing.assert_close(model.model.layers[i].ffn.key.mu,
                                   rwkv6.rwkv.blocks[i].feed_forward.time_maa_k.view(-1))

        print(f"rwkv6.rwkv.blocks{i}.feed_forward.value.weight -> model.model.layers.{i}.ffn.value.weight")
        model.model.layers[i].ffn.value.weight.data.copy_(rwkv6.rwkv.blocks[i].feed_forward.value.weight)
        torch.testing.assert_close(model.model.layers[i].ffn.value.weight,
                                   rwkv6.rwkv.blocks[i].feed_forward.value.weight)

        print(f"rwkv6.rwkv.blocks{i}.feed_forward.receptance.weight -> model.model.layers.{i}.ffn.receptance.linear.weight")
        model.model.layers[i].ffn.receptance.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].feed_forward.receptance.weight)
        torch.testing.assert_close(model.model.layers[i].ffn.receptance.linear.weight,
                                   rwkv6.rwkv.blocks[i].feed_forward.receptance.weight)
        print(f"rwkv6.rwkv.blocks{i}.feed_forward.time_maa_r -> model.model.layers.{i}.ffn.receptance.mu")
        model.model.layers[i].ffn.receptance.mu.data.copy_(rwkv6.rwkv.blocks[i].feed_forward.time_maa_r.view(-1))
        torch.testing.assert_close(model.model.layers[i].ffn.receptance.mu,
                                   rwkv6.rwkv.blocks[i].feed_forward.time_maa_r.view(-1))

    if model.model.norm.weight is not None:
        print("rwkv6.rwkv.ln_out.weight -> model.model.norm.weight")
        model.model.norm.weight.data.copy_(rwkv6.rwkv.ln_out.weight)
        torch.testing.assert_close(model.model.norm.weight, rwkv6.rwkv.ln_out.weight)
    if model.model.norm.bias is not None:
        print("rwkv6.rwkv.ln_out.bias -> model.model.norm.bias")
        model.model.norm.bias.data.copy_(rwkv6.rwkv.ln_out.bias)
        torch.testing.assert_close(model.model.norm.bias, rwkv6.rwkv.ln_out.bias)
    model.model.norm.eps = rwkv6.rwkv.ln_out.eps

    if not model.config.tie_word_embeddings:
        print("rwkv6.rwkv.head.weight -> model.lm_head.weight")
        model.lm_head.weight.data.copy_(rwkv6.head.weight)
        torch.testing.assert_close(model.lm_head.weight, rwkv6.head.weight)

    print(f"Saving converted model \n{model}\n to {output} ...")
    model.save_pretrained(output)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default='RWKV/rwkv-6-world-7b')
    parser.add_argument("--config", default='configs/rwkv6_7B.json')
    parser.add_argument("--output", default='converted/rwkv6-7B')
    args = parser.parse_args()
    convert(args.model, args.config, args.output)