Ahora, entremos en el verdadero significado de este artículo. ¡Analizando matrices (Q, K, V, O) del modelo Llama-3–8B-Instruct a través de sus valores singulares!
El código
Primero importemos todos los paquetes necesarios en este análisis.
import transformers
import torch
import numpy as np
from transformers import AutoConfig, LlamaModel
from safetensors import safe_open
import os
import matplotlib.pyplot as plt
Luego, descarguemos el modelo y guárdelo en nuestro local. /tmpdirectorio.
MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct"
!huggingface-cli download {MODEL_ID} --quiet --local-dir /tmp/{MODEL_ID}
Si dispone de GPU, es posible que el siguiente código no sea relevante para usted. Sin embargo, si tienes poca GPU como yo, el siguiente código será realmente útil para cargar solo capas específicas del modelo LLama-3–8B.
def load_specific_layers_safetensors(model, model_name, layer_to_load):
state_dict = {}
files = [f for f in os.listdir(model_name) if f.endswith('.safetensors')]
for file in files:
filepath = os.path.join(model_name, file)
with safe_open(filepath, framework="pt") as f:
for key in f.keys():
if f"layers.{layer_to_load}." in key:
new_key = key.replace(f"model.layers.{layer_to_load}.", 'layers.0.')
state_dict[new_key] = f.get_tensor(key)missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
if missing_keys:
print(f"Missing keys: {missing_keys}")
if unexpected_keys:
print(f"Unexpected keys: {unexpected_keys}")
La razón por la que hacemos esto es porque el nivel gratuito de la GPU Google Colab no es suficiente para cargar LLama-3–8B incluso con fp16 precisión. Además, este análisis requiere que trabajemos en fp32 precisión debido a cómo np.linalg.svd está construído. A continuación, podemos definir la función principal para obtener valores singulares para un determinado matrix_type , layer_number y head_number.
def get_singular_values(model_path, matrix_type, layer_number, head_number):
"""
Computes the singular values of the specified matrix in the Llama-3 model.Parameters:
model_path (str): Path to the model
matrix_type (str): Type of matrix ('q', 'k', 'v', 'o')
layer_number (int): Layer number (0 to 31)
head_number (int): Head number (0 to 31)
Returns:
np.array: Array of singular values
"""
assert matrix_type in ['q', 'k', 'v', 'o'], "Invalid matrix type"
assert 0 <= layer_number < 32, "Invalid layer number"
assert 0 <= head_number < 32, "Invalid head number"
# Load the model only for that specific layer since we have limited RAM even after using fp16
config = AutoConfig.from_pretrained(model_path)
config.num_hidden_layers = 1
model = LlamaModel(config)
load_specific_layers_safetensors(model, model_path, layer_number)
# Access the specified layer
# Always index 0 since we have loaded for the specific layer
layer = model.layers[0]
# Determine the size of each head
num_heads = layer.self_attn.num_heads
head_dim = layer.self_attn.head_dim
# Access the specified matrix
weight_matrix = getattr(layer.self_attn, f"{matrix_type}_proj").weight.detach().numpy()
if matrix_type in ['q','o']:
start = head_number * head_dim
end = (head_number + 1) * head_dim
else: # 'k', 'v' matrices
# Adjust the head_number based on num_key_value_heads
# This is done since llama3-8b use Grouped Query Attention
num_key_value_groups = num_heads // config.num_key_value_heads
head_number_kv = head_number // num_key_value_groups
start = head_number_kv * head_dim
end = (head_number_kv + 1) * head_dim
# Extract the weights for the specified head
if matrix_type in ['q', 'k', 'v']:
weight_matrix = weight_matrix[start:end, :]
else: # 'o' matrix
weight_matrix = weight_matrix[:, start:end]
# Compute singular values
singular_values = np.linalg.svd(weight_matrix, compute_uv=False)
del model, config
return list(singular_values)
Vale la pena señalar que podemos extraer los pesos para la cabeza especificada en las matrices K, Q y V realizando un corte por filas debido a cómo se implementa mediante AbrazosCara.
(d_out,d_in). Fuente: Imagen del autor.En cuanto a la matriz O, podemos realizar cortes en columnas para extraer los pesos para la cabeza especificada en el peso O gracias al álgebra lineal. Los detalles se pueden ver en la siguiente figura.