p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] experiments/sample.sh (shellscript)

export SLURM_JOB_NAME=tokenizer-no-remat

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=$PWD/checkpoints/$job_name
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

srun python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --ffn_dim=2048 \
    --num_blocks=4 \
    --init_lr=0 \
    --max_lr=1e-4 \
    --log_image_interval=1000 \
    --log_checkpoint_interval=1000 \
    --data_dir $array_records_dir
[tab] experiments/train_dynamics_lr_1e-4_larger_ffn.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=24: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/larger_ffn_dynamics/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/big-runs/larger_ffn_dynamics/%x_%j.log
#SBATCH --job-name=train_dynamics_lr_1e-4_larger_ffn

# 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/dynamics-lrarge-ffn/$job_name
mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaling/train_tokenizer_lr_sweep_1e-4

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --tokenizer_ffn_dim=2048 \
    --tokenizer_num_blocks=4 \
    --lam_ffn_dim=2048 \
    --lam_num_blocks=4 \
    --dyna_ffn_dim=2048 \
    --dyna_num_blocks=6 \
    --init_lr=0 \
    --max_lr=1.5e-4 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-lr-1e-4-larger-ffn \
    --tags dynamics lr-1e-4 larger-ffn mixed-precision flash-attention \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
[content] experiments/train_tokenizer_nnx.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=24: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/larger_ffn_dynamics/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/big-runs/larger_ffn_dynamics/%x_%j.log
#SBATCH --job-name=train_dynamics_lr_1e-4_larger_ffn

# 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/dynamics-lrarge-ffn/$job_name
mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaling/train_tokenizer_lr_sweep_1e-4

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --tokenizer_ffn_dim=2048 \
    --tokenizer_num_blocks=4 \
    --lam_ffn_dim=2048 \
    --lam_num_blocks=4 \
    --dyna_ffn_dim=2048 \
    --dyna_num_blocks=6 \
    --init_lr=0 \
    --max_lr=1.5e-4 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-lr-1e-4-larger-ffn \
    --tags dynamics lr-1e-4 larger-ffn mixed-precision flash-attention \
    --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.nnx as nnx

from utils.nn import STTransformer


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

    def __init__(
        self,
        model_dim: int,
        ffn_dim: int,
        num_latents: int,
        latent_action_dim: int,
        num_blocks: int,
        num_heads: int,
        dropout: float,
        mask_limit: float,
        param_dtype: jnp.dtype,
        dtype: jnp.dtype,
        use_flash_attention: bool,
        rngs: nnx.Rngs,
    ):
        self.model_dim = model_dim
        self.ffn_dim = ffn_dim
        self.num_latents = num_latents
        self.latent_action_dim = latent_action_dim
        self.num_blocks = num_blocks
        self.num_heads = num_heads
        self.dropout = dropout
        self.mask_limit = mask_limit
        self.param_dtype = param_dtype
        self.dtype = dtype
        self.use_flash_attention = use_flash_attention

        self.dynamics = STTransformer(
            self.model_dim,
            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,
            rngs=rngs,
        )
        self.patch_embed = nnx.Embed(self.num_latents, self.model_dim, rngs=rngs)
        self.mask_token = nnx.Param(
            nnx.initializers.lecun_uniform()(rngs.params(), (1, 1, 1, self.model_dim))
        )
        self.action_up = nnx.Linear(
            self.latent_action_dim,
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )

    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:
            batch_size = vid_embed.shape[0]
            _rng_prob, *_rngs_mask = jax.random.split(batch["mask_rng"], batch_size + 1)
            mask_prob = jax.random.uniform(
                _rng_prob, shape=(batch_size,), minval=self.mask_limit
            )
            per_sample_shape = vid_embed.shape[1:-1]
            mask = jax.vmap(
                lambda rng, prob: jax.random.bernoulli(rng, prob, per_sample_shape),
                in_axes=(0, 0),
            )(jnp.asarray(_rngs_mask), mask_prob)
            mask = mask.at[:, 0].set(False)
            vid_embed = jnp.where(
                jnp.expand_dims(mask, -1), self.mask_token.value, 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)
[content] train_tokenizer.py (python)
                        dict(
                            model_state=ocp.args.PyTreeSave(optimizer_state),
                            dataloader_state=grain.checkpoint.CheckpointSave(
                                grain_iterator
                            ),
                        )
[content] train_tokenizer.py (python)
                dict(
                    model_state=ocp.args.PyTreeRestore(abstract_optimizer_state),
                    dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),
                )
[content] train_tokenizer.py (python)
                        model_state=ocp.args.PyTreeSave(optimizer_state),
                        dataloader_state=grain.checkpoint.CheckpointSave(
                            grain_iterator
                        ),
[content] /hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/train_tokenizer.py (python)

    # model_state = nnx.state(optimizer.model)
    # model_shardings = nnx.get_named_sharding(model_state, mesh)
    # model_sharded_state = jax.lax.with_sharding_constraint(model_state, model_shardings)
    # nnx.update(optimizer.model, model_sharded_state)
    # optimizer_state = nnx.state(optimizer, nnx.optimizer.OptState)
    # optimizer_shardings = nnx.get_named_sharding(optimizer_state, mesh)
    # optimizer_sharded_state = jax.lax.with_sharding_constraint(optimizer_state, optimizer_shardings)
    # nnx.update(optimizer, optimizer_sharded_state)
[content] train_tokenizer.py (python)
    model_state = nnx.state(optimizer.model)
    model_sharded_state = jax.device_put(model_state, replicated_sharding)
    nnx.update(optimizer.model, model_sharded_state)
    optimizer_state = nnx.state(optimizer, nnx.optimizer.OptState)
    optimizer_sharded_state = jax.device_put(optimizer_state, replicated_sharding)
    nnx.update(optimizer, optimizer_sharded_state)

    # --- Initialize checkpoint manager ---
    step = 0
    handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
    handler_registry.add(
        "model_state", ocp.args.PyTreeSave, ocp.handlers.PyTreeCheckpointHandler
    )
    handler_registry.add(
        "model_state", ocp.args.PyTreeRestore, ocp.handlers.PyTreeCheckpointHandler
    )
    handler_registry.add("dataloader_state", grain.checkpoint.CheckpointSave, grain.checkpoint.CheckpointHandler)  # type: ignore
    handler_registry.add("dataloader_state", grain.checkpoint.CheckpointRestore, grain.checkpoint.CheckpointHandler)  # type: ignore

    checkpoint_options = ocp.CheckpointManagerOptions(
        save_interval_steps=args.log_checkpoint_interval,
        max_to_keep=3,
        keep_period=args.log_checkpoint_keep_period,
        step_format_fixed_length=6,
        cleanup_tmp_directories=True,
    )

    checkpoint_manager = ocp.CheckpointManager(
        args.ckpt_dir,
        options=checkpoint_options,
        handler_registry=handler_registry,
    )

    # --- Create DataLoaderIterator from dataloader ---
    image_shape = (args.image_height, args.image_width, args.image_channels)
    array_record_files = [
        os.path.join(args.data_dir, x)
        for x in os.listdir(args.data_dir)
        if x.endswith(".array_record")
    ]
    grain_dataloader = get_dataloader(
        array_record_files,
        args.seq_len,
        # NOTE: We deliberately pass the global batch size
        # The dataloader shards the dataset across all processes
        args.batch_size,
        *image_shape,
        num_workers=8,
        prefetch_buffer_size=1,
        seed=args.seed,
    )
    initial_state = grain_dataloader._create_initial_state()
    grain_iterator = grain.DataLoaderIterator(grain_dataloader, initial_state)

    # --- Restore checkpoint ---
    if args.restore_ckpt:
        # FIXME (f.srambical): change all checkpointing/ restore logic to separately checkpoint model and optimizer
        abstract_optimizer = nnx.eval_shape(lambda: optimizer)
        abstract_optimizer_state = nnx.state(abstract_optimizer)
        # FIXME (f.srambical): also checkpoint optimizer state
[content] train_tokenizer.py (python)
    model_sharded_state = jax.lax.with_sharding_constraint(model_state, replicated_sharding)
    nnx.update(optimizer.model, model_sharded_state)
    optimizer_state = nnx.state(optimizer, nnx.optimizer.OptState)
    optimizer_sharded_state = jax.lax.with_sharding_constraint(optimizer_state, replicated_sharding)
[content] experiments/train_dynamics_nnx.sh (shellscript)
#!/usr/bin/env bash

# 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

CHECKPOINT_DIR=$PWD/checkpoints/tokenizer_nnx
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

srun python train_tokenizer.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=48 \
    --log_checkpoint_interval=5 \
    --data_dir $array_records_dir \
[content] experiments/train_dynamics_nnx.sh (shellscript)
srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --tokenizer_ffn_dim=2048 \
    --tokenizer_num_blocks=4 \
    --lam_ffn_dim=2048 \
    --lam_num_blocks=4 \
    --dyna_ffn_dim=2048 \
    --dyna_num_blocks=6 \
    --init_lr=0 \
    --max_lr=1.5e-4 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-lr-1e-4-larger-ffn \
    --tags dynamics lr-1e-4 larger-ffn mixed-precision flash-attention \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
[tab] requirements.txt (pip-requirements)
dm_pix>=0.4.3
einops>=0.8.0
flax>=0.8.5
jax[cuda12]>=0.6.2
optax>=0.2.3
procgen>=0.10.7
tyro>=0.8.5
wandb>=0.15.0
tensorflow>=2.0.0
tensorboard>=2.0.0
pandas>=1.3.0
pathlib2>=2.3.0; python_version < "3.4"
grain>=0.2.10
pre-commit>=4.2.0
array-record>=0.7.2
tqdm>=4.67.1