p-doom/crowd-code-dataset-0.2 huggingface dataset:
[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'

srun python train_tokenizer.py \
    --batch_size 48 \
    --restore_ckpt \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --ckpt_dir checkpoints/tokenizer_openai_grain_checkpointing \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[tab] checkpoints/tokenizer_openai_grain_checkpointing/000005/dataloader_state/process_0-of-4.json (json)
{
    "version": 2,
    "last_seen_indices": {
        "0": 352,
        "1": 356,
        "2": 392,
        "3": 364,
        "4": 368,
        "5": -12,
        "6": -8,
        "7": -4
    },
    "last_worker_index": 4,
    "worker_count": 8,
    "sampler": "IndexSampler(num_records=57126, shard_options=ShardByJaxProcess(shard_index=0, shard_count=4, drop_remainder=True), shuffle=True, num_epochs=None, seed=0)",
    "data_source": "ArrayRecordDataSource(hash_of_paths=7fc1b8e79738369bb97d1274129d39f671dfe08f)"
}
[tab] checkpoints/tokenizer_openai_grain_checkpointing/000005/dataloader_state/process_1-of-4.json (json)
{
    "version": 2,
    "last_seen_indices": {
        "0": 353,
        "1": 357,
        "2": 361,
        "3": 365,
        "4": 369,
        "5": -11,
        "6": -7,
        "7": -3
    },
    "last_worker_index": 4,
    "worker_count": 8,
    "sampler": "IndexSampler(num_records=57126, shard_options=ShardByJaxProcess(shard_index=1, shard_count=4, drop_remainder=True), shuffle=True, num_epochs=None, seed=0)",
    "data_source": "ArrayRecordDataSource(hash_of_paths=7fc1b8e79738369bb97d1274129d39f671dfe08f)"
}
[tab] checkpoints/tokenizer_openai_grain_checkpointing/000005/dataloader_state/process_2-of-4.json (json)
{
    "version": 2,
    "last_seen_indices": {
        "0": 354,
        "1": 390,
        "2": 362,
        "3": 366,
        "4": 370,
        "5": -10,
        "6": -6,
        "7": -2
    },
    "last_worker_index": 4,
    "worker_count": 8,
    "sampler": "IndexSampler(num_records=57126, shard_options=ShardByJaxProcess(shard_index=2, shard_count=4, drop_remainder=True), shuffle=True, num_epochs=None, seed=0)",
    "data_source": "ArrayRecordDataSource(hash_of_paths=7fc1b8e79738369bb97d1274129d39f671dfe08f)"
}
[tab] checkpoints/tokenizer_openai_grain_checkpointing/000005/dataloader_state/process_3-of-4.json (json)
{
    "version": 2,
    "last_seen_indices": {
        "0": 355,
        "1": 359,
        "2": 363,
        "3": 367,
        "4": 371,
        "5": -9,
        "6": -5,
        "7": -1
    },
    "last_worker_index": 4,
    "worker_count": 8,
    "sampler": "IndexSampler(num_records=57126, shard_options=ShardByJaxProcess(shard_index=3, shard_count=4, drop_remainder=True), shuffle=True, num_epochs=None, seed=0)",
    "data_source": "ArrayRecordDataSource(hash_of_paths=7fc1b8e79738369bb97d1274129d39f671dfe08f)"
}
[tab] checkpoints/tokenizer_openai_grain_checkpointing/000010/_CHECKPOINT_METADATA (plaintext)
{"item_handlers": {"dataloader_state": "grain._src.python.checkpoint_handlers.CheckpointHandler", "model_state": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1751986813421941516, "commit_timestamp_nsecs": 1751986814906949563, "custom_metadata": {}}
[tab] checkpoints/tokenizer_openai_grain_checkpointing/000020/dataloader_state/process_0-of-1.json (json)
{
    "version": 2,
    "last_seen_indices": {
        "0": 288,
        "1": 281,
        "2": 298,
        "3": 283,
        "4": 188,
        "5": 189,
        "6": 190,
        "7": 199
    },
    "last_worker_index": 3,
    "worker_count": 8,
    "sampler": "IndexSampler(num_records=57126, shard_options=ShardByJaxProcess(shard_index=0, shard_count=1, drop_remainder=True), shuffle=True, num_epochs=None, seed=0)",
    "data_source": "ArrayRecordDataSource(hash_of_paths=7fc1b8e79738369bb97d1274129d39f671dfe08f)"
}
[tab] checkpoints/tokenizer_openai_grain_checkpointing/000035/dataloader_state/process_0-of-1.json (json)
{
    "version": 2,
    "last_seen_indices": {
        "0": 488,
        "1": 473,
        "2": 490,
        "3": 379,
        "4": 380,
        "5": 381,
        "6": 382,
        "7": 391
    },
    "last_worker_index": 2,
    "worker_count": 8,
    "sampler": "IndexSampler(num_records=57126, shard_options=ShardByJaxProcess(shard_index=0, shard_count=1, drop_remainder=True), shuffle=True, num_epochs=None, seed=0)",
    "data_source": "ArrayRecordDataSource(hash_of_paths=7fc1b8e79738369bb97d1274129d39f671dfe08f)"
}
[content] train_tokenizer.py (python)

    # --- Initialize checkpoint manager ---
    step = 0
    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)
    handler_registry.add('dataloader_state', grain.checkpoint.CheckpointRestore, grain.checkpoint.CheckpointHandler)
    
    checkpoint_options = ocp.CheckpointManagerOptions(
        save_interval_steps=args.log_checkpoint_interval,
        max_to_keep=3,
        step_format_fixed_length=6,
        cleanup_tmp_directories=True,
    )
    
    checkpoint_manager = ocp.CheckpointManager(
        os.path.abspath(args.ckpt_dir),
        options=checkpoint_options,
        handler_registry=handler_registry,
    )
[content] train_tokenizer.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] train_lam.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
    
    checkpoint_options = ocp.CheckpointManagerOptions(
        save_interval_steps=args.log_checkpoint_interval,
        max_to_keep=3,
        step_format_fixed_length=6,
        cleanup_tmp_directories=True,
    )
    
    checkpoint_manager = ocp.CheckpointManager(
        os.path.abspath(args.ckpt_dir),
        options=checkpoint_options,
        handler_registry=handler_registry,
    )
[content] train_lam.py (python)
    if args.checkpoint:
        restore_target = {"model": train_state}
        restore_args = orbax_utils.restore_args_from_target(restore_target)
        train_state.params["params"].update(
            PyTreeCheckpointer()
            .restore(args.checkpoint, item=restore_target, restore_args=restore_args)[
                "model"
            ]
            .params["params"]
        )
        # Assume checkpoint is of the form tokenizer_<timestamp>_<step>
        step += int(args.checkpoint.split("_")[-1])
[content] train_lam.py (python)
            if step % args.log_checkpoint_interval == 0:
                ckpt = {"model": train_state}
                orbax_checkpointer = orbax.checkpoint.PyTreeCheckpointer()
                save_args = orbax_utils.save_args_from_target(ckpt)
                orbax_checkpointer.save(
                    os.path.join(os.getcwd(), args.ckpt_dir, f"lam_{ts}_{step}"),
                    ckpt,
                    save_args=save_args,
                )
[content] train_lam.py (python)
    if args.checkpoint:
        restore_target = {"model": train_state}
        restore_args = orbax_utils.restore_args_from_target(restore_target)
        train_state.params["params"].update(
            PyTreeCheckpointer()
            .restore(args.checkpoint, item=restore_target, restore_args=restore_args)[
                "model"
            ]
            .params["params"]
        )
        # Assume checkpoint is of the form tokenizer_<timestamp>_<step>
        step += int(args.checkpoint.split("_")[-1])

    # --- TRAIN LOOP ---
[content] train_dynamics.py (python)
    initial_state = grain_dataloader._create_initial_state()
    grain_iterator = grain.DataLoaderIterator(grain_dataloader, initial_state)
    
    # --- Restore checkpoint ---
    train_state = restore_genie_components(
        train_state, replicated_sharding, dummy_inputs, rng, args
    )
    
    # --- Restore checkpoint if requested ---
    if args.restore_ckpt:
        abstract_train_state = jax.tree_util.tree_map(ocp.utils.to_shape_dtype_struct, train_state)
        restored = checkpoint_manager.restore(
            checkpoint_manager.latest_step(),
            args=ocp.args.Composite(
                model_state=ocp.args.StandardRestore(abstract_train_state),
                dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),
            )
        )
        train_state = restored["model_state"]
        grain_iterator = restored["dataloader_state"]
        step = checkpoint_manager.latest_step() or 0
        print(f"Restored dataloader and model state from step {step}")

    # --- TRAIN LOOP ---
    dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in grain_iterator) # type: ignore
[content] train_dynamics.py (python)
    
    checkpoint_manager = ocp.CheckpointManager(
        os.path.abspath(args.ckpt_dir),
        options=checkpoint_options,
        handler_registry=handler_registry,
    )

    # --- Create DataLoaderIterator from dataloader ---
[content] train_dynamics.py (python)
    # --- Initialize checkpoint manager ---
    step = 0
    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
    
    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] train_dynamics.py (python)
            # --- Checkpointing ---
            if args.save_ckpt and step % args.log_checkpoint_interval == 0:
                checkpoint_manager.save(
                    step,
                    args=ocp.args.Composite(
                        model_state=ocp.args.StandardSave(train_state),
                        dataloader_state=grain.checkpoint.CheckpointSave(grain_iterator),
                    )
                )
                print(f"Saved checkpoint at step {step}")
[content] train_dynamics.py (python)
            if step % args.log_checkpoint_interval == 0:
                ckpt = {"model": train_state}
                orbax_checkpointer = orbax.checkpoint.PyTreeCheckpointer()
                save_args = orbax_utils.save_args_from_target(ckpt)
                orbax_checkpointer.save(
                    os.path.join(os.getcwd(), args.ckpt_dir, f"genie_{step}"),
                    ckpt,
                    save_args=save_args,
                )
[content] train_dynamics.py (python)
    # --- Initialize checkpoint manager ---
    step = 0
    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
    
    checkpoint_options = ocp.CheckpointManagerOptions(
        save_interval_steps=args.log_checkpoint_interval,
        max_to_keep=3,
        step_format_fixed_length=6,
        cleanup_tmp_directories=True,
    )
    
    checkpoint_manager = ocp.CheckpointManager(
        os.path.abspath(args.ckpt_dir),
        options=checkpoint_options,
        handler_registry=handler_registry,
[content] genie.py (python)
        try:
            # Try new format first
            restored_lam_params = restore_from_new_format(args.lam_checkpoint)
        except Exception:
            # Fall back to old format
            dummy_lam = LatentActionModel(
                in_dim=args.image_channels,
                model_dim=args.lam_dim,
                latent_dim=args.latent_patch_dim,
                num_latents=args.num_latent_actions,
                patch_size=args.lam_patch_size,
                num_blocks=args.lam_num_blocks,
                num_heads=args.lam_num_heads,
                dropout=args.dropout,
                codebook_dropout=args.dropout,
            )
            lam_init_params = dummy_lam.init(_rng, inputs)
            dummy_lam_train_state = TrainState.create(
                apply_fn=dummy_lam.apply, params=lam_init_params, tx=dummy_tx
            )
            abstract_sharded_lam_state = _create_abstract_sharded_pytree(
                dummy_lam_train_state, sharding
            )
            lam_restore_target = {"model": abstract_sharded_lam_state}
            lam_restore_args = orbax_utils.restore_args_from_target(lam_restore_target)
            restored_lam_params = (
                PyTreeCheckpointer()
                .restore(
                    args.lam_checkpoint, item=lam_restore_target, restore_args=lam_restore_args
                )["model"]
                .params["params"]
            )
[content] genie.py (python)
        dummy_lam = LatentActionModel(
            in_dim=args.image_channels,
            model_dim=args.lam_dim,
            latent_dim=args.latent_patch_dim,
            num_latents=args.num_latent_actions,
            patch_size=args.lam_patch_size,
            num_blocks=args.lam_num_blocks,
            num_heads=args.lam_num_heads,
            dropout=args.dropout,
            codebook_dropout=args.dropout,
        )
        lam_init_params = dummy_lam.init(_rng, inputs)
        dummy_lam_train_state = TrainState.create(
            apply_fn=dummy_lam.apply, params=lam_init_params, tx=dummy_tx
        )
        abstract_sharded_lam_state = _create_abstract_sharded_pytree(
            dummy_lam_train_state, sharding
        )
        lam_restore_target = {"model": abstract_sharded_lam_state}
        lam_restore_args = orbax_utils.restore_args_from_target(lam_restore_target)
        restored_lam_params = (
            PyTreeCheckpointer()
            .restore(
                args.lam_checkpoint, item=lam_restore_target, restore_args=lam_restore_args
            )["model"]
            .params["params"]
        )
[content] genie.py (python)
    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_sharded_tokenizer_state = _create_abstract_sharded_pytree(
        dummy_tokenizer_train_state, sharding
    )
    tokenizer_restore_target = {"model": abstract_sharded_tokenizer_state}
    tokenizer_restore_args = orbax_utils.restore_args_from_target(
        tokenizer_restore_target
    )
    restored_tokenizer_params = (
        PyTreeCheckpointer()
        .restore(
            args.tokenizer_checkpoint,
            item=tokenizer_restore_target,
            restore_args=tokenizer_restore_args,
        )["model"]
        .params["params"]
    )
    train_state.params["params"]["tokenizer"].update(restored_tokenizer_params)
[content] genie.py (python)
        return restored["model_state"].params["params"]

    # Restore tokenizer
    if args.tokenizer_checkpoint:
        restored_tokenizer_params = restore_component(args.tokenizer_checkpoint, "tokenizer")
        train_state.params["params"]["tokenizer"].update(restored_tokenizer_params)

    # Restore LAM if checkpoint provided
    if args.lam_checkpoint:
        restored_lam_params = restore_component(args.lam_checkpoint, "lam")
[content] genie.py (python)
        
        # Get the step from the checkpoint path
        checkpoint_name = os.path.basename(checkpoint_path)
        step = int(checkpoint_name.split('_')[-1]) if '_' in checkpoint_name else 0
        
        abstract_train_state = jax.tree_util.tree_map(ocp.utils.to_shape_dtype_struct, dummy_component_train_state)
        restored = checkpoint_manager.restore(
            step,
            args=ocp.args.StandardRestore(abstract_train_state),
[content] genie.py (python)
        
        # Create checkpoint manager to restore
        handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
        handler_registry.add('model_state', ocp.args.StandardRestore, ocp.handlers.StandardCheckpointHandler)
        
        checkpoint_manager = ocp.CheckpointManager(
            os.path.dirname(checkpoint_path),
            options=ocp.CheckpointManagerOptions(),
            handler_registry=handler_registry,
[content] genie.py (python)
    def restore_component(checkpoint_path, component_type):
        """Restore a component from composite checkpoint manager format"""
        if component_type == "tokenizer":
            dummy_component = 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,
            )
        elif component_type == "lam":
            dummy_component = LatentActionModel(
                in_dim=args.image_channels,
                model_dim=args.lam_dim,
                latent_dim=args.latent_patch_dim,
                num_latents=args.num_latent_actions,
                patch_size=args.lam_patch_size,
                num_blocks=args.lam_num_blocks,
                num_heads=args.lam_num_heads,
                dropout=args.dropout,
                codebook_dropout=args.dropout,
            )
        else:
            raise ValueError(f"Unknown component type: {component_type}")
        
        component_init_params = dummy_component.init(_rng, inputs)
        dummy_component_train_state = TrainState.create(
            apply_fn=dummy_component.apply, params=component_init_params, tx=dummy_tx
[content] genie.py (python)
        lam_restore_target = {"model": abstract_sharded_lam_state}
        lam_restore_args = orbax_utils.restore_args_from_target(lam_restore_target)
        restored_lam_params = (
            PyTreeCheckpointer()
            .restore(
                args.lam_checkpoint, item=lam_restore_target, restore_args=lam_restore_args
            )["model"]
            .params["params"]
[content] genie.py (python)
    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_sharded_tokenizer_state = _create_abstract_sharded_pytree(
        dummy_tokenizer_train_state, sharding
    )
    tokenizer_restore_target = {"model": abstract_sharded_tokenizer_state}
    tokenizer_restore_args = orbax_utils.restore_args_from_target(
        tokenizer_restore_target
    )
    restored_tokenizer_params = (
        PyTreeCheckpointer()
        .restore(
            args.tokenizer_checkpoint,
            item=tokenizer_restore_target,
            restore_args=tokenizer_restore_args,
        )["model"]
        .params["params"]
    )
    train_state.params["params"]["tokenizer"].update(restored_tokenizer_params)

    if args.lam_checkpoint:
        dummy_lam = LatentActionModel(
            in_dim=args.image_channels,
            model_dim=args.lam_dim,
            latent_dim=args.latent_patch_dim,
            num_latents=args.num_latent_actions,
            patch_size=args.lam_patch_size,
            num_blocks=args.lam_num_blocks,
            num_heads=args.lam_num_heads,
            dropout=args.dropout,
            codebook_dropout=args.dropout,
        )
        lam_init_params = dummy_lam.init(_rng, inputs)
        dummy_lam_train_state = TrainState.create(
            apply_fn=dummy_lam.apply, params=lam_init_params, tx=dummy_tx
[content] genie.py (python)
        
        # Try to restore using new CheckpointManager format first
        try:
            # Create a checkpoint manager to check if the checkpoint uses the new format
            checkpoint_manager = ocp.CheckpointManager(
                args.lam_checkpoint,
                options=ocp.CheckpointManagerOptions(max_to_keep=1),
            )
            
            # Try to restore using the new format
            abstract_lam_state = jax.tree_util.tree_map(ocp.utils.to_shape_dtype_struct, dummy_lam_train_state)
            restored = checkpoint_manager.restore(
                checkpoint_manager.latest_step(),
                args=ocp.args.Composite(
                    model_state=ocp.args.StandardRestore(abstract_lam_state),
                )
            )
            restored_lam_params = restored["model_state"].params["params"]
        except (ValueError, FileNotFoundError, KeyError):
            # Fall back to old PyTreeCheckpointer format
            abstract_sharded_lam_state = _create_abstract_sharded_pytree(
                dummy_lam_train_state, sharding
            )
            lam_restore_target = {"model": abstract_sharded_lam_state}
            lam_restore_args = orbax_utils.restore_args_from_target(lam_restore_target)
            restored_lam_params = (
                PyTreeCheckpointer()
                .restore(
                    args.lam_checkpoint, item=lam_restore_target, restore_args=lam_restore_args
                )["model"]
                .params["params"]
            )
[content] genie.py (python)
    
    # Try to restore using new CheckpointManager format first
    try:
        # Create a checkpoint manager to check if the checkpoint uses the new format
        checkpoint_manager = ocp.CheckpointManager(
            args.tokenizer_checkpoint,
            options=ocp.CheckpointManagerOptions(max_to_keep=1),
        )
        
        # Try to restore using the new format
        abstract_tokenizer_state = jax.tree_util.tree_map(ocp.utils.to_shape_dtype_struct, dummy_tokenizer_train_state)
        restored = checkpoint_manager.restore(
            checkpoint_manager.latest_step(),
            args=ocp.args.Composite(
                model_state=ocp.args.StandardRestore(abstract_tokenizer_state),
            )
        )
        restored_tokenizer_params = restored["model_state"].params["params"]
    except (ValueError, FileNotFoundError, KeyError):
        # Fall back to old PyTreeCheckpointer format
        abstract_sharded_tokenizer_state = _create_abstract_sharded_pytree(
            dummy_tokenizer_train_state, sharding
        )
        tokenizer_restore_target = {"model": abstract_sharded_tokenizer_state}
        tokenizer_restore_args = orbax_utils.restore_args_from_target(
            tokenizer_restore_target
        )
        restored_tokenizer_params = (
            PyTreeCheckpointer()
            .restore(
                args.tokenizer_checkpoint,
                item=tokenizer_restore_target,
                restore_args=tokenizer_restore_args,
            )["model"]
            .params["params"]
        )
    
    # Update tokenizer parameters
    train_state.params["params"]["tokenizer"] = update_params(
        train_state.params["params"]["tokenizer"], 
        restored_tokenizer_params
    )
[content] genie.py (python)
    # Helper function to update FrozenDict parameters
    def update_params(target_params, source_params):
        return jax.tree_util.tree_map(
            lambda target, source: source if source is not None else target,
            target_params,
            source_params,
            is_leaf=lambda x: x is None
        )

    # Restore tokenizer
[content] genie.py (python)
        abstract_sharded_lam_state = _create_abstract_sharded_pytree(
            dummy_lam_train_state, sharding
        )
        lam_restore_target = {"model": abstract_sharded_lam_state}
        lam_restore_args = orbax_utils.restore_args_from_target(lam_restore_target)
        restored_lam_params = (
            PyTreeCheckpointer()
            .restore(
                args.lam_checkpoint, item=lam_restore_target, restore_args=lam_restore_args
            )["model"]
            .params["params"]
        )
[content] genie.py (python)
    abstract_sharded_tokenizer_state = _create_abstract_sharded_pytree(
        dummy_tokenizer_train_state, sharding
    )
    tokenizer_restore_target = {"model": abstract_sharded_tokenizer_state}
    tokenizer_restore_args = orbax_utils.restore_args_from_target(
        tokenizer_restore_target
    )
    restored_tokenizer_params = (
        PyTreeCheckpointer()
        .restore(
            args.tokenizer_checkpoint,
            item=tokenizer_restore_target,
            restore_args=tokenizer_restore_args,
        )["model"]
        .params["params"]
    )
    train_state.params["params"]["tokenizer"].update(restored_tokenizer_params)
[content] genie.py (python)
    
    # Try to restore using new CheckpointManager format first
    try:
        # Create a checkpoint manager to check if the checkpoint uses the new format
        checkpoint_manager = ocp.CheckpointManager(
            args.tokenizer_checkpoint,
            options=ocp.CheckpointManagerOptions(max_to_keep=1),
        )
        
        # Try to restore using the new format
        abstract_tokenizer_state = jax.tree_util.tree_map(ocp.utils.to_shape_dtype_struct, dummy_tokenizer_train_state)
        restored = checkpoint_manager.restore(
            checkpoint_manager.latest_step(),
            args=ocp.args.Composite(
                model_state=ocp.args.StandardRestore(abstract_tokenizer_state),
            )
        )
        restored_tokenizer_params = restored["model_state"].params["params"]
    except (ValueError, FileNotFoundError, KeyError, OSError):
        # Fall back to old PyTreeCheckpointer format
        abstract_sharded_tokenizer_state = _create_abstract_sharded_pytree(
            dummy_tokenizer_train_state, sharding
        )
        tokenizer_restore_target = {"model": abstract_sharded_tokenizer_state}
        tokenizer_restore_args = orbax_utils.restore_args_from_target(
            tokenizer_restore_target
        )
        restored_tokenizer_params = (
            PyTreeCheckpointer()
            .restore(
                args.tokenizer_checkpoint,
                item=tokenizer_restore_target,
                restore_args=tokenizer_restore_args,
            )["model"]
            .params["params"]
        )
    
    # Update tokenizer parameters
    train_state.params["params"]["tokenizer"] = update_params(
        train_state.params["params"]["tokenizer"], 
        restored_tokenizer_params
    )
[content] genie.py (python)
        abstract_lam_state = jax.tree_util.tree_map(ocp.utils.to_shape_dtype_struct, dummy_lam_train_state)
        
        # Restore LAM using new checkpoint format
        restored_lam = lam_checkpoint_manager.restore(
            os.path.basename(args.lam_checkpoint),
            args=ocp.args.StandardRestore(abstract_lam_state),
        )["model_state"]
[content] genie.py (python)
    abstract_tokenizer_state = jax.tree_util.tree_map(ocp.utils.to_shape_dtype_struct, dummy_tokenizer_train_state)
    
    # Restore tokenizer using new checkpoint format
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
        os.path.basename(args.tokenizer_checkpoint),
        args=ocp.args.StandardRestore(abstract_tokenizer_state),
    )["model_state"]
    
    # 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

    if args.lam_checkpoint:
        # Initialize checkpoint manager for LAM restoration
        lam_checkpoint_manager = ocp.CheckpointManager(
            os.path.dirname(args.lam_checkpoint),
            handler_registry=handler_registry,
        )
[content] genie.py (python)
    # Initialize checkpoint manager for tokenizer restoration
    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(
        os.path.dirname(args.tokenizer_checkpoint),
        handler_registry=handler_registry,
    )
[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] genie.py (python)
    abstract_sharded_tokenizer_state = _create_abstract_sharded_pytree(
        dummy_tokenizer_train_state, sharding
    )
    tokenizer_restore_target = {"model": abstract_sharded_tokenizer_state}
    tokenizer_restore_args = orbax_utils.restore_args_from_target(
        tokenizer_restore_target
    )
    restored_tokenizer_params = (
        PyTreeCheckpointer()
        .restore(
            args.tokenizer_checkpoint,
            item=tokenizer_restore_target,
            restore_args=tokenizer_restore_args,
        )["model"]
        .params["params"]
    )
    train_state.params["params"]["tokenizer"].update(restored_tokenizer_params)

    if args.lam_checkpoint:
[content] genie.py (python)
    # Initialize checkpoint manager for tokenizer restoration
    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.dirname(args.tokenizer_checkpoint),
        handler_registry=handler_registry,
    )
[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] genie.py (python)
    abstract_tokenizer_state = jax.tree_util.tree_map(ocp.utils.to_shape_dtype_struct, dummy_tokenizer_train_state)
    
    # Restore tokenizer using new checkpoint format
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
        os.path.basename(args.tokenizer_checkpoint),
        args=ocp.args.StandardRestore(abstract_tokenizer_state),
    )["model_state"]
    
    # 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

    if args.lam_checkpoint:
        # Initialize checkpoint manager for LAM restoration
        lam_checkpoint_manager = ocp.CheckpointManager(
            directory=os.path.dirname(args.lam_checkpoint),
            handler_registry=handler_registry,
        )