Behavior of dataloader when resuming training from the existing checkpoint

SeoHyeong: <p>Hello,</p>
<p>As the title states, I have a question on the behavior of the torch dataloader when I resume training process from the existing checkpoint.</p>
<p>My codes for loading dataset and dataloader looks like the following. (I’m using Huggingface’s datasets.)</p>
<pre><code class="lang-auto">train_dataset = load_dataset(args.dataset_loading_script_path,
                                     data_files=args.dataset_txt,
                                     split='train')
train_dataset = train_dataset.map(lambda examples: tokenizer(examples['sent_1'],
   max_length=args.max_length, padding='max_length', truncation=True), batched=True)
train_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask',])
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=total_train_batch_size, shuffle=True)
</code></pre>
<p><strong>My question is, when I load the existing checkpoint and optimizer and resume training, is there a way to avoid training on batches (or training examples) that I’ve already trained on?</strong></p>
<p>I’m assuming torch dataloader randomly shuffles data and feeds batches into the model each time the training starts and just simply skips the number of trained steps when iterating through the train dataset using dataloader.</p>
<p>If this is the case, will I be able to achieve what I’m trying to do by shuffling dataset before feeding it into the dataloader (and save dataset to reuse) and drop <code>shuffle=True</code> option?</p>

Epoching: <p>Are you trying to train your model for only 1 epoch because you have so much data and it’ll take too long to do more, or are you possibly trying to do 1 epoch because your machine won’t allow it to finish and everything shuts off, so you’d like to save intermediate progress? (Epoch = single pass through your entire dataset) Just asking out of curiousity, no worries if there’s no reason.</p>
<p>As for for your question, I’d do one of the following:</p>
<ol>
<li>
<p>Drop <code>shuffle=True</code> and as you train keep track of an id (either the step number, which will represent what batch you are on, or just the raw id of current sample you’re on). If you’re using a HuggingFace <code>Trainer</code> instance for your model training, you can use callbacks to do this (add a <code>on_step_end</code> or <code>on_step_begin</code> callback to write out current step # to a file, can be found <a href="https://huggingface.co/transformers/main_classes/callback.html#transformers.TrainerCallback.on_step_end" rel="noopener nofollow ugc">here in the docs</a>). When continuing training, You can slice <code>examples</code> starting from the id you left on, and ending with the last id of the dataset, then append all the samples you’ve already trained with at the end of this slice (essentially shifting the samples you trained with, but putting them at the end). If you don’t care about re-using the samples at the end, you can just use PyTorch’s <code>Subset</code> dataset class.</p>
</li>
<li>
<p>Keep <code>shuffle=True</code> but have a small function call when you fetch a sample that writes-out the id that’s getting fetched/processed. When continuing training, do a similar process as above (option 1) but rather than working with a single slice from <code>shuffle=False</code> you can slice out a subset of your dataset using the ids you’ve saved</p>
</li>
</ol>

SeoHyeong: <p>Thanks for the super helpful response!</p>
<p>Just to be clear, for the first option you’ve suggested, I should randomly shuffle the order of each training example when creating the dataset (but drop <code>shuffle=True</code>), right?</p>
<p>My training data is over a million pairs of sentences extracted from Wikipedia, and if I don’t keep my dataset randomly shuffled, batches will consist of adjacent sentences from the same Wiki document (which I’m trying to avoid). And yes, the reason I’m trying to train my model for only 1 epoch is that I have so much data, which leads to a super long training time with the currently available resource I have.</p>

Epoching: <blockquote>
<p>Just to be clear, for the first option you’ve suggested, I should randomly shuffle the order of each training example when creating the dataset (but drop <code>shuffle=True</code> ), right?</p>
</blockquote>
<p>That’s a nice idea too! Just make sure when you continue training that the pre-shuffled dataset is in the same order!</p>
<blockquote>
<p>My training data is over a million pairs of sentences extracted from Wikipedia</p>
</blockquote>
<p>Woah 0.0</p>
<p>Yea I see the reasoning for this now. Neat problem! <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>

SeoHyeong: <p>Gotcha! I’ll combine randomly shuffling the order of my sentences and your suggested option 1.</p>
<p>Thanks again! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

Dhia680: <p>I’m facing the same challenge. When trying to resume mid-epoch.<br>
Training a few billion-parameters model on ~200 to 500GB of data with FSDP.<br>
I opted for dropping <code>shuffle=True</code>, and saving last batch ID (<code>saved_id</code>).<br>
(Actually I just fix the random seed. Which is similar to shuffling dataset then dropping <code>shuffle=True</code> in dataloader).<br>
And when I resume from checkpoint, I make dataloader loop over all batches until reaching the first batch with id &gt; <code>saved_id</code>. At that moment, I restart training.<br>
However, loading all batches (from cpu), even without doing anything else (no forward-backwards, updates…) takes near 1h !<br>
Is there a hack especially with this new <a href="https://pytorch.org/data/beta/torchdata.stateful_dataloader.html" rel="noopener nofollow ugc">stateful_dataloader</a> class to conserve the ‘<code>state_id</code>’ which should correspond to <code>saved_id</code> and resume without having to loop over already seen data or having to slice the dataset ?<br>
I fear that the remark mentionned in <a href="https://pytorch.org/data/beta/torchdata.stateful_dataloader.html" rel="noopener nofollow ugc">stateful_dataloader</a> saying: “<code>Under the hood, StatefulDataLoader handles aggregation and distribution of state across multiprocess workers (but not across ranks)</code>”.<br>
My setting is distributed across multiple ranks (gpu nodes).</p>
<p>p.s: I want to avoid the solution 1 of cutting off the already seen data given the <code>saved_id</code>, because data gets streamed to my compute instance, not downloaded once.</p>

Dhia680: <p>Actually stateful dataloader is a good solution.<br>
You just have to save its state dict with your model checkpoint.<br>
Then, when resuming, load this saved state dict in your newly initialized Dataloader. <br><br>
This will make it remember latest batch indices for each worker and skip them without loading them. <br><br>
The only compromise is that this requires resuming with the exact same data mixture and compute budget (<code>world_size</code>). <br></p>