Hacker News Board Hacker News, browsed like a classic message board
Live data from Hacker News

MegaTrain: Full Precision Training of 100B+ Parameter LLMs on a Single GPU

arxiv.org

Post reply on HN
Print view View original

21–30 of 71 posts

Re: MegaTrain: Full Precision Training of 100B+ Parameter LLMs on a Single GPU

#21

Post by hirako2000 »

internetguy wrote: post #3

> 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…

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.

Re: MegaTrain: Full Precision Training of 100B+ Parameter LLMs on a Single GPU

#22

Post by kouteiheika »

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.

Re: MegaTrain: Full Precision Training of 100B+ Parameter LLMs on a Single GPU

#26

Post by kouteiheika »

hirako2000 wrote: post #21
internetguy wrote: post #3

> 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…

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 leverag…

> You can only give it a try, but don't get your hopes high on a large context.

You may or may not know this, but: when training off-the-shelf LLMs (i.e. ones which have a huge vocabulary) what consumes a huge amount of memory usage is calculating the cross-entropy loss (which gets worse the more tokens you stuff in your batch), so always use a fused cross-entropy kernel.

For example, for a Gemma 2 model with 2B parameters at a batch size of 8k this consumes 24GB of VRAM by default (!); you can fuse your cross-entropy loss with @torch.compile and that can cut down this memory usage to something like a few gigabytes, but with a dedicated kernel this becomes a few megabytes.

Re: MegaTrain: Full Precision Training of 100B+ Parameter LLMs on a Single GPU

#27

Post by vlovich123 »

kouteiheika wrote: post #22

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…

341 is two orders of magnitude faster than your 1 tok/s so it doesn’t seem like their stuff is all that obvious. I also have no baseline for training to know if 341tok/s is slow but it seems speedy for a 3090.

Re: MegaTrain: Full Precision Training of 100B+ Parameter LLMs on a Single GPU

#28

Post by bastawhiz »

vlovich123 wrote: post #27
kouteiheika wrote: post #22

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…

341 is two orders of magnitude faster than your 1 tok/s so it doesn’t seem like their stuff is all that obvious. I also have no baseline for training to know if 341tok/s is slow but it seems speedy for a 3090.

OP said 1k, not 1

Re: MegaTrain: Full Precision Training of 100B+ Parameter LLMs on a Single GPU

#29

Post by rolandr »

vlovich123 wrote: post #27
kouteiheika wrote: post #22

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…

341 is two orders of magnitude faster than your 1 tok/s so it doesn’t seem like their stuff is all that obvious. I also have no baseline for training to know if 341tok/s is slow but it seems speedy for a 3090.

1k tok/s = 1000 tok/s...

Re: MegaTrain: Full Precision Training of 100B+ Parameter LLMs on a Single GPU

#30

Post by spacebacon »

weitendorf wrote: post #16
internetguy wrote: post #3

> 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…

To make the most of these architectures I think the key is essentially moving more of the knowledge/capabilities out of the "weights" and into the complimentary parts of the system in a way that's proportionate to the capabilities of the hardware. In the past couple months there's been a kind of explosion in small-models that are occupying a niche in this kind of AI-transcoding space. What I'm hoping we're right on t…

You are on the right track. Check out the Semiotic-Reflexive Transformer (SRT) here.

https://open.substack.com/pub/sublius/p/the-semiotic-reflexi...

Post reply on HN