memit_main.py

import os
from copy import deepcopy
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from rome.layer_stats import layer_stats
from util import nethook
from util.generate import generate_fast
from util.globals import *

from .compute_ks import compute_ks
from .compute_z import compute_z, get_module_input_output_at_words, find_fact_lookup_idx
from .memit_hparams import MEMITHyperParams

# Cache variable(s)
CONTEXT_TEMPLATES_CACHE = None
COV_CACHE = {}
flag = 1


def apply_memit2model_modified(
        model: AutoModelForCausalLM,
        tok: AutoTokenizer,
        requests: List[Dict],
        hparams: MEMITHyperParams,
        device: str,
        copy=False,
        return_orig_weights=False,
        cache_template: Optional[str] = None,
) -> Tuple[AutoModelForCausalLM, Dict[str, Any]]:
    """
    Returns a model with the desired changes.
    :param copy: If true, will preserve the original model while creating a new one to edit.
        Note that you are responsible for deallocating the new model's memory to avoid leaks.
    :return: (1) the updated model, (2) an original copy of the weights that changed
    """

    deltas, covs, kkts, ks = execute_memit(
        model, tok, requests, hparams, device, cache_template=cache_template)

    return deltas, covs, kkts, ks


def apply_memit_to_model(
        model: AutoModelForCausalLM,
        tok: AutoTokenizer,
        requests: List[Dict],
        hparams: MEMITHyperParams,
        copy=False,
        return_orig_weights=False,
        cache_template: Optional[str] = None,
) -> Tuple[AutoModelForCausalLM, Dict[str, Any]]:
    """
    Returns a model with the desired changes.
    :param copy: If true, will preserve the original model while creating a new one to edit.
        Note that you are responsible for deallocating the new model's memory to avoid leaks.
    :return: (1) the updated model, (2) an original copy of the weights that changed
    """
    device = "cpu"  # cuda / cpu
    weights_copy = {}
    if copy:
        model = deepcopy(model)
    device = "cuda"
    deltas, _, __, ___ = execute_memit(
        model, tok, requests, hparams, device, cache_template=cache_template)

    with torch.no_grad():
        for w_name, (key_mat, val_mat) in deltas.items():
            key_mat, val_mat = key_mat.to(device), val_mat.to(device)
            upd_matrix = key_mat @ val_mat.T
            w = nethook.get_parameter(model, w_name)
            upd_matrix = upd_matrix_match_shape(upd_matrix, w.shape)

            if return_orig_weights and w_name not in weights_copy:
                weights_copy[w_name] = w.detach().clone()

            w[...] += upd_matrix.float()

    print(f"New weights successfully inserted into {list(deltas.keys())}")

    return model, weights_copy


Covs = None


def execute_memit(
        model: AutoModelForCausalLM,
        tok: AutoTokenizer,
        requests: List[Dict],
        hparams: MEMITHyperParams,
        device: str,
        cache_template: Optional[str] = None,
) -> Dict[str, Tuple[torch.Tensor]]:
    """
    Executes the MEMIT update algorithm for the specified update at the specified layer
    Invariant: model at beginning of function == model at end of function
    """
    # device = "cpu"  # cuda / cpu
    deltas = {}

    # Update target and print info
    requests = deepcopy(requests)
    for i, request in enumerate(requests):
        if request["target_new"]["str"][0] != " ":
            # Space required for correct tokenization
            requests[i]["target_new"]["str"] = " " + \
                request["target_new"]["str"]
    # for request in requests[:10]:
    #     print(
    #         f"MEMIT request sample: "
    #         f"[{request['prompt'].format(request['subject'])}] -> [{request['target_new']['str']}]"
    #     )

    # Retrieve weights that user desires to change
    weights = {
        f"{hparams.rewrite_module_tmp.format(layer)}.weight": nethook.get_parameter(
            model, f"{hparams.rewrite_module_tmp.format(layer)}.weight"
        )
        for layer in hparams.layers
    }
    # Save old weights for future restoration
    weights_copy = {k: v.detach().clone() for k, v in weights.items()}

    # Compute z for final layer
    context_templates = get_context_templates(model, tok)
    z_layer = hparams.layers[-1]
    z_list = []

    for request in requests:
        # Retrieve k/v pair if already stored in cache
        cache_fname = (
            Path(
                str(cache_template).format(
                    z_layer, hparams.clamp_norm_factor, request["case_id"]
                )
            )
            if cache_template is not None
            else None
        )
        data_loaded = False
        if (
                cache_fname is not None  # Require cache template
                and cache_fname.exists()  # Cache file must exist
        ):
            try:
                data = np.load(cache_fname)
                z_list.append(torch.from_numpy(data["v_star"]).to(device))
                data_loaded = True
            except Exception as e:
                print(f"Error reading cache file due to {e}. Recomputing...")

        # Compute k/v pair if not loaded from cache
        if not data_loaded:
            cur_z = compute_z(
                model,
                tok,
                request,
                hparams,
                z_layer,
                context_templates,
                device
            )

            z_list.append(cur_z)

            if cache_fname is not None:
                cache_fname.parent.mkdir(exist_ok=True, parents=True)
                np.savez(
                    cache_fname,
                    **{
                        "v_star": cur_z.detach().cpu().numpy(),
                    },
                )
                # print(f"Cached k/v pair at {cache_fname}")
    zs = torch.stack(z_list, dim=1)
    # print("zs: ",zs)

    # test conflict

    # if zs.shape[1] != 1:
    #     a = zs[:, 0]
    #     b = zs[:, zs.shape[1] - 1]
    #     # mid = torch.add(a, b) / 2
    #     torch.save(a, "./ATensor.pt")
    #     torch.save(b, "./BTensor.pt")
    # else:

    #     B = torch.load("./BTensor.pt")
    #     demozs = deepcopy(zs)
    #     demozs = torch.squeeze(demozs)
    #     r1 = torch.linalg.norm(A - demozs)
    #     print(r1)
    #     r2 = torch.linalg.norm(B - demozs)
    #     print(r2)

    # Insert
    lastz = None
    kkts = {}
    covs = {}
    ks = {}
    for i, layer in enumerate(hparams.layers):
        # print(f"\n\nLAYER {layer}\n")

        # Get current model activations
        layer_ks = compute_ks(model, tok, requests,
                              hparams, layer, context_templates).T
        # print(f"Writing {layer_ks.size(1)} key/value pair(s) into layer {layer}")

        # Compute residual error
        cur_zs = get_module_input_output_at_words(
            model,
            tok,
            z_layer,
            context_templates=[request["prompt"] for request in requests],
            words=[request["subject"] for request in requests],
            module_template=hparams.layer_module_tmp,
            fact_token_strategy=hparams.fact_token,
        )[1].T
        targets = zs - cur_zs
        # print("z error", torch.linalg.norm(targets, dim=0).mean())

        repeat_factor = (layer_ks.size(1) // targets.size(1))
        targets = targets.repeat_interleave(repeat_factor, dim=1)

        # Load covariance matrix
        force_recompute = False
        # force_recompute = layer != hparams.layers[0]
        cov = get_cov(
            model,
            tok,
            hparams.rewrite_module_tmp.format(layer),
            hparams.mom2_dataset,
            hparams.mom2_n_samples
            if not force_recompute
            else hparams.mom2_n_samples // 10,
            hparams.mom2_dtype,
            device,
            force_recompute=force_recompute,
        )

        # Compute update in double precision
        layer_ks, targets = (
            layer_ks.double(),
            targets.double(),
        )

        # if Covs == None:
        #     Covs = deepcopy(cov)
        # else:
        #     Covs = Covs + layer_ks @ layer_ks.T

        adj_k = torch.linalg.solve(
            hparams.mom2_update_weight * cov.double() + layer_ks @ layer_ks.T,
            layer_ks,
        )
        # Distribute residual across layers
        resid = targets / (len(hparams.layers) - i)
        # R_L2 = torch.linalg.norm(resid)
        upd_matrix = resid @ adj_k.T

        # Adjust update matrix shape
        weight_name = f"{hparams.rewrite_module_tmp.format(layer)}.weight"
        upd_matrix = upd_matrix_match_shape(
            upd_matrix, weights[weight_name].shape)
        # orig_norm = torch.linalg.norm(weights[weight_name])
        # upd_norm = torch.linalg.norm(upd_matrix)
        # print("orig norm", orig_norm)
        # print("upd norm", upd_norm)

        # test delta bias
        lastz = cur_zs

        # Update model weights and record desired changes in `delta` variable
        with torch.no_grad():
            # weights[weight_name][...] = weights_copy[weight_name] + upd_matrix.float()
            deltas[weight_name] = (
                adj_k.detach().cpu(),
                resid.detach().cpu(),
            )
            kkts[weight_name] = (layer_ks @ layer_ks.T).cpu()
            covs[weight_name] = (
                hparams.mom2_update_weight * cov.double()).cpu()
            ks[weight_name] = (layer_ks).cpu()

        # Clear GPU memory
        cov.cpu()
        for x in [layer_ks, cur_zs, targets]:
            x.cpu()
            del x
        torch.cuda.empty_cache()

    # test conflict
    # if zs.shape[1] == 1:
    #     A = torch.load("./ATensor.pt")
    #     # torch.save(zs[:,0], "./ATensor.pt")
    #     B = torch.load("./BTensor.pt")
    #     # torch.save(B, "./BTensor.pt")
    #     demozs = deepcopy(lastz)
    #     demozs = torch.squeeze(demozs)
    #     r1 = torch.linalg.norm(A - demozs)
    #     print(r1)
    #     r2 = torch.linalg.norm(B - demozs)
    #     print(r2)

    # Restore state of original model
    with torch.no_grad():
        for k, v in weights.items():
            v[...] = weights_copy[k]
    torch.cuda.empty_cache()
    # print(f"Deltas successfully computed for {list(weights.keys())}")

    return deltas, covs, kkts, ks


def get_cov(
        model: AutoModelForCausalLM,
        tok: AutoTokenizer,
        layer_name: str,
        mom2_dataset: str,
        mom2_n_samples: str,
        mom2_dtype: str,
        device: str,
        inv: bool = False,
        force_recompute: bool = False,
) -> torch.Tensor:
    """
    Retrieves covariance statistics, then computes the algebraic inverse.
    Caches result for future use.
    """
    # device = "cpu"  # cuda / cpu

    model_name = model.config._name_or_path.replace("/", "_")
    key = (model_name, layer_name)

    # print(f"Retrieving covariance statistics for {model_name} @ {layer_name}.")
    if key not in COV_CACHE or force_recompute:
        stat = layer_stats(
            model,
            tok,
            layer_name,
            STATS_DIR,
            mom2_dataset,
            to_collect=["mom2"],
            sample_size=mom2_n_samples,
            precision=mom2_dtype,
            force_recompute=force_recompute,
        )
        COV_CACHE[key] = stat.mom2.moment().float().to("cpu")

    return (
        torch.inverse(COV_CACHE[key].to(device)
                      ) if inv else COV_CACHE[key].to(device)
    )


def upd_matrix_match_shape(matrix: torch.Tensor, shape: torch.Size) -> torch.Tensor:
    """
    GPT-2 and GPT-J have transposed weight representations.
    Returns a matrix that matches the desired shape, else raises a ValueError
    """

    if matrix.shape == shape:
        return matrix
    elif matrix.T.shape == shape:
        return matrix.T
    else:
        raise ValueError(
            "Update matrix computed by MEMIT does not match original weight shape. "
            "Check for bugs in the code?"
        )


def get_context_templates(model, tok):
    global CONTEXT_TEMPLATES_CACHE

    if CONTEXT_TEMPLATES_CACHE is None:
        CONTEXT_TEMPLATES_CACHE = [["{}"]] + [
            [
                f.replace("{", " ").replace("}", " ") + ". {}"
                for f in generate_fast(
                    model,
                    tok,
                    ["The", "Therefore", "Because", "I", "You"],
                    n_gen_per_prompt=n_gen // 5,
                    max_out_len=length,
                )
            ]
            for length, n_gen in [(10, 5)]  # Be careful about changing this.
        ]
        # print(f"Cached context templates {CONTEXT_TEMPLATES_CACHE}")

    return CONTEXT_TEMPLATES_CACHE