p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] test.ipynb (python)
import os
import shutil

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

# The hunch is that parameters like biases in Linear layers are only created if
# they are used (e.g. `use_bias=True`). When checkpointing a model created with
# `use_bias=False`, the bias parameter is not saved.
#
# Attempting to restore this checkpoint into a model instance created with
# `use_bias=True` is expected to fail, as the checkpointed state has a
# different structure (missing the 'bias' leaf in the PyTree) from the target model.

ckpt_dir = '/tmp/flax_nnx_bug_ckpt'
if os.path.exists(ckpt_dir):
    shutil.rmtree(ckpt_dir)
os.makedirs(ckpt_dir)


class MyModel(nnx.Module):
  """A simple model containing a Linear layer."""

  def __init__(self, use_bias: bool, *, rngs: nnx.Rngs):
    self.linear = nnx.Linear(10, 20, use_bias=use_bias, rngs=rngs)

  def __call__(self, x):
    return self.linear(x)


rngs = nnx.Rngs(0)
model_no_bias = MyModel(use_bias=False, rngs=rngs)

dummy_input = jnp.ones((1, 10))
model_no_bias(dummy_input)
[content] README.md (markdown)
Distributed checkpointing & loading
also onto a different hardware topology (e.g. initial training on 1 node --> reinitialization onto 4 nodes)
Asynchronous checkpointing
Optimized dataloading (each GPU loads its shard in parallel)
Model+optimizer states+dataloader checkpointing (just resume your training run and everything is taken care of)
Full reproducibility with exact training curves (seeded dataloading + training)
Automatic checkpoint management (automatic deletion/retention according to specified retention policy)
Mixed precision (with correct full precision casting, e.g. before softmax')
Flash attention (from cudnn, which is more optimized than Tri Dao's)
int8-quantization is on the roadmap (via https://github.com/google/aqt)
KV caching for inference of causal Transformer (still in PR)
Currently working on KV cache resets for accelerated spatiotemporal attention (I don't think this has ever been done before in the open)
Activation checkpointing (even onto host memory if desired)
Distributed Data parallelism (changing to FSDP requires changing one line of code)
wandb, cli and tensorboard logging (tb logging still in PR)
Just-in-time compiled train step
Cosine/ WSD learning rate schedules (no need to retrain if you realize that you want to train for longer)
Index shuffling during dataloading (doesn't require loading the dataset into memory for shuffling)
'google-native' stack
https://github.com/google/orbax for checkpointing
https://github.com/google/grain for dataloading
https://github.com/google-deepmind/dm_pix for image manipulation
https://github.com/google/array_record as the data format
essentially, no other dependencies
I am currently working on migrating to the new flax.nnx API, after which we will also have:
significantly less boilerplate
easy model surgery
model inspection
[content] README.md (markdown)
Jafar is a JAX-based implementation of the DeepMind paper "[Genie: Generative Interactive Environments](https://arxiv.org/abs/2402.15391)" (Bruce et al., 2024).

Jafar supports training of all Genie components and can complete the CoinRun reproducibility experiment (Appendix F) on a single L40S GPU in under a week.
[tab] experiments/dynamics_grain_tok_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"
tokenizer_ckpt_dir="$PWD/checkpoints/tokenizer_openai_grain_checkpointing"

export XLA_FLAGS=--xla_gpu_autotune_level=0
PYTHONUNBUFFERED=1 python train_dynamics.py \
    --restore_ckpt \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --batch_size 12 \
    --tokenizer_checkpoint $tokenizer_ckpt_dir \
    --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] train_dynamics.py (python)
    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
[content] test.ipynb (python)

model_no_dropout(dummy_input)

state_no_dropout = nnx.split(model_no_dropout)
nnx.display(state_no_dropout)

model_with_dropout = MyModel(use_bias=False, dropout=0.5, rngs=rngs)
model_with_dropout(dummy_input)

state_with_dropout = nnx.split(model_with_dropout)
nnx.display(state_with_dropout)
[content] test.ipynb (python)
# To get a summary of which parts are different:
def compare_states(state1, state2):
    """Compare two states and show which parts are different."""
    def compare_leaf(path, x, y):
        if jnp.allclose(x, y):
            return f"{path}: Same"
        else:
            return f"{path}: Different"
    
    return jax.tree_util.tree_map_with_path(compare_leaf, state1, state2)

comparison = compare_states(state_with_dropout, state_no_dropout)
[content] test.ipynb (python)
    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.PRNGKey(args.seed)

    # --- Initialize model ---
    rng, _rng = jax.random.split(rng)
    rngs = nnx.Rngs(_rng)
    genie = Genie(
        # Tokenizer
        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,
        # LAM
        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,
        # Dynamics
        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,
        rngs=rngs,
    )
[content] test.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
    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
    # Tokenizer
    tokenizer_dim: int = 512
    tokenizer_ffn_dim: int = 2048
    latent_patch_dim: int = 32
    num_patch_latents: int = 1024
    patch_size: int = 4
    tokenizer_num_blocks: int = 4
    tokenizer_num_heads: int = 8
    tokenizer_checkpoint: str = ""
    # LAM
    lam_dim: int = 512
    lam_ffn_dim: int = 2048
    latent_action_dim: int = 32
    num_latent_actions: int = 6
    lam_patch_size: int = 16
    lam_num_blocks: int = 4
    lam_num_heads: int = 8
    lam_checkpoint: str = ""
    # Dynamics
    dyna_dim: int = 512
    dyna_ffn_dim: int = 2048
    dyna_num_blocks: int = 6
    dyna_num_heads: int = 8
    dropout: float = 0.0
    mask_limit: float = 0.5
    param_dtype = jnp.float32
    dtype = jnp.bfloat16
    use_flash_attention: bool = True
    # Logging
    log: bool = False
    entity: str = ""
    project: str = ""
    name: str = "train_dynamics"
    tags: list[str] = field(default_factory=lambda: ["dynamics"])
    log_interval: int = 5
    log_image_interval: int = 250
    ckpt_dir: str = ""
    log_checkpoint_interval: int = 25000
    log_checkpoint_keep_period: int = 20000
    log_gradients: bool = False
    wandb_id: str = ""
[content] test.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
    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
    # Tokenizer
    tokenizer_dim: int = 512
    tokenizer_ffn_dim: int = 2048
    latent_patch_dim: int = 32
    num_patch_latents: int = 1024
    patch_size: int = 4
    tokenizer_num_blocks: int = 4
    tokenizer_num_heads: int = 8
    tokenizer_checkpoint: str = ""
    # LAM
    lam_dim: int = 512
    lam_ffn_dim: int = 2048
    latent_action_dim: int = 32
    num_latent_actions: int = 6
    lam_patch_size: int = 16
    lam_num_blocks: int = 4
    lam_num_heads: int = 8
    lam_checkpoint: str = ""
    # Dynamics
    dyna_dim: int = 512
    dyna_ffn_dim: int = 2048
    dyna_num_blocks: int = 6
    dyna_num_heads: int = 8
    dropout: float = 0.0
    mask_limit: float = 0.5
    param_dtype = jnp.float32
    dtype = jnp.bfloat16
    use_flash_attention: bool = True
    # Logging
    log: bool = False
    entity: str = ""
    project: str = ""
    name: str = "train_dynamics"
    tags: list[str] = field(default_factory=lambda: ["dynamics"])
    log_interval: int = 5
    log_image_interval: int = 250
    ckpt_dir: str = ""
    log_checkpoint_interval: int = 25000
    log_checkpoint_keep_period: int = 20000
    log_gradients: bool = False
    wandb_id: str = ""u
[content] test.ipynb (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.dtype,
    )
    optimizer = nnx.Optimizer(genie, tx)
[tab] utils/lr_utils.py (python)
import optax


def get_lr_schedule(
    lr_schedule: str,
    init_lr: float,
    max_lr: float,
    decay_end: float,
    total_steps: int,
    warmup_steps: int,
    wsd_decay_steps: int,
) -> optax.Schedule:
    supported_schedules = ["wsd", "cos"]
    if lr_schedule == "cos":
        assert (
            warmup_steps <= total_steps
        ), "Warmup steps can't be greater than total steps."
        return optax.warmup_cosine_decay_schedule(
            init_value=init_lr,
            peak_value=max_lr,
            warmup_steps=warmup_steps,
            decay_steps=total_steps,  # Note: decay_steps includes the warmup steps, so we need to pass total value
            end_value=decay_end,
        )
    elif lr_schedule == "wsd":
        assert (
            warmup_steps + wsd_decay_steps <= total_steps
        ), "Warmup and decay period is longer than total steps."
        schedules = [
            optax.linear_schedule(
                init_value=init_lr, end_value=max_lr, transition_steps=warmup_steps
            ),
            optax.constant_schedule(value=max_lr),
            optax.linear_schedule(
                init_value=max_lr, end_value=decay_end, transition_steps=wsd_decay_steps
            ),
        ]
        boundaries = [warmup_steps, total_steps - wsd_decay_steps]
        return optax.join_schedules(schedules, boundaries)
    else:
        raise ValueError(
            f"Learning rate schedule not supported. Please use one of {supported_schedules}"
        )
[content] test.ipynb (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)
[content] test.ipynb (python)

    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,
    )
[content] test.ipynb (python)

        optimizer.model.tokenizer.vq.drop.rngs = rngs
        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),
                dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),
            ),
        )
[content] train_dynamics.py (python)

        # FIXME (f.srambical): add link to issue
        # FIXME (f.srambical): quickfix of a bug in flax.nnx
        abstract_optimizer_state.model.tokenizer.vq.drop = nnx.state(nnx.eval_shape(lambda: nnx.Dropout(0.1, rngs=rngs)))
[tab] experiments/sample.sh (shellscript)
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
PYTHONUNBUFFERED=1 python sample.py \
    --batch_size 1 \
    --seq_len 2 \
    --start_frame 1 \
    --checkpoint $ckpt_dir \
    --data_dir $data_dir
[tab] utils/parameter_utils.py (python)
from jax.tree_util import tree_map, tree_reduce


def _count_leaf(x):
    """Count parameters in a single leaf node."""
    if hasattr(x, "size"):
        return x.size
    return 0


def _count_component(component_params):
    """Count total parameters in a component."""
    return tree_reduce(
        lambda x, y: x + y, tree_map(_count_leaf, component_params), initializer=0
    )


def count_parameters_by_component(params):
    """Count parameters for each component of the model.

    Args:
        params: Model parameters from nnx.split(model, nnx.Param, ...)

    Returns:
        Dictionary with parameter counts for each component
    """
    component_names = list(params.keys())
    print(f"Counting all components: {component_names}")

    counts = {}
    total_params = 0

    for name in component_names:
        component_params = params[name]
        count = _count_component(component_params)
        counts[name] = count
        total_params += count

    counts["total"] = total_params
    return counts
[content] /fast/home/franz.srambical/jafar/sample.py (python)
    # @nnx.jit
    # @jax.jit
    def _sampling_fn(model, batch):
        """Runs Genie.sample with pre-defined generation hyper-parameters."""
        return model.sample(
            batch,
            args.seq_len,
            args.maskgit_steps,
            args.temperature,
            args.sample_argmax,
        )


    # --- Define autoregressive sampling loop ---
    def _autoreg_sample(rng, video_batch, action_batch):
        vid = video_batch[:, : args.start_frame + 1]
        rng, _rng = jax.random.split(rng)
        batch = dict(videos=vid, latent_actions=action_batch, rng=_rng)
        generated_vid = genie.sample(batch, args.seq_len, args.maskgit_steps, args.temperature, args.sample_argmax)
        return generated_vid


    # --- Get video + latent actions ---
    array_record_files = [
        os.path.join(args.data_dir, x)
        for x in os.listdir(args.data_dir)
        if x.endswith(".array_record")
    ]
    dataloader = get_dataloader(
        array_record_files,
        args.seq_len,
        args.batch_size,
        args.image_height,
        args.image_width,
        args.image_channels,
        num_workers=8,
        prefetch_buffer_size=1,
        seed=args.seed,
    )
    video_batch = next(iter(dataloader))
    # Get latent actions for all videos in the batch
    batch = dict(videos=video_batch)
    action_batch = genie.vq_encode(batch, training=False)  # type: ignore[arg-type]
    action_batch = jnp.asarray(action_batch).reshape(video_batch.shape[0], args.seq_len - 1, 1)

    # --- Sample + evaluate video ---
    vid = _autoreg_sample(rng, video_batch, action_batch)
    gt = video_batch[:, : vid.shape[1]].clip(0, 1).reshape(-1, *video_batch.shape[2:])
    recon = vid.clip(0, 1).reshape(-1, *vid.shape[2:])
    ssim = pix.ssim(gt[:, args.start_frame + 1 :], recon[:, args.start_frame + 1 :]).mean()
[content] genie.py (python)
        MaskGITLoop = nn.scan(


            MaskGITStep,


            variable_broadcast="params",


            split_rngs={"params": False},


            in_axes=0,


            out_axes=0,


            length=steps,


        )
[content] genie.py (python)
            # Initialize carry for MaskGIT loop
            carry = (step_rng, masked_token_idxs, mask, action_tokens)

            # Run the MaskGIT refinement loop
            for i in range(steps):
                carry, _ = maskgit_step(carry, i)

            # Update token indices with the output from MaskGIT
            current_token_idxs = carry[1]