p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] slurm/jobs/mihir/horeka/modelsize_scaling/dynamics/1_train_dyn_36M.sbatch (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/train_tokenizer_batch_size_scaling_16_node/3321526/tokenizer_22000/

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-modelsize-scaling-36M-$slurm_job_id \
    --tags dynamics modelsize-scaling 36M \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
[tab] slurm/jobs/mihir/horeka/batchsize_scaling/dynamics_cotraining/sqrt_lr/train_dynamics_2_nodes.sbatch (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-batchsize-scaling/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/big-runs/dynamics-cotraining-batchsize-scaling/%x_%j.log
#SBATCH --job-name=train_tokenizer_batch_size_scaling_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-batchsize-scaling/$job_name
mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/train_tokenizer_batch_size_scaling_16_node/3321526/tokenizer_22000/

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --restore_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --min_lr=0 \
    --max_lr=1.5e-5 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-batch-size-scaling-2-node-$slurm_job_id \
    --tags dynamics batch-size-scaling 2-node \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
[content] slurm/jobs/mihir/horeka/lr_tuning/dynamics/train_dyn_default_cos.sbatch (plaintext)
#!/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-batchsize-scaling/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/big-runs/dynamics-cotraining-batchsize-scaling/%x_%j.log
#SBATCH --job-name=train_tokenizer_batch_size_scaling_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-batchsize-scaling/$job_name
mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/train_tokenizer_batch_size_scaling_16_node/3321526/tokenizer_22000/

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --restore_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --min_lr=0 \
    --max_lr=1.5e-5 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-batch-size-scaling-2-node-$slurm_job_id \
    --tags dynamics batch-size-scaling 2-node \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
[content] slurm/jobs/mihir/horeka/lr_tuning/dynamics/train_dyn_default_cos.sbatch (plaintext)
#!/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/train_tokenizer_batch_size_scaling_16_node/3321526/tokenizer_22000/

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-modelsize-scaling-36M-$slurm_job_id \
    --tags dynamics modelsize-scaling 36M \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
[tab] utils/misc_utils.py (python)
from jax.tree_util import tree_map, tree_reduce
import optax



def get_lr_schedule(lr_schedule, init_lr, max_lr, final_lr, total_steps, warmup_steps, wsd_decay_steps):
    supported_schedules = ["wsd", "cos", "const"]
    assert lr_schedule in supported_schedules, f"Learning rate schedule not supported. Please use one of {supported_schedules}"
    if lr_schedule == "const":
        return optax.constant_schedule(max_lr)
    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=final_lr 
        )
    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=final_lr, transition_steps=wsd_decay_steps),
        ]
        boundaries = [warmup_steps, total_steps - wsd_decay_steps]
        return optax.join_schedules(schedules, boundaries)


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 copy.py (python)
from jax.tree_util import tree_map, tree_reduce
import optax



def get_lr_schedule(lr_schedule, init_lr, max_lr, final_lr, total_steps, warmup_steps, wsd_decay_steps):
    supported_schedules = ["wsd", "cos", "const"]
    assert lr_schedule in supported_schedules, f"Learning rate schedule not supported. Please use one of {supported_schedules}"
    if lr_schedule == "const":
        return optax.constant_schedule(max_lr)
    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=final_lr 
        )
    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=final_lr, transition_steps=wsd_decay_steps),
        ]
        boundaries = [warmup_steps, total_steps - wsd_decay_steps]
        return optax.join_schedules(schedules, boundaries)


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/.py (python)
from jax.tree_util import tree_map, tree_reduce
import optax



def get_lr_schedule(lr_schedule, init_lr, max_lr, final_lr, total_steps, warmup_steps, wsd_decay_steps):
    supported_schedules = ["wsd", "cos", "const"]
    assert lr_schedule in supported_schedules, f"Learning rate schedule not supported. Please use one of {supported_schedules}"
    if lr_schedule == "const":
        return optax.constant_schedule(max_lr)
    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=final_lr 
        )
    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=final_lr, transition_steps=wsd_decay_steps),
        ]
        boundaries = [warmup_steps, total_steps - wsd_decay_steps]
        return optax.join_schedules(schedules, boundaries)


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/parameter_utils.py (python)
from jax.tree_util import tree_map, tree_reduce
import optax



def get_lr_schedule(lr_schedule, init_lr, max_lr, final_lr, total_steps, warmup_steps, wsd_decay_steps):
    supported_schedules = ["wsd", "cos", "const"]
    assert lr_schedule in supported_schedules, f"Learning rate schedule not supported. Please use one of {supported_schedules}"
    if lr_schedule == "const":
        return optax.constant_schedule(max_lr)
    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=final_lr 
        )
    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=final_lr, transition_steps=wsd_decay_steps),
        ]
        boundaries = [warmup_steps, total_steps - wsd_decay_steps]
        return optax.join_schedules(schedules, boundaries)


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/lr_utils.py (python)
import optax

def get_lr_schedule(lr_schedule, init_lr, max_lr, final_lr, 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=final_lr 
        )
    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=final_lr, transition_steps=wsd_decay_steps),
        ]
        boundaries = [warmup_steps, total_steps - wsd_decay_steps]
        return optax.join_schedules(schedules, boundaries)
[tab] slurm/jobs/mihir/horeka/lr_tuning/dynamics/train_dyn_default_cos.sbatch (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/yoloruns/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/yoloruns/%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/$job_name/$slurm_job_id
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-lr-schedule-cos-$slurm_job_id \
    --tags dynamics lr-schedule cos \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
[content] slurm/jobs/mihir/horeka/lr_tuning/dynamics/train_dyn_default_wsd.sbatch (plaintext)
#!/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/yoloruns/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/yoloruns/%x_%j.log
#SBATCH --job-name=train_dynamics_lr_schedule_cos

# 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/$job_name/$slurm_job_id
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=40000 \
    --warmup_steps=2500 \
    --lr_schedule=cos \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --init_lr=0 \
    --max_lr=1e-4 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --log_checkpoint_keep_period=10000 \
    --name=dynamics-lr-schedule-cos-$slurm_job_id \
    --tags dynamics lr-schedule cos \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
[content] slurm/jobs/mihir/horeka/lr_tuning/dynamics/train_dyn_default_const.sbatch (plaintext)
#!/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/yoloruns/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/yoloruns/%x_%j.log
#SBATCH --job-name=train_dynamics_lr_schedule_wsd

# 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/$job_name/$slurm_job_id
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=40000 \
    --warmup_steps=2500 \
    --lr_schedule=wsd \
    --wsd_decay_steps=7500 \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --init_lr=0 \
    --max_lr=1e-4 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --log_checkpoint_keep_period=10000 \
    --name=dynamics-lr-schedule-wsd-$slurm_job_id \
    --tags dynamics lr-schedule wsd \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
[content] train_tokenizer.py (python)
<<<<<<< HEAD
    decay_end: float = 0.0
    wsd_decay_steps: int = 20000 # Note: wsd_decay_steps will only be used when using a wsd-schedule
    lr_schedule : str = "wsd" # supported options: wsd, cos 
=======
    final_lr: float = 0.0
    wsd_decay_steps: int = 20000
    lr_schedule: str = "wsd" # supported options: wsd, cos, const 
>>>>>>> 8e9ba6a28bebf86b0428231c0931de375ac23cd5
[tab] slurm/jobs/mihir/horeka/yolo-runs/tester.sh (shellscript)

# 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
ws_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/'

job_name="debug"
slurm_job_id="debug-mihir"

CHECKPOINT_DIR=$ws_dir/checkpoints/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

lam_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/checkpoints/train_lam_action_space_scaling_20/3318547/lam_1751657975_200000/
tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/train_tokenizer_batch_size_scaling_16_node/3321526/tokenizer_22000/

env | grep SLURM

srun python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --min_lr=0 \
    --max_lr=1.4e-4 \
    --log_image_interval=1000 \
    --log \
    --num_latent_actions=20 \
    --log_checkpoint_interval=1000 \
    --name=dynamics-yolorun-tf-records-$slurm_job_id \
    --tags dynamics yolo-run tf_records \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir
    # --lam_checkpoint=$lam_ckpt_dir
[tab] slurm/jobs/mihir/horeka/lr_tuning/dynamics/train_dyn_default_wsd.sbatch (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/yoloruns/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/yoloruns/%x_%j.log
#SBATCH --job-name=train_dynamics_lr_schedule_wsd

# 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/$job_name/$slurm_job_id
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=40000 \
    --warmup_steps=2500 \
    --lr_schedule=wsd \
    --wsd_decay_steps=7500 \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --init_lr=0 \
    --max_lr=1e-4 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --log_checkpoint_keep_period=10000 \
    --name=dynamics-lr-schedule-wsd-$slurm_job_id \
    --tags dynamics lr-schedule wsd \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
[content] slurm/jobs/mihir/horeka/yolo-runs/train_dynamics_yolorun.sbatch (shellscript)
# 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
ws_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/'

job_name="debug"
slurm_job_id="debug-mihir"

CHECKPOINT_DIR=$ws_dir/checkpoints/$job_name/$slurm_job_id
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

XLA_FLAGS=--xla_gpu_autotune_level=0 srun python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --min_lr=0 \
    --grad_clip_threshold=10 \
    --max_lr=1e-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] slurm/jobs/mihir/horeka/yolo-runs/train_dynamics_yolorun.sbatch (shellscript)
# 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
# tf_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/open_ai_minecraft_tfrecord

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/checkpoints/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

lam_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/checkpoints/train_lam_action_space_scaling_20/3318547/lam_1751657975_200000/
tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/train_tokenizer_batch_size_scaling_16_node/3321526/tokenizer_22000/

env | grep SLURM

srun python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --min_lr=0 \
    --max_lr=3e-5 \
    --log_image_interval=1000 \
    --log \
    --num_latent_actions=20 \
    --log_checkpoint_interval=1000 \
    --name=dynamics-yolorun-tf-records-$slurm_job_id \
    --tags dynamics yolo-run tf_records \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir
    # --lam_checkpoint=$lam_ckpt_dir
[content] slurm/jobs/mihir/horeka/yolo-runs/train_dynamics_yolorun.sbatch (shellscript)
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

XLA_FLAGS=--xla_gpu_autotune_level=0 srun python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --min_lr=0 \
    --grad_clip_threshold=10 \
    --max_lr=1e-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