p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] utils/lr_utils.py (python)
import optax

def get_lr_schedule(lr_schedule, init_lr, max_lr, decay_end, total_steps, warmup_steps, wsd_decay_steps):
    supported_schedules = ["wsd", "cos"]
    assert lr_schedule in supported_schedules, f"Learning rate schedule not supported. Please use one of {supported_schedules}"
    if lr_schedule == "cos":
        assert warmup_steps <= total_steps, "Warmup steps can't be greater than total steps."
        return optax.warmup_cosine_decay_schedule(
            init_value=init_lr,
            peak_value=max_lr,
            warmup_steps=warmup_steps,
            decay_steps=total_steps, # Note: decay_steps includes the warmup steps, so we need to pass total value
            end_value=decay_end 
        )
    if lr_schedule == "wsd":
        assert warmup_steps + wsd_decay_steps <= total_steps, "Warmup and decay period is longer than total steps."
        schedules = [
            optax.linear_schedule(init_value=init_lr, end_value=max_lr, transition_steps=warmup_steps),
            optax.constant_schedule(value=max_lr),
            optax.linear_schedule(init_value=max_lr, end_value=decay_end, transition_steps=wsd_decay_steps),
        ]
        boundaries = [warmup_steps, total_steps - wsd_decay_steps]
        return optax.join_schedules(schedules, boundaries)
[content] train_tokenizer.py (python)
    lr_schedule = get_lr_schedule(args.lr_schedule, 
                                  args.init_lr, 
                                  args.max_lr, 
                                  args.decay_end, 
                                  args.num_steps, 
                                  args.warmup_steps, 
                                  args.wsd_decay_steps)
    tx = optax.adamw(learning_rate=lr_schedule, b1=0.9, b2=0.9, weight_decay=1e-4, mu_dtype=args.dtype)
[content] train_lam.py (python)
    lr_schedule = get_lr_schedule(args.lr_schedule, 
                                  args.init_lr, 
                                  args.max_lr, 
                                  args.decay_end, 
                                  args.num_steps, 
                                  args.warmup_steps, 
                                  args.wsd_decay_steps)
    tx = optax.adamw(learning_rate=lr_schedule, b1=0.9, b2=0.9, weight_decay=1e-4, mu_dtype=args.dtype)
[tab] utils/parameter_utils.py (python)
from jax.tree_util import tree_map, tree_reduce


def _count_leaf(x):
    """Count parameters in a single leaf node."""
    if hasattr(x, "size"):
        return x.size
    return 0


def _count_component(component_params):
    """Count total parameters in a component."""
    return tree_reduce(
        lambda x, y: x + y, tree_map(_count_leaf, component_params), initializer=0
    )


def count_parameters_by_component(params):
    """Count parameters for each component of the model.

    Args:
        params: Model parameters dictionary
        component_names: List of component names to count. If None, counts all components.

    Returns:
        Dictionary with parameter counts for each component
    """

    component_names = list(params["params"].keys())
    print(f"Counting all components: {component_names}")

    counts = {}
    total_params = 0

    for name in component_names:
        if "params" in params and name in params["params"]:
            component_params = params["params"][name]
        else:
            component_params = params

        count = _count_component(component_params)
        counts[name] = count
        total_params += count

    counts["total"] = total_params
    return counts
[tab] input_pipeline/download/openai/download_index_files.sh (shellscript)
#!/bin/bash
# Download index files from OpenAI Video-Pre-Training dataset
# https://github.com/openai/Video-Pre-Training
#
# This script downloads specific JSON index files containing metadata
# for the OpenAI Video-Pre-Training dataset from Microsoft Azure blob storage.
#
# Usage:
#   ./download_index_files.sh [output_dir]
#
# Arguments:
#   output_dir: Directory to save downloaded files (default: data/open_ai_index_files)
#
# Example:
#   ./download_index_files.sh /path/to/custom/directory

# Set output directory, use default if not provided
output_dir="${1:-data/open_ai_index_files}" 
mkdir -p $output_dir

# List of index files to download
# These files contain metadata for different video ranges (6xx, 7xx, etc.)
index_files=(
    "all_6xx_Jun_29.json"    
    "all_7xx_Apr_6.json"     
    "all_8xx_Jun_29.json"    
    "all_9xx_Jun_29.json"    
    "all_10xx_Jun_29.json"   
)

# Download each index file from Azure blob storage
for index_file in "${index_files[@]}"; do
    echo "Downloading $index_file..."
    wget https://openaipublic.blob.core.windows.net/minecraft-rl/snapshots/$index_file -O $output_dir/$index_file
done

echo "Download complete. Files saved to: $output_dir"
[tab] input_pipeline/download/download_array_records.sh (shellscript)
#!/bin/bash

# Download and extract array records from Hugging Face
# 
# This script performs a two-step process:
# 1. Downloads compressed array records from a Hugging Face dataset repository
# 2. Extracts the compressed tar files in parallel for better performance
#
# Usage:
#   ./download_array_records.sh [hf_download_dir] [final_dataset_dir]
#
# Arguments:
#   hf_download_dir    - Directory to store compressed downloads (default: data/minecraft_arrayrecords_compressed)
#   final_dataset_dir  - Directory for extracted array records (default: data/minecraft_arrayrecords)

# Set default directories if not provided as arguments
hf_download_dir="${1:-data/minecraft_arrayrecords_compressed}" 
final_dataset_dir="${2:-data/minecraft_arrayrecords}"          

mkdir -p $hf_download_dir
mkdir -p $final_dataset_dir

# Step 1: Download compressed dataset from Hugging Face
echo "Starting download from Hugging Face..."
repo_id=avocadoali/open_ai_minecraft_arrayrecords_chunked
start_time_hf_download=$(date +%s)

HF_HUB_ENABLE_HF_TRANSFER=1 HF_HUB_DISABLE_SYMLINKS=1 \
huggingface-cli download --repo-type dataset $repo_id --local-dir $hf_download_dir

end_time_hf_download=$(date +%s)
echo "Download completed. Time taken: $((end_time_hf_download - start_time_hf_download)) seconds"

# Step 2: Extract compressed array records in parallel
echo "Starting parallel extraction of tar files..."
num_workers=64  # Number of parallel extraction processes
start_time_uncompress=$(date +%s)

# Find all shard tar files and extract them in parallel:
xargs -0 -P $num_workers -I {} bash -c 'echo "Extracting {}"; tar -xf "{}" -C "'$final_dataset_dir'"'

end_time_uncompress=$(date +%s)

# Display timing summary
echo "================================"
echo "Extraction completed successfully!"
echo "Uncompress time: $((end_time_uncompress - start_time_uncompress)) seconds"
echo "Download time: $((end_time_hf_download - start_time_hf_download)) seconds"
echo "Total time: $((end_time_uncompress - start_time_hf_download)) seconds"
echo "Final dataset location: $final_dataset_dir"
[tab] utils/lr_utils.py (python)
import optax

def get_lr_schedule(lr_schedule: str, init_lr: float, max_lr: float, decay_end: float, total_steps: int, warmup_steps: int, wsd_decay_steps: int) -> optax.Schedule:
    supported_schedules = ["wsd", "cos"]
    assert lr_schedule in supported_schedules, f"Learning rate schedule not supported. Please use one of {supported_schedules}"
    if lr_schedule == "cos":
        assert warmup_steps <= total_steps, "Warmup steps can't be greater than total steps."
        return optax.warmup_cosine_decay_schedule(
            init_value=init_lr,
            peak_value=max_lr,
            warmup_steps=warmup_steps,
            decay_steps=total_steps, # Note: decay_steps includes the warmup steps, so we need to pass total value
            end_value=decay_end 
        )
    elif lr_schedule == "wsd":
        assert warmup_steps + wsd_decay_steps <= total_steps, "Warmup and decay period is longer than total steps."
        schedules = [
            optax.linear_schedule(init_value=init_lr, end_value=max_lr, transition_steps=warmup_steps),
            optax.constant_schedule(value=max_lr),
            optax.linear_schedule(init_value=max_lr, end_value=decay_end, transition_steps=wsd_decay_steps),
        ]
        boundaries = [warmup_steps, total_steps - wsd_decay_steps]
        return optax.join_schedules(schedules, boundaries)
    else:
        raise ValueError(f"Learning rate schedule not supported. Please use one of {supported_schedules}")
[tab] experiments/train_dynamics_lr_1e-4_mixed_precision.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/dynamics-cotraining-modelsize-scaling/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/dynamics-cotraining-modelsize-scaling/%x_%j.log
#SBATCH --job-name=train_dynamics_modelsize_scaling_36M_2_node_mixed_precision

# 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=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/dynamics-cotraining-modelsize-scaling/$job_name
mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaling/train_tokenizer_lr_sweep_1e-4

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --min_lr=0 \
    --max_lr=1.5e-4 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-mixed-precision-run-$slurm_job_id \
    --tags dynamics mixed-precision \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
[content] experiments/train_dynamics_lr_1e-4_mixed_precision.sh (shellscript)

#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/dynamics-cotraining-modelsize-scaling/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/dynamics-cotraining-modelsize-scaling/%x_%j.log
#SBATCH --job-name=train_dynamics_modelsize_scaling_36M_2_node_mixed_precision_flash_attention

# Log the sbatch script