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

#SBATCH --nodes=2
#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/dynamics_openai_grain_tok_lam_restore"

export XLA_FLAGS=--xla_gpu_autotune_level=0
export PYTHONUNBUFFERED=1
srun python train_dynamics.py \
    --dyna_type 'causal' \
    --batch_size 12 \
    --restore_ckpt \
    --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] genie.py (python)
        # videos_BTHWC = batch["videos"]
        # FIXME (f.srambical): this is the part that I should omit in order to test
        # whether input-boundness is due to us being in the video generation model regime
        # tokenizer_outputs = self.tokenizer.vq_encode(videos_BTHWC, training=False)
        # token_indices_BTN = tokenizer_outputs["indices"]
        token_indices_BTN = jnp.zeros((12, 4, 920), dtype=jnp.int32)
        # lam_outputs = self.lam.vq_encode(videos_BTHWC, training=False)
        # z_q_BTm11L = lam_outputs["z_q"]
        # action_indices_E = lam_outputs["indices"]
        # latent_actions_BTm11L = jax.lax.cond(
        #     self.lam_co_train,
        #     lambda: z_q_BTm11L,
        #     lambda: jax.lax.stop_gradient(z_q_BTm11L),
        # )
        outputs = dict(
            video_tokens=jax.lax.stop_gradient(token_indices_BTN),
            latent_actions=jnp.zeros((12, 3, 1, self.latent_action_dim), dtype=self.dtype),
        )
        outputs["mask_rng"] = batch["mask_rng"]
        dyna_logits_BTNV, dyna_mask = self.dynamics(outputs, training)
        outputs["token_logits"] = dyna_logits_BTNV
        if dyna_mask is not None:
            outputs["mask"] = dyna_mask
        # mle_indices_BTN = jnp.argmax(outputs["token_logits"], axis=-1)
        # H, W = batch["videos"].shape[2:4]
        # outputs["recon"] = self.tokenizer.decode(mle_indices_BTN, (H, W))
[tab] test/test_nan.ipynb (python)
import torch
import torch.nn as nn

class NaNForwardModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.bias = nn.Parameter(torch.tensor(0.0))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Construct 0/0 to produce NaNs across the output
        z1 = x - x                 # zeros, shape like x
        z2 = self.bias - self.bias # scalar 0.0
        return z1 / z2             # 0.0 / 0.0 => NaN

# Demo
x = torch.randn(2, 3)
model = NaNForwardModel()
y = model(x)
print("Output:
", y)
print("Contains NaN:", torch.isnan(y).any().item())
[content] test/test_nan.ipynb (python)
# Minimal example: flax.nnx model that forces NaNs on forward
import jax
import jax.numpy as jnp
import flax.nnx as nnx


class NaNModel(nnx.Module):
    def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
        # Force NaNs deterministically: log of negative numbers yields NaN
        return jnp.log(-jnp.ones_like(x))


# Small test input
x = jnp.ones((2, 3), dtype=jnp.float32)

# Instantiate and run forward
model = NaNModel()
y = model(x)

print("Output:
", y)
print("Has NaN:", bool(jnp.isnan(y).any()))
[content] test/test_nan.ipynb (python)
import jax
import jax.numpy as jnp
import flax.nnx as nnx

# Shapes
batch, length, features = 1, 2, 4
x = jnp.ones((batch, length, features), dtype=jnp.float32)

# Multi-head attention layer
mha = nnx.MultiHeadAttention(
    num_heads=2,
    in_features=features,
    qkv_features=4,
    out_features=4,
    rngs=nnx.Rngs(0),
)

# All-false mask -> softmax over all -inf -> NaNs
mask = jnp.zeros((batch, mha.num_heads, length, length), dtype=bool)

# Forward pass (forces NaNs)
y = mha(x, mask=mask, sow_weights=True)

print("Output:
", y)
print("Has NaN:", bool(jnp.isnan(y).any()))
[content] test/test_nan.ipynb (python)
import jax.numpy as jnp
import flax.nnx as nnx

batch, length, features = 1, 128, 4

mha = nnx.MultiHeadAttention(
    num_heads=2,
    in_features=features,
    qkv_features=4,
    out_features=4,
    rngs=nnx.Rngs(0),
)
[content] test/test_nan.ipynb (python)

x = jnp.ones((batch, length, features), dtype=jnp.float32)
y = mha(x, sow_weights=True, decode=False)
y = mha(x, sow_weights=True, decode=False)
y = mha(x, sow_weights=True, decode=False)

nnx.display(mha)
[content] test/test_nan.ipynb (python)
# Restore a dynamics checkpoint and enable sowing
import os
from typing import Dict

import jax
import jax.numpy as jnp
import flax.nnx as nnx
import optax
import orbax.checkpoint as ocp

from genie import Genie

# Adjust to your checkpoint directory and correct dynamics type
ckpt_dir = "/home/franz.srambical/jafar/checkpoints/dynamics_openai_grain_tok_restore"  # TODO: set to your run
dyna_type = "causal"  # or "maskgit" to match the checkpoint

# Model hyperparameters (must match the checkpoint)
image_channels = 3
seq_len = 16
image_height = 90
image_width = 160

# Tokenizer
tokenizer_dim = 512
tokenizer_ffn_dim = 2048
latent_patch_dim = 32
num_patch_latents = 1024
patch_size = 4
tokenizer_num_blocks = 4
tokenizer_num_heads = 8
# LAM
lam_dim = 512
lam_ffn_dim = 2048
latent_action_dim = 32
num_latent_actions = 6
lam_patch_size = 16
lam_num_blocks = 4
lam_num_heads = 8
# Dynamics
dyna_dim = 512
dyna_ffn_dim = 2048
dyna_num_blocks = 6
dyna_num_heads = 8
mask_limit = 0.5
dropout = 0.0
param_dtype = jnp.float32
dtype = jnp.bfloat16
use_flash_attention = True

# Build model graph matching the checkpoint
rng = jax.random.key(0)
rngs = nnx.Rngs(rng)
model = Genie(
    # Tokenizer
    in_dim=image_channels,
    tokenizer_dim=tokenizer_dim,
    tokenizer_ffn_dim=tokenizer_ffn_dim,
    latent_patch_dim=latent_patch_dim,
    num_patch_latents=num_patch_latents,
    patch_size=patch_size,
    tokenizer_num_blocks=tokenizer_num_blocks,
    tokenizer_num_heads=tokenizer_num_heads,
    # LAM
    lam_dim=lam_dim,
    lam_ffn_dim=lam_ffn_dim,
    latent_action_dim=latent_action_dim,
    num_latent_actions=num_latent_actions,
    lam_patch_size=lam_patch_size,
    lam_num_blocks=lam_num_blocks,
    lam_num_heads=lam_num_heads,
    lam_co_train=True,
    # Dynamics
    dyna_type=dyna_type,
    dyna_dim=dyna_dim,
    dyna_ffn_dim=dyna_ffn_dim,
    dyna_num_blocks=dyna_num_blocks,
    dyna_num_heads=dyna_num_heads,
    dropout=dropout,
    mask_limit=mask_limit,
    param_dtype=param_dtype,
    dtype=dtype,
    use_flash_attention=use_flash_attention,
    decode=False,
    rngs=rngs,
)

# Optimizer (matches training opt hyperparams; lr value is irrelevant for restore)
tx = optax.adamw(
    learning_rate=1e-4,
    b1=0.9,
    b2=0.9,
    weight_decay=1e-4,
    mu_dtype=param_dtype,
)
optimizer = nnx.Optimizer(model, tx)
[tab] test/test_nan.ipynb (python)
# Restore latest checkpoint into optimizer.model
handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
handler_registry.add(
    "model_state", ocp.args.PyTreeRestore, ocp.handlers.PyTreeCheckpointHandler
)

ckpt_mgr = ocp.CheckpointManager(
    directory=ckpt_dir,
    options=ocp.CheckpointManagerOptions(step_format_fixed_length=6),
    handler_registry=handler_registry,
)

abstract_state = nnx.state(optimizer)
restored = ckpt_mgr.restore(
    ckpt_mgr.latest_step(),
    args=ocp.args.Composite(
        model_state=ocp.args.PyTreeRestore(abstract_state),  # type: ignore
    ),
)
nnx.update(optimizer, restored["model_state"])  # type: ignore
ckpt_mgr.close()

# Convenience handle
model = optimizer.model
print("Restored dynamics checkpoint.")
[tab] test/test_nan.ipynb (python)
import jax.numpy as jnp
import flax.nnx as nnx

batch, length, features = 1, 128, 4

mha = nnx.MultiHeadAttention(
    num_heads=2,
    in_features=features,
    qkv_features=4,
    out_features=4,
    rngs=nnx.Rngs(0),
)
[tab] test/test_nan.ipynb (python)
# Restore a dynamics checkpoint and enable sowing
import os
from typing import Dict

import jax
import jax.numpy as jnp
import flax.nnx as nnx
import optax
import orbax.checkpoint as ocp
import grain

from utils.dataloader import get_dataloader
from genie import Genie

# Adjust to your checkpoint directory, dataset directory, and dynamics type
ckpt_dir = "/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/lam/train_lam_coinrun_reproduction_20067/100000_ckpt"
data_dir = "/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_10m"
train_seed = 0
batch_size = 36

# Model hyperparameters (must match the checkpoint)
image_channels = 3
seq_len = 16
image_height = 90
image_width = 160

# Tokenizer
tokenizer_dim = 512
tokenizer_ffn_dim = 2048
latent_patch_dim = 32
num_patch_latents = 1024
patch_size = 4
tokenizer_num_blocks = 4
tokenizer_num_heads = 8
# LAM
lam_dim = 512
lam_ffn_dim = 2048
latent_action_dim = 32
num_latent_actions = 6
lam_patch_size = 16
lam_num_blocks = 4
lam_num_heads = 8
# Dynamics
dyna_dim = 512
dyna_ffn_dim = 2048
dyna_num_blocks = 6
dyna_num_heads = 8
mask_limit = 0.5
dropout = 0.0
param_dtype = jnp.float32
dtype = jnp.bfloat16
use_flash_attention = False  # set True only if running with cuDNN-backed flash attention

# Build model graph matching the checkpoint
rng = jax.random.key(train_seed)
# One split before constructing modules, matching training
rng, _rng = jax.random.split(rng)
rngs = nnx.Rngs(_rng)
model = Genie(
    # Tokenizer
    in_dim=image_channels,
    tokenizer_dim=tokenizer_dim,
    tokenizer_ffn_dim=tokenizer_ffn_dim,
    latent_patch_dim=latent_patch_dim,
    num_patch_latents=num_patch_latents,
    patch_size=patch_size,
    tokenizer_num_blocks=tokenizer_num_blocks,
    tokenizer_num_heads=tokenizer_num_heads,
    # LAM
    lam_dim=lam_dim,
    lam_ffn_dim=lam_ffn_dim,
    latent_action_dim=latent_action_dim,
    num_latent_actions=num_latent_actions,
    lam_patch_size=lam_patch_size,
    lam_num_blocks=lam_num_blocks,
    lam_num_heads=lam_num_heads,
    lam_co_train=True,
    # Dynamics
    dyna_type=dyna_type,
    dyna_dim=dyna_dim,
    dyna_ffn_dim=dyna_ffn_dim,
    dyna_num_blocks=dyna_num_blocks,
    dyna_num_heads=dyna_num_heads,
    dropout=dropout,
    mask_limit=mask_limit,
    param_dtype=param_dtype,
    dtype=dtype,
    use_flash_attention=use_flash_attention,
    decode=False,
    rngs=rngs,
)

# Optimizer (matches training opt hyperparams; lr value is irrelevant for restore)
tx = optax.adamw(
    learning_rate=1e-4,
    b1=0.9,
    b2=0.9,
    weight_decay=1e-4,
    mu_dtype=param_dtype,
)
optimizer = nnx.Optimizer(model, tx)
[content] test/test_nan.ipynb (python)
    num_steps: int = 200_000
    seed: int = 0
    seq_len: int = 16
    image_channels: int = 3
    image_height: int = 90
    image_width: int = 160
    data_dir: str = ""
    save_ckpt: bool = False
    restore_ckpt: bool = False
    # Optimization
    batch_size: int = 36
    vq_beta: float = 0.25
    init_lr: float = 0.0
    max_lr: float = 3e-5
    decay_end: float = 0.0
    wsd_decay_steps: int = (
        10000  # NOTE: wsd_decay_steps will only be used when using a wsd-schedule
    )
    warmup_steps: int = 5000
    lr_schedule: str = "wsd"  # supported options: wsd, cos
    vq_reset_thresh: int = 50
    # LAM
    model_dim: int = 512
    ffn_dim: int = 2048
    latent_dim: int = 32
    num_latents: int = 6
    patch_size: int = 16
    num_blocks: int = 4
    num_heads: int = 8
    dropout: float = 0.0
    codebook_dropout: float = 0.0
    param_dtype = jnp.float32
    dtype = jnp.bfloat16
    # Logging
    log: bool = False
    entity: str = ""
    project: str = ""
    name: str = "train_lam"
    tags: list[str] = field(default_factory=lambda: ["lam"])
    log_interval: int = 5
    log_image_interval: int = 250
    ckpt_dir: str = ""
    log_checkpoint_interval: int = 10000
    log_checkpoint_keep_period: int = 20000
    wandb_id: str = ""
    use_flash_attention: bool = True
[tab] test/test_nan.ipynb (python)
# Draw the exact next batch and RNG like in training, then run a sow-enabled forward
# Next videos batch from restored Grain iterator
_it = next(loader_iterator)
# Try common layouts from Grain -> ArrayRecord -> Batch
videos = None
try:
    videos = _it["_data"]["video"]["values"][..., :image_channels]
except Exception:
    try:
        videos = _it["video"]
    except Exception:
        # Fallback: look for the first array-like value
        for v in _it.values() if isinstance(_it, dict) else []:
            if hasattr(v, "shape") and v.ndim >= 4:
                videos = v
                break
if videos is None:
    raise RuntimeError("Could not extract 'videos' tensor from the restored dataloader item.")

# Ensure dtype matches training expectations (uint8 frames)
if videos.dtype != jnp.uint8:
    videos = videos.astype(jnp.uint8)

# Reproduce the next mask rng from the training RNG stream
# Training used: rng = jax.random.key(seed); rng, _rng_mask = jax.random.split(rng, 2)
rng, mask_rng = jax.random.split(rng, 2)

batch: Dict[str, jax.Array] = {
    "videos": videos,
    "mask_rng": mask_rng,
}

# Enable sowing by toggling the module flags so that inner calls sow intermediates
tr = model.dynamics.transformer
if hasattr(tr, "sow_logits"): tr.sow_logits = True
if hasattr(tr, "sow_weights"): tr.sow_weights = True
if hasattr(tr, "sow_activations"): tr.sow_activations = True
for blk in getattr(tr, "blocks", []):
    if hasattr(blk, "sow_weights"): blk.sow_weights = True
    if hasattr(blk, "sow_activations"): blk.sow_activations = True

# Forward pass (will populate sowed values)
outputs = model(batch, training=True)
print("Forward with sowing done.")
[content] test/test_nan.ipynb (python)
    lam = LatentActionModel(
        in_dim=args.image_channels,
        model_dim=args.model_dim,
        ffn_dim=args.ffn_dim,
        latent_dim=args.latent_dim,
        num_latents=args.num_latents,
        patch_size=args.patch_size,
        num_blocks=args.num_blocks,
        num_heads=args.num_heads,
        dropout=args.dropout,
        codebook_dropout=args.codebook_dropout,
        param_dtype=args.param_dtype,
        dtype=args.dtype,
        use_flash_attention=args.use_flash_attention,
        rngs=rngs,
    )
[content] test/test_nan.ipynb (python)

if hasattr(decoder, "sow_logits"): decoder.sow_logits = True
if hasattr(decoder, "sow_weights"): decoder.sow_weights = True
if hasattr(decoder, "sow_activations"): decoder.sow_activations = True
for blk in getattr(decoder, "blocks", []):
    if hasattr(blk, "sow_weights"): blk.sow_weights = True
    if hasattr(blk, "sow_activations"): blk.sow_activations = True
[content] test/test_nan.ipynb (python)

    model.sow_activations = True
    model.sow_logits = True
    model.sow_weights = True
    for blk in getattr(model, "blocks", []):
        blk.sow_weights = True
        blk.sow_activations = True