Can a 7B Parameter Large Model Run on 24GB of Memory?

Author(s): Meng Li

Originally published on Towards AI.

Created by Meng Li

Training large language models always presents a significant challenge with memory.

Weights and optimizer states consume a considerable amount of memory.

To save memory, some techniques have been devised, such as Low-Rank Adaptation (LoRA), which involves adding trainable low-rank matrices to pre-trained weights.

This allows for training fewer parameters and saving on optimizer states.

Freezing the parameters of pre-trained models can also speed up training since only the parameters of the new model are updated, and the rest remain unchanged.

However, these methods, while saving memory, might not achieve the same effectiveness as training with full-rank weights.

They limit the parameter search space and change the way training is conducted, sometimes necessitating a full-rank warm-up phase.

Recently, a new training strategy called Gradient Low-Rank Projection (GaLore) has been proposed.

This technique enables full-parameter learning to be more memory-efficient without compromising on performance.

In terms of optimizer state, it can reduce memory usage by up to 65.5%, while still maintaining strong performance.

It has been tested on LLaMA 1B and 7B architectures and has proven effective.

Now, you might want to try pre-training that 7B parameter model on a GPU with 24GB of memory; it might just work!

https://arxiv.org/pdf/2403.03507.pdf

And it doesn’t require complex operations like model parallelism, checkpointing, or offloading… Read the full blog for free on Medium.

Join thousands of data leaders on the AI newsletter. Join over 80,000 subscribers and keep up to date with the latest developments in AI. From research to projects and ideas. If you are building an AI startup, an AI-related product, or a service, we invite you to consider becoming a sponsor.

Published via Towards AI