o }oirP@sddlmZddlmZmZmZmZddlZddl m Z ddl m Z mZddlmZddlmZmZmZddlmZddlmZdd lmZdd lmZdd lmZmZm Z dd l!m"Z"dd l#m$Z$Gddde j%eZ&Gddde"Z'dS)deepcopy)AnyDictLiteralOptionalN)EVAL_DATALOADERSTRAIN_DATALOADERS)parallel_state) WorkerConfigget_savable_loaderget_train_dataset) DataLoader)Self)MultiModalSampleConfig)MultiModalTaskEncoder)IOMixin serializationtrack_io)MegatronDataSampler)loggingcseZdZdZddddddddeddddf deded ed ed ed edBd edededBdeedee deedeedee ddffdd Z de j e fddZd*dedfddZdefdd Zdefd!d"Zd+d#d$Zdeeeffd%d&Zd'eeefddfd(d)ZZS),EnergonMultiModalDataModulea A PyTorch Lightning DataModule for handling multimodal datasets with images and text. This data module is designed to work with multimodal datasets that involve both images and text. It provides a seamless interface to load training and validation data, manage batching, and handle the state of the data pipeline across training epochs. The module integrates with the Megatron-Energon framework for efficient data handling in large-scale distributed training. Attributes: path (str): Path to the energon dataset. tokenizer (Tokenizer): The tokenizer used for processing text. image_processor (ImageProcessor): The image processor used for preprocessing images. seq_length (int): The maximum sequence length for tokenized text. micro_batch_size (int): The batch size for training and validation. num_workers (int): Number of workers for data loading. pin_memory (bool): Whether to pin memory in the DataLoader. multimodal_sample_config (MultiModalSampleConfig): Configuration object for multimodal samples. task_encoder (MultiModalTaskEncoder): Encoder responsible for encoding and batching samples. init_global_step (int): The initial global step for the trainer, used for resuming training. data_sampler (SequentialMegatronSampler): Sampler responsible for generating sequential samples. train_dataloader_object (Optional): The DataLoader object for training data. val_dataloader_object (Optional): The DataLoader object for validation data. iNTdpath seq_lengthmicro_batch_sizeglobal_batch_size num_workersnum_val_workers pin_memoryshuffle_buffer_sizemax_samples_per_sequencemultimodal_sample_config task_encoderdecoder_seq_lengthpacking_buffer_sizevalidation_task_encoderreturnc st||_||_||_||_||_||_||_||_ | |_ | |_ | |_ | |_ | p3t|j|j| d|_d|_t|j|j|j|jd|_d|_d|_||_|pR|j|_|pX|j |_||_dS)a Initialize the EnergonMultiModalDataModule. Parameters: path (str): Path to the dataset. tokenizer (Tokenizer): The tokenizer used for processing text. image_processor (ImageProcessor): The image processor used for preprocessing images. seq_length (int, optional): The maximum sequence length for tokenized text. Defaults to 2048. micro_batch_size (int, optional): The batch size for training and validation. Defaults to 1. num_workers (int, optional): Number of workers for data loading. Defaults to 1. num_val_workers (int, optional): Number of workers for validation data loading. Defaults to num_workers. pin_memory (bool, optional): Whether to pin memory in the DataLoader. Defaults to True. multimodal_sample_config (MultiModalSampleConfig, optional): Configuration object for multimodal samples. Defaults to MultiModalSampleConfig(). shuffle_buffer_size (int, optional): Size of the shuffle buffer. Defaults to 100. max_samples_per_sequence (int, optional): Maximum number of samples per sequence to load from memory. Defaults to None (loads the whole tar file at once). task_encoder (MultiModalTaskEncoder, optional): Encoder responsible for encoding and batching samples. If not provided, a default (MultimodalTaskEncoder) encoder will be created. Defaults to None. decoder_seq_length (int, optional): The max sequence length for the decoder. Used in encoder-decoder models packing_buffer_size (int, optional): Size of the packing buffer for batched samples. Defaults to None. validation_task_encoder (MultiModalTaskEncoder, optional): Encoder responsible for encoding and batching samples for validation. Defaults to None and will be the same as task_encoder. **kwargs: Additional keyword arguments. Will be passed to get_train_dataset() of Energon ) tokenizerimage_processorr#r)seq_lendecoder_seq_lenrrN)super__init__rr)r*rr%rrrr r#r!r"rr$init_global_stepSequentialMegatronSampler data_samplertrain_dataloader_objectval_dataloader_objectr&r'rkwargs)selfrr)r*rrrrrr r!r"r#r$r%r&r'r4 __class__a/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/multimodal/data/energon/base.pyr.:s> .   z$EnergonMultiModalDataModule.__init__cKsTdd|D}|D]}tt|stt|q tjt|fi|}|S)NcSs"i|] \}}|dvr|t|qS))r*r$r'r).0kvr8r8r9 s z7EnergonMultiModalDataModule.io_init..)itemsvaluesrfind_node_traversertyperfdlConfig)r5r4 cfg_kwargsvalcfgr8r8r9io_inits  z#EnergonMultiModalDataModule.io_initrEsplit)trainrEc CsV|dvrtd|dkr|j}n|j}t|jf|j|||j||j|jd|j }|S)a Provide the dataset for training or validation. This method retrieves the dataset for the specified split (either 'train' or 'val') and configures it according to the worker configuration. Parameters: worker_config: Configuration for the data loader workers. split (Literal['train', 'val'], optional): The data split to retrieve ('train' or 'val'). Defaults to 'val'. Returns: Dataset: The dataset configured for the specified split. >rErIz=Invalid value for split. Allowed values are 'train' or 'val'.rI) batch_sizer$ worker_configr& split_partr!r") ValueErrorr$r'r rrr&r!r"r4)r5rKrHr$_datasetr8r8r9datasets_providers&  z-EnergonMultiModalDataModule.datasets_providerc Cs|jr|jj|_|j|j_td|j|jr|jSts0td|j t |j }n&t }t }t}td|d|d|dt |||j |ddd }|j|d d }t||d }||_|jS) a Initialize and return the training DataLoader. This method initializes the DataLoader for the training dataset. It uses the global step from the trainer to configure the data sampler and ensures that the parallel state is initialized correctly for distributed training. Returns: TRAIN_DATALOADERS: The DataLoader for the training dataset. z?Multimodal train dataloader initializing with init_global_step zfMuiltimodal data loader parallel state is not initialized,using default worker config with no_workers z4 Multimodal train dataloader initializing withrank world_size  data_parallel_group z ****** Nrrank world_sizerdata_parallel_groupworker_debug_pathworker_log_levelrIrHrK)trainer global_stepr/r1rinfor2r is_initializedrr default_worker_configget_data_parallel_rankget_data_parallel_world_sizeget_data_parallel_grouprOr )r5rKrSrTrU train_datasetenergon_dataloaderr8r8r9train_dataloadersJ   z,EnergonMultiModalDataModule.train_dataloadercCs|jr|jStstd|jt|j}n%t }t }t }td|d|d|t|||j|ddd}|j |dd }t ||d }||_|jS) an Initialize and return the validation DataLoader. This method initializes the DataLoader for the validation dataset. It ensures that the parallel state is initialized correctly for distributed training and returns a configured DataLoader object. Returns: EVAL_DATALOADERS: The DataLoader for the validation dataset. zjMuiltimodal val data loader parallel state is not initialized,using default worker config with no_workers zrank rPrQNrrRrErXrY)r3r r]rr\rr r^rr_r`rarOr )r5rKrSrTrU val_datasetenergon_loaderr8r8r9val_dataloaders2  z*EnergonMultiModalDataModule.val_dataloadercCstddS)z Return None as test dataset split does not exist. This method overrides the test_dataloader method and returns None since the test dataset split is not defined or used in this module. Returns: None z7Multimodal dataloader test dataset split does not existN)rwarningr5r8r8r9test_dataloader s z+EnergonMultiModalDataModule.test_dataloadercCs|jr?|jj}g}tptptptdkr!|jdd}|j |jj |j }|dur2g}t d|||dSt diS)aT Save the state of the data module. This method is called when saving a checkpoint. It generates and saves the state of the data module, including the state of the dataloader and the number of consumed samples. Returns: Dict[str, Any]: A dictionary containing the state of the data module. r)global_dst_rankNzEMultimodal data loader saving dataloader state dict consumed samples )dataloader_stateconsumed_sampleszHtrainer object not connected to data module object returning empty state)rZrdr get_context_parallel_rank get_pipeline_model_parallel_rankget_tensor_model_parallel_rankget_expert_model_parallel_ranksave_state_globalr1compute_consumed_samplesr[r/rr\rh)r5dataloader_objstatermr8r8r9 state_dicts*     z&EnergonMultiModalDataModule.state_dictrvc Csd|vrtd|dS|d}z |jr'|jj|tdn td|t dWnt yN}ztd|WYd}~nd}~wwzdd l m }Wnt tfymtd dd lm }Ynw|d }||j_||j_td |||d ddS)az Load the state of the data module from a checkpoint. This method is called when loading a checkpoint. It restores the state of the data module, including the state of the dataloader and the number of consumed samples. Parameters: state_dict (Dict[str, Any]): The state dictionary containing the saved state of the data module. rlzpData loader state cannot be resumed from state_dict, it does not have the required key dataloader_state. It has Nz$Multimodal dataloader state restoredz%Cannot restore state from state_dict zhCannot restore state from state_dict: Is the trainer object is initialized and attached to datamodule???zFailed to dataloader restore state due to [Please ensure you are using same version of energon while saving and loading, Continuing without restoring data loader] : r)update_num_microbatcheszCMegatron num_microbatches_calculator not found, using Apex version.rmzs$         V