import argparse import os import jax import copy import json import torch import optax import numpy as np from tqdm import tqdm import jax.numpy as jnp from jax import random, vmap from typing import Any, Dict, List from flax.core import freeze, unfreeze from flax.serialization import from_bytes from utils import flatten_params, lerp, unflatten_params from flax.traverse_util import flatten_dict, unflatten_dict from lmc_model import LMCFlaxGPT2LMHeadModel, print_model from transformers.models.gpt2.modeling_flax_gpt2 import GPT2Config, FlaxGPT2LMHeadModel from flax.training import checkpoints, train_state from flax.training.common_utils import get_metrics, onehot, shard from matching_utils import weight_matching_attn, all_matching_attn from data_utils import get_lm_corpus import matplotlib.pyplot as plt import numpy as np # def check_params_nan(params, name="params"): # flat_params = flatten_dict(unfreeze(params)) # found_nan = False # for path, val in flat_params.items(): # if isinstance(val, jnp.ndarray): # if jnp.isnan(val).any(): # print(f"🚫 NaN detected in {name}: {'/'.join(path)} | shape={val.shape}") # found_nan = True # if not found_nan: # print(f"✅ No NaNs found in {name}.") # else: # print(f"⚠️ NaNs found in {name}.") def load_flax_params(checkpoint_dir, target): """ Loads Flax parameters from a msgpack file. Args: checkpoint_dir (str): path to folder containing flax_model.msgpack target (PyTree): a target object matching the structure of your model's parameters Returns: PyTree: deserialized parameters """ msgpack_path = os.path.join(checkpoint_dir, "flax_model.msgpack") with open(msgpack_path, "rb") as f: packed_bytes = f.read() params = from_bytes(target, packed_bytes) return params def prepare_lm_batch(data: torch.Tensor, target: torch.Tensor) -> Dict[str, Any]: """ Convert and shard a language modeling batch from PyTorch to JAX. Args: data (torch.Tensor): Input data of shape (seq_len, batch) target (torch.Tensor): Target data of shape (seq_len, batch) Returns: Dict[str, jnp.ndarray]: Dict with 'data' and 'target', both sharded with shape (n_devices, batch_per_device, seq_len) """ # Transpose to (batch, seq_len), then convert to jnp arrays input_ids = jnp.array(data.T) target = jnp.array(target.T) # Shard across devices return {'input_ids': input_ids,'target': target} def make_stuff(model): apply_fn = model.__call__ @jax.jit def batch_eval(params, batch): labels = batch.pop("target") # destructive, just like your train_step logits = apply_fn(**batch, params=params, train=False)[0] loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1])).mean() return logits, labels, loss @jax.jit def step(state: train_state.TrainState, batch, dropout_rng): dropout_rng, new_dropout_rng = jax.random.split(dropout_rng, 2) def loss_fn(params): labels = batch.pop("target") logits = apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0] loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1])).mean() return loss, logits (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params) new_state = state.apply_gradients(grads=grads) metrics = {"batch_loss": loss,"logits": logits,} return new_state, metrics, new_dropout_rng def dataset_loss_and_ppl(params, dataloader): """ Iterate once over `dataloader`, pplumulate token-level CE, and return (mean_loss, perplexity). Works on **one device**. """ total_loss, total_tok = 0.0, 0 pbar = tqdm(dataloader, desc="Evaluating", leave=False) for eval_data, eval_target, _ in pbar: eval_batch = prepare_lm_batch(eval_data, eval_target) _, labels, loss = batch_eval(params,eval_batch) ntok = jnp.sum(labels != -100) total_loss += loss * ntok total_tok += ntok pbar.set_postfix(loss=f"{loss:.4f}", ppl=f"{np.exp(loss):.2f}") mean_loss = (total_loss / total_tok).item() ppl = jnp.exp(mean_loss).item() return mean_loss, ppl return {"batch_eval": batch_eval,"step": step,"dataset_loss_and_ppl": dataset_loss_and_ppl,} def compute_interpolation(params_a, params_b_target, lambdas, stuff, val_ds, test_ds, desc="Interpolation"): train_loss_interp, test_loss_interp = [], [] train_ppl_interp, test_ppl_interp = [], [] for lam in tqdm(lambdas, desc=desc): p_interp = freeze(lerp(lam, unfreeze(params_a), unfreeze(params_b_target))) train_loss, train_ppl = stuff["dataset_loss_and_ppl"](p_interp, val_ds) # test_loss, test_ppl = stuff["dataset_loss_and_ppl"](p_interp, test_ds) test_loss, test_ppl = train_loss, train_ppl train_loss_interp.append(train_loss) test_loss_interp.append(test_loss) train_ppl_interp.append(train_ppl) test_ppl_interp.append(test_ppl) return { "Val Loss": [float(f"{x:.4f}") for x in train_loss_interp], "Test Loss": [float(f"{x:.4f}") for x in test_loss_interp], "Val PPL": [float(f"{x:.4f}") for x in train_ppl_interp], "Test PPL": [float(f"{x:.4f}") for x in test_ppl_interp] } def main(): parser = argparse.ArgumentParser() parser.add_argument("--model-a", type=str, required=True, help="Path to first fine-tuned GPT2 model checkpoint") parser.add_argument("--model-b", type=str, required=True, help="Path to second fine-tuned GPT@ model checkpoint") parser.add_argument("--seed", type=int, default =0) parser.add_argument("--data-path", type=str, default="./data/lm1b", help="train datset paths (multiple paths)") parser.add_argument('--dataset', type=str, default='lm1b',choices=['wt103', 'lm1b', 'enwik8', 'text8'],help='dataset name') parser.add_argument("--batch-size", type=int, default=24, help="train, eval batch size (batch size will be devided by device count)") parser.add_argument('--tgt_len', type=int, default=256,help='number of tokens to predict') parser.add_argument('--eval_tgt_len', type=int, default=256,help='number of tokens to predict for evaluation') parser.add_argument('--ext_len', type=int, default=0,help='length of the extended context') parser.add_argument('--mem_len', type=int, default=0,help='length of the retained previous heads') parser.add_argument("--save-path", type=str, default="/", help="Path to plot directory") parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype") args = parser.parse_args() corpus = get_lm_corpus(args.data_path, args.dataset) ntokens = len(corpus.vocab) eval_batch_size = 12 tr_iter = corpus.get_iterator('train', args.batch_size, args.tgt_len, ext_len=args.ext_len) va_iter = corpus.get_iterator('valid', eval_batch_size, args.eval_tgt_len, ext_len=args.ext_len) te_iter = corpus.get_iterator('test', eval_batch_size, args.eval_tgt_len, ext_len=args.ext_len) val_ds = va_iter test_ds = te_iter config = GPT2Config.from_json_file(os.path.join(os.path.dirname(args.model_a).rstrip("/"),'config.json')) lmc_config = GPT2Config(**config.lmc_config) config.lmc_config = lmc_config model = LMCFlaxGPT2LMHeadModel(config,input_shape=(1, args.tgt_len),seed=args.seed,dtype=jnp.dtype(args.dtype),) print_model(model.params) if os.path.exists(args.model_a) and os.path.exists(args.model_b): params_a = load_flax_params(args.model_a,copy.deepcopy(model.params)) # check_params_nan(params_a, name="params_a") params_b = load_flax_params(args.model_b,copy.deepcopy(model.params)) # check_params_nan(params_b, name="params_b") else: raise FileNotFoundError(f"Checkpoint path does not exist") stuff = make_stuff(model = model) lambdas = jnp.linspace(0, 1, num=25) rng = random.PRNGKey(args.seed) # Compute naive interpolation naive_results = compute_interpolation(params_a, params_b, lambdas, stuff, val_ds, test_ds, desc="Naive Interpolation") print(json.dumps({"Naive": naive_results}, indent=2)) all_results = {"Naive": naive_results} # Compute weight matching interpolations for each method aligned_models = weight_matching_attn(rng, params_a, params_b, None, config) for method, params_b_aligned in aligned_models.items(): method_results = compute_interpolation(params_a, params_b_aligned, lambdas, stuff, val_ds, test_ds, desc=f"{method} Interpolation") all_results[method] = method_results print(json.dumps({method: method_results}, indent=2)) # permutations_results = {} # aligned_models = all_matching_attn(rng, params_a, params_b, config) # for method, params_b_aligned in aligned_models.items(): # method_results = compute_interpolation(params_a, params_b_aligned, lambdas, stuff, val_ds, test_ds, desc=f"{method} Interpolation") # permutations_results[method] = method_results # print(json.dumps({method: method_results}, indent=2)) # Save directories os.makedirs(f"./plots/{args.dataset}", exist_ok=True) os.makedirs(f"./results/{args.dataset}", exist_ok=True) # Save results JSON print("Save List of Values...") name_a = os.path.basename(os.path.dirname(args.model_a).rstrip("/")) name_b = os.path.basename(os.path.dirname(args.model_b).rstrip("/")) result_path = f'results/{args.dataset}/[{name_a}+{name_b}].json' with open(result_path, 'w') as f: json.dump(all_results, f, indent=2) # permute_path = f'results/{args.dataset}/permute[{name_a}+{name_b}].json' # with open(permute_path, 'w') as f: # json.dump(permutations_results, f, indent=2) # Plot # print("Generating plots...") # plot_path = f"./plots/{args.dataset}/[{name_a}+{name_b}].pdf" # plt.rcParams.update({ # "font.family": "serif", # 'legend.frameon': False, # 'lines.linewidth': 2, # 'font.size': 13, # 'axes.labelsize': 16, # 'xtick.labelsize': 11, # 'ytick.labelsize': 11, # 'legend.fontsize': 11, # }) # plt.style.use('tableau-colorblind10') # num_points = len(all_results["Naive"]["Val Loss"]) # lambda_values = np.linspace(0, 1, num_points) # fig, axs = plt.subplots(1, 2, figsize=(12, 5)) # 1 row, 2 columns # metrics = ["Val Loss", "Test Loss"] # custom_colors = [ # "#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd", # "#8c564b", "#e377c2", "#7f7f7f", "#bcbd22", "#17becf" # ] # for idx, metric in enumerate(metrics): # ax = axs[idx] # for j, method in enumerate(all_results): # Naive + aligned methods # ax.plot(lambda_values, # all_results[method][metric], # label=method, # color=custom_colors[j % len(custom_colors)]) # ax.set_xticks([0, 0.5, 1]) # ax.set_xticklabels(["Model 1", r"$\lambda$", "Model 2"]) # ax.set_ylabel(metric) # ax.legend(loc='best') # plt.tight_layout() # plt.savefig(plot_path.replace(".pdf", "_row0.pdf")) # plt.close() # fig, axs = plt.subplots(2, 2, figsize=(12, 10)) # metrics = ["Val Loss", "Test Loss", "Val Loss", "Test Loss"] # positions = [(0, 0), (0, 1), (1, 0), (1, 1)] # custom_colors = [ # "#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd", # "#8c564b", "#e377c2", "#7f7f7f", "#bcbd22", "#17becf" # ] # for metric, pos in zip(metrics, positions): # row, col = pos # ax = axs[row, col] # if row == 0: # for idx, method in enumerate(all_results): # ax.plot(lambda_values, all_results[method][metric], label=method,color=custom_colors[idx % len(custom_colors)]) # if row == 1: # for idx, method in enumerate(permutations_results): # ax.plot(lambda_values, permutations_results[method][metric], label=method) # ax.set_xticks([0, 0.5, 1]) # ax.set_xticklabels(["Model 1", r"$\lambda$", "Model 2"]) # ax.set_ylabel(metric) # ax.legend(loc='best') # plt.tight_layout() # plt.savefig(plot_path) # plt.close() if __name__ == "__main__": main()