Optimizing Data Loading for Large-Scale Distributed Training

ouhenio: <p>I’m currently working on training a GPT model using <a href="https://pile.eleuther.ai/" rel="noopener nofollow ugc">The Pile</a> dataset on a single node with 8 A100 GPUs. This dataset is quite large, weighing in at around 950 GB. To achieve efficient distributed training, I’m leveraging <code>torchrun</code> for its ease of use and seamless integration. However, I’ve encountered an issue where the data loading process seems to be triggered 8 times in parallel, which, I suspect, leads to excessive disk read overhead and consequently, a slow loading process.</p>
<p>When running the training on a single GPU, the dataset takes about 10 minutes to load. However, with all 8 GPUs in action, the loading time increases significantly, taking over an hour.</p>
<p>Here’s the command I’m using to initiate the training:</p>
<p><code>torchrun -m --nproc_per_node=8 gpt train</code></p>
<p>And this is the general structure of the relevant training code:</p>
<pre><code class="lang-python">def train(cfg: DictConfig):

    # Setup data and dataloader
    data_module = prepare_data_module(**cfg.data)

    train_dataloader = data_module.train_dataloader()
    val_dataloader = data_module.val_dataloader()

    # Extract tokenizer from datamodule
    tokenizer = data_module.tokenizer

    # Setup model and optimizer
    model = GPT(**cfg.model)

    # Setup data collator
    data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)

    train_args = TrainingArguments(**cfg.train_args)

    trainer = Trainer(
        model=model,
        tokenizer=tokenizer,
        train_args=args,
        data_collator=data_collator,
        train_dataset=train_dataloader,
        eval_dataset=val_dataloader,
    )

    trainer.train()
</code></pre>
<p>I am seeking insights from the community on potential optimizations to resolve this data loading bottleneck. Are there any recommendations or best practices for efficiently handling data loading and distribution when training large-scale models like GPT with PyTorch?</p>
<p>Looking forward to your thoughts and suggestions!</p>

kumpera: <p>cc <a class="mention" href="/u/vitalyfedyunin">@VitalyFedyunin</a> for DataLoader questions</p>

suraj.pt: <p>MosaicML’s <code>streaming</code> API tries to address this exact usecase: <a href="https://docs.mosaicml.com/projects/streaming/en/stable" class="inline-onebox">Streaming</a><br>
Since it’s a drop-in replacement for <code>Dataset</code> it might be appropriate for your program</p>

ouhenio: <p>Hi <a class="mention" href="/u/suraj.pt">@suraj.pt</a>!</p>
<p>No way, you replied to my post! I learned about this stuff though your videos. It’s amazing to get advice  straight from you, thanks!</p>
<p>I looked into <code>Streaming</code>, but I’m not sure if its useful in my case.</p>
<p>It’s documentation states:</p>
<blockquote>
<p>StreamingDataset helps to make training on large datasets from <strong>cloud storage</strong> as fast, cheap, and scalable as possible.</p>
</blockquote>
<p>I’m reading the data locally, so if I’m not mistaken this library would not work in my case, right?</p>

ouhenio: <p>I just realized it has a <code>local</code> option. I’ll look into it further.</p>
<p>Thanks a lot <a class="mention" href="/u/suraj.pt">@suraj.pt</a>!</p>

kissakos: <p>What does the dataset initialisation do? I assume you don’t preload the dataset, because it would mean loading 8 x num_workers x 950GB in memory. If you only generate some metadata for the loader, you could cache this data in a temporary file and load the cache in the dataset.</p>