If you hold (mini) batch size constant (as well as everything else) but increase the number of examples (and therefore the number of training iterations), should you expect a (significant) increase in GPU memory since?

My instinct is to say no, because only a (mini) batch is uploaded on to the GPU at a time. However this is not what I am seeing in practice on a project I am currently working on. Moving from a very small dataset (1000 examples) to a large one (~3 million examples) caused me to have to decrease batch size 32-fold to avoid OOM errors (from 64 to 2). I don't understand why increasing the size of the dataset should cause such a spike in memory usage- shouldn't the dataset be held primarily in CPU memory?

Since this might depend on the specific code packages being used, I am using the transformers library for training with pytorch models. If model architecture matters, I'm using OPT (open version of GPT). If modality matters, my dataset examples are text sentences, tokenized by whitespace.

EDIT (2): AFAIK, the entire dataset should not be loaded on to the GPU at once. I am using the Trainer and DataCollatorForLanguageModeling classes from the transformers library to train and collate data, respectively; and the datasets library to hold the data (both libraries by HuggingFace)- if I am wrong and these libraries do load the entire dataset into GPU memory, please correct me.

EDIT (1): To clarify, by "use more/less GPU memory," I mean a greater/lesser percentage of available GPU memory being in use at any given moment.


1 Answer 1


It depends on how you actually load your data on the GPU: if you load your whole dataset on the GPU, then increasing the dataset size will certainly increase the GPU memory consumption.

If the project started with a very small dataset, as it seems to be the case, it is likely that the person who created the code did not bother to load the data in batches or use data loaders that handle that, and just loaded the whole dataset into the GPU memory.

In conclusion: check the data loading source code to see if the data is loaded all at once or not.

  • $\begingroup$ AFAIK, the entire dataset is not be loaded on to the GPU at once. I believe this was implied but I've edited my post to make it explicit. Is that the only reason that I would see such behavior? $\endgroup$
    – ubadub
    Nov 22, 2023 at 18:52
  • $\begingroup$ Yes, that's the typical reason. If you are just loading the current batch in the GPU memory, there is no reason for the GPU memory consumption to increase. $\endgroup$
    – noe
    Nov 22, 2023 at 19:36

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.