p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining/linear_lr/train_dynamics_1_nodes.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=48: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_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --job-name=train_dynamics_batch_size_scaling_linear_lr_1_node

# 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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/batchsize_scaling/dynamics_cotraining/linear-lr/$job_name
mkdir -p $CHECKPOINT_DIR

# tokenizer with the new structure supporting larger ffn_dim
tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/checkpoints/train_tokenizer_lr_sweep_1e-4_larger_ffn/

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=48 \
    --init_lr=0 \
    --max_lr=1e-5 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-batch-size-scaling-linear-lr-1-node-$slurm_job_id \
    --tags dynamics batch-size-scaling linear-lr 1-node \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir
[tab] slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/linear_lr/train_dynamics_1_nodes.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=48: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_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --job-name=train_dynamics_batch_size_scaling_linear_lr_1_node

# 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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/batchsize_scaling/dynamics_cotraining/linear-lr/$job_name
mkdir -p $CHECKPOINT_DIR

# tokenizer with the new structure supporting larger ffn_dim
tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/checkpoints/train_tokenizer_lr_sweep_1e-4_larger_ffn/

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=48 \
    --init_lr=0 \
    --max_lr=1e-5 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-batch-size-scaling-linear-lr-1-node-$slurm_job_id \
    --tags dynamics batch-size-scaling linear-lr 1-node \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir
[tab] slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/linear_lr/start_runs.sh (shellscript)
sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining/linear_lr/train_dynamics_16_nodes.sbatch
sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining/linear_lr/train_dynamics_1_nodes.sbatch
sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining/linear_lr/train_dynamics_2_nodes.sbatch
sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining/linear_lr/train_dynamics_4_nodes.sbatch
sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining/linear_lr/train_dynamics_8_nodes.sbatch

# sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining/linear_lr/train_dynamics_32_nodes.sbatch
[tab] slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/sqrt_lr/start_runs.sh (shellscript)
sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining/linear_lr/train_dynamics_16_nodes.sbatch
sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining/linear_lr/train_dynamics_1_nodes.sbatch
sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining/linear_lr/train_dynamics_2_nodes.sbatch
sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining/linear_lr/train_dynamics_4_nodes.sbatch
sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining/linear_lr/train_dynamics_8_nodes.sbatch

# sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining/linear_lr/train_dynamics_32_nodes.sbatch
[tab] slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/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_chunked

job_name="debug"
slurm_job_id="0000"

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/checkpoints/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

# tokenizer with the new structure supporting larger ffn_dim
tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/checkpoints/train_tokenizer_lr_sweep_1e-4_larger_ffn/

env | grep SLURM

srun python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --init_lr=0 \
    --max_lr=3e-4 \
    --num_latent_actions=20 \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
[tab] slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/linear_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_chunked

job_name="debug"
slurm_job_id="0000"

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/checkpoints/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

# tokenizer with the new structure supporting larger ffn_dim
tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/checkpoints/train_tokenizer_lr_sweep_1e-4_larger_ffn/

env | grep SLURM

srun python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --init_lr=0 \
    --max_lr=3e-4 \
    --num_latent_actions=20 \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
[tab] slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/linear_lr/train_dynamics_2_nodes.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=48: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_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --job-name=train_dynamics_batch_size_scaling_linear_lr_2_node

# 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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/batchsize_scaling/dynamics_cotraining/linear-lr/$job_name
mkdir -p $CHECKPOINT_DIR

# tokenizer with the new structure supporting larger ffn_dim
tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/checkpoints/train_tokenizer_lr_sweep_1e-4_larger_ffn/

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --init_lr=0 \
    --max_lr=1.5e-5 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-batch-size-scaling-linear-lr-2-node-$slurm_job_id \
    --tags dynamics batch-size-scaling linear-lr 2-node \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
[tab] slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/linear_lr/train_dynamics_4_nodes.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=4
#SBATCH --ntasks-per-node=4
#SBATCH --time=48: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_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --job-name=train_dynamics_batch_size_scaling_linear_lr_4_node

# 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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/batchsize_scaling/dynamics_cotraining/linear-lr/$job_name
mkdir -p $CHECKPOINT_DIR

# tokenizer with the new structure supporting larger ffn_dim
tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/checkpoints/train_tokenizer_lr_sweep_1e-4_larger_ffn/

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=192 \
    --init_lr=0 \
    --max_lr=2.1e-5 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-batch-size-scaling-linear-lr-4-node-$slurm_job_id \
    --tags dynamics batch-size-scaling linear-lr 4-node \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir
[tab] slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/linear_lr/train_dynamics_8_nodes.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=8
#SBATCH --ntasks-per-node=4
#SBATCH --time=48: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_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --job-name=train_dynamics_batch_size_scaling_linear_lr_8_node

# 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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/batchsize_scaling/dynamics_cotraining/linear-lr/$job_name
mkdir -p $CHECKPOINT_DIR

# tokenizer with the new structure supporting larger ffn_dim
tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/checkpoints/train_tokenizer_lr_sweep_1e-4_larger_ffn/

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=384 \
    --init_lr=0 \
    --max_lr=3e-5 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-batch-size-scaling-linear-lr-8-node-$slurm_job_id \
    --tags dynamics batch-size-scaling linear-lr 8-node \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir
[tab] slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/linear_lr/train_dynamics_16_nodes.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=16
#SBATCH --ntasks-per-node=4
#SBATCH --time=48: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_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --job-name=train_dynamics_batch_size_scaling_linear_lr_16_node

# 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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/batchsize_scaling/dynamics_cotraining/linear-lr/$job_name
mkdir -p $CHECKPOINT_DIR

# tokenizer with the new structure supporting larger ffn_dim
tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/checkpoints/train_tokenizer_lr_sweep_1e-4_larger_ffn/

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=768 \
    --init_lr=0 \
    --max_lr=4.2e-5 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-batch-size-scaling-linear-lr-16-node-$slurm_job_id \
    --tags dynamics batch-size-scaling linear-lr 16-node \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir
[tab] slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/sqrt_lr/train_dynamics_2_nodes.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=48: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_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --job-name=train_dynamics_batch_size_scaling_linear_lr_2_node

# 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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/batchsize_scaling/dynamics_cotraining/linear-lr/$job_name
mkdir -p $CHECKPOINT_DIR

# tokenizer with the new structure supporting larger ffn_dim
tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/checkpoints/train_tokenizer_lr_sweep_1e-4_larger_ffn/

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --init_lr=0 \
    --max_lr=1.5e-5 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-batch-size-scaling-linear-lr-2-node-$slurm_job_id \
    --tags dynamics batch-size-scaling linear-lr 2-node \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_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
    ffn_dim: int
    num_latents: int
    num_blocks: int
    num_heads: int
    dropout: float
    mask_limit: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool
    use_gt_actions: bool = False

    def setup(self):
        self.dynamics = STTransformer(
            self.model_dim,
            self.ffn_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
        )
        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),
        )
        if self.use_gt_actions:
            self.action_embed = nn.Embed(self.num_actions, self.model_dim)
        else:
            self.action_up = nn.Dense(
                self.model_dim,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
            )

    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)
            vid_embed = jnp.where(jnp.expand_dims(mask, -1), self.mask_token, vid_embed)
        else:
            mask = None

        # --- Predict transition ---
        if self.use_gt_actions:
            act_embed = self.action_embed(batch["actions"])
        else:
            act_embed = self.action_up(batch["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)
[content] train_dynamics.py (python)
                actions = jax.random.randint(
                    _rng, 
                    shape=(videos.shape[0], videos.shape[1]), 
                    minval=0, 
                    maxval=args.num_actions,
                    dtype=jnp.int32
                )
[content] train_dynamics.py (python)
                    args=ocp.args.Composite(
                        model_state=ocp.args.StandardSave(train_state),
                        dataloader_state=grain.checkpoint.CheckpointSave(
                            grain_iterator
                        ),
[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
    ffn_dim: int
    num_latents: int
    num_blocks: int
    num_heads: int
    dropout: float
    mask_limit: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool
    use_gt_actions: bool = False
    num_actions: int = 6

    def setup(self):
        self.dynamics = STTransformer(
            self.model_dim,
            self.ffn_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
        )
        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),
        )
        if self.use_gt_actions:
            self.action_embed = nn.Embed(self.num_actions, self.model_dim)
        else:
            self.action_up = nn.Dense(
                self.model_dim,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
            )

    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)
            vid_embed = jnp.where(jnp.expand_dims(mask, -1), self.mask_token, vid_embed)
        else:
            mask = None

        # --- Predict transition ---
        if self.use_gt_actions:
            act_embed = self.action_embed(batch["actions"])
        else:
            act_embed = self.action_up(batch["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/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/sqrt_lr/start_runs.sh (shellscript)
sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/sqrt_lr/train_dynamics_16_nodes.sbatch
sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/sqrt_lr/train_dynamics_1_nodes.sbatch
sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/sqrt_lr/train_dynamics_2_nodes.sbatch
sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/sqrt_lr/train_dynamics_4_nodes.sbatch
sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/sqrt_lr/train_dynamics_8_nodes.sbatch

# sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/sqrt_lr/train_dynamics_32_nodes.sbatch
[tab] slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/linear_lr/start_runs.sh (shellscript)
sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/linear_lr/train_dynamics_16_nodes.sbatch
sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/linear_lr/train_dynamics_1_nodes.sbatch
sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/linear_lr/train_dynamics_2_nodes.sbatch
sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/linear_lr/train_dynamics_4_nodes.sbatch
sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/linear_lr/train_dynamics_8_nodes.sbatch

# sbatch slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/linear_lr/train_dynamics_32_nodes.sbatch
[tab] slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/linear_lr/train_dynamics_2_nodes.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=48: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_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --job-name=train_dynamics_batch_size_scaling_linear_lr_2_node

# 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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/batchsize_scaling/dynamics_cotraining/linear-lr/$job_name
mkdir -p $CHECKPOINT_DIR

# tokenizer with the new structure supporting larger ffn_dim
tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/checkpoints/train_tokenizer_lr_sweep_1e-4_larger_ffn/

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --init_lr=0 \
    --max_lr=2e-5 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-batch-size-scaling-linear-lr-2-node-$slurm_job_id \
    --tags dynamics batch-size-scaling linear-lr 2-node \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
[tab] slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/linear_lr/train_dynamics_4_nodes.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=4
#SBATCH --ntasks-per-node=4
#SBATCH --time=48: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_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --job-name=train_dynamics_batch_size_scaling_linear_lr_4_node

# 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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/batchsize_scaling/dynamics_cotraining/linear-lr/$job_name
mkdir -p $CHECKPOINT_DIR

# tokenizer with the new structure supporting larger ffn_dim
tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/checkpoints/train_tokenizer_lr_sweep_1e-4_larger_ffn/

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=192 \
    --init_lr=0 \
    --max_lr=4e-5 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-batch-size-scaling-linear-lr-4-node-$slurm_job_id \
    --tags dynamics batch-size-scaling linear-lr 4-node \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir
[tab] slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/linear_lr/train_dynamics_8_nodes.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=8
#SBATCH --ntasks-per-node=4
#SBATCH --time=48: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_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --job-name=train_dynamics_batch_size_scaling_linear_lr_8_node

# 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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/batchsize_scaling/dynamics_cotraining/linear-lr/$job_name
mkdir -p $CHECKPOINT_DIR

# tokenizer with the new structure supporting larger ffn_dim
tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/checkpoints/train_tokenizer_lr_sweep_1e-4_larger_ffn/

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=384 \
    --init_lr=0 \
    --max_lr=8e-5 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-batch-size-scaling-linear-lr-8-node-$slurm_job_id \
    --tags dynamics batch-size-scaling linear-lr 8-node \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir
[tab] slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/linear_lr/train_dynamics_16_nodes.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=16
#SBATCH --ntasks-per-node=4
#SBATCH --time=48: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_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --job-name=train_dynamics_batch_size_scaling_linear_lr_16_node

# 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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/batchsize_scaling/dynamics_cotraining/linear-lr/$job_name
mkdir -p $CHECKPOINT_DIR

# tokenizer with the new structure supporting larger ffn_dim
tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/checkpoints/train_tokenizer_lr_sweep_1e-4_larger_ffn/

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=768 \
    --init_lr=0 \
    --max_lr=16e-5 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-batch-size-scaling-linear-lr-16-node-$slurm_job_id \
    --tags dynamics batch-size-scaling linear-lr 16-node \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir
[tab] slurm/jobs/franz/horeka/batchsize_scaling/dynamics_cotraining_new_arch/sqrt_lr/train_dynamics_1_nodes.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=48: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_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --job-name=train_dynamics_batch_size_scaling_linear_lr_1_node

# 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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/batchsize_scaling/dynamics_cotraining/linear-lr/$job_name
mkdir -p $CHECKPOINT_DIR

# tokenizer with the new structure supporting larger ffn_dim
tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/checkpoints/train_tokenizer_lr_sweep_1e-4_larger_ffn/

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=48 \
    --init_lr=0 \
    --max_lr=1e-5 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-batch-size-scaling-linear-lr-1-node-$slurm_job_id \
    --tags dynamics batch-size-scaling linear-lr 1-node \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir