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

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=36: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_tokenizer_batch_size_scaling_2_node

# 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

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

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 $tf_records_dir
[tab] scripts_horeka/batchsize_scaling/sqrt_lr/train_tokenizer_1_nodes.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=36: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_tokenizer_batch_size_scaling_1_node

# 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

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

srun python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=48 \
    --min_lr=0 \
    --max_lr=1e-4 \
    --log_image_interval=1000 \
    --log_checkpoint_interval=1000 \
    --log \
    --name=tokenizer-batch-size-scaling-1-node-sqrt-lr-$slurm_job_id \
    --tags tokenizer batch-size-scaling 1-node sqrt-lr \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[tab] scripts_horeka/batchsize_scaling/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
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=12 \
    --min_lr=0 \
    --max_lr=1e-4 \
    --log_image_interval=1000 \
    --log_checkpoint_interval=1000 \
    --log \
    --name=tokenizer-batch-size-scaling-1-node-sqrt-lr-$slurm_job_id \
    --tags tokenizer batch-size-scaling 1-node sqrt-lr \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir
[content] scripts_horeka/batchsize_scaling/sqrt_lr/tester.sh (shellscript)
    --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 $tf_records_dir
[content] utils/nn.py (python)
class FlashAttention(nn.Module):
    dim: int
    num_heads: int
    is_causal: bool = False

    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        B, T, C = x.shape
        H = self.num_heads
        D = C // H

        # Linear projections for Q, K, V
        qkv = nn.Dense(3 * C, use_bias=False)(x)
        qkv = qkv.reshape(B, T, 3, H, D)
        q, k, v = jnp.moveaxis(qkv, 2, 0)  # q, k, v: [B, T, H, D]

        # Apply Flash Attention
        out = flash_mha(q, k, v, is_causal=self.is_causal)  # [B, T, H, D]

        # Recombine heads
        out = out.reshape(B, T, C)
        out = nn.Dense(C)(out)
        return out
[content] utils/nn.py (python)
class STBlock(nn.Module):
    dim: int
    num_heads: int
    dropout: float

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm()(z)
        z = FlashAttention(self.dim, self.num_heads)(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)  # Swap spatial and temporal
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm()(z)
        z = FlashAttention(self.dim, self.num_heads, is_causal=True)(z)
        x = x + z
        x = x.swapaxes(1, 2)  # Swap back

        # --- Feedforward ---
        z = nn.LayerNorm()(x)
        z = nn.Dense(self.dim)(z)
        z = nn.gelu(z)
        x = x + z

        return x
[content] utils/nn.py (python)
class FlashMHA(nn.Module):
    dim: int
    num_heads: int
    causal: bool = False

    def setup(self):
        self.head_dim = self.dim // self.num_heads
        self.q_proj = nn.Dense(self.dim)
        self.k_proj = nn.Dense(self.dim)
        self.v_proj = nn.Dense(self.dim)
        self.out_proj = nn.Dense(self.dim)
        self.flash_attention = jfa.FlashAttention(
            jfa.AttentionConfig(
                platform=jfa.Platform.TRITON,
                backend=jfa.Backend.GPU,
            )
        )

    def __call__(self, x):
        B, T, C = x.shape
        q = self.q_proj(x).reshape(B, T, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
        k = self.k_proj(x).reshape(B, T, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
        v = self.v_proj(x).reshape(B, T, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)

        # Apply FlashAttention
        out = self.flash_attention(query=q, key=k, value=v, causal=self.causal)

        # Combine heads
        out = out.transpose(0, 2, 1, 3).reshape(B, T, C)
        return self.out_proj(out)


class STBlock(nn.Module):
    dim: int
    num_heads: int
    dropout: float

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm()(z)
        z = FlashMHA(dim=self.dim, num_heads=self.num_heads, causal=False)(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm()(z)
        z = FlashMHA(dim=self.dim, num_heads=self.num_heads, causal=True)(z)
        x = x + z
        x = x.swapaxes(1, 2)

        # --- Feedforward ---
        z = nn.LayerNorm()(x)
        z = nn.Dense(self.dim)(z)
        z = nn.gelu(z)
        x = x + z

        return x
[content] utils/nn.py (python)
class FlashMHA(nn.Module):
    dim: int
    num_heads: int
    causal: bool = False

    def setup(self):
        self.head_dim = self.dim // self.num_heads
        self.q_proj = nn.Dense(self.dim)
        self.k_proj = nn.Dense(self.dim)
        self.v_proj = nn.Dense(self.dim)
        self.out_proj = nn.Dense(self.dim)
        self.flash_attention = jfa.FlashAttention(
            jfa.AttentionConfig(
                platform=jfa.Platform.TRITON,
                backend=jfa.Backend.GPU,
            )
        )

    def __call__(self, x):
        orig_shape = x.shape
        if x.ndim != 3:
            # Collapse to 3D: (B*T, N, C) or similar
            x = x.reshape(-1, orig_shape[-2], orig_shape[-1])  # (B*T, N, C)

        B, T, C = x.shape
        q = self.q_proj(x).reshape(B, T, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
        k = self.k_proj(x).reshape(B, T, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
        v = self.v_proj(x).reshape(B, T, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)

        out = self.flash_attention(query=q, key=k, value=v, causal=self.causal)
        out = out.transpose(0, 2, 1, 3).reshape(B, T, C)

        if orig_shape != x.shape:
            # Reshape back to original shape
            out = out.reshape(*orig_shape)

        return self.out_proj(out)
[tab] utils/visualizer.py (python)
import matplotlib.pyplot as plt
import numpy as np
import jax

def visualize_and_save_batch(batch, filename="batch_visualization.png", max_samples=8):
    """
    Visualizes a batch of videos (or images) and saves as a PNG.
    Args:
        batch: jax/numpy array of shape (B, T, H, W, C) or (B, H, W, C)
        filename: output filename
        max_samples: maximum number of batch elements to visualize
    """

    arr = np.asarray(batch)
    if arr.ndim == 5:
        # (B, T, H, W, C): show first frame of each sample
        arr = arr[:, 0]  # (B, H, W, C)
    arr = arr[:max_samples]
    B = arr.shape[0]
    fig, axes = plt.subplots(1, B, figsize=(B * 2.5, 2.5))
    if B == 1:
        axes = [axes]
    for i in range(B):
        axes[i].imshow(arr[i].clip(0, 1))
        axes[i].axis("off")
    plt.tight_layout()
    plt.savefig(filename)
    plt.close(fig)
[tab] scripts_horeka/action_space_scaling/train_lam_actionspace_6.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=36: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_lam_action_space_scaling_6

# 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

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

srun python train_lam.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --min_lr=0 \
    --max_lr=1e-4 \
    --log_image_interval=1000 \
    --log_checkpoint_interval=1000 \
    --log \
    --num_latents=6 \
    --name=lam-action-space-scaling-6-$slurm_job_id \
    --tags lam action-space-scaling 6 \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[tab] scripts_horeka/action_space_scaling/train_lam_actionspace_8.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=36: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_lam_action_space_scaling_8

# 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

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

srun python train_lam.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --min_lr=0 \
    --max_lr=1e-4 \
    --log_image_interval=1000 \
    --log_checkpoint_interval=1000 \
    --log \
    --num_latents=8 \
    --name=lam-action-space-scaling-8-$slurm_job_id \
    --tags lam action-space-scaling 8 \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[tab] scripts_horeka/action_space_scaling/train_lam_actionspace_10.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=36: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_lam_action_space_scaling_10

# 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

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

srun python train_lam.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --min_lr=0 \
    --max_lr=1e-4 \
    --log_image_interval=1000 \
    --log_checkpoint_interval=1000 \
    --log \
    --num_latents=10 \
    --name=lam-action-space-scaling-10-$slurm_job_id \
    --tags lam action-space-scaling 10 \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[tab] scripts_horeka/action_space_scaling/train_lam_actionspace_12.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=36: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_lam_action_space_scaling_12

# 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

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

srun python train_lam.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --min_lr=0 \
    --max_lr=1e-4 \
    --log_image_interval=1000 \
    --log_checkpoint_interval=1000 \
    --log \
    --num_latents=12 \
    --name=lam-action-space-scaling-12-$slurm_job_id \
    --tags lam action-space-scaling 12 \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[tab] scripts_horeka/action_space_scaling/train_lam_actionspace_20.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=36: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_lam_action_space_scaling_20

# 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

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

srun python train_lam.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --min_lr=0 \
    --max_lr=1e-4 \
    --log_image_interval=1000 \
    --log_checkpoint_interval=1000 \
    --log \
    --num_latents=20 \
    --name=lam-action-space-scaling-20-$slurm_job_id \
    --tags lam action-space-scaling 20 \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[tab] scripts_horeka/action_space_scaling/train_lam_actionspace_50.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=36: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_lam_action_space_scaling_50

# 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

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

srun python train_lam.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --min_lr=0 \
    --max_lr=1e-4 \
    --log_image_interval=1000 \
    --log_checkpoint_interval=1000 \
    --log \
    --num_latents=50 \
    --name=lam-action-space-scaling-50-$slurm_job_id \
    --tags lam action-space-scaling 50 \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[content] scripts_horeka/action_space_scaling/start_runs.sh (shellscript)
scripts_horeka/action_space_scaling/train_lam_actionspace_10.sbatch
scripts_horeka/action_space_scaling/train_lam_actionspace_12.sbatch
scripts_horeka/action_space_scaling/train_lam_actionspace_20.sbatch
scripts_horeka/action_space_scaling/train_lam_actionspace_50.sbatch
scripts_horeka/action_space_scaling/train_lam_actionspace_6.sbatch
scripts_horeka/action_space_scaling/train_lam_actionspace_8.sbatch
[tab] scripts_horeka/batchsize_scaling/sqrt_lr/run_sbatch.sh (shellscript)
sbatch /home/hk-project-p0023960/tum_cte0515/Projects/jafar/scripts_horeka/batchsize_scaling/adjusted_lr/train_tokenizer_1_nodes.sbatch
sbatch /home/hk-project-p0023960/tum_cte0515/Projects/jafar/scripts_horeka/batchsize_scaling/adjusted_lr/train_tokenizer_2_nodes.sbatch
sbatch /home/hk-project-p0023960/tum_cte0515/Projects/jafar/scripts_horeka/batchsize_scaling/adjusted_lr/train_tokenizer_4_nodes.sbatch
[tab] scripts_horeka/batchsize_scaling/sqrt_lr/start_runs.sh (shellscript)
sbatch /home/hk-project-p0023960/tum_cte0515/Projects/jafar/scripts_horeka/batchsize_scaling/adjusted_lr/train_tokenizer_1_nodes.sbatch
sbatch /home/hk-project-p0023960/tum_cte0515/Projects/jafar/scripts_horeka/batchsize_scaling/adjusted_lr/train_tokenizer_2_nodes.sbatch
sbatch /home/hk-project-p0023960/tum_cte0515/Projects/jafar/scripts_horeka/batchsize_scaling/adjusted_lr/train_tokenizer_4_nodes.sbatch
[content] scripts_horeka/batchsize_scaling/sqrt_lr/start_runs.sh (shellscript)
scripts_horeka/batchsize_scaling/sqrt_lr/train_tokenizer_16_nodes.sbatch
scripts_horeka/batchsize_scaling/sqrt_lr/train_tokenizer_1_nodes.sbatch
scripts_horeka/batchsize_scaling/sqrt_lr/train_tokenizer_2_nodes.sbatch
scripts_horeka/batchsize_scaling/sqrt_lr/train_tokenizer_32_nodes.sbatch
scripts_horeka/batchsize_scaling/sqrt_lr/train_tokenizer_4_nodes.sbatch
scripts_horeka/batchsize_scaling/sqrt_lr/train_tokenizer_8_nodes.sbatch
[tab] scripts_horeka/batchsize_scaling/sqrt_lr/train_tokenizer_4_nodes.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=4
#SBATCH --ntasks-per-node=4
#SBATCH --time=36: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_tokenizer_batch_size_scaling_4_node

# 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

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

srun python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=192 \
    --min_lr=0 \
    --max_lr=2.00e-4 \
    --log_image_interval=1000 \
    --log_checkpoint_interval=1000 \
    --log \
    --name=tokenizer-batch-size-scaling-4-node-sqrt-lr-$slurm_job_id \
    --tags tokenizer batch-size-scaling 4-node sqrt-lr \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[tab] scripts_horeka/batchsize_scaling/sqrt_lr/train_tokenizer_8_nodes.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=8
#SBATCH --ntasks-per-node=4
#SBATCH --time=36: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_tokenizer_batch_size_scaling_8_node

# 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

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

srun python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=384 \
    --min_lr=0 \
    --max_lr=3e-4 \
    --log_image_interval=1000 \
    --log_checkpoint_interval=1000 \
    --log \
    --name=tokenizer-batch-size-scaling-8-node-sqrt-lr-$slurm_job_id \
    --tags tokenizer batch-size-scaling 8-node sqrt-lr-scaling \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[tab] scripts_horeka/batchsize_scaling/sqrt_lr/train_tokenizer_16_nodes.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=16
#SBATCH --ntasks-per-node=4
#SBATCH --time=36: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_tokenizer_batch_size_scaling_16_node

# 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/'
tf_records_dir=$ws_dir/data/open_ai_minecraft_tfrecord

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=768 \
    --min_lr=0 \
    --max_lr=4e-4 \
    --log_image_interval=1000 \
    --log_checkpoint_interval=1000 \
    --log \
    --name=tokenizer-batch-size-scaling-16-node-sqrt-lr-$slurm_job_id \
    --tags tokenizer batch-size-scaling 16-node sqrt-lr-scaling \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[tab] scripts_horeka/batchsize_scaling/sqrt_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

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

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-sqrt-lr-$slurm_job_id \
    --tags tokenizer batch-size-scaling 32-node sqrt-lr-scaling \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[tab] scripts_horeka/batchsize_scaling/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
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] scripts_horeka/batchsize_scaling/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
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=10 \
    --log_checkpoint_interval=1000 \
    --loggers console local \
    --log_dir local-logs/$slurm_job_id/ \
    --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] 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
torch
triton
[content] /hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/utils/logger.py (python)
            metrics = {k: (v if not hasattr(v, "to_py") and not hasattr(v, "numpy") and not hasattr(v, "__array__") else (v.to_py() if hasattr(v, "to_py") else (v.numpy() if hasattr(v, "numpy") else v.__array__()))) for k, v in metrics.items()}
[content] /hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/utils/logger.py (python)
        log_dir = os.path.join(config["log_dir"], "tb_logger", config["name"])
        self.log_dir = log_dir
        self.writer = SummaryWriter(log_dir=log_dir)

    def log_metrics(self, metrics, step):
        for k, v in metrics.items():
            self.writer.add_scalar(f"metrics/{k}", v, step)

    def log_images(self, images, step):
        for k, v in images.items():
            self.writer.add_image(f"media/{k}", v, step, dataformats='HWC')

class LocalLogger(BaseLogger):
    def __init__(self, config):
        log_dir = os.path.join(config["log_dir"], "local_logger", config["name"])
[content] utils/logger.py (python)
    def __init__(self, cfg):
        pprint(cfg, compact=True)

    def log_metrics(self, metrics, step):
        print(f"[Step {step}] Metrics: " + ", ".join(f"{k}: {v}" for k, v in metrics.items()))
[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
torch
[content] train_tokenizer.py (python)
                    log_images = dict(
                        image=np.asarray(gt_seq[0] * 255.).astype(np.uint8),
                        recon=np.asarray(recon_seq[0] * 255.).astype(np.uint8),
                        true_vs_recon=np.asarray(comparison_seq.astype(np.uint8)
                        ),
                    )
[content] train_lam.py (python)
                    if jax.process_index() == 0:
                        log_images = dict(
                            image=wandb.Image(np.asarray(gt_seq[0])),
                            recon=wandb.Image(np.asarray(recon_seq[0])),
                            true_vs_recon=wandb.Image(
                                np.asarray(comparison_seq.astype(np.uint8))
                            ),
                        )
                        wandb.log(log_images)
[content] train_lam.py (python)
    if args.log and jax.process_index() == 0:
        wandb.init(
            entity=args.entity,
            project=args.project,
            name=args.name,
            tags=args.tags,
            group="debug",
            config=args,
        )
        wandb.config.update({"model_param_count": param_counts})
[content] train_dynamics.py (python)
                    if jax.process_index() == 0:
                        log_images = dict(
                            image=wandb.Image(np.asarray(gt_seq[args.seq_len - 1])),
                            recon=wandb.Image(np.asarray(recon_seq[args.seq_len - 1])),
                            true_vs_recon=wandb.Image(
                                np.asarray(comparison_seq.astype(np.uint8))
                            ),
                        )
                        wandb.log(log_images)
[content] train_dynamics.py (python)
        wandb.init(
            entity=args.entity,
            project=args.project,
            name=args.name,
            tags=args.tags,
            group="debug",
            config=args,
        )
        wandb.config.update({"model_param_count": param_counts})