I’m training a Transformer language model with 5000 vocabularies using a single M60 GPU (w/ actually usable memory about 7.5G).
The number of tokens per batch is about 8000, and the hidden dimension to the softmax layer is 512. In other words, the input to nn.Linear(256, 5000) is of size [256, 32, 256]. So, if I understand correctly, a fully-connected layer theoretically consumes 5000x8000x512x4=81.92GB of GPU memory for a forward pass (4 is for float32). But the GPU performed the forward and backward passes without any problem, and it says the GPU memory usage is less than 7GB in total.
What’s causing this?