Lightning Talk: Trinity Large - Torchtitan on 2000+ B300s - Matej Sirovatka, Prime Intellect

This talk focuses on the journey of pre-training a 400 billion parameter model from scratch, utilizing 2,000 Nvidia B300 GPUs and the Torch Titan framework. The speaker details how the model was trained on 17 trillion tokens, balancing both synthetic and real data while maintaining a low sparsity ratio, which enhances its efficiency for inference. He discusses challenges faced during the training process, such as load balancing issues and GPU stability, as well as the solutions implemented, including their own load balancing algorithm and fault tolerance strategies. Despite some hardware constraints and compatibility issues with different Torch versions, the resulting model achieved competitive performance metrics, ranking among the top models for open-core usage shortly after training.

Hello, everyone. Uh my name is Matej. I work at Prime Intellect, where I do everything related to performance, basically from training to inference. And today I'm going to talk to you about how we pre-trained a 400B uh model from scratch on up to 2,000 B300 GPUs only using Torch Titan. Uh so, up here you can see like the result of it, which is a loss curve of

our model. It's trained on 17 trillion tokens from scratch. Uh half of them were synthetic, half of them not. And it's a 400B uh MoE with only 13 billion active parameters, which, as you might know, is quite friendly for inference because only the 13 billion are getting activated. Uh this is like a very small sparsity ratio, where only four out of 256 experts get activated, which is

a which comes with a bunch of issues. We trained this model on uh in 1 month on 2K B300 GPUs, which is like the newest one of the newest pieces of hardware from Nvidia. We did this in December of the last year. And our aim was to make this model very friendly for local inference and also for post-training. As you know, uh half of post-training is doing

inference, so we wanted to make sure that this model is very friendly to infra. Uh we did a bunch of decisions to uh allow for this. One of them being like the low sparsity ratio, where only four experts get activated. Another one is for long context, where only doing local attention in 3/4 of the layers. we did this for 20 trillion tokens, which we didn't finish. We

did only 17 trillion, as you can see. And we did this only Torch with Torch Titan only, which, if you're not familiar, is a pre-training framework that is very easily extensible and really easily hackable from Torch. Uh Uh, you can see the loss curve on the right where we There are no wall spikes, which is, I think, very cool. Uh, as I was saying, it's an MoE

architecture with only four experts activated. This makes it very friendly to inference. Also, long context capability is uh, is fast because of the local attention. Uh, the architecture you can see there. And like the result is that our model is way faster for inference than uh, lot of the other frontier models such as Deep Sea V3 or GLM 4.7 at that time. Uh, you can see the

output throughput is like way higher and uh, latency way lower across like the whole spectrum. So, what problems does this entail? Like, what problems can we have training that large of a model? Uh, one of them being we have a lot of B300 which you might say is not a problem, but it actually is. The MoE is very sparse and very thin. The hidden dimension of the

math walls is only like 3K, and the GPUs are very hard to get saturated. This also This allows you from doing parallelism like tensor parallelism, which just shards it up a bit more. um, another issue is, of course, the sparsity. Like, load balancing these experts comes with a lot of issues that arise and can actually break your uh, in a lot of unexpected ways. We've actually done

our own load balancing algorithm for this. And yeah, for the first 7 days, our model was basically we were just like trying to figure out how to make this stable. Another issue is like how to parallelize this model. As I said, the hidden dimension is very small. So, this leaves you without like the classic solutions as tensor parallelism. And also, a pipeline parallelism is a bit too

difficult to implement in PyTorch properly for it to be fast. So, yeah, the parallelization is we ended up with like the simplest possible solution. Too simple to be true, I guess. And that is just doing HSDP, so uh FSDP I want to play 120 GPUs with I guess like 16 replicas if my math is correct. And EP inside of these only EP8, which is like the best

solution that PyTorch offers right now for our use case. And it's very simple to do this. Scaling it up to more GPUs could and will give us would give us way more problems to do that. another issue is we used B300s. We used it in the December of last year, I guess, or October December November December. And B300s were new at the time. So, setting up the

environment to actually work was not so easy. There was a lot of It was between torch 2.9 and two torch 2.10 release as well. And there were a lot of weird quirks uh like try to not being compatible on our environment. Uh there was also a very weird uh that you can see it there like I we upgraded torch and suddenly we lost 10% of MFU because

it uh the matmul dispatched to a different matmul implementation. And suddenly the whole training was slower. So, we had to like pick up all of pieces and do a lot of like compatibility checks between different uh packages to ensure the highest throughput and not to break anything, I guess. Also, there were a lot of correctness issues between versions, which I can talk to you about a bit

more later. I have a pretty big PTSD about that. And yeah, next one. It's B300s, right? It's a cool piece of hardware, but it was new hardware as well in the data center. So, they were not so fault tolerant per se and they were crashing fairly often. if I remember correctly, we've had one GPU crash like every 4 hours at the beginning. Then like the the faulty

ones got replaced, so it went okay after that. But how to solve this? Like the obvious solution would be do like fault tolerance. For example, Torch FT, if you're familiar with, has offers this out of the box for HSDP solution, which we did use. Well, not so simple. Uh Torch uh FT replaces an AFSDP. There's an all scatter reduce scatter in back backward. Torch FT replaces this

with a custom one uh didn't work on our hardware. So, it only fall back to like CPU implementation, which was insanely slow. So, we couldn't do fault tolerance as is, right? So, how to solve this after? the next obvious solution we came up with is have one follower node that we just always have someone on call. That someone on call just picks up the sees like a

notification. We used a platform called Better Stack. when the platform when like we got this notification, the person actually just replaced the node and restarted the training and pinged the data center to replace the node. This worked fairly well, even though we had like I think like the whole training was actually like this is take it with a grain of salt. I think it was like 300

uh restarts all together. So, yeah, now we have all the solutions. We have parallelism, performance environments, everything. We can just train, right? Well, not so simple. Then we actually like people say pre-training is a solved uh area. Well, it's not. You actually now have to solve like the issues that come with pre-training, which is like expert collapse, grad norm spikes, a lot of weird quirks. And yeah,

this was like the second most difficult part of the whole pre-training. We had to come up with our experts. Because the one that used was used by Deep Seek was just crashing for us. We ended up with the solution after staring to on one DB graphs for like 3 hours, 4 hours at night. And yeah, it just worked. Uh so, you can see the loss curve like

this one is insanely smooth now. There are no loss spikes anywhere. And it's a very fairly pretty loss curve I I would say. And yeah, what is the conclusion? Like from this you would expect that I was just complaining about Torch having a lot of issues for 10 minutes, I guess. Uh but that's not the conclusion. We just took a framework that we haven't written. The framework

that we were not authors of. We just took this framework on 2K of new GPUs and we trained a model. This model is actually fairly good. And it was a surprisingly stable experience. Except of the issues we had, right? Uh and yeah, like the results are uh speak for itself. You can see that our model actually became a uh top four model on open core usage on

open router in like first 50 days. People used I think like 3 trillion tokens for it. And it became a somewhat of a sort of open model. You can see it here that some of the benchmarks were competing together with like models that are way larger than our our model like Opus. Uh also GLM 5 Kimmy K 1.2. And yeah, that's it. Uh thank you very much.

And if you have any question, I guess we have like 1 minute, so you can ask. Still. Yeah. Yes. Uh back back then would did decide for different GPU than the V100 if you would do it again? It wasn't my decision to make to do this. And even back then, yes, I would decide for a different GPU. It was a big complaint of mine during the training.

Like, why are we doing V300s? Yes. Yes. Did you think about using Megatron or Megatron? Uh I I'm not going to I don't really like Megatron, to be fair. I don't like setting it up. I assume that it's faster in a lot of use cases, but from my older experiments, the FSDP, which is what we ended up using, is way faster in Torch than it is in

Megatron. Where Megatron shines is like the more like 4D 5D parallelisms, which we don't have to do on this much hardware. I guess that's it. Yep. Thank you very much.