p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] experiments/train_tokenizer_lr_1e-4_larger_ffn.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=24:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/big-runs/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/big-runs/%x_%j.log
#SBATCH --job-name=train_tokenizer_lr_sweep_1e-4_larger_ffn

# Log the sbatch script
cat $0

module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate

array_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrecords_chunked

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=$PWD/checkpoints/$job_name
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

srun python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --restore_ckpt \
    --save_ckpt \
    --batch_size=96 \
    --ffn_dim=2048 \
    --num_blocks=4 \
    --init_lr=0 \
    --max_lr=1e-4 \
    --decay_end=0.0 \
    --wsd_decay_steps=10000 \
    --num_steps=53000 \
    --log_image_interval=1000 \
    --log_checkpoint_interval=1000 \
    --log \
    --name=mixed-precision-tokenizer-lr-1e-4-larger-ffn \
    --tags tokenizer mixed-precision 1e-4 mixed-precision flash-attention larger-ffn \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir
[content] fetch_wandb_run.py (python)
    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]
[content] plot.ipynb (plaintext)
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 = "instant-uv"
PROJECT = "jafar"

# 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 = "dynamics-batch-size-scaling" 

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


# --- 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

    # --- 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")
[content] plot.ipynb (python)
# --- 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

    # --- 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")
[content] plot.ipynb (python)
# --- 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

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

    # 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
    })
    
    print(f"Processed run: {run.name} (Batch Size: {config['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")
[content] plot.ipynb (python)
# --- 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
    
    # Debug: Print column names to see what's actually available
    print(f"Available columns: {history_df.columns.tolist()}")
    print(f"First few rows: {history_df.head()}")

    # --- THE CORE NORMALIZATION STEP ---
    # Calculate total tokens seen at each step
    # Use the correct column name based on what we see above
    step_column = '_step' if '_step' in history_df.columns else 'Step'
    history_df['tokens_seen'] = history_df[step_column] * config['batch_size'] * 900 * 16

    # 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
    })
    
    print(f"Processed run: {run.name} (Batch Size: {config['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")
[content] plot.ipynb (python)
# --- 5. PROCESS AND PLOT DATA ---
plt.style.use('seaborn-v0_8-whitegrid')
fig, ax = plt.subplots(figsize=(12, 8))

all_plot_data = []

for i, run in enumerate(runs):
    if i >= 3:  # Only check first 3 runs to avoid too much output
        break
        
    print(f"
=== Run {i+1}: {run.name} ===")
    print(f"Config: {run.config}")
    
    # Check what keys are available in the history
    try:
        # Get a sample of history to see what keys exist
        sample_history = list(run.scan_history())
        if sample_history:
            print(f"Available keys in history: {list(sample_history[0].keys())}")
            print(f"Sample history entry: {sample_history[0]}")
        else:
            print("No history data found")
    except Exception as e:
        print(f"Error getting history: {e}")
    
    # Try to get the full history with all keys
    try:
        history = run.scan_history()
        history_df = pd.DataFrame(history)
        print(f"Full history columns: {history_df.columns.tolist()}")
        print(f"History shape: {history_df.shape}")
        if not history_df.empty:
            print(f"First few rows: {history_df.head()}")
    except Exception as e:
        print(f"Error getting full history: {e}")

print("
=== End of debugging ===")
[content] plot.ipynb (python)
# Check run status and available metrics
for i, run in enumerate(runs):
    if i >= 5:  # Check first 5 runs
        break
        
    print(f"
=== Run {i+1}: {run.name} ===")
    print(f"State: {run.state}")
    print(f"Created: {run.created_at}")
    print(f"Updated: {run.updated_at}")
    
    # Check if there are any logged metrics
    try:
        # Try to get summary stats
        summary = run.summary
        print(f"Summary keys: {list(summary.keys()) if summary else 'No summary'}")
        
        # Try different ways to get history
        try:
            # Try without specifying keys first
            history = list(run.scan_history())
            print(f"History entries: {len(history)}")
            if history:
                print(f"First entry keys: {list(history[0].keys())}")
        except Exception as e:
            print(f"Error with scan_history: {e}")
            
        # Try with history() method
        try:
            history_df = run.history()
            print(f"History shape: {history_df.shape}")
            if not history_df.empty:
                print(f"History columns: {history_df.columns.tolist()}")
        except Exception as e:
            print(f"Error with history(): {e}")
            
    except Exception as e:
        print(f"Error accessing run data: {e}")

print("
=== End of status check ===")
[content] plot.ipynb (python)
# Check run status and available metrics
for i, run in enumerate(runs):
    if i >= 10:  # Check first 10 runs
        break
        
    print(f"
=== Run {i+1}: {run.name} ===")
    print(f"State: {run.state}")
    print(f"Created: {run.created_at}")
    
    # Always try to get data, regardless of run state
    try:
        # Try to get summary stats
        summary = run.summary
        print(f"Summary keys: {list(summary.keys()) if summary else 'No summary'}")
        
        # Try different ways to get history
        try:
            # Try without specifying keys first
            history = list(run.scan_history())
            print(f"History entries: {len(history)}")
            if history:
                print(f"First entry keys: {list(history[0].keys())}")
                print(f"Last entry keys: {list(history[-1].keys())}")
        except Exception as e:
            print(f"Error with scan_history: {e}")
            
        # Try with history() method
        try:
            history_df = run.history()
            print(f"History shape: {history_df.shape}")
            if not history_df.empty:
                print(f"History columns: {history_df.columns.tolist()}")
                print(f"First few rows:")
                print(history_df.head())
        except Exception as e:
            print(f"Error with history(): {e}")
            
    except Exception as e:
        print(f"Error accessing run data: {e}")

print("
=== End of status check ===")
[content] plot.ipynb (python)
# --- 5. PROCESS AND PLOT DATA ---
plt.style.use('seaborn-v0_8-whitegrid')
fig, ax = plt.subplots(figsize=(12, 8))

all_plot_data = []

# Filter runs that have actual data
runs_with_data = []
for run in 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)} runs with actual data out of {len(runs)} total runs")

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'] * config['image_width'] * config['image_height']

    # 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
    })
    
    print(f"Processed run: {run.name} (Batch Size: {config['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")
[content] plot.ipynb (python)
# --- 5. PROCESS AND PLOT DATA ---
plt.style.use('seaborn-v0_8-whitegrid')
fig, ax = plt.subplots(figsize=(12, 8))

all_plot_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'] * config['image_width'] * config['image_height']

    # 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
    })
    
    print(f"Processed run: {run.name} (Batch Size: {config['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")