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

data_dir='data_arrayrecord/dummy'
ckpt_dir='checkpoints/tokenizer_openai_grain_checkpointing'

export XLA_FLAGS=--xla_gpu_autotune_level=0
srun python train_tokenizer.py \
    --batch_size 48 \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[tab] experiments/dynamics_grain_tok_restore.sh (shellscript)
#!/usr/bin/env bash
source .venv/bin/activate


data_dir='data_arrayrecord/dummy'
ckpt_dir='checkpoints/dynamics_openai_grain_tok_lam_restore'
tokenizer_ckpt_dir='checkpoints/tokenizer_openai_grain_checkpointing'

export XLA_FLAGS=--xla_gpu_autotune_level=0
srun python train_dynamics.py \
    --batch_size 48 \
    --tokenizer_checkpoint $tokenizer_ckpt_dir \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[tab] experiments/dynamics_grain_full_restore.sh (shellscript)
#!/usr/bin/env bash
source .venv/bin/activate

data_dir='data_arrayrecord/dummy'
ckpt_dir='checkpoints/dynamics_openai_grain_full_restore'

srun python train_dynamics.py \
    --batch_size 48 \
    --restore_ckpt \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[tab] experiments/dynamics_grain_tok_lam_restore.sh (shellscript)
#!/usr/bin/env bash
source .venv/bin/activate

data_dir='data_arrayrecord/dummy'
ckpt_dir='checkpoints/dynamics_openai_grain_tok_lam_restore'
tokenizer_ckpt_dir='checkpoints/tokenizer_openai_grain_checkpointing'
lam_ckpt_dir='checkpoints/lam_openai_grain_checkpointing'

srun python train_dynamics.py \
    --batch_size 48 \
    --tokenizer_checkpoint $tokenizer_ckpt_dir \
    --lam_checkpoint $lam_ckpt_dir \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[content] genie.py (python)
            args=ocp.args.Composite(
                model_state=ocp.args.StandardRestore(abstract_lam_state),
                dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),
            ),
        )["model_state"]
        restored_lam_params = restored_lam.params["params"]
[content] genie.py (python)
        dummy_lam_train_state = TrainState.create(
            apply_fn=dummy_lam.apply, params=lam_init_params, tx=dummy_tx
        )
        abstract_lam_state = jax.tree_util.tree_map(ocp.utils.to_shape_dtype_struct, dummy_lam_train_state)
        restored_lam = lam_checkpoint_manager.restore(
[content] genie.py (python)
    # dummy values since we only use tx to initialize the dummy train states
    dummy_tx = optax.adamw(
        learning_rate=optax.constant_schedule(args.max_lr),
        b1=0.9,
        b2=0.9,
        weight_decay=1e-4,
    )
    handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
    handler_registry.add('model_state', ocp.args.StandardRestore, ocp.handlers.StandardCheckpointHandler)
    handler_registry.add('dataloader_state', grain.checkpoint.CheckpointRestore, grain.checkpoint.CheckpointHandler)
    
    tokenizer_checkpoint_manager = ocp.CheckpointManager(
        directory=args.tokenizer_checkpoint,
        handler_registry=handler_registry,
    )
    dummy_tokenizer = TokenizerVQVAE(
        in_dim=args.image_channels,
        model_dim=args.tokenizer_dim,
        latent_dim=args.latent_patch_dim,
        num_latents=args.num_patch_latents,
        patch_size=args.patch_size,
        num_blocks=args.tokenizer_num_blocks,
        num_heads=args.tokenizer_num_heads,
        dropout=args.dropout,
        codebook_dropout=args.dropout,
    )
    tokenizer_init_params = dummy_tokenizer.init(_rng, inputs)
    dummy_tokenizer_train_state = TrainState.create(
        apply_fn=dummy_tokenizer.apply, params=tokenizer_init_params, tx=dummy_tx
    )
    abstract_tokenizer_state = jax.tree_util.tree_map(ocp.utils.to_shape_dtype_struct, dummy_tokenizer_train_state)
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
        step=tokenizer_checkpoint_manager.latest_step(),
        args=ocp.args.Composite(
            model_state=ocp.args.StandardRestore(abstract_tokenizer_state),
            dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),
        ),
    )["model_state"]
    restored_tokenizer_params = restored_tokenizer.params["params"]
    train_state.params["params"]["tokenizer"].update(restored_tokenizer_params)
    tokenizer_checkpoint_manager.close()
[content] genie.py (python)
        
        # Create dummy train state for restoration
        dummy_tx = optax.adamw(
            learning_rate=optax.constant_schedule(args.max_lr),
            b1=0.9,
            b2=0.9,
            weight_decay=1e-4,
        )
[content] genie.py (python)
        # Create handler registry like in train_tokenizer.py
        handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
        handler_registry.add('model_state', ocp.args.StandardRestore, ocp.handlers.StandardCheckpointHandler)
        handler_registry.add('dataloader_state', grain.checkpoint.CheckpointRestore, grain.checkpoint.CheckpointHandler)
[content] genie.py (python)
    # Restore tokenizer
    if args.tokenizer_checkpoint:
        # Create handler registry like in train_tokenizer.py
        handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
        handler_registry.add('model_state', ocp.args.StandardRestore, ocp.handlers.StandardCheckpointHandler)
        handler_registry.add('dataloader_state', grain.checkpoint.CheckpointRestore, grain.checkpoint.CheckpointHandler)
        
        tokenizer_checkpoint_manager = ocp.CheckpointManager(
            directory=args.tokenizer_checkpoint,
            handler_registry=handler_registry,
        )
        
        # Create dummy tokenizer to get the parameter structure
        dummy_tokenizer = TokenizerVQVAE(
            in_dim=args.image_channels,
            model_dim=args.tokenizer_dim,
            latent_dim=args.latent_patch_dim,
            num_latents=args.num_patch_latents,
            patch_size=args.patch_size,
            num_blocks=args.tokenizer_num_blocks,
            num_heads=args.tokenizer_num_heads,
            dropout=args.dropout,
            codebook_dropout=args.dropout,
        )
        tokenizer_init_params = dummy_tokenizer.init(_rng, inputs)
        
        # Create dummy train state for restoration
        dummy_tx = optax.adamw(
            learning_rate=optax.constant_schedule(args.max_lr),
            b1=0.9,
            b2=0.9,
            weight_decay=1e-4,
        )
        dummy_tokenizer_train_state = TrainState.create(
            apply_fn=dummy_tokenizer.apply, params=tokenizer_init_params, tx=dummy_tx
        )
        abstract_tokenizer_state = jax.tree_util.tree_map(ocp.utils.to_shape_dtype_struct, dummy_tokenizer_train_state)
        
        # Restore using the same approach as train_tokenizer.py
        restored = tokenizer_checkpoint_manager.restore(
            step=tokenizer_checkpoint_manager.latest_step(),
            args=ocp.args.Composite(
                model_state=ocp.args.StandardRestore(abstract_tokenizer_state),
                dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),
            )
        )
        
        # Update tokenizer parameters in the train state
        restored_tokenizer_params = restored["model_state"].params["params"]
        train_state.params["params"]["tokenizer"].update(restored_tokenizer_params)
        tokenizer_checkpoint_manager.close()
[content] genie.py (python)
        # Create handler registry exactly like in train_tokenizer.py
        handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
        handler_registry.add('model_state', ocp.args.StandardSave, ocp.handlers.StandardCheckpointHandler)
        handler_registry.add('model_state', ocp.args.StandardRestore, ocp.handlers.StandardCheckpointHandler)
        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] genie.py (python)
    # Restore tokenizer
    if args.tokenizer_checkpoint:
        # Create handler registry exactly like in train_tokenizer.py
        handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
        handler_registry.add('model_state', ocp.args.StandardSave, ocp.handlers.StandardCheckpointHandler)
        handler_registry.add('model_state', ocp.args.StandardRestore, ocp.handlers.StandardCheckpointHandler)
        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
        
        tokenizer_checkpoint_manager = ocp.CheckpointManager(
            directory=args.tokenizer_checkpoint,
            handler_registry=handler_registry,
        )
        
        # Create dummy tokenizer to get the parameter structure
        dummy_tokenizer = TokenizerVQVAE(
            in_dim=args.image_channels,
            model_dim=args.tokenizer_dim,
            latent_dim=args.latent_patch_dim,
            num_latents=args.num_patch_latents,
            patch_size=args.patch_size,
            num_blocks=args.tokenizer_num_blocks,
            num_heads=args.tokenizer_num_heads,
            dropout=args.dropout,
            codebook_dropout=args.dropout,
        )
        tokenizer_init_params = dummy_tokenizer.init(_rng, inputs)
        
        # Create dummy train state for restoration
        dummy_tx = optax.adamw(
            learning_rate=optax.constant_schedule(args.max_lr),
            b1=0.9,
            b2=0.9,
            weight_decay=1e-4,
        )
        dummy_tokenizer_train_state = TrainState.create(
            apply_fn=dummy_tokenizer.apply, params=tokenizer_init_params, tx=dummy_tx
        )
        abstract_tokenizer_state = jax.tree_util.tree_map(ocp.utils.to_shape_dtype_struct, dummy_tokenizer_train_state)
        
        # Restore using the exact same approach as train_tokenizer.py
        restored = tokenizer_checkpoint_manager.restore(
            tokenizer_checkpoint_manager.latest_step(),
            args=ocp.args.Composite(
                model_state=ocp.args.StandardRestore(abstract_tokenizer_state),
                dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),
            )
        )
        
        # Update tokenizer parameters in the train state
        restored_tokenizer_params = restored["model_state"].params["params"]
        train_state.params["params"]["tokenizer"].update(restored_tokenizer_params)
        tokenizer_checkpoint_manager.close()
[content] genie.py (python)
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
        step=tokenizer_checkpoint_manager.latest_step(),
        args=ocp.args.Composite(
            model_state=ocp.args.StandardRestore(abstract_tokenizer_state),
            dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),
        ),
    )["model_state"]
    restored_tokenizer_params = restored_tokenizer.params["params"]
[content] genie.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
    
    checkpoint_options = ocp.CheckpointManagerOptions(
        save_interval_steps=10000,  # Not used for restoration, but needed for compatibility
        max_to_keep=3,
        step_format_fixed_length=6,
        cleanup_tmp_directories=True,
    )
[content] genie.py (python)

    checkpoint_options = ocp.CheckpointManagerOptions(
        save_interval_steps=args.log_checkpoint_interval,
        max_to_keep=3,
        step_format_fixed_length=6,
        cleanup_tmp_directories=True,
    )
[content] genie.py (python)
def _create_abstract_sharded_pytree(pytree_template, sharding_spec):
    """Replaces arrays in a pytree with ShapeDtypeStructs having the given sharding."""

    def map_fn(leaf_template):
        if hasattr(leaf_template, "shape") and hasattr(leaf_template, "dtype"):
            return jax.ShapeDtypeStruct(
                leaf_template.shape, leaf_template.dtype, sharding=sharding_spec
            )
        return leaf_template

    return jax.tree_util.tree_map(map_fn, pytree_template)
[content] checkpoints/lam_grain_checkpointing.sh (shellscript)
#!/usr/bin/env bash
source .venv/bin/activate

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

export XLA_FLAGS=--xla_gpu_autotune_level=0
srun python train_tokenizer.py \
    --batch_size 48 \
    --restore_ckpt \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[tab] experiments/lam_grain_checkpointing.sh (shellscript)
#!/usr/bin/env bash
source .venv/bin/activate

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

export XLA_FLAGS=--xla_gpu_autotune_level=0
srun python train_lam.py \
    --batch_size 48 \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[content] experiments/tokenizer_optimal_batch_size.sh (shellscript)
#!/usr/bin/env bash
source .venv/bin/activate

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

export XLA_FLAGS=--xla_gpu_autotune_level=0
srun python train_tokenizer.py \
    --batch_size 48 \
    --restore_ckpt \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir