p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] slurm/dev/alfred/overfit_sample/train_tokenizer_overfit_sample_size_0.6_mio.sbatch (shellscript)
#!/usr/bin/env bash
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=5:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:1
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_alfred/logs_training_tokenizer/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_alfred/logs_training_tokenizer/%x_%j.log
#SBATCH --job-name=train_tokenizer_overfit_sample_size_0.6_mio

# Log the sbatch script
cat $0

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

ws_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/'
tf_records_dir=$ws_dir/data/knoms_tfrecords_200_shards

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

tags=("overfit_sample", "tokenizer", "model_scaling", "alfred")

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

env | grep SLURM

srun python train_tokenizer_single_sample.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=4.3e-5 \
    --max_lr=4.3e-4 \
    --log_image_interval=10 \
    --log \
    --entity instant-uv \
    --project jafar \
    --name $job_name \
    --tags $tags \
    --model_dim 64 \
    --data_dir $tf_records_dir
[tab] slurm/dev/alfred/overfit_sample/train_tokenizer_overfit_sample_size_9_mio.sbatch (shellscript)
#!/usr/bin/env bash
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=5:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:1
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_alfred/logs_training_tokenizer/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_alfred/logs_training_tokenizer/%x_%j.log
#SBATCH --job-name=train_tokenizer_overfit_sample_size_9_mio

# Log the sbatch script
cat $0

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

ws_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/'
tf_records_dir=$ws_dir/data/knoms_tfrecords_200_shards

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

tags=("overfit_sample", "tokenizer", "model_scaling", "alfred")

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

env | grep SLURM

srun python train_tokenizer_single_sample.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=4.3e-5 \
    --max_lr=4.3e-4 \
    --log_image_interval=10 \
    --log \
    --entity instant-uv \
    --project jafar \
    --name $job_name \
    --tags $tags \
    --model_dim 256 \
    --data_dir $tf_records_dir
[tab] slurm/dev/alfred/overfit_sample/train_tokenizer_overfit_sample_size_0_5.sh (shellscript)
# Log the sbatch script
cat $0

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

ws_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/'
tf_records_dir=$ws_dir/data/knoms_tfrecords_200_shards

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

tags=("overfit_sample", "tokenizer", "size_0_5", "alfred")

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

env | grep SLURM

python train_tokenizer_single_sample.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=4.3e-5 \
    --max_lr=4.3e-4 \
    --log_image_interval=10 \
    --log \
    --entity instant-uv \
    --project jafar \
    --name $job_name \
    --tags $tags \
    --model_dim 64 \
    --data_dir $tf_records_dir
[tab] slurm/dev/alfred/overfit_sample/train_tokenizer_overfit_sample_size_21_mio.sbatch (shellscript)
#!/usr/bin/env bash
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=5:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:1
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_alfred/logs_training_tokenizer/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_alfred/logs_training_tokenizer/%x_%j.log
#SBATCH --job-name=train_tokenizer_overfit_sample_size_21_mio

# Log the sbatch script
cat $0

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

ws_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/'
tf_records_dir=$ws_dir/data/knoms_tfrecords_200_shards

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

tags=("overfit_sample", "tokenizer", "size_21_mio", "alfred")

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

env | grep SLURM

srun python train_tokenizer_single_sample.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=4.3e-5 \
    --max_lr=4.3e-4 \
    --log_image_interval=10 \
    --log \
    --entity instant-uv \
    --project jafar \
    --name $job_name \
    --tags $tags \
    --model_dim 384 \
    --data_dir $tf_records_dir
[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] slurm/dev/alfred/overfit_sample/train_tokenizer_overfit_sample_size_2_mio.sbatch (shellscript)
#!/usr/bin/env bash
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=5:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:1
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_alfred/logs_training_tokenizer/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_alfred/logs_training_tokenizer/%x_%j.log
#SBATCH --job-name=train_tokenizer_overfit_sample_size_2_mio

# Log the sbatch script
cat $0

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

ws_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/'
tf_records_dir=$ws_dir/data/knoms_tfrecords_200_shards

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

tags=("overfit_sample", "tokenizer", "model_scaling", "alfred")

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

env | grep SLURM

srun python train_tokenizer_single_sample.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=4.3e-5 \
    --max_lr=4.3e-4 \
    --log_image_interval=10 \
    --log \
    --entity instant-uv \
    --project jafar \
    --name $job_name \
    --tags $tags \
    --model_dim 128 \
    --data_dir $tf_records_dir
[tab] sbatch_scripts/train_lam_dev/train_lam_single_batch.sh (shellscript)
module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv_jafar/bin/activate

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID


CHECKPOINT_DIR=$ws_dir/checkpoints/lam/$job_name_$slurm_job_id
mkdir -p $CHECKPOINT_DIR


# data_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_tfrecords_1_shard_overfit'
data_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_tfrecords_500_shards'
# data_dir = '/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_tfrecords_500_shards_overfit_1'

srun python train_lam.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=48 \
    --log_checkpoint_interval=1000 \
    --log_image_interval=10 \
    --seed=0 \
    --min_lr=0.0000433 \
    --max_lr=0.0000433 \
    --log \
    --entity instant-uv \
    --data_dir $data_dir \
    --project jafar
[content] slurm/dev/alfred/overfit_sample/train_tokenizer_overfit_sample_size_0_5.sh (shellscript)
 "--model_dim",
    "384",
    "--latent_dim",
    "32",
    "--num_latents",
    "1024",
    "--patch_size",
    "4",
    "--num_blocks",
    "8",
    "--num_heads",
    "8",
    "--codebook_dropout",
    "0.0"
[content] slurm/dev/alfred/overfit_sample/train_tokenizer_overfit_sample_size_0_5.sh (shellscript)
"--model_dim",
    "384",
    "--latent_dim",
    "32",
    "--num_latents",
    "1024",
    "--patch_size",
    "4",
    "--num_blocks",
    "8",
    "--num_heads",
    "8",
    "--codebook_dropout",
    "0.0"