| """carga de modelos y extracción por-cabeza de (w_v, w_o). |
| |
| reimplementación autocontenida, para este repo, de la extracción ya |
| verificada en el paper 1 (angular-separation-vit): gpt-neox fusiona |
| q,k,v por cabeza intercalado (``query_key_value``); timm concatena |
| [q,k,v] por bloque y v ocupa el último tercio (``qkv``). convención |
| de forma en todo el repo: w_v^(h) es [d_h, d], w_o^(h) es [d, d_h]. |
| """ |
|
|
| import torch |
|
|
|
|
| def cargar_pythia( |
| model_id: str = "EleutherAI/pythia-410m", |
| dtype: torch.dtype = torch.float32, |
| ): |
| """carga pythia-410m y verifica su anatomía contra el config real. |
| |
| args: |
| model_id: identificador del modelo en hugging face hub. |
| dtype: precisión de carga. |
| |
| returns: |
| el modelo en modo evaluación. |
| """ |
| from transformers import AutoModelForCausalLM |
|
|
| modelo = AutoModelForCausalLM.from_pretrained( |
| model_id, dtype=dtype, attn_implementation="eager") |
| modelo.eval() |
| cfg = modelo.config |
| dim_cabeza = cfg.hidden_size // cfg.num_attention_heads |
| assert dim_cabeza == 64, f"dim_cabeza inesperada: {dim_cabeza}" |
| assert cfg.num_attention_heads == 16, "num_attention_heads != 16" |
| assert cfg.hidden_size == 1024, "hidden_size != 1024" |
| assert len(modelo.gpt_neox.layers) == 24, "num capas != 24" |
| return modelo |
|
|
|
|
| @torch.no_grad() |
| def w_v_w_o_pythia( |
| modelo, |
| capa: int, |
| n_cabezas: int = 16, |
| dim_cabeza: int = 64, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """(w_v, w_o) por cabeza de una capa de pythia, en float64. |
| |
| query_key_value fusiona q, k y v por cabeza: cada cabeza ocupa |
| ``3*dim_cabeza`` filas contiguas ``[q_h|k_h|v_h]``; v_h es el |
| último tercio. dense es la proyección de salida (columnas por |
| cabeza). |
| |
| args: |
| modelo: modelo gpt-neox (pythia) cargado. |
| capa: índice de capa. |
| n_cabezas: cabezas h. |
| dim_cabeza: d_h. |
| |
| returns: |
| (w_v [h, d_h, d], w_o [h, d, d_h]), float64. |
| """ |
| attn = modelo.gpt_neox.layers[capa].attention |
| w_qkv = attn.query_key_value.weight |
| w_dense = attn.dense.weight |
| w_v, w_o = [], [] |
| for h in range(n_cabezas): |
| base = h * 3 * dim_cabeza + 2 * dim_cabeza |
| w_v.append(w_qkv[base:base + dim_cabeza, :].double()) |
| w_o.append(w_dense[:, h * dim_cabeza:(h + 1) * dim_cabeza] |
| .double()) |
| return torch.stack(w_v), torch.stack(w_o) |
|
|
|
|
| def cargar_vitb( |
| ckpt_path: str, |
| model_name: str = "vit_base_patch16_224", |
| num_classes: int = 100, |
| img_size: int = 224, |
| ): |
| """carga la columna vit-b con los pesos de un checkpoint de a. |
| |
| args: |
| ckpt_path: ruta al checkpoint (.pt con clave 'model', |
| prefijo 'backbone.' en las claves). |
| model_name: identificador timm de la columna. |
| num_classes: clases de la cabeza. |
| img_size: resolución de entrada. |
| |
| returns: |
| el vit interno de timm (el que expone .blocks), en eval. |
| """ |
| import timm |
|
|
| modelo = timm.create_model( |
| model_name, pretrained=False, num_classes=num_classes, |
| img_size=img_size) |
| blob = torch.load(ckpt_path, map_location="cpu", weights_only=False) |
| |
| |
| |
| |
| sd = {k.replace("backbone.model.", "", 1): v |
| for k, v in blob["model"].items()} |
| modelo.load_state_dict(sd) |
| return modelo.eval() |
|
|
|
|
| @torch.no_grad() |
| def w_v_w_o_vitb( |
| modelo, |
| capa: int, |
| n_cabezas: int = 12, |
| dim_cabeza: int = 64, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """(w_v, w_o) por cabeza de una capa de vit-b, en float64. |
| |
| en timm el qkv concatena [q, k, v] por bloque; v ocupa el último |
| tercio y se reordena por cabeza. proj es la proyección de salida. |
| |
| args: |
| modelo: el vit interno de timm. |
| capa: índice de capa. |
| n_cabezas: cabezas h. |
| dim_cabeza: d_h. |
| |
| returns: |
| (w_v [h, d_h, d], w_o [h, d, d_h]), float64. |
| """ |
| attn = modelo.blocks[capa].attn |
| w_qkv = attn.qkv.weight |
| w_proj = attn.proj.weight |
| base_v = 2 * w_qkv.shape[1] |
| w_v, w_o = [], [] |
| for h in range(n_cabezas): |
| fil = slice(base_v + h * dim_cabeza, base_v + (h + 1) * dim_cabeza) |
| col = slice(h * dim_cabeza, (h + 1) * dim_cabeza) |
| w_v.append(w_qkv[fil, :].double()) |
| w_o.append(w_proj[:, col].double()) |
| return torch.stack(w_v), torch.stack(w_o) |
|
|
|
|
| @torch.no_grad() |
| def escribe_v_o_vitb( |
| modelo, |
| capa: int, |
| h: int, |
| w_v_h: torch.Tensor, |
| w_o_h: torch.Tensor, |
| dim_cabeza: int = 64, |
| ) -> None: |
| """escribe (w_v, w_o) de una cabeza de vuelta en el modelo, in situ. |
| |
| inversa de `w_v_w_o_vitb` para una sola cabeza: usada por la |
| validación end-to-end para materializar un modelo con los pesos |
| ya gaugeados-y-cuantizados antes del forward real. |
| |
| args: |
| modelo: el vit interno de timm, modificado in situ. |
| capa: índice de capa. |
| h: índice de cabeza. |
| w_v_h: tensor [d_h, d] (cualquier dtype, se castea al del |
| modelo). |
| w_o_h: tensor [d, d_h]. |
| dim_cabeza: d_h. |
| """ |
| attn = modelo.blocks[capa].attn |
| w_qkv = attn.qkv.weight |
| w_proj = attn.proj.weight |
| base_v = 2 * w_qkv.shape[1] |
| fil = slice(base_v + h * dim_cabeza, base_v + (h + 1) * dim_cabeza) |
| col = slice(h * dim_cabeza, (h + 1) * dim_cabeza) |
| w_qkv[fil, :] = w_v_h.to(w_qkv.dtype) |
| w_proj[:, col] = w_o_h.to(w_proj.dtype) |
|
|
|
|
| @torch.no_grad() |
| def b_v_vitb( |
| modelo, |
| capa: int, |
| n_cabezas: int = 12, |
| dim_cabeza: int = 64, |
| ) -> torch.Tensor | None: |
| """sesgo de valor por cabeza de una capa de vit-b, en float64. |
| |
| el sesgo de valor entra en la invariancia de gauge: como |
| v = x w_v^t + b_v, la compensación de w_o exige que b_v gire igual |
| que la weight (b_v <- r^t b_v). omitirlo en una columna con |
| qkv_bias rompe la invariancia, y el forward deja de medir el |
| gauge para medir otro modelo. |
| |
| args: |
| modelo: el vit interno de timm. |
| capa: índice de capa. |
| n_cabezas: cabezas h. |
| dim_cabeza: d_h. |
| |
| returns: |
| tensor [h, d_h] en float64, o None si la columna no lleva |
| qkv_bias. |
| """ |
| attn = modelo.blocks[capa].attn |
| if attn.qkv.bias is None: |
| return None |
| b_qkv = attn.qkv.bias |
| base_v = 2 * (b_qkv.shape[0] // 3) |
| return torch.stack([ |
| b_qkv[base_v + h * dim_cabeza: |
| base_v + (h + 1) * dim_cabeza].double() |
| for h in range(n_cabezas)]) |
|
|
|
|
| @torch.no_grad() |
| def escribe_b_v_vitb( |
| modelo, |
| capa: int, |
| h: int, |
| b_v_h: torch.Tensor, |
| dim_cabeza: int = 64, |
| ) -> None: |
| """escribe el sesgo de valor de una cabeza, in situ. |
| |
| inversa de `b_v_vitb` para una sola cabeza. el sesgo no se |
| cuantiza: el objeto cuantizado de la nota es el par de weights |
| (w_v, w_o), y el sesgo solo viaja para que el gauge siga siendo |
| un cambio de coordenadas y no un cambio de función. |
| |
| args: |
| modelo: el vit interno de timm, modificado in situ. |
| capa: índice de capa. |
| h: índice de cabeza. |
| b_v_h: tensor [d_h], se castea al dtype del modelo. |
| dim_cabeza: d_h. |
| """ |
| attn = modelo.blocks[capa].attn |
| if attn.qkv.bias is None: |
| raise ValueError("la columna no lleva qkv_bias") |
| b_qkv = attn.qkv.bias |
| base_v = 2 * (b_qkv.shape[0] // 3) |
| fil = slice(base_v + h * dim_cabeza, base_v + (h + 1) * dim_cabeza) |
| b_qkv[fil] = b_v_h.to(b_qkv.dtype) |
|
|
