p-doom/crowd-code-dataset-0.2 huggingface dataset:
[terminal_command] TERMINAL ()
less /fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/lam/lam_coinrun_nan_investigation_100k_to_107k_20629.logfast/(jafar) [franz.srambical@hai-login2.haicore.berlin:~/jafar] $ tail -f /fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/lam/lam_coinrun_nan_investigation_100k_to_107k_20629.log
[terminal_command] TERMINAL ()
ls /fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/lam/train_lam_coinrun_reproduction_20067/100000_ckpt/100000/_CHECKPOINT_METADATA  dataloader_state/     model_state/          (jafar) [franz.srambical@hai-login2.haicore.berlin:~/jafar] $ ls /fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/lam/train_lam_coinrun_reproduction_20067/100000_ckpt/100000/_CHECKPOINT_METADATA
[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.lr_utils import get_lr_schedule
from utils.dataloader import get_dataloader
from models.lam import LatentActionModel

# 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"
num_steps: int = 200_000
seed: int = 0
seq_len: int = 16
image_channels: int = 3
image_height: int = 64
image_width: int = 64
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_interval: int = 5
log_image_interval: int = 250
use_flash_attention: bool = False

# Build model graph matching the checkpoint
rng = jax.random.key(seed)
rng, _rng = jax.random.split(rng)
rngs = nnx.Rngs(_rng)
lam = LatentActionModel(
    in_dim=image_channels,
    model_dim=model_dim,
    ffn_dim=ffn_dim,
    latent_dim=latent_dim,
    num_latents=num_latents,
    patch_size=patch_size,
    num_blocks=num_blocks,
    num_heads=num_heads,
    dropout=dropout,
    codebook_dropout=codebook_dropout,
    param_dtype=param_dtype,
    dtype=dtype,
    use_flash_attention=use_flash_attention,
    rngs=rngs,
)

# Optimizer (matches training opt hyperparams)
lr_schedule_fn = get_lr_schedule(
    lr_schedule,
    init_lr,
    max_lr,
    decay_end,
    num_steps,
    warmup_steps,
    wsd_decay_steps,
)

tx = optax.adamw(
    learning_rate=lr_schedule_fn,
    b1=0.9,
    b2=0.9,
    weight_decay=1e-4,
    mu_dtype=param_dtype,
)
optimizer = nnx.Optimizer(lam, tx)
[content] train_lam.py (python)
    if args.restore_ckpt:
        abstract_optimizer = nnx.eval_shape(lambda: optimizer)
        abstract_optimizer_state = nnx.state(abstract_optimizer)
        restored = checkpoint_manager.restore(
            checkpoint_manager.latest_step(),
            args=ocp.args.Composite(
                model_state=ocp.args.PyTreeRestore(abstract_optimizer_state),  # type: ignore
                dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),  # type: ignore
            ),
        )
        restored_optimizer_state = restored["model_state"]
        nnx.update(optimizer, restored_optimizer_state)
        grain_iterator = restored["dataloader_state"]
        step = checkpoint_manager.latest_step() or 0
        print(f"Restored dataloader and model state from step {step}")
[content] train_lam.py (python)
    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)
[content] train_lam.py (python)
    model_state = nnx.state(optimizer.model)
    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
    )
    nnx.update(optimizer, optimizer_sharded_state)
[content] train_lam.py (python)
    device_mesh_arr = create_device_mesh((num_devices,))
    mesh = Mesh(devices=device_mesh_arr, axis_names=("data",))

    replicated_sharding = NamedSharding(mesh, PartitionSpec())
    videos_sharding = NamedSharding(mesh, PartitionSpec("data", None, None, None, None))
[content] train_lam.py (python)
    lr_schedule = get_lr_schedule(
        args.lr_schedule,
        args.init_lr,
        args.max_lr,
        args.decay_end,
        args.num_steps,
        args.warmup_steps,
        args.wsd_decay_steps,
    )
    tx = optax.adamw(
        learning_rate=lr_schedule,
        b1=0.9,
        b2=0.9,
        weight_decay=1e-4,
        mu_dtype=args.param_dtype, # moments in full precision
    )
    optimizer = nnx.Optimizer(lam, tx)
[content] train_lam.py (python)
    rng, _rng = jax.random.split(rng)
    rngs = nnx.Rngs(_rng)
    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] train_lam.py (python)
    return step, optimizer, grain_iterator


def enable_sowing(lam: LatentActionModel) -> None:
    for model in [lam.encoder, lam.decoder]:
        setattr(model, "sow_logits", True)
        for blk in getattr(model, "blocks", []):
            setattr(blk, "sow_weights", True)
            setattr(blk, "sow_activations", True)


if __name__ == "__main__":
    jax.distributed.initialize()
    num_devices = jax.device_count()
    if num_devices == 0:
        raise ValueError("No JAX devices found.")
    print(f"Running on {num_devices} devices.")

    if args.batch_size % num_devices != 0:
        raise ValueError(
            f"Global batch size {args.batch_size} must be divisible by "
            f"number of devices {num_devices}."
        )

    per_device_batch_size_for_init = args.batch_size // num_devices

    rng = jax.random.key(args.seed)

    # --- Initialize model ---
    lam = build_model(args)

    # Count parameters
    _, params, _ = nnx.split(lam, nnx.Param, ...)
    param_counts = count_parameters_by_component(params)

    if args.log and jax.process_index() == 0:
        wandb_init_kwargs = {
            "entity": args.entity,
            "project": args.project,
            "name": args.name,
            "tags": args.tags,
            "group": "debug",
            "config": args,
        }

        if args.wandb_id:
            wandb_init_kwargs.update(
                {
                    "id": args.wandb_id,
                    "resume": "allow",
                }
            )
        wandb.init(**wandb_init_kwargs)

        wandb.config.update({"model_param_count": param_counts})

    print("Parameter counts:")
    print(param_counts)

    # --- Initialize optimizer ---
    optimizer, lr_schedule = build_optimizer(lam, args)

    # FIXME: switch to create_hybrid_device_mesh for runs spanning multiple nodes
    mesh, replicated_sharding, videos_sharding = build_mesh_and_sharding(num_devices)

    shard_optimizer_states(optimizer, replicated_sharding)

    # --- Initialize checkpoint manager ---
    step = 0
    checkpoint_manager = build_checkpoint_manager(args)

    # --- Create DataLoaderIterator from dataloader ---
    grain_dataloader, grain_iterator = build_dataloader(args)

    # --- Restore checkpoint ---
    step, optimizer, grain_iterator = restore_checkpoint_if_needed(
        args, checkpoint_manager, optimizer, grain_iterator
    )
[content] train_lam.py (python)
    handler_registry.add("model_state", ocp.args.PyTreeSave, ocp.handlers.PyTreeCheckpointHandler)
    handler_registry.add("model_state", ocp.args.PyTreeRestore, ocp.handlers.PyTreeCheckpointHandler)
[content] train_lam.py (python)
    optimizer_sharded_state = jax.lax.with_sharding_constraint(optimizer_state, replicated_sharding)
    nnx.update(optimizer, optimizer_sharded_state)


def build_dataloader(a: Args):
    image_shape = (a.image_height, a.image_width, a.image_channels)
    array_record_files = [
        os.path.join(a.data_dir, x)
        for x in os.listdir(a.data_dir)
        if x.endswith(".array_record")
    ]
    grain_dataloader = get_dataloader(
        array_record_files,
        a.seq_len,
        a.batch_size,
        *image_shape,
        num_workers=8,
        prefetch_buffer_size=1,
        seed=a.seed,
[content] train_lam.py (python)
    return LatentActionModel(
        in_dim=a.image_channels,
        model_dim=a.model_dim,
        ffn_dim=a.ffn_dim,
        latent_dim=a.latent_dim,
        num_latents=a.num_latents,
        patch_size=a.patch_size,
        num_blocks=a.num_blocks,
        num_heads=a.num_heads,
        dropout=a.dropout,
        codebook_dropout=a.codebook_dropout,
        param_dtype=a.param_dtype,
        dtype=a.dtype,
        use_flash_attention=a.use_flash_attention,
[content] train_dynamics.py (python)
    handler_registry.add(
        "model_state", ocp.args.PyTreeSave, ocp.handlers.PyTreeCheckpointHandler
    )
    handler_registry.add(
        "model_state", ocp.args.PyTreeRestore, ocp.handlers.PyTreeCheckpointHandler
    )
[content] train_dynamics.py (python)
        dyna_type=args.dyna_type,
        dyna_dim=args.dyna_dim,
        dyna_ffn_dim=args.dyna_ffn_dim,
        dyna_num_blocks=args.dyna_num_blocks,
        dyna_num_heads=args.dyna_num_heads,
        dropout=args.dropout,
        mask_limit=args.mask_limit,
        param_dtype=args.param_dtype,
        dtype=args.dtype,
        use_flash_attention=args.use_flash_attention,
[content] train_dynamics.py (python)
        lam_dim=args.lam_dim,
        lam_ffn_dim=args.lam_ffn_dim,
        latent_action_dim=args.latent_action_dim,
        num_latent_actions=args.num_latent_actions,
        lam_patch_size=args.lam_patch_size,
        lam_num_blocks=args.lam_num_blocks,
        lam_num_heads=args.lam_num_heads,
        lam_co_train=not args.lam_checkpoint,
[content] train_dynamics.py (python)
        in_dim=args.image_channels,
        tokenizer_dim=args.tokenizer_dim,
        tokenizer_ffn_dim=args.tokenizer_ffn_dim,
        latent_patch_dim=args.latent_patch_dim,
        num_patch_latents=args.num_patch_latents,
        patch_size=args.patch_size,
        tokenizer_num_blocks=args.tokenizer_num_blocks,
        tokenizer_num_heads=args.tokenizer_num_heads,
[content] train_dynamics.py (python)
if __name__ == "__main__":
    jax.distributed.initialize()
    num_devices = jax.device_count()
    if num_devices == 0:
        raise ValueError("No JAX devices found.")
    print(f"Running on {num_devices} devices.")
[content] train_dynamics.py (python)
        # NOTE: We have to remove the (unused) tokenizer vq dropout due flax.nnx lazily initializing modules.
        # Specifically, the first dynamics model checkpoint will contain the vq dropout module,
        # but the first full restore will fail due to nnx not initializing the module when
        # dropout is set to 0.0.
[tab] test/coinrun_lam_nan_investigation.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.lr_utils import get_lr_schedule
from utils.dataloader import get_dataloader
from models.lam import LatentActionModel

# 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"
num_steps: int = 200_000
seed: int = 0
seq_len: int = 16
image_channels: int = 3
image_height: int = 64
image_width: int = 64
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_interval: int = 5
log_image_interval: int = 250
use_flash_attention: bool = False

# Build model graph matching the checkpoint
rng = jax.random.key(seed)
rng, _rng = jax.random.split(rng)
rngs = nnx.Rngs(_rng)
lam = LatentActionModel(
    in_dim=image_channels,
    model_dim=model_dim,
    ffn_dim=ffn_dim,
    latent_dim=latent_dim,
    num_latents=num_latents,
    patch_size=patch_size,
    num_blocks=num_blocks,
    num_heads=num_heads,
    dropout=dropout,
    codebook_dropout=codebook_dropout,
    param_dtype=param_dtype,
    dtype=dtype,
    use_flash_attention=use_flash_attention,
    rngs=rngs,
)

# Optimizer (matches training opt hyperparams)
lr_schedule_fn = get_lr_schedule(
    lr_schedule,
    init_lr,
    max_lr,
    decay_end,
    num_steps,
    warmup_steps,
    wsd_decay_steps,
)

tx = optax.adamw(
    learning_rate=lr_schedule_fn,
    b1=0.9,
    b2=0.9,
    weight_decay=1e-4,
    mu_dtype=param_dtype,
)
optimizer = nnx.Optimizer(lam, tx)
[content] test/coinrun_lam_nan_investigation.ipynb (python)

from train_lam import (
    Args,
    build_model,
    build_optimizer,
    build_dataloader,
    build_checkpoint_manager,
    restore_checkpoint_if_needed,
    enable_sowing,
)

# Configure
args = Args(
    seed=0,
    seq_len=16,
    image_channels=3,
    image_height=64,
    image_width=64,
    data_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_10m",
    ckpt_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/lam/train_lam_coinrun_reproduction_20067/100000_ckpt",
    batch_size=36,
    vq_beta=0.25,
    init_lr=0.0,
    max_lr=3e-5,
    decay_end=0.0,
    wsd_decay_steps=10000,
    warmup_steps=5000,
    lr_schedule="wsd",
    vq_reset_thresh=50,
    model_dim=512,
    ffn_dim=2048,
    latent_dim=32,
    num_latents=6,
    patch_size=16,
    num_blocks=4,
    num_heads=8,
    dropout=0.0,
    codebook_dropout=0.0,
    use_flash_attention=False,
)

# Build components
lam = build_model(args)
optimizer, lr_schedule_fn = build_optimizer(lam, args)
ckpt_mgr = build_checkpoint_manager(args)
_, loader_iterator = build_dataloader(args)

# Restore latest checkpoint
step, optimizer, loader_iterator = restore_checkpoint_if_needed(
    args, ckpt_mgr, optimizer, loader_iterator
)
lam = optimizer.model
print(f"Restored optimizer and dataloader at step {step}.")

# RNG for batch forward
rng = jax.random.key(args.seed)
[tab] test/coinrun_lam_nan_investigation.ipynb (python)
videos = next(loader_iterator)

gt = jnp.asarray(videos, dtype=jnp.float32) / 255.0
videos = gt.astype(args.dtype)

rng, _rng = jax.random.split(rng, 2)

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

outputs = lam(batch, training=True)
[tab] test/coinrun_lam_nan_investigation.ipynb (python)
# Use modular helpers from train_lam to restore a checkpoint and enable sowing
from typing import Dict
import os

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

from train_lam import (
    Args,
    build_model,
    build_optimizer,
    build_dataloader,
    build_checkpoint_manager,
    restore_checkpoint_if_needed,
    enable_sowing,
)

# Configure
args = Args(
    seed=0,
    seq_len=16,
    image_channels=3,
    image_height=64,
    image_width=64,
    data_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_10m",
    ckpt_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/lam/train_lam_coinrun_reproduction_20067/100000_ckpt",
    batch_size=36,
    vq_beta=0.25,
    init_lr=0.0,
    max_lr=3e-5,
    decay_end=0.0,
    wsd_decay_steps=10000,
    warmup_steps=5000,
    lr_schedule="wsd",
    vq_reset_thresh=50,
    model_dim=512,
    ffn_dim=2048,
    latent_dim=32,
    num_latents=6,
    patch_size=16,
    num_blocks=4,
    num_heads=8,
    dropout=0.0,
    codebook_dropout=0.0,
    use_flash_attention=False,
)

# Build components
lam = build_model(args)
optimizer, lr_schedule_fn = build_optimizer(lam, args)
ckpt_mgr = build_checkpoint_manager(args)
_, loader_iterator = build_dataloader(args)

# Restore latest checkpoint
step, optimizer, loader_iterator = restore_checkpoint_if_needed(
    args, ckpt_mgr, optimizer, loader_iterator
)
lam = optimizer.model
print(f"Restored optimizer and dataloader at step {step}.")

# RNG for batch forward
rng = jax.random.key(args.seed)
[tab] test/coinrun_lam_nan_investigation.ipynb (python)
# Use modular helpers from train_lam to restore a checkpoint and enable sowing
from typing import Dict
import os

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

from train_lam import (
    Args,
    build_model,
    build_optimizer,
    build_dataloader,
    build_checkpoint_manager,
    restore_checkpoint_if_needed,
    enable_sowing,
)

# Configure
args = Args(
    seed=0,
    seq_len=16,
    image_channels=3,
    image_height=64,
    image_width=64,
    data_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_10m",
    ckpt_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/lam/train_lam_coinrun_reproduction_20067/100000_ckpt",
    batch_size=36,
    vq_beta=0.25,
    init_lr=0.0,
    max_lr=3e-5,
    decay_end=0.0,
    wsd_decay_steps=10000,
    warmup_steps=5000,
    lr_schedule="wsd",
    vq_reset_thresh=50,
    model_dim=512,
    ffn_dim=2048,
    latent_dim=32,
    num_latents=6,
    patch_size=16,
    num_blocks=4,
    num_heads=8,
    dropout=0.0,
    codebook_dropout=0.0,
    use_flash_attention=False,
)

# Build components
rng = jax.random.key(args.seed)
lam, rng = build_model(args, rng)
optimizer, lr_schedule_fn = build_optimizer(lam, args)
ckpt_mgr = build_checkpoint_manager(args)
_, loader_iterator = build_dataloader(args)

# Restore latest checkpoint
step, optimizer, loader_iterator = restore_checkpoint_if_needed(
    args, ckpt_mgr, optimizer, loader_iterator
)
lam = optimizer.model
print(f"Restored optimizer and dataloader at step {step}.")

# RNG for batch forward
rng = jax.random.key(args.seed)
[tab] test/coinrun_lam_nan_investigation.ipynb (python)
# Use modular helpers from train_lam to restore a checkpoint and enable sowing
from typing import Dict
import os

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

from train_lam import (
    Args,
    build_model,
    build_optimizer,
    build_dataloader,
    build_checkpoint_manager,
    restore_checkpoint_if_needed,
    enable_sowing,
)

args = Args(
    # Experiment
    num_steps=200_000
    seed=0
    seq_len=16
    image_channels=3
    image_height=64
    image_width=64
    data_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_10m",
    save_ckpt=False
    restore_ckpt=True
    # Optimization
    batch_size=36
    vq_beta=0.25
    init_lr=0.0
    max_lr=3e-5
    decay_end=0.0
    wsd_decay_steps=(
        10000  # NOTE: wsd_decay_steps will only be used when using a wsd-schedule
    )
    warmup_steps=5000
    lr_schedule="wsd"  # supported options: wsd, cos
    vq_reset_thresh=50
    # LAM
    model_dim=512
    ffn_dim=2048
    latent_dim=32
    num_latents=6
    patch_size=16
    num_blocks=4
    num_heads=8
    dropout=0.0
    codebook_dropout=0.0
    param_dtype=jnp.float32
    dtype=jnp.bfloat16
    # Logging
    log=False
    entity=""
    project=""
    name="train_lam"
    tags=["lam"]
    log_interval=5
    log_image_interval=250
    ckpt_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/lam/train_lam_coinrun_reproduction_20067/100000_ckpt",
    log_checkpoint_interval=10000
    log_checkpoint_keep_period=20000
    wandb_id=""
    use_flash_attention=True
)

# Build components
rng = jax.random.key(args.seed)
lam, rng = build_model(args, rng)
optimizer, lr_schedule_fn = build_optimizer(lam, args)
ckpt_mgr = build_checkpoint_manager(args)
_, loader_iterator = build_dataloader(args)

# Restore latest checkpoint
step, optimizer, loader_iterator = restore_checkpoint_if_needed(
    args, ckpt_mgr, optimizer, loader_iterator
)
lam = optimizer.model
print(f"Restored optimizer and dataloader at step {step}.")

# RNG for batch forward
rng = jax.random.key(args.seed)
[tab] test/coinrun_lam_nan_investigation.ipynb (python)
# Use modular helpers from train_lam to restore a checkpoint and enable sowing
from typing import Dict
import os

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

from train_lam import (
    Args,
    build_model,
    build_optimizer,
    build_dataloader,
    build_checkpoint_manager,
    restore_checkpoint_if_needed,
    enable_sowing,
)

args = Args(
    num_steps=200_000,
    seed=0,
    seq_len=16,
    image_channels=3,
    image_height=64,
    image_width=64,
    data_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_10m",
    save_ckpt=False,
    restore_ckpt=True,
    # Optimization
    batch_size=36,
    vq_beta=0.25,
    init_lr=0.0,
    max_lr=3e-5,
    decay_end=0.0,
    wsd_decay_steps=(
        10000  # NOTE: wsd_decay_steps will only be used when using a wsd-schedule
    ),
    warmup_steps=5000,
    lr_schedule="wsd",
    vq_reset_thresh=50,
    # LAM
    model_dim=512,
    ffn_dim=2048,
    latent_dim=32,
    num_latents=6,
    patch_size=16,
    num_blocks=4,
    num_heads=8,
    dropout=0.0,
    codebook_dropout=0.0,
    # Logging
    log=False,
    entity="",
    project="",
    name="train_lam",
    tags=["lam"],
    log_interval=5,
    log_image_interval=250,
    ckpt_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/lam/train_lam_coinrun_reproduction_20067/100000_ckpt",
    log_checkpoint_interval=10000,
    log_checkpoint_keep_period=20000,
    wandb_id="",
    use_flash_attention=True,
)

# Build components
rng = jax.random.key(args.seed)
lam, rng = build_model(args, rng)
optimizer, lr_schedule_fn = build_optimizer(lam, args)
ckpt_mgr = build_checkpoint_manager(args)
_, loader_iterator = build_dataloader(args)

# Restore latest checkpoint
step, optimizer, loader_iterator = restore_checkpoint_if_needed(
    args, ckpt_mgr, optimizer, loader_iterator
)
lam = optimizer.model
print(f"Restored optimizer and dataloader at step {step}.")

# RNG for batch forward
rng = jax.random.key(args.seed)
[tab] demo/coinrun_lam_nan_investigation.ipynb (python)
# Use modular helpers from train_lam to restore a checkpoint and enable sowing
from typing import Dict
import os

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

from train_lam import (
    Args,
    build_model,
    build_optimizer,
    build_dataloader,
    build_checkpoint_manager,
    restore_checkpoint_if_needed,
    enable_sowing,
)

args = Args(
    num_steps=200_000,
    seed=0,
    seq_len=16,
    image_channels=3,
    image_height=64,
    image_width=64,
    data_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_10m",
    save_ckpt=False,
    restore_ckpt=True,
    # Optimization
    batch_size=36,
    vq_beta=0.25,
    init_lr=0.0,
    max_lr=3e-5,
    decay_end=0.0,
    wsd_decay_steps=(
        10000  # NOTE: wsd_decay_steps will only be used when using a wsd-schedule
    ),
    warmup_steps=5000,
    lr_schedule="wsd",
    vq_reset_thresh=50,
    # LAM
    model_dim=512,
    ffn_dim=2048,
    latent_dim=32,
    num_latents=6,
    patch_size=16,
    num_blocks=4,
    num_heads=8,
    dropout=0.0,
    codebook_dropout=0.0,
    # Logging
    log=False,
    entity="",
    project="",
    name="train_lam",
    tags=["lam"],
    log_interval=5,
    log_image_interval=250,
    ckpt_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/lam/train_lam_coinrun_reproduction_20067/100000_ckpt",
    log_checkpoint_interval=10000,
    log_checkpoint_keep_period=20000,
    wandb_id="",
    use_flash_attention=True,
)

# Build components
rng = jax.random.key(args.seed)
lam, rng = build_model(args, rng)
optimizer, lr_schedule_fn = build_optimizer(lam, args)
ckpt_mgr = build_checkpoint_manager(args)
_, loader_iterator = build_dataloader(args)
[tab] demo/nan_investigation.ipynb (python)
# Use modular helpers from train_lam to restore a checkpoint and enable sowing
from typing import Dict
import os

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

from train_lam import (
    Args,
    build_model,
    build_optimizer,
    build_dataloader,
    build_checkpoint_manager,
    restore_checkpoint_if_needed,
    enable_sowing,
)

data_dir = "/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_10m"
ckpt_dir = "/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/lam/train_lam_coinrun_reproduction_20067/100000_ckpt"

args = Args(
    num_steps=200_000,
    seed=0,
    seq_len=16,
    image_channels=3,
    image_height=64,
    image_width=64,
    data_dir=data_dir,
    save_ckpt=False,
    restore_ckpt=True,
    # Optimization
    batch_size=36,
    vq_beta=0.25,
    init_lr=0.0,
    max_lr=3e-5,
    decay_end=0.0,
    wsd_decay_steps=(
        10000  # NOTE: wsd_decay_steps will only be used when using a wsd-schedule
    ),
    warmup_steps=5000,
    lr_schedule="wsd",
    vq_reset_thresh=50,
    # LAM
    model_dim=512,
    ffn_dim=2048,
    latent_dim=32,
    num_latents=6,
    patch_size=16,
    num_blocks=4,
    num_heads=8,
    dropout=0.0,
    codebook_dropout=0.0,
    # Logging
    log=False,
    entity="",
    project="",
    name="train_lam",
    tags=["lam"],
    log_interval=5,
    log_image_interval=250,
    ckpt_dir=ckpt_dir,
    log_checkpoint_interval=10000,
    log_checkpoint_keep_period=20000,
    wandb_id="",
    use_flash_attention=True,
)

# Build components
rng = jax.random.key(args.seed)
lam, rng = build_model(args, rng)
optimizer, lr_schedule_fn = build_optimizer(lam, args)
ckpt_mgr = build_checkpoint_manager(args)
_, loader_iterator = build_dataloader(args)
[tab] demo/nan_investigation.ipynb (python)
videos = next(loader_iterator)

gt = jnp.asarray(videos, dtype=jnp.float32) / 255.0
videos = gt.astype(args.dtype)

rng = jax.random.key(args.seed)
rng, _rng = jax.random.split(rng, 2)

batch: Dict[str, jax.Array] = {
    "videos": videos,
    "rng": _rng,
}
[tab] demo/nan_investigation.ipynb (python)
# Restore latest checkpoint
step, optimizer, loader_iterator = restore_checkpoint_if_needed(
    args, ckpt_mgr, optimizer, loader_iterator
)
lam = optimizer.model
print(f"Restored optimizer and dataloader at step {step}.")

enable_sowing(lam)
print("Sowing enabled on encoder/decoder.")
outputs = lam(batch, training=True)

nnx.display(lam.encoder)
[tab] demo/nan_investigation.ipynb (python)
# checkpoint @ 20k
lam, rng = build_model(args, rng)
optimizer, lr_schedule_fn = build_optimizer(lam, args)
step, optimizer, loader_iterator = restore_checkpoint_if_needed(
    args, ckpt_mgr, optimizer, loader_iterator, 20000
)
lam = optimizer.model
print(f"Restored optimizer and dataloader at step {step}.")

enable_sowing(lam)
print("Sowing enabled on encoder/decoder.")

outputs = lam(batch, training=True)

nnx.display(lam.encoder)