[tab] plot.ipynb (python)
# --- 5. PROCESS AND PLOT DATA ---
plt.style.use('seaborn-v0_8-whitegrid')

# Create two subplots: one for the original loss curves, one for variance analysis
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(20, 8))

all_plot_data = []
variance_data = []

# Filter to only include the specific runs you want
target_runs = [
    "dynamics-batch-size-scaling-1-node-3335740",
    "dynamics-batch-size-scaling-2-node-3335741", 
    "dynamics-batch-size-scaling-16-node-3335739",
    "dynamics-batch-size-scaling-4-node-3335742",
    "dynamics-batch-size-scaling-8-node-3335743"
]

runs_with_data = []
for run in runs:
    if run.name in target_runs:
        try:
            history = list(run.scan_history())
            if history:  # Only include runs with history data
                runs_with_data.append(run)
        except:
            continue

print(f"Found {len(runs_with_data)} target runs with actual data")

for run in runs_with_data:
    # Get run history (the metrics over time) as a pandas DataFrame
    history = run.scan_history(keys=["step", METRIC_NAME])
    history_df = pd.DataFrame(history)

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

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

    # 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={config['batch_size']}" # Create a nice label for the legend
    })
    
    # Calculate variance for this run
    loss_variance = history_df[METRIC_NAME].var()
    variance_data.append({
        'batch_size': config['batch_size'],
        'variance': loss_variance,
        'run_name': run.name
    })
    
    print(f"Processed run: {run.name} (Batch Size: {config['batch_size']}, Variance: {loss_variance:.6f})")

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

# Plot 1: Original loss curves
for data in all_plot_data:
    ax1.plot(data['x'], data['y'], label=data['label'])

ax1.set_title("Token-Normalized Loss vs. Batch Size", fontsize=16)
ax1.set_xlabel("Total Tokens Seen", fontsize=12)
ax1.set_ylabel(METRIC_NAME, fontsize=12)
ax1.set_yscale('log')  # Loss is often best viewed on a log scale
ax1.legend(title="Batch Size")
ax1.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))
ax1.xaxis.set_major_formatter(formatter)

# Plot 2: Variance analysis
variance_df = pd.DataFrame(variance_data)
variance_df = variance_df.sort_values('batch_size')

ax2.plot(variance_df['batch_size'], variance_df['variance'], 'o-', linewidth=2, markersize=8)
ax2.set_title("Loss Variance vs. Batch Size", fontsize=16)
ax2.set_xlabel("Batch Size", fontsize=12)
ax2.set_ylabel("Loss Variance", fontsize=12)
ax2.set_xscale('log')  # Log scale for batch size
ax2.set_yscale('log')  # Log scale for variance
ax2.grid(True, which="both", linestyle='--', linewidth=0.5)

# Add batch size labels on points
for _, row in variance_df.iterrows():
    ax2.annotate(f"BS={row['batch_size']}", 
                (row['batch_size'], row['variance']), 
                xytext=(5, 5), textcoords='offset points', fontsize=10)

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

print("
Plot saved as loss_and_variance_analysis.png")

# Print summary statistics
print("
=== Variance Analysis Summary ===")
print(variance_df[['batch_size', 'variance']].to_string(index=False))