# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md) # SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project. # # SPDX-License-Identifier: MIT """ Paper-style plots: memory usage scaling with N for different lengths (L). - Rows = config (SoftPrompt, Flamingo) - Cols = sequence lengths (L) [excluding L=1] - Hue = base model - Y-axis sharing logic: * Flamingo: all panels share y-axis * SoftPrompt: all panels have independent y-axes - OOM cases (peak_cuda_reserved_gb > OOM_THRESHOLD) are shown by extending the line upward and marking with a red X + "OOM". """ import pandas as pd import seaborn as sns import matplotlib.pyplot as plt import matplotlib import re from matplotlib.lines import Line2D OOM_THRESHOLD = 180 # GB def parse_model_name(llm_id, model_type): """Return base_model, config (SoftPrompt or Flamingo).""" if llm_id.startswith("meta-llama/"): base_name = llm_id.replace("meta-llama/", "") elif llm_id.startswith("google/"): base_name = llm_id.replace("google/", "") else: base_name = llm_id # Normalize base model names to match expected order if "Llama-3.2-1B" in base_name: base_name = "Llama-3.2-1B" elif "Llama-3.2-3B" in base_name: base_name = "Llama-3.2-3B" elif "gemma-3-270m" in base_name: base_name = "Gemma-3-270M" elif "gemma-3-1b-pt" in base_name: base_name = "Gemma-3-1B-pt" if model_type == "EmbedHealthSP": type_name = "SoftPrompt" elif model_type == "EmbedHealthFlamingo": type_name = "Flamingo" else: type_name = model_type return base_name, type_name def parse_simulation_dataset(name): """Parse Simulation dataset name like 'Simulation-L10-N5' → (L=10, N=5).""" match = re.match(r"Simulation-L(\d+)-N(\d+)", name) if match: return int(match.group(1)), int(match.group(2)) return None, None def plot_memory_usage_paper(csv_file="memory_simulation.csv"): # Publication style plt.style.use("seaborn-v0_8-white") matplotlib.rcParams.update({ "font.family": "serif", "font.serif": ["Palatino", "Times New Roman", "DejaVu Serif"], "font.size": 18, "axes.labelsize": 20, "axes.titlesize": 20, "legend.fontsize": 17, "xtick.labelsize": 17, "ytick.labelsize": 17, "axes.linewidth": 0.6, "axes.edgecolor": "0.15", }) # Load & preprocess df = pd.read_csv(csv_file) df["peak_cuda_reserved_gb"] = df["peak_cuda_reserved_gb"].replace(-1, pd.NA) df = df[df["dataset"].str.startswith("Simulation-")] df[["base_model", "config"]] = df.apply( lambda row: pd.Series(parse_model_name(row["llm_id"], row["model"])), axis=1 ) df[["L", "N"]] = df["dataset"].apply( lambda s: pd.Series(parse_simulation_dataset(s)) ) df = df.dropna(subset=["L", "N"]) df["L"] = df["L"].astype(int) df["N"] = df["N"].astype(int) df = df[df["L"] != 1] df = df.sort_values(by=["base_model", "config", "L", "N"]) # Palette + markers - use consistent colors with the other script # Define consistent order and colors for each model model_order = ["Gemma-3-270M", "Gemma-3-1B-pt", "Llama-3.2-1B", "Llama-3.2-3B"] color_map = { "Gemma-3-270M": "#4477AA", "Gemma-3-1B-pt": "#66CCEE", "Llama-3.2-1B": "#228833", "Llama-3.2-3B": "#CC6677" } # Get base models in the specified order base_models = [bm for bm in model_order if bm in df["base_model"].unique()] custom_palette = [color_map.get(bm, "#888888") for bm in base_models] markers_dict = dict(zip( base_models, ["o", "s", "^", "D", "p", "X", "*"] )) # Unique sequence lengths unique_L = sorted(df["L"].unique()) # Create subplot grid manually: 2 rows (SoftPrompt, Flamingo) fig, axes = plt.subplots( 2, len(unique_L), figsize=(3.2 * len(unique_L), 6), sharex="col", ) # Row mapping row_map = {"SoftPrompt": 0, "Flamingo": 1} # Precompute Flamingo y-lims flamingo_df = df[df["config"] == "Flamingo"] flamingo_ymin, flamingo_ymax = None, None if not flamingo_df.empty: flamingo_ymin = flamingo_df["peak_cuda_reserved_gb"].min(skipna=True) flamingo_ymax = flamingo_df["peak_cuda_reserved_gb"].max(skipna=True) flamingo_ymin = 0 flamingo_ymax = max(flamingo_ymax if flamingo_ymax else 0, OOM_THRESHOLD * 1.1) # Iterate configs for cfg in ["SoftPrompt", "Flamingo"]: cfg_df = df[df["config"] == cfg] for j, L in enumerate(unique_L): ax = axes[row_map[cfg], j] subdf = cfg_df[cfg_df["L"] == L] for bm, sdf in subdf.groupby("base_model"): sdf = sdf.sort_values("N") # Split successful vs OOM runs ok_df = sdf[sdf["peak_cuda_reserved_gb"] <= OOM_THRESHOLD] oom_df = sdf[sdf["peak_cuda_reserved_gb"] > OOM_THRESHOLD] # Normal line if not ok_df.empty: ax.plot( ok_df["N"], ok_df["peak_cuda_reserved_gb"], label=bm, color=custom_palette[base_models.index(bm)], marker=markers_dict[bm], linewidth=2.2, markersize=5, alpha=0.9, ) # OOM handling if not oom_df.empty: first_oom = oom_df.iloc[0] last_ok_y = ok_df["peak_cuda_reserved_gb"].iloc[-1] if not ok_df.empty else OOM_THRESHOLD * 0.9 # extend line upward ax.plot( [ok_df["N"].iloc[-1] if not ok_df.empty else first_oom["N"], first_oom["N"]], [last_ok_y, OOM_THRESHOLD * 1.05], color=custom_palette[base_models.index(bm)], linestyle="--", linewidth=1.5, alpha=0.8, ) # red X marker at OOM ax.scatter( first_oom["N"], OOM_THRESHOLD * 1.05, color="red", marker="x", s=70, linewidth=2, zorder=5, ) ax.text( first_oom["N"], OOM_THRESHOLD * 1.05, "OOM", color="red", fontsize=9, fontweight="bold", ha="center", va="bottom" ) # Titles if row_map[cfg] == 0: ax.set_title(f"L = {L}", fontsize=19, fontweight="bold") # Y labels only leftmost col if j == 0: ax.set_ylabel( f"{cfg}", fontsize=18, fontweight="bold" ) else: ax.set_ylabel("") # Styling ax.set_facecolor("#F8F9FA") ax.grid(True, which="major", linestyle="-", linewidth=0.4, alpha=0.5) ax.grid(True, which="minor", linestyle=":", linewidth=0.3, alpha=0.3) ax.minorticks_on() # Remove individual x-axis labels - will add global one later ax.set_xlabel("") ax.set_xticks([1, 2, 3, 4, 5]) ax.tick_params(axis="x", which="both", labelbottom=True) # Y-axis rules if cfg == "Flamingo": ax.set_ylim(0, 65) elif j <= 2: ax.set_ylim(0, 25) # Legend (global, right side) - create in specified order # Get handles and labels from first subplot handles_dict = {} labels_dict = {} for handle, label in zip(*axes[0, 0].get_legend_handles_labels()): handles_dict[label] = handle labels_dict[label] = label # Create ordered handles and labels ordered_handles = [] ordered_labels = [] for bm in base_models: if bm in handles_dict: ordered_handles.append(handles_dict[bm]) ordered_labels.append(labels_dict[bm]) # Add OOM handle oom_handle = Line2D([0], [0], color="red", marker="x", linestyle="--", markersize=8, label="Out of Memory (OOM)") ordered_handles.append(oom_handle) ordered_labels.append("Out of Memory (OOM)") # Legend in top left plot only top_left_ax = axes[0, 0] top_left_ax.legend( ordered_handles, ordered_labels, title=None, loc="upper right", frameon=True, framealpha=0.95, edgecolor="0.3", fontsize=12, ) # Add main vertical title fig.text(0.02, 0.5, "Peak Memory (GB)", rotation=90, fontsize=20, fontweight="bold", ha="center", va="center") # Add global x-axis label spanning the bottom row fig.text(0.5, 0.01, "Number of Time Series (N)", fontsize=16, fontweight="bold", ha="center", va="center") # Layout - no longer need space for bottom legend plt.tight_layout(pad=0.5) plt.subplots_adjust(left=0.08) # Save for fmt in ["png", "pdf"]: plt.savefig( f"memory_usage_paper.{fmt}", dpi=300 if fmt == "png" else None, bbox_inches="tight", pad_inches=0, facecolor="white", format=fmt, ) plt.show() if __name__ == "__main__": plot_memory_usage_paper()