p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] experiments/tokenizer_knoms_time_measurement_grain.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/knoms_arrayrecords_500_shards_optimized_layout'

srun python train_tokenizer.py \
    --time_measurement_interval 10 \
    --batch_size 48 \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[content] experiments/tokenizer_openai_time_measurement_grain.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/knoms_arrayrecords_500_shards_optimized_layout0'

srun python train_tokenizer.py \
    --time_measurement_interval 10 \
    --batch_size 48 \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[content] utils/dataloader.py (python)
            # Instead of returning None, we'll pad the episode to the required length
            # This is a simple approach that repeats the last frame
            padding_needed = self.seq_len - current_episode_len
            last_frame = episode_tensor[-1:]
            padding_frames = np.repeat(last_frame, padding_needed, axis=0)
            episode_tensor = np.concatenate([episode_tensor, padding_frames], axis=0)
            current_episode_len = episode_tensor.shape[0]
[content] utils/dataloader.py (python)
class EpisodeLengthFilter(grain.transforms.FilterTransform):
    """
    A Grain FilterTransform that filters out episodes shorter than the required sequence length.
    """
    
    def __init__(self, seq_len: int, image_h: int, image_w: int, image_c: int):
        """Initializes the filter with processing parameters."""
        self.seq_len = seq_len
        self.image_h = image_h
        self.image_w = image_w
        self.image_c = image_c
    
    def filter(self, element: bytes) -> bool:
        """
        Filters episodes based on their length.
        
        Args:
            element: A bytes object representing one record from the DataSource.
            
        Returns:
            True if the episode should be kept, False if it should be filtered out.
        """
        try:
            record = pickle.loads(element)
            video_shape = (
                record["sequence_length"],
                self.image_h,
                self.image_w,
                self.image_c,
            )
            episode_tensor = np.frombuffer(record["raw_video"], dtype=np.uint8)
            episode_tensor = episode_tensor.reshape(video_shape)
            
            current_episode_len = episode_tensor.shape[0]
            if current_episode_len < self.seq_len:
                print(f"Filtering out episode with length {current_episode_len}, which is "
                      f"shorter than the requested sequence length {self.seq_len}.")
                return False
            return True
        except Exception as e:
            print(f"Error filtering episode: {e}")
            return False
[content] utils/dataloader.py (python)
class EpisodeLengthFilter(grain.transforms.MapTransform):
    """
    A Grain MapTransform that filters out episodes shorter than the required sequence length.
    Returns None for episodes that should be filtered out.
    """
    
    def __init__(self, seq_len: int, image_h: int, image_w: int, image_c: int):
        """Initializes the filter with processing parameters."""
        self.seq_len = seq_len
        self.image_h = image_h
        self.image_w = image_w
        self.image_c = image_c
    
    def map(self, element: bytes) -> bytes | None:
        """
        Filters episodes based on their length.
        
        Args:
            element: A bytes object representing one record from the DataSource.
            
        Returns:
            The original element if it should be kept, None if it should be filtered out.
        """
        try:
            record = pickle.loads(element)
            video_shape = (
                record["sequence_length"],
                self.image_h,
                self.image_w,
                self.image_c,
            )
            episode_tensor = np.frombuffer(record["raw_video"], dtype=np.uint8)
            episode_tensor = episode_tensor.reshape(video_shape)
            
            current_episode_len = episode_tensor.shape[0]
            if current_episode_len < self.seq_len:
                print(f"Filtering out episode with length {current_episode_len}, which is "
                      f"shorter than the requested sequence length {self.seq_len}.")
                return None
            return element
        except Exception as e:
            print(f"Error filtering episode: {e}")
            return None
[content] utils/dataloader.py (python)
class EpisodeLengthFilter:
    """
    A custom filter that filters out episodes shorter than the required sequence length.
    This is used as a pre-processing step before the main transformation.
    """
    
    def __init__(self, seq_len: int, image_h: int, image_w: int, image_c: int):
        """Initializes the filter with processing parameters."""
        self.seq_len = seq_len
        self.image_h = image_h
        self.image_w = image_w
        self.image_c = image_c
    
    def __call__(self, element: bytes) -> bool:
        """
        Filters episodes based on their length.
        
        Args:
            element: A bytes object representing one record from the DataSource.
            
        Returns:
            True if the episode should be kept, False if it should be filtered out.
        """
        try:
            record = pickle.loads(element)
            video_shape = (
                record["sequence_length"],
                self.image_h,
                self.image_w,
                self.image_c,
            )
            episode_tensor = np.frombuffer(record["raw_video"], dtype=np.uint8)
            episode_tensor = episode_tensor.reshape(video_shape)
            
            current_episode_len = episode_tensor.shape[0]
            if current_episode_len < self.seq_len:
                print(f"Filtering out episode with length {current_episode_len}, which is "
                      f"shorter than the requested sequence length {self.seq_len}.")
                return False
            return True
        except Exception as e:
            print(f"Error filtering episode: {e}")
            return False
[content] utils/dataloader.py (python)
class EpisodeLengthFilter:
    """
    A custom filter that filters out episodes shorter than the required sequence length.
    This is used as a pre-processing step before the main transformation.
    """
    
    def __init__(self, seq_len: int, image_h: int, image_w: int, image_c: int):
        """Initializes the filter with processing parameters."""
        self.seq_len = seq_len
        self.image_h = image_h
        self.image_w = image_w
        self.image_c = image_c
    
    def __call__(self, element: bytes) -> bool:
        """
        Filters episodes based on their length.
        
        Args:
            element: A bytes object representing one record from the DataSource.
            
        Returns:
            True if the episode should be kept, False if it should be filtered out.
        """
        try:
            record = pickle.loads(element)
            video_shape = (
                record["sequence_length"],  # type: ignore
                self.image_h,
                self.image_w,
                self.image_c,
            )
            episode_tensor = np.frombuffer(record["raw_video"], dtype=np.uint8)  # type: ignore
            episode_tensor = episode_tensor.reshape(video_shape)
            
            current_episode_len = episode_tensor.shape[0]
            if current_episode_len < self.seq_len:
                print(f"Filtering out episode with length {current_episode_len}, which is "
                      f"shorter than the requested sequence length {self.seq_len}.")
                return False
            return True
        except Exception as e:
            print(f"Error filtering episode: {e}")
            return False
[content] utils/dataloader.py (python)
            raise ValueError(f"Episode length {current_episode_len} is shorter than "
                           f"requested sequence length {self.seq_len}. This should "
                           f"have been filtered out.")
[content] utils/dataloader.py (python)
class EpisodeLengthFilter(grain.transforms.Filter):
    """
    A Grain Filter that keeps only episodes with sufficient length.
    """
    
    def __init__(self, seq_len: int, image_h: int, image_w: int, image_c: int):
        """Initializes the filter with sequence length requirements."""
        self.seq_len = seq_len
        self.image_h = image_h
        self.image_w = image_w
        self.image_c = image_c
    
    def filter(self, element: Any) -> bool:
        """
        Filters episodes based on length.
        
        Args:
            element: A dictionary representing one record from the DataSource.
                     Expected to contain 'raw_video' (bytes) and 'sequence_length' (int)
        
        Returns:
            True if the episode has sufficient length, False otherwise.
        """
        assert isinstance(element, bytes)
        element = pickle.loads(element)
        
        current_episode_len = element["sequence_length"]
        if current_episode_len < self.seq_len:
            print(f"Filtering out episode with length {current_episode_len}, which is "
                f"shorter than the requested sequence length {self.seq_len}.")
            return False
        
        return True
[content] train_tokenizer.py (python)
                orbax_checkpointer.save(ckpt_path, ckpt, save_args=save_args)
                
                # Save dataloader state
                grain_checkpoint_handler.save(
                    os.path.join(ckpt_path, "dataloader"), 
                    item=dataloader
[content] train_tokenizer.py (python)
    
    # Create Grain checkpoint handler
    grain_checkpoint_handler = CheckpointHandler()
    
    # Restore dataloader state if resuming from checkpoint
    if args.checkpoint and os.path.exists(os.path.join(args.checkpoint, "dataloader")):
        dataloader = grain_checkpoint_handler.restore(
            os.path.join(args.checkpoint, "dataloader"), 
            item=dataloader
        )
[content] train_tokenizer.py (python)
                orbax_checkpointer.save(ckpt_path, ckpt, save_args=save_args)
                
                # Save dataloader state
                grain_checkpoint_handler.save(
                    Path(os.path.join(ckpt_path, "dataloader")), 
                    item=grain_dataloader
[content] train_tokenizer.py (python)
    
    # Create Grain checkpoint handler
    grain_checkpoint_handler = CheckpointHandler()
    
    # Restore dataloader state if resuming from checkpoint
    if args.checkpoint and os.path.exists(os.path.join(args.checkpoint, "dataloader")):
        grain_dataloader = grain_checkpoint_handler.restore(
            Path(os.path.join(args.checkpoint, "dataloader")), 
            item=grain_dataloader
        )
    
    # Convert to JAX arrays
    dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in grain_dataloader) # type: ignore
[content] train_tokenizer.py (python)
                orbax_checkpointer.save(ckpt_path, ckpt, save_args=save_args)
                
                # Save dataloader state
                grain_checkpoint_handler.save(
                    os.path.join(ckpt_path, "dataloader"), 
                    item=grain_dataloader
[content] train_tokenizer.py (python)
    
    # Create Grain checkpoint handler
    grain_checkpoint_handler = CheckpointHandler()
    
    # Restore dataloader state if resuming from checkpoint
    if args.checkpoint and os.path.exists(os.path.join(args.checkpoint, "dataloader")):
        grain_dataloader = grain_checkpoint_handler.restore(
            os.path.join(args.checkpoint, "dataloader"), 
            item=grain_dataloader
        )
    
    # Convert to JAX arrays
    dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in grain_dataloader) # type: ignore
[content] train_tokenizer.py (python)
                orbax_checkpointer.save(ckpt_path, ckpt, save_args=save_args)
                
                # Save dataloader state
                grain_checkpoint_handler.save(
                    GrainPath(os.path.join(ckpt_path, "dataloader")), 
                    item=grain_dataloader
[content] train_tokenizer.py (python)
    
    # Create Grain checkpoint handler
    grain_checkpoint_handler = CheckpointHandler()
    
    # Restore dataloader state if resuming from checkpoint
    if args.checkpoint and os.path.exists(os.path.join(args.checkpoint, "dataloader")):
        grain_dataloader = grain_checkpoint_handler.restore(
            GrainPath(os.path.join(args.checkpoint, "dataloader")), 
            item=grain_dataloader
        )
    
    # Convert to JAX arrays
    dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in grain_dataloader) # type: ignore
[content] train_tokenizer.py (python)
                orbax_checkpointer.save(ckpt_path, ckpt, save_args=save_args)
                
                # Save dataloader state
                grain_checkpoint_handler.save(
                    epath.Path(os.path.join(ckpt_path, "dataloader")), 
                    item=grain_dataloader
[content] train_tokenizer.py (python)
    
    # Create Grain checkpoint handler
    grain_checkpoint_handler = CheckpointHandler()
    
    # Restore dataloader state if resuming from checkpoint
    if args.checkpoint and os.path.exists(os.path.join(args.checkpoint, "dataloader")):
        grain_dataloader = grain_checkpoint_handler.restore(
            epath.Path(os.path.join(args.checkpoint, "dataloader")), 
            item=grain_dataloader
        )
    
    # Convert to JAX arrays
    dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in grain_dataloader) # type: ignore
[content] train_tokenizer.py (python)
    if args.checkpoint and checkpoint_manager.latest_step() is not None:
        try:
            restored = checkpoint_manager.restore(checkpoint_manager.latest_step())
            if hasattr(restored, 'dataloader'):
                grain_dataloader = restored.dataloader
        except Exception as e:
            print(f"Could not restore dataloader state: {e}")
[content] train_tokenizer.py (python)
    # --- Initialize checkpoint manager ---
    checkpoint_dir = epath.Path(args.ckpt_dir) if args.ckpt_dir else epath.Path(f"checkpoints/tokenizer_{ts}")
    checkpoint_dir.mkdir(parents=True, exist_ok=True)
    
    # Create handler registry for Grain dataloader
    handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
    handler_registry.add('dataloader', ocp.args.StandardSave)
    handler_registry.add('dataloader', ocp.args.StandardRestore)
    
    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(
        checkpoint_dir,
        options=options,
        handler_registry=handler_registry,
    )
[content] train_tokenizer.py (python)
                # Save both model state and dataloader state using CompositeCheckpointHandler
                checkpoint_manager.save(
                    step,
                    args=ocp.args.Composite(
                        model=ocp.args.StandardSave(train_state),
                        dataloader=ocp.args.StandardSave(grain_dataloader),
                    )
                )
                checkpoint_manager.wait_until_finished()
[content] train_tokenizer.py (python)
                ckpt_path = os.path.join(os.getcwd(), args.ckpt_dir, f"tokenizer_{ts}_{step}")
                ckpt = {"model": train_state}
                orbax_checkpointer = PyTreeCheckpointer()
                save_args = orbax_utils.save_args_from_target(ckpt)
                orbax_checkpointer.save(ckpt_path, ckpt, save_args=save_args)
                
                # Save dataloader state
                grain_checkpoint_handler.save(
                    epath.Path(os.path.join(ckpt_path, "dataloader")), 
                    item=grain_dataloader
                )
[content] train_tokenizer.py (python)
    if args.checkpoint and os.path.exists(os.path.join(args.checkpoint, "dataloader")):
        grain_dataloader = grain_checkpoint_handler.restore(
            epath.Path(os.path.join(args.checkpoint, "dataloader")), 
            item=grain_dataloader
        )
[content] train_tokenizer.py (python)
    
    # Restore dataloader state if resuming from checkpoint
    if args.checkpoint and checkpoint_manager.latest_step() is not None:
        try:
            restored = checkpoint_manager.restore(checkpoint_manager.latest_step())
            if hasattr(restored, 'dataloader'):
                grain_dataloader = restored.dataloader
        except Exception as e:
            print(f"Could not restore dataloader state: {e}")
    
    # Convert to JAX arrays
    dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in grain_dataloader) # type: ignore
[content] train_tokenizer.py (python)
                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"tokenizer_{ts}_{step}"),
                    ckpt,
                    save_args=save_args,
                )
[content] train_tokenizer.py (python)
    
    # Restore dataloader state if resuming from checkpoint
    if args.checkpoint and checkpoint_manager.latest_step() is not None:
        try:
            restored = checkpoint_manager.restore(checkpoint_manager.latest_step())
            if 'dataloader' in restored:
                grain_dataloader = restored['dataloader']
        except Exception as e:
            print(f"Could not restore dataloader state: {e}")
    
    # Convert to JAX arrays
    dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in grain_dataloader) # type: ignore
[content] train_tokenizer.py (python)
                # Save both model state and dataloader state using CompositeCheckpointHandler
                checkpoint_manager.save(
                    step,
                    args=ocp.args.Composite(
                        state=ocp.args.StandardSave(train_state),
                        dataloader=ocp.args.StandardSave(grain_dataloader),
                    )
                )
                checkpoint_manager.wait_until_finished()
[content] train_tokenizer.py (python)
                    args=ocp.args.StandardSave(train_state)
                )
                model_checkpoint_manager.wait_until_finished()
                
                # Save dataloader state using Grain checkpoint handler
                grain_checkpoint_handler.save(
                    dataloader_checkpoint_dir, 
                    item=grain_dataloader
[content] train_tokenizer.py (python)
                    args=ocp.args.Composite(
                        state=ocp.args.StandardSave(train_state),
                        dataloader=ocp.args.StandardSave(grain_dataloader),
                    )
[content] train_tokenizer.py (python)
    # Create handler registry for Grain dataloader
    handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
    handler_registry.add('dataloader', ocp.args.StandardSave)
    handler_registry.add('dataloader', ocp.args.StandardRestore)
[content] train_tokenizer.py (python)
                model_checkpoint_manager.wait_until_finished()
                
                # Save dataloader state using Grain checkpoint handler
                grain_checkpoint_handler.save(
                    dataloader_checkpoint_dir, 
                    item=grain_dataloader
                )
[content] train_tokenizer.py (python)
    
    # Restore dataloader state if resuming from checkpoint
    if args.checkpoint and os.path.exists(dataloader_checkpoint_dir):
        try:
            grain_dataloader = grain_checkpoint_handler.restore(
                dataloader_checkpoint_dir, 
                item=grain_dataloader
            )
        except Exception as e:
            print(f"Could not restore dataloader state: {e}")
    
    # Convert to JAX arrays
    dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in grain_dataloader) # type: ignore
[content] train_tokenizer.py (python)
    # --- Initialize checkpoint manager ---
    checkpoint_dir = epath.Path(args.ckpt_dir) if args.ckpt_dir else epath.Path(f"checkpoints/tokenizer_{ts}")
    checkpoint_dir.mkdir(parents=True, exist_ok=True)
    
    # Create separate checkpoint managers for model and dataloader
    model_checkpoint_dir = checkpoint_dir / "model"
    dataloader_checkpoint_dir = checkpoint_dir / "dataloader"
    
    model_options = ocp.CheckpointManagerOptions(
        save_interval_steps=args.log_checkpoint_interval,
        max_to_keep=3,
        step_format_fixed_length=6,
        cleanup_tmp_directories=True,
    )
    
    model_checkpoint_manager = ocp.CheckpointManager(
        model_checkpoint_dir,
        options=model_options,
    )
    
    # Create Grain checkpoint handler for dataloader
    grain_checkpoint_handler = CheckpointHandler()
[content] train_tokenizer.py (python)
                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"tokenizer_{ts}_{step}"),
                    ckpt,
                    save_args=save_args,
[content] train_tokenizer.py (python)
                    args=ocp.args.Composite(
                        model_state=ocp.args.StandardSave(train_state),
                        dataloader_state=ocp.args.StandardSave(grain_dataloader),
                    )
[content] train_tokenizer.py (python)
            # Try to restore dataloader state from the same checkpoint
            restored = checkpoint_manager.restore(
                step,
                args=ocp.args.Composite(
                    model_state=None,  # Skip model state
                    dataloader_state=None,  # Will use default handler
                )
            )
            if hasattr(restored, 'dataloader_state') and restored.dataloader_state is not None:
                grain_dataloader = restored.dataloader_state
                print(f"Restored dataloader state from step {step}")
[content] train_tokenizer.py (python)
        # Restore both model and dataloader state using CompositeCheckpointHandler
        try:
            # Create abstract state for restoration
            abstract_train_state = jax.tree.map(ocp.tree.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=None,  # Will be restored separately if exists
                )
            )
            
            # Update train state with restored parameters
            train_state = restored.model_state
            
            # Get step from checkpoint
            step = checkpoint_manager.latest_step()
            print(f"Restored model state from step {step}")
            
        except Exception as e:
            print(f"Could not restore from checkpoint: {e}")
            print("Starting from scratch...")
[content] train_tokenizer.py (python)
    # Create handler registry for proper restoration
    handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
    handler_registry.add('model_state', ocp.args.StandardSave)
    handler_registry.add('model_state', ocp.args.StandardRestore)
    handler_registry.add('dataloader_state', ocp.args.StandardSave)
    handler_registry.add('dataloader_state', ocp.args.StandardRestore)
[content] train_tokenizer.py (python)
                    train_state
                )
                model_checkpoint_manager.wait_until_finished()
                
                # Save dataloader state using Grain checkpoint handler
                grain_checkpoint_handler.save(
                    dataloader_checkpoint_dir, 
                    item=grain_dataloader
[content] train_tokenizer.py (python)
        # Restore model state
        restore_target = {"model": train_state}
        restore_args = orbax_utils.restore_args_from_target(restore_target)
        train_state.params["params"].update(
            ocp.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_tokenizer.py (python)
                # Save both model and dataloader state using CompositeCheckpointHandler
                checkpoint_manager.save(
                    step,
                    args=ocp.args.Composite(
                        model_state=ocp.args.StandardSave(train_state),
                        dataloader_state=ocp.args.StandardSave(grain_dataloader),
                    )
                )
                checkpoint_manager.wait_until_finished()
                print(f"Saved checkpoint at step {step}")
[content] train_tokenizer.py (python)
    
    # Restore dataloader state if resuming from checkpoint
    if args.checkpoint and step > 0:
        try:
            # Try to restore dataloader state from the same checkpoint
            restored = checkpoint_manager.restore(
                step,
                args=ocp.args.Composite(
                    model_state=None,  # Skip model state
                    dataloader_state=None,  # Will use default handler
                )
            )
            if hasattr(restored, 'dataloader_state') and restored.dataloader_state is not None:
                grain_dataloader = restored.dataloader_state
                print(f"Restored dataloader state from step {step}")
        except Exception as e:
            print(f"Could not restore dataloader state: {e}")
            print("Starting with fresh dataloader...")
    
    # Convert to JAX arrays
    dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in grain_dataloader) # type: ignore
[content] train_tokenizer.py (python)
    # --- Initialize checkpoint manager with CompositeCheckpointHandler ---
    checkpoint_dir = epath.Path(args.ckpt_dir) if args.ckpt_dir else epath.Path(f"checkpoints/tokenizer_{ts}")
    checkpoint_dir.mkdir(parents=True, exist_ok=True)
    
    # Create handler registry for proper restoration
    handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
    handler_registry.add('model_state', ocp.args.StandardSave)
    handler_registry.add('model_state', ocp.args.StandardRestore)
    handler_registry.add('dataloader_state', ocp.args.StandardSave)
    handler_registry.add('dataloader_state', ocp.args.StandardRestore)
    
    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(
        checkpoint_dir,
        options=checkpoint_options,
        handler_registry=handler_registry,
    )
[content] train_tokenizer.py (python)
        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_tokenizer.py (python)

    if args.restore_ckpt:
        # Restore both model and dataloader state using CompositeCheckpointHandler
        # Create abstract state for restoration
        abstract_train_state = jax.tree.map(ocp.tree.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=None,  # Will be restored separately if exists
            )
        )
        
        # Update train state with restored parameters
        train_state = restored.model_state
        
        # Get step from checkpoint
        step = checkpoint_manager.latest_step()
        print(f"Restored model state from step {step}")
[content] train_tokenizer.py (python)

    # --- Initialize checkpoint manager with CompositeCheckpointHandler ---
    if args.ckpt_dir:
        checkpoint_dir = os.path.join(args.ckpt_dir, f"tokenizer_{ts}")
        
        # Create handler registry for proper restoration
        handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
        handler_registry.add('model_state', ocp.args.StandardSave)
        handler_registry.add('model_state', ocp.args.StandardRestore)
        handler_registry.add('dataloader_state', ocp.args.StandardSave)
        handler_registry.add('dataloader_state', ocp.args.StandardRestore)
        
        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(
            checkpoint_dir,
            options=checkpoint_options,
            handler_registry=handler_registry,
        )

        if args.restore_ckpt:
            abstract_train_state = jax.tree.map(ocp.tree.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=None,  # Will be restored separately if exists
                )
            )
            
            # Update train state with restored parameters
            train_state = restored.model_state
            
            # Get step from checkpoint
            step = checkpoint_manager.latest_step()
            print(f"Restored model state from step {step}")
[content] train_tokenizer.py (python)

    
    # Create Grain dataloader
    grain_dataloader = get_dataloader(
        # NOTE: We deliberately pass the global batch size
        # The dataloader shards the dataset across all processes
        array_record_files,
        args.seq_len,
        args.batch_size,
        *image_shape,
        num_workers=8,
        prefetch_buffer_size=1,
        seed=args.seed,
    )
[content] train_tokenizer.py (python)
        try:
            abstract_train_state = jax.tree.map(ocp.tree.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=None,  # Will be restored if exists
                )
            )
            
            train_state = restored.model_state
            step = checkpoint_manager.latest_step()
            print(f"Restored model state from step {step}")
            
            # Try to restore dataloader state separately
            try:
                restored_dataloader = checkpoint_manager.restore(
                    step,
                    args=ocp.args.Composite(
                        model_state=None,  # Skip model state
                        dataloader_state=None,  # Will use default handler
                    )
                )
                if hasattr(restored_dataloader, 'dataloader_state') and restored_dataloader.dataloader_state is not None:
                    dataloader = restored_dataloader.dataloader_state
                    print(f"Restored dataloader state from step {step}")
                else:
                    print("No dataloader state found in checkpoint, using fresh dataloader")
            except Exception as e:
                print(f"Could not restore dataloader state: {e}")
                print("Using fresh dataloader...")
                
        except Exception as e:
            print(f"Could not restore from checkpoint: {e}")
            print("Starting from scratch...")
            step = 0
[content] train_tokenizer.py (python)
        abstract_train_state = jax.tree.map(ocp.tree.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),
                # FIXME (f.srambical): also restore dataloader state
                dataloader_state=None,  # Will be restored separately if exists
            )
        )
        
        train_state = restored.model_state
        dataloader = restored.dataloader_state
        
        step = checkpoint_manager.latest_step()
        print(f"Restored model state from step {step}")
[content] train_tokenizer.py (python)
        try:
            abstract_train_state = jax.tree.map(ocp.tree.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=None,  # Will be restored if exists
                )
            )
            
            train_state = restored["model_state"]
            step = checkpoint_manager.latest_step()
            print(f"Restored model state from step {step}")
            
            # Try to restore dataloader state separately
            try:
                restored_dataloader = checkpoint_manager.restore(
                    step,
                    args=ocp.args.Composite(
                        model_state=None,  # Skip model state
                        dataloader_state=None,  # Will use default handler
                    )
                )
                if "dataloader_state" in restored_dataloader and restored_dataloader["dataloader_state"] is not None:
                    dataloader = restored_dataloader["dataloader_state"]
                    print(f"Restored dataloader state from step {step}")
                else:
                    print("No dataloader state found in checkpoint, using fresh dataloader")
            except Exception as e:
                print(f"Could not restore dataloader state: {e}")
                print("Using fresh dataloader...")
                
        except Exception as e:
            print(f"Could not restore from checkpoint: {e}")
            print("Starting from scratch...")
            step = 0
[tab] convert_knoms_to_arrayrecords.py (python)
from utils.dataloader import convert_tfrecords_to_arrayrecords

convert_tfrecords_to_arrayrecords(
    tfrecord_folder="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_tfrecords_500_shards",
    output_folder="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_arrayrecords_500_shards_group_size_1",
)
[content] train_tokenizer.py (python)
        try:
            abstract_train_state = jax.tree.map(ocp.tree.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=None,  # Will be restored if exists
                )
            )
            
            train_state = restored["model_state"]
            step = checkpoint_manager.latest_step()
            print(f"Restored model state from step {step}")
            
            # Check if dataloader state was also restored
            if "dataloader_state" in restored and restored["dataloader_state"] is not None:
                dataloader = restored["dataloader_state"]
                print(f"Restored dataloader state from step {step}")
            else:
                print("No dataloader state found in checkpoint, using fresh dataloader")
                
        except Exception as e:
            print(f"Could not restore from checkpoint: {e}")
            print("Starting from scratch...")
            step = 0
[content] train_tokenizer.py (python)
        try:
            train_state_restore_args = orbax_utils.restore_args_from_target(train_state)
            
            restored = checkpoint_manager.restore(
                checkpoint_manager.latest_step(),
                args=ocp.args.Composite(
                    model_state=train_state_restore_args,
                    dataloader_state=None,
                )
            )
            train_state = restored["model_state"]
            step = checkpoint_manager.latest_step()
            print(f"Restored model state from step {step}")
            
            # Check if dataloader state was also restored
            if "dataloader_state" in restored and restored["dataloader_state"] is not None:
                dataloader = restored["dataloader_state"]
                print(f"Restored dataloader state from step {step}")
            else:
                print("No dataloader state found in checkpoint, using fresh dataloader")
                
        except Exception as e:
            print(f"Could not restore from checkpoint: {e}")
            print("Starting from scratch...")
            step = 0
[content] train_tokenizer.py (python)
        train_state_restore_args = orbax_utils.restore_args_from_target(train_state)
        abstract_train_state = jax.tree.map(ocp.tree.to_shape_dtype_struct, train_state)
        
        restored = checkpoint_manager.restore(
            checkpoint_manager.latest_step(),
            args=ocp.args.Composite(
                model_state=train_state_restore_args,
                dataloader_state=None,
            )
        )
        train_state = restored.model_state
        dataloader = restored.dataloader_state
        step = checkpoint_manager.latest_step()
        print(f"Restored model state from step {step}")
[content] train_tokenizer.py (python)
        try:
            latest_step = checkpoint_manager.latest_step()
            if latest_step is not None:
                train_state_restore_args = orbax_utils.restore_args_from_target(train_state)
                
                restored = checkpoint_manager.restore(
                    latest_step,
                    args=ocp.args.Composite(
                        model_state=train_state_restore_args,
                        dataloader_state=None,
                    )
                )
                train_state = restored["model_state"]
                step = latest_step
                print(f"Restored model state from step {step}")
                
                # Check if dataloader state was also restored
                if "dataloader_state" in restored and restored["dataloader_state"] is not None:
                    dataloader = restored["dataloader_state"]
                    print(f"Restored dataloader state from step {step}")
                else:
                    print("No dataloader state found in checkpoint, using fresh dataloader")
            else:
                print("No checkpoint found, starting from scratch...")
                step = 0
                
        except Exception as e:
            print(f"Could not restore from checkpoint: {e}")
            print("Starting from scratch...")
            step = 0
[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 \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --ckpt_dir checkpoints/tokenizer_openai_time_measurement_grain \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[content] train_tokenizer.py (python)
    # Convert to JAX arrays - this creates a generator that cannot be checkpointed
    jax_dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in dataloader) # type: ignore
[content] train_dynamics.py (python)
                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_{ts}_{step}"),
                    ckpt,
                    save_args=save_args,
                )
[content] train_dynamics.py (python)
        tokenizer_dim=args.tokenizer_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,
        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,
        # Dynamics
        dyna_dim=args.dyna_dim,
        dyna_num_blocks=args.dyna_num_blocks,
        dyna_num_heads=args.dyna_num_heads,
[content] train_dynamics.py (python)
        cross_entropy_loss=ce_loss,
        masked_token_accuracy=acc,
        select_logit=outputs["token_logits"].max(-1).mean(),
        select_p=select_probs.max(-1).mean(),
        entropy=jax.scipy.special.entr(select_probs).sum(-1).mean(),
[content] train_dynamics.py (python)
    mask = outputs["mask"]
    ce_loss = optax.softmax_cross_entropy_with_integer_labels(
        outputs["token_logits"], outputs["video_tokens"]
    )
    ce_loss = (mask * ce_loss).sum() / mask.sum()
    acc = outputs["token_logits"].argmax(-1) == outputs["video_tokens"]
    acc = (mask * acc).sum() / mask.sum()
    select_probs = jax.nn.softmax(outputs["token_logits"])
[content] train_dynamics.py (python)
    batch_size: int = 36
    min_lr: float = 3e-6
    max_lr: float = 3e-5
    warmup_steps: int = 5000
    # Tokenizer
    tokenizer_dim: int = 512
    latent_patch_dim: int = 32
    num_patch_latents: int = 1024
    patch_size: int = 4
    tokenizer_num_blocks: int = 8
    tokenizer_num_heads: int = 8
    tokenizer_checkpoint: str = ""
    # LAM
    lam_dim: int = 512
    latent_action_dim: int = 32
    num_latent_actions: int = 6
    lam_patch_size: int = 16
    lam_num_blocks: int = 8
    lam_num_heads: int = 8
    lam_checkpoint: str = ""
    # Dynamics
    dyna_dim: int = 512
    dyna_num_blocks: int = 12
    dyna_num_heads: int = 8
[content] train_lam.py (python)
                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)
            inputs = dict(videos=videos, rng=_rng)
            train_state, loss, recon, action_last_active, metrics = train_step(
                train_state, inputs, action_last_active
            )
[content] train_lam.py (python)
        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)

    # --- Reset inactive latent actions ---
    codebook = state.params["params"]["vq"]["codebook"]
    num_codes = len(codebook)
    active_codes = idx_counts != 0.0
    action_last_active = jnp.where(active_codes, 0, action_last_active + 1)
    p_code = active_codes / active_codes.sum()
    reset_idxs = jax.random.choice(rng, num_codes, shape=(num_codes,), p=p_code)
    do_reset = action_last_active >= args.vq_reset_thresh
    new_codebook = jnp.where(
        jnp.expand_dims(do_reset, -1), codebook[reset_idxs], codebook
    )
    state.params["params"]["vq"]["codebook"] = new_codebook
    action_last_active = jnp.where(do_reset, 0, action_last_active)
    return state, loss, recon, action_last_active, metrics
[content] train_lam.py (python)
    gt_future_frames = inputs["videos"][:, 1:]
    mse = jnp.square(gt_future_frames - outputs["recon"]).mean()
    q_loss = jnp.square(jax.lax.stop_gradient(outputs["emb"]) - outputs["z"]).mean()
    commitment_loss = jnp.square(
        outputs["emb"] - jax.lax.stop_gradient(outputs["z"])
    ).mean()
    loss = mse + q_loss + args.vq_beta * commitment_loss
[content] train_tokenizer.py (python)
    # Create iterator from the dataloader
    grain_dataloader_iter = iter(grain_dataloader)
    dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in grain_dataloader_iter) # type: ignore
[content] train_tokenizer.py (python)
        except StopIteration:
            # Iterator exhausted, create a new one from the dataloader
            grain_dataloader_iter = iter(grain_dataloader)
            dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in grain_dataloader_iter) # type: ignore
            continue
[content] train_tokenizer.py (python)
    # Create DataLoaderIterator for checkpointing
    dataloader_iterator = grain.DataLoaderIterator(grain_dataloader, None)
    dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in dataloader_iterator) # type: ignore
[content] train_tokenizer.py (python)
            args=ocp.args.Composite(
                model_state=train_state_restore_args,
                dataloader_state=None,
            )
        )
        train_state = restored.model_state
        grain_dataloader = restored.dataloader_state
[content] train_tokenizer.py (python)
    handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
    handler_registry.add('model_state', ocp.args.StandardSave)
    handler_registry.add('model_state', ocp.args.StandardRestore)
    handler_registry.add('dataloader_state', ocp.args.StandardSave)
    handler_registry.add('dataloader_state', ocp.args.StandardRestore)