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

# Create subplots for different variance analyses
fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(20, 16))

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
    })
    
    # === DIFFERENT VARIANCE CALCULATIONS ===
    
    # 1. Overall variance (original method - potentially biased)
    overall_variance = history_df[METRIC_NAME].var()
    
    # 2. Variance in the last 20% of training (more stable phase)
    last_20_percent = int(len(history_df) * 0.8)
    late_variance = history_df.iloc[last_20_percent:][METRIC_NAME].var()
    
    # 3. Rolling variance (window of 100 steps) - average of local variances
    rolling_variance = history_df[METRIC_NAME].rolling(window=100, min_periods=50).var().mean()
    
    # 4. Variance at similar token counts (e.g., last 10M tokens seen)
    min_tokens = min(history_df['tokens_seen'].max() for data in all_plot_data) if all_plot_data else history_df['tokens_seen'].max()
    similar_tokens_data = history_df[history_df['tokens_seen'] >= min_tokens * 0.8]
    similar_tokens_variance = similar_tokens_data[METRIC_NAME].var() if len(similar_tokens_data) > 0 else np.nan
    
    variance_data.append({
        'batch_size': config['batch_size'],
        'overall_variance': overall_variance,
        'late_variance': late_variance,
        'rolling_variance': rolling_variance,
        'similar_tokens_variance': similar_tokens_variance,
        'run_name': run.name,
        'total_steps': len(history_df)
    })
    
    print(f"Processed run: {run.name} (Batch Size: {config['batch_size']})")
    print(f"  - Overall variance: {overall_variance:.6f}")
    print(f"  - Late training variance: {late_variance:.6f}")
    print(f"  - Rolling variance: {rolling_variance:.6f}")
    print(f"  - Similar tokens variance: {similar_tokens_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')
ax1.legend(title="Batch Size")
ax1.grid(True, which="both", linestyle='--', linewidth=0.5)

# Plot 2: Overall variance
variance_df = pd.DataFrame(variance_data)
variance_df = variance_df.sort_values('batch_size')

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

# Plot 3: Late training variance
ax3.plot(variance_df['batch_size'], variance_df['late_variance'], 'o-', linewidth=2, markersize=8)
ax3.set_title("Late Training Loss Variance vs. Batch Size
(Last 20% of steps)", fontsize=16)
ax3.set_xlabel("Batch Size", fontsize=12)
ax3.set_ylabel("Late Training Variance", fontsize=12)
ax3.set_xscale('log')
ax3.set_yscale('log')
ax3.grid(True, which="both", linestyle='--', linewidth=0.5)

# Plot 4: Rolling variance
ax4.plot(variance_df['batch_size'], variance_df['rolling_variance'], 'o-', linewidth=2, markersize=8)
ax4.set_title("Rolling Loss Variance vs. Batch Size
(100-step window average)", fontsize=16)
ax4.set_xlabel("Batch Size", fontsize=12)
ax4.set_ylabel("Rolling Variance", fontsize=12)
ax4.set_xscale('log')
ax4.set_yscale('log')
ax4.grid(True, which="both", linestyle='--', linewidth=0.5)

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

print("
Plot saved as comprehensive_variance_analysis.png")

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