p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] requirements-franz.txt (pip-requirements)
absl-py==2.2.2
annotated-types==0.7.0
astunparse==1.6.3
certifi==2025.4.26
cffi==1.17.1
charset-normalizer==3.4.2
chex==0.1.89
click==8.2.1
cloudpickle==3.1.1
dm-pix==0.4.4
docker-pycreds==0.4.0
docstring-parser==0.16
einops==0.8.1
etils==1.12.2
filelock==3.18.0
flatbuffers==25.2.10
flax==0.10.6
fsspec==2025.5.0
gast==0.6.0
gitdb==4.0.12
gitpython==3.1.44
glcontext==3.0.0
glfw==1.12.0
google-pasta==0.2.0
grpcio==1.71.0
gym==0.26.2
gym-notices==0.0.8
gym3==0.3.3
h5py==3.13.0
humanize==4.12.3
idna==3.10
imageio==2.37.0
imageio-ffmpeg==0.3.0
importlib-resources==6.5.2
jax==0.6.1
jax-cuda12-pjrt==0.6.1
jax-cuda12-plugin==0.6.1
jaxlib==0.6.1
keras==3.10.0
libclang==18.1.1
markdown==3.8
markdown-it-py==3.0.0
markupsafe==3.0.2
mdurl==0.1.2
ml-dtypes==0.5.1
moderngl==5.12.0
msgpack==1.1.0
namex==0.0.9
nest-asyncio==1.6.0
numpy==1.26.4
nvidia-cublas-cu12==12.8.4.1
nvidia-cuda-cupti-cu12==12.9.19
nvidia-cuda-nvcc-cu12==12.9.41
nvidia-cuda-runtime-cu12==12.9.37
nvidia-cudnn-cu12==9.10.1.4
nvidia-cufft-cu12==11.4.0.6
nvidia-cusolver-cu12==11.7.4.40
nvidia-cusparse-cu12==12.5.9.5
nvidia-nccl-cu12==2.26.5
nvidia-nvjitlink-cu12==12.9.41
nvidia-nvshmem-cu12==3.2.5
opt-einsum==3.4.0
optax==0.2.4
optree==0.15.0
orbax-checkpoint==0.11.13
packaging==25.0
pillow==11.2.1
platformdirs==4.3.8
procgen==0.10.7
protobuf==5.29.4
psutil==7.0.0
pycparser==2.22
pydantic==2.11.5
pydantic-core==2.33.2
pygments==2.19.1
pyyaml==6.0.2
requests==2.32.3
rich==14.0.0
scipy==1.15.3
sentry-sdk==2.29.1
setproctitle==1.3.6
setuptools==80.8.0
shtab==1.7.2
simplejson==3.20.1
six==1.17.0
smmap==5.0.2
tensorboard==2.19.0
tensorboard-data-server==0.7.2
tensorflow==2.19.0
# tensorflow[and-cuda]
tensorflow-io-gcs-filesystem==0.37.1
tensorstore==0.1.75
termcolor==3.1.0
toolz==1.0.0
treescope==0.1.9
typeguard==4.4.2
typing-extensions==4.13.2
typing-inspection==0.4.1
tyro==0.9.21
urllib3==2.4.0
wandb==0.19.11
werkzeug==3.1.3
wheel==0.45.1
wrapt==1.17.2
zipp==3.21.0
ffmpeg-python==0.2.0
tqdm==4.67.1
[tab] scripts/train_tokenizer.sh (shellscript)
#!/usr/bin/env bash

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


# ws_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/'
tf_records_dir="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_tfrecords_tmp"
ws_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/'
# tf_records_dir='/home/hk-project-p0023960/tum_cte0515/Projects/jafar/data_tfrecord_duplicated'

#job_name=$SLURM_JOB_NAME
#slurm_job_id=$SLURM_JOB_ID
job_name="debug"
slurm_job_id="0000"

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

python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --num_steps=75000 \
    --warmup_steps=2500 \
    --log_checkpoint_interval=18750 \
    --log \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[tab] scripts_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] scripts/train_tokenizer.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=10:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:4
#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_tokenizer_minecraft_overfit
#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


ws_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data'
tf_records_dir="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/tf_records_test"
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

srun python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=192 \
    --num_steps=75000 \
    --warmup_steps=2500 \
    --log_checkpoint_interval=18750 \
    --log \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[tab] scripts_horeka/train_tokenizer.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=10:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:4
#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_tokenizer_minecraft_overfit
#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


ws_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data'
tf_records_dir="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/tf_records_test"
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

srun python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=192 \
    --num_steps=75000 \
    --warmup_steps=2500 \
    --log_checkpoint_interval=18750 \
    --log \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[tab] scripts_cremers/train_tokenizer_overfit_sample.sbatch (shellscript)
#!/bin/bash
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=05:00:00
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1,VRAM:24G
#SBATCH --mem=50G
#SBATCH --output=logs/logs_training/%x_%j.log
#SBATCH --error=logs/logs_training/%x_%j.log
#SBATCH --job-name=train_tokenizer_minecraft_overfit_sample

# Log the sbatch script
cat $0

tf_records_dir="/storage/user/mahajanm/Projects/world-modeling/knoms_tfrecords_500/"
ws_dir='/storage/user/mahajanm/Projects/world-modeling'
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

SLURM_STEP_NODELIST=$SLURM_NODELIST python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=4.3e-5 \
    --max_lr=4.3e-5 \
    --log_image_interval=3 \
    --log \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[tab] scripts_horeka/train_tokenizer.sh (shellscript)
#!/usr/bin/env bash

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


# ws_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/'
tf_records_dir="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_tfrecords_tmp"
ws_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/'
# tf_records_dir='/home/hk-project-p0023960/tum_cte0515/Projects/jafar/data_tfrecord_duplicated'

#job_name=$SLURM_JOB_NAME
#slurm_job_id=$SLURM_JOB_ID
job_name="debug"
slurm_job_id="0000"

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

python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --num_steps=75000 \
    --warmup_steps=2500 \
    --log_checkpoint_interval=18750 \
    --log \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[tab] scripts_cremers/train_tokenizer.sh (shellscript)
#!/usr/bin/env bash

# Log the sbatch script
cat $0

tf_records_dir="/storage/user/mahajanm/Projects/world-modeling/knoms_tfrecords_500/"
ws_dir='/storage/user/mahajanm/Projects/world-modeling'
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

SLURM_STEP_NODELIST=$SLURM_NODELIST python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=4.3e-5 \
    --max_lr=4.3e-5 \
    --log_image_interval=3 \
    --log \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[content] scripts_horeka/train_tokenizer.sh (shellscript)
# Log the sbatch script
cat $0

tf_records_dir="/storage/user/mahajanm/Projects/world-modeling/knoms_tfrecords_500/"
ws_dir='/storage/user/mahajanm/Projects/world-modeling'
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

SLURM_STEP_NODELIST=$SLURM_NODELIST python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=4.3e-5 \
    --max_lr=4.3e-5 \
    --log_image_interval=3 \
    --log \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[content] scripts_horeka/train_tokenizer.sh (shellscript)
# ws_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/'
tf_records_dir="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_tfrecords_tmp"
ws_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/'
# tf_records_dir='/home/hk-project-p0023960/tum_cte0515/Projects/jafar/data_tfrecord_duplicated'

#job_name=$SLURM_JOB_NAME
#slurm_job_id=$SLURM_JOB_ID
job_name="debug"
slurm_job_id="0000"

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

python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --num_steps=75000 \
    --warmup_steps=2500 \
    --log_checkpoint_interval=18750 \
    --log \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[content] scripts_horeka/train_tokenizer.sh (shellscript)
python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=4.3e-5 \
    --max_lr=4.3e-5 \
    --log_image_interval=3 \
    --log \
    --entity instant-uv \
    --project jafar \
    --data_dir $t
[content] scripts_horeka/train_tokenizer.sbatch (shellscript)
CHECKPOINT_DIR=$ws_dir/checkpoints/$job_name_$slurm_job_id
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=4.3e-5 \
    --max_lr=4.3e-5 \
    --log_image_interval=3 \
    --log \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[content] scripts_horeka/train_tokenizer.sbatch (shellscript)
CHECKPOINT_DIR=$ws_dir/checkpoints/$job_name_$slurm_job_id
mkdir -p $CHECKPOINT_DIR

srun python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=192 \
    --num_steps=75000 \
    --warmup_steps=2500 \
    --log_checkpoint_interval=18750 \
    --log \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[tab] scripts_horeka/train_tokenizer_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_tokenizer_minecraft_overfit_sample
#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

python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=4.3e-5 \
    --max_lr=4.3e-5 \
    --log_image_interval=3 \
    --log \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[tab] scripts_horeka/train_lam.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

python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=4.3e-5 \
    --max_lr=4.3e-5 \
    --log_image_interval=3 \
    --log \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[content] scripts_horeka/train_lam.sh (shellscript)
python train_lam.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=3 \
    --log \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[tab] scripts_horeka/train_lam_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_tokenizer_minecraft_overfit_sample
#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

python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=4.3e-5 \
    --max_lr=4.3e-5 \
    --log_image_interval=3 \
    --log \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[content] scripts_horeka/train_lam_overfit_sample.sbatch (shellscript)
python train_lam.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=3 \
    --log \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[content] logs/logs_training/train_tokenizer_minecraft_overfit_sample_3290283.log (log)
2025-06-25 11:34:09.092446: F external/xla/xla/pjrt/distributed/client.h:80] Terminating process because the JAX distributed service detected fatal errors. This most likely indicates that another task died; see the other task logs for more details. Disable Python buffering, i.e. `python -u`, to be sure to see all the previous output. absl::Status: DEADLINE_EXCEEDED: Deadline Exceeded

RPC: /tensorflow.CoordinationService/RegisterTask
/var/spool/slurmd/job3290283/slurm_script: line 42: 1623047 Aborted                 (core dumped) python train_tokenizer.py --ckpt_dir $CHECKPOINT_DIR --batch_size=1 --min_lr=4.3e-5 --max_lr=4.3e-5 --log_image_interval=3 --log --entity instant-uv --project jafar --data_dir $tf_records_dir

============================= JOB FEEDBACK =============================

Job ID: 3290283
Cluster: hk
User/Group: tum_cte0515/hk-project-p0023960
Account: hk-project-p0023960
State: FAILED (exit code 134)
Partition: accelerated
Nodes: 1
Cores per node: 6
Nodelist: hkn0821
CPU Utilized: 00:00:05
CPU Efficiency: 0.21% of 00:39:00 core-walltime
Job Wall-clock time: 00:06:30
Starttime: Wed Jun 25 11:27:39 2025
Endtime: Wed Jun 25 11:34:09 2025
Memory Utilized: 654.56 MB
Memory Efficiency: 0.00% of 0.00 MB
Energy Consumed: 242131 Joule / 67.2586111111111 Watthours
Average node power draw: 620.848717948718 Watt
[content] logs/logs_training/train_tokenizer_minecraft_overfit_sample_3290295.log (log)
2025-06-25 11:40:11.758859: F external/xla/xla/pjrt/distributed/client.h:80] Terminating process because the JAX distributed service detected fatal errors. This most likely indicates that another task died; see the other task logs for more details. Disable Python buffering, i.e. `python -u`, to be sure to see all the previous output. absl::Status: DEADLINE_EXCEEDED: Deadline Exceeded

RPC: /tensorflow.CoordinationService/RegisterTask
/var/spool/slurmd/job3290295/slurm_script: line 40: 2931932 Aborted                 (core dumped) python train_tokenizer.py --ckpt_dir $CHECKPOINT_DIR --batch_size=1 --min_lr=4.3e-5 --max_lr=4.3e-5 --log_image_interval=3 --log --entity instant-uv --project jafar --data_dir $tf_records_dir
[tab] scripts_horeka/train_tokenizer_overfit_batch.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_tokenizer_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

srun python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=4.3e-5 \
    --max_lr=4.3e-5 \
    --log_image_interval=3 \
    --log \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[tab] scripts_horeka/train_lam_overfit_batch.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_tokenizer_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

srun python train_lam.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=5e-7 \
    --max_lr=5e-6 \
    --warmup_steps=125 \
    --log_image_interval=3 \
    --log \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[tab] scripts_cremers/train_tokenizer_overfit_batch.sbatch (shellscript)
#!/bin/bash
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=00:30:00
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1,VRAM:24G
#SBATCH --mem=50G
#SBATCH --output=logs/logs_training/%x_%j.log
#SBATCH --error=logs/logs_training/%x_%j.log
#SBATCH --job-name=train_tokenizer_minecraft_overfit_batch

# Log the sbatch script
cat $0

tf_records_dir="/storage/user/mahajanm/Projects/world-modeling/knoms_tfrecords_500/"
ws_dir='/storage/user/mahajanm/Projects/world-modeling'
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

SLURM_STEP_NODELIST=$SLURM_NODELIST python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=1e-4 \
    --max_lr=1e-4 \
    --log_image_interval=3 \
    --log \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir