p-doom/crowd-code-dataset-0.2 huggingface dataset:
[terminal_command] TERMINAL ()
mv /home/hk-project-p0023960/tum_dbd0378/jafar/checkpoints/train_tokenizer_lr_sweep_1e-4_larger_ffn  /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/train_tokenizer_lr_sweep_1e-4_larger_ffn
[terminal_command] TERMINAL ()
ls /home/hk-project-p0023960/tum_dbd0378/jafar/checkpoints/train_tokenizer_lr_sweep_1e-4train_tokenizer_lr_sweep_1e-4/                                 train_tokenizer_lr_sweep_1e-4_larger_ffn_5_blocks/             train_tokenizer_lr_sweep_1e-4_mixed_precision_flash_attention/ (jafar) [tum_dbd0378@hkn1993 jafar]$ ls /home/hk-project-p0023960/tum_dbd0378/jafar/checkpoints/train_tokenizer_lr_sweep_1e-4_mixed_precision_flash_attention/
[tab] experiments/train_tokenizer_lr_1e-4_mixed_precision.sbatch (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

SLURM_JOB_NAME=train_tokenizer_lr_sweep_1e-4_mixed_precision_flash_attention

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 \
    --init_lr=0 \
    --max_lr=1e-4 \
    --log_image_interval=1000 \
    --log_checkpoint_interval=1000 \
    --log \
    --name=mixed-precision-tokenizer-lr-1e-4-flash-attention \
    --tags tokenizer mixed-precision 1e-4 flash-attention \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir
[terminal_command] TERMINAL ()
sacct --format="JobID%15,JobName%30,Partition%16,AllocCPUS%3,State%12,Elapsed%10,Timelimit%10" --starttime $(date -d "last week" +%Y-%m-%d) | grep -vE ".batch|.extern|*.inter|bash|python|CANCELLED|echo

^C
[terminal_command] TERMINAL ()
sacct --format="JobID%15,JobName%30,Partition%16,AllocCPUS%3,State%12,Elapsed%10,Timelimit%10" --starttime $(date -d "last week" +%Y-%m-%d) | grep -vE ".batch|.extern|*.inter|bash|python|CANCELLED|echo|
^C
[terminal_command] TERMINAL ()
sacct --format="JobID%15,JobName%30,Partition%16,AllocCPUS%3,State%12,Elapsed%10,Timelimit%10" --starttime $(date -d "last week" +%Y-%m-%d) | grep -vE ".batch|.extern|*.inter|bash|python|CANCELLED|echo"
[terminal_command] TERMINAL ()
sacct --format="JobID%15,JobName%30,Partition%16,AllocCPUS%3,State%12,Elapsed%10,Timelimit%10" --starttime $(date -d "last week" +%Y-%m-%d) | grep -vE ".batch|.extern|interactive" | grep "accelerated"
[content] utils/nn.py (python)
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
            rngs=rngs
        )
        
        # Feedforward components
        self.ffn_norm = nnx.LayerNorm(param_dtype=self.param_dtype, dtype=self.dtype, rngs=rngs)
        self.ffn_dense1 = nnx.Linear(self.ffn_dim, param_dtype=self.param_dtype, dtype=self.dtype, rngs=rngs)
        self.ffn_dense2 = nnx.Linear(self.dim, param_dtype=self.param_dtype, dtype=self.dtype, rngs=rngs)
[content] utils/nn.py (python)
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=False),
            rngs=rngs
        )
        
        # Temporal attention components
        self.temporal_pos_enc = PositionalEncoding(self.dim, rngs=rngs)
        self.temporal_norm = nnx.LayerNorm(param_dtype=self.param_dtype, dtype=self.dtype, rngs=rngs)
        self.temporal_attention = nnx.MultiHeadAttention(
[content] utils/nn.py (python)
class STBlock(nnx.Module):
    def __init__(self, dim: int, ffn_dim: int, num_heads: int, dropout: float, 
                 param_dtype: jnp.dtype, dtype: jnp.dtype, use_flash_attention: bool, rngs: nnx.Rngs):
        self.dim = dim
        self.ffn_dim = ffn_dim
        self.num_heads = num_heads
        self.dropout = dropout
        self.param_dtype = param_dtype
        self.dtype = dtype
        self.use_flash_attention = use_flash_attention
        
        # Spatial attention components
        self.spatial_pos_enc = PositionalEncoding(self.dim, rngs=rngs)
        self.spatial_norm = nnx.LayerNorm(param_dtype=self.param_dtype, dtype=self.dtype, rngs=rngs)
        self.spatial_attention = nnx.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
[content] utils/nn.py (python)
    def __init__(self, d_model: int, max_len: int = 5000, rngs: nnx.Rngs):
        # Hidden dimensionality of the input.
        self.d_model = d_model
        # Maximum length of a sequence to expect.
        self.max_len = max_len
[content] utils/nn.py (python)
class VectorQuantizer(nnx.Module):
    def __init__(self, latent_dim: int, num_latents: int, dropout: float, rngs: nnx.Rngs):
        self.latent_dim = latent_dim
        self.num_latents = num_latents
        self.dropout = dropout
        
        # Initialize codebook
        self.codebook = nnx.Param(
            normalize(
                jax.random.lecun_uniform(rngs.params(), (self.num_latents, self.latent_dim))
[content] utils/nn.py (python)
                rngs=rngs
            ))
        
        # Output projection
        self.output_dense = nnx.Linear(self.out_dim, param_dtype=self.param_dtype, dtype=self.dtype, rngs=rngs)

    def __call__(self, x: jax.Array) -> jax.Array:
        # Input projection
        x = self.input_norm1(x)
        x = self.input_dense(x)
        x = self.input_norm2(x)
        
        # Transformer blocks
        for block in self.blocks:
            x = block(x)
            
        # Output projection
        x = self.output_dense(x)
[content] utils/nn.py (python)
class STTransformer(nnx.Module):
    def __init__(self, model_dim: int, ffn_dim: int, out_dim: int, num_blocks: int, 
                 num_heads: int, dropout: 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.out_dim = out_dim
        self.num_blocks = num_blocks
        self.num_heads = num_heads
        self.dropout = dropout
        self.param_dtype = param_dtype
        self.dtype = dtype
        self.use_flash_attention = use_flash_attention
        
        # Input projection
        self.input_norm1 = nnx.LayerNorm(param_dtype=self.param_dtype, dtype=self.dtype, rngs=rngs)
        self.input_dense = nnx.Linear(self.model_dim, param_dtype=self.param_dtype, dtype=self.dtype, rngs=rngs)
        self.input_norm2 = nnx.LayerNorm(param_dtype=self.param_dtype, dtype=self.dtype, rngs=rngs)
        
        # Transformer blocks
        self.blocks = []
[content] utils/nn.py (python)

    @nnx.remat
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = self.spatial_pos_enc(x)
        z = self.spatial_norm(z)
        z = self.spatial_attention(z)
[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, 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.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.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(
            jax.random.uniform(rngs.params(), (1, 1, 1, self.model_dim), minval=-1.0, maxval=1.0)
        )
        self.action_up = nnx.Linear(
            self.num_latents, 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:
            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.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] genie.py (python)
        scanned_maskgit_loop = nnx.scan(
            maskgit_loop_fn,
            in_axes=(nnx.Carry, None, 0),
            out_axes=(nnx.Carry, 0),
            length=steps # Explicitly specify the length for JAX
        )
[content] train_tokenizer.py (python)
    if args.log_gradients:
        metrics["encoder_gradients_std/"] = jax.tree.map(
            lambda x: x.std(), grads["params"]["encoder"]
        )
        metrics["vq_gradients_std/"] = jax.tree.map(
            lambda x: x.std(), grads["params"]["vq"]
        )
        metrics["decoder_gradients_std/"] = jax.tree.map(
            lambda x: x.std(), grads["params"]["decoder"]
        )
[content] train_lam.py (python)
    # --- Reset inactive latent actions ---
    codebook = state.params["params"]["vq"]["codebook"]
    num_codes = len(codebook)
    active_codes = idx_counts != 0.0
    action_last_active = jnp.where(active_codes, 0, action_last_active + 1)
    p_code = active_codes / active_codes.sum()
    reset_idxs = jax.random.choice(rng, num_codes, shape=(num_codes,), p=p_code)
    do_reset = action_last_active >= args.vq_reset_thresh
    new_codebook = jnp.where(
        jnp.expand_dims(do_reset, -1), codebook[reset_idxs], codebook
    )
    state.params["params"]["vq"]["codebook"] = new_codebook
    action_last_active = jnp.where(do_reset, 0, action_last_active)
[content] train_lam.py (python)
            
            # Reconstruct model from TrainState for training
            model = nnx.merge(graphdef, train_state.params, train_state.other_variables)
            model, loss, recon, metrics = train_step(model, inputs, action_last_active)
            
            # Update TrainState
            _, new_params, new_other_variables = nnx.split(model, nnx.Param, ...)
            train_state = train_state.replace(
                params=new_params,
                other_variables=new_other_variables
            )
[content] train_lam.py (python)
    
    # Create TrainState with NNX model
    graphdef, params, other_variables = nnx.split(lam, nnx.Param, ...)
    
    class CustomTrainState(TrainState):
        other_variables: nnx.State
    
    train_state = CustomTrainState.create(
        apply_fn=graphdef.apply,
        params=params,
        other_variables=other_variables,
        tx=tx
    )
[content] train_lam.py (python)
    grads = nnx.grad(loss_fn)(model)
    _, params, rest = nnx.split(model, nnx.Param, ...)
    params = jax.tree.map(lambda p, g: p - 0.1 * g, params, grads)
    nnx.update(model, nnx.merge(params, rest))
    
    # Get loss for logging
    loss, (recon, metrics) = lam_loss_fn(model, inputs)
    return model, loss, recon, metrics
[content] train_lam.py (python)
    
    # Update the model's codebook
    new_params = params.copy()
    new_params["vq"]["codebook"] = new_codebook
    nnx.update(lam, nnx.merge(lam.graphdef, new_params, lam.state))
    
    # Reset action_last_active for codes that were reset
[content] train_lam.py (python)
    
    # Calculate reset probabilities
    p_code = active_codes / (active_codes.sum() + 1e-8)  # Add small epsilon to avoid division by zero
    
    # Generate reset indices
    rng, reset_rng = jax.random.split(rng)
    reset_idxs = jax.random.choice(reset_rng, num_codes, shape=(num_codes,), p=p_code)
    
    # Determine which codes to reset
[content] train_lam.py (python)
    # Get the current model outputs to check codebook usage
    outputs = lam(inputs, training=False)
    count_fn = jax.vmap(lambda i: (outputs["indices"] == i).sum())
    index_counts = count_fn(jnp.arange(args.num_latents))
    
    # Get codebook from the model
    _, params, _ = nnx.split(lam, nnx.Param, ...)
    codebook = params["vq"]["codebook"]
[content] train_lam.py (python)
    # --- Update model ---
    (loss, (recon, metrics)), grads = nnx.value_and_grad(loss_fn, has_aux=True)(lam)
    optimizer.update(grads)

    # --- Reset inactive latent actions ---
    # Get the current model outputs to check codebook usage
    outputs = lam(inputs, training=False)
    count_fn = jax.vmap(lambda i: (outputs["indices"] == i).sum())
    index_counts = count_fn(jnp.arange(args.num_latents))
    
    # Get codebook from the model
    _, params, _ = nnx.split(lam, nnx.Param, ...)
    codebook = params["vq"]["codebook"]
    num_codes = len(codebook)
    active_codes = index_counts != 0.0
    
    # Update action_last_active tracking
    action_last_active = jnp.where(active_codes, 0, action_last_active + 1)
    
    # Calculate reset probabilities
    p_code = active_codes / (active_codes.sum() + 1e-8)  # Add small epsilon to avoid division by zero
    
    # Generate reset indices
    rng, reset_rng = jax.random.split(rng)
    reset_idxs = jax.random.choice(reset_rng, num_codes, shape=(num_codes,), p=p_code)
    
    # Determine which codes to reset
    do_reset = action_last_active >= args.vq_reset_thresh
    
    # Apply reset
    new_codebook = jnp.where(
        jnp.expand_dims(do_reset, -1), codebook[reset_idxs], codebook
    )
    
    # Update the model's codebook
    new_params = params.copy()
    new_params["vq"]["codebook"] = new_codebook
    nnx.update(lam, nnx.merge(lam.graphdef, new_params, lam.state))
    
    # Reset action_last_active for codes that were reset
    action_last_active = jnp.where(do_reset, 0, action_last_active)
[content] train_lam.py (python)
    grads = nnx.grad(loss_fn)(model)
    _, params, rest = nnx.split(model, nnx.Param, ...)
    params = jax.tree.map(lambda p, g: p - 0.1 * g, params, grads)
    nnx.update(model, nnx.merge(params, rest))
[content] train_lam.py (python)
    # For now, we'll skip the codebook reset logic since it requires more complex
    # parameter manipulation that doesn't work well with the current NNX API
    # TODO: Implement codebook reset using a different approach
[content] train_lam.py (python)
    # Calculate reset probabilities
    p_code = active_codes / (active_codes.sum() + 1e-8)  # Add small epsilon to avoid division by zero
    
    # Generate reset indices
    rng, reset_rng = jax.random.split(rng)
    reset_idxs = jax.random.choice(reset_rng, num_codes, shape=(num_codes,), p=p_code)
    
    # Determine which codes to reset
    do_reset = action_last_active >= args.vq_reset_thresh
    
    # Apply reset
    new_codebook = jnp.where(
        jnp.expand_dims(do_reset, -1), codebook[reset_idxs], codebook
    )
    
    # Update the model's codebook
    new_params = params.copy()
    new_params["vq"]["codebook"] = new_codebook
    nnx.update(lam, nnx.merge(lam.graphdef, new_params, lam.state))
    
    # Reset action_last_active for codes that were reset
    action_last_active = jnp.where(do_reset, 0, action_last_active)
[content] train_lam.py (python)
    # For now, we'll skip the codebook reset logic since it requires more complex
    # parameter manipulation that doesn't work well with the current NNX API
    # TODO: Implement codebook reset using a different approach
    
    return loss, recon, action_last_active, metrics, rng
[content] train_lam.py (python)
    # --- Update model ---
    (loss, (recon, metrics)), grads = nnx.value_and_grad(loss_fn, has_aux=True)(lam)
    optimizer.update(grads)

    # --- Reset inactive latent actions ---
    # Get the current model outputs to check codebook usage
    outputs = lam(inputs, training=False)
    count_fn = jax.vmap(lambda i: (outputs["indices"] == i).sum())
    index_counts = count_fn(jnp.arange(args.num_latents))
    
    # Update action_last_active tracking
    active_codes = index_counts != 0.0
    action_last_active = jnp.where(active_codes, 0, action_last_active + 1)
[content] train_lam.py (python)
    # --- Update model ---
    (loss, (recon, metrics)), grads = nnx.value_and_grad(loss_fn, has_aux=True)(lam)
    optimizer.update(grads)

    # --- Reset inactive latent actions ---
    # Get the current model outputs to check codebook usage
    outputs = lam(inputs, training=False)
    count_fn = jax.vmap(lambda i: (outputs["indices"] == i).sum())
    index_counts = count_fn(jnp.arange(args.num_latents))
    
    # Update action_last_active tracking
    active_codes = index_counts != 0.0
    action_last_active = jnp.where(active_codes, 0, action_last_active + 1)
    
    # For now, we'll skip the codebook reset logic since it requires more complex
    # parameter manipulation that doesn't work well with the current NNX API
    # TODO: Implement codebook reset using a different approach
[content] train_lam.py (python)
    reset_idxs = jax.random.choice(rng, num_codes, shape=(num_codes,), p=p_code)
    do_reset = action_last_active >= args.vq_reset_thresh
    new_codebook = jnp.where(
        jnp.expand_dims(do_reset, -1), codebook[reset_idxs], codebook
    )