p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] slurm/jobs/mihir/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_yolo_lr
#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 overfit yolo-lr \
    --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] experiments/dynamics_grain_tok_lam_restore.sh (shellscript)
#!/usr/bin/env bash
source .venv/bin/activate

data_dir="$PWD/data_arrayrecord/dummy"
ckpt_dir="$PWD/checkpoints/dynamics_openai_grain_tok_lam_restore"
tokenizer_ckpt_dir="$PWD/checkpoints/tokenizer_openai_grain_checkpointing"
lam_ckpt_dir="$PWD/checkpoints/lam_openai_grain_checkpointing"

export XLA_FLAGS=--xla_gpu_autotune_level=0
srun python train_dynamics.py \
    --batch_size 48 \
    --tokenizer_checkpoint $tokenizer_ckpt_dir \
    --lam_checkpoint $lam_ckpt_dir \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[tab] experiments/dynamics_grain_tok_restore.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=8
#SBATCH --time=12:00:00
#SBATCH --cpus-per-task=1
#SBATCH --gres=gpu:8
#SBATCH --mem=100GB

source .venv/bin/activate

data_dir="$PWD/data_arrayrecord/dummy"
ckpt_dir="$PWD/checkpoints/dynamics_openai_grain_tok_lam_restore"
tokenizer_ckpt_dir="$PWD/checkpoints/tokenizer_openai_grain_checkpointing"

export XLA_FLAGS=--xla_gpu_autotune_level=0
srun python train_dynamics.py \
    --batch_size 192 \
    --tokenizer_checkpoint $tokenizer_ckpt_dir \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --init_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[tab] slurm-13771.out (plaintext)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
[content] slurm-13772.out (plaintext)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
[content] slurm-13772.out (plaintext)
2025-07-21 18:28:53.334491: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 13 configs for gemm_fusion_dot on a single thread.
2025-07-21 18:28:53.580317: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 13 configs for gemm_fusion_dot on a single thread.
2025-07-21 18:28:53.580318: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 13 configs for gemm_fusion_dot on a single thread.
2025-07-21 18:28:53.613710: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 13 configs for gemm_fusion_dot on a single thread.
2025-07-21 18:28:53.619240: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 13 configs for gemm_fusion_dot on a single thread.
2025-07-21 18:28:53.844724: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 13 configs for gemm_fusion_dot on a single thread.
2025-07-21 18:28:53.895209: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 13 configs for gemm_fusion_dot on a single thread.
2025-07-21 18:28:53.954865: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 13 configs for gemm_fusion_dot on a single thread.
2025-07-21 18:28:56.907284: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 3 configs for gemm_fusion_dot on a single thread.
2025-07-21 18:28:57.186827: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 3 configs for gemm_fusion_dot on a single thread.
2025-07-21 18:28:57.186827: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 3 configs for gemm_fusion_dot on a single thread.
2025-07-21 18:28:57.205879: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 3 configs for gemm_fusion_dot on a single thread.
2025-07-21 18:28:57.209851: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 3 configs for gemm_fusion_dot on a single thread.
2025-07-21 18:28:57.449794: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 3 configs for gemm_fusion_dot on a single thread.
2025-07-21 18:28:57.488386: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 3 configs for gemm_fusion_dot on a single thread.
[content] slurm-13772.out (plaintext)
2025-07-21 18:28:57.553489: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 3 configs for gemm_fusion_dot on a single thread.
2025-07-21 18:28:58.943144: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 11 configs for gemm_fusion_dot on a single thread.
2025-07-21 18:28:59.239381: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 11 configs for gemm_fusion_dot on a single thread.
2025-07-21 18:28:59.272806: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 11 configs for gemm_fusion_dot on a single thread.
2025-07-21 18:28:59.272806: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 11 configs for gemm_fusion_dot on a single thread.
2025-07-21 18:28:59.279890: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 11 configs for gemm_fusion_dot on a single thread.
2025-07-21 18:28:59.531914: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 11 configs for gemm_fusion_dot on a single thread.
2025-07-21 18:28:59.566290: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 11 configs for gemm_fusion_dot on a single thread.
2025-07-21 18:28:59.603794: W external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1154] Compiling 11 configs for gemm_fusion_dot on a single thread.
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 37989616, 'lam': 19349472, 'dynamics': 29734912, 'total': 87074000}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 37989616, 'lam': 19349472, 'dynamics': 29734912, 'total': 87074000}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 37989616, 'lam': 19349472, 'dynamics': 29734912, 'total': 87074000}
Parameter counts:
{'tokenizer': 37989616, 'lam': 19349472, 'dynamics': 29734912, 'total': 87074000}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 37989616, 'lam': 19349472, 'dynamics': 29734912, 'total': 87074000}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 37989616, 'lam': 19349472, 'dynamics': 29734912, 'total': 87074000}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 37989616, 'lam': 19349472, 'dynamics': 29734912, 'total': 87074000}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 37989616, 'lam': 19349472, 'dynamics': 29734912, 'total': 87074000}
[tab] experiments/tokenizer_grain_checkpointing.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=8
#SBATCH --time=12:00:00
#SBATCH --cpus-per-task=1
#SBATCH --gres=gpu:8
#SBATCH --mem=100GB

source .venv/bin/activate

data_dir="$PWD/data_arrayrecord/dummy"
ckpt_dir="$PWD/checkpoints/tokenizer_openai_grain_checkpointing"

export XLA_FLAGS=--xla_gpu_autotune_level=0
srun python train_tokenizer.py \
    --batch_size 96 \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --init_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[content] utils/nn.py (python)
        )(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),