[content] fetch_wandb_run.py (python)
import wandb
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker

# --- 1. SET YOUR PROJECT DETAILS ---
# Replace with your entity (username or team name) and project name
ENTITY = "your-entity"
PROJECT = "your-project-name"

# Optional: If your runs have a specific tag, you can use it to filter them.
# This is highly recommended to isolate your ablation runs.
RUN_TAG = "batch_size_ablation" 

# --- 2. DEFINE METRIC AND CONFIG KEYS ---
# The metric you want to plot (e.g., 'train/loss', 'validation_loss')
METRIC_NAME = "train/loss" 

# The keys in your wandb config for batch size and sequence length
# These MUST match what you logged in your training script.
BATCH_SIZE_KEY = "batch_size"
SEQ_LENGTH_KEY = "sequence_length" # Common in language/world models

# --- 3. INITIALIZE WANDB API ---
api = wandb.Api()

# --- 4. FETCH RUNS ---
# Construct a filter to get only the runs you care about
filters = {"tags": RUN_TAG} if RUN_TAG else {}
runs = api.runs(f"{ENTITY}/{PROJECT}", filters=filters)

print(f"Found {len(runs)} runs to process...")

# --- 5. PROCESS AND PLOT DATA ---
plt.style.use('seaborn-v0_8-whitegrid')
fig, ax = plt.subplots(figsize=(12, 8))

all_plot_data = []

for run in runs:
    # Get run history (the metrics over time) as a pandas DataFrame
    # Using scan_history is more memory-efficient for very long runs.
    history = run.scan_history(keys=["_step", METRIC_NAME])
    history_df = pd.DataFrame(history)

    # Get hyperparameters from the run's config
    config = run.config

    if BATCH_SIZE_KEY not in config or SEQ_LENGTH_KEY not in config:
        print(f"Skipping run {run.name}: missing '{BATCH_SIZE_KEY}' or '{SEQ_LENGTH_KEY}' in config.")
        continue

    batch_size = config[BATCH_SIZE_KEY]
    seq_length = config[SEQ_LENGTH_KEY]

    # --- THE CORE NORMALIZATION STEP ---
    # Calculate total tokens seen at each step
    history_df['tokens_seen'] = history_df['_step'] * batch_size * seq_length

    # Clean up data: drop rows where the metric might be NaN
    history_df.dropna(subset=[METRIC_NAME, 'tokens_seen'], inplace=True)
    
    # Store data for plotting
    all_plot_data.append({
        'x': history_df['tokens_seen'],
        'y': history_df[METRIC_NAME],
        'label': f"BS={batch_size}" # Create a nice label for the legend
    })
    
    print(f"Processed run: {run.name} (Batch Size: {batch_size})")

# Sort the plots by batch size for a cleaner legend
all_plot_data.sort(key=lambda d: int(d['label'].split('=')[1]))

for data in all_plot_data:
    ax.plot(data['x'], data['y'], label=data['label'])

# --- 6. FORMAT THE PLOT ---
ax.set_title("Token-Normalized Loss vs. Batch Size", fontsize=16)
ax.set_xlabel("Total Tokens Seen", fontsize=12)
ax.set_ylabel(METRIC_NAME, fontsize=12)
ax.set_yscale('log')  # Loss is often best viewed on a log scale
ax.legend(title="Batch Size")
ax.grid(True, which="both", linestyle='--', linewidth=0.5)

# Use scientific notation for the x-axis for better readability
formatter = mticker.ScalarFormatter(useMathText=True)
formatter.set_powerlimits((-1, 6))
ax.xaxis.set_major_formatter(formatter)

plt.tight_layout()
plt.savefig("token_normalized_loss_plot.png", dpi=300)
plt.show()

print("
Plot saved as token_normalized_loss_plot.png")