analyze.py

import matplotlib.pyplot as plt
import torch
from clize import run
from outliers.functional.utils import add_global_plot_styles
from transformers import AutoModelForCausalLM, AutoTokenizer

def check_module_hidden_states(model, tokenizer, test_text, layer_path, module_name, input_or_output="output", plot_fname=None, spike_threshold=100):
    if input_or_output not in ["input", "output"]:
        raise ValueError("input_or_output should be 'input' or 'output', instead of", input_or_output)
    
    all_activations = {}

    def get_activations(layer_index):
        def hook(model, inputs, outputs):
            hidden_states = inputs if input_or_output == "input" else outputs
            all_activations.setdefault(layer_index, {})[f"{module_name}_{input_or_output}_hidden_states"] = hidden_states
        return hook   

    all_hooks = []

    def get_layers(model, layer_path):
        attributes = layer_path.split('.')
        layers = model
        for attr in attributes:
            layers = getattr(layers, attr)
        return layers

    attributes = module_name.split('.') if module_name != "layer" else []
    layers = get_layers(model, layer_path)

    for layer_index, layer in enumerate(layers):
        current_attr = layer
        valid = True
        for attr in attributes:
            if hasattr(current_attr, attr):
                current_attr = getattr(current_attr, attr)
            else:
                valid = False
                break
        
        if valid:
            hook = current_attr.register_forward_hook(get_activations(layer_index))
            all_hooks.append(hook)

    inputs = tokenizer(test_text, return_tensors='pt').to(model.device)
    model.eval()
    with torch.no_grad():
        model(**inputs)

    for hook in all_hooks:
        hook.remove()

    top1_values_all_layers = []
    top1_indexes_all_layers = []
    for layer_index, outputs in all_activations.items():
        values = outputs[f'{module_name}_{input_or_output}_hidden_states']
        tensor = values[0] if isinstance(values, tuple) else values
        tensor = tensor.detach().cpu()
        tensor_abs = tensor.view(-1).abs().float()

        max_value, max_index = torch.max(tensor_abs, 0)
        max_index = torch.unravel_index(max_index, tensor.shape)
        top1_values_all_layers.append(tensor[max_index])
        top1_indexes_all_layers.append(max_index)

    return top1_values_all_layers, top1_indexes_all_layers

def plot_down_proj_input_output(pretrained="allenai/OLMo-7B-0724-hf", module_name="mlp.down_proj"):
    model = AutoModelForCausalLM.from_pretrained(pretrained, trust_remote_code=True, torch_dtype=torch.float16)
    tokenizer = AutoTokenizer.from_pretrained(pretrained, trust_remote_code=True)

    test_text = "Apple Inc. is a worldwide tech company."
    layer_path = "model.layers"

    for name in ("input", "output"):
        magnitude, index = check_module_hidden_states(
            model, tokenizer, test_text, layer_path, module_name, input_or_output=name, spike_threshold=50)

        # Report any spikes
        spikes_input = [i for i, value in enumerate(magnitude) if abs(value) > 50]
        print(f"Activation spikes for {module_name} {name}:")
        for i in spikes_input:
            spike_index = index[i]
            print(f" - layer {i}, value {magnitude[i]}, index {tuple(i.item() for i in spike_index)}")

        # Plot input activations
        plt.figure(figsize=(5,3.5))
        add_global_plot_styles()
        plt.plot(range(len(magnitude)), magnitude, color='blue', marker='o', markersize=5)
        plt.xlabel('Layer Number')
        plt.ylabel('Max Activation Value')
        plt.title(f"OLMo-7B Max down_proj {name}")
        plt.yticks(rotation=90, va='center')
        plt.savefig(f"outputs/figures/{name}_down_proj.pdf", bbox_inches='tight')
        print(f"Plot saved to 'outputs/figures/{name}_down_proj.pdf'")

        # Print output magnitudes
        print(f"largest_activations_down_proj_{name}={list(map(float, magnitude))}")

def record_SO(model, pretrained):
    '''Record SO values for original models'''
    SUPER_WEIGHTS_MAP = {
        "Mistral-7B-v0.1": [(1, 2070, 7310)],
        "llama-7B": [(2, 3968, 7003)],
        "llama-13B": [(2, 2231, 2278), (2, 2231, 6939)],
        "llama-30B": [(3, 5633, 12817), (3, 5633, 17439), (10, 5633, 14386)],
        "Meta-Llama-3-8B": [(1, 788, 2427), (1, 1384, 2427), (1, 4062, 2427)],
        "OLMo-1B-0724-hf": [(1, 1764, 1710), (2, 1764, 8041)],
        "OLMo-7B-0724-hf": [(1, 269, 7467), (2, 269, 8275), (7, 269, 453), (24, 269, 2300)],
        "Phi-3-mini-4k-instruct": [(2, 525, 808), (2, 1693, 808), (2, 1113, 808), (4, 525, 2723),  (4, 1113, 2723), (4, 1693, 2723)],
    }
    
    def _record_SO(SO_map, layer, row, col):
        if pretrained in [
            "tiiuae/falcon-7b",
        ]:
            SO_map[(layer, row, col)] = model.transformer.h[layer].mlp.dense_4h_to_h.weight.data[row, col].item()
        else:
            SO_map[(layer, row, col)] = model.model.layers[layer].mlp.down_proj.weight.data[row, col].item()

    SO_values = {}
    for model_name, coordinates in SUPER_WEIGHTS_MAP.items():
        if model_name in pretrained:
            for layer, row, col in coordinates:
                _record_SO(SO_values, layer, row, col)
            break
    return SO_values

def scale_SO(model, pretrained, SO_values, scaling_factor):
    if pretrained in [
        "huggyllama/llama-30B", 
        "huggyllama/llama-13B", 
        "huggyllama/llama-7B", 
        "mistralai/Mistral-7B-v0.1",
        "meta/Meta-Llama-3-8B",
        "allenai/OLMo-1B-0724-hf",
        "allenai/OLMo-7B-0724-hf",
        "microsoft/Phi-3-mini-4k-instruct"
        ]:
        for (layer, row, col), value in SO_values.items():  
            old_value = model.model.layers[layer].mlp.down_proj.weight.data[row, col].item()
            new_value = value * scaling_factor
            model.model.layers[layer].mlp.down_proj.weight.data[row, col] = new_value
            print(f"Layer {layer}, Index [{row}, {col}], Old value: {old_value}, New value: {new_value}")

def remove_outliers(model, pretrained, percentage_threshold):

    num_selected_elements = []
    for name, param in model.named_parameters():
        if not name.endswith("weight"):
            continue

        weight = param.data
        num_top_elements = int(weight.numel() * percentage_threshold)
        # print("# Total params:", weight.numel(), "# top params", num_top_elements)
        if num_top_elements < 1: # too few elements to apply on
            continue 
        threshold = torch.topk(weight.view(-1).abs(), num_top_elements).values[-1]
        mask = weight.abs() >= threshold
        true_indices = mask.nonzero(as_tuple=False)
        num_selected_elements.append(len(true_indices))

        weight[mask] = 0.
        param = torch.nn.Parameter(weight)   
    


def plot_max_activation_ablation(pretrained="allenai/OLMo-7B-0724-hf"):
    model = AutoModelForCausalLM.from_pretrained(pretrained, trust_remote_code=True, torch_dtype=torch.float16)
    tokenizer = AutoTokenizer.from_pretrained(pretrained, trust_remote_code=True)

    
    test_text = "Apple Inc. is a worldwide tech company."
    layer_path = "model.layers"
    module_name = "layer"
    name = "output"
    # original
    magnitude, index = check_module_hidden_states(
            model, tokenizer, test_text, layer_path, module_name, input_or_output=name, spike_threshold=50)
    # Print output magnitudes
    print(f"original={list(map(float, magnitude))}")
    
    # remove SO
    SO_values = record_SO(model, pretrained)
    scale_SO(model, pretrained, SO_values, 0)
    magnitude, index = check_module_hidden_states(
            model, tokenizer, test_text, layer_path, module_name, input_or_output=name, spike_threshold=50)
    # Report any spikes
    spikes_input = [i for i, value in enumerate(magnitude) if abs(value) > 50]
    print(f"Activation spikes for {module_name} {name}:")
    for i in spikes_input:
        spike_index = index[i]
        print(f" - layer {i}, value {magnitude[i]}, index {tuple(i.item() for i in spike_index)}")
    # Print output magnitudes
    print(f"super_weight_removed={list(map(float, magnitude))}")

    # remove outliers
    percentage = 5e-7
    remove_outliers(model, pretrained, percentage)
    magnitude, index = check_module_hidden_states(
            model, tokenizer, test_text, layer_path, module_name, input_or_output=name, spike_threshold=50)
    # Print output magnitudes
    print(f"all_outliers_removed={list(map(float, magnitude))}")

    # restore SO
    scale_SO(model, pretrained, SO_values, 1)
    magnitude, index = check_module_hidden_states(
            model, tokenizer, test_text, layer_path, module_name, input_or_output=name, spike_threshold=50)
    # Print output magnitudes
    print(f"all_other_outliers_removed={list(map(float, magnitude))}")

def plot_token_probs(pretrained="mistralai/Mistral-7B-v0.1"):

    model_name = pretrained
    MODEL_ID = model_name.split('/')[1]

    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float16,
        device_map="auto",
        output_hidden_states=True,
        low_cpu_mem_usage=True,
    )

    model_map = {
        "Mistral-7B-v0.1": [(1, 2070, 7310)],
        "llama-7B": [(2, 3968, 7003)],
        "llama-13B": [(2, 2231, 2278), (2, 2231, 6939)],
        "llama-3v0B": [(3, 5633, 12817), (3, 5633, 17439), (10, 5633, 14386)],
        "Meta-Llama-3-8B": [(1, 788, 2427), (1, 1384, 2427), (1, 4062, 2427)],
        "OLMo-1B-0724-hf": [(1, 1764, 1710), (2, 1764, 8041)],
        "OLMo-7B-0724-hf": [(1, 269, 7467), (2, 269, 8275), (7, 269, 453), (24, 269, 2300)],
        "gemma-7b": [(0, 1995, 21041)], # not sufficient
        "Phi-3-mini-4k-instruct": [(2, 525, 808), (2, 1693, 808), (2, 1113, 808), (4, 525, 2723),  (4, 1113, 2723), (4, 1693, 2723)],
        # "tiiuae/falcon-7b": [(3, 2002, 10708), (4, 2002, 5921)]
    }
    sw_map = {}

    def remove_SO(model):
        sw_map[MODEL_ID] = []
        for (layerno, y, x) in model_map[MODEL_ID]:
            weight = model.model.layers[layerno].mlp.down_proj.weight.data
            sw_map[MODEL_ID].append(float(weight[y, x]))
            weight[y, x] = 0.
            model.model.layers[layerno].mlp.down_proj.weight = torch.nn.Parameter(weight)

    def restore_SO(model):
        assert sw_map.get(MODEL_ID, None), "Run remove_SO before running restore_SO"
        for value, (layerno, y, x) in zip(sw_map[MODEL_ID], model_map[MODEL_ID]):
            weight = model.model.layers[layerno].mlp.down_proj.weight.data
            weight[y, x] = value
            model.model.layers[layerno].mlp.down_proj.weight = torch.nn.Parameter(weight)

    def print_SO(model):
        for weight, (layerno, y, x) in zip(sw_map[MODEL_ID], model_map[MODEL_ID]):
            weight = model.model.layers[layerno].mlp.down_proj.weight.data
            print(weight[y, x])


    def get_next_token_probs(model, tokenizer, input_text):
        input_ids = tokenizer.encode(input_text, return_tensors="pt").to(model.device)

        # Get model outputs
        with torch.no_grad():
            outputs = model(input_ids)
            
            # Get the logits for the next token prediction
            next_token_logits = outputs.logits[:, -1, :]
            next_token_probs = torch.softmax(next_token_logits, dim=-1)
        return next_token_probs


    from datasets import load_dataset
    import json
    from tqdm import tqdm

    N_SAMPLES = 500
    dataset = load_dataset("EleutherAI/lambada_openai", "en", split="test")

    all_difference = []

    all_probs_SO_removed = []
    all_probs_Original = []


    # Original model
    for text in tqdm(dataset[:N_SAMPLES]["text"]):
        prompt = ' '.join(text.split(' ')[:-1])
        target = text.split(' ')[-1]
        next_token_probs = get_next_token_probs(model, tokenizer, prompt)
        all_probs_Original.append(next_token_probs)

    avg_probs_Original = (sum(all_probs_Original)  / len(all_probs_Original))[0] 
    avg_probs_Original = avg_probs_Original.to('cpu')
    sorted_probs_Original, sorted_indices_Original = torch.sort(avg_probs_Original, descending=True)
    top_n = 100
    top_n_probs_Original = sorted_probs_Original[:top_n].tolist()
    top_n_indices_Original = sorted_indices_Original[:top_n]
    top_tokens = [tokenizer.decode(i) for i in top_n_indices_Original]

    # Remove super weight
    remove_SO(model)
    print(sw_map[MODEL_ID])

    for text in tqdm(dataset[:N_SAMPLES]["text"]):
        prompt = ' '.join(text.split(' ')[:-1])
        target = text.split(' ')[-1]
        next_token_probs = get_next_token_probs(model, tokenizer, prompt)
        all_probs_SO_removed.append(next_token_probs)

    # Average probabilities acorss all samples
    avg_probs_SO_removed = (sum(all_probs_SO_removed) / len(all_probs_SO_removed))[0]
    avg_probs_SO_removed = avg_probs_SO_removed.to('cpu')
    selected_token_probs_SO_removed = []
    for i in top_n_indices_Original:
        selected_token_probs_SO_removed.append(avg_probs_SO_removed[i].item())

    print("Top n tokens:")
    print(top_tokens)
    print("Original")
    print(top_n_probs_Original)
    print("SW removed")
    print(selected_token_probs_SO_removed)
    



if __name__ == '__main__':
    run([plot_down_proj_input_output, plot_max_activation_ablation, plot_token_probs])