mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
338 lines
13 KiB
Python
338 lines
13 KiB
Python
if __name__ == "__main__":
|
||
# New imports for analysis
|
||
from collections import defaultdict
|
||
from functools import partial
|
||
|
||
import torch
|
||
from peft import get_peft_model
|
||
from torch.utils.data import DataLoader
|
||
from tqdm import tqdm
|
||
|
||
try:
|
||
from sklearn.manifold import TSNE
|
||
except ImportError:
|
||
TSNE = None
|
||
try:
|
||
import umap
|
||
except ImportError:
|
||
umap = None
|
||
# Added plotting imports
|
||
try:
|
||
import matplotlib.pyplot as plt
|
||
except ImportError:
|
||
plt = None
|
||
try:
|
||
import seaborn as sns
|
||
except ImportError:
|
||
sns = None
|
||
|
||
from ctx_to_lora.data.collator import generation_collator
|
||
from ctx_to_lora.data.definitions import CTX_AFFIXES
|
||
from ctx_to_lora.data.processing import get_tokenized_dataset
|
||
from ctx_to_lora.data.self_gen_template import SELF_QA_INTX
|
||
from ctx_to_lora.model_loading import get_lora_config, get_model_and_tokenizer
|
||
|
||
# from ctx_to_lora.modeling.context_distillation import CtxDistillModel
|
||
from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
|
||
|
||
model_name_or_path = "google/gemma-2-2b-it"
|
||
base_model, tokenizer = get_model_and_tokenizer(
|
||
model_name_or_path, train=False, requires_grad=False
|
||
)
|
||
ds = get_tokenized_dataset(
|
||
ds_name="squad",
|
||
split="test",
|
||
max_qas_len=-1,
|
||
max_qas_per_sample=1,
|
||
base_model_max_len=8192,
|
||
tokenizer=tokenizer,
|
||
ctx_model_max_len=8192,
|
||
ctx_tokenizer=tokenizer,
|
||
max_ctx_chunk_len=-1,
|
||
min_ctx_chunk_len=-1,
|
||
num_chunk_probs=None,
|
||
max_ctx_chunk_num=None,
|
||
add_ctx_to_chat=False,
|
||
use_kl_loss=True,
|
||
max_new_tokens=0,
|
||
add_self_distill_template=False,
|
||
set_format="pt",
|
||
)
|
||
ds = ds.take(100)
|
||
|
||
peft_config = get_lora_config(
|
||
model_name_or_path,
|
||
lora_r=8,
|
||
lora_dropout=0,
|
||
target_modules=["down_proj"],
|
||
)
|
||
peft_config.lora_alpha = 16
|
||
peft_model = get_peft_model(base_model, peft_config)
|
||
sep_seq = (
|
||
tokenizer(
|
||
SELF_QA_INTX.strip("\n"),
|
||
add_special_tokens=False,
|
||
return_tensors="pt",
|
||
)
|
||
.input_ids[0]
|
||
.to(base_model.device)
|
||
)
|
||
ctx_distill_kwargs = dict(
|
||
prefix_tokens=torch.tensor(
|
||
CTX_AFFIXES[model_name_or_path]["prefix"], device=base_model.device
|
||
),
|
||
ctx_inp_sep_seq=sep_seq,
|
||
pad_token_id=tokenizer.pad_token_id,
|
||
update_iterations=100,
|
||
tokenizer=tokenizer,
|
||
reprompt_ctx=False,
|
||
)
|
||
|
||
# if args.cd_use_gen_q:
|
||
# q_model, q_tokenizer = get_model_and_tokenizer(
|
||
# "google/gemma-3-4b-it",
|
||
# train=False,
|
||
# requires_grad=False,
|
||
# )
|
||
# ctx_distill_kwargs["q_model"] = q_model
|
||
# ctx_distill_kwargs["q_tokenizer"] = q_tokenizer
|
||
# model = CtxDistillModel(peft_model, **ctx_distill_kwargs)
|
||
checkpoint_path = "train_outputs/runs/Sep06_12-27-02_slurm0-a3nodeset-11_84523_bbcb67ca/checkpoint-40000/pytorch_model.bin"
|
||
model = ModulatedPretrainedModel.from_state_dict(
|
||
state_dict=torch.load(checkpoint_path, weights_only=False),
|
||
train=False,
|
||
base_model_kwargs=dict(attn_implementation="flash_attention_2"),
|
||
use_flash_attn=True,
|
||
use_sequence_packing=False, # for generation
|
||
)
|
||
dataloader = DataLoader(
|
||
ds,
|
||
batch_size=1,
|
||
shuffle=False,
|
||
sampler=None,
|
||
batch_sampler=None,
|
||
num_workers=0,
|
||
collate_fn=partial(generation_collator, tokenizer=tokenizer),
|
||
pin_memory=False,
|
||
drop_last=False,
|
||
timeout=0,
|
||
persistent_workers=False,
|
||
)
|
||
|
||
# Control / config variables for analysis
|
||
collect_max_samples = None # set int to limit samples processed
|
||
modules_filter_substr = None # e.g. "down_proj" to restrict; or None for all
|
||
perform_analysis = True # set False to only collect
|
||
save_dir = "./lora_weight_analysis" # will be created if not exists
|
||
use_fp32_storage = False # store ΔW in float32 (else float16 if original)
|
||
apply_pca_before_tsne = True
|
||
pca_dim = 64
|
||
random_seed = 42
|
||
perform_plotting = True # NEW: set False to skip figure generation
|
||
|
||
# Data structures:
|
||
# delta_w_flat[module] -> list[Tensor(num_params,)] (one per sample)
|
||
delta_w_flat = defaultdict(list)
|
||
sample_ids = [] # overall sample index order
|
||
# Keep raw shapes (out_features, in_features)
|
||
delta_w_shapes = {}
|
||
|
||
# Precompute scaling factor (α / r)
|
||
lora_alpha = peft_config.lora_alpha
|
||
lora_r = peft_config.r
|
||
lora_scaling = lora_alpha / lora_r
|
||
|
||
for sample_idx, sample in tqdm(enumerate(dataloader)):
|
||
if collect_max_samples is not None and sample_idx >= collect_max_samples:
|
||
break
|
||
for k, v in sample.items():
|
||
sample[k] = v.to(model.device)
|
||
with torch.no_grad():
|
||
lora_state, _ = model.generate_weights(
|
||
sample["ctx_ids"], sample["ctx_attn_mask"]
|
||
)
|
||
|
||
# # NEW: lora_state is now a nested dict:
|
||
# # { module_name: { "A": [1, n_layers, r, d_in], "B": [1, n_layers, r, d_out] } }
|
||
for module_name, parts in lora_state.items():
|
||
if "A" not in parts or "B" not in parts:
|
||
continue
|
||
A_all = parts["A"].squeeze(0).float() # [n_layers, r, d_in]
|
||
B_all = parts["B"].squeeze(0).float() # [n_layers, r, d_out]
|
||
n_layers_local = A_all.size(0)
|
||
for layer_idx in range(n_layers_local):
|
||
A = A_all[layer_idx] # [r, d_in]
|
||
B = B_all[layer_idx] # [r, d_out]
|
||
# Original expectation was B: [d_out, r]; now transpose
|
||
delta_w = B.T @ A # [d_out, d_in]
|
||
mod_key = f"{module_name}.layer{layer_idx}"
|
||
if modules_filter_substr and modules_filter_substr not in mod_key:
|
||
continue
|
||
delta_w_shapes.setdefault(mod_key, delta_w.shape)
|
||
if use_fp32_storage:
|
||
delta_w = delta_w.float()
|
||
delta_w_flat[mod_key].append(delta_w.reshape(-1).cpu())
|
||
sample_ids.append(sample_idx)
|
||
|
||
if perform_analysis and len(delta_w_flat):
|
||
import os
|
||
|
||
os.makedirs(save_dir, exist_ok=True)
|
||
torch.manual_seed(random_seed)
|
||
|
||
# Optional PCA helper
|
||
def run_pca(x, k):
|
||
# x: (n, d)
|
||
if k >= x.shape[1]:
|
||
return x
|
||
# Center
|
||
mean = x.mean(0, keepdim=True)
|
||
x0 = x - mean
|
||
# SVD economy
|
||
try:
|
||
U, S, Vt = torch.linalg.svd(x0, full_matrices=False)
|
||
except RuntimeError:
|
||
# fallback CPU
|
||
U, S, Vt = torch.linalg.svd(x0.cpu(), full_matrices=False)
|
||
U, S, Vt = U.to(x0.device), S.to(x0.device), Vt.to(x0.device)
|
||
comp = Vt[:k]
|
||
return torch.matmul(x0, comp.T)
|
||
|
||
analysis_summary = {}
|
||
|
||
for module_name, vec_list in delta_w_flat.items():
|
||
X = torch.stack(vec_list) # (num_samples, num_params)
|
||
# Cosine similarity matrix
|
||
X_norm = torch.nn.functional.normalize(X, dim=1)
|
||
cosine_sim = X_norm @ X_norm.T # (N,N)
|
||
print(cosine_sim)
|
||
|
||
# Dimensionality reduction inputs
|
||
X_dr = X
|
||
if apply_pca_before_tsne and TSNE is not None:
|
||
# Keep at most pca_dim or (num_samples - 1, num_params)
|
||
k = min(pca_dim, X.shape[1])
|
||
X_dr = run_pca(X.to(torch.float32), k)
|
||
|
||
embeddings = {}
|
||
if TSNE is not None and X.shape[0] > 2:
|
||
try:
|
||
perplexity = max(2, min(30, X.shape[0] - 1))
|
||
tsne_2d = TSNE(
|
||
n_components=2,
|
||
perplexity=perplexity,
|
||
random_state=random_seed,
|
||
init="pca",
|
||
learning_rate="auto",
|
||
).fit_transform(X_dr.numpy())
|
||
embeddings["tsne"] = torch.from_numpy(tsne_2d)
|
||
except Exception as e:
|
||
print(f"[WARN] TSNE failed for {module_name}: {e}")
|
||
|
||
if umap is not None and X.shape[0] > 2:
|
||
try:
|
||
reducer = umap.UMAP(
|
||
n_components=2,
|
||
n_neighbors=min(15, max(2, X.shape[0] - 1)),
|
||
random_state=random_seed,
|
||
metric="cosine",
|
||
)
|
||
umap_2d = reducer.fit_transform(X_dr.numpy())
|
||
embeddings["umap"] = torch.from_numpy(umap_2d)
|
||
except Exception as e:
|
||
print(f"[WARN] UMAP failed for {module_name}: {e}")
|
||
|
||
# Save tensors
|
||
base_fname = module_name.replace(".", "_")
|
||
torch.save(
|
||
{
|
||
"module": module_name,
|
||
"sample_ids": sample_ids,
|
||
"delta_W_flat": X, # shape (N, P)
|
||
"shape_matrix": delta_w_shapes[module_name],
|
||
"cosine_sim": cosine_sim, # (N, N)
|
||
"embeddings": embeddings, # dict
|
||
"config": {
|
||
"lora_alpha": lora_alpha,
|
||
"lora_r": lora_r,
|
||
"scaling": lora_scaling,
|
||
"pca_applied": apply_pca_before_tsne and TSNE is not None,
|
||
},
|
||
},
|
||
os.path.join(save_dir, f"{base_fname}.pt"),
|
||
)
|
||
|
||
analysis_summary[module_name] = {
|
||
"num_samples": X.shape[0],
|
||
"num_params": X.shape[1],
|
||
"saved_file": os.path.join(save_dir, f"{base_fname}.pt"),
|
||
"have_tsne": "tsne" in embeddings,
|
||
"have_umap": "umap" in embeddings,
|
||
}
|
||
|
||
# -------- NEW PLOTTING SECTION --------
|
||
if perform_plotting and plt is not None:
|
||
# Cosine similarity heatmap
|
||
try:
|
||
fig, ax = plt.subplots(figsize=(4, 4))
|
||
|
||
im = ax.imshow(cosine_sim.numpy(), cmap="viridis")
|
||
ax.set_xticks(range(len(sample_ids)))
|
||
ax.set_yticks(range(len(sample_ids)))
|
||
ax.set_xticklabels(sample_ids, fontsize=6)
|
||
ax.set_yticklabels(sample_ids, fontsize=6)
|
||
plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
|
||
ax.set_title(
|
||
f"Cosine Sim: {module_name.split('.')[-3:]}".replace(",", "")
|
||
)
|
||
ax.set_xlabel("Sample")
|
||
ax.set_ylabel("Sample")
|
||
fig.tight_layout()
|
||
fig.savefig(
|
||
os.path.join(save_dir, f"{base_fname}_cosine.png"), dpi=300
|
||
)
|
||
plt.close(fig)
|
||
except Exception as e:
|
||
print(f"[WARN] Cosine plot failed for {module_name}: {e}")
|
||
|
||
# TSNE / UMAP scatter
|
||
def _scatter(emb_key):
|
||
emb = embeddings.get(emb_key)
|
||
if emb is None:
|
||
return
|
||
try:
|
||
emb_np = emb.numpy()
|
||
fig, ax = plt.subplots(figsize=(4, 4))
|
||
ax.scatter(
|
||
emb_np[:, 0],
|
||
emb_np[:, 1],
|
||
s=30,
|
||
c=range(len(emb_np)),
|
||
cmap="tab10",
|
||
)
|
||
if len(emb_np) <= 20:
|
||
for i, (xv, yv) in enumerate(emb_np):
|
||
ax.text(xv, yv, str(sample_ids[i]), fontsize=6)
|
||
ax.set_title(f"{emb_key.upper()} {module_name.split('.')[-3:]}")
|
||
ax.set_xticks([])
|
||
ax.set_yticks([])
|
||
fig.tight_layout()
|
||
fig.savefig(
|
||
os.path.join(save_dir, f"{base_fname}_{emb_key}.png"),
|
||
dpi=150,
|
||
)
|
||
plt.close(fig)
|
||
except Exception as e:
|
||
print(f"[WARN] {emb_key} plot failed for {module_name}: {e}")
|
||
|
||
_scatter("tsne")
|
||
_scatter("umap")
|
||
# -------- END NEW PLOTTING SECTION --------
|
||
|
||
# Print a concise summary
|
||
print("LoRA ΔW analysis summary:")
|
||
for m, info in analysis_summary.items():
|
||
print(
|
||
f"{m}: samples={info['num_samples']} params={info['num_params']} "
|
||
f"tsne={info['have_tsne']} umap={info['have_umap']} -> {info['saved_file']}"
|
||
)
|