p-doom/crowd-code-dataset-0.2 huggingface dataset:
[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

tf_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/open_ai_minecraft_tfrecord
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 $tf_records_dir \
    --lam_checkpoint=$lam_ckpt_dir
[tab] slurm/jobs/mihir/horeka/yolo-runs/train_dynamics_yolorun.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=10: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/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/%x_%j.log
#SBATCH --job-name=train_dyn_yolorun

# 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=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 $tf_records_dir \
    --lam_checkpoint=$lam_ckpt_dir
[tab] models/dynamics.py (python)
from typing import Dict, Any

import jax
import jax.numpy as jnp
import flax.linen as nn

from utils.nn import STTransformer


class DynamicsMaskGIT(nn.Module):
    """MaskGIT dynamics model"""

    model_dim: int
    num_latents: int
    num_blocks: int
    num_heads: int
    dropout: float
    mask_limit: float

    def setup(self):
        self.dynamics = STTransformer(
            self.model_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
        )
        self.patch_embed = nn.Embed(self.num_latents, self.model_dim)
        self.mask_token = self.param(
            "mask_token",
            nn.initializers.lecun_uniform(),
            (1, 1, 1, self.model_dim),
        )
        self.action_up = nn.Dense(self.model_dim)

    def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
        # --- Mask videos ---
        vid_embed = self.patch_embed(batch["video_tokens"])
        if training:
            rng1, rng2 = jax.random.split(batch["mask_rng"])
            mask_prob = jax.random.uniform(rng1, minval=self.mask_limit)
            mask = jax.random.bernoulli(rng2, mask_prob, vid_embed.shape[:-1])
            mask = mask.at[:, 0].set(False)
            # before: with mask token
            # vid_embed = jnp.where(jnp.expand_dims(mask, -1), self.mask_token, vid_embed)

            # my change, with gaussian noise
            rng1, _rng = jax.random.split(rng1)
            noise = jax.random.normal(_rng, self.mask_token.shape) 
            vid_embed = jnp.where(jnp.expand_dims(mask, -1), noise, vid_embed)
            

        else:
            mask = None

        # --- Predict transition ---
        act_embed = self.action_up(batch["latent_actions"])
        vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
        logits = self.dynamics(vid_embed)
        return dict(token_logits=logits, mask=mask)
[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

tf_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_tfrecord_uncurrupted-2
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 $tf_records_dir \
    --lam_checkpoint=$lam_ckpt_dir
[tab] slurm/jobs/mihir/horeka/batchsize_scaling/tokenizer/sqrt_lr/tester.sh (shellscript)
#!/usr/bin/env bash

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/10fps_160x90
array_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrecords_chunked_subset
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

srun python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --min_lr=0 \
    --max_lr=1.4e-4 \
    --log_image_interval=1000 \
    --log_checkpoint_interval=1000 \
    --log \
    --name=tokenizer-batch-size-scaling-2-node-sqrt-lr-$slurm_job_id \
    --tags tokenizer batch-size-scaling 2-node sqrt-lr-scaling \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir
[tab] slurm/jobs/mihir/horeka/batchsize_scaling/tokenizer/sqrt_lr/tester.sh (shellscript)
#!/usr/bin/env bash

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/10fps_160x90
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="0000"

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=96 \
    --min_lr=0 \
    --max_lr=1.4e-4 \
    --log_image_interval=1000 \
    --log_checkpoint_interval=1000 \
    --log \
    --name=tokenizer-batch-size-scaling-2-node-sqrt-lr-$slurm_job_id \
    --tags tokenizer batch-size-scaling 2-node sqrt-lr-scaling \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir