p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] experiments/tokenizer_grain_checkpointing.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=8
#SBATCH --time=12:00:00
#SBATCH --cpus-per-task=1
#SBATCH --gres=gpu:8
#SBATCH --mem=100GB

source .venv/bin/activate

data_dir="$PWD/data_arrayrecord/dummy"
ckpt_dir="$PWD/checkpoints/tokenizer_openai_grain_checkpointing"

mkdir -p $ckpt_dir

export XLA_FLAGS=--xla_gpu_autotune_level=0
export PYTHONUNBUFFERED=1
srun python train_tokenizer.py \
    --batch_size 96 \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --init_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[content] dynamics.py (python)
            # Insert the generated tokens into the sequence
            token_idxs_full = token_idxs_full.at[:, t, n].set(next_token)
            
            new_carry = (rng, token_idxs_full, action_tokens)
            return new_carry, None

        # --- Run the autoregressive generation using jax.lax.scan ---
        # Create flattened step indices: [(T, 0), (T, 1), ..., (T, N-1), (T+1, 0), ...]
        future_frames = seq_len - T
        total_tokens = future_frames * N
        step_indices = jnp.arange(total_tokens)
        
        initial_carry = (batch["rng"], token_idxs_full, action_tokens)
        final_carry, _ = jax.lax.scan(
            token_step_fn, initial_carry, step_indices
        )
        final_token_idxs = final_carry[1]
[content] dynamics.py (python)
            if sample_argmax:
                next_token = jnp.argmax(next_token_logits, axis=-1)  # (B,)
            else:
                rng, step_rng = jax.random.split(rng)
                next_token = jax.random.categorical(
                    step_rng, next_token_logits / temperature, axis=-1
                )  # (B,)
[content] dynamics.py (python)

        def token_step_fn(
            carry: tuple[jax.Array, jax.Array, jax.Array], step_info: jax.Array
        ) -> tuple[tuple[jax.Array, jax.Array, jax.Array], None]:
            rng, token_idxs_full, action_tokens = carry
            # step_info contains (t, n) flattened index
            t = step_info // N
            n = step_info % N
            
            jax.debug.print("Sampling token {} from frame {}", n, t)
            dyna_inputs = {
                "video_tokens": token_idxs_full,
                "latent_actions": action_tokens,
            }
            next_token_logits, _ = self.dynamics(dyna_inputs, training=False)
            next_token_logits = next_token_logits[:, t, n, :].astype(
                jnp.float32
            )  # (B, vocab_size)
[tab] models/causal.py (python)
from typing import Dict

import jax
import jax.numpy as jnp
import flax.nnx as nnx

from utils.nn import STTransformer


class DynamicsCausal(nnx.Module):
    """Causal 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,
        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.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,
            spatial_causal=True,
            rngs=rngs,
        )
        self.patch_embed = nnx.Embed(self.num_latents, self.model_dim, rngs=rngs)
        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, jax.Array], training: bool = True
    ) -> tuple[jax.Array, jax.Array | None]:
        vid_embed = self.patch_embed(batch["video_tokens"])
        act_embed = self.action_up(batch["latent_actions"])
        vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
        vid_embed_padded = jnp.pad(vid_embed, ((0, 0), (0, 0), (1, 0), (0, 0)))

        logits = self.dynamics(vid_embed_padded)[:, :, :-1]

        mask = jnp.ones(vid_embed.shape[:-1])
        return logits, mask
[tab] models/maskgit.py (python)
from typing import Dict

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,
            spatial_causal=False,
            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, jax.Array], training: bool = True
    ) -> tuple[jax.Array, jax.Array | None]:
        # --- 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 logits, mask
[tab] models/dynamics_causal.py (python)
from typing import Dict

import jax
import jax.numpy as jnp
import flax.nnx as nnx

from utils.nn import STTransformer


class DynamicsCausal(nnx.Module):
    """Causal 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,
        param_dtype: jnp.dtype,
        dtype: jnp.dtype,
        use_flash_attention: bool,
        decode: 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.param_dtype = param_dtype
        self.dtype = dtype
        self.use_flash_attention = use_flash_attention
        self.decode = decode

        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,
            spatial_causal=True,
            decode=self.decode,
            rngs=rngs,
        )
        self.patch_embed = nnx.Embed(self.num_latents, self.model_dim, rngs=rngs)
        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, jax.Array], training: bool = True
    ) -> tuple[jax.Array, jax.Array | None]:
        vid_embed = self.patch_embed(batch["video_tokens"])
        act_embed = self.action_up(batch["latent_actions"])
        vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
        vid_embed_padded = jnp.pad(vid_embed, ((0, 0), (0, 0), (1, 0), (0, 0)))

        logits = self.dynamics(vid_embed_padded)[:, :, :-1]

        mask = jnp.ones(vid_embed.shape[:-1])
        return logits, mask
[tab] models/dynamics_maskgit.py (python)
from typing import Dict

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,
            spatial_causal=False,
            decode=False,
            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, jax.Array], training: bool = True
    ) -> tuple[jax.Array, jax.Array | None]:
        # --- 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 logits, mask
[content] sample.py (python)
        # For temporal attention, we need to initialize cache with the shape that will be used
        # after swapaxes(1, 2) in STBlock. The temporal attention expects (batch, spatial_dim, seq_len, model_dim)
        # but the cache expects (batch, seq_len, num_heads, head_dim)
        B, T, H, W, C = video_batch.shape
        num_patches = (H * W) // (4 * 4)  # patch_size=4
        # For temporal attention, the "sequence length" is the spatial dimension (num_patches)
        # and we need to reshape to (batch, seq_len, num_heads, head_dim)
        temporal_cache_shape = (B, num_patches, 8, 64)  # (batch, seq_len, num_heads, head_dim)
        block.temporal_attention.init_cache(temporal_cache_shape)
[content] sample.py (python)
        # For temporal attention, we need to initialize cache with the shape that will be used
        # after swapaxes(1, 2) in STBlock. The temporal attention expects (batch, spatial_dim, seq_len, model_dim)
        B, T, H, W, C = video_batch.shape
        num_patches = (H * W) // (4 * 4)  # patch_size=4
        # The temporal attention input shape after swapaxes is (B, num_patches, T, model_dim)
        # But the cache expects (batch, seq_len, num_heads, head_dim)
        # For temporal attention, the sequence length is num_patches, not T
        temporal_input_shape = (B, num_patches, T, 512)  # (batch, seq_len, spatial_dim, model_dim)
        block.temporal_attention.init_cache(temporal_input_shape)
[content] sample.py (python)
        # For temporal attention, we need to initialize cache with the shape that will be used
        # after swapaxes(1, 2) in STBlock. The temporal attention expects (batch, spatial_dim, seq_len, model_dim)
        B, T, H, W, C = video_batch.shape
        num_patches = (H * W) // (4 * 4)  # patch_size=4
        # The temporal attention input shape after swapaxes is (B, num_patches, T, model_dim)
        # The cache expects (batch, seq_len, num_heads, head_dim)
        # For temporal attention, the sequence length is T (time dimension), not num_patches
        temporal_input_shape = (B, T, num_patches, 512)  # (batch, seq_len, spatial_dim, model_dim)
        block.temporal_attention.init_cache(temporal_input_shape)
[content] sample.py (python)
    # The autoregressive cache needs to be initialized with the shape of the tokenized inputs, not the raw video.
    # The number of spatial tokens is derived from the image dimensions and patch size.
    # It appears the 90x160 image is padded to 92x160, and a CLS token is added.
    # (92 // args.patch_size) * (160 // args.patch_size) + 1 = 23 * 40 + 1 = 921
    num_patches = ((args.image_height + 3) // 4 * 4 // args.patch_size) * (
        args.image_width // args.patch_size
    ) + 1
    # Shape for spatial attention: (batch, time, patches, dim)
    spatial_token_shape = (
        args.batch_size,
        args.seq_len,
        num_patches,
        args.dyna_dim,
    )
    # Shape for temporal attention: (batch, patches, time, dim)
    temporal_token_shape = (
        args.batch_size,
        num_patches,
        args.seq_len,
        args.dyna_dim,
    )
[content] sample.py (python)
    # The autoregressive cache needs to be initialized with the shape of the tokenized inputs, not the raw video.
    # The number of spatial tokens is derived from the image dimensions and patch size.
    # It appears the 90x160 image is padded to 92x160, and a CLS token is added.
    # (92 // args.patch_size) * (160 // args.patch_size) + 1 = 23 * 40 + 1 = 921
    num_patches = ((args.image_height + 3) // 4 * 4 // args.patch_size) * (
        args.image_width // args.patch_size
    ) + 1
    head_dim = args.dyna_dim // args.dyna_num_heads
    # Shape for spatial attention: (batch, time, patches, num_heads, head_dim)
    spatial_token_shape = (
        args.batch_size,
        args.seq_len,
        num_patches,
        args.dyna_num_heads,
        head_dim,
    )
    # Shape for temporal attention: (batch, patches, time, num_heads, head_dim)
    temporal_token_shape = (
        args.batch_size,
        num_patches,
        args.seq_len,
        args.dyna_num_heads,
        head_dim,
    )
[content] sample.py (python)
    # The autoregressive cache needs to be initialized with the shape of the tokenized inputs, not the raw video.
    # The number of spatial tokens is derived from the image dimensions and patch size.
    # It appears the 90x160 image is padded to 92x160, and a CLS token is added.
    # (92 // args.patch_size) * (160 // args.patch_size) + 1 = 23 * 40 + 1 = 921
    num_patches = ((args.image_height + 3) // 4 * 4 // args.patch_size) * (
        args.image_width // args.patch_size
    ) + 1
    head_dim = args.dyna_dim // args.dyna_num_heads
    # Shape for spatial attention: (batch, time, patches, num_heads, head_dim)
    spatial_token_shape = (
        args.batch_size,
        args.seq_len,
        num_patches,
        args.dyna_num_heads,
        head_dim,
    )
    # Shape for temporal attention: (batch, patches, time, 1, num_heads, head_dim)
    temporal_token_shape = (
        args.batch_size,
        num_patches,
        args.seq_len,
        1,
        args.dyna_num_heads,
        head_dim,
    )
[content] sample.py (python)
    # The autoregressive cache needs to be initialized with the shape of the tokenized inputs, not the raw video.
    # The number of spatial tokens is derived from the image dimensions and patch size.
    # It appears the 90x160 image is padded to 92x160, and a CLS token is added.
    # (92 // args.patch_size) * (160 // args.patch_size) + 1 = 23 * 40 + 1 = 921
    num_patches = ((args.image_height + 3) // 4 * 4 // args.patch_size) * (
        args.image_width // args.patch_size
    ) + 1
    head_dim = args.dyna_dim // args.dyna_num_heads
    # Shape for spatial attention: (batch, time, patches, num_heads, head_dim)
    spatial_token_shape = (
        args.batch_size,
        args.seq_len,
        num_patches,
        args.dyna_num_heads,
        head_dim,
    )
    # Shape for temporal attention: (batch, patches, time, 1, 1, num_heads, head_dim)
    temporal_token_shape = (
        args.batch_size,
        num_patches,
        args.seq_len,
        1,
        1,
        args.dyna_num_heads,
        head_dim,
    )
[content] sample.py (python)
# The autoregressive cache needs to be initialized with the shape of the tokenized inputs, not the raw video.
    # The number of spatial tokens is derived from the image dimensions and patch size.
    # It appears the 90x160 image is padded to 92x160, and a CLS token is added.
    # (92 // args.patch_size) * (160 // args.patch_size) + 1 = 23 * 40 + 1 = 921
    num_patches = ((args.image_height + 3) // 4 * 4 // args.patch_size) * (
        args.image_width // args.patch_size
    ) + 1
    # Shape for spatial attention: (batch, time, patches, dim)
    spatial_token_shape = (
        args.batch_size,
        args.seq_len,
        num_patches,
        args.dyna_dim,
    )
    # Shape for temporal attention: (batch, patches, time, dim)
    temporal_token_shape = (
        args.batch_size,
        num_patches,
        args.seq_len,
        args.dyna_dim,
    )
[content] sample.py (python)
    # Only initialize autoregressive cache for causal dynamics (not MaskGIT)
    if args.dynamics_type == "causal":
        transformer_blocks = jasmine.dynamics.transformer.blocks
        for block in transformer_blocks:
            block.spatial_attention.init_cache(spatial_token_shape)
            block.temporal_attention.init_cache(temporal_token_shape)
[content] sample.py (python)
    transformer_blocks = jasmine.dynamics.transformer.blocks
    for block in transformer_blocks:
        block.spatial_attention.init_cache(spatial_token_shape)
        block.temporal_attention.init_cache(temporal_token_shape)
[content] sample.py (python)
    # Note: Causal dynamics processes full sequences, so no cache initialization needed
    # The cache is only needed for true autoregressive token-by-token generation
    # which is not how the current causal sampling works
[content] sample.py (python)
    if args.dynamics_type == "causal":
        transformer_blocks = jasmine.dynamics.transformer.blocks
        for block in transformer_blocks:
            block.spatial_attention.init_cache(spatial_token_shape)
            block.temporal_attention.init_cache(temporal_token_shape)
[content] sample.py (python)
    # Shape for spatial attention: (batch, time, patches, num_heads, head_dim)
    spatial_token_shape = (
        args.batch_size,
        args.seq_len,
        num_patches,
        args.dyna_dim,
    )
    # Shape for temporal attention: (batch, patches, time, num_heads, head_dim)
    temporal_token_shape = (
        args.batch_size,
        num_patches,
        args.seq_len,
        args.dyna_dim,
    )
    if args.dynamics_type == "causal":
        transformer_blocks = jasmine.dynamics.transformer.blocks
        for block in transformer_blocks:
            block.spatial_attention.init_cache(spatial_token_shape)
            block.temporal_attention.init_cache(temporal_token_shape)
[content] jasmine.py (python)

            # For autoregressive decoding, we only need to pass the token from the previous step.
            # The model internally uses a KV cache to remember previous tokens.
            current_token_sequence = token_idxs_full[:, t : t + 1, :]
[content] jasmine.py (python)
            # Use `lax.dynamic_slice` for JIT compatibility with dynamic indices.
            current_token_sequence = jax.lax.dynamic_slice(
                token_idxs_full, (0, t, 0), (B, 1, N)
            )
[content] utils/nn.py (python)
        if decode:
            return jax.nn.dot_product_attention(
                query=query,
                key=key,
                value=value,
                bias=bias,
                mask=mask,
                is_causal=is_causal,
                **kwargs,
            )