Feature Request: add LazyPromptDataset to DeepSpeedChat

[0xDing]

Currently, in our multi-GPU training process, every GPU process loads the train and evaluation datasets using the torch.load function. However, when the dataset is large, this can lead to an issue with CPU memory constraints.

To address this issue, I would like to suggest introducing the LazyPromptDataset functionality. This functionality will enable us to load the data on demand instead of pre-loading it into memory, which will save memory and improve overall performance.

[conglongli]

@0xDing what do you mean by LazyPromptDataset? If it's an existing thing, please provide a pointer. If it's a new feature that doesn't exist anywhere else, it'd be helpful if you could clarify what exact feature you want inside LazyPromptDataset.

[conglongli]

@0xDing and in case you haven't git pull, this recently merged PR can help avoid "every GPU process loads the train and evaluation datasets" #409

[0xDing]

@0xDing what do you mean by LazyPromptDataset? If it's an existing thing, please provide a pointer. If it's a new feature that doesn't exist anywhere else, it'd be helpful if you could clarify what exact feature you want inside LazyPromptDataset.

@conglongli

I apologize for any confusion caused by my previous description. I am trying to load a large dataset (the original file is 1.5GB with approximately 2 million data entries, and after tokenization and other preprocessing, the resulting traindata_xxx.pt file is 160GB). In the current implementation of DeepSpeedChat, if multiple GPUs are used, each process will call torch.load on this file. This causes a memory error (exit with return code = -9) even if there is 1TB of CPU memory available.

Therefore, I believe that implementing a LazyPromptDataset class (which performs tokenization and other preprocessing only when getitem is called, similar to Vicuna) could be helpful for loading large datasets.

Do you have any other suggestions for using a large dataset in DeepSpeedChat? Finally, thank you for creating this awesome project.

[conglongli]

@0xDing no worries and thank you for the clarifications. Yes we agree that our data processing/loading can be improved to reduce memory consumption. For the LazySupervisedDataset you cited from Vicuna, my understanding is that it still won't solve all memory consumption problem because the self.cached_data_dict can still explode. So currently our plan is to switch to memmap-based data processing/loading/management, something similar to the way we do pre-training in https://github.com/microsoft/Megatron-DeepSpeed/blob/main/megatron/data/indexed_dataset.py. This should prevent any CPU memory issue related to data. Please understand that it would take some time to do this transition due to the difference in terms of code and data structure between the two repos. But hopefully we can release this feature sooner than later.