to_hf.py

import torch
import argparse
from tqdm  import tqdm
from copy import deepcopy
from transformers import AutoTokenizer
from save_and_load import load_lr_tensors
from modelutils import get_opt, find_layers

base_model_name = "facebook/opt-1.3b"
base_model = get_opt(base_model_name)

target_model = deepcopy(base_model)

tensors = load_lr_tensors("outputs/model.safetensors")

target_layers = target_model.model.decoder.layers

for i in tqdm(range(len(target_layers))):
    subset = find_layers(target_layers[i])
    for name in subset:
        layer_id = f"decoder.layers.{i}.{name}"
        layer = find_layers(target_model.model)
        R = tensors[f"<R>.model.{layer_id}"]
        L = tensors[f"<L>.model.{layer_id}"]
        new_weight = torch.matmul(L, R)
        layer[layer_id].weight.data = new_weight

# save target model as HF
target_model.save_pretrained("outputs/lnair-opt-1.3b-wikitext2-r32-w8")
tokenizer = AutoTokenizer.from_pretrained(base_model_name)
tokenizer.save_pretrained("outputs/lnair-opt-1.3b-wikitext2-r32-w8")