# 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 """ Plot memory usage on simulation datasets from memory_simulation.csv. - Only uses datasets starting with 'Simulation-'. - Extracts time series length (L) and number of series (N). - Computes total_length = N * L. - Plots memory vs total_length per base model, comparing SoftPrompt vs Flamingo. - OOM runs (> 180GB) are shown with a dashed line, red X, and "OOM" label. - Always shows panels in order: gemma-270m, gemma-1b, llama-1b, llama-3b. """ import pandas as pd import matplotlib.pyplot as plt import seaborn as sns import matplotlib import re 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_sim(csv_file="memory_simulation.csv"): # --- Paper-style settings --- plt.style.use("seaborn-v0_8-white") matplotlib.rcParams.update({ "font.family": "serif", "font.serif": ["Palatino", "Times New Roman", "DejaVu Serif"], "font.size": 16, "axes.labelsize": 18, "axes.titlesize": 18, "legend.fontsize": 15, "xtick.labelsize": 15, "ytick.labelsize": 15, "axes.linewidth": 0.6, "axes.edgecolor": "0.15", }) df = pd.read_csv(csv_file) # Replace -1 with NaN (ignore failed runs) df["peak_cuda_reserved_gb"] = df["peak_cuda_reserved_gb"].replace(-1, pd.NA) # Keep only simulation datasets df = df[df["dataset"].str.startswith("Simulation-")] # Parse model name and dataset details 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) # Compute total sequence length df["total_length"] = df["L"] * df["N"] # Sort df = df.sort_values(by=["base_model", "config", "total_length"]) # Fixed base_model order base_model_order = ["Gemma-3-270M", "Gemma-3-1B-pt", "Llama-3.2-1B", "Llama-3.2-3B"] # One subplot per model (always 4) n_models = len(base_model_order) fig, axes = plt.subplots(1, n_models, figsize=(3.2 * n_models, 4.5), sharey=True) if n_models == 1: axes = [axes] # Muted palette for configs - order matters for legend palette = {"SoftPrompt": "#4477AA", "Flamingo": "#CC6677"} config_order = ["SoftPrompt", "Flamingo"] for ax, base_model in zip(axes, base_model_order): subdf = df[df["base_model"] == base_model] if subdf.empty: ax.set_title(base_model, fontsize=13, fontweight="bold") ax.set_facecolor("#F8F9FA") ax.text( 0.5, 0.5, "No data", ha="center", va="center", fontsize=10, color="gray" ) ax.set_xticks([]) ax.set_yticks([]) continue for cfg in config_order: cfg_df = subdf[subdf["config"] == cfg] if cfg_df.empty: continue cfg_df = cfg_df.sort_values("total_length") color = palette[cfg] # Successful runs (≤ threshold) ok_df = cfg_df[cfg_df["peak_cuda_reserved_gb"] <= OOM_THRESHOLD] ax.plot( ok_df["total_length"], ok_df["peak_cuda_reserved_gb"], label=cfg, color=color, linewidth=3.0, alpha=0.9, ) # First OOM run (if any) oom_df = cfg_df[cfg_df["peak_cuda_reserved_gb"] > OOM_THRESHOLD] if not oom_df.empty and not ok_df.empty: first_oom = oom_df.iloc[0] last_ok = ok_df.iloc[-1] # dashed line up to OOM ax.plot( [last_ok["total_length"], first_oom["total_length"]], [last_ok["peak_cuda_reserved_gb"], OOM_THRESHOLD * 1.05], color=color, linestyle="--", linewidth=1.5, alpha=0.8, ) # red X marker ax.scatter( first_oom["total_length"], OOM_THRESHOLD * 1.05, color="red", marker="x", s=70, linewidth=2, zorder=5, ) ax.text( first_oom["total_length"], OOM_THRESHOLD * 1.05, "OOM", color="red", fontsize=9, fontweight="bold", ha="center", va="bottom" ) # Titles & labels ax.set_title(base_model, fontsize=17, fontweight="bold") # Only show axis labels on specific subplots if ax == axes[0]: # Leftmost subplot ax.set_ylabel("Peak CUDA Reserved (GB)", fontsize=16, fontweight="bold") ax.set_xlabel("Total Sequence Length (N × L)", fontsize=16, fontweight="bold") else: ax.set_ylabel("") ax.set_xlabel("") 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() ax.tick_params(axis="both", labelsize=15) # Legend only in first subplot if ax == axes[0]: leg = ax.legend(title=None, fontsize=15, loc="best", frameon=True, framealpha=0.95, edgecolor="0.3") for text in leg.get_texts(): text.set_fontweight("bold") plt.tight_layout(pad=0.5) for fmt in ["png", "pdf"]: plt.savefig( f"memory_usage_simulation.{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_sim()