Package: litdata
Blazingly fast, distributed streaming of training data from cloud storage⚡ Welcome to Lightning DataWe developedStreamingDatasetto optimize training of large datasets stored on the cloud while prioritizing speed, affordability, and scalability.Specifically crafted for multi-gpu & multi-node (withDDP,FSDP, etc...), distributed training with large models, it enhances accuracy, performance, and user-friendliness. Now, training efficiently is possible regardless of the data's location. Simply stream in the required data when needed.TheStreamingDatasetis compatible with any data type, includingimages, text, video, audio, geo-spatial, and multimodal dataand it is a drop-in replacement for your PyTorchIterableDatasetclass. For example, it is used byLit-GPTto pretrain LLMs.BenchmarksImagenet-1.2Mis a commonly used dataset to compare computer vision models. Its training dataset contains1,281,167 images.In this benchmark, we measured the streaming speed (images per second) loaded fromAWS S3for several frameworks.Find the reproducibleStudio Benchmark.Imagenet-1.2M Streaming from AWS S3FrameworkImages / sec  1st Epoch (float32)Images / sec   2nd Epoch (float32)Images / sec 1st Epoch (torch16)Images / sec 2nd Epoch (torch16)PL Data5800.346589.986282.177221.88Web Dataset3134.423924.953343.404424.62Mosaic ML2898.615099.932809.695158.98Higher is better.Imagenet-1.2M ConversionFrameworkTrain Conversion TimeVal Conversion TimeDataset Size# FilesPL Data10:05 min00:30 min143.1 GB2.339Web Dataset32:36 min01:22 min147.8 GB1.144Mosaic ML49:49 min01:04 min143.1 GB2.298The dataset needs to be converted into an optimized format for cloud streaming. We measured how fast the 1.2 million images are converted.Faster is better.Real World ExamplesWe have built end-to-end freeStudiosshowing all the steps to prepare the following datasets:DatasetData typeStudioLAION-400MImage & descriptionUse or explore LAION-400MILLION datasetChesapeake Roads Spatial ContextImage & MaskConvert GeoSpatial data to Lightning StreamingImagenet 1MImage & LabelBenchmark cloud data-loading librariesSlimPajama&StartCoderTextPrepare the TinyLlama 1T token datasetEnglish WikepediaTextEmbed English Wikipedia under 5 dollarsGeneratedParquet FilesConvert parquets to Lightning StreamingLightning Studiosare fully reproducible cloud IDE with data, code, dependencies, etc...Getting StartedInstallationLightning Data can be installed withpip:pipinstall-UlitdataQuick Start1. Prepare Your DataConvert your raw dataset into Lightning Streaming format using theoptimizeoperator. More formats are coming...importnumpyasnpfromlitdataimportoptimizefromPILimportImage# Store random images into the chunksdefrandom_images(index):data={"index":index,"image":Image.fromarray(np.random.randint(0,256,(32,32,3),np.uint8)),"class":np.random.randint(10),}returndata# The data is serialized into bytes and stored into chunks by the optimize operator.if__name__=="__main__":optimize(fn=random_images,# The function applied over each input.inputs=list(range(1000)),# Provide any inputs. The fn is applied on each item.output_dir="my_dataset",# The directory where the optimized data are stored.num_workers=4,# The number of workers. The inputs are distributed among them.chunk_bytes="64MB"# The maximum number of bytes to write into a chunk.)Theoptimizeoperator supports any data structures and types. Serialize whatever you want.2. Upload Your Data to Cloud StorageCloud providers such asAWS,Google Cloud,Azure, etc.. provide command line client to upload your data to their storage.Here is an example withAWS S3.⚡awss3cp--recursivemy_datasets3://my-bucket/my_dataset3. Use StreamingDataset and DataLoaderfromlitdataimportStreamingDatasetfromtorch.utils.dataimportDataLoader# Remote path where full dataset is persistently storedinput_dir='s3://pl-flash-data/my_dataset'# Create streaming datasetdataset=StreamingDataset(input_dir,shuffle=True)# Check any elementssample=dataset[50]img=sample['image']cls=sample['class']# Create PyTorch DataLoaderdataloader=DataLoader(dataset)Here is an illustration showing how theStreamingDatasetworks under the hood.Transform dataSimilar tooptimize, themapoperator can be used to transform data by applying a function over a list of item and persist all the files written inside the output directory.1. Put some images on a cloud storageWe generates 1000 images and upload them to AWS S3.importosfromPILimportImageimportnumpyasnpdata_dir="my_images"os.makedirs(data_dir,exist_ok=True)foriinrange(1000):width=np.random.randint(224,320)height=np.random.randint(224,320)image_path=os.path.join(data_dir,f"{i}.JPEG")Image.fromarray(np.random.randint(0,256,(width,height,3),np.uint8)).save(image_path,format="JPEG",quality=90)⚡awss3cp--recursivemy_imagess3://my-bucket/my_images2. Resize the imagesimportosfromlitdataimportmapfromPILimportImageinput_dir="s3://my-bucket/my_images"inputs=[os.path.join(input_dir,f)forfinos.listdir(input_dir)]defresize_image(image_path,output_dir):output_image_path=os.path.join(output_dir,os.path.basename(image_path))Image.open(image_path).resize((224,224)).save(output_image_path)if__name__=="__main__":map(fn=resize_image,inputs=inputs,output_dir="s3://my-bucket/my_resized_images",num_workers=4,)Easily scale data processingTo scale data processing, create a free account onlightning.aiplatform. With the platform, theoptimizeandmapcan start multiple machines to make data processing drastically faster as follows:fromlitdataimportoptimize,Machineoptimize(...num_nodes=32,machine=Machine.DATA_PREP,# You can select between dozens of optimized machines)ORfromlitdataimportmap,Machinemap(...num_nodes=32,machine=Machine.DATA_PREP,# You can select between dozens of optimized machines)The Data Prep Job UI from theLAION 400M Studiowhere we used 32 machines with 32 CPU each to download 400 million images in only 2 hours.Key FeaturesMulti-GPU / Multi-NodeTheStreamingDatasetandStreamingDataLoadertakes care of everything for you. They automatically make sure each rank receives different batch of data. There is nothing for you to do if you use them.Easy data mixingYou can easily experiment with dataset mixtures using the CombinedStreamingDataset.fromlitdataimportStreamingDataset,CombinedStreamingDatasetfromlitdata.streaming.item_loaderimportTokensLoaderfromtqdmimporttqdmimportosfromtorch.utils.dataimportDataLoadertrain_datasets=[StreamingDataset(input_dir="s3://tinyllama-template/slimpajama/train/",item_loader=TokensLoader(block_size=2048+1),# Optimized loader for tokens used by LLMsshuffle=True,drop_last=True,),StreamingDataset(input_dir="s3://tinyllama-template/starcoder/",item_loader=TokensLoader(block_size=2048+1),# Optimized loader for tokens used by LLMsshuffle=True,drop_last=True,),]# Mix SlimPajama data and Starcoder data with these proportions:weights=(0.693584,0.306416)combined_dataset=CombinedStreamingDataset(datasets=train_datasets,seed=42,weights=weights)train_dataloader=DataLoader(combined_dataset,batch_size=8,pin_memory=True,num_workers=os.cpu_count())# Iterate over the combined datasetsforbatchintqdm(train_dataloader):passStateful StreamingDataLoaderLightning Data provides a statefulStreamingDataLoader. This simplifies resuming training over large datasets.Note: TheStreamingDataLoaderis used byLit-GPTto pretrain LLMs. The statefulness still works when using a mixture of datasets with theCombinedStreamingDataset.importosimporttorchfromlitdataimportStreamingDataset,StreamingDataLoaderdataset=StreamingDataset("s3://my-bucket/my-data",shuffle=True)dataloader=StreamingDataLoader(dataset,num_workers=os.cpu_count(),batch_size=64)# Restore the dataLoader state if it existsifos.path.isfile("dataloader_state.pt"):state_dict=torch.load("dataloader_state.pt")dataloader.load_state_dict(state_dict)# Iterate over the dataforbatch_idx,batchinenumerate(dataloader):# Store the state every 1000 batchesifbatch_idx%1000==0:torch.save(dataloader.state_dict(),"dataloader_state.pt")ProfilingTheStreamingDataLoadersupports profiling your data loading. Simply use theprofile_batchesargument as follows:fromlitdataimportStreamingDataset,StreamingDataLoaderStreamingDataLoader(...,profile_batches=5)This generates a Chrome trace calledresult.json. You can visualize this trace by opening Chrome browser at thechrome://tracingURL and load the trace inside.Random accessAccess the data you need when you need it.fromlitdataimportStreamingDatasetdataset=StreamingDataset(...)print(len(dataset))# display the length of your dataprint(dataset[42])# show the 42th element of the datasetUse data transformsfromlitdataimportStreamingDataset,StreamingDataLoaderimporttorchvision.transforms.v2.functionalasFclassImagenetStreamingDataset(StreamingDataset):def__getitem__(self,index):image=super().__getitem__(index)returnF.resize(image,(224,224))dataset=ImagenetStreamingDataset(...)dataloader=StreamingDataLoader(dataset,batch_size=4)forbatchindataloader:print(batch.shape)# Out: (4, 3, 224, 224)Disk usage limitsLimit the size of the cache holding the chunks.fromlitdataimportStreamingDatasetdataset=StreamingDataset(...,max_cache_size="10GB")Support yieldWhen processing large files like compressedparquet files, you can use python yield to process and store one item at the time.frompathlibimportPathimportpyarrow.parquetaspqfromlitdataimportoptimizefromtokenizerimportTokenizerfromfunctoolsimportpartial# 1. Define a function to convert the text within the parquet files into tokensdeftokenize_fn(filepath,tokenizer=None):parquet_file=pq.ParquetFile(filepath)# Process per batch to reduce RAM usageforbatchinparquet_file.iter_batches(batch_size=8192,columns=["content"]):fortextinbatch.to_pandas()["content"]:yieldtokenizer.encode(text,bos=False,eos=True)# 2. Generate the inputsinput_dir="/teamspace/s3_connections/tinyllama-template"inputs=[str(file)forfileinPath(f"{input_dir}/starcoderdata").rglob("*.parquet")]# 3. Store the optimized data wherever you want under "/teamspace/datasets" or "/teamspace/s3_connections"outputs=optimize(fn=partial(tokenize_fn,tokenizer=Tokenizer(f"{input_dir}/checkpoints/Llama-2-7b-hf")),# Note: You can use HF tokenizer or any othersinputs=inputs,output_dir="/teamspace/datasets/starcoderdata",chunk_size=(2049*8012),)Network Drive On-Prem SupportOn-prem compute nodes can mount and use network drive. In order to reduce their network overload, theStreamingDatasetsupportscachingthe chunks.fromlightning.dataimportStreamingDatasetdataset=StreamingDataset(input_dir="local:/data/shared-drive/some-data")⚡ ContributorsWe welcome any contributions, pull requests, or issues. If you use the Streaming Dataset for your own project, please reach out to us on Slack or Discord.