p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] scripts_horeka/batchsize_scaling/adjusted_lr/train_tokenizer_1_nodes.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=05:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=logs/logs_training/%x_%j.log
#SBATCH --error=logs/logs_training/%x_%j.log
#SBATCH --job-name=train_tokenizer_batch_size_scaling_1_node
#SBATCH --mem=100G

# Log the sbatch script
cat $0

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

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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

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

env | grep SLURM

srun python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=48 \
    --min_lr=3e-4 \
    --max_lr=3e-4 \
    --log_image_interval=250 \
    --log \
    --name=tokenizer-batch-size-scaling-1-node-$slurm_job_id \
    --tags tokenizer batch-size-scaling 1-node time-step \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[content] scripts_horeka/modelsize_scaling/dynamics/train_dyn.sh (shellscript)
#!/usr/bin/env bash

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


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


job_name="debug"
slurm_job_id="0000"

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

env | grep SLURM

tokenizer_ckpt_dir=$ws_dir/checkpoints/3290391/tokenizer_1750845012_40000
lam_ckpt_dir=$ws_dir/checkpoints/3290392/lam_1750845133_130000
python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=100 \
    --log \
    --log_checkpoint_interval=500 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --lam_checkpoint=$lam_ckpt_dir \
    --data_dir $tf_records_dir
[tab] scripts_horeka/batchsize_scaling/adjusted_lr/train_tokenizer_32_nodes.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=32
#SBATCH --ntasks-per-node=4
#SBATCH --time=05:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=logs/logs_training/%x_%j.log
#SBATCH --error=logs/logs_training/%x_%j.log
#SBATCH --job-name=train_tokenizer_batch_size_scaling_32_node
#SBATCH --mem=400G

# Log the sbatch script
cat $0

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

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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

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

env | grep SLURM

srun python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1536 \
    --min_lr=3e-4 \
    --max_lr=3e-4 \
    --log_image_interval=250 \
    --log \
    --name=tokenizer-batch-size-scaling-32-node-$slurm_job_id \
    --tags tokenizer batch-size-scaling 32-node \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[content] scripts_horeka/modelsize_scaling/dynamics/model_sizes.md (markdown)
| Config | dyna_dim | dyna_num_blocks | dyna_num_heads | Approx. Params |
|--------|----------|-----------------|---------------|----------------|
| 1 | 512 | 12 | 8 | ~36M |
| 2 | 768 | 16 | 12 | ~110M |
| 3 | 1024 | 16 | 16 | ~180M |
| 4 | 1024 | 24 | 16 | ~270M |
| 5 | 1536 | 24 | 24 | ~500M |
[tab] scripts_horeka/overfit_sample/train_dynamics_overfit_sample.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=05:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --output=logs/logs_training/%x_%j.log
#SBATCH --error=logs/logs_training/%x_%j.log
#SBATCH --mail-user=mihir.mahajan2002@gmail.com
#SBATCH --job-name=train_dynamics_minecraft_overfit_sample
#SBATCH --mem=50G
#SBATCH --mail-type=ALL

# Log the sbatch script
cat $0

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

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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

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

env | grep SLURM

echo Running dynamics model overfit run. Slurm id: $slurm_job_id

tokenizer_ckpt_dir=$ws_dir/checkpoints/3290391/tokenizer_1750845012_40000
lam_ckpt_dir=$ws_dir/checkpoints/3290392/lam_1750845133_130000

python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=100 \
    --log \
    --log_checkpoint_interval=500 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --lam_checkpoint=$lam_ckpt_dir \
    --data_dir $tf_records_dir
[content] scripts_horeka/modelsize_scaling/dynamics/1_train_dyn_36M.sbatch (plaintext)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=05:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --output=logs/logs_training/%x_%j.log
#SBATCH --error=logs/logs_training/%x_%j.log
#SBATCH --mail-user=mihir.mahajan2002@gmail.com
#SBATCH --job-name=train_dynamics_minecraft_overfit_sample
#SBATCH --mem=50G
#SBATCH --mail-type=ALL
[content] scripts_horeka/modelsize_scaling/dynamics/1_train_dyn_36M.sbatch (plaintext)
# Log the sbatch script
cat $0

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

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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

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

env | grep SLURM

echo Running dynamics model overfit run. Slurm id: $slurm_job_id

tokenizer_ckpt_dir=$ws_dir/checkpoints/3290391/tokenizer_1750845012_40000
lam_ckpt_dir=$ws_dir/checkpoints/3290392/lam_1750845133_130000

python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=100 \
    --log \
    --log_checkpoint_interval=500 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --lam_checkpoint=$lam_ckpt_dir \
    --data_dir $tf_records_dir
[content] scripts_horeka/modelsize_scaling/dynamics/1_train_dyn_36M.sbatch (plaintext)
tokenizer_ckpt_dir=$ws_dir/checkpoints/3290391/tokenizer_1750845012_40000
lam_ckpt_dir=$ws_dir/checkpoints/3290392/lam_1750845133_130000

python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=100 \
    --log \
    --log_checkpoint_interval=500 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --lam_checkpoint=$lam_ckpt_dir \
    --data_dir $tf_records_dir
[content] scripts_horeka/overfit_sample/train_dynamics_overfit_sample.sbatch (shellscript)
tokenizer_ckpt_dir=$ws_dir/checkpoints/3290391/tokenizer_1750845012_40000
lam_ckpt_dir=$ws_dir/checkpoints/3290392/lam_1750845133_130000

python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=100 \
    --log \
    --log_checkpoint_interval=500 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --lam_checkpoint=$lam_ckpt_dir \
    --data_dir $tf_records_dir
[content] scripts_horeka/modelsize_scaling/dynamics/1_train_dyn_36M.sbatch (plaintext)
tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290391/tokenizer_1750845012_50000/
lam_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290392/lam_1750845133_180000/

python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=100 \
    --log \
    --name=dynamics-model-size-scaling-36M-$slurm_job_id \
    --tags dynamics model-size-scaling 36M debug \
    --log_checkpoint_interval=500 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --lam_checkpoint=$lam_ckpt_dir \
    --data_dir $tf_records_dir
[content] scripts_horeka/overfit_sample/train_dynamics_overfit_sample.sbatch (shellscript)
tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290391/tokenizer_1750845012_50000/
lam_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290392/lam_1750845133_180000/

python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=100 \
    --log \
    --name=dynamics-model-size-scaling-36M-$slurm_job_id \
    --tags dynamics model-size-scaling 36M debug \
    --log_checkpoint_interval=500 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --lam_checkpoint=$lam_ckpt_dir \
    --data_dir $tf_records_dir
[tab] scripts_horeka/modelsize_scaling/dynamics/1_train_dyn_36M copy.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=05:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --output=logs/logs_training/%x_%j.log
#SBATCH --error=logs/logs_training/%x_%j.log
#SBATCH --mail-user=mihir.mahajan2002@gmail.com
#SBATCH --job-name=train_dynamics_minecraft_overfit_sample
#SBATCH --mem=50G
#SBATCH --mail-type=ALL

# Log the sbatch script
cat $0

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

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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

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

env | grep SLURM

echo Running dynamics model overfit run. Slurm id: $slurm_job_id

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290391/tokenizer_1750845012_50000/
lam_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290392/lam_1750845133_180000/

python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=100 \
    --log \
    --name=dynamics-model-size-scaling-36M-$slurm_job_id \
    --tags dynamics model-size-scaling 36M \
    --log_checkpoint_interval=500 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --lam_checkpoint=$lam_ckpt_dir \
    --data_dir $tf_records_dir
[tab] scripts_horeka/modelsize_scaling/dynamics/2_train_dyn_110M.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=05:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --output=logs/logs_training/%x_%j.log
#SBATCH --error=logs/logs_training/%x_%j.log
#SBATCH --mail-user=mihir.mahajan2002@gmail.com
#SBATCH --job-name=train_dynamics_minecraft_overfit_sample
#SBATCH --mem=50G
#SBATCH --mail-type=ALL

# Log the sbatch script
cat $0

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

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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

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

env | grep SLURM

echo Running dynamics model overfit run. Slurm id: $slurm_job_id

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290391/tokenizer_1750845012_50000/
lam_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290392/lam_1750845133_180000/

python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=100 \
    --log \
    --name=dynamics-model-size-scaling-36M-$slurm_job_id \
    --tags dynamics model-size-scaling 36M \
    --log_checkpoint_interval=500 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --lam_checkpoint=$lam_ckpt_dir \
    --data_dir $tf_records_dir
[content] scripts_horeka/modelsize_scaling/dynamics/2_train_dyn_110M.sbatch (shellscript)
python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=100 \
    --log \
    --name=dynamics-model-size-scaling-110M-$slurm_job_id \
    --tags dynamics model-size-scaling 110M debug \
    --log_checkpoint_interval=500 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --lam_checkpoint=$lam_ckpt_dir \
    --data_dir $tf_records_dir \
    --dyna_dim=768 \
    --dyna_num_blocks=16 \
    --dyna_num_heads=12
[tab] scripts_horeka/modelsize_scaling/dynamics/2_train_dyn_110M copy.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=05:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --output=logs/logs_training/%x_%j.log
#SBATCH --error=logs/logs_training/%x_%j.log
#SBATCH --mail-user=mihir.mahajan2002@gmail.com
#SBATCH --job-name=train_dynamics_minecraft_overfit_sample
#SBATCH --mem=50G
#SBATCH --mail-type=ALL

# Log the sbatch script
cat $0

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

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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

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

env | grep SLURM

echo Running dynamics model overfit run. Slurm id: $slurm_job_id

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290391/tokenizer_1750845012_50000/
lam_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290392/lam_1750845133_180000/

python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=100 \
    --log \
    --name=dynamics-model-size-scaling-110M-$slurm_job_id \
    --tags dynamics model-size-scaling 110M debug \
    --log_checkpoint_interval=500 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --lam_checkpoint=$lam_ckpt_dir \
    --data_dir $tf_records_dir \
    --dyna_dim=768 \
    --dyna_num_blocks=16 \
    --dyna_num_heads=12
[tab] scripts_horeka/modelsize_scaling/dynamics/2_train_dyn_180M.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=05:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --output=logs/logs_training/%x_%j.log
#SBATCH --error=logs/logs_training/%x_%j.log
#SBATCH --mail-user=mihir.mahajan2002@gmail.com
#SBATCH --job-name=train_dynamics_minecraft_overfit_sample
#SBATCH --mem=50G
#SBATCH --mail-type=ALL

# Log the sbatch script
cat $0

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

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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

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

env | grep SLURM

echo Running dynamics model overfit run. Slurm id: $slurm_job_id

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290391/tokenizer_1750845012_50000/
lam_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290392/lam_1750845133_180000/

python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=100 \
    --log \
    --name=dynamics-model-size-scaling-110M-$slurm_job_id \
    --tags dynamics model-size-scaling 110M debug \
    --log_checkpoint_interval=500 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --lam_checkpoint=$lam_ckpt_dir \
    --data_dir $tf_records_dir \
    --dyna_dim=768 \
    --dyna_num_blocks=16 \
    --dyna_num_heads=12
[content] scripts_horeka/modelsize_scaling/dynamics/2_train_dyn_180M.sbatch (shellscript)
python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=100 \
    --log \
    --name=dynamics-model-size-scaling-180M-$slurm_job_id \
    --tags dynamics model-size-scaling 180M debug \
    --log_checkpoint_interval=500 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --lam_checkpoint=$lam_ckpt_dir \
    --data_dir $tf_records_dir \
    --dyna_dim=1024 \
    --dyna_num_blocks=16 \
    --dyna_num_heads=16
[tab] scripts_horeka/modelsize_scaling/dynamics/3_train_dyn_180M.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=05:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --output=logs/logs_training/%x_%j.log
#SBATCH --error=logs/logs_training/%x_%j.log
#SBATCH --mail-user=mihir.mahajan2002@gmail.com
#SBATCH --job-name=train_dynamics_minecraft_overfit_sample
#SBATCH --mem=50G
#SBATCH --mail-type=ALL

# Log the sbatch script
cat $0

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

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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

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

env | grep SLURM

echo Running dynamics model overfit run. Slurm id: $slurm_job_id

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290391/tokenizer_1750845012_50000/
lam_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290392/lam_1750845133_180000/

python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=100 \
    --log \
    --name=dynamics-model-size-scaling-180M-$slurm_job_id \
    --tags dynamics model-size-scaling 180M \
    --log_checkpoint_interval=500 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --lam_checkpoint=$lam_ckpt_dir \
    --data_dir $tf_records_dir \
    --dyna_dim=1024 \
    --dyna_num_blocks=16 \
    --dyna_num_heads=16
[tab] scripts_horeka/modelsize_scaling/dynamics/3_train_dyn_180M copy.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=05:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --output=logs/logs_training/%x_%j.log
#SBATCH --error=logs/logs_training/%x_%j.log
#SBATCH --mail-user=mihir.mahajan2002@gmail.com
#SBATCH --job-name=train_dynamics_minecraft_overfit_sample
#SBATCH --mem=50G
#SBATCH --mail-type=ALL

# Log the sbatch script
cat $0

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

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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

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

env | grep SLURM

echo Running dynamics model overfit run. Slurm id: $slurm_job_id

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290391/tokenizer_1750845012_50000/
lam_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290392/lam_1750845133_180000/

python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=100 \
    --log \
    --name=dynamics-model-size-scaling-180M-$slurm_job_id \
    --tags dynamics model-size-scaling 180M \
    --log_checkpoint_interval=500 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --lam_checkpoint=$lam_ckpt_dir \
    --data_dir $tf_records_dir \
    --dyna_dim=1024 \
    --dyna_num_blocks=16 \
    --dyna_num_heads=16
[tab] scripts_horeka/modelsize_scaling/dynamics/4_train_dyn_270M.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=05:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --output=logs/logs_training/%x_%j.log
#SBATCH --error=logs/logs_training/%x_%j.log
#SBATCH --mail-user=mihir.mahajan2002@gmail.com
#SBATCH --job-name=train_dynamics_minecraft_overfit_sample
#SBATCH --mem=50G
#SBATCH --mail-type=ALL

# Log the sbatch script
cat $0

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

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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

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

env | grep SLURM

echo Running dynamics model overfit run. Slurm id: $slurm_job_id

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290391/tokenizer_1750845012_50000/
lam_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290392/lam_1750845133_180000/

python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=100 \
    --log \
    --name=dynamics-model-size-scaling-180M-$slurm_job_id \
    --tags dynamics model-size-scaling 180M \
    --log_checkpoint_interval=500 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --lam_checkpoint=$lam_ckpt_dir \
    --data_dir $tf_records_dir \
    --dyna_dim=1024 \
    --dyna_num_blocks=16 \
    --dyna_num_heads=16
[content] scripts_horeka/modelsize_scaling/dynamics/4_train_dyn_270M.sbatch (shellscript)
python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=100 \
    --log \
    --name=dynamics-model-size-scaling-270M-$slurm_job_id \
    --tags dynamics model-size-scaling 270M debug \
    --log_checkpoint_interval=500 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --lam_checkpoint=$lam_ckpt_dir \
    --data_dir $tf_records_dir \
    --dyna_dim=1024 \
    --dyna_num_blocks=24 \
    --dyna_num_heads=16
[tab] scripts_horeka/modelsize_scaling/dynamics/4_train_dyn_270M copy.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=05:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --output=logs/logs_training/%x_%j.log
#SBATCH --error=logs/logs_training/%x_%j.log
#SBATCH --mail-user=mihir.mahajan2002@gmail.com
#SBATCH --job-name=train_dynamics_minecraft_overfit_sample
#SBATCH --mem=50G
#SBATCH --mail-type=ALL

# Log the sbatch script
cat $0

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

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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

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

env | grep SLURM

echo Running dynamics model overfit run. Slurm id: $slurm_job_id

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290391/tokenizer_1750845012_50000/
lam_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290392/lam_1750845133_180000/

python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=100 \
    --log \
    --name=dynamics-model-size-scaling-270M-$slurm_job_id \
    --tags dynamics model-size-scaling 270M \
    --log_checkpoint_interval=500 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --lam_checkpoint=$lam_ckpt_dir \
    --data_dir $tf_records_dir \
    --dyna_dim=1024 \
    --dyna_num_blocks=24 \
    --dyna_num_heads=16
[tab] scripts_horeka/modelsize_scaling/dynamics/5_train_dyn_500M.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=05:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --output=logs/logs_training/%x_%j.log
#SBATCH --error=logs/logs_training/%x_%j.log
#SBATCH --mail-user=mihir.mahajan2002@gmail.com
#SBATCH --job-name=train_dynamics_minecraft_overfit_sample
#SBATCH --mem=50G
#SBATCH --mail-type=ALL

# Log the sbatch script
cat $0

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

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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

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

env | grep SLURM

echo Running dynamics model overfit run. Slurm id: $slurm_job_id

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290391/tokenizer_1750845012_50000/
lam_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290392/lam_1750845133_180000/

python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=100 \
    --log \
    --name=dynamics-model-size-scaling-270M-$slurm_job_id \
    --tags dynamics model-size-scaling 270M \
    --log_checkpoint_interval=500 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --lam_checkpoint=$lam_ckpt_dir \
    --data_dir $tf_records_dir \
    --dyna_dim=1024 \
    --dyna_num_blocks=24 \
    --dyna_num_heads=16
[content] scripts_horeka/modelsize_scaling/dynamics/5_train_dyn_500M.sbatch (shellscript)
python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=100 \
    --log \
    --name=dynamics-model-size-scaling-500M-$slurm_job_id \
    --tags dynamics model-size-scaling 500M debug \
    --log_checkpoint_interval=500 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --lam_checkpoint=$lam_ckpt_dir \
    --data_dir $tf_records_dir \
    --dyna_dim=1536 \
    --dyna_num_blocks=24 \
    --dyna_num_heads=24
[tab] /home/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/scripts_horeka/modelsize_scaling/dynamics/2_train_dyn_110M.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=10:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --output=logs/logs_training/%x_%j.log
#SBATCH --error=logs/logs_training/%x_%j.log
#SBATCH --mail-user=mihir.mahajan2002@gmail.com
#SBATCH --job-name=train_dynamics_minecraft_overfit_sample_110M
#SBATCH --mem=50G
#SBATCH --mail-type=ALL

# Log the sbatch script
cat $0

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

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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

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

env | grep SLURM

echo Running dynamics model overfit run. Slurm id: $slurm_job_id

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290391/tokenizer_1750845012_50000/
lam_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3290392/lam_1750845133_180000/

srun python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=100 \
    --log \
    --name=dynamics-model-size-scaling-110M-$slurm_job_id \
    --tags dynamics model-size-scaling 110M \
    --log_checkpoint_interval=500 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --lam_checkpoint=$lam_ckpt_dir \
    --data_dir $tf_records_dir \
    --dyna_dim=768 \
    --dyna_num_blocks=16 \
    --dyna_num_heads=12
[tab] scripts_horeka/get_lrs.py (python)
import math

GPUS_PER_NODE = 4
VRAM_PER_GPU_GB = 40

def sqrt_scaling_lr(base_lr, base_bs, target_bs):
    return base_lr * math.sqrt(target_bs / base_bs)

def get_nodes_gpus(bs):
    """
    Returns (num_nodes, num_gpus) needed for a given batch size.
    - batch size <= 12: 1 GPU, 1 node
    - batch size <= 48: up to 4 GPUs, 1 node
    - batch size > 48: spread across multiple nodes
    """
    if bs <= 12:
        return 1, 1
    elif bs <= 48:
        # Each GPU can take up to 12, so ceil(bs/12) GPUs, all in 1 node
        gpus = math.ceil(bs / 12)
        return 1, gpus
    else:
        # Each node can take up to 48 (4 GPUs * 12)
        nodes = math.ceil(bs / 48)
        gpus = math.ceil(bs / 12)
        return nodes, gpus

def get_batch_sizes_for_nodes(base_per_gpu, nodes_list):
    """
    For each node count, compute the total batch size if each GPU gets base_per_gpu.
    """
    batch_sizes = []
    for nodes in nodes_list:
        total_gpus = nodes * GPUS_PER_NODE
        bs = total_gpus * base_per_gpu
        batch_sizes.append(bs)
    return batch_sizes

# Reference values from train_tokenizer.py, train_lam.py, train_dynamics.py
models = [
    {
        "name": "Tokenizer",
        "base_bs": 48,
        "base_lr": 3e-4,
        "batches": [1, 6, 12, 24, 48, 96],
        "per_gpu_bs": 12,
    },
    {
        "name": "LAM",
        "base_bs": 36,
        "base_lr": 3e-5,
        "batches": [1, 6, 12, 24, 36, 48, 96],
        "per_gpu_bs": 12,
    },
    {
        "name": "Dynamics",
        "base_bs": 36,
        "base_lr": 3e-5,
        "batches": [1, 6, 12, 24, 36, 48, 96],
        "per_gpu_bs": 12,
    },
]

NODES_LIST = [1, 2, 4, 8, 16, 32]

def print_table(model):
    print(f"
## {model['name']}")
    print("| Batch Size | Learning Rate (sqrt scaling) | # Nodes | # GPUs |")
    print("|------------|-----------------------------|---------|--------|")
    # Print original batch sizes
    for bs in model["batches"]:
        lr = sqrt_scaling_lr(model["base_lr"], model["base_bs"], bs)
        nodes, gpus = get_nodes_gpus(bs)
        print(f"| {bs:<10} | {lr:.2e}                  | {nodes:^7} | {gpus:^6} |")
    # Print expanded node/gpu/batch size table
    print("
 Expanded for nodes: 1, 2, 4, 8, 16, 32 (assuming per-GPU batch size = {})".format(model["per_gpu_bs"]))
    print("| Batch Size | Learning Rate (sqrt scaling) | # Nodes | # GPUs |")
    print("|------------|-----------------------------|---------|--------|")
    batch_sizes = get_batch_sizes_for_nodes(model["per_gpu_bs"], NODES_LIST)
    for nodes, bs in zip(NODES_LIST, batch_sizes):
        total_gpus = nodes * GPUS_PER_NODE
        lr = sqrt_scaling_lr(model["base_lr"], model["base_bs"], bs)
        print(f"| {bs:<10} | {lr:.2e}                  | {nodes:^7} | {total_gpus:^6} |")

if __name__ == "__main__":
    print(f"
# Scaling requirements (sqrt lr scaling)")
    for model in models:
        print_table(model)
[tab] scripts_cremers/read_tf_record.py (python)
import tensorflow as tf
import sys
import matplotlib.pyplot as plt
import numpy as np

def get_tfrecord_shape(tfrecord_path):
    raw_dataset = tf.data.TFRecordDataset(tfrecord_path)
    for raw_record in raw_dataset.take(1):
        example = tf.train.Example()
        example.ParseFromString(raw_record.numpy())
        feature = example.features.feature
        height = feature['height'].int64_list.value[0]
        width = feature['width'].int64_list.value[0]
        channels = feature['channels'].int64_list.value[0]
        sequence_length = feature['sequence_length'].int64_list.value[0]
        print(f"Shape: ({sequence_length}, {height}, {width}, {channels})")
        break

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python read_tf_record.py <tfrecord_file_path>")
    else:
        tfrecord_file_path = sys.argv[1]
        get_tfrecord_shape(tfrecord_file_path)

        def visualize_first_frame(tfrecord_path):
            raw_dataset = tf.data.TFRecordDataset(tfrecord_path)
            for raw_record in raw_dataset.take(1):
                example = tf.train.Example()
                example.ParseFromString(raw_record.numpy())
                feature = example.features.feature
                raw_video = feature['raw_video'].bytes_list.value[0]
                height = feature['height'].int64_list.value[0]
                width = feature['width'].int64_list.value[0]
                channels = feature['channels'].int64_list.value[0]
                sequence_length = feature['sequence_length'].int64_list.value[0]

                # Decode the raw video bytes
                video_array = np.frombuffer(raw_video, dtype=np.uint8)
                video_array = video_array.reshape((sequence_length, height, width, channels))

                # Extract the first frame
                first_frame = video_array[0]

                # Save the first frame as an image
                plt.imshow(first_frame)
                plt.axis('off')
                plt.savefig('frame.png', bbox_inches='tight', pad_inches=0)
                plt.close()
                break

        visualize_first_frame(tfrecord_file_path)
[tab] slurm/utils/mihir/read_tf_record.py (python)
import tensorflow as tf
import sys
import matplotlib.pyplot as plt
import numpy as np

def get_tfrecord_shape(tfrecord_path):
    raw_dataset = tf.data.TFRecordDataset(tfrecord_path)
    for raw_record in raw_dataset.take(1):
        example = tf.train.Example()
        example.ParseFromString(raw_record.numpy())
        feature = example.features.feature
        height = feature['height'].int64_list.value[0]
        width = feature['width'].int64_list.value[0]
        channels = feature['channels'].int64_list.value[0]
        sequence_length = feature['sequence_length'].int64_list.value[0]
        print(f"Shape: ({sequence_length}, {height}, {width}, {channels})")
        break

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python read_tf_record.py <tfrecord_file_path>")
    else:
        tfrecord_file_path = sys.argv[1]
        get_tfrecord_shape(tfrecord_file_path)

        def visualize_first_frame(tfrecord_path):
            raw_dataset = tf.data.TFRecordDataset(tfrecord_path)
            for raw_record in raw_dataset.take(1):
                example = tf.train.Example()
                example.ParseFromString(raw_record.numpy())
                feature = example.features.feature
                raw_video = feature['raw_video'].bytes_list.value[0]
                height = feature['height'].int64_list.value[0]
                width = feature['width'].int64_list.value[0]
                channels = feature['channels'].int64_list.value[0]
                sequence_length = feature['sequence_length'].int64_list.value[0]

                # Decode the raw video bytes
                video_array = np.frombuffer(raw_video, dtype=np.uint8)
                video_array = video_array.reshape((sequence_length, height, width, channels))

                # Extract the first frame
                first_frame = video_array[0]

                # Save the first frame as an image
                plt.imshow(first_frame)
                plt.axis('off')
                plt.savefig('frame.png', bbox_inches='tight', pad_inches=0)
                plt.close()
                break

        visualize_first_frame(tfrecord_file_path)
[tab] scripts_horeka/preprocess_dataset.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=03:00:00
#SBATCH --partition=cpuonly
#SBATCH --cpus-per-task=8
#SBATCH --output=logs/logs_preprocessing/%x_%j.log
#SBATCH --error=logs/logs_preprocessing/%x_%j.log
#SBATCH --mail-user=mihir@pdoom.org
#SBATCH --job-name=preprocess_dataset
#SBATCH --mail-type=ALL

# Log the sbatch script
cat $0

source .venv/bin/activate

input_path="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_npy_tmp/10fps_160x90/"
output_path="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_tfrecords_tmp_500_shards"

start_time=$(date +%s)  

python utils/preprocess_dataset.py --source_data_dir=$input_path --output_tfrecords_dir=$output_path --num_shards=500

end_time=$(date +%s)
echo "Time taken: $((end_time - start_time)) seconds"
[tab] slurm/dev/mihir/cremers/batch_sizes.md (markdown)
## Batchsizes for jafar with Minecraft dataset

### Tokenizer
| GPU VRAM | Batch Size | Learning Rate |
|----------|------------|---------------|
| -        | 1          | 4.3e-5        | 
| 12 GB    | 6          | 1e-4          | 
| 24 GB    | 12         | 1.5e-4        | ?? 
| 40 GB    | 12         |               |
| 48 GB    | 24         |               | ?? why does this work on cremers
| 80 GB    | 48         | 3e-4          |

### LAM
| GPU VRAM | Batch Size | Learning Rate |
|----------|------------|---------------|
| -        | 1          | 5e-6/5e-7     | 
| 12 GB    | 6          |               | 
| 24 GB    | 12         |               |
| 48 GB    | 24         |               |
| -        | 36         |  3e-5/3e-6    |
| 80 GB    | 48         |               |
[tab] slurm/dev/mihir/cremers/read_tf_record.py (python)
import tensorflow as tf
import sys
import matplotlib.pyplot as plt
import numpy as np

def get_tfrecord_shape(tfrecord_path):
    raw_dataset = tf.data.TFRecordDataset(tfrecord_path)
    for raw_record in raw_dataset.take(1):
        example = tf.train.Example()
        example.ParseFromString(raw_record.numpy())
        feature = example.features.feature
        height = feature['height'].int64_list.value[0]
        width = feature['width'].int64_list.value[0]
        channels = feature['channels'].int64_list.value[0]
        sequence_length = feature['sequence_length'].int64_list.value[0]
        print(f"Shape: ({sequence_length}, {height}, {width}, {channels})")
        break

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python read_tf_record.py <tfrecord_file_path>")
    else:
        tfrecord_file_path = sys.argv[1]
        get_tfrecord_shape(tfrecord_file_path)

        def visualize_first_frame(tfrecord_path):
            raw_dataset = tf.data.TFRecordDataset(tfrecord_path)
            for raw_record in raw_dataset.take(1):
                example = tf.train.Example()
                example.ParseFromString(raw_record.numpy())
                feature = example.features.feature
                raw_video = feature['raw_video'].bytes_list.value[0]
                height = feature['height'].int64_list.value[0]
                width = feature['width'].int64_list.value[0]
                channels = feature['channels'].int64_list.value[0]
                sequence_length = feature['sequence_length'].int64_list.value[0]

                # Decode the raw video bytes
                video_array = np.frombuffer(raw_video, dtype=np.uint8)
                video_array = video_array.reshape((sequence_length, height, width, channels))

                # Extract the first frame
                first_frame = video_array[0]

                # Save the first frame as an image
                plt.imshow(first_frame)
                plt.axis('off')
                plt.savefig('frame.png', bbox_inches='tight', pad_inches=0)
                plt.close()
                break

        visualize_first_frame(tfrecord_file_path)