diff --git a/logs/step_size_analysis/generate_step_size_table_sfadamw.py b/logs/step_size_analysis/generate_step_size_table_sfadamw.py deleted file mode 100644 index f45e13de..00000000 --- a/logs/step_size_analysis/generate_step_size_table_sfadamw.py +++ /dev/null @@ -1,148 +0,0 @@ -import os -import glob -import pandas as pd -import numpy as np -from pathlib import Path -import json - -# Define the target mapping of submission keys to display names -OPTIMIZERS = { - 'schedule_free_adamw_jax': 'JAX Schedule-Free v1', - 'schedule_free_adamw_jax_v2': 'JAX Schedule-Free v2', - 'schedule_free_adamw': 'PyTorch Schedule-Free v1', - 'schedule_free_adamw_v2': 'PyTorch Schedule-Free v2' -} - -base_log_dir = Path('~/submissions_algorithms/logs/self_tuning').expanduser() - -# Find all workloads present in any of the target folders -workloads = set() -for sub in OPTIMIZERS.keys(): - pattern = os.path.join(base_log_dir, sub, 'study_*', '*') - dirs = glob.glob(pattern) - for d in dirs: - if os.path.isdir(d): - dirname = os.path.basename(d) - base_name = dirname.replace('_pytorch', '').replace('_jax', '') - workloads.add(base_name) - -workloads = sorted(list(workloads)) -print(f"Found workloads: {workloads}") - -# Data structure to hold results: results[opt_name][workload] = list of step_times -results = {opt: {wl: [] for wl in workloads} for opt in OPTIMIZERS.values()} - -for sub_key, opt_display in OPTIMIZERS.items(): - print(f"\nProcessing {opt_display} ({sub_key})...") - for wl in workloads: - # Find all trials for this workload and optimizer - pattern = os.path.join(base_log_dir, sub_key, 'study_*', f"{wl}*", 'trial_*', 'measurements.csv') - files = glob.glob(pattern) - - trial_times = [] - for f in files: - try: - df = pd.read_csv(f) - # Ensure we have the necessary columns and drop rows that don't have them - if 'accumulated_submission_time' in df.columns and 'global_step' in df.columns: - df_valid = df.dropna(subset=['accumulated_submission_time', 'global_step']) - if len(df_valid) >= 2: - first_row = df_valid.iloc[0] - last_row = df_valid.iloc[-1] - - delta_t = last_row['accumulated_submission_time'] - first_row['accumulated_submission_time'] - delta_s = last_row['global_step'] - first_row['global_step'] - - if delta_s > 0: - avg_step_time_ms = (delta_t / delta_s) * 1000.0 - trial_times.append(avg_step_time_ms) - except Exception as e: - print(f"Error processing trial file {f}: {e}") - - results[opt_display][wl] = trial_times - -# Format table values as "mean ± std" or "-" if no data -formatted_table = {} -raw_table = {} # Store numeric mean for optional sorting/analysis - -for opt_display in OPTIMIZERS.values(): - formatted_table[opt_display] = {} - raw_table[opt_display] = {} - for wl in workloads: - times = results[opt_display][wl] - if times: - mean_val = np.mean(times) - std_val = np.std(times) - if len(times) > 1 and std_val > 0.01: - formatted_table[opt_display][wl] = f"{mean_val:.1f} ± {std_val:.1f}" - else: - formatted_table[opt_display][wl] = f"{mean_val:.1f}" - raw_table[opt_display][wl] = mean_val - else: - formatted_table[opt_display][wl] = "N/A" - raw_table[opt_display][wl] = None - -# Generate Markdown table manually -headers = ["Optimizer"] + workloads -md_lines = [] -md_lines.append("| " + " | ".join(headers) + " |") -md_lines.append("| " + " | ".join(["---"] * len(headers)) + " |") -for opt in OPTIMIZERS.values(): - row_vals = [formatted_table[opt][wl] for wl in workloads] - md_lines.append("| " + opt + " | " + " | ".join(row_vals) + " |") -markdown_table = "\n".join(md_lines) - -# Generate LaTeX table -latex_lines = [] -latex_lines.append("\\begin{table*}[t]") -latex_lines.append("\\centering") -latex_lines.append("\\caption{Step Execution Time Comparison (milliseconds per step) across different workloads.}") -latex_lines.append("\\label{tab:step_time_comparison}") - -# Column alignment: l followed by r for each workload column -col_align = "l" + "r" * len(workloads) -latex_lines.append(f"\\begin{{tabular}}{{{col_align}}}") -latex_lines.append("\\toprule") - -# Header row (clean workload names formatted for LaTeX) -escaped_workloads = [wl.replace('_', '\\_') for wl in workloads] -latex_lines.append("Optimizer & " + " & ".join(escaped_workloads) + " \\\\") -latex_lines.append("\\midrule") - -# Data rows -for opt in OPTIMIZERS.values(): - row_vals = [] - for wl in workloads: - val = formatted_table[opt][wl] - # Replace ± with \pm for LaTeX math mode - if "±" in val: - parts = val.split(" ± ") - row_vals.append(f"${parts[0]} \\pm {parts[1]}$") - elif val == "N/A": - row_vals.append("---") - else: - row_vals.append(f"${val}$") - latex_lines.append(f"{opt} & " + " & ".join(row_vals) + " \\\\") - -latex_lines.append("\\bottomrule") -latex_lines.append("\\end{tabular}") -latex_lines.append("\\end{table*}") - -latex_table = "\n".join(latex_lines) - -# Output results to files and console -output_dir = Path('~/submissions_algorithms/logs/step_size_analysis/sfadamw_step_size').expanduser() -output_dir.mkdir(exist_ok=True, parents=True) - -with open(output_dir / 'step_time_comparison.md', 'w') as f: - f.write("# Step Execution Time Comparison of Schedule Free Adamw in Markdown (ms/step)\n\n") - f.write(markdown_table) - f.write("\n\n## LaTeX Source Code\n\n```latex\n") - f.write(latex_table) - f.write("\n```\n") - -print("\n=================== MARKDOWN TABLE ===================") -print(markdown_table) -print("\n==================== LAATEX TABLE ====================") -print(latex_table) -print(f"\nSaved tables to {output_dir / 'step_time_comparison.md'}") diff --git a/logs/step_size_analysis/sfadamw_step_size/step_time_comparison.md b/logs/step_size_analysis/sfadamw_step_size/step_time_comparison.md deleted file mode 100644 index 99f1b547..00000000 --- a/logs/step_size_analysis/sfadamw_step_size/step_time_comparison.md +++ /dev/null @@ -1,28 +0,0 @@ -# Step Execution Time Comparison of Schedule Free Adamw in Markdown (ms/step) - -| Optimizer | criteo1tb | fastmri | finewebedu_lm | imagenet_resnet | imagenet_vit | librispeech_conformer | librispeech_deepspeech | ogbg | wmt | -| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | -| JAX Schedule-Free v1 | 1507.4 ± 402.5 | 252.4 ± 30.2 | 224.7 ± 0.2 | 288.4 ± 16.8 | 618.4 ± 185.1 | 2591.7 ± 78.4 | 1285.0 ± 7.6 | 199.5 ± 2.4 | 131.9 ± 0.2 | -| JAX Schedule-Free v2 | 903.6 ± 6.6 | 195.4 ± 9.7 | 225.7 ± 1.0 | 457.5 ± 42.2 | 512.4 ± 25.6 | 2731.3 ± 50.3 | 2781.9 ± 3.6 | 196.4 ± 1.1 | 132.2 ± 0.3 | -| PyTorch Schedule-Free v1 | 1248.2 ± 375.1 | 90.8 ± 16.8 | 390.8 ± 0.4 | 755.1 ± 0.4 | 797.9 ± 4.0 | 640.6 ± 5.9 | 442.4 ± 3.4 | 228.5 ± 1.3 | 148.2 ± 0.6 | -| PyTorch Schedule-Free v2 | 896.2 ± 33.2 | 78.7 ± 6.8 | 395.6 ± 0.3 | 780.5 ± 76.9 | 876.3 ± 30.6 | 669.2 ± 1.7 | 514.2 ± 27.7 | 217.4 ± 5.8 | 140.4 ± 0.1 | - -## LaTeX Source Code - -```latex -\begin{table*}[t] -\centering -\caption{Step Execution Time Comparison (milliseconds per step) across different workloads.} -\label{tab:step_time_comparison} -\begin{tabular}{lrrrrrrrrr} -\toprule -Optimizer & criteo1tb & fastmri & finewebedu\_lm & imagenet\_resnet & imagenet\_vit & librispeech\_conformer & librispeech\_deepspeech & ogbg & wmt \\ -\midrule -JAX Schedule-Free v1 & $1507.4 \pm 402.5$ & $252.4 \pm 30.2$ & $224.7 \pm 0.2$ & $288.4 \pm 16.8$ & $618.4 \pm 185.1$ & $2591.7 \pm 78.4$ & $1285.0 \pm 7.6$ & $199.5 \pm 2.4$ & $131.9 \pm 0.2$ \\ -JAX Schedule-Free v2 & $903.6 \pm 6.6$ & $195.4 \pm 9.7$ & $225.7 \pm 1.0$ & $457.5 \pm 42.2$ & $512.4 \pm 25.6$ & $2731.3 \pm 50.3$ & $2781.9 \pm 3.6$ & $196.4 \pm 1.1$ & $132.2 \pm 0.3$ \\ -PyTorch Schedule-Free v1 & $1248.2 \pm 375.1$ & $90.8 \pm 16.8$ & $390.8 \pm 0.4$ & $755.1 \pm 0.4$ & $797.9 \pm 4.0$ & $640.6 \pm 5.9$ & $442.4 \pm 3.4$ & $228.5 \pm 1.3$ & $148.2 \pm 0.6$ \\ -PyTorch Schedule-Free v2 & $896.2 \pm 33.2$ & $78.7 \pm 6.8$ & $395.6 \pm 0.3$ & $780.5 \pm 76.9$ & $876.3 \pm 30.6$ & $669.2 \pm 1.7$ & $514.2 \pm 27.7$ & $217.4 \pm 5.8$ & $140.4 \pm 0.1$ \\ -\bottomrule -\end{tabular} -\end{table*} -``` diff --git a/logs/step_time_analysis/generate_step_time_table.ipynb b/logs/step_time_analysis/generate_step_time_table.ipynb new file mode 100644 index 00000000..5b54e5cc --- /dev/null +++ b/logs/step_time_analysis/generate_step_time_table.ipynb @@ -0,0 +1,224 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "7a0f6ac8", + "metadata": {}, + "source": [ + "# Step Execution Time Table Generator\n", + "\n", + "This notebook imports the helper functions from `generate_step_size_table.py` to generate and display publication-grade step execution time comparison tables across algorithms and workloads.\n", + "\n", + "You can select which algorithms (`'sfadamw'`, `'muon'`, `'nadamw'`, `'ademamix'`, `'cautious_nadamw'`, `'lion'`, `'diloco'`, or `'all'`) to include in the table, and view both the Markdown and LaTeX source code directly in the notebook console." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "add3f30b", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Available algorithms in registry:\n", + " - 'sfadamw' (Schedule-Free AdamW): 4 submissions -> ['JAX Schedule-Free v1', 'JAX Schedule-Free v2', 'PyTorch Schedule-Free v1', 'PyTorch Schedule-Free v2']\n", + " - 'muon' (Muon): 5 submissions -> ['JAX Muon v1', 'PyTorch Muon v1', 'PyTorch Muon v2 (JAX HPS)', 'PyTorch Muon Replicated (JAX HPS)', 'PyTorch Muon Replicated (Torch HPS)']\n", + " - 'nadamw' (NAdamW): 3 submissions -> ['JAX NAdamW v1', 'JAX NAdamW Baseline v0.5', 'JAX NAdamW ResNet']\n", + " - 'ademamix' (AdemaMix): 1 submissions -> ['PyTorch AdemaMix']\n", + " - 'cautious_nadamw' (Cautious NAdamW): 1 submissions -> ['JAX Cautious NAdamW']\n", + " - 'lion' (Lion): 1 submissions -> ['PyTorch Lion']\n", + " - 'diloco' (DiLoCo): 2 submissions -> ['PyTorch DiLoCo v1', 'PyTorch DiLoCo v2']\n" + ] + } + ], + "source": [ + "import sys\n", + "from pathlib import Path\n", + "\n", + "# Ensure the directory containing generate_step_size_table.py is in sys.path\n", + "script_dir = Path('./').resolve()\n", + "if str(script_dir) not in sys.path:\n", + " sys.path.append(str(script_dir))\n", + "\n", + "import generate_step_size_table\n", + "from generate_step_size_table import (\n", + " ALGO_CONFIGS,\n", + " resolve_selected_algorithms,\n", + " get_selected_submissions,\n", + " find_workloads,\n", + " collect_step_times,\n", + " format_table_values,\n", + " generate_markdown_table,\n", + " generate_latex_table,\n", + " generate_tables\n", + ")\n", + "\n", + "# Display available algorithms in registry\n", + "print(\"Available algorithms in registry:\")\n", + "for algo_key, algo_info in ALGO_CONFIGS.items():\n", + " versions = list(algo_info['submissions'].values())\n", + " print(f\" - '{algo_key}' ({algo_info['display_name']}): {len(versions)} submissions -> {versions}\")" + ] + }, + { + "cell_type": "markdown", + "id": "739f86a7", + "metadata": {}, + "source": [ + "## Select Algorithms & Generate Tables\n", + "\n", + "You can set `ALGO_NAMES` to:\n", + "- `'all'` to generate a single comprehensive table containing all algorithms, versions, and languages.\n", + "- A single algorithm like `'sfadamw'` or `'muon'`.\n", + "- A list of algorithms like `['sfadamw', 'muon']` or comma-separated string `'sfadamw,muon'` to compare multiple specific algorithms side by side.\n", + "\n", + "All multiple versions and frameworks (JAX, PyTorch, v1/v2, etc.) for every selected algorithm will automatically be included in the table." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "946104c2", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "==================== LATEX TABLE =====================\n", + "\\begin{table*}[t]\n", + "\\centering\n", + "\\caption{Step Execution Time Comparison (Normalized Ratios relative to Schedule-Free AdamW v2) across different workloads.}\n", + "\\label{tab:step_time_comparison}\n", + "\\begin{tabular}{lrrrrrrrrr}\n", + "\\toprule\n", + "Optimizer & criteo1tb & fastmri & finewebedu\\_lm & imagenet\\_resnet & imagenet\\_vit & librispeech\\_conformer & librispeech\\_deepspeech & ogbg & wmt \\\\\n", + "\\midrule\n", + "JAX Schedule-Free v1 & $1.67 \\pm 0.45$ & $1.29 \\pm 0.15$ & $1.00$ & $0.63 \\pm 0.04$ & $1.21 \\pm 0.36$ & $0.95 \\pm 0.03$ & $0.46$ & $1.02 \\pm 0.01$ & $1.00$ \\\\\n", + "JAX Schedule-Free v2 & $1.00$ & $1.00 \\pm 0.05$ & $1.00$ & $1.00 \\pm 0.09$ & $1.00 \\pm 0.05$ & $1.00 \\pm 0.02$ & $1.00$ & $1.00$ & $1.00$ \\\\\n", + "PyTorch Schedule-Free v1 & $1.39 \\pm 0.42$ & $1.15 \\pm 0.21$ & $0.99$ & $0.97$ & $0.91$ & $0.96$ & $0.86$ & $1.05$ & $1.06$ \\\\\n", + "PyTorch Schedule-Free v2 & $1.00 \\pm 0.04$ & $1.00 \\pm 0.09$ & $1.00$ & $1.00 \\pm 0.10$ & $1.00 \\pm 0.03$ & $1.00$ & $1.00 \\pm 0.05$ & $1.00 \\pm 0.03$ & $1.00$ \\\\\n", + "JAX Muon v1 & $1.96 \\pm 0.05$ & $1.00 \\pm 0.31$ & $2.24$ & $0.89 \\pm 0.20$ & $1.12 \\pm 0.20$ & $0.94 \\pm 0.02$ & $0.97$ & $0.98 \\pm 0.01$ & $1.32$ \\\\\n", + "PyTorch Muon v1 & $1.25 \\pm 0.39$ & $1.22 \\pm 0.16$ & $0.99$ & $0.91 \\pm 0.07$ & $0.97$ & $1.12$ & $1.58 \\pm 0.01$ & $1.08$ & $1.04$ \\\\\n", + "PyTorch Muon v2 (JAX HPS) & $1.80 \\pm 0.02$ & $3.51 \\pm 1.45$ & $0.99$ & $0.88$ & $0.94$ & $1.13$ & $1.57 \\pm 0.02$ & $1.03 \\pm 0.02$ & $0.97$ \\\\\n", + "PyTorch Muon Replicated (JAX HPS) & $0.98$ & $2.27 \\pm 0.88$ & $1.06$ & $0.88 \\pm 0.01$ & $0.95 \\pm 0.03$ & $1.11$ & $1.57 \\pm 0.02$ & $1.09 \\pm 0.01$ & $1.18$ \\\\\n", + "PyTorch Muon Replicated (Torch HPS) & $1.91 \\pm 0.04$ & $1.70 \\pm 0.14$ & $1.08$ & $0.86$ & $1.00 \\pm 0.04$ & $1.14$ & $1.55 \\pm 0.02$ & $1.08 \\pm 0.03$ & $1.28 \\pm 0.02$ \\\\\n", + "JAX NAdamW v1 & $1.01 \\pm 0.02$ & $0.76 \\pm 0.13$ & $1.94$ & $0.60$ & $0.91 \\pm 0.04$ & $0.91 \\pm 0.03$ & $0.89 \\pm 0.05$ & $0.97$ & $0.95 \\pm 0.01$ \\\\\n", + "JAX NAdamW Baseline v0.5 & $1.00$ & $1.04 \\pm 0.15$ & $1.93$ & $0.60$ & $0.91 \\pm 0.03$ & $0.94 \\pm 0.02$ & $0.97$ & $0.97 \\pm 0.02$ & $0.95$ \\\\\n", + "JAX NAdamW ResNet & $1.53 \\pm 0.38$ & $0.92 \\pm 0.29$ & $1.93$ & $0.77 \\pm 0.17$ & $1.63 \\pm 0.86$ & $0.84 \\pm 0.03$ & $0.84 \\pm 0.03$ & $0.98$ & $0.95$ \\\\\n", + "PyTorch AdemaMix & $1.46 \\pm 0.62$ & $1.59 \\pm 0.22$ & $1.02$ & $0.86 \\pm 0.02$ & $0.91$ & $1.10 \\pm 0.01$ & $1.57 \\pm 0.02$ & $1.04 \\pm 0.02$ & $1.10$ \\\\\n", + "JAX Cautious NAdamW & $1.03 \\pm 0.02$ & $1.11 \\pm 0.10$ & $1.96$ & $0.62$ & $0.99 \\pm 0.03$ & $0.97 \\pm 0.02$ & $1.00 \\pm 0.01$ & $1.02$ & $0.95$ \\\\\n", + "PyTorch Lion & $1.03 \\pm 0.02$ & $1.16 \\pm 0.31$ & $0.98$ & $0.86 \\pm 0.03$ & $0.92 \\pm 0.02$ & $1.09$ & $1.54$ & $0.99$ & $1.01$ \\\\\n", + "PyTorch DiLoCo v1 & $2.00 \\pm 0.69$ & $2.89 \\pm 0.38$ & $1.17$ & $0.41 \\pm 0.05$ & $0.64 \\pm 0.10$ & $4.08 \\pm 0.08$ & $5.17 \\pm 0.14$ & $0.90$ & $1.01$ \\\\\n", + "PyTorch DiLoCo v2 & $2.23 \\pm 0.05$ & $2.94 \\pm 0.23$ & $1.18 \\pm 0.03$ & $0.62 \\pm 0.18$ & $0.57$ & $3.91 \\pm 0.26$ & $5.07 \\pm 0.11$ & $0.89 \\pm 0.02$ & $0.98 \\pm 0.02$ \\\\\n", + "\\bottomrule\n", + "\\end{tabular}\n", + "\\end{table*}\n" + ] + } + ], + "source": [ + "# User interactive variables - feel free to change these\n", + "ALGO_NAMES = 'all' # Options: 'all', 'sfadamw', 'muon', ['sfadamw', 'muon'], etc.\n", + "LOG_DIR = '~/submissions_algorithms/logs/self_tuning' # Base log directory containing study folders\n", + "SAVE_DIR = None # Optional: set to path like './output_tables' to save markdown & latex files\n", + "\n", + "# Generate the table across all selected algorithms, versions, and languages\n", + "output = generate_tables(\n", + " algo_args=ALGO_NAMES,\n", + " base_log_dir=LOG_DIR,\n", + " save_dir=SAVE_DIR\n", + ")\n", + "\n", + "# Display the Markdown Table\n", + "# print(\"=================== MARKDOWN TABLE ===================\")\n", + "# print(output['markdown_table'])\n", + "\n", + "# Display the LaTeX Table\n", + "print(\"\\n==================== LATEX TABLE =====================\")\n", + "print(output['latex_table'])" + ] + }, + { + "cell_type": "markdown", + "id": "31ebbf3f", + "metadata": {}, + "source": [ + "## Fine-Grained Step-by-Step Customization (Optional)\n", + "\n", + "If you need to inspect raw step times, filter specific workloads, or customize the table generation step-by-step, you can directly invoke the modular helper functions:" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "b9e242ff", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Workloads found: ['criteo1tb', 'fastmri', 'finewebedu_lm', 'imagenet_resnet', 'imagenet_vit', 'librispeech_conformer', 'librispeech_deepspeech', 'ogbg', 'wmt']\n", + "\n", + "Raw mean step times for 'finewebedu_lm':\n", + " JAX Schedule-Free v1: 224.69 ms\n", + " JAX Schedule-Free v2: 225.71 ms\n", + " PyTorch Schedule-Free v1: 390.75 ms\n", + " PyTorch Schedule-Free v2: 395.64 ms\n", + " JAX Muon v1: 505.06 ms\n", + " PyTorch Muon v1: 392.11 ms\n", + " PyTorch Muon v2 (JAX HPS): 390.16 ms\n", + " PyTorch Muon Replicated (JAX HPS): 421.15 ms\n", + " PyTorch Muon Replicated (Torch HPS): 427.71 ms\n" + ] + } + ], + "source": [ + "# 1. Resolve selected algorithms and gather submissions\n", + "selected_algos = resolve_selected_algorithms(['sfadamw', 'muon'])\n", + "submissions_map = get_selected_submissions(selected_algos)\n", + "\n", + "# 2. Discover available workloads across these submissions\n", + "base_path = Path(LOG_DIR).expanduser()\n", + "workloads = find_workloads(base_path, submissions_map)\n", + "print(f\"Workloads found: {workloads}\\n\")\n", + "\n", + "# 3. Collect step times and compute raw/formatted statistics\n", + "results, raw_table = collect_step_times(base_path, submissions_map, workloads)\n", + "formatted_table = format_table_values(results, workloads)\n", + "\n", + "# 4. Inspect raw mean step times (ms/step) for a specific workload\n", + "print(\"Raw mean step times for 'finewebedu_lm':\")\n", + "for opt_display in [disp for disp, _ in submissions_map.values()]:\n", + " raw_val = raw_table[opt_display]['finewebedu_lm']\n", + " if raw_val is not None:\n", + " print(f\" {opt_display}: {raw_val:.2f} ms\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "myenv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.20" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/logs/step_time_analysis/generate_step_time_table.py b/logs/step_time_analysis/generate_step_time_table.py new file mode 100644 index 00000000..39895e06 --- /dev/null +++ b/logs/step_time_analysis/generate_step_time_table.py @@ -0,0 +1,436 @@ +import os +import glob +import json +import argparse +import pandas as pd +import numpy as np +from pathlib import Path +from collections import OrderedDict + +# Registry of algorithms mapping algorithm key -> display name & submission versions/languages +ALGO_CONFIGS = { + 'sfadamw': { + 'display_name': 'Schedule-Free AdamW', + 'submissions': OrderedDict([ + ('schedule_free_adamw_jax', 'JAX Schedule-Free v1'), + ('schedule_free_adamw_jax_v2', 'JAX Schedule-Free v2'), + ('schedule_free_adamw', 'PyTorch Schedule-Free v1'), + ('schedule_free_adamw_v2', 'PyTorch Schedule-Free v2') + ]) + }, + 'muon': { + 'display_name': 'Muon', + 'submissions': OrderedDict([ + ('muon', 'JAX Muon v1'), + ('muon_torch', 'PyTorch Muon v1'), + ('muon_torch_jax_hps_lr_fix', 'PyTorch Muon v2 (JAX HPS)'), + ('muon_torch_replicated_jax_hps', 'PyTorch Muon Replicated (JAX HPS)'), + ('muon_torch_replicated_torch_hps', 'PyTorch Muon Replicated (Torch HPS)') + ]) + }, + 'nadamw': { + 'display_name': 'NAdamW', + 'submissions': OrderedDict([ + ('nadamw', 'JAX NAdamW v1'), + ('nadamw_baselinev05', 'JAX NAdamW Baseline v0.5'), + ('nadamw_resnet', 'JAX NAdamW ResNet') + ]) + }, + 'ademamix': { + 'display_name': 'AdemaMix', + 'submissions': OrderedDict([ + ('ademamix', 'PyTorch AdemaMix') + ]) + }, + 'cautious_nadamw': { + 'display_name': 'Cautious NAdamW', + 'submissions': OrderedDict([ + ('cautious_nadamw', 'JAX Cautious NAdamW') + ]) + }, + 'lion': { + 'display_name': 'Lion', + 'submissions': OrderedDict([ + ('lion', 'PyTorch Lion') + ]) + }, + 'diloco': { + 'display_name': 'DiLoCo', + 'submissions': OrderedDict([ + ('single_worker_diloco', 'PyTorch DiLoCo v1'), + ('single_worker_dilocov2', 'PyTorch DiLoCo v2') + ]) + } +} + +def resolve_selected_algorithms(algo_args): + """ + Parses and validates requested algorithm flags. + Supports single strings, list of strings, or comma-separated lists. + If 'all' is requested, returns all available algorithm keys. + """ + if isinstance(algo_args, str): + algo_args = [algo_args] + + resolved = [] + for item in algo_args: + for part in item.split(','): + part_clean = part.strip().lower() + if not part_clean: + continue + if part_clean == 'all': + return list(ALGO_CONFIGS.keys()) + elif part_clean in ALGO_CONFIGS: + if part_clean not in resolved: + resolved.append(part_clean) + else: + valid_keys = ", ".join(list(ALGO_CONFIGS.keys()) + ['all']) + raise ValueError(f"Unknown algorithm '{part_clean}'. Available choices: {valid_keys}") + + if not resolved: + return list(ALGO_CONFIGS.keys()) + return resolved + +def get_selected_submissions(selected_algos): + """ + Gathers all submission keys and their display names across selected algorithms. + Returns an OrderedDict mapping: sub_key -> (opt_display, algo_key). + Ensures every version and language of the selected algorithms is included. + """ + submissions_map = OrderedDict() + for algo in selected_algos: + if algo not in ALGO_CONFIGS: + continue + for sub_key, opt_display in ALGO_CONFIGS[algo]['submissions'].items(): + submissions_map[sub_key] = (opt_display, algo) + return submissions_map + +def find_workloads(base_log_dir, selected_submissions_map): + """ + Finds distinct workload names present in any of the selected submission folders. + """ + workloads = set() + for sub_key in selected_submissions_map.keys(): + pattern = os.path.join(base_log_dir, sub_key, 'study_*', '*') + dirs = glob.glob(pattern) + for d in dirs: + if os.path.isdir(d): + dirname = os.path.basename(d) + base_name = dirname.replace('_pytorch', '').replace('_jax', '') + workloads.add(base_name) + return sorted(list(workloads)) + +def collect_step_times(base_log_dir, selected_submissions_map, workloads): + """ + Scans trial measurements.csv files and computes step execution times (ms/step). + Returns `results` dict: results[opt_display][workload] = list of trial step times. + And `raw_table` dict: raw_table[opt_display][workload] = mean step time across trials (or None). + """ + opt_displays = [disp for disp, _ in selected_submissions_map.values()] + results = {disp: {wl: [] for wl in workloads} for disp in opt_displays} + raw_table = {disp: {wl: None for wl in workloads} for disp in opt_displays} + + for sub_key, (opt_display, _) in selected_submissions_map.items(): + for wl in workloads: + pattern = os.path.join(base_log_dir, sub_key, 'study_*', f"{wl}*", 'trial_*', 'measurements.csv') + files = glob.glob(pattern) + + trial_times = [] + for f in files: + try: + df = pd.read_csv(f) + if 'accumulated_submission_time' in df.columns and 'global_step' in df.columns: + df_valid = df.dropna(subset=['accumulated_submission_time', 'global_step']) + if len(df_valid) >= 2: + df_valid = df_valid.sort_values('global_step') + first_row = df_valid.iloc[0] + last_row = df_valid.iloc[-1] + + delta_t = last_row['accumulated_submission_time'] - first_row['accumulated_submission_time'] + delta_s = last_row['global_step'] - first_row['global_step'] + + if delta_s > 0: + avg_step_time_ms = (delta_t / delta_s) * 1000.0 + trial_times.append(avg_step_time_ms) + except Exception as e: + print(f"Warning: error processing trial file {f}: {e}") + + results[opt_display][wl] = trial_times + if trial_times: + raw_table[opt_display][wl] = np.mean(trial_times) + + return results, raw_table + +def normalize_step_times(results, base_log_dir, workloads, submissions_map, normalize_choice='sfadamw_v2'): + """ + Normalizes step execution times (trial times) relative to a baseline algorithm. + - If normalize_choice == 'sfadamw_v2': divides JAX by 'JAX Schedule-Free v2' and PyTorch by 'PyTorch Schedule-Free v2'. + - If normalize_choice == 'sfadamw_v1': divides JAX by 'JAX Schedule-Free v1' and PyTorch by 'PyTorch Schedule-Free v1'. + - If normalize_choice == 'nadamw': divides by 'JAX NAdamW v1'. + Returns (normalized_results, raw_normalized_table, caption). + """ + if normalize_choice == 'none' or not normalize_choice: + return results, None, "Step Execution Time Comparison (milliseconds per step) across different workloads." + + # Map normalize_choice to baseline submission keys + if normalize_choice == 'sfadamw_v2': + jax_base_key = 'schedule_free_adamw_jax_v2' + pt_base_key = 'schedule_free_adamw_v2' + caption = "Step Execution Time Comparison (Normalized Ratios relative to Schedule-Free AdamW v2) across different workloads." + elif normalize_choice == 'sfadamw_v1': + jax_base_key = 'schedule_free_adamw_jax' + pt_base_key = 'schedule_free_adamw' + caption = "Step Execution Time Comparison (Normalized Ratios relative to Schedule-Free AdamW v1) across different workloads." + elif normalize_choice == 'nadamw': + jax_base_key = 'nadamw' + pt_base_key = 'nadamw' + caption = "Step Execution Time Comparison (Normalized Ratios relative to JAX NAdamW v1) across different workloads." + else: + raise ValueError(f"Unknown normalization choice '{normalize_choice}'") + + base_keys = {'JAX': jax_base_key, 'PyTorch': pt_base_key} + base_means = {'JAX': {}, 'PyTorch': {}} + + reverse_map = {sub_k: disp for sub_k, (disp, _) in submissions_map.items()} + + for fw, b_key in base_keys.items(): + if b_key in reverse_map and reverse_map[b_key] in results: + disp = reverse_map[b_key] + for wl in workloads: + times = results[disp].get(wl, []) + base_means[fw][wl] = np.mean(times) if times else None + else: + for wl in workloads: + pattern = os.path.join(base_log_dir, b_key, 'study_*', f"{wl}*", 'trial_*', 'measurements.csv') + files = glob.glob(pattern) + trial_times = [] + for f in files: + try: + df = pd.read_csv(f) + if 'accumulated_submission_time' in df.columns and 'global_step' in df.columns: + df_valid = df.dropna(subset=['accumulated_submission_time', 'global_step']) + if len(df_valid) >= 2: + df_valid = df_valid.sort_values('global_step') + first_row = df_valid.iloc[0] + last_row = df_valid.iloc[-1] + delta_t = last_row['accumulated_submission_time'] - first_row['accumulated_submission_time'] + delta_s = last_row['global_step'] - first_row['global_step'] + if delta_s > 0: + trial_times.append((delta_t / delta_s) * 1000.0) + except Exception: + pass + base_means[fw][wl] = np.mean(trial_times) if trial_times else None + + normalized_results = {} + raw_normalized_table = {} + for opt_display, wl_dict in results.items(): + normalized_results[opt_display] = {} + raw_normalized_table[opt_display] = {} + is_jax = opt_display.startswith('JAX') + fw = 'JAX' if is_jax else 'PyTorch' + + for wl in workloads: + times = wl_dict.get(wl, []) + base_m = base_means[fw].get(wl) + if times and base_m and base_m > 0: + norm_times = [t / base_m for t in times] + normalized_results[opt_display][wl] = norm_times + raw_normalized_table[opt_display][wl] = np.mean(norm_times) + else: + normalized_results[opt_display][wl] = [] + raw_normalized_table[opt_display][wl] = None + + return normalized_results, raw_normalized_table, caption + +def format_table_values(results, workloads, is_normalized=False): + """ + Formats numeric step times as 'mean ± std' (if std > 0.01 and multiple trials) + or 'mean' or 'N/A'. Ratios use 2 decimal places, raw ms/step uses 1 decimal place. + Returns formatted_table dict: formatted_table[opt_display][workload] = string. + """ + formatted_table = {} + for opt_display, wl_dict in results.items(): + formatted_table[opt_display] = {} + for wl in workloads: + times = wl_dict.get(wl, []) + if times: + mean_val = np.mean(times) + std_val = np.std(times) + if is_normalized: + if len(times) > 1 and std_val > 0.01: + formatted_table[opt_display][wl] = f"{mean_val:.2f} ± {std_val:.2f}" + else: + formatted_table[opt_display][wl] = f"{mean_val:.2f}" + else: + if len(times) > 1 and std_val > 0.01: + formatted_table[opt_display][wl] = f"{mean_val:.1f} ± {std_val:.1f}" + else: + formatted_table[opt_display][wl] = f"{mean_val:.1f}" + else: + formatted_table[opt_display][wl] = "N/A" + return formatted_table + +def generate_markdown_table(formatted_table, workloads, opt_displays): + """ + Constructs a Markdown table representing the step execution time comparison. + """ + headers = ["Optimizer"] + workloads + md_lines = [ + "| " + " | ".join(headers) + " |", + "| " + " | ".join(["---"] * len(headers)) + " |" + ] + for opt in opt_displays: + if opt in formatted_table: + row_vals = [formatted_table[opt][wl] for wl in workloads] + md_lines.append("| " + opt + " | " + " | ".join(row_vals) + " |") + return "\n".join(md_lines) + +def generate_latex_table(formatted_table, workloads, opt_displays, + caption="Step Execution Time Comparison (milliseconds per step) across different workloads.", + label="tab:step_time_comparison"): + """ + Constructs a publication-ready LaTeX table representing the step execution time comparison. + """ + latex_lines = [ + "\\begin{table*}[t]", + "\\centering", + f"\\caption{{{caption}}}", + f"\\label{{{label}}}", + f"\\begin{{tabular}}{{{'l' + 'r' * len(workloads)}}}", + "\\toprule" + ] + + escaped_workloads = [wl.replace('_', '\\_') for wl in workloads] + latex_lines.append("Optimizer & " + " & ".join(escaped_workloads) + " \\\\") + latex_lines.append("\\midrule") + + for opt in opt_displays: + if opt not in formatted_table: + continue + row_vals = [] + for wl in workloads: + val = formatted_table[opt][wl] + if "±" in val: + parts = val.split(" ± ") + row_vals.append(f"${parts[0]} \\pm {parts[1]}$") + elif val == "N/A": + row_vals.append("---") + else: + row_vals.append(f"${val}$") + escaped_opt = opt.replace('_', '\\_') + latex_lines.append(f"{escaped_opt} & " + " & ".join(row_vals) + " \\\\") + + latex_lines.extend([ + "\\bottomrule", + "\\end{tabular}", + "\\end{table*}" + ]) + return "\n".join(latex_lines) + +def generate_tables(algo_args='all', base_log_dir='~/submissions_algorithms/logs/self_tuning', save_dir=None, normalize='sfadamw_v2'): + """ + High-level orchestrator function exposed for both CLI and Jupyter/Colab notebooks. + Takes algorithm flags/names, loads step times across all versions/languages, normalizes them (if requested), + and generates both Markdown and LaTeX tables. + """ + base_log_path = Path(base_log_dir).expanduser() + selected_algos = resolve_selected_algorithms(algo_args) + submissions_map = get_selected_submissions(selected_algos) + + workloads = find_workloads(base_log_path, submissions_map) + results, raw_table = collect_step_times(base_log_path, submissions_map, workloads) + + is_normalized = (normalize != 'none' and normalize is not False and normalize is not None) + if is_normalized: + active_results, active_raw, caption = normalize_step_times( + results, base_log_path, workloads, submissions_map, normalize_choice=normalize + ) + else: + active_results, active_raw = results, raw_table + caption = "Step Execution Time Comparison (milliseconds per step) across different workloads." + + formatted_table = format_table_values(active_results, workloads, is_normalized=is_normalized) + + opt_displays = [disp for disp, _ in submissions_map.values()] + markdown_table = generate_markdown_table(formatted_table, workloads, opt_displays) + latex_table = generate_latex_table(formatted_table, workloads, opt_displays, caption=caption) + + if save_dir: + save_path = Path(save_dir).expanduser() + save_path.mkdir(exist_ok=True, parents=True) + md_file = save_path / 'step_time_comparison.md' + tex_file = save_path / 'step_time_comparison.tex' + + title_text = f"# {caption}\n\n" + with open(md_file, 'w') as f: + f.write(title_text) + f.write(markdown_table) + f.write("\n\n## LaTeX Source Code\n\n```latex\n") + f.write(latex_table) + f.write("\n```\n") + print(f"\nSaved tables to {md_file} and {tex_file}") + + return { + 'markdown_table': markdown_table, + 'latex_table': latex_table, + 'formatted_table': formatted_table, + 'raw_table': active_raw, + 'raw_ms_table': raw_table, + 'workloads': workloads, + 'opt_displays': opt_displays, + 'results': active_results, + 'raw_ms_results': results, + 'selected_algos': selected_algos, + 'normalization': normalize + } + +def parse_arguments(): + parser = argparse.ArgumentParser( + description="Generate publication-grade step execution time comparison tables across algorithms and workloads." + ) + parser.add_argument( + '--algo', + type=str, + nargs='+', + default=['all'], + help="Algorithm(s) to include in the table. Options: 'all' or one/more of: " + ", ".join(ALGO_CONFIGS.keys()) + ) + parser.add_argument( + '--log-dir', + type=str, + default='~/submissions_algorithms/logs/self_tuning', + help="Path to the self-tuning log directory." + ) + parser.add_argument( + '--normalize', + type=str, + choices=['sfadamw_v2', 'sfadamw_v1', 'nadamw', 'none'], + default='sfadamw_v2', + help="Normalize step execution times relative to a baseline ('sfadamw_v2' by default, or 'none' for raw ms/step)." + ) + parser.add_argument( + '--save-dir', + type=str, + default=None, + help="Optional directory to save generated markdown and LaTeX table files." + ) + return parser.parse_args() + +def main(): + args = parse_arguments() + print(f"Generating step size table for algorithms: {args.algo} (normalization: {args.normalize}) ...") + output = generate_tables( + algo_args=args.algo, + base_log_dir=args.log_dir, + save_dir=args.save_dir, + normalize=args.normalize + ) + + print("\n=================== MARKDOWN TABLE ===================") + print(output['markdown_table']) + print("\n==================== LATEX TABLE =====================") + print(output['latex_table']) + +if __name__ == "__main__": + main() +