p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] train_dynamics.py (python)
import optax
import orbax.checkpoint as ocp
import numpy as np
import jax
import jax.numpy as jnp
import tyro
import wandb
import grain

from genie import Genie, restore_genie_components
from utils.dataloader import get_dataloader
from utils.parameter_utils import count_parameters_by_component


@dataclass
class Args:
    # Experiment
    num_steps: int = 200_000
    seed: int = 0
    seq_len: int = 16
    image_channels: int = 3
    image_height: int = 90
    image_width: int = 160
    data_dir: str = ""
    save_ckpt: bool = False
    restore_ckpt: bool = False
    # Optimization
    batch_size: int = 36
    min_lr: float = 0.0
    max_lr: float = 3e-5
    warmup_steps: int = 5000
    # Tokenizer
    tokenizer_dim: int = 512
    latent_patch_dim: int = 32
    num_patch_latents: int = 1024
    patch_size: int = 4
    tokenizer_num_blocks: int = 8
    tokenizer_num_heads: int = 8
    tokenizer_checkpoint: str = ""
    # LAM
    lam_dim: int = 512
    latent_action_dim: int = 32
    num_latent_actions: int = 6
    lam_patch_size: int = 16
    lam_num_blocks: int = 8
    lam_num_heads: int = 8
    lam_checkpoint: str = ""
    # Dynamics
    dyna_dim: int = 512
    dyna_num_blocks: int = 12
    dyna_num_heads: int = 8
    dropout: float = 0.0
    mask_limit: float = 0.5
    # Logging
    log: bool = False
    entity: str = ""
    project: str = ""
    name: str = "train_dynamics"
    tags: list[str] = field(default_factory=lambda: ["dynamics"])
    log_interval: int = 5
    log_image_interval: int = 250
    ckpt_dir: str = ""
    log_checkpoint_interval: int = 25000
    log_checkpoint_keep_period: int = 20000
    log_gradients: bool = False


args = tyro.cli(Args)


def dynamics_loss_fn(params, state, inputs):
    """Compute masked dynamics loss"""
    inputs["videos"] = inputs["videos"].astype(jnp.float32) / 255.0
[tab] slurm/jobs/mihir/horeka/modelsize_scaling/runner.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=48: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_mihir/big-runs/dynamics-cotraining-modelsize-scaling/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/big-runs/dynamics-cotraining-modelsize-scaling/%x_%j.log
#SBATCH --job-name=train_dynamics_modelsize_scaling_36M_2_node

# 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-debug-run-$slurm_job_id \
    --tags dynamics debug \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
[content] train_dynamics.py (python)
    schedules = [
        optax.linear_schedule(init_value=0.0, end_value=max_lr, transition_steps=warmup_steps),
        optax.constant_schedule(value=max_lr),
        optax.linear_schedule(init_value=max_lr, end_value=final_lr, transition_steps=decay_steps),
    ]
    boundaries = [warmup_steps, warmup_steps + stable_steps]
    return optax.join_schedules(schedules, boundaries)
[tab] utils/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] utils/misc_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
[content] utils/misc_utils.py (python)
def _get_lr_schedule():
    supported_schedules = ["wsd", "cos", "const"]
    assert args.lr_schedule in supported_schedules, f"Learning rate schedule not supported. Please use one of {supported_schedules}"
    if args.lr_schedule == "const":
        return optax.constant_schedule(args.max_lr)
    if args.lr_schedule == "cos":
        return optax.warmup_cosine_decay_schedule(
            args.min_lr, args.max_lr, args.warmup_steps, args.num_steps
        )
    if args.lr_schedule == "wsd":
        schedules = [
            optax.linear_schedule(init_value=args.min_lr, end_value=args.max_lr, transition_steps=args.warmup_steps),
            optax.constant_schedule(value=args.max_lr),
            optax.linear_schedule(init_value=args.max_lr, end_value=args.final_lr, transition_steps=args.wsd_decay_steps),
        ]
        boundaries = [args.warmup_steps, args.num_steps - args.wsd_decay_steps]
        return optax.join_schedules(schedules, boundaries)
[content] train_lam.py (python)
    # Optimization
    batch_size: int = 36
    init_lr: float = 0.0
    max_lr: float = 3e-5
    final_lr: float = 0.0
    wsd_decay_steps: int = 10000
    warmup_steps: int = 5000
    lr_schedule : str = "wsd" # supported options: wsd, cos, const
[content] train_tokenizer.py (python)
    # Optimization
    batch_size: int = 36
    init_lr: float = 0.0
    max_lr: float = 3e-5
    final_lr: float = 0.0
    wsd_decay_steps: int = 10000
    warmup_steps: int = 5000
    lr_schedule : str = "wsd" # supported options: wsd, cos, const
[content] train_lam.py (python)
    lr_schedule = get_lr_schedule(args.lr_schedule, 
                                  args.init_lr, 
                                  args.max_lr, 
                                  args.final_lr, 
                                  args.num_steps, 
                                  args.warmup_steps, 
                                  args.wsd_decay_steps)
[content] train_tokenizer.py (python)
    lr_schedule = get_lr_schedule(args.lr_schedule, 
                                  args.init_lr, 
                                  args.max_lr, 
                                  args.final_lr, 
                                  args.num_steps, 
                                  args.warmup_steps, 
                                  args.wsd_decay_steps)
[tab] slurm/jobs/mihir/horeka/sbatch_dir.sh (shellscript)
#!/usr/bin/env bash

# Check if directory argument is provided
if [ $# -lt 1 ]; then
    echo "Usage: $0 <directory>"
    exit 1
fi

DIR="$1"

# Check if the directory exists
if [ ! -d "$DIR" ]; then
    echo "Error: Directory '$DIR' does not exist."
    exit 1
fi

echo "The following files will be sbatched from directory: $DIR"
ls "$DIR"

read -p "Proceed to sbatch all files in $DIR? [y/N] " confirm
confirm=${confirm,,} # tolower

if [[ "$confirm" != "y" ]]; then
    echo "Aborted."
    exit 0
fi

for file in "$DIR"/*; do
    if [ -f "$file" ]; then
        echo "Submitting $file"
        sbatch "$file"
    fi
done
[tab] slurm/jobs/mihir/horeka/modelsize_scaling/runner.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=48: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_mihir/big-runs/dynamics-cotraining-modelsize-scaling/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/big-runs/dynamics-cotraining-modelsize-scaling/%x_%j.log
#SBATCH --job-name=train_dynamics_modelsize_scaling_36M_2_node

# 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 \
    --num_steps=100 \
    --warmup_steps=10 \
    --lr_schedule="wsd" \
    --wsd_decay_steps=20 \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --init_lr=0 \
    --max_lr=1.5e-4 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-debug-run-$slurm_job_id \
    --tags dynamics debug \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \