Datapipe optimisation

knoriy: <p>Hi,</p>
<p>I’m working with tar file streamed from S3 bucket. I find that the current Datapipe is slow at fetching all the data. The delay grows with more URLs (<code>data_dir</code>).</p>
<p>I’ve tried moving the <code>sharding_filter</code> before <code>load_from_tar</code>, but this results in undesired behaviour, such that when more GPUs are added, the sharding fails, and each GPU receives the full data instead of <code>batch/GPUs</code></p>
<pre><code class="lang-js">datapipe = torchdata.datapipes.iter.IterableWrapper(data_dir)\
			.shuffle()\
			.open_files_by_fsspec(mode='rb')\
			.load_from_tar() \
			.batch(2) \
			.sharding_filter()\
			.shuffle()\
			.map(self.to_sampels) \
			.batch(self.batch_size) \
			.map(self.collate_fn)

service = [DistributedReadingService(),
		MultiProcessingReadingService(num_workers=self.num_workers)]
reading_service = SequentialReadingService(*service)
dataloader =  DataLoader2(datapipe, reading_service=reading_service)
</code></pre>

AbdulsalamBande: <p>Try this config, lets see</p>
<pre><code class="lang-auto">datapipe = (
    torchdata.datapipes.iter.IterableWrapper(data_dir)
    .shuffle()
    .open_files_by_fsspec(mode='rb')
    .load_from_tar()
    .enumerate()  # Assign an index to each sample
    .sharding_filter()  # Filter samples based on their index and the GPU rank
    .shuffle()
    .map(self.to_samples)
    .batch(self.batch_size)
    .map(self.collate_fn)
)

service = [
    DistributedReadingService(),
    MultiProcessingReadingService(num_workers=self.num_workers),
]
reading_service = SequentialReadingService(*service)
dataloader = DataLoader2(datapipe, reading_service=reading_service)

</code></pre>

knoriy: <p>Tested your solution of adding <code>.enumerate()</code> but this doubled the loading time.</p>
<p>The process that is taking long is the <code>load_from_tar</code> because this is done on the main node i assume and not split across workers.</p>

AbdulsalamBande: <p>Try creating a custom datapipe that fetches tar files in parallel, extracts the content, and yields the extracted files.</p>

knoriy: <p>I don’t think this will work. The Datapipe is loading the full dataset before <code>sharding_filter</code> is applied. ideally <code>load_from_tar</code> should be performed after sharding</p>
<p>This is likely a bug or an order of operation issue.</p>

knoriy: <p>I also tried this approach, but it is still slow.</p>

nivek: <p>Hi <a class="mention" href="/u/knoriy">@knoriy</a></p>
<p>You almost certainly want to use <code>.sharding_filter()</code> immediately after your first <code>shuffle()</code>. Otherwise, you will be duplicating work across workers (such as <code>open_files_by_fsspec</code>, <code>load_from_tar</code>, <code>batch(2)</code>).</p>
<blockquote>
<p>but this results in undesired behaviour, such that when more GPUs are added, the sharding fails, and each GPU receives the full data instead of <code>batch/GPUs</code></p>
</blockquote>
<p>As for this, can you further elaborate? Perhaps provide something reproducible?</p>
<p>For example, it would be useful to check <code>torch.distributed.get_world_size()</code> and <code>torch.distributed.get_rank()</code> as those will be used for sharding by <code>DistributedReadingService</code> within <code>DataLoader2</code>.</p>

knoriy: <p>Hi <a class="mention" href="/u/nivek">@nivek</a>, Thank you for your comment.</p>
<p>Re:</p>
<aside class="quote no-group" data-username="nivek" data-post="7" data-topic="176758">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/nivek/48/39672_2.png" class="avatar"> nivek:</div>
<blockquote>
<p>You almost certainly want to use <code>.sharding_filter()</code> immediately after your first <code>shuffle()</code>. Otherwise, you will be duplicating work across workers (such as <code>open_files_by_fsspec</code>, <code>load_from_tar</code>, <code>batch(2)</code>).</p>
</blockquote>
</aside>
<p>I agree with your statement. Unfortunately, this results in unexpected behaviour; these are:</p>
<ol>
<li>when <code>.shardfilter</code> is above each <code>batch(2)</code>, GPU received the full dataset and not <code>batch/gpu</code>, I fixed this by changing the order of the reading service (<code>[MPRS, DRS] &gt;&gt; [DRS, MPRS]</code>) and moving <code>.sharding_filter</code> below <code>.batch(2)</code>.</li>
<li>Using <code>MPRS</code> and them <code>DRS</code> resulted in getting <code>can't pickle: ExFileObject</code></li>
</ol>
<p>Re:</p>
<aside class="quote no-group" data-username="nivek" data-post="7" data-topic="176758">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/nivek/48/39672_2.png" class="avatar"> nivek:</div>
<blockquote>
<p>As for this, can you further elaborate? Perhaps provide something reproducible?</p>
</blockquote>
</aside>
<p>I fixed this by reordering the reading services. What is the significance of the order by which reading services are passed to <code>dataloader2</code>?</p>
<p>Here is a simple working example: using <code>pytorch==2.1.0</code> and <code>pytorch_lightning=2.0</code>. Where <code>"path/to/*.tar"</code> (line 21) is a path to an s3 bucket holding tar files containing <code>*.wav</code> and its corresponding<code>*.json</code> label.</p>
<pre data-code-wrap="python"><code class="lang-python">import io
import soundfile
import json

import torch
import torch.nn as nn
import torchdata
from torchdata.dataloader2 import DataLoader2, DistributedReadingService, MultiProcessingReadingService, SequentialReadingService 

import pytorch_lightning as pl

class PLModel(pl.LightningModule):
	def __init__(self):
		super().__init__()
		self.layer = nn.Linear(32, 2)

	def forward(self, x):
		return x

	def training_step(self, batch, batch_idx):
		out = self(batch)

	def configure_optimizers(self):
		return torch.optim.SGD(self.layer.parameters(), lr=0.1)

class PL_datamodule(pl.LightningDataModule):
	def __init__(self):
		super().__init__()
		self.urls = ["path/to/0.tar", "path/to/1.tar"]

		self.batch_size = 16
		self.num_workers = 12

	def get_pipeline(self, data_dir):
		datapipe = torchdata.datapipes.iter.IterableWrapper(data_dir)\
			.shuffle()\
			.open_files_by_fsspec(mode='rb')\
			.load_from_tar() \
			.batch(2) \
			.sharding_filter()\
			.shuffle(buffer_size=32)\
			.map(self.to_sampels) \
			.batch(self.batch_size) \
			# .map(self.collate_fn)
		
		return datapipe
	
	def setup(self, stage=None):
		self.train = self.get_pipeline(self.urls)

	def train_dataloader(self):
		service = [
			DistributedReadingService(),
			MultiProcessingReadingService(num_workers=self.num_workers),
		]
		reading_service = SequentialReadingService(*service)
		return DataLoader2(self.train, reading_service=reading_service)

	def to_sampels(self, data):
		a, t = data
		return soundfile.read(io.BytesIO(a[1].read())), json.loads(t[1].read().decode('utf-8'))


def main(devices=1, max_epochs=1, strategy='ddp'):
	model = PLModel()
	datamodule = PL_datamodule()
	trainer = pl.Trainer(devices=devices, max_epochs=max_epochs, strategy=strategy)
	trainer.fit(model, datamodule=datamodule)

if __name__ == '__main__':
	main(devices=2, max_epochs=1, strategy='ddp')
</code></pre>

knoriy: <p>Furthermore, regarding what was causing the slowdown in my original post, This was caused by the large default <code>buffer_size</code> of the <code>shuffle</code> after the <code>sharding_filter</code>. Setting this to a smaller number helped.</p>
<aside class="quote no-group" data-username="knoriy" data-post="1" data-topic="176758">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/knoriy/48/35651_2.png" class="avatar"> knoriy:</div>
<blockquote>
<pre><code class="lang-auto">datapipe = torchdata.datapipes.iter.IterableWrapper(data_dir)\
			.shuffle()\
			.open_files_by_fsspec(mode='rb')\
			.load_from_tar() \
			.batch(2) \
			.sharding_filter()\
			.shuffle()\
			.map(self.to_sampels) \
			.batch(self.batch_size) \
			.map(self.collate_fn)

service = [DistributedReadingService(),
		MultiProcessingReadingService(num_workers=self.num_workers)]
reading_service = SequentialReadingService(*service)
dataloader =  DataLoader2(datapipe, reading_service=reading_service)
</code></pre>
</blockquote>
</aside>

cinjon: <p>I am doing almost exactly the same thing but running into an error when loading from tar - “cannot pickle ‘ExFileObject’ object”. Did you run into this knoriy?</p>

knoriy: <p>Yes, I have. Can you share your pipeline here?</p>

cinjon: <p>I am no longer running into that error, but I am running into another issue. Help here would be greatly appreciated.</p>
<p>The new problem is that sometimes the fetch time is horrendous. I think what’s happening is that there are times when one worker has to do all of the work. The reason I think that is because the mean/median/max of a single worker (num_workers = 1) operating is 30.3s / 31.0s / 40s. When I use num_workers=4, then mean=4-5s, median~0.001s, and it will spike to max ~40-50.</p>
<p>I’m <em>pretty sure</em> that the big issue is the load_from_tar(). When I use just one worker, that and groupby (plus the batch size at the bottom) is enough to spike the timing.</p>
<p>Do you have any ideas for fixing this? I tried adding a sharding_filter just before the load_from_tar but that made it worse →  median=0.0018754005432128906 mean=5.149842675139264. I tried adding it just after the groupby and then it slows down precipitously.</p>
<pre><code class="lang-auto">def get_chunked_dataloader(datapipe, num_workers):
  # The datapipe is the result of get_chunked_dataset below.
  mp = MultiProcessingReadingService(num_workers=num_workers)
  return DataLoader2(datapipe, reading_service=mp)  

def get_chunked_dataset(split, batch_size=100):
  assert split in ['train', 'val', 'test']
  bucket = "s3://&lt;BUCKET&gt;/%s" % split
  dp = IterableWrapper([bucket]).list_files_by_fsspec()
  dp = dp.shuffle()
  dp = dp.open_files_by_fsspec(mode="rb")
  dp = dp.prefetch(5)
  dp = dp.load_from_tar() # "r"
  dp = dp.groupby(groupby, group_size=2, guaranteed_group_size=2)
  dp = dp.map(process)
  # ExpandDataPipe is an IterDataPipe that takes in a source data pipe and 
  # makes ~3-4 examples from each small wav + label combo it receives. 
  # It's reasonably fast.
  dp = ExpandDataPipe(dp)
  dp = dp.shuffle(buffer_size=20)
  dp = dp.batch(batch_size)
  dp = dp.map(collate)
  return dp 

def process(row):
    if row[0][0].endswith('json'):
        stream_json_wrapper = row[0][1]
        stream_wav_wrapper = row[1][1]
    else:
        stream_json_wrapper = row[1][1]
        stream_wav_wrapper = row[0][1]
    labels = stream_json_wrapper.read()
    labels = json.loads(labels.decode('utf-8'))
    wav = io.BytesIO(stream_wav_wrapper.read())
    wav, _ = torchaudio.load(wav)    
    return wav, labels

def groupby(row):
  return os.path.basename(row[0]).split(".")[0]
</code></pre>

cinjon: <p><a class="mention" href="/u/knoriy">@knoriy</a> , mind taking a look? It looks similar to this one → <a href="https://discuss.pytorch.org/t/slow-aws-s3-data-loading-using-torchdata-open-files-by-fsspec/181129" class="inline-onebox">Slow aws S3 data loading using TorchData open_files_by_fsspec</a></p>

knoriy: <p>It looks like you’re not using <code>.sharding_filter()</code>. This could mean all the loading is happening in one place and not distributed appropriately.</p>
<p>See <a href="https://pytorch.org/data/main/dp_tutorial.html#:~:text=The%20reason%20why,add%20the%20following." rel="noopener nofollow ugc">HERE</a></p>
<p>The docs mention you should place it as early into your pipeline as possible. In your instance, this would be after the first <code>.shuffle()</code></p>
<p>Give that a try <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

cinjon: <p>I tried that and it didn’t work. GPU power usage was between 60% and 80% the whole time. Without the sharding_filter, it’s much higher and bounces down to the 60s only when I’m loading in the val set. Any other ideas?</p>

kchoi: <p>hi all. this thread is not very young but… (hi <a class="mention" href="/u/cinjon">@cinjon</a>! aren’t you … the Cinjon with audio research, also KC’s student? <img src="https://discuss.pytorch.org/images/emoji/apple/wave.png?v=12" title=":wave:" class="emoji" alt=":wave:" loading="lazy" width="20" height="20"> )</p>
<p>did <a class="mention" href="/u/cinjon">@cinjon</a> or <a class="mention" href="/u/knoriy">@knoriy</a> figure out how to configure a datapipe with fast S3 loading? i would appreciate it very, very much if you can share your working setup <img src="https://discuss.pytorch.org/images/emoji/apple/pray.png?v=12" title=":pray:" class="emoji" alt=":pray:" loading="lazy" width="20" height="20"></p>

cinjon: <p>hehe hi! nice to see you here.</p>
<p>i gave up on this approach and went with mosaic’s streamer. it’s quite good and did/does the job well → <a href="https://github.com/mosaicml/streaming" class="inline-onebox" rel="noopener nofollow ugc">GitHub - mosaicml/streaming: A Data Streaming Library for Efficient Neural Network Training</a>.</p>

kchoi: <p>haha <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20"><br>
gotcha, i’ve been debating so much about Torchdata… thanks for another datapoint, <a class="mention" href="/u/cinjon">@cinjon</a>!</p>