p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] sample.py (python)
# Create a dummy TrainState for checkpoint structure
dummy_train_state = TrainState.create(
    apply_fn=genie.apply,
    params=params["params"],  # or params if your model expects that
    tx=None,  # No optimizer needed for eval
)

# Set up Orbax CheckpointManager
checkpoint_manager = ocp.CheckpointManager(
    args.checkpoint,  # Directory containing the checkpoint
    options=ocp.CheckpointManagerOptions(max_to_keep=1),
)

# Prepare abstract state for restore
abstract_train_state = jax.tree_util.tree_map(
    ocp.utils.to_shape_dtype_struct, dummy_train_state
)

# Restore the checkpoint (only model_state needed for sampling)
restored = checkpoint_manager.restore(
    checkpoint_manager.latest_step(),
    args=ocp.args.Composite(
        model_state=ocp.args.StandardRestore(abstract_train_state),
    ),
)
restored_train_state = restored["model_state"]

# Extract model parameters for inference
params = restored_train_state.params
[content] sample.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)
    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] sample.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] sample.py (python)
# Set up Orbax CheckpointManager
checkpoint_manager = ocp.CheckpointManager(
    args.checkpoint,  # Directory containing the checkpoint
    options=ocp.CheckpointManagerOptions(max_to_keep=1),
    handler_registry=handler_registry
)
[content] sample.py (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,
    )
    
    checkpoint_manager = ocp.CheckpointManager(
        args.ckpt_dir,
        options=checkpoint_options,
        handler_registry=handler_registry,
    )
[content] sample.py (python)
Set up Orbax CheckpointManager
checkpoint_manager = ocp.CheckpointManager(
    args.checkpoint,  # Directory containing the checkpoint
    options=ocp.CheckpointManagerOptions(max_to_keep=1),
    handler_registry=handler_registry
)

#
[tab] scripts_horeka/train_lam.sh (shellscript)
#!/usr/bin/env bash

module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate


tf_records_dir=$ws_dir/knoms_tfrecords_500_shards
ws_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/'

job_name="debug"
slurm_job_id="0000"

CHECKPOINT_DIR=$ws_dir/checkpoints/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

python train_lam.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=1 \
    --min_lr=3e-4 \
    --max_lr=3e-4 \
    --warmup_steps=125 \
    --log_image_interval=100 \
    --log \
    --log_checkpoint_interval=5 \
    --name=lam-tiny-overfit-$slurm_job_id \
    --tags lam overfit tiny \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir \
    --model_dim=128 \
    --latent_dim=8 \
    --num_latents=4 \
    --patch_size=16 \
    --num_blocks=2 \
    --num_heads=2 \
    --codebook_dropout=0.0
[content] sample.py (python)
    # --- Create DataLoaderIterator from dataloader ---
    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)