Hi,
I don’t know if this is the right place, and the right moment, to post this, in case not I apologise and kindly ask to redirect me to another post, thank you.
I’m facing also a problem of memory consumption when training a model on a single GPU.
The dataset I use is quite small with respect to most of the dataset used nowadays: roughly 13000 sentences for training, 1250 for development and 3000 for testing.
I think it is not important to know what the task is, just I have sentences and I have to tag them with some labels, one label per token (basically like in POS tagging for those who are familiar).
I’m padding sentences with ‘__’ (3 underscores) so that they have all the same length of 96 tokens, and I split the training data into batches of 25 sentences each.
Since I want to use also character features, I also pad tokens with '’ so that all tokens have the same length of 17.
Finally, since I’m testing a variant of a Jordan RNN where also predicted labels are re-embedded and given as input to the network, I also use batches of labels as input to the network.
So, to summarise what the input to the network is:

the batch of tokens 96 x 25 indices (LongTensor)
the batch of labels 96 x 25 indices (LongTensor)
the batch of character sequences 96 x 25 x 17 indices (LongTensor)

Now the model:
I use character embeddings of size 50, an LSTM over characters with an output size of 100, token and label embeddings of size 200.
The hidden layer is another bidirectional-LSTM, which takes thus as input token and label embeddings, and the status output by the character-level LSTM. The output size of the second LSTM is 200 as well.
Finally I use a linear layer to map the hidden layer into the label space and I compute a log-softmax on that.
Since I re-inject predicted labels at each position in a sequence as input to the network, I perform the bi-LSTM forward and backward passes by hand (I checked, computations are correct, just it is slower than letting LSTM performs forward and backward over the whole sequence in one shot).
So, why all these details ?
Because I made few simple computations, and I found out that, assuming a double takes 10 bytes, the model should take roughly 9.5 MB of memory, while the input to the network (tokens + characters + labels) should take roughly 54 MB.
Instead, when I run the training on the GPU, the process takes 6.5GB of memory.
I know there a couple of more buffers that should be taken into account for gradients, states etc. but the difference here is more than 2 orders of magnitude.
If I run the process on CPU, the memory consumption is 1GB in total, so this accounts for the model, the whole dataset and anything else.
I’ve not a lot of experience with pytorch, does anybody can tell me if the memory consumption on GPU is normal given the model and the data I use ? In case I can give more details…
Any help would be more than appreciated
Thank you in advance