> MegaTrain stores parameters and optimizer states in host memory (CPU memory) and treats GPUs as transient compute engines. For each layer, we stream parameters in and compute gradients out, minimizing persistent device state
This is pretty awesome. The only compute I have at home is an RTX 3080 with 10 GB of VRAM, so I struggle with training larger models (>40M, 50M params). I get OOM errors and have to optimize a lot.
I have a lot more CPU RAM in my PC, and this would likely increase the size of models I can train locally.
The claims of the article assumes far more compute and far more VRAM..while the trick enables less back and forth, they don't eliminate it.
I doubt you meant 50M. Rather 50B?
You can only give it a try, but don't get your hopes high on a large context. If their technique works I would guess 8096k context limits would still OOM. 2048 maybe.
I'm extrapolating based on my experiment without this paper's trick to leverage the system memory.
Anything that can run on a AMD395+ w/128GB or whatever the apple equivelent would break things wide open. Training a model on my frameworks of choice or our business info would be awesome.
This isn't really anything new; I've been doing something like this for quite a while, I just haven't bothered writing a paper. (: Probably anyone who would seriously tackle the problem of "how do I train a huge model on a tiny amount of VRAM?" would come up with something similar.
However, most people in the field don't, because the actual practical utility of training huge models on a single GPU is quite low. (e.g they got 341 tok/s for a 14B model on a single 3090 while with my method I was getting ~1k tok/s on a single 4090; that's still very slow)
Also, there are more tricks one can use to speed up training/lower VRAM usage which they're not using. For example, you don't need any gradient offloading (you can just accumulate the gradients directly into the optimizers' states if you modify your optimizer), you can use Muon instead of Adam (which needs only half of VRAM of Adam), you can use quantization (both for parameters and for the optimizer states; e.g. I found Muon quantized into 4-bit working relatively well), etc.
The GPU is no longer the brain, it's the hand. The brain is your RAM. Suddenly that 256GB DDR5 build your wife questioned is 'research infrastructure.'
That’s a pretty non-standard H200 configuration. In the regular HGX configurations, a node with 8xH200 has that much CPU DRAM. That makes the title of the paper somewhat arguable imo.
I’m curious how this technique works, or not, with unified memory architectures such as Apple’s M series. It seems like it’s relying on using overlapping processes to help speed things up, but I would assume that having everything unified in main memory such that you don’t have to transfer everything back and forth to the GPU would also have some advantages. Can someone wiser explain this to me?
This is a fantastic step toward democratizing large model training. Making 100B+ parameter training accessible on a single GPU could open the door to a lot more independent research. Really impressive work!
Having just started to dabble with training LLMs, it seems training a model if you have a training and validation data set is fairly trivial. Creating a good and sufficiently large training and validation data set seems to be the hard part.
Sourcing, cleaning, curating, labeling, generating and quality controlling training data is hard and a lot of work, at least has been for the projects I've dabbled with.
32 comments
[ 3.9 ms ] story [ 57.7 ms ] threadThis is pretty awesome. The only compute I have at home is an RTX 3080 with 10 GB of VRAM, so I struggle with training larger models (>40M, 50M params). I get OOM errors and have to optimize a lot.
I have a lot more CPU RAM in my PC, and this would likely increase the size of models I can train locally.
I doubt you meant 50M. Rather 50B?
You can only give it a try, but don't get your hopes high on a large context. If their technique works I would guess 8096k context limits would still OOM. 2048 maybe.
I'm extrapolating based on my experiment without this paper's trick to leverage the system memory.
However, most people in the field don't, because the actual practical utility of training huge models on a single GPU is quite low. (e.g they got 341 tok/s for a 14B model on a single 3090 while with my method I was getting ~1k tok/s on a single 4090; that's still very slow)
Also, there are more tricks one can use to speed up training/lower VRAM usage which they're not using. For example, you don't need any gradient offloading (you can just accumulate the gradients directly into the optimizers' states if you modify your optimizer), you can use Muon instead of Adam (which needs only half of VRAM of Adam), you can use quantization (both for parameters and for the optimizer states; e.g. I found Muon quantized into 4-bit working relatively well), etc.
While yes it's one GPU.. It's not exactly a slim one.
I'm talking about the training set.
Sure there are some open sets out there.
But my guess is they are nowhere near what OpenAI, Google and Anthropic are actually using.
Happy to be proven wrong.
Sourcing, cleaning, curating, labeling, generating and quality controlling training data is hard and a lot of work, at least has been for the projects I've dabbled with.