diff --git a/tmp/compare_fp16_vs_bf16.py b/tmp/compare_fp16_vs_bf16.py new file mode 100644 index 0000000..8048384 --- /dev/null +++ b/tmp/compare_fp16_vs_bf16.py @@ -0,0 +1,59 @@ +# filepath: /home/tan/research/ctx-to-lora/tmp/compare_fp16_vs_bf16.py + +import torch + + +def compare_fp16_vs_bf16(): + """ + Compare operations between float16 and bfloat16 tensors. + Computes difference and multiplication to see output dtypes. + """ + + # Create sample tensors with the same values but different dtypes + values = [[1.5, 2.3, 3.7], [4.1, 5.9, 6.2]] + + # Create tensors with different dtypes + tensor_fp16 = torch.tensor(values, dtype=torch.float16) + tensor_bf16 = torch.tensor(values, dtype=torch.bfloat16) + + print("Original tensors:") + print(f"FP16 tensor: {tensor_fp16}") + print(f"FP16 dtype: {tensor_fp16.dtype}") + print(f"BF16 tensor: {tensor_bf16}") + print(f"BF16 dtype: {tensor_bf16.dtype}") + print() + + # Compute difference (subtraction) + diff_result = tensor_fp16 - tensor_bf16 + print("Difference operation (FP16 - BF16):") + print(f"Result: {diff_result}") + print(f"Result dtype: {diff_result.dtype}") + print() + + # Compute multiplication + mult_result = tensor_fp16 * tensor_bf16 + print("Multiplication operation (FP16 * BF16):") + print(f"Result: {mult_result}") + print(f"Result dtype: {mult_result.dtype}") + print() + + # Test the reverse order to see if it matters + diff_result_reverse = tensor_bf16 - tensor_fp16 + mult_result_reverse = tensor_bf16 * tensor_fp16 + + print("Reverse order operations:") + print(f"BF16 - FP16 dtype: {diff_result_reverse.dtype}") + print(f"BF16 * FP16 dtype: {mult_result_reverse.dtype}") + print() + + # Show the actual values to check precision differences + print("Value comparison:") + print(f"Original values: {values}") + print(f"FP16 values: {tensor_fp16.tolist()}") + print(f"BF16 values: {tensor_bf16.tolist()}") + print(f"Difference values: {diff_result.tolist()}") + print(f"Multiplication values: {mult_result.tolist()}") + + +if __name__ == "__main__": + compare_fp16_vs_bf16() diff --git a/tmp/display_self_gen_data.py b/tmp/display_self_gen_data.py new file mode 100644 index 0000000..a507c6b --- /dev/null +++ b/tmp/display_self_gen_data.py @@ -0,0 +1,27 @@ +import sys + +from datasets import load_dataset +from transformers import AutoTokenizer + + +def display_sample(ds, tk, i): + sample = ds[i] + print(f"ctx: {tk.decode(sample['ctx_ids'])}") + for ids, (start, end) in zip(sample["input_ids"], sample["response_start_end"]): + print(f"ids: {ids}") + print(f"input: {tk.decode(ids)}") + print(f"response: {tk.decode(ids[start:])}") + len_input_ids = len(ids) + pad_len_right = len_input_ids - end + print(f"{pad_len_right=}") + + +data_files = sys.argv[1] + +ds = load_dataset("parquet", data_files=data_files, split="train") +tk = AutoTokenizer.from_pretrained("google/gemma-2-2b-it") + +for i in range(5): + display_sample(ds, tk, i) + +breakpoint() diff --git a/tmp/fit_sigmoid_curve.py b/tmp/fit_sigmoid_curve.py new file mode 100644 index 0000000..de6c26e --- /dev/null +++ b/tmp/fit_sigmoid_curve.py @@ -0,0 +1,78 @@ +import matplotlib.pyplot as plt +import numpy as np +from scipy.optimize import curve_fit + + +def sigmoid(x, k, x0): + """ + Sigmoid function with L fixed at 1. + + Parameters: + x (array-like): Input data (independent variable). + k (float): The steepness or growth rate of the curve. + x0 (float): The x-value of the sigmoid's midpoint. + """ + return 1 / (1 + np.exp(-k * (x - x0))) + + +# --- User Input --- +# Replace these with your actual x and y data points. +# You can have 3 or more points. +x_data = np.array([10, 20, 30]) +y_data = np.array([0.503, 0.554, 0.590]) +# -------------------- + +# Check if the data is suitable for a sigmoid curve (monotonically increasing) +if not np.all(np.diff(y_data) > 0): + print("Warning: The y-values of your data are not strictly increasing.") + print("A sigmoid fit may not be appropriate for this data.") + +try: + # Use curve_fit to find the best parameters + # p0 is an initial guess for the parameters [k, x0] + # We can estimate x0 as the mean of x_data as a starting point. + initial_guess = [1, np.mean(x_data)] + + # The 'bounds' argument constrains the parameters. We'll keep k positive. + popt, pcov = curve_fit( + sigmoid, x_data, y_data, p0=initial_guess, bounds=(0, np.inf) + ) + + # Extract the optimal parameters + k_opt, x0_opt = popt + print("Optimal Parameters (L=1):") + print(f"k = {k_opt:.4f}") + print(f"x0 = {x0_opt:.4f}") + + # --- Visualization --- + # Generate a smooth curve for plotting + x_curve = np.linspace(0, 300, 200) + y_curve = sigmoid(x_curve, k_opt, x0_opt) + + # Plot the data and the fitted curve + plt.figure(figsize=(10, 6)) + plt.plot(x_data, y_data, "o", label="Data Points", markersize=8, color="red") + plt.plot( + x_curve, + y_curve, + label=f"Fitted Sigmoid Curve\n(k={k_opt:.2f}, x₀={x0_opt:.2f})", + color="blue", + ) + plt.plot([0, 300], [0.823, 0.823], label="Ours") + plt.plot([0, 300], [0.988, 0.988], label="CD (oracle)") + plt.plot([40, 40], [0, 1.2], "--", label="OOM") + + plt.title("Overdetermined Sigmoid Curve Fit (L=1)") + plt.xlabel("x") + plt.ylabel("y") + plt.legend() + plt.grid(True) + plt.ylim(0, 1.1) # Set y-axis limits since L=1 + plt.show() + plt.savefig("sigmoid_fit.png") + +except RuntimeError as e: + print(f"Error: Could not find optimal parameters. {e}") + print( + "This can happen if the data is not sigmoid-shaped or if the initial guess is poor." + ) diff --git a/tmp/needle_haystack_accuracy_plot.pdf b/tmp/needle_haystack_accuracy_plot.pdf new file mode 100644 index 0000000..da5b25b Binary files /dev/null and b/tmp/needle_haystack_accuracy_plot.pdf differ diff --git a/tmp/needle_haystack_accuracy_plot.png b/tmp/needle_haystack_accuracy_plot.png new file mode 100644 index 0000000..d3ea29e Binary files /dev/null and b/tmp/needle_haystack_accuracy_plot.png differ diff --git a/tmp/performance_comparison.png b/tmp/performance_comparison.png new file mode 100644 index 0000000..bcc4cf9 Binary files /dev/null and b/tmp/performance_comparison.png differ diff --git a/tmp/performance_ctx_latency.pdf b/tmp/performance_ctx_latency.pdf new file mode 100644 index 0000000..85257b2 Binary files /dev/null and b/tmp/performance_ctx_latency.pdf differ diff --git a/tmp/performance_ctx_latency.png b/tmp/performance_ctx_latency.png new file mode 100644 index 0000000..26f456d Binary files /dev/null and b/tmp/performance_ctx_latency.png differ diff --git a/tmp/performance_ctx_latency_mem.pdf b/tmp/performance_ctx_latency_mem.pdf new file mode 100644 index 0000000..8a78913 Binary files /dev/null and b/tmp/performance_ctx_latency_mem.pdf differ diff --git a/tmp/performance_ctx_latency_mem.png b/tmp/performance_ctx_latency_mem.png new file mode 100644 index 0000000..74cbcdc Binary files /dev/null and b/tmp/performance_ctx_latency_mem.png differ diff --git a/tmp/performance_ctx_latency_mem_bar.pdf b/tmp/performance_ctx_latency_mem_bar.pdf new file mode 100644 index 0000000..26ddb44 Binary files /dev/null and b/tmp/performance_ctx_latency_mem_bar.pdf differ diff --git a/tmp/performance_ctx_latency_mem_bar.png b/tmp/performance_ctx_latency_mem_bar.png new file mode 100644 index 0000000..dc28c1f Binary files /dev/null and b/tmp/performance_ctx_latency_mem_bar.png differ diff --git a/tmp/performance_gen_mem.pdf b/tmp/performance_gen_mem.pdf new file mode 100644 index 0000000..ec9793b Binary files /dev/null and b/tmp/performance_gen_mem.pdf differ diff --git a/tmp/performance_gen_mem.png b/tmp/performance_gen_mem.png new file mode 100644 index 0000000..d2f039d Binary files /dev/null and b/tmp/performance_gen_mem.png differ diff --git a/tmp/performance_vs_latency.png b/tmp/performance_vs_latency.png new file mode 100644 index 0000000..4a6fccb Binary files /dev/null and b/tmp/performance_vs_latency.png differ diff --git a/tmp/plot_bar_perf_ctx_latency_mem.py b/tmp/plot_bar_perf_ctx_latency_mem.py new file mode 100644 index 0000000..a988e55 --- /dev/null +++ b/tmp/plot_bar_perf_ctx_latency_mem.py @@ -0,0 +1,375 @@ +from dataclasses import dataclass # added + +import matplotlib as mpl +import matplotlib.pyplot as plt +import numpy as np + +# Styling (reused) +latte_style = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle" +plt.style.use(["ggplot", latte_style]) +plt.rcParams["axes.prop_cycle"] = mpl.cycler( + color=[ + "#FFB83D", + # "#555555", + "#9BC750", + # "#35775A", + "#5571EA", + # "#6CACFF", + # "#E84494", + # "#E47FB0", + # "#785EF0", + # "#FF924E", + # "#5BC5DB", + # "#AD6F50", + # "#A1A9AD", + ] +) +plt.rcParams["axes.facecolor"] = "white" +plt.rcParams["figure.facecolor"] = "white" +plt.rcParams["savefig.facecolor"] = "white" + +# plt.rcParams["font.family"] = "Ubuntu" +# plt.rcParams["font.size"] = 14 +# plt.rcParams["font.weight"] = "bold" +plt.rcParams["axes.labelweight"] = "bold" +plt.rcParams["axes.facecolor"] = "F7FBFC" +plt.rcParams["figure.facecolor"] = "white" +plt.rcParams["savefig.facecolor"] = "white" +plt.rcParams["grid.color"] = "cccccc" +plt.rcParams["grid.linewidth"] = 1 +plt.rcParams["axes.edgecolor"] = "ccd0da" +plt.rcParams["legend.facecolor"] = "white" +plt.rcParams["legend.fontsize"] = 16 +plt.rcParams["xtick.labelsize"] = 13 +plt.rcParams["ytick.labelsize"] = 13 + + +@dataclass +class Point: + performance: float + context_ratio: float | None = None + latency: float | None = None + latency_std: float | None = None + peak_mem: float | None = None # added + peak_mem_std: float | None = None # added + + +def combine_std_for_sum(std_list): + return np.sqrt(sum(std**2 for std in std_list)) + + +# Unified data structure (merged context + latency via Point) +datasets = { + "2WikiMultihopQA\n(Long Context)": { + "Base model w/ context": [ + Point( + performance=0.3387, + context_ratio=1.0, + ), + ], + "CD": [ + Point( + performance=0.2385, + context_ratio=0.0, + latency=72.537, + latency_std=10, + peak_mem=79.371, + ), + ], + "Ours": [ + Point( + performance=0.2857, + context_ratio=0.0, + latency=0.551, + latency_std=combine_std_for_sum([43.423 * 1e-3, 29.673 * 1e-6]), + peak_mem=3.791, + ), + ], + }, + # "SQuAD": { + # "Base model w/ context": [ + # Point(performance=0.8692, context_ratio=1.0), + # ], + # # "Base model w/o context": [ + # # Point( + # # performance=0.1866, + # # context_ratio=0.0, + # # latency=0.0, + # # latency_std=0.0, + # # peak_mem=0.0, + # # ), + # # ], + # # "CD (oracle)": [ + # # Point( + # # performance=0.859, + # # context_ratio=0.0, + # # latency=39.974 + 456.989 * 1e-3, + # # latency_std=combine_std_for_sum([965.326 * 1e-3, 259.130 * 1e-3]), + # # peak_mem=855.762 / 1024, + # # ), + # # ], + # "CD": [ + # Point( + # performance=0.5986, + # context_ratio=0.0, + # latency=21.254 * 4 + 49.358 + 7.450, + # latency_std=10, + # peak_mem=41.098, + # ), + # ], + # # "Ours (batched)": [ + # # Point( + # # performance=0.7174, + # # # context_ratio=0.0, + # # latency=85.954 * 1e-3 + 259.920 * 1e-6, + # # latency_std=combine_std_for_sum([31.057 * 1e-3, 740.286 * 1e-6]), + # # peak_mem=1.151, + # # ), + # # ], + # "Ours": [ + # Point( + # performance=0.715, + # context_ratio=0.0, + # latency=426.849 * 1e-3 + 260.073 * 1e-6, + # latency_std=combine_std_for_sum([43.423 * 1e-3, 29.673 * 1e-6]), + # peak_mem=328.134 / 1024, + # ), + # ], + # # "T2L": [ + # # Point( + # # performance=0.1761, + # # context_ratio=0.0, + # # latency=35.480 * 1e-3, + # # latency_std=22.193 * 1e-3, + # # peak_mem=1.793, + # # ), + # # ], + # # "LLMLingua-2": [ + # # Point(performance=0.8492, context_ratio=0.9), + # # Point(performance=0.8278, context_ratio=0.8), + # # Point(performance=0.7749, context_ratio=0.6), + # # Point(performance=0.6624, context_ratio=0.4), + # # Point(performance=0.4271, context_ratio=0.2), + # # Point(performance=0.3192, context_ratio=0.1), + # # ], + # }, + # "DROP": { + # "Base model w/ context": [ + # Point(performance=0.4541, context_ratio=1.0), + # ], + # "Base model w/o context": [ + # Point( + # performance=0.1417, + # context_ratio=0.0, + # latency=0.0, + # latency_std=0.0, + # peak_mem=0.0, + # ), + # ], + # "CD (oracle)": [ + # Point( + # performance=0.443, + # context_ratio=0.0, + # latency=39.974 + 422.820 * 1e-3, + # latency_std=combine_std_for_sum([965.326 * 1e-3, 241.419 * 1e-3]), + # peak_mem=1.381, + # ), + # ], + # "CD": [ + # Point( + # performance=0.2289, + # context_ratio=0.0, + # latency=23.883 * 4 + 9.755 + 41.301, + # latency_std=14.8852904910, + # peak_mem=44.305, + # ), + # ], + # "Ours (batched)": [ + # Point( + # performance=0.2972, + # # context_ratio=0.0, + # latency=84.132 * 1e-3 + 251.876 * 1e-6, + # latency_std=combine_std_for_sum([30.503 * 1e-3, 22.829 * 1e-6]), + # peak_mem=1.976, + # ), + # ], + # "Ours": [ + # Point( + # performance=0.3002, + # context_ratio=0.0, + # latency=419.462 * 1e-3 + 254.899 * 1e-6, + # latency_std=combine_std_for_sum([40.211 * 1e-3, 29.784 * 1e-6]), + # peak_mem=429.881 / 1024, + # ), + # ], + # "T2L": [ + # Point( + # performance=0.1468, + # context_ratio=0.0, + # latency=35.036 * 1e-3, + # latency_std=18.144 * 1e-3, + # peak_mem=3.004, + # ), + # ], + # "LLMLingua-2": [ # + # Point(performance=0.4513, context_ratio=0.9), + # Point(performance=0.4488, context_ratio=0.8), + # Point(performance=0.4076, context_ratio=0.6), + # Point(performance=0.3488, context_ratio=0.4), + # Point(performance=0.2617, context_ratio=0.2), + # Point(performance=0.1712, context_ratio=0.1), + # ], + # }, + # "ROPES": { + # "Base model w/ context": [ + # Point(performance=0.7457, context_ratio=1.0), + # ], + # "Base model w/o context": [ + # Point( + # performance=0.4545, + # context_ratio=0.0, + # latency=0.0, + # latency_std=0.0, + # peak_mem=0.0, + # ), + # ], + # "CD (oracle)": [ + # Point( + # performance=0.7349, + # context_ratio=0.0, + # latency=39.633 + 285.467 * 1e-3, + # latency_std=combine_std_for_sum([1.752, 72.791 * 1e-3]), + # peak_mem=527.312 / 1024, + # ), + # ], + # "CD": [ + # Point( + # performance=0.5786, + # context_ratio=0.0, + # latency=23.714 * 4 + 11.473 + 59.660, + # latency_std=11.7329040736, + # peak_mem=43.271, + # ), + # ], + # "Ours (batched)": [ + # Point( + # performance=0.6759, + # # context_ratio=0.0, + # latency=85.652 * 1e-3 + 254.029 * 1e-6, + # latency_std=combine_std_for_sum([30.852 * 1e-3, 51.127 * 1e-6]), + # peak_mem=496.898 / 1024, + # ), + # ], + # "Ours": [ + # Point( + # performance=0.6786, + # context_ratio=0.0, + # latency=417.860 * 1e-3 + 257.422 * 1e-6, + # latency_std=combine_std_for_sum([45.514 * 1e-3, 31.419 * 1e-6]), + # peak_mem=273.867 / 1024, + # ), + # ], + # "T2L": [ + # Point( + # performance=0.4661, + # context_ratio=0.0, + # latency=37.913 * 1e-3, + # latency_std=17.313 * 1e-3, + # peak_mem=697.217 / 1024, + # ), + # ], + # "LLMLingua-2": [ + # Point(performance=0.708, context_ratio=0.9), + # Point(performance=0.7126, context_ratio=0.8), + # Point(performance=0.7291, context_ratio=0.6), + # Point(performance=0.677, context_ratio=0.4), + # Point(performance=0.6188, context_ratio=0.2), + # Point(performance=0.4894, context_ratio=0.1), + # ], + # }, +} + + +def build_color_map(datasets): + methods = [] + for data in datasets.values(): + for m in data.keys(): + if m not in methods: + methods.append(m) + palette = plt.rcParams["axes.prop_cycle"].by_key()["color"] + return {m: palette[i % len(palette)] for i, m in enumerate(methods)} + + +def plot_bar_methods(datasets): + # Use first dataset found + dataset_name, methods = next(iter(datasets.items())) + color_map = build_color_map(datasets) + + ref_perf = methods["Base model w/ context"][0].performance + # Aggregate metrics per method + stats_by_method = {} + for method, points in methods.items(): + if method == "Base model w/ context": + continue + if not points: + continue + + perf_vals = [p.performance for p in points if p.performance is not None] + lat_vals = [p.latency for p in points if p.latency is not None] + mem_vals = [p.peak_mem for p in points if p.peak_mem is not None] + perf = max(perf_vals) / ref_perf if perf_vals else np.nan + lat = float(np.mean(lat_vals)) if lat_vals else np.nan + mem = float(np.mean(mem_vals)) if mem_vals else np.nan + stats_by_method[method] = {"perf": perf, "latency": lat, "memory": mem} + + # Ordering by performance (descending), NaNs last + with_perf = [m for m, s in stats_by_method.items() if not np.isnan(s["perf"])] + without_perf = [m for m, s in stats_by_method.items() if np.isnan(s["perf"])] + order = ( + sorted(with_perf, key=lambda m: stats_by_method[m]["perf"], reverse=True) + + without_perf + ) + + # Build 3 horizontal bar plots: perf, latency, memory + fig, axes = plt.subplots(3, 1, figsize=(6, 4), sharex=False) + metric_specs = [ + ("perf", "Normalized\nPerformance"), + ("latency", "Latency\n(s)"), + ("memory", "Memory\n(GB)"), + ] + for ax, (metric_key, y_label) in zip(axes, metric_specs): + names = [m for m in order if not np.isnan(stats_by_method[m][metric_key])] + values = [stats_by_method[m][metric_key] for m in names] + if not names: + ax.axis("off") + continue + y = np.arange(len(names)) + colors = [color_map.get(m, "#888888") for m in names] + ax.barh(y, values, color=colors, edgecolor="black", alpha=0.9) + ax.set_yticks(y, labels=names) + ax.invert_yaxis() # highest at top to match ordering + ax.set_ylabel(y_label, fontweight="bold", fontsize=12) + ax.grid(True, axis="x", alpha=0.3) + # Make latency and memory log-scale + if metric_key in ("latency", "memory"): + ax.set_xscale("log") + + fig.suptitle(dataset_name, fontweight="bold", fontsize=20, y=0.91) + plt.tight_layout(rect=[0, 0, 1, 0.96]) + return fig + + +if __name__ == "__main__": + # fig = plot_ctx_and_latency(datasets) + fig = plot_bar_methods(datasets) + plt.show() + fig.savefig( + "/home/tan/research/ctx-to-lora/tmp/performance_ctx_latency_mem_bar.png", + dpi=300, + bbox_inches="tight", + ) + fig.savefig( + "/home/tan/research/ctx-to-lora/tmp/performance_ctx_latency_mem_bar.pdf", + dpi=300, + bbox_inches="tight", + ) diff --git a/tmp/plot_niah_mem_latency_vs_ctx_len.py b/tmp/plot_niah_mem_latency_vs_ctx_len.py new file mode 100644 index 0000000..5027d13 --- /dev/null +++ b/tmp/plot_niah_mem_latency_vs_ctx_len.py @@ -0,0 +1,251 @@ +from dataclasses import dataclass # added + +import matplotlib as mpl +import matplotlib.pyplot as plt +import numpy as np + +# Styling (reused) +latte_style = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle" +plt.style.use(["ggplot", latte_style]) +plt.rcParams["axes.prop_cycle"] = mpl.cycler( + color=[ + "#FFB83D", + "#5571EA", + ] +) +plt.rcParams["axes.facecolor"] = "white" +plt.rcParams["figure.facecolor"] = "white" +plt.rcParams["savefig.facecolor"] = "white" + +# plt.rcParams["font.family"] = "Ubuntu" +# plt.rcParams["font.size"] = 14 +# plt.rcParams["font.weight"] = "bold" +plt.rcParams["axes.labelweight"] = "bold" +plt.rcParams["axes.facecolor"] = "F7FBFC" +plt.rcParams["figure.facecolor"] = "white" +plt.rcParams["savefig.facecolor"] = "white" +plt.rcParams["grid.color"] = "cccccc" +plt.rcParams["grid.linewidth"] = 1 +plt.rcParams["axes.edgecolor"] = "ccd0da" +plt.rcParams["legend.facecolor"] = "white" +plt.rcParams["legend.fontsize"] = 16 +plt.rcParams["xtick.labelsize"] = 16 +plt.rcParams["ytick.labelsize"] = 16 + + +@dataclass +class Point: + context_length: int + latency: float | None = None + latency_std: float | None = None + peak_mem: float | None = None + + +def combine_std_for_sum(std_list): + return np.sqrt(sum(std**2 for std in std_list)) + + +# Unified data structure (merged context + latency via Point) +datasets = { + "NIAH": { + "Base model w/ ICL": [ + Point( + context_length=1024, + latency=375.689 * 1e-3, + latency_std=19.718 * 1e-3, + peak_mem=210.722 / 1024, + ), + Point( + context_length=2**11, + latency=375.428 * 1e-3, + latency_std=794.507 * 1e-6, + peak_mem=235.718 / 1024, + ), + Point( + context_length=2**12, + latency=384.812 * 1e-3, + latency_std=4.376 * 1e-3, + peak_mem=429.681 / 1024, + ), + Point( + context_length=2**13, + latency=438.886 * 1e-3, + latency_std=17.096 * 1e-3, + peak_mem=1.165, + ), + Point( + context_length=2**14, + # latency=13.877, + # latency_std=394.682, + peak_mem=1.740, + ), + Point( + context_length=2**15, + latency=13.985, + latency_std=66.566 * 1e-3, + peak_mem=3.223, + ), + Point( + context_length=2**16, + latency=14.953, + latency_std=131.882 * 1e-3, + peak_mem=5.278, + ), + Point( + context_length=2**17, + latency=18.489, + latency_std=227.430 * 1e-3, + peak_mem=12.072, + ), + ], + "Ours": [ + Point( + context_length=1024, + latency=300.995 * 1e-3, + latency_std=24.009 * 1e-3, + peak_mem=30.988 / 1024, + ), + Point( + context_length=2048, + latency=296.856 * 1e-3, + latency_std=568.892 * 1e-6, + peak_mem=3.207 / 1024, + ), + Point( + context_length=2**12, + latency=297.037 * 1e-3, + latency_std=548.641 * 1e-6, + peak_mem=4.363 / 1024, + ), + Point( + context_length=2**13, + latency=296.507 * 1e-3, + latency_std=18.850 * 1e-3, + peak_mem=33.449 / 1024, + ), + Point( + context_length=2**14, + latency=297.774 * 1e-3, + latency_std=788.634 * 1e-6, + peak_mem=8.792 / 1024, + ), + Point( + context_length=2**15, + latency=297.553 * 1e-3, + latency_std=2.016 * 1e-3, + peak_mem=14.457 / 1024, + ), + Point( + context_length=2**16, + latency=296.913 * 1e-3, + latency_std=13.901 * 1e-3, + peak_mem=25.003 / 1024, + ), + Point( + context_length=2**17, + latency=274.299 * 1e-3, + latency_std=53.295 * 1e-3, + peak_mem=47.696 / 1024, + ), + ], + } +} + + +def build_color_map(datasets): + methods = [] + for data in datasets.values(): + for m in data.keys(): + if m not in methods: + methods.append(m) + palette = plt.rcParams["axes.prop_cycle"].by_key()["color"] + return {m: palette[i % len(palette)] for i, m in enumerate(methods)} + + +def plot_memory_and_latency(datasets): + """Plot memory and latency vs context length in two subplots.""" + color_map = build_color_map(datasets) + + fig, (ax1) = plt.subplots(1, 1, figsize=(9, 6)) + + # Plot 1: Memory vs Context Length + for dataset_name, methods in datasets.items(): + for method_name, points in methods.items(): + # Extract data with valid memory values + context_lengths = [ + p.context_length for p in points if p.peak_mem is not None + ] + memories = [p.peak_mem for p in points if p.peak_mem is not None] + + if context_lengths and memories: + ax1.plot( + context_lengths, + memories, + marker="o", + linewidth=2, + markersize=10, + alpha=1.0, + markeredgecolor="black", + markeredgewidth=1, + markerfacecolor=None, + label=method_name, + color=color_map[method_name], + ) + + ax1.set_xlabel("Haystack Length (tokens)", fontweight="bold", fontsize=20) + ax1.set_ylabel("Additional Memory (GB)", fontweight="bold", fontsize=20) + ax1.set_title( + "Additional Memory Needed for Inference", fontweight="bold", fontsize=24 + ) + ax1.set_xscale("log", base=2) + # ax1.set_yscale("log") + ax1.grid(True, alpha=0.3) + # ax1.legend() + + # # Plot 2: Latency vs Context Length + # for dataset_name, methods in datasets.items(): + # for method_name, points in methods.items(): + # # Extract data with valid latency values + # context_lengths = [ + # p.context_length for p in points if p.latency is not None + # ] + # latencies = [p.latency for p in points if p.latency is not None] + # latency_stds = [ + # p.latency_std if p.latency_std is not None else 0 + # for p in points + # if p.latency is not None + # ] + + # if context_lengths and latencies: + # ax2.errorbar( + # context_lengths, + # latencies, + # yerr=latency_stds, + # marker="o", + # linewidth=2, + # markersize=6, + # capsize=5, + # label=method_name, + # color=color_map[method_name], + # ) + + # ax2.set_xlabel("Context Length", fontweight="bold") + # ax2.set_ylabel("Latency (seconds)", fontweight="bold") + # ax2.set_title("Latency vs Context Length", fontweight="bold", fontsize=16) + # ax2.set_xscale("log", base=2) + # ax2.grid(True, alpha=0.3) + # ax2.legend() + + plt.tight_layout() + return fig + + +# Generate the plots +if __name__ == "__main__": + fig = plot_memory_and_latency(datasets) + plt.show() + + # Optionally save the figure + fig.savefig( + "niah_memory_latency_vs_context_length.pdf", dpi=300, bbox_inches="tight" + ) diff --git a/tmp/plot_perf_ctx_latency_mem.py b/tmp/plot_perf_ctx_latency_mem.py new file mode 100644 index 0000000..6a26059 --- /dev/null +++ b/tmp/plot_perf_ctx_latency_mem.py @@ -0,0 +1,786 @@ +from dataclasses import dataclass # added + +import matplotlib as mpl +import matplotlib.pyplot as plt +import numpy as np + +# Styling (reused) +latte_style = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle" +plt.style.use(["ggplot", latte_style]) +plt.rcParams["axes.prop_cycle"] = mpl.cycler( + color=[ + "#FFB83D", + "#555555", + "#9BC750", + "#35775A", + "#5571EA", + "#6CACFF", + # "#E84494", + # "#E47FB0", + # "#785EF0", + # "#FF924E", + # "#5BC5DB", + "#AD6F50", + "#A1A9AD", + ] +) +plt.rcParams["axes.facecolor"] = "white" +plt.rcParams["figure.facecolor"] = "white" +plt.rcParams["savefig.facecolor"] = "white" + +# plt.rcParams["font.family"] = "Ubuntu" +# plt.rcParams["font.size"] = 14 +# plt.rcParams["font.weight"] = "bold" +plt.rcParams["axes.labelweight"] = "bold" +plt.rcParams["axes.facecolor"] = "F7FBFC" +plt.rcParams["figure.facecolor"] = "white" +plt.rcParams["savefig.facecolor"] = "white" +plt.rcParams["grid.color"] = "cccccc" +plt.rcParams["grid.linewidth"] = 1 +plt.rcParams["axes.edgecolor"] = "ccd0da" +plt.rcParams["legend.facecolor"] = "white" +plt.rcParams["legend.fontsize"] = 16 +plt.rcParams["xtick.labelsize"] = 13 +plt.rcParams["ytick.labelsize"] = 13 + + +@dataclass +class Point: + performance: float + context_ratio: float | None = None + latency: float | None = None + latency_std: float | None = None + peak_mem: float | None = None # added + peak_mem_std: float | None = None # added + + +def combine_std_for_sum(std_list): + return np.sqrt(sum(std**2 for std in std_list)) + + +# Unified data structure (merged context + latency via Point) +datasets = { + "SQuAD": { + "Base model w/ context": [ + Point(performance=0.8692, context_ratio=1.0), + ], + "Base model w/o context": [ + Point( + performance=0.1866, + context_ratio=0.0, + latency=0.0, + latency_std=0.0, + peak_mem=0.0, + ), + ], + "CD (oracle)": [ + Point( + performance=0.859, + context_ratio=0.0, + latency=39.974 + 456.989 * 1e-3, + latency_std=combine_std_for_sum([965.326 * 1e-3, 259.130 * 1e-3]), + peak_mem=855.762 / 1024, + ), + ], + "CD (generated queries)": [ + Point( + performance=0.5986, + context_ratio=0.0, + latency=21.254 * 4 + 49.358 + 7.450, + latency_std=10, + peak_mem=41.098, + ), + ], + "Ours (batched)": [ + Point( + # performance=0.7174, + performance=0.7077, + # context_ratio=0.0, + latency=85.954 * 1e-3 + 259.920 * 1e-6, + latency_std=combine_std_for_sum([31.057 * 1e-3, 740.286 * 1e-6]), + peak_mem=1.151, + ), + ], + "Ours (iterative)": [ + Point( + performance=0.715, + context_ratio=0.0, + latency=426.849 * 1e-3 + 260.073 * 1e-6, + latency_std=combine_std_for_sum([43.423 * 1e-3, 29.673 * 1e-6]), + peak_mem=328.134 / 1024, + ), + ], + "T2L": [ + Point( + performance=0.1761, + context_ratio=0.0, + latency=35.480 * 1e-3, + latency_std=22.193 * 1e-3, + peak_mem=1.793, + ), + ], + "LLMLingua-2": [ + Point(performance=0.8492, context_ratio=0.9), + Point(performance=0.8278, context_ratio=0.8), + Point(performance=0.7749, context_ratio=0.6), + Point(performance=0.6624, context_ratio=0.4), + Point(performance=0.4271, context_ratio=0.2), + Point(performance=0.3192, context_ratio=0.1), + ], + }, + "DROP": { + "Base model w/ context": [ + Point(performance=0.4541, context_ratio=1.0), + ], + "Base model w/o context": [ + Point( + performance=0.1417, + context_ratio=0.0, + latency=0.0, + latency_std=0.0, + peak_mem=0.0, + ), + ], + "CD (oracle)": [ + Point( + performance=0.443, + context_ratio=0.0, + latency=39.974 + 422.820 * 1e-3, + latency_std=combine_std_for_sum([965.326 * 1e-3, 241.419 * 1e-3]), + peak_mem=1.381, + ), + ], + "CD (generated queries)": [ + Point( + performance=0.2289, + context_ratio=0.0, + latency=23.883 * 4 + 9.755 + 41.301, + latency_std=14.8852904910, + peak_mem=44.305, + ), + ], + "Ours (batched)": [ + Point( + performance=0.2972, + # context_ratio=0.0, + latency=84.132 * 1e-3 + 251.876 * 1e-6, + latency_std=combine_std_for_sum([30.503 * 1e-3, 22.829 * 1e-6]), + peak_mem=1.976, + ), + ], + "Ours (iterative)": [ + Point( + # performance=0.3002, + performance=0.3387, + context_ratio=0.0, + latency=419.462 * 1e-3 + 254.899 * 1e-6, + latency_std=combine_std_for_sum([40.211 * 1e-3, 29.784 * 1e-6]), + peak_mem=429.881 / 1024, + ), + ], + "T2L": [ + Point( + performance=0.1468, + context_ratio=0.0, + latency=35.036 * 1e-3, + latency_std=18.144 * 1e-3, + peak_mem=3.004, + ), + ], + "LLMLingua-2": [ # + Point(performance=0.4513, context_ratio=0.9), + Point(performance=0.4488, context_ratio=0.8), + Point(performance=0.4076, context_ratio=0.6), + Point(performance=0.3488, context_ratio=0.4), + Point(performance=0.2617, context_ratio=0.2), + Point(performance=0.1712, context_ratio=0.1), + ], + }, + "ROPES": { + "Base model w/ context": [ + Point(performance=0.7457, context_ratio=1.0), + ], + "Base model w/o context": [ + Point( + performance=0.4545, + context_ratio=0.0, + latency=0.0, + latency_std=0.0, + peak_mem=0.0, + ), + ], + "CD (oracle)": [ + Point( + performance=0.7349, + context_ratio=0.0, + latency=39.633 + 285.467 * 1e-3, + latency_std=combine_std_for_sum([1.752, 72.791 * 1e-3]), + peak_mem=527.312 / 1024, + ), + ], + "CD (generated queries)": [ + Point( + performance=0.5786, + context_ratio=0.0, + latency=23.714 * 4 + 11.473 + 59.660, + latency_std=11.7329040736, + peak_mem=43.271, + ), + ], + "Ours (batched)": [ + Point( + performance=0.6759, + # context_ratio=0.0, + latency=85.652 * 1e-3 + 254.029 * 1e-6, + latency_std=combine_std_for_sum([30.852 * 1e-3, 51.127 * 1e-6]), + peak_mem=496.898 / 1024, + ), + ], + "Ours (iterative)": [ + Point( + # performance=0.6786, + performance=0.6088, + context_ratio=0.0, + latency=417.860 * 1e-3 + 257.422 * 1e-6, + latency_std=combine_std_for_sum([45.514 * 1e-3, 31.419 * 1e-6]), + peak_mem=273.867 / 1024, + ), + ], + "T2L": [ + Point( + performance=0.4661, + context_ratio=0.0, + latency=37.913 * 1e-3, + latency_std=17.313 * 1e-3, + peak_mem=697.217 / 1024, + ), + ], + "LLMLingua-2": [ + Point(performance=0.708, context_ratio=0.9), + Point(performance=0.7126, context_ratio=0.8), + Point(performance=0.7291, context_ratio=0.6), + Point(performance=0.677, context_ratio=0.4), + Point(performance=0.6188, context_ratio=0.2), + Point(performance=0.4894, context_ratio=0.1), + ], + }, +} + + +def build_color_map(datasets): + methods = [] + for data in datasets.values(): + for m in data.keys(): + if m not in methods: + methods.append(m) + palette = plt.rcParams["axes.prop_cycle"].by_key()["color"] + return {m: palette[i % len(palette)] for i, m in enumerate(methods)} + + +def plot_ctx_and_latency(datasets): + n = len(datasets) + # detect memory usage presence + memory_present = any( + p.peak_mem is not None + for methods in datasets.values() + for plist in methods.values() + for p in plist + ) + num_cols = 3 if memory_present else 2 + # Share x-axes across rows only when there are multiple datasets (n > 1) + if n > 1: + fig, axes = plt.subplots( + n, + num_cols, + figsize=(6 * num_cols, 5 * n), + sharey=True, + sharex="col", + ) + else: + fig, axes = plt.subplots( + n, num_cols, figsize=(6 * num_cols, 5 * n), sharey=True + ) + if n == 1: + axes = np.array([axes]) + + color_map = build_color_map(datasets) + global_methods = set() + + label_text = "Efficient and Effective\nInternalization" + ours_methods = {"Ours (batched)", "Ours (iterative)"} # added + annot_name = { + "Ours (batched)": "Ours (b)", + "Ours (iterative)": "Ours (i)", + } + + # Keep track of axes groups per row for later separator lines (always store as list) + row_axes_groups = [] + for row_idx, (dataset_name, methods) in enumerate(datasets.items()): + row_axes = axes[row_idx] + if memory_present: + ctx_ax, lat_ax, mem_ax = row_axes + else: + ctx_ax, lat_ax = row_axes + # Ensure each row group is a plain Python list of Axes objects + if isinstance(row_axes, (list, tuple)): + row_axes_groups.append(list(row_axes)) + elif isinstance(row_axes, np.ndarray): + row_axes_groups.append(list(row_axes.ravel())) + else: # single axis fallback (should not occur with current layout) + row_axes_groups.append([row_axes]) + # Row subtitle (dataset name) centered across the row using the middle axis + title_ax = lat_ax # use latency axis (middle) for centering + title_ax.set_title(dataset_name, fontweight="bold", fontsize=24, pad=20) + + # Reference performance from ICL + ref_perf = methods["Base model w/ context"][0].performance + base_wo_ctx_perf_rel = ( + next( + p.performance + for p in methods["Base model w/o context"] + if p.context_ratio is not None + ) + / ref_perf + ) + + # --- Context subplot --- + ours_ctx_present = [ # added + m + for m in ours_methods + if any(p.context_ratio is not None for p in methods.get(m, [])) + ] + for method, points in methods.items(): + ctx_points = [p for p in points if p.context_ratio is not None] + if not ctx_points: + continue + xs = [p.context_ratio for p in ctx_points] + ys = [p.performance / ref_perf for p in ctx_points] + marker = ( + "D" + if method == "CD (oracle)" + else ("*" if method in {"Ours (batched)", "Ours (iterative)"} else "o") + ) + markersize = 17 if "Ours" in method else 12 + + if len(xs) == 1: + ctx_ax.scatter( + xs, + ys, + marker=marker, + s=markersize**2, # scatter uses area, so square the size + label=method, + color=color_map[method], + edgecolors="black", + linewidths=1.5, + alpha=0.9, + ) + else: + ctx_ax.plot( + xs, + ys, + marker=marker, + linewidth=2, + markersize=markersize, + label=method, + color=color_map[method], + markeredgecolor="black", + markeredgewidth=1.5, + alpha=0.9, + ) + if method in {"Ours (batched)", "Ours (iterative)"}: + yoff = ( + (10 if method == "Ours (batched)" else -10) + if len(ours_ctx_present) == 2 + else 6 + ) + + ctx_ax.annotate( # fixed to ctx_ax + annot_name[method], + (xs[-1], ys[-1]), + xytext=(8, yoff), + textcoords="offset points", + fontsize=12, + color=color_map[method], + fontweight="bold", + bbox=None, + ) + global_methods.add(method) + + # Add x-axis label on every context subplot (reverted per-user preference) + ctx_ax.set_xlabel("Context Length Ratio", fontweight="bold", fontsize=16) + # if row_idx == 0: + ctx_ax.set_ylabel( + "Normalized Performance", + fontweight="bold", + fontsize=16, + ) + ctx_ax.grid(True, alpha=0.3) + ctx_ax.set_xlim(-0.05, 1.05) + ctx_ax.axhline(1.0, color="gray", linestyle="--", alpha=0.5, zorder=-1) + ctx_ax.axhline( + base_wo_ctx_perf_rel, color="gray", linestyle="--", alpha=0.5, zorder=-1 + ) + # Removed highlight rectangle and label for context subplot + # --- Latency subplot --- + ours_lat_present = [ # added + m + for m in ours_methods + if any(p.latency is not None for p in methods.get(m, [])) + ] + all_latencies = [] + for method, points in methods.items(): + lat_points = [p for p in points if p.latency is not None] + if not lat_points: + continue + lat_x = [p.latency for p in lat_points] + perf_rel = [p.performance / ref_perf for p in lat_points] + xerr = [p.latency_std for p in lat_points] + marker = ( + "D" + if method == "CD (oracle)" + else ("*" if method in {"Ours (batched)", "Ours (iterative)"} else "o") + ) + markersize = 17 if "Ours" in method else 12 + + if len(lat_x) == 1 and all(err == 0 for err in xerr): + lat_ax.scatter( + lat_x, + perf_rel, + marker=marker, + s=markersize**2, + label=method, + color=color_map[method], + edgecolors="black", + linewidths=1.5, + alpha=0.9, + ) + else: + lat_ax.errorbar( + lat_x, + perf_rel, + xerr=xerr, + fmt=marker + "-", + linewidth=2, + markersize=markersize, + capsize=4, + capthick=1, + label=method, + color=color_map[method], + markeredgecolor="black", + markeredgewidth=1.5, + alpha=0.9, + ) + # Add floating annotation for Ours + if method in {"Ours (batched)", "Ours (iterative)"}: + yoff = ( + (-20 if method == "Ours (batched)" else 10) + if len(ours_lat_present) == 2 + else 6 + ) + lat_ax.annotate( + annot_name[method], + (lat_x[-1], perf_rel[-1]), + xytext=(-25, yoff), + textcoords="offset points", + fontsize=12, + color=color_map[method], + fontweight="bold", + bbox=None, + ) + + all_latencies.extend(lat_x) + global_methods.add(method) + lat_ax.set_xscale("log") + if all_latencies: + positive = [x for x in all_latencies if x > 0] + if positive: + min_pos = min(positive) + max_x = max(all_latencies) + lat_ax.set_xlim(min_pos * 0.8, max_x * 1.2) + lat_ax.set_xlabel("Update Latency (seconds)", fontweight="bold", fontsize=16) + lat_ax.grid(True, which="both", alpha=0.3) + lat_ax.axhline(1.0, color="gray", linestyle="--", alpha=0.5, zorder=-1) + lat_ax.axhline( + base_wo_ctx_perf_rel, color="gray", linestyle="--", alpha=0.5, zorder=-1 + ) + # Add vertical line at 1 second + lat_ax.axvline(1.0, color="gray", linestyle="--", alpha=0.7, zorder=-1) + lat_ax.text( + 1.1, + 0.5, + "Sub-second internalization", + transform=lat_ax.get_xaxis_transform(), + ha="left", + va="center", + fontsize=8, + color="gray", + fontstyle="italic", + rotation=90, + ) + # # Highlight rectangle + # mpl.patches.Rectangle( + # (0.0, 0.5), + # 0.5, + # 0.5, + # transform=lat_ax.transAxes, + # facecolor="green", + # edgecolor="none", + # alpha=0.15, + # zorder=-2, + # ) + # ) + # lat_ax.text( + # 0.02, + # 0.62, + # label_text, + # transform=lat_ax.transAxes, + # fontsize=8.5, + # ha="left", + # va="bottom", + # color="green", + # fontweight="bold", + # ) + + # --- Memory subplot (new) --- + if memory_present: + ours_mem_present = [ # added + m + for m in ours_methods + if any(p.peak_mem is not None for p in methods.get(m, [])) + ] + all_mems = [] + for method, points in methods.items(): + mem_points = [p for p in points if p.peak_mem is not None] + if not mem_points: + continue + mem_x = [p.peak_mem for p in mem_points] + perf_rel = [p.performance / ref_perf for p in mem_points] + xerr = [ + p.peak_mem_std if p.peak_mem_std is not None else 0 + for p in mem_points + ] + marker = ( + "D" + if method == "CD (oracle)" + else ( + "*" if method in {"Ours (batched)", "Ours (iterative)"} else "o" + ) + ) + markersize = 17 if "Ours" in method else 12 + + if len(mem_x) == 1 and all(err == 0 for err in xerr): + mem_ax.scatter( + mem_x, + perf_rel, + marker=marker, + s=markersize**2, + label=method, + color=color_map[method], + edgecolors="black", + linewidths=1.5, + alpha=0.9, + ) + else: + mem_ax.errorbar( + mem_x, + perf_rel, + xerr=xerr, + fmt=marker + "-", + linewidth=2, + markersize=markersize, + capsize=4, + capthick=1, + label=method, + color=color_map[method], + markeredgecolor="black", + markeredgewidth=1.5, + alpha=0.9, + ) + # Add floating annotation for Ours + if method in {"Ours (batched)", "Ours (iterative)"}: + yoff = ( + (-20 if method == "Ours (batched)" else 15) + if len(ours_mem_present) == 2 + else 6 + ) + mem_ax.annotate( + annot_name[method], + (mem_x[-1], perf_rel[-1]), + xytext=(-15, yoff), + textcoords="offset points", + fontsize=12, + color=color_map[method], + fontweight="bold", + bbox=None, + ) + + all_mems.extend(mem_x) + if all_mems: + max_mem = max(all_mems) + min_mem = min(all_mems) + span = ( + max_mem - min_mem + if max_mem > min_mem + else max_mem + if max_mem != 0 + else 1 + ) + mem_ax.set_xlim(min_mem - 0.05 * span, max_mem * 1.05) + mem_ax.set_xlabel( + "Additional Memory Needed\nfor Model Updates (GB)", + fontweight="bold", + fontsize=16, + ) + mem_ax.grid(True, alpha=0.3) + mem_ax.axhline(1.0, color="gray", linestyle="--", alpha=0.5, zorder=-1) + mem_ax.axhline( + base_wo_ctx_perf_rel, color="gray", linestyle="--", alpha=0.5, zorder=-1 + ) + # # Highlight rectangle (memory) fixed 20% width in axes coords + # mem_ax.add_patch( + # mpl.patches.Rectangle( + # (0.0, 0.6), + # 0.4, + # 0.4, + # transform=mem_ax.transAxes, + # facecolor="green", + # edgecolor="none", + # alpha=0.15, + # zorder=-2, + # ) + # ) + # mem_ax.text( + # 0.02, + # 0.62, + # label_text, + # transform=mem_ax.transAxes, + # fontsize=8.5, + # ha="left", + # va="bottom", + # color="green", + # fontweight="bold", + # ) + + # REMOVE old single legend block + # (deleted fig.legend(...) that was on the right) + + # Build grouped legends + in_context = [ + m for m in ["Base model w/ context", "LLMLingua-2"] if m in global_methods + ] + in_param_order = [ + "CD (oracle)", + "CD (generated queries)", + "Ours (iterative)", + "Ours (batched)", + "T2L", + "Base model w/o context", + ] + in_param = [m for m in in_param_order if m in global_methods] + + def make_handles(names): + return [ + mpl.lines.Line2D( + [], + [], + label=m, + marker=( + "D" + if m == "CD (oracle)" + else ("*" if m in {"Ours (batched)", "Ours (iterative)"} else "o") + ), + linestyle="-", + linewidth=2, + color=plt.rcParams["axes.prop_cycle"].by_key()["color"][ + list(build_color_map(datasets).keys()).index(m) + % len(plt.rcParams["axes.prop_cycle"].by_key()["color"]) + ], + markeredgecolor="black", + markeredgewidth=1.5, + markersize=20 if "Ours" in m else 10, + ) + for m in names + ] + + ic_handles = make_handles(in_context) + ip_handles = make_handles(in_param) + + # Extra bottom margin for legends + # (reduced wspace & hspace to tighten gaps) + # Increased hspace to add more vertical separation between dataset rows + plt.subplots_adjust(right=0.88, wspace=0.12, hspace=0.8, bottom=0.25) + + # Re-enable x tick labels for all rows (matplotlib hides them for shared axes not on bottom) + if n > 1: + for ax in axes.reshape(-1): + ax.tick_params(labelbottom=True) + + # Set legend title fontweight before creating legends + plt.rcParams["legend.title_fontsize"] = 14 + + fig.legend( + ic_handles, + [h.get_label() for h in ic_handles], + title="In-context knowledge", + loc="upper center", + bbox_to_anchor=(0.25, 0.1), + ncol=len(ic_handles), + frameon=True, + fancybox=True, + fontsize=12, + title_fontproperties={"weight": "bold", "size": 14}, + ) + fig.legend( + ip_handles, + [h.get_label() for h in ip_handles], + title="In-parameter knowledge", + loc="upper center", + bbox_to_anchor=(0.6, 0.1), + ncol=min(len(ip_handles), 3), + frameon=True, + fancybox=True, + fontsize=12, + title_fontproperties={"weight": "bold", "size": 14}, + ) + + # Removed global suptitle referencing only the last dataset; per-row titles now provided. + + # Add horizontal separator lines between dataset rows (if more than one dataset) + if len(row_axes_groups) > 1: + # Use figure coordinate system + for upper_group, lower_group in zip(row_axes_groups[:-1], row_axes_groups[1:]): + # Reference axes (first axis in each group) + upper_ax = upper_group[0] + lower_ax = lower_group[0] + upper_bottom = upper_ax.get_position().y0 + lower_top = lower_ax.get_position().y1 + line_y = (upper_bottom + lower_top) / 2.0 + combined = upper_group + lower_group + x0 = min(ax.get_position().x0 for ax in combined) + x1 = max(ax.get_position().x1 for ax in combined) + fig.add_artist( + mpl.lines.Line2D( + [x0, x1], + [line_y, line_y], + transform=fig.transFigure, + color="#cccccc", + linewidth=1.2, + alpha=0.8, + ) + ) + return fig + + +if __name__ == "__main__": + fig = plot_ctx_and_latency(datasets) + plt.show() + fig.savefig( + "/home/tan/research/ctx-to-lora/tmp/performance_ctx_latency_mem.png", + dpi=300, + bbox_inches="tight", + ) + fig.savefig( + "/home/tan/research/ctx-to-lora/tmp/performance_ctx_latency_mem.pdf", + dpi=300, + bbox_inches="tight", + ) diff --git a/tmp/plot_perf_gen_mem.py b/tmp/plot_perf_gen_mem.py new file mode 100644 index 0000000..9a83f6d --- /dev/null +++ b/tmp/plot_perf_gen_mem.py @@ -0,0 +1,453 @@ +from dataclasses import dataclass + +import matplotlib as mpl +import matplotlib.pyplot as plt +import numpy as np + +# Styling (mirroring the other plot for visual consistency) +latte_style = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle" +plt.style.use(["ggplot", latte_style]) +plt.rcParams["axes.prop_cycle"] = mpl.cycler( + color=[ + "#FFB83D", + "#EEE684", + "#555555", + "#9BC750", + "#35775A", + "#5571EA", + "#6CACFF", + # "#E84494", + # "#E47FB0", + # "#785EF0", + # "#FF924E", + # "#5BC5DB", + "#AD6F50", + "#A1A9AD", + ] +) +plt.rcParams["axes.facecolor"] = "white" +plt.rcParams["figure.facecolor"] = "white" +plt.rcParams["savefig.facecolor"] = "white" + +# plt.rcParams["font.family"] = "Ubuntu" +# plt.rcParams["font.size"] = 14 +# plt.rcParams["font.weight"] = "bold" +plt.rcParams["axes.labelweight"] = "bold" +plt.rcParams["axes.facecolor"] = "F7FBFC" +plt.rcParams["figure.facecolor"] = "white" +plt.rcParams["savefig.facecolor"] = "white" +plt.rcParams["grid.color"] = "cccccc" +plt.rcParams["grid.linewidth"] = 1 +plt.rcParams["axes.edgecolor"] = "ccd0da" +plt.rcParams["legend.facecolor"] = "white" +# plt.rcParams["legend.fontsize"] = 20 +plt.rcParams["xtick.labelsize"] = 20 +plt.rcParams["ytick.labelsize"] = 20 + + +@dataclass +class Point: + performance: float + peak_mem: float | None = None # in GB + + +# Dataset focused only on memory vs performance (relative performance computed later) +datasets = { + "2WikiMultihopQA": { + # We keep Base model w/ context only as a performance reference (no point plotted if memory absent) + "Base model w/ context": [ + Point(performance=0.242, peak_mem=2.099 * 1024), + ], + "Base model w/ truncated context": [ + Point(performance=0.3387, peak_mem=1.194 * 1024), + ], + "Base model w/o context": [ + Point(performance=0.2138, peak_mem=65.162), + ], + "CD (oracle)": [ + Point( + performance=0.305, + peak_mem=18.196, + ), + ], + "CD (5 generated queries)": [ + Point( + performance=0.2385, + peak_mem=32.501, + ), + ], + "Ours": [ + Point( + performance=0.2901, + peak_mem=35.777, + ), + ], + "Ours + truncated context": [ + Point( + performance=0.3585, + peak_mem=1.367 * 1024, + ) + ], + # "Ours (8k chunks)": [ + # Point( + # performance=0.2901, + # peak_mem=35.777, + # ), + # ], + "T2L": [ + Point( + performance=0.2474, + peak_mem=37.776, + ), + ], + "LLMLingua-2": [ + # Point(performance=0.3141, peak_mem=1.111 * 1024), + Point(performance=0.3169, peak_mem=1.018 * 1024), + Point(performance=0.3558, peak_mem=807.366), + Point(performance=0.3591, peak_mem=592.584), + Point(performance=0.2767, peak_mem=305.298), + Point(performance=0.2343, peak_mem=158.760), + ], + }, + "MultiFieldQA": { + "Base model w/ context": [ + Point(performance=0.303, peak_mem=1.211 * 1024), + ], + "Base model w/ truncated context": [ + Point(performance=0.3938, peak_mem=1.084 * 1024), + ], + "Base model w/o context": [ + Point(performance=0.0948, peak_mem=65.853), + ], + "CD (oracle)": [ + Point( + performance=0.4101, + peak_mem=7.836, + ), + ], + "CD (5 generated queries)": [ + Point( + performance=0.165, + peak_mem=31.587, + ), + ], + "Ours": [ + Point( + performance=0.1895, + peak_mem=35.769, + ), + ], + "Ours + truncated context": [ + Point( + performance=0.4293, + peak_mem=1.233 * 1024, + ), + ], + "T2L": [ + Point( + performance=0.083, + peak_mem=38.907, + ), + ], + "LLMLingua-2": [ + Point(performance=0.4222, peak_mem=909.663), + Point(performance=0.3898, peak_mem=716.981), + Point(performance=0.3773, peak_mem=490.478), + Point(performance=0.2885, peak_mem=256.169), + Point(performance=0.2298, peak_mem=146.394), + ], + }, + "QASPER": { + # We keep Base model w/ context only as a performance reference (no point plotted if memory absent) + "Base model w/ context": [ + Point(performance=0.346, peak_mem=2.217 * 1024), + ], + "Base model w/ truncated context": [ + Point(performance=0.3839, peak_mem=788.697), + ], + "Base model w/o context": [ + Point(performance=0.0859, peak_mem=64.989), + ], + "CD (oracle)": [ + Point( + performance=0.3627, + peak_mem=19.611, + ), + ], + "CD (5 generated queries)": [ + Point( + performance=0.1712, + peak_mem=36.767, + ), + ], + "Ours": [ + Point( + performance=0.2051, + peak_mem=46.130, + ), + ], + "Ours + truncated context": [ + Point( + performance=0.3751, + peak_mem=873.099, + ), + ], + "T2L": [ + Point( + performance=0.1259, + peak_mem=54.176, + ), + ], + "LLMLingua-2": [ + # Point(performance=0.3734, peak_mem=710.176), + Point(performance=0.3849, peak_mem=647.385), + Point(performance=0.3838, peak_mem=487.781), + Point(performance=0.3328, peak_mem=386.447), + Point(performance=0.3329, peak_mem=219.062), + Point(performance=0.2614, peak_mem=132.576), + ], + }, +} + + +def build_color_map(datasets): + # Collect all unique methods in a deterministic order + methods = [] + for data in datasets.values(): + for m in data.keys(): + if m not in methods: + methods.append(m) + + # Sort methods to ensure consistent ordering across runs + # methods = sorted(methods) + + palette = plt.rcParams["axes.prop_cycle"].by_key()["color"] + return {m: palette[i % len(palette)] for i, m in enumerate(methods)} + + +def plot_memory_vs_performance(datasets): + n = len(datasets) + # If multiple datasets, arrange horizontally; otherwise single standard figure + if n == 1: + fig, single_ax = plt.subplots(1, 1, figsize=(7, 5)) + axes = [single_ax] + else: + fig, axes_arr = plt.subplots(1, n, figsize=(7 * n, 5)) + # Ensure axes is a simple list + axes = list(np.atleast_1d(axes_arr)) + + color_map = build_color_map(datasets) + global_methods = set() + ours_methods = {"Ours", "Ours + truncated context"} # added + annot_name = { # added + "Ours": "Ours", + "Ours + truncated context": "Ours+ctx", + } + + highlight_label = "Efficient and Effective\nInternalization" + + for idx, (dataset_name, methods) in enumerate(datasets.items()): + ax = axes[idx] + ref_perf = methods["Base model w/ truncated context"][0].performance + base_wo_ctx_perf_rel = ( + methods["Base model w/o context"][0].performance / ref_perf + if "Base model w/o context" in methods + else None + ) + + all_mems = [] + all_rel_perfs = [] + ours_present = [ # added: detect both ours points for staggering + m + for m in ours_methods + if m in methods and any(p.peak_mem is not None for p in methods[m]) + ] + for method, points in methods.items(): + mem_points = [p for p in points if p.peak_mem is not None] + if not mem_points: + continue # skip reference method if no memory value + xs = [p.peak_mem for p in mem_points] + ys = [p.performance / ref_perf for p in mem_points] + # marker mapping (oracle=diamond, ours=star, others=circle) + marker = ( + "D" + if method == "CD (oracle)" + else ("*" if method in ours_methods else "o") + ) + size = 20 if "Ours" in method else 12 + ax.plot( + xs, + ys, + linestyle="-", + marker=marker, + linewidth=2, + markersize=size, + label=method, + color=color_map[method], + markeredgecolor="black", + markeredgewidth=1.5, + alpha=0.9, + ) + # Add annotation for Ours, stagger if both are present + if method in ours_methods: + yoff = ( + (10 if method == "Ours" else -10) if len(ours_present) == 2 else 6 + ) + place_left = method == "Ours + truncated context" and dataset_name in { + "2WikiMultihopQA", + "MultiFieldQA", + } + if place_left: + yoff = 10 # move annotation above the marker for these two plots + xoff = -5 if place_left else 8 + ha = "right" if place_left else "left" + ax.annotate( + annot_name[method], + (xs[-1], ys[-1]), + xytext=(xoff, yoff), + textcoords="offset points", + fontsize=12, + color=color_map[method], + fontweight="bold", + ha=ha, + bbox=None, + ) + + all_mems.extend(xs) + all_rel_perfs.extend(ys) + global_methods.add(method) + + if all_mems: + max_mem = max(all_mems) + min_mem = min(all_mems) + span = ( + (max_mem - min_mem) + if max_mem > min_mem + else (max_mem if max_mem != 0 else 1) + ) + ax.set_xlim(5, max_mem * 1.25) + + # Set individual y-axis limits for each subplot + if all_rel_perfs: + min_y = min(all_rel_perfs) + max_y = max(all_rel_perfs) + y_range = max_y - min_y + margin = y_range * 0.1 if y_range > 0 else 0.1 + ax.set_ylim(min_y - margin, max_y + margin) + + ax.set_xscale("log") + ax.set_xlabel( + "Additional Memory Needed\nfor Generation (MB)", + fontweight="bold", + fontsize=22, + ) + if idx == 0: # Only show y-label on first plot + ax.set_ylabel("Normalized Performance", fontweight="bold", fontsize=22) + ax.grid(True, alpha=0.3) + ax.axhline(1.0, color="gray", linestyle="--", alpha=0.5, zorder=-1) + if base_wo_ctx_perf_rel is not None: + ax.axhline( + base_wo_ctx_perf_rel, color="gray", linestyle="--", alpha=0.5, zorder=-1 + ) + + ax.set_title(dataset_name, fontweight="bold", fontsize=30) + + # Build grouped legends + in_context = [ + m + for m in [ + "Base model w/ context", + "Base model w/ truncated context", + "LLMLingua-2", + ] + if m in global_methods + ] + in_param_order = [ + "CD (oracle)", + "CD (5 generated queries)", + "Ours", + "Ours + truncated context", + "T2L", + "Base model w/o context", + ] + in_param = [m for m in in_param_order if m in global_methods] + + def make_handles(names): + return [ + mpl.lines.Line2D( + [], + [], + label=m, + marker=( + "D" + if m == "CD (oracle)" + else ("*" if m in {"Ours", "Ours + truncated context"} else "o") + ), + linestyle="-", + linewidth=2, + color=color_map[m], + markeredgecolor="black", + markeredgewidth=1.5, + markersize=22 if "Ours" in m else 14, + ) + for m in names + ] + + ic_handles = make_handles(in_context) + ip_handles = make_handles(in_param) + + # Extra bottom margin for legends + # (reduced wspace & hspace to tighten gaps) + # Increased hspace to add more vertical separation between dataset rows + plt.subplots_adjust(right=0.88, wspace=0.12, hspace=1.0, bottom=0.25) + + # Set legend title fontweight before creating legends + # plt.rcParams["legend.title_fontsize"] = 14 + + fig.legend( + ic_handles, + [h.get_label() for h in ic_handles], + title="In-context knowledge", + loc="upper center", + bbox_to_anchor=(0.2, -0.05), + ncol=2, + frameon=True, + fancybox=True, + fontsize=18, + title_fontproperties={"weight": "bold", "size": 20}, + ) + fig.legend( + ip_handles, + [h.get_label() for h in ip_handles], + title="In-parameter knowledge", + loc="upper center", + bbox_to_anchor=(0.7, -0.05), + ncol=min(len(ip_handles), 3), + frameon=True, + fancybox=True, + fontsize=18, + title_fontproperties={"weight": "bold", "size": 20}, + ) + + # Adjust spacing: horizontal layout uses wspace, single layout unaffected + if n == 1: + plt.subplots_adjust(bottom=0.18, left=0.12, right=0.98) + else: + # Tighter horizontal spacing and margins + plt.subplots_adjust(bottom=0.18, left=0.06, right=0.99, wspace=0.15) + return fig + + +if __name__ == "__main__": + fig = plot_memory_vs_performance(datasets) + plt.show() + fig.savefig( + "/home/tan/research/ctx-to-lora/tmp/performance_gen_mem.png", + dpi=300, + bbox_inches="tight", + ) + fig.savefig( + "/home/tan/research/ctx-to-lora/tmp/performance_gen_mem.pdf", + dpi=300, + bbox_inches="tight", + ) diff --git a/tmp/plot_perf_vs_ctx_len.py b/tmp/plot_perf_vs_ctx_len.py new file mode 100644 index 0000000..6edb4c0 --- /dev/null +++ b/tmp/plot_perf_vs_ctx_len.py @@ -0,0 +1,272 @@ +import matplotlib as mpl +import matplotlib.pyplot as plt +import numpy as np + +latte_style = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle" +plt.style.use(["ggplot", latte_style]) +plt.rcParams["axes.prop_cycle"] = mpl.cycler( + color=[ + "#648FFF", + "#E84494", + "#9BC750", + # "#DC4BDC", + "#229487", + "#785EF0", + "#FF924E", + "#5BC5DB", + "#F0434F", + "#AD6F50", + # "#1982C4", + # "#6A4C93", + # "#56CCF2", + # "#2F52E0", + # "#FF006E", + # "#FF8A00", + # "#FFD300", + # "#00A6ED", + # "#FF7A5A", + # "#FFAA92", + # "#FFB7B2", + # "#FFDAC1", + # "#E2F0CB", + # "#B5EAD7", + # "#C7CEEA", + # "#D4A5A5", + # "#E9B0DF", + # "#F8D3DA", + # "#FBE0C3", + # "#FFF0B8", + # "#E6E2AF", + # "#A2D2FF", + # "#B4F8C8", + # "#FFD6E0", + # "#FFEFB5", + ] +) +plt.rcParams["axes.facecolor"] = "white" +plt.rcParams["figure.facecolor"] = "white" +plt.rcParams["savefig.facecolor"] = "white" + +# plt.rcParams["font.family"] = "Ubuntu" +# plt.rcParams["font.size"] = 14 +# plt.rcParams["font.weight"] = "bold" +plt.rcParams["axes.labelweight"] = "bold" +plt.rcParams["axes.facecolor"] = "F7FBFC" +plt.rcParams["figure.facecolor"] = "white" +plt.rcParams["savefig.facecolor"] = "white" +plt.rcParams["grid.color"] = "cccccc" +plt.rcParams["grid.linewidth"] = 1 +plt.rcParams["axes.edgecolor"] = "ccd0da" +plt.rcParams["legend.facecolor"] = "white" + +# Reproducibility +rng = np.random.default_rng(42) + +# X-axis: context length as a fraction of full ICL +x_inparam = np.array([0.0]) # in-parameter updates +x_icl = np.array([1.0]) # full ICL + +# Define mock performances (F1) per dataset, consistent with the paper narrative: +# - ICL highest +# - Compression approaches ICL as ratio increases + +datasets = { + "SQuAD": { + "Base model w/ ICL": [(1.0, 0.8692)], + "Base model w/o context": [(0.0, 0.1866)], + "CD (oracle)": [(0.0, 0.859)], + "CD (generated queries)": [(0.0, 0.5986)], + "Ours": [(0.0, 0.7174)], + "T2L": [(0.0, 0.1761)], + "LLMLingua-2": [ + (0.9, 0.8492), + (0.8, 0.8278), + (0.6, 0.7749), + (0.4, 0.6624), + (0.2, 0.4271), + (0.1, 0.3192), + ], + }, + # "DROP": { + # "Base model w/ ICL": [(1.0, 0.4541)], + # "Base model w/o context": [(0.0, 0.1417)], + # "CD (oracle)": [(0.0, 0.443)], + # "CD (generated queries)": [(0.0, X)], + # "Ours": [(0.0, 0.7174)], + # "T2L": [(0.0, 0.1761)], + # "LLMLingua-2": [ + # (0.9, 0.8492), + # (0.8, 0.8278), + # (0.6, 0.7749), + # (0.4, 0.6624), + # (0.2, 0.4271), + # (0.1, 0.3192), + # ], + # }, +} + + +def plot_performance_comparison(datasets): + """Plot performance of methods relative to Base model w/ ICL reference.""" + n_datasets = len(datasets) + fig, axes = plt.subplots(1, n_datasets, figsize=(8 * n_datasets, 6)) + + # Handle single dataset case + if n_datasets == 1: + axes = [axes] + + global_methods = set() # collect methods for global legends + + for idx, (dataset_name, methods) in enumerate(datasets.items()): + ax = axes[idx] + + # Get reference performance from "Base model w/ ICL" + ref_performance = methods["Base model w/ ICL"][0][1] + + # Get base model w/o context performance for reference line + no_context_performance = ( + methods["Base model w/o context"][0][1] / ref_performance + ) + + # Define method groups + in_context_methods = {"Base model w/ ICL", "LLMLingua-2"} + in_param_methods = { + "Base model w/o context", + "Ours", + "T2L", + "CD (oracle)", + "CD (generated queries)", + } + + # Plot each method + for method_name, data_points in methods.items(): + x_vals = [point[0] for point in data_points] + y_vals = [ + point[1] / ref_performance for point in data_points + ] # Relative to reference + + ax.plot( + x_vals, + y_vals, + label=method_name, + linewidth=2, + markersize=8, + marker="o", + alpha=1, + markeredgecolor="black", + markeredgewidth=1, + markerfacecolor=None, + ) + global_methods.add(method_name) + + # Formatting + ax.set_xlabel("Context Length Ratio vs ICL", fontweight="bold") + ax.set_ylabel("Relative Performance vs ICL)", fontweight="bold") + ax.grid(True, alpha=0.3) + ax.set_xlim(-0.05, 1.05) + + # Add horizontal line at y=1 (reference performance) + ax.axhline( + y=1.0, + color="gray", + linestyle="--", + alpha=0.5, + zorder=-1, + ) + + # Add horizontal line for base model w/o context + ax.axhline( + y=no_context_performance, + color="red", + linestyle=":", + alpha=0.7, + zorder=-1, + ) + + # Add highlight rectangle: fixed 20% of axis width (axes coords), y from 0.8 to 1.0 + ax.add_patch( + mpl.patches.Rectangle( + (0.0, 0.8), + 0.2, + 0.2, + transform=ax.transAxes, + facecolor="green", + edgecolor="none", + alpha=0.15, + zorder=-2, + ) + ) + + # Global legends (after all subplots) + in_context = [ + m for m in ["Base model w/ ICL", "LLMLingua-2"] if m in global_methods + ] + in_param_order = [ + "CD (oracle)", + "CD (generated queries)", + "Ours", + "T2L", + "Base model w/o context", + ] + in_param = [m for m in in_param_order if m in global_methods] + + def make_handles(names): + return [ + mpl.lines.Line2D( + [], + [], + label=m, + marker="o", + linestyle="-", + linewidth=2, + markeredgecolor="black", + markeredgewidth=1, + markersize=8, + ) + for m in names + ] + + ic_handles = make_handles(in_context) + ip_handles = make_handles(in_param) + + # Adjust bottom for legends + plt.subplots_adjust(right=0.75, bottom=0.22) + + # Place legends below figure, centered in two groups + fig.legend( + ic_handles, + [h.get_label() for h in ic_handles], + title="In-context knowledge", + loc="upper center", + bbox_to_anchor=(0.25, -0.05), + ncol=len(ic_handles), + frameon=True, + fancybox=True, + fontsize=10, + title_fontsize=11, + ) + fig.legend( + ip_handles, + [h.get_label() for h in ip_handles], + title="In-parameter knowledge", + loc="upper center", + bbox_to_anchor=(0.72, -0.05), + ncol=len(ip_handles), + frameon=True, + fancybox=True, + fontsize=10, + title_fontsize=11, + ) + return fig + + +# Create the plot +fig = plot_performance_comparison(datasets) +plt.show() + +# Optionally save the figure +plt.savefig( + "/home/tan/research/ctx-to-lora/tmp/performance_comparison.png", + dpi=300, + bbox_inches="tight", +) diff --git a/tmp/plot_perf_vs_update_latency.py b/tmp/plot_perf_vs_update_latency.py new file mode 100644 index 0000000..ed6340e --- /dev/null +++ b/tmp/plot_perf_vs_update_latency.py @@ -0,0 +1,204 @@ +import matplotlib as mpl +import matplotlib.pyplot as plt +import numpy as np + +latte_style = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle" +plt.style.use(["ggplot", latte_style]) +plt.rcParams["axes.prop_cycle"] = mpl.cycler( + color=[ + "#648FFF", + "#E84494", + "#9BC750", + "#229487", + "#785EF0", + "#FF924E", + "#5BC5DB", + "#F0434F", + "#AD6F50", + ] +) +plt.rcParams["axes.facecolor"] = "white" +plt.rcParams["figure.facecolor"] = "white" +plt.rcParams["savefig.facecolor"] = "white" + +# plt.rcParams["font.family"] = "Ubuntu" +# plt.rcParams["font.size"] = 14 +# plt.rcParams["font.weight"] = "bold" +plt.rcParams["axes.labelweight"] = "bold" +plt.rcParams["axes.facecolor"] = "F7FBFC" +plt.rcParams["figure.facecolor"] = "white" +plt.rcParams["savefig.facecolor"] = "white" +plt.rcParams["grid.color"] = "cccccc" +plt.rcParams["grid.linewidth"] = 1 +plt.rcParams["axes.edgecolor"] = "ccd0da" +plt.rcParams["legend.facecolor"] = "white" +plt.rcParams["legend.fontsize"] = 16 +plt.rcParams["xtick.labelsize"] = 13 +plt.rcParams["ytick.labelsize"] = 13 + + +# Reproducibility +rng = np.random.default_rng(42) + +# Define mock performances vs update latency (in seconds) +# X-axis: update latency in seconds (log scale) +# Y-axis: F1 performance + + +def combine_std_for_sum(std_list): + """ + Compute combined standard deviation when adding independent random variables. + + For independent random variables X1, X2, ..., Xn: + Var(X1 + X2 + ... + Xn) = Var(X1) + Var(X2) + ... + Var(Xn) + Therefore: std(sum) = sqrt(std1² + std2² + ... + stdn²) + + Args: + std_list: List of standard deviations + + Returns: + Combined standard deviation for the sum + """ + return np.sqrt(sum(std**2 for std in std_list)) + + +datasets = { + "SQuAD": { + "Base model w/o context": [(0.0, 0.0, 0.1866)], # No update needed + "CD (oracle)": [ + ( + 39.974 + 456.989 * 1e-3, + combine_std_for_sum([965.326 * 1e-3, 259.130 * 1e-3]), + 0.859, + ) + ], + "CD (generated queries)": [(21.254 * 4 + 49.358 + 7.450, 10, 0.5986)], + "Ours": [(85.954 * 1e-3, 31.057 * 1e-3, 0.7174)], + "T2L": [(35.480 * 1e-3, 22.193 * 1e-3, 0.1761)], + }, + # "DROP": { + # "Base model w/o context": [(0.0, 0.1417)], + # "CD (oracle)": [(300, 0.443)], + # "CD (generated queries)": [(600, 0.3986)], + # "Ours": [(120, 0.4174)], + # "T2L": [(180, 0.1561)], + # }, +} + + +def plot_performance_comparison(datasets): + """Plot performance of methods vs update latency.""" + n_datasets = len(datasets) + fig, axes = plt.subplots(1, n_datasets, figsize=(8 * n_datasets, 6)) + + # Handle single dataset case + if n_datasets == 1: + axes = [axes] + + for idx, (dataset_name, methods) in enumerate(datasets.items()): + ax = axes[idx] + + # Use a fixed reference performance for ICL (since we removed it from the plot) + ref_performance = 0.8692 if dataset_name == "SQuAD" else 0.4541 + + # Get base model w/o context performance for reference line + no_context_performance = ( + methods["Base model w/o context"][0][2] / ref_performance + ) + + # All methods are now in-parameter methods + in_param_methods = { + "Base model w/o context", + "Ours", + "T2L", + "CD (oracle)", + "CD (generated queries)", + } + + # Plot each method + for method_name, data_points in methods.items(): + x_vals = [point[0] for point in data_points] + y_vals = [ + point[2] / ref_performance for point in data_points + ] # Relative to reference (third element) + std_vals = [ + point[1] / ref_performance for point in data_points + ] # Standard deviation (second element) + + ax.errorbar( + x_vals, + y_vals, + xerr=std_vals, + label=method_name, + linewidth=2, + markersize=8, + marker="o", + alpha=1, + markeredgecolor="black", + markeredgewidth=1, + markerfacecolor=None, + capsize=4, + capthick=1, + ) + + # Calculate max x value for dynamic xlim + max_x = max( + point[0] for data_points in methods.values() for point in data_points + ) + + # Formatting + ax.set_xlabel("Update Latency (seconds)", fontweight="bold") + ax.set_ylabel("Relative Performance vs ICL", fontweight="bold") + ax.set_title(f"{dataset_name}", fontweight="bold", fontsize=14) + ax.grid(True, alpha=0.3) + ax.set_xlim(-5, max_x * 1.2) + + # Add horizontal line at y=1 (reference performance) + ax.axhline( + y=1.0, + color="gray", + linestyle="--", + alpha=0.5, + zorder=-1, + ) + + # Add horizontal line for base model w/o context + ax.axhline( + y=no_context_performance, + color="red", + linestyle="--", + alpha=0.7, + zorder=-1, + ) + + # Create single legend for in-parameter methods + handles, labels = ax.get_legend_handles_labels() + + ax.legend( + handles, + labels, + title="In-parameter knowledge", + bbox_to_anchor=(1.05, 0.6), + loc="upper left", + ncol=1, + fontsize=10, + title_fontsize=11, + frameon=True, + fancybox=True, + ).get_title().set_fontweight("bold") + + plt.tight_layout() + plt.subplots_adjust(right=0.75) # Make room for legend on the right + return fig + + +# Create the plot +fig = plot_performance_comparison(datasets) +plt.show() + +# Optionally save the figure +plt.savefig( + "/home/tan/research/ctx-to-lora/tmp/performance_vs_latency.png", + dpi=300, + bbox_inches="tight", +) diff --git a/tmp/test_ce.py b/tmp/test_ce.py new file mode 100644 index 0000000..3758e2c --- /dev/null +++ b/tmp/test_ce.py @@ -0,0 +1,26 @@ +#!/usr/bin/env python +""" +Test cross entropy loss shape with reduction=None +""" + +import torch +from torch import nn + +# Create sample data +batch_size, seq_len, vocab_size = 2, 5, 1000 +logits = torch.randn(batch_size, seq_len, vocab_size) +labels = torch.randint(0, vocab_size, (batch_size, seq_len)) + +# Flatten for cross entropy (expects 2D input) +logits_flat = logits.view(-1, vocab_size) # (batch_size * seq_len, vocab_size) +labels_flat = labels.view(-1) # (batch_size * seq_len,) + +print(f"Input shapes:") +print(f" logits_flat: {logits_flat.shape}") +print(f" labels_flat: {labels_flat.shape}") + +loss = nn.functional.cross_entropy(logits_flat, labels_flat, reduction="none") + +print(f"Output shape:") +print(f" loss: {loss.shape}") +print(f" loss values (first 10): {loss[:10]}") diff --git a/tmp/test_distillation_sparse_vs_dense.py b/tmp/test_distillation_sparse_vs_dense.py new file mode 100644 index 0000000..7c48ae3 --- /dev/null +++ b/tmp/test_distillation_sparse_vs_dense.py @@ -0,0 +1,305 @@ +#!/usr/bin/env python +""" +Compare the current dense KL (teacher top-k scatter into full vocab) vs a sparse gather-based +implementation used in distillation. + +Metrics: +- Exact numerical equality (max |diff|) across many random seeds. +- Runtime (per iteration, averaged). +- Peak memory (GPU if available, else best-effort RSS delta on CPU). + +By default we mimic the current code semantics (NO renormalization of the provided teacher +log-probs). That is, we treat absent vocab items as probability 0 and do: + loss_i = - sum_j exp(target_logp[i,j]) * log q(indices[i,j]) +which matches the existing dense scatter approach that sets -inf everywhere else. + +We also optionally show (flag --normalized) a variant where teacher log-probs are re-normalized +over the provided top-k indices (useful if you want a true cross-entropy on the partial dist). + +Usage (from repo root): + python tmp/test_distillation_sparse_vs_dense.py --seeds 100 --n 2048 --vocab 32000 --k 8 --device cuda + +Adjust sizes to fit your GPU memory. The dense method allocates an (N, V) float tensor twice +(outputs_logits + teacher_logp). For large (N,V) this can OOM. +""" + +from __future__ import annotations + +import argparse +import math +import os +import statistics +import time +import tracemalloc +from dataclasses import dataclass + +import torch + +try: + import psutil # optional for CPU RSS +except ImportError: # pragma: no cover + psutil = None + + +def dense_loss( + outputs_logits: torch.Tensor, indices: torch.Tensor, target_logp: torch.Tensor +) -> torch.Tensor: + """Replicates current implementation: build full teacher_logp filled with -inf, scatter top-k logp. + outputs_logits: (N, V) + indices: (N, K) long + target_logp: (N, K) float (NOT assumed normalized) + Returns: (N,) per-position losses. + """ + N, V = outputs_logits.shape + # Create full teacher log-probs + teacher_logp = torch.full_like(outputs_logits, -torch.inf) + teacher_logp.scatter_(1, indices, target_logp) + p = teacher_logp.exp() # zeros for -inf + # logq = torch.log_softmax(outputs_logits, dim=-1) + logq_full_denom = torch.logsumexp(outputs_logits, dim=-1, keepdim=True) # (N,1) + logq = outputs_logits - logq_full_denom + loss = -(p * logq).sum(dim=-1) # (N,) + return loss + + +def sparse_loss( + outputs_logits: torch.Tensor, + indices: torch.Tensor, + target_logp: torch.Tensor, + *, + normalize=False, +) -> torch.Tensor: + """Sparse equivalent to dense_loss without constructing a full (N,V) teacher matrix. + If normalize=True we renormalize teacher log-probs over provided indices. + Returns (N,) per-position losses with SAME semantics as dense_loss when normalize=False. + """ + # Gather student log-probs only at teacher indices. + logq_full_denom = torch.logsumexp(outputs_logits, dim=-1, keepdim=True) # (N,1) + selected_logits = outputs_logits.gather(1, indices) # (N,K) + logq_selected = selected_logits - logq_full_denom # log softmax at selected indices + + if normalize: + # True cross-entropy over provided subset. + teacher_logp_norm = target_logp - torch.logsumexp( + target_logp, dim=-1, keepdim=True + ) + teacher_p = teacher_logp_norm.exp() + else: + teacher_p = target_logp.exp() + + # NOT EFFICIENT? + # logq = torch.nn.functional.log_softmax(outputs_logits, dim=-1) + # logq_selected = logq.gather(1, indices) # (N,K) + # teacher_p = target_logp.exp() + + loss = -(teacher_p * logq_selected).sum(dim=-1) + return loss + + +@dataclass +class Metrics: + time_ms: float + peak_mem_mb: float | None + max_abs_diff: float | None = None + + +def measure( + method_fn, *args, repeat=1, device="cpu", measure_memory=True, **kwargs +) -> tuple[torch.Tensor, Metrics]: + """Run method_fn (returns tensor), measure avg wall time and peak memory. + For GPU memory we use torch.cuda.reset_peak_memory_stats / max_memory_allocated. + For CPU fallback we sample RSS before/after (approx).""" + is_cuda = device.startswith("cuda") and torch.cuda.is_available() + + if is_cuda: + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + else: + if measure_memory and psutil is not None: + process = psutil.Process(os.getpid()) + rss_before = process.memory_info().rss + elif measure_memory: + tracemalloc.start() + + start = time.perf_counter() + out = None + for _ in range(repeat): + out = method_fn(*args, **kwargs) + if is_cuda: + torch.cuda.synchronize() + elapsed = (time.perf_counter() - start) / repeat * 1000.0 + + peak_mb = None + if is_cuda: + peak_mb = torch.cuda.max_memory_allocated(device=device) / (1024**2) + else: + if measure_memory and psutil is not None: + rss_after = process.memory_info().rss + peak_mb = (rss_after - rss_before) / (1024**2) + elif measure_memory: + current, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + peak_mb = peak / (1024**2) + + return out, Metrics(time_ms=elapsed, peak_mem_mb=peak_mb) + + +def run_experiment(args): + device = args.device + torch.manual_seed(0) + + # Pre-generate shapes once per seed to reduce variance from dynamic graph allocations. + dense_times = [] + dense_mems = [] + sparse_times = [] + sparse_mems = [] + diffs = [] + allclose_fail_seeds = [] + + N = args.n + V = args.vocab + K = args.k + + if K > V: + raise ValueError("k cannot exceed vocab size") + + dtype = torch.float16 if args.fp16 else torch.float32 + + for seed in range(args.seeds): + torch.manual_seed(seed) + # Student logits (simulate model outputs for label positions) + outputs_logits = 1 + 10 * torch.randn(N, V, device=device, dtype=dtype) + # Ensure uniqueness per row: sample without replacement via randperm + indices = torch.stack([torch.randperm(V, device=device)[:K] for _ in range(N)]) + # Teacher log-probs (could be any real numbers). Simulate by sampling logits then subtract logsumexp to get log-probs, + # then optionally add random temperature scaling so they are not normalized (matching current dense semantics). + teacher_logits = 1 + 10 * torch.randn( + N, K, device=device, dtype=torch.float32 + ) # keep higher precision for stability + target_logp = ( + teacher_logits - torch.logsumexp(teacher_logits, dim=-1, keepdim=True) + ).to(dtype) + + # Measure dense + dense_out, dense_metrics = measure( + dense_loss, outputs_logits, indices, target_logp, device=device + ) + # Measure sparse (semantics-matching version) + sparse_out, sparse_metrics = measure( + sparse_loss, + outputs_logits, + indices, + target_logp, + device=device, + normalize=args.normalize_sparse, + ) + + # Compare + # Move to float32 for diff to reduce underflow when fp16 + diff = (dense_out.float() - sparse_out.float()).abs().max().item() + diffs.append(diff) + dense_times.append(dense_metrics.time_ms) + dense_mems.append(dense_metrics.peak_mem_mb or math.nan) + sparse_times.append(sparse_metrics.time_ms) + sparse_mems.append(sparse_metrics.peak_mem_mb or math.nan) + + # Per-seed torch.allclose check when semantics should match (no normalization flags) + if not args.normalize_sparse: + if not torch.allclose( + dense_out, sparse_out, rtol=args.rtol, atol=args.atol + ): + allclose_fail_seeds.append(seed) + + print("=== Distillation Dense vs Sparse Comparison ===") + print(f"Device: {device} dtype: {dtype} seeds: {args.seeds}") + print(f"Shape: N={N} positions, V={V} vocab, K={K} top-k per position") + print(f"Sparse renormalize flag (normalize_sparse): {args.normalize_sparse}") + print("-- Timing (ms per iteration) --") + print( + f"Dense mean: {statistics.mean(dense_times):.3f} median: {statistics.median(dense_times):.3f}" + ) + print( + f"Sparse mean: {statistics.mean(sparse_times):.3f} median: {statistics.median(sparse_times):.3f}" + ) + speedup = statistics.mean(dense_times) / statistics.mean(sparse_times) + print(f"Speedup (dense / sparse): {speedup:.2f}x") + + if all(not math.isnan(m) for m in dense_mems + sparse_mems): + print("-- Peak Memory (MB) --") + print( + f"Dense mean: {statistics.mean(dense_mems):.1f} median: {statistics.median(dense_mems):.1f}" + ) + print( + f"Sparse mean: {statistics.mean(sparse_mems):.1f} median: {statistics.median(sparse_mems):.1f}" + ) + mem_reduction = statistics.mean(dense_mems) / max( + 1e-9, statistics.mean(sparse_mems) + ) + print(f"Memory reduction (dense / sparse): {mem_reduction:.2f}x") + else: + print("(Memory stats not fully available on this platform / configuration)") + + print("-- Numerical Differences --") + max_diff = max(diffs) + mean_diff = statistics.mean(diffs) + print(f"Max |dense - sparse| over seeds: {max_diff:.3e}") + print(f"Mean |dense - sparse|: {mean_diff:.3e}") + + tol = 1e-5 + if args.normalize_sparse: + print( + "NOTE: Normalization flags enabled; numerical results may intentionally differ from dense baseline." + ) + else: + # Summary of allclose results + if not allclose_fail_seeds: + print( + f"SUCCESS: Sparse matches dense within tolerance (max diff {max_diff:.3e} < {tol}) and allclose passed for all seeds (rtol={args.rtol}, atol={args.atol})." + ) + else: + if allclose_fail_seeds: + print( + f"WARNING: torch.allclose failed for {len(allclose_fail_seeds)} / {args.seeds} seeds. Failed seeds: {allclose_fail_seeds[:10]}{'...' if len(allclose_fail_seeds) > 10 else ''}" + ) + # if max_diff >= tol: + # print( + # f"WARNING: Max diff {max_diff:.3e} exceeds heuristic dtype tolerance {tol}." + # ) + + +def parse_args(): + p = argparse.ArgumentParser() + p.add_argument( + "--seeds", type=int, default=100, help="Number of random seeds / trials" + ) + p.add_argument("--n", type=int, default=1024, help="Number of label positions (N)") + p.add_argument("--vocab", type=int, default=32000, help="Vocabulary size (V)") + p.add_argument("--k", type=int, default=8, help="Top-k teacher tokens (K)") + p.add_argument( + "--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu" + ) + p.add_argument( + "--fp16", action="store_true", help="Use float16 for student logits/log-probs" + ) + p.add_argument( + "--normalize-sparse", + dest="normalize_sparse", + action="store_true", + help="Renormalize teacher log-probs inside sparse path", + ) + return p.parse_args() + + +if __name__ == "__main__": + torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False + torch.backends.cudnn.benchmark = False + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + + args = parse_args() + + # Set default tolerances if not provided + args.rtol = 1.3e-6 + args.atol = 1e-5 + run_experiment(args) diff --git a/tmp/test_element_wise_mult_einsum.py b/tmp/test_element_wise_mult_einsum.py new file mode 100644 index 0000000..bfd8da1 --- /dev/null +++ b/tmp/test_element_wise_mult_einsum.py @@ -0,0 +1,33 @@ +import torch + +# Setup from your example +x = torch.rand(16, 26, 8, 2048) +a = torch.rand(1, 26, 8, 1) + +# Original operation for comparison +y = x * a + +# --- Einsum Solutions --- + +# 1. Explicit form (best for clarity) +# We assign letters to each dimension and specify the exact output format. +y_einsum_explicit = torch.einsum("ijkl,ijkl->ijkl", x, a) + +# 2. Implicit form (more concise) +# For element-wise products, you can omit the output part. +# Einsum infers the output should contain all unique indices. +y_einsum_implicit = torch.einsum("ijkl,ijkl", x, a) + + +# --- Verification --- +print("Original shape:", y.shape) +print("Einsum shape:", y_einsum_explicit.shape) + +# Check if the results are numerically the same +print("\nAre the results the same?") +print("Explicit form match:", torch.allclose(y, y_einsum_explicit)) +print("Implicit form match:", torch.allclose(y, y_einsum_implicit)) + +print(f"{y}=") +print(f"{y_einsum_explicit=}") +print(f"{y_einsum_implicit=}") diff --git a/tmp/test_llm_lingua.py b/tmp/test_llm_lingua.py new file mode 100644 index 0000000..787a03d --- /dev/null +++ b/tmp/test_llm_lingua.py @@ -0,0 +1,19 @@ +from llmlingua import PromptCompressor + +with open("src/ctx_to_lora/modeling/test_ctx.txt") as f: + prompt = f.read() + +llm_lingua = PromptCompressor( + model_name="microsoft/llmlingua-2-xlm-roberta-large-meetingbank", + use_llmlingua2=True, # Whether to use llmlingua-2 +) +compressed_prompt = llm_lingua.compress_prompt( + prompt, rate=0.25, force_tokens=["\n", "?"] +) +print(compressed_prompt) + +## Or use LLMLingua-2-small model +# llm_lingua = PromptCompressor( +# model_name="microsoft/llmlingua-2-bert-base-multilingual-cased-meetingbank", +# use_llmlingua2=True, # Whether to use llmlingua-2 +# ) diff --git a/tmp/test_logits_processor_vllm.py b/tmp/test_logits_processor_vllm.py new file mode 100644 index 0000000..6b3e2e8 --- /dev/null +++ b/tmp/test_logits_processor_vllm.py @@ -0,0 +1,39 @@ +import os + +import torch +from vllm import LLM, SamplingParams + +os.environ["VLLM_USE_V1"] = "0" # vLLM V1 does not support logits processors. + + +class LogitsSpy: + def __init__(self): + self.processed_logits: list[torch.Tensor] = [] + + def __call__(self, token_ids: list[int], logits: torch.Tensor): + self.processed_logits.append(logits) + return logits + + +if __name__ == "__main__": + llm = LLM(model="google/gemma-2-2b-it", trust_remote_code=True) + prompts = [ + "What is the capital of France?", + "Explain the theory of relativity in simple terms.", + "Write a short story about a robot learning to love.", + ] + logits_spy = LogitsSpy() + sampling_params = SamplingParams( + temperature=0.7, top_p=0.95, max_tokens=5, logits_processors=[logits_spy] + ) + outputs = llm.generate(prompts, sampling_params) + logits = logits_spy.processed_logits + print(f"Logits: {logits}") + print(len(logits)) + print(logits[0].shape) + print(logits[1].shape) + print(logits[2].shape) + print(logits[3].shape) + print(logits[4].shape) + + breakpoint() diff --git a/tmp/test_needle_haystack_plot.py b/tmp/test_needle_haystack_plot.py new file mode 100644 index 0000000..4c81a50 --- /dev/null +++ b/tmp/test_needle_haystack_plot.py @@ -0,0 +1,209 @@ +#!/usr/bin/env python3 +""" +Test script to create a line plot showing needle-in-a-haystack task accuracy vs input length. +""" + +import matplotlib as mpl +import matplotlib.pyplot as plt + +latte_style = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle" +plt.style.use(["ggplot", latte_style]) +plt.rcParams["axes.prop_cycle"] = mpl.cycler( + color=[ + "#FFB83D", + "#5571EA", + ] +) +plt.rcParams["axes.facecolor"] = "white" +plt.rcParams["figure.facecolor"] = "white" +plt.rcParams["savefig.facecolor"] = "white" + +# plt.rcParams["font.family"] = "Ubuntu" +# plt.rcParams["font.size"] = 14 +# plt.rcParams["font.weight"] = "bold" +plt.rcParams["axes.labelweight"] = "bold" +plt.rcParams["axes.facecolor"] = "F7FBFC" +plt.rcParams["figure.facecolor"] = "white" +plt.rcParams["savefig.facecolor"] = "white" +plt.rcParams["grid.color"] = "cccccc" +plt.rcParams["grid.linewidth"] = 1 +plt.rcParams["axes.edgecolor"] = "ccd0da" +plt.rcParams["legend.facecolor"] = "white" +plt.rcParams["legend.fontsize"] = 16 +plt.rcParams["xtick.labelsize"] = 16 +plt.rcParams["ytick.labelsize"] = 16 + + +def main(): + # Token bins from the user's specification + tok_bins = [1024 * (i + 1) for i in range(16)] + tok_bins += [2**14 + 2**12 * (i + 1) for i in range(4)] + tok_bins += [2**15 + 2**13 * (i + 1) for i in range(12)] + x_values = tok_bins + print(x_values) + + y_values = { + "Base model w/ context": [ + 1.00, + 1.00, + 1.00, + 1.00, + 1.00, + 1.00, + 1.00, + 0.99, + 0.80, + 0.00, + 0.00, + 0.00, + 0.00, + 0.00, + 0.00, + 0.00, + 0.00, + 0.00, + 0.00, + 0.00, + 0.00, + 0.00, + 0.00, + 0.00, + 0.00, + 0.00, + 0.00, + 0.00, + 0.00, + 0.00, + 0.00, + 0.00, + ], + "Ours": [ + 1.00, + 1.00, + 1.00, + 1.00, + 1.00, + 1.00, + 1.00, + 1.00, + 1.00, + 1.00, + 1.00, + 1.00, + 1.00, + 1.00, + 1.00, + 1.00, + 1.00, + 1.00, + 1.00, + 1.00, + 1.00, + 0.95, + 0.87, + 0.80, + 0.65, + 0.53, + 0.43, + 0.37, + 0.28, + 0.24, + 0.17, + 0.15, + ], + } + + # Create the plot + cleaner_x_vals = [ + 2**10, + 2**11, + 2**12, + 2**13, + 2**13 + 2**10, + 2**13 + 2**11, + 2**14, + 2**15, + 2**15 + 2**13, + ] + plt.figure(figsize=(9, 6)) + for name, ys in y_values.items(): + xys = [ + (x, y) + for x, y in zip(x_values, ys) + if x in cleaner_x_vals or x > 2**15 + 2**13 + ] + plt.plot( + [xy[0] for xy in xys], + [xy[1] for xy in xys], + marker="o", + linewidth=2, + markersize=8, + alpha=1.0, + label=name, + markeredgecolor="black", + markeredgewidth=1, + markerfacecolor=None, + ) + + # Customize the plot + plt.xlabel("Haystack Length (tokens)", fontsize=20) + plt.ylabel("ROUGE-L F1 Score", fontsize=20) + plt.title( + "NIAH Retrieval Performance", + fontsize=24, + fontweight="bold", + ) + plt.grid(True, alpha=0.3) + + # Set log scale for x-axis + plt.xscale("log", base=2) + + # Format x-axis labels with K suffix (representing 1024 multiples) + # x_labels = [ + # f"{x // 1024}K" if x >= 1024 else str(x) + # for x in x_values + # if x in cleaner_x_vals or x > 2**15 + 2**13 + # ] + # plt.xticks( + # [x for x in x_values if x in cleaner_x_vals or x > 2**15 + 2**13], + # x_labels, + # rotation=90, + # fontsize=8, + # ) + + # Remove minor ticks for cleaner appearance + plt.gca().xaxis.set_minor_locator(plt.NullLocator()) + + # Add some styling + plt.ylim(-0.05, 1.05) + plt.xlim(min(x_values) * 0.8, max(x_values) * 1.2) + + # Add horizontal lines for reference + + # Add legend + plt.legend(loc="lower left") + + # Tight layout for better appearance + plt.tight_layout() + + # Save the plot + plt.savefig( + "/home/tan/research/ctx-to-lora/tmp/needle_haystack_accuracy_plot.pdf", + format="pdf", + dpi=300, + bbox_inches="tight", + ) + + # Show statistics + print(f"Generated plot with {len(x_values)} data points") + print(f"Input length range: {min(x_values):.0f} - {max(x_values):.0f} tokens") + print( + f"Plot saved to: /home/tan/research/ctx-to-lora/tmp/needle_haystack_accuracy_plot.pdf" + ) + + # Display the plot + plt.show() + + +if __name__ == "__main__": + main() diff --git a/tmp/test_per_ctx_loss.py b/tmp/test_per_ctx_loss.py new file mode 100644 index 0000000..578ba0a --- /dev/null +++ b/tmp/test_per_ctx_loss.py @@ -0,0 +1,243 @@ +"""Test / diagnostic script for comparing current per_ctx_loss implementation +against a proposed fixed version that detects label spans via mask transitions. + +Run: + python tmp/test_per_ctx_loss.py + +It will: + * Construct synthetic batches with known label span structure. + * Call the existing per_ctx_loss (imported from ctx_to_lora.trainer). + * Call a fixed implementation (fixed_per_ctx_loss) using label_mask transitions. + * Show detected spans, per-query losses, per-context losses, and scaling ratios. + * Highlight cases where the current implementation fails (empty spans / mismatch). + +Note: The current per_ctx_loss assumes the provided `loss` vector is ordered +exactly as the concatenation of all label-token losses (no unlabeled tokens), +matching DistillationTrainer's usage, but NOT CrossEntropyTrainer's raw CE output. +""" + +from __future__ import annotations + +import math + +# Ensure "src" is on path so we can import project modules when running from project root +import sys +import traceback +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path + +import torch + +ROOT = Path(__file__).resolve().parents[1] +SRC = ROOT / "src" +if str(SRC) not in sys.path: + sys.path.insert(0, str(SRC)) + +from ctx_to_lora.trainer import per_ctx_loss # type: ignore + + +@dataclass +class SyntheticCase: + name: str + label_mask: list[bool] + n_queries_per_ctx: list[int] + description: str + + +def per_ctx_loss(inputs, labels, loss): + n_queries_per_ctx = inputs["n_queries"].tolist() + + position_ids = inputs["position_ids"].squeeze(0) + # account only label positions + label_mask = labels.squeeze(0) != -100 + label_pos_ids = label_mask * position_ids + label_pos_ids_diff = label_pos_ids.diff( + append=torch.tensor([0], device=position_ids.device) + ) + # assumes the input starts with non-assistant tokens + start_label_pos = torch.where((label_pos_ids_diff > 0) * ~label_mask)[0] + end_label_pos = torch.where((label_pos_ids_diff < 0) * label_mask)[0] + + label_seq_lens = end_label_pos - start_label_pos + cu_label_seq_lens = torch.cumsum(label_seq_lens, dim=0) + start_indices = torch.cat( + ( + torch.tensor([0], device=cu_label_seq_lens.device), + cu_label_seq_lens[:-1], + ) + ) + + # these stack and split can be optimized but let's keep it simple + # mean across tokens of each q + qa_losses = torch.stack( + [loss[start:end].mean() for start, end in zip(start_indices, cu_label_seq_lens)] + ) + + # mean across queries of each ctx + per_ctx_losses = [ql.mean() for ql in torch.split(qa_losses, n_queries_per_ctx)] + + # per-ctx loss + loss = torch.stack(per_ctx_losses) + return loss + + +def fixed_per_ctx_loss(inputs, labels, loss_vec: torch.Tensor) -> torch.Tensor: + """Reference implementation using label_mask transitions. + + Args: + inputs: dict with keys: n_queries, position_ids + labels: [1, L] + loss_vec: tensor of shape (#label_tokens,) in label-token order. + Returns: + Tensor of shape (#contexts,) containing mean loss per context. + """ + n_queries_per_ctx = inputs["n_queries"].tolist() + label_mask = labels.squeeze(0) != -100 + + lm = label_mask.to(torch.int) + if lm.numel() == 0: + return torch.empty(0) + + if lm.numel() == 1: + spans = [(0, 1)] if lm[0] == 1 else [] + else: + transition = lm[1:] - lm[:-1] + starts = (transition == 1).nonzero(as_tuple=False).flatten() + 1 + ends = (transition == -1).nonzero(as_tuple=False).flatten() + 1 + if lm[0] == 1: + starts = torch.cat([torch.tensor([0], device=starts.device), starts]) + if lm[-1] == 1: + ends = torch.cat([ends, torch.tensor([lm.numel()], device=ends.device)]) + assert starts.numel() == ends.numel() + spans = list(zip(starts.tolist(), ends.tolist())) + + # Extract per-query mean losses (loss_vec is already label-token sequential). + lengths = [e - s for s, e in spans] + assert sum(lengths) == loss_vec.numel(), ( + "Mismatch: total span label tokens != loss vector length." + ) + cu_lengths = torch.tensor(lengths).cumsum(0) + starts_loss = torch.cat([torch.tensor([0]), cu_lengths[:-1]]) + qa_means = [loss_vec[a:b].mean() for a, b in zip(starts_loss, cu_lengths)] + qa_means_t = torch.stack(qa_means) + + # Group into contexts + split_q = torch.split(qa_means_t, n_queries_per_ctx) + per_ctx = torch.stack([ctx.mean() for ctx in split_q]) + return per_ctx + + +def build_labels_from_mask(mask: Sequence[bool]) -> torch.Tensor: + # Assign dummy token id 1 where True, -100 elsewhere. + data = [1 if m else -100 for m in mask] + return torch.tensor([data]) # shape [1, L] + + +def build_position_ids(L: int) -> torch.Tensor: + return torch.arange(L).unsqueeze(0) + + +def make_loss_vector(mask: Sequence[bool]) -> torch.Tensor: + # Provide a smoothly varying loss value per label token to visualize averaging. + # Example: losses increase linearly → weighting differences are obvious. + values = [] + counter = 0 + for m in mask: + if m: + # synthetic harder tokens later + values.append(0.5 + 0.1 * counter) + counter += 1 + return torch.tensor(values) + + +def run_case(case: SyntheticCase): + print(f"\n=== Case: {case.name} ===") + print(case.description) + mask = case.label_mask + labels = build_labels_from_mask(mask) + position_ids = build_position_ids(len(mask)) + loss_vec = make_loss_vector(mask) + inputs = { + "n_queries": torch.tensor(case.n_queries_per_ctx), + "position_ids": position_ids, + } + print(f"Label mask: {''.join('1' if m else '.' for m in mask)}") + print(f"Loss vector (per label token order): {loss_vec.tolist()}") + + # Attempt current implementation + try: + current = per_ctx_loss(inputs, labels, loss_vec) + print( + "Current per_ctx_loss output:", + current.tolist(), + "shape=", + list(current.shape), + ) + except Exception as e: + print("Current per_ctx_loss raised exception:") + traceback.print_exc() + current = None + + # Fixed implementation + try: + fixed = fixed_per_ctx_loss(inputs, labels, loss_vec) + print( + "Fixed per_ctx_loss output: ", fixed.tolist(), "shape=", list(fixed.shape) + ) + except Exception as e: + print("Fixed implementation raised exception:") + traceback.print_exc() + fixed = None + + # Compare magnitudes if both succeeded + if current is not None and fixed is not None and current.numel() == fixed.numel(): + ratio = ( + (current.mean() / fixed.mean()).item() if fixed.mean() != 0 else math.nan + ) + print(f"Mean(current)/Mean(fixed) ratio: {ratio:.4f}") + print("---") + + +def main(): + cases = [ + # SyntheticCase( + # name="Two spans, single context", + # label_mask=[False, False, True, True, True, False, False, True, True], + # n_queries_per_ctx=[2], + # description="Two answer/query spans of lengths 3 and 2 in one context.", + # ), + SyntheticCase( + name="Single span", + label_mask=[False, True, True, True, False], + n_queries_per_ctx=[1], + description="One continuous labeled span.", + ), + # SyntheticCase( + # name="Three contexts (2,1 queries)", + # label_mask=[False, True, True, False, True, False, True, True, True, False], + # n_queries_per_ctx=[1, 1, 1], + # description="Three separate single-query contexts.", + # ), + # SyntheticCase( + # name="Edge case: starts labeled", + # label_mask=[True, True, False, True, False], + # n_queries_per_ctx=[2], + # description="Sequence begins with a labeled span then another later.", + # ), + # SyntheticCase( + # name="Edge case: ends labeled", + # label_mask=[False, True, False, True, True], + # n_queries_per_ctx=[2], + # description="Sequence ends with a labeled span.", + # ), + ] + + for case in cases: + run_case(case) + + print("\nDone. Inspect above for discrepancies.\n") + + +if __name__ == "__main__": + main() diff --git a/tmp/test_perf_vs_latency_plot.py b/tmp/test_perf_vs_latency_plot.py new file mode 100644 index 0000000..20cebd5 --- /dev/null +++ b/tmp/test_perf_vs_latency_plot.py @@ -0,0 +1,89 @@ +import matplotlib.patches as patches +import matplotlib.pyplot as plt +import numpy as np + +# --- Data Points --- +# Each tuple represents (Update Latency, Performance) +data = {"CD": (27, 92), "TTT": (28, 65), "T2L": (4, 25), "P2L": (6, 80)} + +labels = list(data.keys()) +x_coords = [val[0] for val in data.values()] +y_coords = [val[1] for val in data.values()] + +# --- Plotting Setup --- +# Using a modern and clean style for the plot +plt.style.use("seaborn-v0_8-whitegrid") +fig, ax = plt.subplots(figsize=(10, 7), dpi=100) + +# --- Create the Scatter Plot --- +# Using a single color for points and making them large for visibility +colors = plt.cm.viridis(np.linspace(0.3, 0.9, len(labels))) +scatter = ax.scatter( + x_coords, y_coords, s=200, c=colors, edgecolors="black", alpha=0.8, zorder=5 +) + +# --- Annotate Each Point --- +# Add text labels next to each point for clarity +for i, label in enumerate(labels): + # Adjust text position slightly for better aesthetics + ax.annotate( + label, + (x_coords[i], y_coords[i]), + textcoords="offset points", + xytext=(15, 0), # Offset text to the right of the point + ha="left", + fontsize=12, + fontweight="bold", + zorder=6, + ) + +# --- Add the "Efficient Internalization" Zone --- +# Create a transparent rectangle in the top-left corner +efficient_zone = patches.Rectangle( + (0, 70), # (x, y) coordinate of the bottom-left corner + 12, # Width of the rectangle + 30, # Height of the rectangle + edgecolor="green", + facecolor="green", + alpha=0.15, # Transparency level + linewidth=1.5, + linestyle="--", +) +ax.add_patch(efficient_zone) + +# Add text label for the zone +ax.text( + 1, + 85, # (x, y) position for the text + "Efficient Internalization", + fontsize=14, + fontweight="bold", + color="darkgreen", + rotation=0, # Can be rotated if needed, e.g., rotation=90 +) + + +# --- Final Touches & Labels --- +# Set plot title and axis labels +ax.set_title( + "Performance vs. Update Latency Analysis", fontsize=18, fontweight="bold", pad=20 +) +ax.set_xlabel("Update Latency (seconds)", fontsize=14, labelpad=15) +ax.set_ylabel("Performance Score", fontsize=14, labelpad=15) + +# Set axis limits to provide some padding +ax.set_xlim(0, 32) +ax.set_ylim(0, 105) + +# Customize tick parameters for a cleaner look +ax.tick_params(axis="both", which="major", labelsize=12) + +# Add a subtle grid +ax.grid(True, which="both", linestyle="--", linewidth=0.5) + +# Ensure the layout is tight and clean +plt.tight_layout() + +# --- Display the Plot --- +plt.show() +plt.savefig("perf_vs_latency_plot.png", dpi=300) # Save as high-res PNG diff --git a/tmp/viz_gen_lora.py b/tmp/viz_gen_lora.py new file mode 100644 index 0000000..eed6b37 --- /dev/null +++ b/tmp/viz_gen_lora.py @@ -0,0 +1,113 @@ +from collections import defaultdict + +import torch + +# Store activations +activations = defaultdict(list) + + +def apply_forward_hook(module, name): + def hook_fn(m, input_args, output_tensor): + activations[name].append(output_tensor.cpu().detach()) + + module.register_forward_hook(hook_fn) + + +def last_token_pool( + outputs: dict[str, torch.Tensor], attention_mask: torch.Tensor +) -> torch.Tensor: + last_hidden_states = ( + outputs["hidden_states"][-1].detach() + if "hidden_states" in outputs + else outputs["last_hidden_state"].detach() + ) + left_padding = attention_mask[:, -1].sum() == attention_mask.shape[0] + if left_padding: + return last_hidden_states[:, -1] + else: + sequence_lengths = attention_mask.sum(dim=1) - 1 + batch_size = last_hidden_states.shape[0] + return last_hidden_states[ + torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths + ] + + +if __name__ == "__main__": + from functools import partial + + import torch + from peft import get_peft_model + from torch.utils.data import DataLoader + from tqdm import tqdm + + from ctx_to_lora.data.collator import generation_collator + from ctx_to_lora.data.processing import get_tokenized_dataset + from ctx_to_lora.model_loading import get_lora_config, get_model_and_tokenizer + 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) + + 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, + ) + + apply_forward_hook(model.ctx_encoder, "ctx_encoder") + apply_forward_hook(model.layers, "hypernet_layers") + + for sample_idx, sample in tqdm(enumerate(dataloader)): + 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"] + ) diff --git a/tmp/viz_lora_weights.py b/tmp/viz_lora_weights.py new file mode 100644 index 0000000..9ce075f --- /dev/null +++ b/tmp/viz_lora_weights.py @@ -0,0 +1,338 @@ +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']}" + ) diff --git a/tmp/viz_token_dist.py b/tmp/viz_token_dist.py new file mode 100644 index 0000000..e69de29 diff --git a/tmp/vllm_logprobs.py b/tmp/vllm_logprobs.py new file mode 100644 index 0000000..e141dd1 --- /dev/null +++ b/tmp/vllm_logprobs.py @@ -0,0 +1,36 @@ +from datasets import load_dataset + +from ctx_to_lora.data.processing import get_tokenized_dataset +from ctx_to_lora.model_loading import get_tokenizer + +if __name__ == "__main__": + ds = load_dataset( + "parquet", + data_files="./data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact/train/ds.parquet", + split="train", + ) + + # tokenizer = ctx_tokenizer = AutoTokenizer.from_pretrained( + # "google/gemma-2-2b-it", + # ) + tokenizer = ctx_tokenizer = get_tokenizer("google/gemma-2-2b-it", train=True) + + tokenized_ds = get_tokenized_dataset( + "self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact", + split="train", + base_model_max_len=2**13, + tokenizer=tokenizer, + tokenizer_kwargs={}, + ctx_model_max_len=2**13, + ctx_tokenizer=ctx_tokenizer, + ctx_tokenizer_kwargs={}, + max_qas_len=2048, + max_qas_per_sample=1, + add_ctx_to_chat=False, + add_repeat_prompt=False, + add_negative_prompt=False, + use_kl_loss=True, + ) + print(ds[0]) + print(tokenized_ds[0]) + breakpoint()