Memory mapped dataloader

Another restart overhead stems from data loading: the training cluster remains idle while the dataloader initializes, downloads data from remote file systems, and processes it into batches.

To address this, we introduce the Memory Mapped DataLoader(MMAP) Dataloader, which caches prefetched batches in persistent memory, ensuring they remain available even after a fault-induced restart. This approach eliminates dataloader setup time and enables training to resume immediately using cached batches, while the dataloader concurrently reinitializes and fetches subsequent data in the background. The data cache resides on each rank that requires training data and maintains two types of batches: recently consumed batches that have been used for training, and prefetched batches ready for immediate use.

MMAP dataloader offers two following features:

Using the cache, the training job will benefit from:

MMAP configurations

To use MMAP, simply pass in the your original data module into MMAPDataModule

data_module=MMAPDataModule(
    data_module=MY_DATA_MODULE(...),
    mmap_config=CacheResumeMMAPConfig(
        cache_dir=self.cfg.mmap.cache_dir,
        checkpoint_frequency=self.cfg.mmap.checkpoint_frequency),
)

CacheResumeMMAPConfig: MMAP Dataloader parameters control cache directory location, size limits, and data fetching delegation. By default, only TP rank 0 per node fetches data from the source, while other ranks in the same data replication group read from the shared cache, eliminating redundant transfers.

MMAPDataModule: It wraps the original data module and returns the mmap dataloader for both train and validation.

See the example for enabling MMAP.

API reference

CacheResumeMMAPConfig

class hyperpod_checkpointless_training.dataloader.config.CacheResumeMMAPConfig(
  cache_dir='/dev/shm/pdl_cache',
  prefetch_length=10,
  val_prefetch_length=10,
  lookback_length=2,
  checkpoint_frequency=None,
  model_parallel_group=None,
  enable_batch_encryption=False)

Configuration class for cache-resume memory-mapped (MMAP) dataloader functionality in HyperPod checkpointless training.

This configuration enables efficient data loading with caching and prefetching capabilities, allowing training to resume quickly after failures by maintaining cached data batches in memory-mapped files.

Parameters

Methods

create(dataloader_init_callable,
    parallel_state_util,
   step,
    is_data_loading_rank,
   create_model_parallel_group_callable,
    name='Train',
   is_val=False,
   cached_len=0)

Creates and returns a configured MMAP dataloader instance.

Parameters

Returns CacheResumePrefetchedDataLoader or CacheResumeReadDataLoader – Configured MMAP dataloader instance

Raises ValueError if the step parameter is None.

Example

from hyperpod_checkpointless_training.dataloader.config import CacheResumeMMAPConfig

# Create configuration
config = CacheResumeMMAPConfig(
    cache_dir="/tmp/training_cache",
    prefetch_length=20,
    checkpoint_frequency=100,
    enable_batch_encryption=False
)

# Create dataloader
dataloader = config.create(
    dataloader_init_callable=my_dataloader_init,
    parallel_state_util=parallel_util,
    step=current_step,
    is_data_loading_rank=lambda: rank == 0,
    create_model_parallel_group_callable=create_mp_group,
    name="TrainingData"
)

Notes

MMAPDataModule

class hyperpod_checkpointless_training.dataloader.mmap_data_module.MMAPDataModule(  
    data_module,  
    mmap_config,  
    parallel_state_util=MegatronParallelStateUtil(),  
    is_data_loading_rank=None)

A PyTorch Lightning DataModule wrapper that applies memory-mapped (MMAP) data loading capabilities to existing DataModules for checkpointless training.

This class wraps an existing PyTorch Lightning DataModule and enhances it with MMAP functionality, enabling efficient data caching and fast recovery during training failures. It maintains compatibility with the original DataModule interface while adding checkpointless training capabilities.

Parameters

Attributes

Methods

setup(stage=None)

Setup the underlying data module for the specified training stage.

train_dataloader()

Create the training DataLoader with MMAP wrapping.

Returns: DataLoader – MMAP-wrapped training DataLoader with caching and prefetching capabilities

val_dataloader()

Create the validation DataLoader with MMAP wrapping.

Returns: DataLoader – MMAP-wrapped validation DataLoader with caching capabilities

test_dataloader()

Create the test DataLoader if the underlying data module supports it.

Returns: DataLoader or None – Test DataLoader from the underlying data module, or None if not supported

predict_dataloader()

Create the predict DataLoader if the underlying data module supports it.

Returns: DataLoader or None – Predict DataLoader from the underlying data module, or None if not supported

load_checkpoint(checkpoint)

Load checkpoint information to resume training from a specific step.

get_underlying_data_module()

Get the underlying wrapped data module.

Returns: pl.LightningDataModule – The original data module that was wrapped

state_dict()

Get the state dictionary of the MMAP DataModule for checkpointing.

Returns: dict – Dictionary containing cached dataloader lengths

load_state_dict(state_dict)

Load the state dictionary to restore MMAP DataModule state.

Properties

data_sampler

Expose the underlying data module's data sampler to NeMo framework.

Returns: object or None – The data sampler from the underlying data module

Example

from hyperpod_checkpointless_training.dataloader.mmap_data_module import MMAPDataModule  
from hyperpod_checkpointless_training.dataloader.config import CacheResumeMMAPConfig  
from my_project import MyLLMDataModule  

# Create MMAP configuration  
mmap_config = CacheResumeMMAPConfig(  
    cache_dir="/tmp/training_cache",  
    prefetch_length=20,  
    checkpoint_frequency=100  
)  

# Create original data module  
original_data_module = MyLLMDataModule(  
    data_path="/path/to/data",  
    batch_size=32  
)  

# Wrap with MMAP capabilities  
mmap_data_module = MMAPDataModule(  
    data_module=original_data_module,  
    mmap_config=mmap_config  
)  

# Use in PyTorch Lightning Trainer  
trainer = pl.Trainer()  
trainer.fit(model, data=mmap_data_module)  

# Resume from checkpoint  
checkpoint = {"global_step": 1000}  
mmap_data_module.load_checkpoint(checkpoint)

Notes