get_llama_layers_statistics.py

import os
import torch
import warnings
import pandas as pd

from tqdm import tqdm
from huggingface_hub import login
from itertools import combinations
from transformers import AutoModel

from utils import _get_nodes, _get_layer_kinds

warnings.filterwarnings("ignore")


def _save_llama_layers_locally():
    """ Save llama layers locally """
    login()  # you should provide your own token
    nodes_ = _get_nodes(llama=True)

    for i, (node_name, node) in enumerate(nodes_):
        model = AutoModel.from_pretrained(node)
        for name, layer in model.state_dict().items():
            splits = name.split('layers.')
            if len(splits) == 1:
                continue

            print(name, '->', splits[-1].split('.')[0])
            block_idx_ = int(splits[-1].split('.')[0])

            if not os.path.exists(f'llama_blocks/{block_idx_}'):
                os.makedirs(f'llama_blocks/{block_idx_}')

            torch.save(layer, f'llama_blocks/{block_idx_}/{os.path.basename(node)}-{name}.pt')


if __name__ == '__main__':
    nodes = _get_nodes(llama=True)
    layer_kinds = _get_layer_kinds(llama=True)
    res = pd.DataFrame(columns=['block_idx', 'layer_kind', 'mean', 'std', 'min', 'max'])

    for block_idx in tqdm(range(32)):
        files = os.listdir(f'llama_blocks/{block_idx}')
        for kind in layer_kinds:
            layer_data = [torch.load(os.path.join(f'llama_blocks/{block_idx}', f)) for f in files if kind in f]
            dist_layer_data = torch.stack([(l1 - l2) for l1, l2 in combinations(layer_data, 2)])

            res.loc[res.shape[0]] = {'block_idx': block_idx,
                                     'layer_kind': kind,
                                     'mean': dist_layer_data.mean().item(),
                                     'std': dist_layer_data.std().item(),
                                     'min': dist_layer_data.min().item(),
                                     'max': dist_layer_data.max().item()}

    res.to_csv('llama_layers_statistics.csv')