p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] genie.py (python)
    abstract_tokenizer_state = jax.tree_util.tree_map(ocp.utils.to_shape_dtype_struct, dummy_tokenizer_train_state)
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
        tokenizer_checkpoint_manager.latest_step(),
        args=ocp.args.StandardRestore(abstract_tokenizer_state),
    )["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)
    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)
    
    tokenizer_checkpoint_manager = ocp.CheckpointManager(
        directory=os.path.abspath(args.tokenizer_checkpoint),
        handler_registry=handler_registry,
    )
[content] genie.py (python)
        abstract_lam_state = jax.tree_util.tree_map(ocp.utils.to_shape_dtype_struct, dummy_lam_train_state)
        restored_lam = lam_checkpoint_manager.restore(
            lam_checkpoint_manager.latest_step(),
            args=ocp.args.StandardRestore(abstract_lam_state),
        )["model_state"]
        restored_lam_params = restored_lam.params["params"]
[content] genie.py (python)
    
    # Convert FrozenDict to regular dict for updating
    current_tokenizer_params = dict(train_state.params["params"]["tokenizer"])
    restored_tokenizer_params = dict(restored_tokenizer.params["params"])
    current_tokenizer_params.update(restored_tokenizer_params)
    train_state.params["params"]["tokenizer"] = current_tokenizer_params
[content] genie.py (python)
        
        # Convert FrozenDict to regular dict for updating
        current_lam_params = dict(train_state.params["params"]["lam"])
        current_lam_params.update(restored_lam_params)
        train_state.params["params"]["lam"] = current_lam_params
[content] genie.py (python)

        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,
    ))
[content] genie.py (python)
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,
    )
[tab] experiments/tokenizer_openai_grain_checkpointing.sh (shellscript)
#!/usr/bin/env bash
module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate


data_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrecords/10fps_160x90'
ckpt_dir='checkpoints/tokenizer_openai_grain_checkpointing'

srun python train_tokenizer.py \
    --batch_size 12 \
    --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
[content] experiments/dynamics_openai_grain_tok_lam_restore.sh (shellscript)
#!/usr/bin/env bash
module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate


data_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrecords/10fps_160x90'
ckpt_dir='checkpoints/tokenizer_openai_grain_checkpointing'

srun python train_tokenizer.py \
    --batch_size 12 \
    --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