"Loss: nan" in training about mgpt HOT 3 CLOSED

Recent experiment shows that model with gradient checkpointing and SGD optim fits into 16Gb of GPU RAM.
Try this model initialization:
model = GPT2LMHeadModel.from_pretrained("sberbank-ai/mGPT", gradient_checkpointing=True, use_cache=False), but remove further conversion model = model.half()
For me it consumes 15774Mb of RAM with batch size 1, seq length 512.

from mgpt.

Thank you! It really worked.

from mgpt.

Hi! Unfortunately, the amount of GPU memory of colab instances is very limited, so to be able to fit tuning example in it we convert model to an fp16 there and use SGD optimizer that not requires additional memory for parameters. But this tuning mode is definitely far from stable and can lead to a NaN loss due to limitations of fp16 numeric range. We're planning to make an updated version of colab tuning with more advanced techniques to reduce memory consumption. This could possibly allows us to use mixed-precision training with this memory constrains.
For now you can try to tune model using notebook from this repo on your own GPU, it uses float32 training and AdamW optimizer but requires 32Gb of GPU RAM.

from mgpt.