Why does my datapipe based dataloader take up so much memory?

Vedant_Roy: <p>I am training on 2 GPUs, each of which uses a dataloader with 16 workers. Due to the nature of my training objective (contrastive loss), each GPU must receive different samples, in order to prevent a batch with duplicate samples.</p>
<p>I’m noticing that my dataloader takes up a very large amount of memory (~64 GB &amp; then my computer crashes), and I don’t know why.</p>
<p>For context, my images are size 240x240, and my batch size is 512. Thus, I expect my 2 dataloaders to take up 2 * 16 * 3 * 240 * 240 * 512 = 2831155200 bytes = 2.831 GB of memory at peak usage.</p>
<p>Here’s what the code looks like:</p>
<pre><code class="lang-auto"># friendly wrapper over torch.distributed
import composer.utils.dist as dist
def load_parquet(path):
    table = pq.read_table(path)
    imgs = table["img"]
    imgs = [load_tensor(img.as_py()) for img in imgs]
    del table

    shard_idx = dist.get_global_rank()
    worker_id = get_worker_info().id
    print(f"Loaded {len(imgs)} imgs (shape={imgs[0].shape}) from {path} with shard={shard_idx} &amp; worker_id={worker_id}")
    return imgs

def sharded_parquet_dataloader(dir, **kwargs):
    print(f"Creating dl with kwargs: {kwargs}")
    datapipe = sharded_local_datapipe(dir)
    datapipe = datapipe.sharding_filter()
    datapipe = datapipe.map(load_parquet)
    datapipe = datapipe.flatmap()
    dl = DataLoader(
        datapipe, **kwargs
    )
    return dl
</code></pre>
<p>where the <code>sharder_local_datapipe</code> is a tool to ensure that each GPU gets unique samples:</p>
<pre><code class="lang-auto">def sharded_local_datapipe(dir: Path):
    world_size = dist.get_world_size()
    shard_idx = dist.get_global_rank()
    total_files = count_files(dir)
    # All data loaders should have the same size
    shard_size = total_files // world_size

    # For this to work, we assume the Lister is deterministic
    pipe = dp.iter.FSSpecFileLister(str(dir))
    pipe = pipe.enumerate()

    def in_shard(idx_and_name):
        idx, _ = idx_and_name
        return (idx % world_size) == shard_idx

    def take_last(args):
        return args[-1]

    pipe = pipe.filter(in_shard)
    pipe = pipe.map(take_last)
    return pipe.header(shard_size)
</code></pre>
<p>Does anyone know what might be going on here? Print statements indicate that not that many images are being loaded, so I’m trying to figure out where the extra memory is coming from:</p>
<pre><code class="lang-auto">contrastive_train-contrastive_train-1  | Train dataset: /data/train
contrastive_train-contrastive_train-1  | Val dataset: None
contrastive_train-contrastive_train-1  | t_warmup: 100ba
contrastive_train-contrastive_train-1  | t_max: 3600ba
contrastive_train-contrastive_train-1  | Creating dl with kwargs: {'batch_size': 512, 'num_workers': 16, 'pin_memory': True, 'drop_last': True}
contrastive_train-contrastive_train-1  | Checking dataloader...
contrastive_train-contrastive_train-1  | Loaded 243 imgs (shape=torch.Size([1, 3, 240, 240])) from /data/train/102.parquet with shard=0 &amp; worker_id=4
contrastive_train-contrastive_train-1  | flat-mapping: 243, shard=0, worker_id=4
contrastive_train-contrastive_train-1  | Loaded 242 imgs (shape=torch.Size([1, 3, 240, 240])) from /data/train/97.parquet with shard=0 &amp; worker_id=3
contrastive_train-contrastive_train-1  | Loaded 228 imgs (shape=torch.Size([1, 3, 240, 240])) from /data/train/143.parquet with shard=0 &amp; worker_id=5
contrastive_train-contrastive_train-1  | Loaded 228 imgs (shape=torch.Size([1, 3, 240, 240])) from /data/train/135.parquet with shard=0 &amp; worker_id=12
contrastive_train-contrastive_train-1  | Loaded 228 imgs (shape=torch.Size([1, 3, 240, 240])) from /data/train/124.parquet with shard=0 &amp; worker_id=4
contrastive_train-contrastive_train-1  | flat-mapping: 228, shard=0, worker_id=4
contrastive_train-contrastive_train-1  | Loaded 249 imgs (shape=torch.Size([1, 3, 240, 240])) from /data/train/123.parquet with shard=0 &amp; worker_id=14
contrastive_train-contrastive_train-1  | Loaded 254 imgs (shape=torch.Size([1, 3, 240, 240])) from /data/train/117.parquet with shard=0 &amp; worker_id=15
contrastive_train-contrastive_train-1  | flat-mapping: 254, shard=0, worker_id=15
contrastive_train-contrastive_train-1  | Loaded 200 imgs (shape=torch.Size([1, 3, 240, 240])) from /data/train/129.parquet with shard=0 &amp; worker_id=3
contrastive_train-contrastive_train-1  | Loaded 214 imgs (shape=torch.Size([1, 3, 240, 240])) from /data/train/110.parquet with shard=0 &amp; worker_id=4
contrastive_train-contrastive_train-1  | Loaded 258 imgs (shape=torch.Size([1, 3, 240, 240])) from /data/train/116.parquet with shard=0 &amp; worker_id=3
contrastive_train-contrastive_train-1  | Loaded 1000 imgs (shape=torch.Size([1, 3, 240, 240])) from /data/train/48.parquet with shard=0 &amp; worker_id=10
contrastive_train-contrastive_train-1  | Loaded 1000 imgs (shape=torch.Size([1, 3, 240, 240])) from /data/train/45.parquet with shard=0 &amp; worker_id=0
^CGracefully stopping... (press Ctrl+C again to force)
</code></pre>

ejguan: <p>Wondering what would be the memory usage without using <code>pin_memory</code>?</p>

Vedant_Roy: <p>I’ll take a look. AFAIK, the memory usage is roughly the same.<br>
What would be a good way to debug something like “too much memory usage”?</p>
<p>Is there anything worth using here besides <a href="https://docs.python.org/3/library/tracemalloc.html#module-tracemalloc" rel="noopener nofollow ugc"><code>tracemalloc</code></a>?</p>
<p>Also, is there a way to check how much memory is being pinned?</p>

Vedant_Roy: <p>Without any dataloaders, the memory usage is  22 GB.<br>
With a pinned_memory dataloader, the usage is 49 GB.<br>
Without pinned_memory, the memory usage is 40.5 GB.</p>
<p>While not pinning memory, the average usage is still a few orders of magnitude greater than I expected.</p>

ejguan: <p>cc: <a class="mention" href="/u/nivek">@nivek</a> to see if you have any recommendation on the tools to monitor emory usage.</p>
<p>I think <a href="https://github.com/janestreet/magic-trace" class="inline-onebox" rel="noopener nofollow ugc">GitHub - janestreet/magic-trace: magic-trace collects and displays high-resolution traces of what a process is doing</a> is a good tool to do profiling.</p>
<p>BTW, have you ever tried to use this DataPipe to load from parquet?</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/data/blob/daa2fc9e9d6f5965618cf026a6e2feda93a16690/torchdata/datapipes/iter/util/dataframemaker.py#L94">
  <header class="source">

      <a href="https://github.com/pytorch/data/blob/daa2fc9e9d6f5965618cf026a6e2feda93a16690/torchdata/datapipes/iter/util/dataframemaker.py#L94" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/data/blob/daa2fc9e9d6f5965618cf026a6e2feda93a16690/torchdata/datapipes/iter/util/dataframemaker.py#L94" target="_blank" rel="noopener nofollow ugc">pytorch/data/blob/daa2fc9e9d6f5965618cf026a6e2feda93a16690/torchdata/datapipes/iter/util/dataframemaker.py#L94</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="84" style="counter-reset: li-counter 83 ;">
          <li>                "Please visit https://github.com/facebookresearch/torcharrow/ to install it."</li>
          <li>            )</li>
          <li>        # In this version, DF tracing is not available, which would allow DataPipe to run DataFrame operations</li>
          <li>        batch_dp = source_dp.batch(dataframe_size)</li>
          <li>        df_dp = batch_dp.map(</li>
          <li>            partial(_construct_dataframe, dtype=dtype, dtype_generator=dtype_generator, columns=columns, device=device)</li>
          <li>        )</li>
          <li>        return df_dp</li>
          <li>
          </li>
<li>
          </li>
<li class="selected">@functional_datapipe("load_parquet_as_df")</li>
          <li>class ParquetDFLoaderIterDataPipe(IterDataPipe):  # IterDataPipe[torcharrow.IDataFrame[T_co]]</li>
          <li>    r"""</li>
          <li>    Takes in paths to Parquet files and return a `TorchArrow` DataFrame for each row group</li>
          <li>    within a Parquet file (functional name: ``load_parquet_as_df``).</li>
          <li>
          </li>
<li>    Args:</li>
          <li>        source_dp: source DataPipe containing paths to the Parquet files</li>
          <li>        columns: List of `str` that specifies the column names of the DataFrame</li>
          <li>        use_threads: if ``True``, Parquet reader will perform multi-threaded column reads</li>
          <li>        dtype: specify the `TorchArrow` dtype for the DataFrame, use ``torcharrow.dtypes.DType``</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>
<p>
I believe this is kind of streaming way to read parquet file, which should reduce memory footprint in theory.</p>

nivek: <p><a href="https://github.com/plasma-umass/scalene" rel="noopener nofollow ugc">scalene</a> is another option. It can give you memory usage line by line.</p>