Lightning Talk: Fast, Scalable Distributed Training with StreamingDataset - Saaketh Narayan

Key Takeaways

so I'm s the Ryan uh today I'm just going to talk a little bit how um our open source streaming data set enables uh fast scalable distributed training so yeah so I guess this is just what we're going to talk about um go through these four and then hopefully we'll have some time for Q&A at the end um cool so streaming data set is a pytorch iterable data set um it is built purpose built for fast distributed training um and it makes uh basically stream streaming your data from cloud storage as fast and cheap and easy as possible um so some examples of how we've been able to train with streaming as well as our customers are um we've been able to scale with streaming to thousands of gpus um dbrx for example was trained on uh more than 3,000 gpus and is a 132 billion parameter model um we've also internally used huge data sets like hundreds of terabytes and streaming has performed well on all of these um across many data modalities and many model types um so why is streaming needed so without streaming there's a bunch of challenges that you see especially related to your data and and model training honestly um is oftentimes blocked by data challenges so without streaming you have you know your data can't be ingested by your models your training is slowed down because of data loader bottlenecks um you can't combine multiple data sources training resumption becomes very hard Etc and and with streaming as you can see I'm not just going to go I'm not going to go through all these but um we've been able to successfully address address a lot of these challenges so um one thing we really focus on is getting the most throughput out of your training uh out of your data loading so on the left here you see some an example of having really bad throughput so at sporadic intervals during your training process you'll be dead loader bottlenecked and you'll see like frequent throughput drops just like on the left um if you really have good throughput you should see something like on the right um the second thing here is data mixing and shuffling so in this example we have a bunch of data sources so C4 GitHub and like an internal data set um streaming lets you just create stream objects over these data sources specify some sort of proportion and seamlessly mix these data sets um during your training um another feature that is really important to us is elastically deterministic training so with streaming data set as you see on the top you get fully deterministic loss curves whether you use 1 8 16 32 or 64 gpus um this is really important because it lets you automatically scale whatever runs you did at small scale to very large GPU scales even in the middle of training you can just take your run on 32 gpus say and scale it up to 128 um and you will get deterministic resumption every single time so here's another example of where deterministic resumption comes into play if I have data that's stored among many Cloud providers then I can just continue this training run uh let's say you know I have a oci failure or something I want to trade from S3 then I can just resume my training run straight from s three with the same data set um so how does this all work uh the first step is you need to convert your data into a streaming data set inestable format and the second thing is there's uh some mechanisms that happen during training that sort of enable all these features that I just talked about so the first step is data set conversion so you take your raw data set samples shown here in Gray um let's say they're like a CSV file whatever you want and you pass these through MDS writer objects so what these do is these serialize your raw data set samples into binary um using these uh MDS writer objects and write out um row oriented Mosaic data shards so that's shown here in blue basically each Shard has your serialized sample as well as a header that describes where to seek to to access that particular sample so for example I'll I'll write out maybe an image and it'll say Okay image 3 is at you know um byte let's say 500 and and we can very easily very quickly just uh access the sample in the header and then seek to where that sample lives in The Shard and that allows for really fast low latency random sample access during training what happens is we store basically you can store all of your data sets or data sources in in the cloud and so here we have two data sources or streams um each one of them has a bunch of shards as well as an Associated index. Json file so this index. Json file describes all of your sample uh data set metadata so for example how many shards are in there how many samples are in each Shard uh how many bytes do I expect from each sample stuff like that um and dur when the streaming data set is initialized these index. Json files are downloaded so we have an idea of what exactly your data is composed of then we deterministically partition your samples among nodes gpus and data loaded workers so here's a sort of minimized view of this um my total sample space is composed of samples across these two data sources so that's shown here in blue and in green I have I'm training on two nodes each with eight gpus so I want to split up all of these samples across workers nodes and gpus such that when I increase my number of nodes this partitioning stays the same so what this allows me to do is basically I need to maintain the invariant that each Global batch regardless of the number of nodes or gpus or workers I'm using is composed of the same number of samples um not just the same number of samples but the same sample IDs um so we make sure that this determinist uh this partitioning is deterministic and then we can also apply some partial shuffling algorithms that are going to take the samples that are assigned to each node and uh distribute them among among the nodes such that your data set egress costs are still minimized so you still have the same sample IDs and shards assigned to each node but you have really good uh Shuffle quality so this enables a lot of different uh things that you can see in our documentation such as different patching methods various different shuffles and we've done a lot of work to make sure that these partial shuffles are equivalent in quality and model conversions to like fully shuffling your data which um if you're trying to download shards is is highly inefficient so during training what happens is uh your data lives in in cloud storage so each datal loader worker is assigned some portion of the sample partition tensor and dat loader workers check if The Shard that contains the sample that they need is present on disk and if it's not they go ahead and download it and this happens sort of uh preemptively this happens ahead of time and we're able to prefetch shards um these shards are downloaded from the cloud and then because of the row oriented very low latency uh MDS format we're able to just um very quickly retrieve these samples each data loader worker retrieves these samples from the shards that are on disk and loads them into GPU for training so yeah that's sort of an overview of of streaming um and yeah would love to feel any questions from you guys thank you