Training Zamba: A Hybrid Model Master Class with Zyphra's Quentin Anthony

The future of AGI will involve a combination of cloud and on-device deployment.

These large monolithic model companies doing like Anthropic or OpenAI just can't really specialize to every single person on the planet. We think that you need to have your own set of weights and changing a system prompt per person is not enough, right? We want to actually bake into the weights. You can make the model simulate learning faster than it really is by doing activation steering. Like if the user tells the model, you're being too dry. Then you can very quickly steer the activation to be a bit more fun. Until tonight, when you can bake into the model, I think it's got to be continual learning, and it's got to be per user. And the only way to do that is with weights on the phone.

Hello and welcome to The Cognitive Revolution, where we interview visionary researchers, entrepreneurs, and builders working on the frontier of artificial intelligence. Each week, we'll explore their revolutionary ideas, and together we'll build a picture of how AI technology will transform work, life, and society in the coming years. I'm Nathan Labenz, joined by my co-host, Eric Thornberg. Hello and welcome back to The Cognitive Revolution. Today we're once again going down the state-space model rabbit hole with returning co-host Jason Moh, who regular listeners will remember from our MambaPalooza literature review and Albert Gu interview episodes, and Quentin Anthony, head of model training at Xiphra, a large language model startup that's just released their Zamba-2 7B model. Which is built on a hybrid architecture that uses both the selective state-space mechanism and the traditional attention mechanism, albeit with some notable tweaks relative to the standard implementation. In addition to sharing Zephyr's high-level vision for highly personalized on-device AI, Quentin was super generous with both his time and knowledge, sharing a wealth of practical lessons learned from the front lines of model training. Over the next 2 hours, we will cover the delicate architectural choices that balance efficiency and capability, The many practical challenges of training at scale, including choosing the right learning schedules for different phases of training. The nitty-gritty details of training hybrid architectures, including why Zamba models don't need positional embeddings. The Zamba model's use of shared attention blocks and internal LoRA adapters to maximize performance on the edge. The not-so-simple relationship between loss metrics and model quality and capabilities, as well as the challenges of context-length extension. The Cypher team's experiments with different optimizers and why they're sticking with Adam for now. Quentin's intuitions about the relationship between model scale and loss landscapes. And finally, even their recent published work on tree attention, which offers important advantages over ring attention for multi-node training. I have to say, I got a lot from this episode. And while it's technical enough that I wouldn't necessarily call it entertainment, I am confident that you will too. If so, we always appreciate it when folks take a moment to share the show with friends or write an online review on Apple Podcasts or Spotify. And we welcome your feedback via our website, cognitiverevolution.ai, or by DMing me on your favorite social network. Now, I hope you enjoy this highly technical conversation with co-host Jason Moe and guest Quentin Anthony, model training lead at Zyphra. Jason Moe, returning guest and co-host and chronicler of state space models at statespace.info. And Quentin Anthony, model training lead at Zyphra, which has just released the new Zamba 7B SSM hybrid model. Welcome both of you to The Cognitive Revolution.

Thanks a ton. Great to be here.

Jason, regular listeners will know as a partner in crime who's also obsessed with state-space models and the potential that they have to unlock new capabilities in terms of potentially long-term memory, extreme efficiency, all these kind of interesting things. And we've been down the rabbit hole together on that a couple of times. So I was excited to have him back to help me with this conversation about the new Zomba model and all the ins and outs of that. So, Jason, I'm going to ask you to lead the questioning today, and I'll be in the supporting role, but listening intently and probably jumping in with a few of my own follow-up questions along the way as well. How's that sound?

Sounds good. Let's do it. So let's start with maybe a broad question first. So on the Zephyr website, it says that, quote, the future of AGI will involve a combination of cloud and on-device deployment with an increasing shift towards local inference. So why is local inference going to be important, and how does this idea overall influence the direction of the company?

Yeah, there's a few aspects, a few angles of this. The main one is personalizability. We think that doing system prompt hacking and things like these large monolithic model companies are doing, like Anthropic or OpenAI, just can't really specialize to every single person on the planet. We think that you need to have your own set of weights, and changing a system prompt per person is not enough, right? We want to actually bake into the weights You like to be talked to this way. Your favorite restaurant is X, that sort of thing. A second thing here is privacy. So there's a lot of data about your laptops, your phones that you don't really want to be communicating to the cloud. There's a lot of enterprises out there who don't want to share all of their data with OpenAI, all of their proprietary code and things. So keeping that all on device, on-prem, in the organization is where we want to keep it. And then there's just some very practical challenges here of OpenAI spends— how, what was their CapEx realistically, right, in order to infer all these models? That's not really something that can keep up forever. It's being fueled by a lot of VC money. But if everyone's device is able to efficiently run their own models, then we offload that to the users. And then it's also just, it's much faster. You can run it offline. All these small benefits come out from having the model locally in the device, but it's all, it's steered by this global personalizability, privacy, these sort of north stars. In terms of— Yeah, go ahead. The cloud, uh, there are some model capabilities that are more challenging. Like, they're too challenging to have on a 1B or a 2B or something like that. Most people we see actually, when they communicate with models, they just need these higher-level, more simplistic tasks, just like basic chat, like factual recall. They're not asking to make these O-type Math questions, these Olympiad questions or anything to like their day-to-day on-device model. But sometimes you will need it. And that's why we have also a cloud offering, like in the 7B range. The 7B range is a turning point where you can fit it on like more powerful laptops, powerful on-prem, but it's able to think through these more difficult problems more clearly. And we want those on-device models to be able to know like when they need to go back and ask like an oracle model, something a bit more powerful. And then this is all seamless to the user. And then you can Pick and choose how often your model wants to talk to the cloud.

It sounds like it's— when I hear local inference, it's often a privacy bullet point is the main bullet point that gets hit. But it sounds like the overall idea is it's some mix of capability and privacy, that if we had super efficient, super personalized localized models, that this is going to increase capability of the full stack. Is that correct?

Definitely.

Yeah.

You can't forget personalizability. It's really hard to bake in personalized outputs, personalized how the model likes to talk to you in the same voice. It's just a hard problem.

I would love to dig into that a little bit further.

Yeah.

Because I'm working on a project right now to try to get a model to write as me. It's been a white whale project for me for a while. I figure if a model can write as me in a way that's compelling to me on a consistent basis, it seems like that's a pretty big threshold. So I've been trying to see, can I make that happen? And the best readily available fine-tunable model today is GPT-4o. And so I've been working mostly with the OpenAI stack. They make it easy to do fine-tuning. They don't tell you too much about exactly how it works. One presumes that it's a LoRA type process that they're running. I think there's a lot of reasons to believe it's something like that. It doesn't seem to learn facts very well. It does seem to learn patterns and stylistic layer, if you will, reasonably well.

Okay.

I wonder what you think is the right way to even approach personalisation in the first place. Is it like when you envision doing this personalisation for individuals on their machines, are we talking LoRAs? Are we talking full training of a 1 or 2 or 3B model? How many tokens does that take? Do we still need to stuff huge amounts of context in there? Because that's another approach I've been taking where I'm like, maybe if I give it 50,000 tokens of everything it needs to know about me as the system prompt, then it won't need to hallucinate. And I still haven't quite figured that out. I'm still getting hallucinations. So what is the overall paradigm or path to this deep personalization for individuals?

Yeah, so the exact process will not be entirely clear until we have a bunch of people having Zamba on their phones. But we have a few bets. One is that we have a lot of expertise on continual learning and continual pre-training. So we envision doing a weight update overnight, maybe while your phone is plugged in, your laptop's not really actively being used on consumer hardware. It's really cheap to do this when the model is really small. Updating LoRAs, as you mentioned, is another, like, that makes it even cheaper. But just continual learning, I think, is the way to go. In terms of your GPT-4o problem, I think I would need to know more details about the fine-tuning dataset and the model, which I'm just never going to get. We find continual learning on a small model per person is really great. There's some sort of, you can make the model simulate learning faster than it really is by doing things like activation steering. Like, if you remember Golden Gate Bridge Claude and stuff, if the user tells the model, you're being too dry. Or something like that, then you can very quickly steer the activation to be a bit more fun until tonight when you can bake into the model, okay, this person likes generally to be a bit less dry and to talk in a more fun way. I think it's got to be continual learning and it's got to be per user. And the only way to do that is with weights on the phone.

But how much of that exists today? Can I get a Zamba? First of all, which scale of model would you recommend that I Bring locally and can I do this continual learning? Can I like start to feed it my emails and my Slack messages? I assume I would have to mix in some like general purpose pre-training mix as well. I wouldn't just go all in on like my content. So tell me a little bit more about what that looks like practically. I'm taking notes.

Yeah, absolutely. So right now we have the 1.2B, the 2.7B, and the 7B that we just dropped. All of the— our Zamba2 series of models. The 1.2B is really good if you're like really edge constrained, like you don't want your phone to use much battery, or you want it on a phone in the first place. In terms of whether you can do this yourself, right now you can. It's just a little less user-friendly than it will be someday. But right now you— we have our own Hugging Face fork for Zamba2. Uh, so you would just download the model weights and then you would just have to tokenize and actually just train On whatever personal emails or whatever else you want baked into the model. The product vision here though for Maya is that you have like a Maya cloud across all of your devices. And that's— there's some like social aspects here. You can upload pictures, it'll be multimodal. Your conversations and such, you can decide how much you want to share, how much you want Maya to be personalized to you. But training on that cloud of your own data, just as a continual pre-train. On whatever edge hardware you have is the high-level north star here.

Are you doing this today? Like, how far have you guys taken this internally? Do you have a Quentin model that drafts your emails, or what's the frontier in application?

So earlier I said the 1.2B is good for edge constraint. The 2.7B is really good for quality. I'm always surprised by it. And we do have some fine-tunes for that. So we have some roleplay models, some summarization models, if you want to summarize emails, summarize meetings and such. We have some very early results on audio as well. So generating audio in the style of people you talk to, note-taking audio from the meetings you're attending. It's definitely more mature in-house, but we really want this to be really seamless before we give it to everybody else. So I would say that the 2.7B, I've seen really amazing results on in terms of summarization and like these tasks that I was mentioning. And it's all from this Hugging Face, just fine-tuning our DPO. on whatever my personal data is. Cool.

I'll probably have some more questions about that as we go, but maybe a good transition to, as we think about this on-edge paradigm, Jason, you want to take the baton and talk about why state-space models are the, or hybrids at least, are the attractive approach to take there?

Yeah, absolutely. I guess when you think of what kind of models can run locally, You get hit by transformers in 2 places pretty hard. The first place is just memory. Anyone who's done their own work with a limited compute budget knows the classic out-of-memory error.

So you quickly—

whether that's too large of a model in terms of parameter count, or if your KV cache is growing, and at some point you run out of memory. So transformers... Because of the way they work, they're famously pretty memory intensive. The KV cache scales as you scale your sequence. And then of course, secondly, is just computation. Transformer has amazing properties, but it is pretty computationally intensive, especially when compared to other architectures. So when you start looking at Mamba, that's one of the exciting things about it. It has properties that are much more memory efficient, and computationally, it's able to be more efficient as it processes through sequences due to its fixed state size. And so you even see that play out in the literature. There's been some early work with Mamba models with a couple papers on robotics. These are models that would have to run locally. You could not do cloud API calls for split-second decisions that would be made in robotics. So you can just imagine, finally, maybe we have an architecture that helps inform what the ultimate models running locally could look like for these kinds of applications. So how does that vision look like for Zyphra and the Zamba architecture? What does the ideal local inference strategy look like as it relates to Mamba?

Yeah, so we love Mamba for its systems properties. As a guiding principle in terms of model training, we try to train the models that are perfect for inference. Because that's where we're going to need them to perform the most. We can eat some training costs, but at the end of the day, it needs to be blazing fast for inference. This translates to a few different things for us. For one, attention for like dense transformers, like a Pi-type model, is just not going to happen for us. We don't think that this is efficient on a phone. Earlier I mentioned like we want text and email summarization. Sometimes you want news article summarization. Those inputs are too long for the KV cache. Like it's just going to grow outside of memory for most phones today. Attention though is really nice. So you get these exact cross-sequence dependencies that we think are just required for these, some specific tasks. Remember I say specific, so not every task needs attention, but there are some tasks like in-context learning and long sequence dependencies for which attention is just necessary for you to be performant on. So other, eventually if you train on several trillion tokens, then eventually Mamba becomes somewhat performant. at this on benchmarks, but I've never played with a pure SSM or pure RNN model that's able to really speak well in an in-context learning regime. I think that's a common intuition that people have when they actually sit down and play with the models. So pure SSMs have quality issues from what I've seen, and then pure attention dense transformers have performance issues. So you can't get Phi on a phone realistically. to run the amount of time with the fastest time to first token that the users actually want. So there's some other architectural changes here. Instead of having independent attention blocks, we have this global shared attention block. And the reason for this is that us and a few other groups have independently found that attention blocks are highly correlated across depth. So they're telling the MLPs in a dense transformer, these are the important tokens, these are still the important tokens, these are still the important tokens. There's some small changes. But most of the specialization across depth is coming from the MLPs. So if you tie all of those weights together, then you're able to allocate more parameters towards the, for example, the Mamba block for us. If you were on a Dense Transformer, you could apply more to the MLP blocks or something. So now we just have, like, for the Zomba-1 series of models, you just have a single attention block and you apply it every 6 Mamba blocks across depth of the network. We improved on this a bit in Zamba 2 by seeing, okay, a little bit of specialization across depth is helpful. So we added LoRA back in, and those LoRAs are independent across depth. So now you have this single global attention block in the 1.2B case, and then the 2.7B and 7B case for Zamba 2, you have inter— uh, interleaved. So you have attention block 1, attention block 2, attention block 1, attention block 2. All across depth, uh, and then individual LoRs across depth as well. So you get a little bit of specialization. But back to the systems angle, you're not really applying very much attention. So your slope of like your KV cache memory curve— you'll see this in all of our plots— is something more like 10 or 13 invocations rather than 30, 32 in like a dense Transformer, because we're just invoking far fewer attention blocks. Now, the very first criticism people should be thinking is that, okay, you're applying for the same parameter budget, you're applying parameters over and over, the total number of flops is going to increase. And this is true, we're applying blocks more often than like a typical 7B, but the Mamba blocks and now the Mamba 2 blocks are so high throughput on parallel hardware in general that this totally makes up for it. And we end up with a net improvement of about 20 to 30%. In time to first token and time per output token. So they make up for our sins of sharing these attention blocks. But yeah, this, both of these translates to significantly lower memory overhead, significantly lower time to first token at the prefill stage. So the model responds really fast initially, and then time per output token is also 20-30% beat compared to similar dense transformer models.

Hey, we'll continue our interview in a moment after a word from our sponsors.

The choice for using global shared attention blocks is really interesting to build the hybrid model that way. I guess it's in alignment. I saw on Zypher's site the desire to maximize performance per parameter. That seems like a choice that obviously is doing that as well as maybe some other architectural choices. How important is it in the Zamba family of models to focus on performance per parameter?

It's pretty important, right? At the end of the day, most of your memory overhead is coming from just storing the parameters. We can get away with a lot of quantization, but each parameter is really painful when you're deploying it on a phone. So trying to make all of those parameters maximally useful for the attention block, as in tying it all together and making that attention block really well trained.

Yeah.

Maximizing the performance of each one of those parameters, we've got to do. We don't want this really sparse model where most of what we're storing is not doing much. But yeah, very important.

I think there's like a broader kind of trend here of like dense attention being the original thing. And then there's been obviously many attempts to make that more efficient with like linear attention or things that are like not quadratic. My sense is that those have all been pretty good, but somehow still fall short. Then you have these other techniques of shared Query attention. Maybe you can help us understand in a technical sense exactly what that is. And then I'm also aware of shared query key attention, and now you're going all the way up to shared full attention blocks. So help us map that space and just understand what are the different options and trade-offs. And is this the end of that paradigm, or is there even still further that we could go on the spectrum from full dense to— Full shared, is there even further on that continuum?

So we use grouped query attention for our upcoming attention blocks. I would say I haven't seen like a non-quadratic attention that really performs the same and does what we want it to. Frankly, I would say that there is a very interesting emerging area of work where people try to remove the attention from attention models so that maybe the theory there is This is more of a play from, for example, a Liquid AI sort of player. You take a model that's already been trained with attention, and then you try to linearize that attention at a final post-training step. And the hope there is that the model's already baked in the ability to do these long sequence dependencies in context learning, and then you can just sort of copy it over to much lower cost blocks. It's a potential play.

Can you give any— help us develop an intuition for— this is like really a question I want to I'm asking for me, and hopefully other people find it useful.

Yeah.

I have an intuition for why the full attention would be qualitatively better versus, like, a state space mechanism or some approximation of attention. For me, that kind of boils down to we compute a relationship between every token and every other token. We don't forget about any tokens. It's all there. And it's all when a new token comes in, We can see how that relates to everything that came before with no compromises. When we have that, we try to map that onto a linear thing. What's the intuition for why that would work or wouldn't work? Maybe there is no intuition and people are just trying stuff and seeing what works, but I'm grasping for understanding of what that transformation is and why we would— what should we expect and what should surprise us?

Okay. My first order intuition here is that attention is really good because the point of the attention block is to note-take as much as possible about the incoming sequence. So you have this really long sequence, let's say, and if you're taking the dependency from every token to every other token, you're giving MLPs or whatever it is that you're projecting as much information as possible in the sequence. And you have multiple heads, each of them saying, okay, they're each looking at the sentence in a slightly different way. And then if you have large heads, they're keeping a lot of detail about their own specific angle of the sequence. Any sort of attempt to linearize this is saying, okay, how much of that detail do we actually need? So the original attention paper was so useful because a lot of that— we found that MLPs could make use of a lot of that detail because MLPs are just like trying to find the deeper meaning on whatever those notes that the attention took. So Attention takes notes about this, as many notes as possible, but very computationally heavy. And then MLPs try to find deeper meaning from those relationships. They sort of mix the heads, they do cross-heads. So for linear attention, it's all just trying to say, okay, are all those heads doing something important? Are we able to cut away a lot of them? Do you really need every relationship to every other token? And this is like a very theoretical question, right? So It's not super clear how many nodes, whatever you need. You also have different ways of projecting across these heads. So MLPs are a very powerful way to mix across heads and make use of these nodes that you're being— that are being taken. But like a Mamba block has much smaller projectors. It's a very different story there of like how much they can really do from the nodes taken by like the selective scan versus an attention block. Maybe if you had like bigger MLPs for a linearized attention, maybe that would make up for it. Maybe you still need like more notes if you had a Mamba block without like powerful projectors and then an MLP. This is really like something I can't really say. If I did, then I would make the model. But I think all of these people are motivated by the fact of how much can we throw away from these notes and still get performance because Do you really need n squared dependencies? This makes sense. Yeah.

And so coming to the shared attention, I guess I'm not 100% sure what exactly all is shared, but if I understand correctly, it's like you are literally using the same attention weights over and over again at every layer, but you are recomputing the effect of that attention At each or every 6 layers because the computation has proceeded. And so each time the attention block is used, you're able to use the same parameters. That saves space, but it doesn't save runtime computation per se, because you're still doing that attention on the input as it's been processed to that stage through the model so far.

Is that right? Exactly. Yes. Okay, cool.

And that's different from— so when you reuse an entire attention block, that's different from just reusing the queries or just reusing the queries and keys. You're reusing queries, keys, and value.

Yeah. So yeah, we have taken it to the extreme of saying these attention blocks are pretty much just reminding MLPs anyway. So let's just stop storing all of these independent parameters that eventually learn to do the same thing anyway. So let's just get them— when you're reapplying, you also get really accurate gradients, right, the training time, because it's applied 6 times instead of, you know, 1 if they're all independent. So we get this really highly trained, highly saturated attention block, and it's completely tied to query key value.

Mm-hmm.

And then we add a little bit more specialization back in for Zamba 2 by having alternating layers. So now you have 2 attention blocks, they're each tied, but then we apply them one after another. And then we just have some depth-wise LoRAs that are themselves independent. So you get a little bit of specialization across depth, but not too much.

Yeah. Interesting. The fact that it's alternating is somewhat counterintuitive when you say there's specialization across depth. I would have thought you'd like maybe segment it into halves and have a first half and a second half or a first, second, third blocks, but you're still staggering them. And then I guess with the LoRA adapter, you have some kind of depth-specific thing that is fully localized to that place in the model.

Yep.

Cool.

Fascinating.

All right.

I'm talking more than I intended to. Jason, let's get back to your main line of questioning on the technical details.

Yeah, I guess if I could just get, set the history straight real quick, we can go back to February of this year. You actually released, cuz this is how Nathan and I actually became familiar with Zyphra was Not the release of Zamba, but the release of Black Mamba, a mixture of experts model that used Mamba. This was really early in the literature. I think Nathan and I were counting papers at that point. This may have been among the first ten Mamba papers published.

We've since lost count.

Yeah, well over three hundred now. I really don't know the the full count. So what was it about Mamba in particular that got Zyphra making this early bet and? Going forward with Black Mamba?

Yeah, so you're taking me back. When I first joined Zyphra, we were thinking we were all about MOEs. We're like a mixture of experts with very vanilla attention blocks, no exotic shared blocks or anything like that. It was the way forward for on-device. We had some ideas for like reducing the inference time, memory, and compute for these, and we were all in on training. We trained some good models actually at the time for MOEs. Then Mamba came out and we were like, okay, wow, this is amazing for on-device. But the question is, what does quality look like? Because we already knew that it's like attention-free models had this sort of problem of quality. Our motivating figure was in the original Mamba paper. There's a figure where they try pure Mamba versus Mamba plus attention versus Mamba plus MLP, and they have the loss per flop. graph for all of them. And they don't show the absolute loss. So you don't know what the actual quality is at the very end. So the MLP plus Mamba hybrid, they saw was really expensive. It had a lot of flops for equivalent loss. And we thought, we don't know what the actual quality is here. So maybe the quality is actually really good. And that you in fact don't need all of these notes from attention, all these really rich cross-sequence dependencies, which ended up being the wrong, the exact wrong thing. And we replaced these really expensive attention blocks in our current MOEs because they're the ones that's really giving us so much trouble. And then just do Mamba and then a mixture of experts MLP block and just do that over and over. That would be really inference efficient for what we were preparing for. So we trained it. So we set everything up, we trained the model. We did indeed see that in terms of long sequences, it was amazing for overhead, but Mamba on its own can't really make up Without these, the attention, cross-sequence dependencies, in-context learning, like those exact same things we learned. So the model was not great on MMLEU, for example, which is like your early indicator, not that you want to saturate MMLEU, but that's an alarm bell of, uh-oh, your model is not really able to learn in context. So that was the early prep. That was our early lesson. That's why we didn't really proceed with doing Mamba without attention anymore. MOEs, maybe they'll come back in the future. But yeah, that was the initial model. Any questions or thoughts on it? It was an experiment that we shared with everyone else rather than something that we wanted to actually deploy.

Hey, we'll continue our interview in a moment after a word from our sponsors.

It's a huge change in strategy to pursue, to go from an MOE to something like XOM. I mean, something that's MOE models, really hard to run. The parameter count's huge. If you scale that up in terms of just memory, you've got to store a lot more weights. And then you go to Zamba-1, way more parameter efficient. So I guess just the thinking of what drove that, drove you to go the other way. I understand the attention part with in-context learning, that makes sense. But I guess like the local inference part sort of as well, a little bit.

Oh, definitely. Like in terms, there's a reason all of our Black Mamba plots are like in terms of efficiency. It is such an efficient model because there's no attention anywhere. There's no quadratic complexities. There's no KV cache at all. Like, it's just the constant state that Mamba has to require. So your memory overhead is on the floor. Oh, I also forgot to mention that motivating figure. So yeah, the reason we were doing MOEs was that we were trying— we knew that like the, the FLOPS per loss of like MOEs is much more amenable because you're only routing to a few of these experts at inference time. So we were trying to hope that we would like dip that line. of loss per flop quite a bit lower than both the attention plus Mamba and the MLP plus Mamba. And in terms of pure loss, oh yeah, we did. We definitely did that. But in terms of model quality, it just wasn't where we wanted it to be.

Can you talk more about that divergence between pure loss and model quality? I think that is one of the biggest picture questions that the entire field is struggling with, right? Like from a sort of super big picture, like, when does AGI arrive? And also people worry about this in the context of AI safety with emergence. We believe, or we feel pretty confident we can predict a relatively narrow range of what the loss value is going to be as we extrapolate scaling laws out. But then everybody's, but what does that turn into? And so you're highlighting here one moment where you have, or one example where you have seen this divergence between—

Yeah.

loss and quality in a negative way where the quality is less than you might have hoped based on the loss. Presumably that could go the other way too, right? It doesn't seem like these are— I think more guidelines than actual laws is probably the right way to think about the scaling laws. So yeah, tell us basically, I'd be interested in everything about how these divergences between loss and quality open up.

Yeah. So Loss is not a very useful indicator, like at all. So early, I use it to detect spikes, obviously. So it tells you, okay, your model died. And we definitely saw those, especially with Black Mamba. It's a very unstable model to train, actually. In terms of model quality though, it's not a great, like generally you want it to go down. Do you want it to go down to 3 or 2? It's dataset dependent. The noise of your data strongly determines this. I would say one thing I check it for is that it's a very early indicator of whether your model is saturating. Or doesn't have a high enough learning rate, whether it's progressing. So your loss decrease is obviously going to slow way down as you hit a stable point of training. You've reached the slope of the optimizer landscape and you're going down. So that's great. But if your learning rate, for example, was not aggressive enough, then you might just completely plateau in your loss. And this is kind of like danger, like you've done something wrong. So that's really all it can tell you. Everything else is coming from evaluations. Now, evaluations are also tough because if you maximize evaluations at the cost of everything else, which if I wanted to do that, I would just do synthetic textbooks for my entire training dataset and I would blow up evals. Great. But when you actually talk to the model, it's super dry. You can't really fine-tune it very well because it's assumed this exact, very clean distribution of data. It doesn't really respond to unclean data very well. So what you need is to understand what each eval is actually doing and providing for you. For example, I believe it's HELASwag is one of the evals. And that one is like general language capability we're finding. Like, does your model just model very simple sequences? And that's like a very smooth eval. It's like a better loss almost, where it just like linearly increases across training. Then you have weirder evals like MMLEU that are important indicators. For one, they tell you how contaminated your data is. Yeah. MLU explodes, you've probably done something wrong. You've probably downloaded some data that includes some MLU questions. And then obviously you also, like we do very careful test overlap with our training sets as well to also indicate this. But MLU is a really good indicator of whether your model has reached this emergence of in-context learning. And it really is an emergence. So it doesn't show signal for a very long time. And then near the end it starts to grow up because your model is starting to grok this, this in-context learning task, which is difficult for it. Finally, and then there's a bunch of us, like we test like factual understanding with the Arc Easy and Arc Challenging or Arc Challenge evals. So those really show us how well the model like is understanding factual recall, how much Wikipedia we put in there, how many textbooks. And then lastly, you really just need vibes, man. We have just like a set of prompts that show us really key, like, is the model fun? Like things like, is the model able to roleplay like as a pirate? Is legitimately something that's really important to see because if you give a model like Phi, you tell it to act like a pirate, it'll ignore you. It'll stick with it a little bit and then just immediately drop it. And that's important to test. We also test things like not quite strawberry, but how many times a flower blooms, if it blooms this many times in a year. There's a lot of vibes questions that we ask the model across training. And at the end of the day, you want your model to perform well on things you actually want it to do, not just benchmarks. And those are the questions that we're asking it, like personal assistant questions and this sort of thing. And that's what actually indicates whether the model is becoming quality. Anything there you want me to drill down on? That's kind of our process.

That's fascinating. Moments of, I think more anecdotes, if any, on like moments of emergence are interesting. It sounded like you were saying. MMLU kind of, you understand to be a reflection of in-context learning? Is that because you're like doing few-shot testing? Like, I'm a little bit unclear on what basis you would equate MMLU with in-context learning. And also just interested in what the 2 curves look like as you go through training. I'm taking away, I think, that like the loss curve is just doing its smooth decay thing. But you're seeing these relatively sudden spikes in performance on key tasks. Is that right?

This is exactly true. All of the vibes questions are also like this, where it's like, oh, the model suddenly gets it early on. And even halfway through, the model keeps getting tripped up on silly things. But when the model gets it, we're like, okay, we've reached this key milestone. So yeah, loss is always, you just zoom in and you're like, okay, is there a slope? Is it negative? Yes, move on. The reason MLOU is hard and the reason it tests in-context learning is that its format requires a very strict format. That's really it. And also very weakly tests factual recall because these are all college-based questions, but it's answer as a letter, multi-choice, A, B, C, or D. So the model is used to filling in the next token. It's not really used to just this big long question, oh, format your answer like this. The model's not used to just saying— C. It's used to saying, oh, astrophysics is this field with all these things. It's used to actually using text. And for the model to understand you are supposed to answer with a single letter, and that letter is supposed to encapsulate all of the background knowledge of whatever the question is trying to get you to answer, that is something that shows up very late. So yeah, MMLU is like random for most of training. And then after— It depends on the architecture, actually. So for A dense transformer, you need to be later on in your learning rate schedule. So your learning rate has to have settled down a little bit, and then you have to have seen something like several hundred billion tokens. And then suddenly, MLU just starts to climb pretty linearly until you saturate. So it's flat, and then it grows up, and then it flattens again once your learning rate is too low and you've seen enough tokens already. For pure SSMs, actually, this number is Falcon and Mamba, this was like 3 or 4 trillion tokens. I think it's multiple trillion before. And we theorize this is because of in-context learning. Like, it takes a long time. Like, you need a little bit of attention to kick the model into understanding. Early on in the question, it told me answer with a letter, and that sort of relationship gets lost in a pure SSM block. But yeah, this is what— this is all it looks like.

So I often think back to that paper that was given Paper of the Year at one of the major conferences. Emergence is a mirage, basically, was the sort of counter narrative to the earlier narrative of emergence is a big deal. It sounds like you're on the side of emergence is a big deal, and I'm basically with you.

Yeah.

But to try to steelman the idea of emergence as mirage, I think one of the more powerful points that they made was like, A lot of these graphs that we see are logarithmic, obviously. And so when a thing is presented as a spike, that is still often over actually an order of magnitude of data, right? You're like 10 to the whatever flops, 10 to the n plus 1, 10 to the n plus 2 flops. And you're like, oh my, that came up fast. But it's actually, that was 10 times more than everything that came before in that plot.

Yeah.

Are you seeing things emerge, is that consistent with what you're seeing in terms of emergence? Does that change how you think about it at all? Or are you perhaps seeing things that are spiking even faster as you really drill into what's happening? And I guess one more, if I could tack on even one more question there. You said the term grokking. Is that basically your mental model for this, that these things are learning the right algorithm or the right conceptual representation to make a phase change?

Yeah, it's hard. Everything I'm about to say is speculation, as just a very empirical one in my case. So it's definitely something to do with learning rate. So there's papers like Mini-CPM where they just keep the constant learning rate for a long time, and then all of a sudden they anneal it really fast to zero. And you get all of the quality benefit during that annealing phase at the very end when you're annealing it very close to zero. And also Mini-Token. It requires a lot of tokens, and then also something about their stage in the learning rate schedule. My sort of theory here is that the model needs to see a lot of tokens. Like, that's how I'm saying grokking. Like, it needs to grok, okay, this kind of type of data that requires me to follow instructions exists. That requires a lot of tokens for the model to learn. So that's the grokking part. In terms of learning rate, I think it might be you need to settle more into the loss basin in order to, like, make these really fine-grained changes to how the model views the world. It really makes me feel gross to say something that generic, but it's some sort of relationship between the two of these. But both of them I would call grokking together.

Yeah, the learning rate I do think is a— or the learning schedule is a probably generally underappreciated part of this. It's not something I'm like super deep on either, but I've been watching for Sophia. I don't know if you've looked into that, like alternative to Adam. Do you guys use that?

We do not.

Interesting. Have you tried?

Yes. We think that this is partially because of the learning rate schedule. So more aggressive learning rate schedules make a mirage of increased quality when you're really just playing with learning rate. So we stuck with Adam for now. We might be looking at like Shampoo here soon, but not Sophia.

Okay. I want to understand that a little bit better, but just to try to set it up for you, if we envision a— of course, we're all going to do this in different ways in our own minds. For me, it's a 3D landscape, even though of course this is huge numbers of dimensions in reality.

Yeah.

But I subscribe to the school of thought where we visualize things in 3D and then say n dimensions really hard in our head to allude to the famous cartoon about that.

Sure.

So these are very strange landscapes, right? Where we're seeing an observation, computing a gradient. How would we have done better on this observation? Let's move the weights a little bit in that direction and we'll be at a slightly lower loss as a result of having made that change. And then in theory, that gives us better performance. Although of course it's a super noisy thing because an update that made us work better on one particular observation might not work better on another observation. There's batching that goes into this and aggregating a bunch of observations at once and making an update. And this is basically a very black magic sort of thing still, where nobody really has a great sense for what the true nature of these landscapes are or how best to navigate them. But when you talk about aggressive learning rate versus annealing, that's basically saying, how big is the step size that we're going to make, right? From one moment in weight space to the next moment, like how big is that delta? And if you find that the value is really happening most in this annealing phase where you're shrinking the step size. This sort of suggests to me like a, almost like an Antarctic landscape where there's like lots of little crevices that you could be like around a good space. And then you really have to drill into this very, very local space in the landscape to find something that really works well.

Mm-hmm.

Then that would suggest, because obviously we see a million things happening in the broader ecosystem, that would suggest that there's probably a lot of them. And that's where I get to this crevices. Notion. But then I wonder there, I'm like, I also see papers like the Platonic Representation Hypothesis, where at the same time that's happening, there's this other line of thinking that's at scale, things seem to be converging. The basic notion of that paper, I'm sure you've at least come across it a little bit, is with greater scale, with more modalities, it seems like the internal representations of concepts are converging in some high-level statistical sense. So how do you make sense of— can you begin to make sense of all that where there's seemingly a lot of different local points that you could drill into, but there's also this sense of global convergence? Can those be compatible? Does one have to be wrong? I find myself confused.

For the pocket idea, I don't— it's definitely not a flat area with lots of pockets. If we were to go with that, it would definitely be a general slope downwards with pockets along the way, and then you can settle into one of them as you go. This is just because even in the constant phase, when you're taking those big steps, people like MiniCPM out there still saw improvements, just less improvement. I definitely know that your loss landscape is much smoother for larger models. You can more richly model the space, and therefore just increasing the scale makes it much easier for your optimizer to make progress. This is why things like novel optimizers like AdaFactor or Sophia are second-order optimizers. It's like the bumpier the landscape is, the more accurate your step should be so that you don't get kicked out of the gradual slopes that you want to reach. That's also why we're seeing A lot of the conventional wisdom is assuming big models because big tech assumes big models. So the reason people all do dense transformers with large-scale optimizers, the reason why people say, oh, Shampoo doesn't scale, and all this kind of thing, is because when you're at large, your step is going to be broadly accurate anyway when you're in this giant smooth loss basin. Why bother with any of these things? And I agree with them. You shouldn't bother. But when you have a 3D novel architecture model with a really bumpy loss landscape, that accuracy really matters. That also translates into things like batch size. So I trained the Pythia models back with Eleuther, and we were seeing that, for example, like large batch sizes was causing instabilities. And it really confused us for a long time because the higher your batch size is, the more accurate your gradient should be, right?

Mm-hmm.

So if you have these super accurate gradients, this is even at the cost of seeing less data, right? As you increase batch size, you'll see less data. But now I'm fairly certain, and by looking back at the data, if you have bad samples, and as we were increasing the batch size, we were also increasing the learning rate. And if there's just a few outliers in that batch that kick your gradient in the wrong direction, and then you have a large learning rate, you take a big step in that wrong direction. Then larger batch sizes were actually bad because you were reaching too many of those outliers per batch. And you reached too many bad batches in a row, and now you're going completely the wrong direction. You just diverge entirely. So we actually had to use lower batch sizes, but it's complex. You also need to— this is why people clip their gradients also, because if you take one step in the wrong direction, make sure it's not a big step and then hope that big step can be reversed by the next couple steps. So there's a lot of like weird alchemy going on to get models to actually converge. I've noticed that every run has its own cursed way, and the reason people really want to use Adam, the reason we're using Adam, is because Adam is extremely robust. That's the main reason. So that's why people can publish papers like Sophia and AdaFactor. If you look enough, then you'll find cases where they're better. But Adam is good at everything. It's not the best at some things, but it's good. It doesn't have any— Glaring weak points, like optimizers like AdaFactor and such that might just not work. And then you spend 3 weeks trying to get your initial stability for this weird combination of hyperparameters and model size to actually work. That's how I look at it though.

Somewhat novel intuitions for me that I want to just confirm with you that they match yours. Overall, you're like, it's easier to get big models to work. It's harder with small models. That seems to be maybe a reflection of just how overloaded the small models are. If I think about just a super dense, everything packed in super tight sort of network, then I would imagine that it's just easier to mess things up if I take a step in the wrong direction. Whereas if I go to a trillion parameters at the very far end of reportedly a GPT-4 type architecture, then I can spread this out over many experts and things aren't so overloaded. And one wrong step in one place might not even affect most of the network because it's localised to experts or what have you. And so then I can be more aggressive in my step size, and I can get something like a Sophia to work well in those bigger contexts because I just have more room for error, essentially, in my training process that I just don't have that luxury if I'm training a smaller model. Am I interpreting you correctly there?

You're close. So for larger models, you want actually lower learning rates because even though it's a smoother loss landscape, you don't need the small model ability to really jump over these crazy hurdles and stuff. Think of Adam, right? You have a variance term, you have a momentum term, and that's to tell you if there are bumps along the way downwards, Then you still make progress downwards because you still have momentum. You have faith that this way is probably the best for the loss landscape despite some bad gradients. For larger models, despite the smooth loss landscape, you want a lower learning rate just because a lot of instabilities just in general tend to come up. It's because you have so many parameters that go beyond the point of no return.

So I heard the Adam one, the momentum is there to say, We're not just going to look at this particular batch of updates and calculate naively what would have made us better here, because we know there's a lot of randomness. And so we're going to keep track of the general direction we're going, which is the momentum term, to make sure that we're not getting too randomly sent off in weird directions. We're going to try to have some general sense of what direction we're going. I'd be interested to hear the Sort of similar narrative form for Sophia, if you have one there that would explain why that's different and maybe why it's better in some places than others. And then I was also going to ask about distillation. It seems like the big trend with the frontier developers has been go huge and then distill. I wonder if that's something that you guys might also be thinking about, if there's like a Zamba 70B or even 400 and however many B in the background somewhere.

So in terms of Sophia, I don't have a good intuition for this actually. I've explained why Adam is great and why I think Adam is good at everything. It's very robust, but I can't really explain why everything else is brittle in its own way. And this is not really just to say that everything else is bad. It's just to say I haven't put the time into any other optimizer just because I've been burned enough times by alternative training strategies that just end up not Being generalizable. I don't want a bunch of, like, a toolkit that's specialized to each training run. That's not the point, because then I have to spend as much compute as a training run to try and figure out how to get this specific thing to converge. I want it to just converge, and even if it doesn't converge quite as well, I'll train on more tokens because I was able to kick it off day 1, and I'll still end up at a lower loss. On your MOE point, so MOEs are themselves, like, a really tricky thing to optimize. Because like, for one, the attention block is not divided. The attention block sees everything. But for the expert blocks, you're subdividing the batch across everybody. So like, they're effectively seeing like a different gradient noise scale than the attention block. So trying to find the right batch size is clearly much higher. But do you scale it by number of experts higher for the batch size because you subdivide? Like, that intuitively makes sense. But the problem with this is that now the attention block is getting this massive batch size. This is part of why MoEs are super unstable, I think, to train. Like Black Mamba was like a pain to train because of this. Like attention blocks are seeing more accurate gradients is one way you can think of it than like the experts. Trying to port mu transfer is something that I'm really looking into here. Actually, the people at the team at Cerebras is really deep on this. They're really good. So we're trying to crack this, mostly them, but yeah, we're working together. But in general, trying to find for these hyperarchitectures that we're seeing, for example, that's kind of true for us too, because we have this shared attention block, which is going to have more accurate gradients and see more data. Finding something that converges is really difficult. And being able to navigate the hyperparameter soup really determines if you live or die.

Fascinating. What about— did we get to distillation? Are you guys going to pursue a sort of big-to-small strategy as well?

Yeah. Yes, that is the plan. The problem here is that you need a lot of compute to train those large models. So it's something we're scaling up for now. In terms of the actual distillation step itself, unfortunately, big labs are very tight-lipped about this. So like even annealing, it started emerging— oh, annealing, by the way, it's like a continual pre-training step that we picked up early. Where you first do like a standard cosine decay, you warm up the learning rates, and then you slowly decay it over time to some high value, higher than normal. And that's on like your phase 1 data. So this is like more like web data, it's a bit noisier. It's still important to have so that the model can like— knows what this data looks like, knows how to respond. Talk like a pirate. Yeah, exactly, exactly. But later on you want to bake in more intensely A higher quality subset of your data. The higher quality doesn't mean textbooks, again, doesn't mean you want your model to be dry, but you can only look closely at so much of your data. Everyone says, look at your data, but you can't look at 3 trillion tokens. You can look at 50 billion, 100 billion tokens and apply some really detailed filtering and deduplication to those and make sure it's the highest quality, even if it's a distribution that represents all of your data. Anyway, you decay down to a relatively high learning rate on your first phase, and that's where most of your tokens are coming from. Then in the second phase, you very quickly rewarm your learning rate. You say, okay, I want the model to respond really quickly again. We're in a new loss landscape. We have a new distribution of data. Then very aggressively decay down to zero. You say, okay, I want my model to respond quickly to this data. And then settle very quickly into the new loss basin that emerges from the high-quality data. And also, you're able to stomach this new high learning rate again because your data is very clean. You don't have this problem I was mentioning earlier with Pythia where you have these crazy outliers that kick your gradient in some different direction. So you're able to go to— with a model that's already settled into a good place, you're able to increase the learning rate a lot on this high-quality data and then really Bake it in so it's top of mind to the model. This also increases the flexibility. I see it as a branching path. The base model is the base, the trunk, and then you can construct-tune and then fine-tune so your model's becoming progressively narrower along one branch. But if you have multiple anneals, that's like multiple sub-trunks, right? So you can anneal on roleplay kind of data. You can anneal on, like, if you want to try, uh, like very Factual model. You can anneal on textbooks and such, but then from there you can fine-tune so that the roleplay model is roleplayed on like specific characters from the annealed roleplay checkpoint. And that is a much more flexible scheme than just having a single base model, a bunch of fine-tunes from that base model.

Can you define anneal? I've learned this in the context of like biology, and I'm looking it up online. It's also used in glass and metalwork. Maybe this is wrong, but I guess I think of it as hardening into its final shape, or with DNA, it's actually combining, coming back together and reaching some lower energy stable point, basically. Is that the mental model here too?

Yeah, we put very detailed instructions and explanations on— it's called Letter Z and then Cookbook, Z Cookbook. If you look it up on GitHub, we have a whole section on what is annealing, what it means for us, what it means for your model. But yeah, annealing for us is like trying to harden the model a little bit, the base model, into a pre-fine-tuned checkpoint that has some type of data that we want top of mind. And practically, this means rewarming your learning rate on the new distribution. It means having some replay. So you still have to feed some of the tokens from your batch from the phase 1 data so that you don't have catastrophic forgetting. And then that replay plus new data distribution, then you very aggressively anneal the learning rate down so that your model settles very definitively into a loss basin for you. This all getting back to the reason I brought up annealing was you were asking about distillation. Since the GPT-3 paper, the big labs have learned, be a bit less open with people because then they'll actually take compute and try and do what you're doing. That's why annealing for a long time was like a hidden trade secret. Like the LLaMA team cracked this with LLaMA 2. They split it into the Mistral team. They split off. They knew this sort of secret. And then it became more and more public over time. And now everyone publishes it. Distillation has been the same sort of regime where the big labs have cracked it, but they're staying tight-lipped about it so that they can try and extract all the value before their people get poached somewhere else. And then it ends up being public anyway. We don't really know the secret to distillation, actually, but synthetic data from large models is a proxy for this, right? So we're trying to, to look at this at least so that we can steer our annealing and post-training strategies to, to more approximate the output of large models via synthetic data. And then trying to actually use the logits during training to stabilize. So remember earlier we were talking about— Can distillation help you navigate the bumpy loss landscape by having basically a big atom that's holding your hand and saying, okay, no, this is the right direction because this is what the big model would do. The big model is able to see much better that this is just a little bump and you should keep going down. Using those logits at training time is something that we're still trying to crack. We see generally positive things from a model conversion standpoint, but it's still really inefficient. to try and have a live inferred big model at training time. So training, pre-training scale distillation, we have not cracked yet, but I'm confident that we will.

So the basic idea that I have of distillation is training a small model to reproduce the outputs of a big model, not so much at the token level, but at either the logits level or potentially if you had like direct access to the large model, then you could perhaps tie it, like, multiple places, right? You could say, I've got my trillion parameter model here, it has 100 layers here. I've got my 7B, it's got however many layers, but I can map. Okay, at this point, I want it to tie up to the big model and try to figure out a way to imitate the internal representations. Am I off base conceptually there at all?

That's exactly right. And maybe I'll expand on what I was saying to tie it to what you are. So when, when I'm saying synthetic data is like a poor man's way of doing distillation, that's, you're just choosing the max of the logits, right? So like you're only training on what the big model is paying attention to. And you can do that also with closed source models. We could do that from GPT-4 or whatever else. Now logits are important, uh, because those are like the mo— the large model saying exactly how it would respond to the current batch of data. That's why I'm saying it's kind of like holding your hand in terms of the optimizer. Logits are also super important because you also want to know what the large model is not paying attention to. So those tiny 0 values in the logits, like the whole distribution helps the small model, not just the large values that you would get if you were to approximate with just the synthetic data. But yes, you need access to the model for these logits, which is a big hurdle. It's helped now by big LLaMA 3, 4, or 5B. But when you're running inference of Llama 4 or 5B every single training step, that's why I was saying that at pre-training scale, this becomes prohibitively expensive. And trying to get around that expense hurdle while still getting logits per step is what we're trying to figure out. We think there are some cheaper ways to do this, and I think we will get them figured out. But this is why it's not popping up all over for open-source people, because it's either prohibitively expensive. Or synthetic data is just not like rich enough of a representation to really help you steer in the right angle.

And so the, just to make sure I understand what makes it prohibitively expensive, it's basically that running 405B is like some fraction as expensive as training 405B. So if you want to train your small model, but also have to run 405B, it's, you're going to end up spending as much compute as if you were just training a much bigger model in the first place.

Exactly. Yeah. In terms of practical sense, if I have a 7B, then I'm already doing parallelism across GPUs to get it to fit in VRAM. And then if I have 400 billion parameters that I also need to store, even if it's for inference, you don't need to store optimizer states and gradients and all this for the big model. You still need to fit it somewhere. You still need to use the FLOPS to infer it. You need to synchronize its forward pass with the forward pass of your own small model. Why don't I just train like a 4B from the start? And why even deal with distillation? Is it really worth it? So far, these practical trade-offs have not been cracked for us. That's what I'm saying.

Yeah, gotcha.

It's not practical. It's probably a slight improvement, but a lot of things are slight improvements. And just seeing more data is always a very, like, direct way to improve your model. And I could see more data or have a bigger, more parameters, and that's a more direct way to improve model quality right now. than trying to get logits from Llama 4 or 5B. That's just the practicality.

Do you think it's possible that the big labs maybe don't necessarily have a huge secret, but what they maybe have is just a different cost-benefit analysis where they're running so much inference that it's worth it for them to do this?

Oh, I think this is likely. Yeah, I think this is likely. I think that if you're Microsoft working on Phi, then just use GPT-4. It's, uh, on our own systems. We have infinite funding. So making it— there's 2 possibilities, right? Either they've cracked a cheaper way to do it and we have to find it, or they have more money and we have to find the cheaper way, crack it, and then do it. And both of those cases require us to crack or find some cheaper way, unfortunately.

Are there other things that you would put in the bucket of Frontier Labs are sort of generally understood to have a solution that nobody else hasn't diffused to the rest of the community yet?

Data is a big one. Data and data cleaning requires time, patience, and really just manpower that open source does not have. So when you're open source, then you're just beholden to whatever dataset other people are able to like completely focus on, and not a lot of people focus on data. So it will always lag the big tech companies. Also, Big tech has access to their own data, which is also, you know, very clean. It's, you know, in infinite abundance, and we're definitely behind. So yeah, clean data, clean data pipelines. This is what we talked about. We just released Zyda 2, which is our Zephyr Dataset 2. There's some new datasets in the open source regime called DCLM and FineWeb and Dalma and Zyda 1, and cleaning them up and producing like an actual, like, at-scale dataset with, uh, Nemo Data Curator, which is NVIDIA's, like, data processing tool. We found has really lowered the bar for entry for data. But yeah, data's one. Another one is things like hyperparameter transfer, unfortunately, has not been picked up a lot by the community. You might recognize this as MuPy or MuTransfer. There's like, I think it's Modula was a new paper that's also found similar things. But it's like 30, 40 pages of math for these papers. And it's not approachable for open source people or small labs and such. So only big labs have the compute and funding, like, to put one person, okay, go read the MuP paper for a month and try and port it to our ginormous monolithic codebase and get it to actually work. So there's a lot of trade secrets, I would say. Distillation is another one.

Give us a little bit more on that one, because I honestly don't have a lot of intuition for what that is adding to which we're all stacked.

MuP?

Yeah, exactly.

Oh yeah. Oh, so mu transfer is a way of initialization. Were across— it's like width scaling. So you do all of your hyperparameter search on some really low-width model that's really cheap to run training on, and you do all of your learning rate sweeps, for example, all of your batch size sweeps. And then if you literally just divide by a scaling factor of your low width by the high width of the actual model that you want to train, both in terms of the learning rates of your hidden layers, like your embedding scaling, you can multiply by some a new hyperparameter, then you just get perfect transfer from the small model to the big model. And as in the learning rate that I found was optimal for my small model will also be the optimal learning rate for my big model. And this is all just making sure activations don't explode by all of these small scaling factors of width. This helps you in terms of both— you're able to more accurately search the loss landscape because the compute is so cheap. On the small model that you're able to really search a higher amount, search space of learning rates and batch sizes and stuff. It also helps just the general stability of your model because your activations aren't exploding and your model's not learning on the fly how to account for its own explosions. So you're able to also deal with a much higher learning rate. Alternative architectures and optimizers are also much easier for this because you can just directly, for example, the Mixture of Experts, If mu transfer is cracked, then you know exactly how to ensure both the attention block and the expert blocks are both seeing relatively equal magnitude weight updates, which is what you want for stability. These sort of details tend to be hoarded in, in big tech right now and slowly diffuse out to everyone else just by people going to different organizations in big tech, people publishing, people in open source kind of cracking it, that sort of thing.

If there's others in that category, keep running them down.

It's endless, man. So this is why EleutherAI is so important. I mentioned them before, an open source collective where everything start to finish is open source. The dataset, the training framework, the actual produced models. If you don't see all of it, like for example, Llama 1 was great because they shared everything. They shared all of their data pipeline and people made Red Pajama, which is a full reproduction. Llama 2. Was they stopped sharing the dataset. And as I told you before, this is because of annealing. It's just too tempting to keep the important parts to yourself and publish a paper that's just vague and glosses over it. But keep going. What else? So for example, the second-order optimizer thing that I told you for small models specifically is definitely something hoarded. So that's like a trade secret. I mean, Gemma also did it. So they just say we use a second-order optimizer.

Which one?

How do you distribute it is a big one, because if you have a second-order optimizer, you have a bunch more states that you have to store in memory, and how you split those across GPUs is non-trivial. So that's another engineering detail that they want other people to play catch-up on while they extract the information from it. So second-order optimizers for small models that are distributed in a specific way. Now we're cracking it with Shampoo. And then all of these things that I'm mentioning all interplay. If you have a second-order optimizer, How do you port MuP to it, MuTransfer to it, right? Because now it's not Adam, you have to reapply the theory. Do you spend another month trying to figure out the new theory? I would say, curiously, alternative architectures like Zamba are not hoarded by big tech because big tech saw the quality issue with pure SSMs in a lot of ways and said, okay, quality is most important to us and we have infinite inference time compute. So let's not bother with it at all and ensure we have the highest quality model. And they— that's more for us, right? But I think they were much slower to catch on to SSMs and alternative architectures just because it was not super clear to them at the time that you could solve this accuracy problem by some exotic architectures and by being really data-driven. What else? Parallelism schemes for all of these, like I was starting to mention with the optimizers, but parallelism schemes for new model architectures and optimizers and combinations, those tend to be proprietary. So LLaMA shared that they had like a 4D parallelism sort of scheme, but how you set up that topology on your own, that's all like empirical knowledge that you have to get from training a model. Same thing with like arcane things like checkpoint restart. So like a lot of these bigger runs, they use like dedicated nodes that are just being used to save checkpoints. And then when you have thousands of GPUs and one of them dies, uh, all the time. So you have to keep slotting in the checkpoints from those backup nodes. That's something that's really easy conceptually, but hard to set up and in some sort of proprietary big tech stack. It's everything is what I'm getting to. There's a lot of points, especially about gritty details, engineering details, scalability details that are very closely kept tied to whatever organization has them. And they also know that even if you say worked on one of those big tech teams and you know the secret and you go somewhere else, Porting it to their stack is a very non-trivial thing. It's completely built off of the proprietary stack of whatever company came up with the idea. So do you remember it well enough to port it to this new stack? Maybe not. So there's really a hardening of ideas and training processes that exist in each organization and that is increasing over time.

So I think we have a few more technical questions. Obvious kind of pulling up question for a second would be like, the story you're telling sounds pretty consistent with the allegedly leaked Anthropic fundraising deck in which they said that the leading companies in '25, '26 might get so far ahead that nobody can catch up. And I think we've understood that, or I've understood that broadly as sort of an accumulation of these secrets and then also just like having the best models. We now have this, I think, pretty credible reporting that OpenAI is using 01 to generate the data for the next scale up again. And that seems like there's not just a moat, but there's multiple moats in depth there.

Yeah.

And yet you guys at Zyphra are a small company and trying to figure out how to compete in this space. So what is your sense of how you and other companies that are not infinitely resourced, even if you're like reasonably well-resourced, how can you compete? Is it about finding a niche that they're not going to play in, or is it about some sort of breakthrough insight that is going to be needed? How do they not just come to dominate the entire space?

Yeah, I will say small models is overlooked by big companies because like OpenAI, xAI, Anthropic, all of them, their job is to Producing the highest quality model. And yeah, sure, they wanted to distill into a smaller model for inference reasons, but if it came down to cracking distillation or making an even better, bigger model, that— because it only needs to be a little bit better. Claude Sonnet 3.5 only had to be a little bit better than ChatGPT to take a large portion of market share. So accuracy really takes them. That's also why they don't want to like experiment too much with these weird hybrid architectures, because if it affects model quality and you only determine that When you train a big model, it's a lot of cost when you know that dense transformers work the way that you expect. Yeah, in terms of how we can compete with them, I would say it's a mixture of this. They're focused on large-scale inference, large-scale training. Definitely makes them— even if they have things like distillation, that's only one part of the process, right? So like, we had to crack annealing, we had to crack architecture, we had to crack data, the entire like— training process from start to finish. It wasn't just like the distillation step, which is only helped when you're at a big company, right? We're very nimble in that way. We're definitely in a unique time where if you're nimble enough and you're focused on small models, like, you can very clearly have an advantage over big tech over time as more of them use distillation and such. That if we were to not move at all, then we would definitely vaporize compared to their resources. So the way you're staying ahead is by continually pushing onto what the next free lunch is for small models. Uh, next, like, second-order optimizers is one. Getting a cheap way to do distillation is one that will protect us from that angle. Like, we're clearly very willing to make big bets on the alternative ways of training, and big tech just is not really looking very closely at this sort of regime. And that's why we have the best model now. The 7B is the best because—

Are you studying the Entropix project that has recently been catching a lot of buzz or similar?

Very loosely. I'm not an expert. I would kind of categorize it under the definitely doesn't scale. So I believe it was Tim Detmers tried this at scale with like larger parameter counts and it was like no better than Adam. Another thing is that I kind of view it as kind of specific. So whether it's robust or not, or whether it's very specific to the training setup. On like NanoGPT is yet to be determined. It's true with this for a lot of things. You can show like huge speedups on even the, what is it, like crowdsourced compute, like pre-training compute. But it's hard to show things that are robust. Like that is the actual test. Is this robust? And I'm waiting for someone else to spend the compute to determine if this Anthropic's is robust.

You're describing it in a different way than I expected.

Okay.

The one that I have been at least momentarily fascinated by recently is essentially an improved sampler. So it is like trying to differentiate between when the model is confident and should just pick the most likely next token and everything's kind of hunky-dory versus when it's out of distribution and then when it might need to— it does interesting things like kind of dynamically inject a additional clarifying question. If it detects that generally we're confused, let's maybe just burn some more cycles.

How can we do that?

Let's artificially input a clarifying question. And then sometimes it'll even back up and say, we seem to have kind of lost the thread here. Let's go back to an earlier point and rerun inference again. So I've thought about this entirely as an inference thing. You were going toward the training side. Are we talking about the same project? And if so, I missed the training side of it.

You're making me question myself. Maybe I'm confusing it.

Well, for what it's worth, in terms of a possible next free lunch, the one thing you did say that matched my intuition, even though I've read a little bit that it might not be working as well at larger scale models, I wouldn't say that's well established in my mind. But if there was a story as to why it might be, it would be like the larger models are maybe like, Grokking into some version of that on their own, especially through all the reinforcement learning that they're doing in late-stage training, whereas the small models are maybe not getting there. And so layering on these sampling heuristics could be just much more beneficial at a small scale. This naturally has me thinking too, what about learning a sampling function as part of the overall architecture? That feels like a candidate for, if any of this is real, I haven't quite got to the point where I'm convinced that there's a real there, but it seems like there is. Perhaps learning a sampling function could be another relatively small amount of parameters that could lead to some qualitative change.

I agree. It definitely could. I also agreed on the point you made about larger models doing this on their own. Larger models, for example, data quality matters a little less because larger models learn to self-select what is important, what's not. There's a reason that, like, when you go to a large enough model scale, you just scale it up and just churn it out, churn out more and more data on more and more parameters. And the model kind of is so robust that it's able to learn these properties that we would need explicit training properties and processes for at the smaller scale. So that is one thing that big tech kind of relies on as a crutch, that we are able to find the explicit process that is able to emulate the same behavior from the model a bit faster because they can just rely on the model figuring out on its own. But yeah, I do agree that it could help us, but we haven't looked too deeply into this just yet.

I'll ask one more and then Jason, let's come back to you and you can map some of this, uh, stuff onto the specifics of Zamba and And state-space hybrids in particular. How would you describe like where we are in this overall process? We always hear we're still so early. That basically seemed, that would be like my vibe level takeaway from your comments. There's a lot still that maybe hasn't even been properly investigated, let alone understood.

I would agree with this. I definitely don't think this is a solved problem. I don't think we found the limits of scaling. We clearly have not found the limits of architecture. Whether there's some groundbreaking thing tomorrow where like Attention 2 comes out and everything is thrown away, I can't really say that. But I can definitely say that we've moved very quickly into a search space that's massive. There's a lot of very low-hanging fruit that is just obvious, like it's everywhere, that we have not picked up. And these training processes and stability, like there, there's still so much out there. Even if you were to just take pre-training model papers and compare them. Like, I'm sure there's a lot of things that people are reinventing. There's a lot of trends like across scale and training process that are to be found that people have like implicitly found, but either not known or purposely not shared with everyone else that are just lying around.

That's my sense too. It doesn't seem like we're at the end of this cycle or anywhere particularly close, which is Pretty crazy to contemplate because it's already come pretty far.

Yeah, one really low-hanging fruit example of this is like the GPT-3 paper had like really specific like hidden dimensions, and this was both true in terms of quality, but it was also in terms of like kernel paths. So for example, there's like lots of powers of 2 in their hidden dimensions, but not quite enough. And everyone else copied the GPT-3 model configurations for a long time without understanding that, oh, These are more efficient on GPUs in addition to being approximately the right size. So this is another thing we bake into our Zamba models, is that since we're trying to make the most inference-efficient thing on any parallel hardware, if you do very simple tricks, put lots of powers of 2 in your hidden dimension, like do all sizing of these different blocks just by rounding very slightly to good kernel sizes, will give you speedups on any parallel hardware. So whether it's multi-core CPU on a phone or some like NPU on Intel, like an M1 for Apple or something. If you have parallelizable model sizes, then for the lifetime, the entire lifetime of the model, you're more efficient. And this is something that big tech just says Llama 3 has this size, but they don't show you that they went per block and found the efficient sizes for kernels for each of these and then baked it into the model. Same with the vocab size. The vocabs are all rounded to some factor of 64 for the same reason. Like, there are just more efficient kernel paths you can go down with that.

And my very simple intuition for that would be just like, if you want to break something out into 2, 4, or 8 GPUs, then if you have a 9th thing, then you're probably like, if you go 8 wide, the difference between 8 and 9 is 1 and 2 in terms of how long you have to wait. Are there a lot of different versions of that you see, or is it really pretty much that simple?

It's not that simple. So you can decide how baked in your model size is going to be. So for example, this, I mentioned a few times, this powers of 2 in your hidden dimension. This is true for any parallel hardware. Like if you train on an AMD MI300X GPU and you have lots of powers of 2, it's probably also going to be efficient on NVIDIA H100. It's going to be efficient on your phone. But you can get really deep with this. There's something called wave quantization where let's say you have, I don't know, 100 streaming multiprocessors on the H100. If you have 100 of 1 units of work, then you need to do 2 time steps, right? First one with 100 and the second one with 1. And the throughput will be roughly halved because that second time step is pretty much empty. About the same amount of work, but 2 total time steps, throughput is halved. Whereas if you had 100 units of work, you just have to do 1 time step. So then on your graph of throughput on the y-axis and problem size on the x-axis, you're gonna get waves. And each time you kick down on the wave, like you have lower throughput, you are on like a new time step. This is baked into your models. If I was NVIDIA and I wanted to train a model that's only efficient on NVIDIA hardware, I would bake in that this model has very specific sizes to fit on the number of streaming multiprocessors on the H100, has the amount of SRAM, whereas on an AMD competitor, like an AMD MI300X, On that one, like, it'll require 2 waves or something like that, or SRAM will be like undersaturated with the model sizes I put in. And like, the bad version of this feature is that if everyone is building their own chips and everyone is building their own models, then everyone would produce a model that's only efficient on their own chip and is like specifically inefficient for everyone else. So like, there's all— it's unclear whether we'll go this direction of like perfect specialization per company, or if companies will say, okay, I can't crack all of model training. And I need my hardware to be able to run competitor models that we run on our own data. So maybe in that sort of case, you would want everyone— the good future of everyone provides similar hardware that can run everyone else. It's unclear which way this goes, but this is just another example of there's a lot of low-hanging fruit for one, and for two, there's a lot of return that you can get by just getting out a magnifying glass and looking at every single step of your training process. Extract everything you can from it. Yeah.

You follow Greg Brockman on Twitter and you realize how much of that he has personally done over the last couple of years at OpenAI. And you're like, yeah, there's a reason that they're obviously extremely good at what they do. And you can see that the sort of the rough outlines of how that work is unfolding there in his Twitter feed in particular. So recommend that as a not very open window into what's going on, but at least a tiny little peek.

Yeah.

Jason, what's on your mind?

Yes. Just a pleasure. Super interesting conversation to hear the ins and outs, the inside scoop on training. I guess I would really be remiss not to do a couple deep dives with Quentin while we have him on Zomba 1, Zomba 2. Maybe let's just start with a deep dive on Zomba 1, if that's okay. And in particular, you look at the architecture, seems pretty straightforward. There's 6 Mamba layers to start off that the input sequence goes into. And then this interesting thing happens that I'd like to talk about. So there's this step before the attention block of concatenating the residual stream coming out of the Mamba block with the original input embedding. So this is very interesting. So the attention block gets to see the Mamba residual stream plus completely unmodified Input embedding, it feels right that if you're running the attention computation, you'd want to do this. But at the same time, it's interesting. How does one arrive at a feature like that? Is this something that your team is, okay, it's intuitive, so let's try it? Or is this something that through many, many different experiments that just rises to the top and it's more experimentally found?

Okay, on the specific concat, if I recall, we did run a lot of experiments, but it also was intuitively right. Our experiments matched our intuition, which is where we want to be, that we want the attention block to be able to see like the entire input and residual. There's really not much more to that. Empirically works and it intuitively makes sense, so we stuck with it. In terms of Zombo-1 in general, we kind of crystallized more and more of our thoughts on how an inference-optimal model should look like. So that's when we discovered that yes, indeed, attention blocks are correlated. Yes, indeed, you can replace them with just one and get pretty much all of the same benefits of attention. Yes, indeed, hybrid models with the rough 6 Mamba blocks to 1 attention block was like a good split. Let's see other things. Everyone likes to talk about model architecture, but it's one step of the 3 holy— the holy trinity, right? So model architecture is a big one, and we really optimized that for inference efficiency and like performance per parameter, as we discussed earlier. And number 2 is training framework and process. We learned from the mistakes of Black Mamba of how to affect for optimizers and stability, how to tune hyperparameters for exotic model architectures. We baked all those in. And then number 3 is data. Zombo-1 also came out with Zaita-1, which was our training dataset at the time. Really getting quality baked in at scale for this size of model was really important. We have some nice graphs where at 1 trillion tokens, we're able to Perform, you know, similar to models trained on 5 trillion, 10 trillion, 15 trillion. If you really optimize for quality at this smaller model scale, it really has effects. Annealing was first found with Zamba-1. So this process of 2-phase training, we discovered it with this. So this is like the first pass of everything in our overall training process. Any thoughts or anything you want me to drill down on more?

Yeah, I guess if we could just Walk through the layers, and they're pretty straightforward, 6 Mamba layers. And then just a follow-up question on the concatenation. So I guess first thing is, because you're using concatenation, so you have to pay for that computationally. I guess the dimension is expanding. Obviously, the trade-off's worth it if you're getting a bump in performance.

Yeah.

And then I guess at some point, it's not entirely clear, at some point, I guess you're projecting back to the original dimension. So if you could just Step through, imagine there was a visual on the screen, step through the Zamba-1 architecture from beginning to end.

So yeah, no positional encodings. Then you have an embedding block, and then you have this alternating between Mamba-1 blocks, because it was only Mamba-1 at the time was out. So 6 Mamba-1 blocks, global attention plus MLP. So there's also an MLP in that global block. That's important for people to know. 6 Mamba-1s, attention plus MLP. 6 Mamba 1s than the same attention MLP together. In terms of the concat, yes, empirically it is better to do a concat than to try and like— I think we tried to like project it back down and not eat that computational cost. Nothing really worked very well, so it's very empirical.

So I guess in the literature there's a range of results of Mamba 2 attention layers. Zamba lands on 6 to 1. That's consistent with Zamba 2. So it seems that don't fix what's already working. But some of the literature suggests wide ranges up to 7 to 1, 10 to 1. You know, how did the Zypher team arrive at the 6 to 1 ratio? Was that experimental or?

That was purely experimental. I also suspect that we have to apply our attention block more often because it's shared. So there's some sort of interplay there. Yeah, I mean, 7 and 5 are slightly worse. That's really all it is for us. I think it's similar to everyone else. Jamba, mixture of experts, clearly has some sort of interplay with how often you apply attention. I'm sure there's some, there's some consistent rules here that are being applied to everyone, but the way it's manifesting is we all are empiric— we don't have time to make these, like, these consistent rules play nicely. So empirically, we find that, yeah, for us it's every 6.

Excellent. So I guess jumping from Zamba 1 to Zamba 2, so the Mamba 2 paper comes out in May. Both of the original authors published that paper, Albert Gu and Tri Dao. The main result being the SSD algorithm is adding a restriction, which is adding structure to the A matrix. That's the part of the algorithm that influences how the state changes over time. And so now we can do matrix multiplication.

Mm-hmm.

And take advantage of modern GPUs, get those tensor cores going. So now we can train much faster and as a consequence, much larger state sizes. When that paper came out, was this immediately on the minds of the team at Cypher and the kind of impact it could have? Was Zamba2 immediately planned? And what were the features of the Mamba2, the SSD that— gave your team confidence it was worth coming up with the follow-up model?

Yeah, after Mamba-1, we evaluated a lot of architectures at the time. And since Mamba-1 gained our confidence initially, Mamba-2 was something that we initially immediately wanted to test out. We ran a lot of ablations of pure Mamba. We found— pure Mamba-2, that is— we found, okay, it has the same problems with in-context learning and long sequence dependencies. So we can't just make a pure Mamba-2 model. Still. But we also were able to verify that we get significantly higher throughput on H100s with our training hardware for no model quality degradation compared to our initial Mamba-1 tests. We also found something a little different from Jamba. So Jamba 1.5, or it was either 1.5 or the initial one, I don't remember which, but they found Mamba-1 was better in a hybrid case. So they did like Mamba-1 plus attention. We didn't find that, so we found they really performed the same. So we stuck with Mamba 2 plus attention. But yeah, the Mamba 1 to Mamba 2 change in from Zamba 1 to Zamba 2 is literally just throughput. So we get much faster models both at training and inference time.

I'm not sure if this is the best time to interject, but I also, my ears perked up around no positional embeddings and I was wondering How you understand that working. Is that an artifact of the, or a sort of consequence of the inherently sequential nature of the Mamba blocks that they create sequential understanding where the transformer doesn't have it natively or the attention mechanism doesn't have it as natively?

That is our exact intuition. Yeah. So if you, we found that with and without them, with Mamba, wasn't much of a performance difference for the first model. So we're just okay without them. For the second model, for Zamba 2, we actually put rotary positional embeddings in. This is because we found there are some cases where you can make rotary help you a little bit. So we kind of had to figure out how to ask the right questions. If I recall correctly, rotary helps us a bit with long-range dependencies, but it hurts us in terms of context-length Rotary is very inflexible with its context length. And so that definitely did bite us. In the future, whether we continue with Rotary, whether we get rid of embedding or positional embeddings again, whether we go with some like Alibi or learn positional embeddings instead, I'm very uncertain with which one we do. But I will say that they're all very slight effects. And this is because you correctly state that Mamba kind of handles a lot of the position, like it already encodes the position. So like you don't live and die by the positional embeddings, that is.

Yeah. Context extension is another thing that I think is really interesting. And that was one of the biggest things that kind of inspired me about the Mamba architecture in the first place. At the simplest level, the state size doesn't grow with sequence. So that means we can, in theory, have an arbitrarily long sequence and have some sort of more, I don't want to say human-like, but on the grand scale of possibility space of memory, At least one step toward more human-like in the sense that obviously my memory, our memories are not all tokens to all tokens, but rather some sort of lossy compression of everything that came before. And I wonder what would be your outlook for that now? We've been almost a year since the original Mamba paper. I would say it's definitely shown itself to be a big deal. But I haven't seen, outside of maybe Magic, I don't know what they're doing, but I haven't seen many millions of tokens contributing to a single long-lived state. You talk about on-device, constant, continued pre-training, learning about me, baking that into the weights. That sort of suggests a disbelief in the idea that we can Compress my full history into some long-lived state. So what would be your outlook for training, not training necessarily, but running my whole Gmail history through a Zamba 3, let's say that, do we have a line of sight to something that could handle 100 million tokens and have this holistic fuzzy sense of who I am, what I care about, how I respond to things? Or does that seem just still like— Too far to envision a path to?

Yeah, in terms of engineering, there's some work to be done. For example, we have ring attention and tree attention for attention blocks for these million-plus context lengths. We don't really have that for Mamba, and this is just because the memory of training time is not high enough to support million-context Mamba, right? You need some sort of sequence parallelism for the Mamba blocks that just doesn't exist right yet, right now. And we're working on it. We hopefully will be the first. In terms of from the modeling angle, one thing I'll say is that I don't really want my models to be brain to model. I forget why I'm walking into a room. I can't reverse a string or anything else. The cross-sequence dependencies of attention are more powerful than the way your brain works in a good way. So I still want them to stick around. In terms of how Mamba behaves as super long context, I don't I don't think anyone knows this. I agree that it intuitively makes sense, but I don't have a specific bet either way whether Mamba will do super well or super badly because there's something also to be said about losing cross-sequence dependencies for Mamba at these really large contexts that are definitely noticeable right now. And whether attention will make up for those sins at even larger scales, I don't know. We might need more attention, for example. If we want a million contexts. So I want it to be there, but whether— well, I think there's a lot of challenges that need to be figured out first.

Yeah, yeah.

I'll just pick up on the— okay, on one detail you mentioned, uh, having tried to get longer sequence lengths trained into the original Mamba, you know, there you reach a point, you go through all the steps that you normally would to do long sequences and to save memory, and then you memory out, and The next big leap you'd have to do is figure out sequence parallelism. Much less straightforward, it seems like, for the original Mamba. Actually beat my head against it. I was like, okay, cool, I'll do a pull request. I think there's, for DeepSpeed, there's the Ulysses library. I don't know if you're familiar with it. Okay, cool. Let's get a pull request going.

Yes.

Could not even crack 1% of what it would've taken. With Mamba 2, it seems a little bit more tractable. They even mention it. I think in the paper that it's going to be much more transformer-like. I mean, any insights to— are we going to get the same kind of sequence parallelism that we're used to working with transformers?

It will look similar, and I don't want to say exactly how because I want to be the first. I do think this has an interesting implication for hybrids though. As I was mentioning before, whether attention can make up for it, I see a mixture of depths on the attention block being the best way to move forward where The model decides how much attention it actually needs at inference time and also at training time. So like, maybe if you want a 2 million context of all of Nathan's emails for his entire life, then maybe the model wants attention every single block. Can we also make the model decide like how many LoRAs or what, like how aggressive the LoRAs need to be at every block? I feel like dynamic attention is what I'm getting at. It's gonna be a big determining factor on how long context for Mamba really will work. But yeah, in terms of SP, it's unknown. It's still, still being figured out. Mamba 2 definitely made things easier, I will say that.

Excellent. Looking forward to that. Yeah, no worries. I understand you wouldn't want to share details at this time.

For anybody who might be a little confused about the whole sequence parallelism thing, can you maybe just give a little bit more intuition for the problem and how this is even conceptually possible to solve? Because Again, a naive understanding would be like, the Mamba block is sequential. So how are we ever going to get around that? Isn't it generating one token at a time? And like, how do we change the perspective or change the frame on the computation to go from the naive understanding to at least opening up the space of possibility for parallelism in that respect?

Sure. So like sequence parallelism is purely a necessary evil to get around memory constraints. So even though Mamba is generating one token at a time, it needs to store somewhere the activations and just the input sequence of a million tokens. And that is, even if you only require a constant hidden state, yeah, activations plus the actual input sequence are just too large to fit on even an H100 at training time. Training time is even worse because you have gradients and optimizer states those have to be in higher precision, for example, you can't fit it. And if you don't train on a longer sequence length, then your model doesn't know how to generalize to a longer sequence length, then you have accuracy problems. So at training time, there will definitely be less sequence parallelism needed just because there's less pressure on your device memory. But at training time, those are the reasons why we need it. How it actually looks is the thing that, like, we're still cracking, but that's the reason we require memory.

Got it.

That's all.

Yeah.

Okay. We will stay tuned for Zomba3.

For that. Absolutely. Yeah. Yeah.

I guess just one last breakdown, if you could go step by step. The Zomba2 architecture, it's not just, hey, let's swap in the SSD algorithm rather than the original SSM and let's call it a day. You're making these other changes. It makes sense that if you're going to make the new architecture, hey, let's update a few other things. If you could talk about exactly what you've updated, the LoRA modules on the MLPs, and then what you've done with attention and collectively all of that together. I think the highlight result from the Zamba 2 release is it's sort of at this Pareto frontier of efficiency and performance in certain respects against the other open models.

Mm-hmm.

So with your intuition, You already said it. I think you had 3 trillion training tokens rather than just 1. So a lot goes into this performance boost, but yeah, as well as architecturally, what do you feel is giving the uplift, the most significant uplift for Zamba 2?

Yeah. Okay. So it's different per model. So let's start with the 1.2B, let's say. So the 1.2B says there is still this global transformer block of attention plus MLP. But if you only have one block, there is correlations across depth, but maybe we want a little bit of more representation. So that's why we have LoRAs on both the MLP and the attention block. And these do different things. So remember, I meant like attention is correlated across depth, but MLPs are not. So we wanted that MLP to do a bit more heavy lifting. We thought there was still room for it to specialize across depth, which is why we introduced that LoRA. I think that one does a bit more heavy lifting than the attention. LoRA, which is why we only have the MLP-LoRA on the 2.7B and the 7B Mamba-2 models. We stick with a single global block instead of the ABAB for the 1.2B just because we were parameter constrained. We thought those parameters were better used on Mamba blocks. That's the 1.2B. So that's— the 1.2B is really good if you have one specific thing you want the model to do. 1.2B is enough for summarization on its own, and you want it to just 4-bit quantize, and you want it to be very low overhead on any Raspberry Pi or Jetson Nano in the world. The 2.7B has this ABAB, so now there's 2 global blocks. Only LoRAs on the MLP for the reason I mentioned, that it specializes more across depth, and we wanted to kind of, for a very small overhead in FLOPS and params, we could get a bit more expressivity. That was also when we put in the rotary, the positional encodings. That was just because we got like the slight boost, like I was mentioning, along context. Whether it was a mistake or not in terms of context length extension is still to be determined. That was the main change though for the 2.7B. And then the 7B is the exact same architecture as the 2.7B. So we just scaled that one up. But the main differences from Zamba 1 to Zamba 2, if I were to summarize it, is that we were trying to really investigate The depth relationship with the shared block and try and find out how much rep do you want the model to be able to be flexible across depth? What is actually being flexible? Is it the MLP or the attention? And then the other 2 pillars of our training stack, right? Like our data, we improved a lot and then scaled it way up. So we doubled the size of our annealing set. We tripled the size of our phase 1. This helps the model like be more, more broad in what it's good at. So first model was a bit brittle. In terms of what, like, the trunk I was talking about was not super wide, but on 3 trillion tokens, like, we can anneal and fine-tune the model to be much more flexible across a lot of different tasks that our customers and partners wanted. Because just because it's seen more data, like, if you can look at this actually in the weights themselves, if you look at the Llama 3 weights, you can see they're much more evenly distributed across the representation. You've filled out that weight more. You can, and you can see it. And going from 1 trillion to 3 trillion, we were able to get most of the way there as well. The model's much more flexible. But yeah, just improving end-to-end the stack of our model production, but with a focus on depth, was Zamba 1 to Zamba 2. Yeah.

And I think you mentioned on the first read of the blog post, I didn't see the state size of the Mamba blocks in Zamba 2.

64. So yeah, all of those details are in our Hugging Face.

Excellent. Okay. Yeah, I didn't check out the model card yet. So 64, that's interesting. So you could have trained larger state sizes. Of course, inference efficiency being governed by state size. Could you talk through experimentally how your team landed at 64? And is it actually a little bit lower than my expectation because the attention blocks are doing so much heavy lifting? There's less pressure for the state size to be larger for The sake of memory.

This is pretty much right. So if I recall, larger states— larger state sizes is mostly for improving your long context ability. For, uh, like you can more accurately like store longer range states. And we have attention, so that's— it's a less clear benefit for us when we try to increase it. We've tested this, but it's a very clear damage at inference time. So 64 we found was enough for our ZomboBase models. Yeah, just empirically. And this intuitively makes sense for the reasons I just mentioned.

Yeah, fascinating. That makes total sense. I guess if we could just go deep on— I haven't seen that many papers or architectures where we're applying LoRA projectors to MLP blocks. It like intuitively makes sense that, wow, okay, a really certainly memory efficient way because you're not storing a whole nother set of weights to add some specialization into the MLP blocks. It sort of like is in alignment with that sort of performance per parameter, as you know, idea. I guess if you could just talk about that idea, what you've seen like in the ablations or the experiments you've run, it's a very interesting feature.

Yeah, it's going to boil down to what I was mentioning earlier of like the MLP specifically definitely are not correlated across depth. Very strongly. So this even boils down to what we've learned with CNNs, where like the first few layers learn very general representations, and then the deeper layers learn more specific representations. In the CNN world, this is a face, this face is a child, this child is crying. You get more and more abstract and specific as you go deeper. And in the language world, this is all manifested in the MLPs. So it's important that the MLPs learn the same notes contextualize themselves across depth. So that's why attention is learning that anyway. Oh, remind the MLPs that these are all the notes that we have across all tokens of the sequence. But those MLPs are like learning different deep relationships across depth. But all of this to say, that is the intuitive reason why we have MLPs, or why the MLP that we were carrying around anyway in this shared block should probably have some expressivity across depth. And a LoRA is a good way to get some of that expression, but with lower memory and compute costs, which is really what we want. If I was not bound by memory and compute, then I would just probably have independent MLPs and a shared attention. But this is a good middle ground. And then empirically, it improved accuracy across depth. So we stuck with it.

And is it the kind of thing where you've, you also ran the experiment with independent MLPs or you just know that's not feasible for the constraints you want for the Zamba model, so you don't even bother running that ablation?

Yeah, unfortunately that's just not on the table. Like, that model is not in the cards for us. So I would hope someone does it, but yeah, ideally I'm sure we would get some pretty big boosts from doing that. And if I had more parameter budget, I would go to full MLPs.

Yeah, yeah, makes sense.

But the specific ablation of no MLP LoRA to MLP LoRA is worth the jump in params and flops for sure, in terms of loss.

Very good. Interesting. I guess a little bit more of a general question is, if you read the literature, it seems like almost every time someone tries some type of Mamba attention hybrid, it's like this unreasonable effectiveness, at least across the metrics that they're measuring. We talked about how loss can be misleading. That's not always a measure of how well that model will actually perform, but it maybe tells you at least something. And then you can go down the list, you can look at other evaluations. You know, it's not surprising to see that a hybrid Mamba hybrid model would beat Mamba by itself. If like you think about the limitations of a fixed state size, the way information would propagate through that network. But it's always surprising to me when I see—

Hmm.

All other factors held equal, that it's beating something that has, that's like a full transform with full attention blocks. Do you have any insight around that, how these dynamics are playing together? Almost like maybe it's greater than the sum of its parts?

I really think it's powerful when you're doing hybrids to boil down, to almost apply qualitative, track qualitative differences in how well your model is encoding cross-sequence dependencies and projecting them to mixed heads and to find deeper relationships from those notes of the cross-sequence dependencies. Mamba does both. MLPs just does taking the notes and finding deeper relationships, like they're projecting onto some abstract feature space. And then attention is just finding the dependencies themselves and not doing much else with them. At least in terms of how we approach things at Cyphra, it's really just boiling down to this critical intuition of things. And for example, that's how we can explain why Black Mamba had struggled with the long sequence because Mamba notes on themselves about sequence dependencies are not rich enough. So you need some attention. The unreasonable effectiveness comes from reaching like the critical point of where everything is balanced, where you are not spending all of your flops in memory and everything on taking notes and then doing nothing with them, like with tiny MLPs or something and big attentions. It's not with like a weaker cross-sequence relationship from Mamba and then big MLP projections where you're like, trying to find really deep things about very minimal notes. I really just— if you find a balance of these 2, that's when you get really strong model architectures. And eventually I think we'll probably converge to like a better way to do both, where maybe we boil down attention to like— maybe you can have a retrieval head on a MambaBlock, whatever that may look like, to help you patch the problem with long sequence dependencies. Maybe there's some unified block that finds the exact mixture of projections and like note-taking across sequence dependencies, and you just scale it up and you're done. That seems to be the case, right, based on this intuition that you just need a balance. But yeah, I would still tie it back to you need balance between these 2 effects, and these 2 effects are what you need for language modeling. It differs per modality.

Yeah, very interesting. Do you get the sense— Nathan alluded to this— we, we see several iterations of, of Mamba hybrids out there. But at the same time, Mm-hmm. it seems like people are, like there's some lock-in maybe of even some of the private labs. Are there lock-in effects once you commit to a certain type of architecture where even though something begins to be promising, you almost ignore it purposefully to some degree just because you've invested so many resources in a particular path? And do those lock-in effects make it more unlikely that a lab that's been committed to a completely different tech tree for years is going to just jump ship and start— are those lock-in effects there?

The lock-in effects are stronger here than in most places. One thing is that as your lab continues, you have more and more ablations that you can look back on. And when you jump ship to a new architecture, all of those ablations, you have to retake a bunch of stuff, right? Um, another thing is that we were talking about evaluations. It's really hard to evaluate specific tasks, um, on very few tokens, uh, meaning that, for example, if you have 2 architectures, um, how are you going to test how well they do on MMLU, on in-context learning? You have to train for hundreds of billions of tokens, remember? And I can't just take a pre-trained checkpoint and then test both architectures because you have to train these architectures from scratch. So to predict how well a model is going to do on those emergent evaluations is a real leap of faith on new architectures that a lot of people in labs just don't want to do. Like I mentioned, LLaMA CPP and Ollama, the whole hacker space and every serving framework is built around transformers. It's slowly changing to support hybrids and SSMs and such, but do I really want to relearn fine-tuning? Do I really want to relearn efficient CPU inference or what efficient set of fused kernels will work at inference time? Do I really want to get it in VLM and TensorRT? We're going through that pain, but I can definitely see why a big company with even more technical debt and more ablations and interpretability is another one. Big labs have really finely honed interpretability, especially at somewhere like Anthropic or OpenAI. They can interpret really well what these models are doing under the hood, way beyond my intuitions that I'm— that I have right now. And if you jump into a new model architecture, you start from scratch with all of that in some cases. I'm sure there's some— an attention block might act similarly when it's in a hybrid versus when it's in a dense transformer, but there's still like a lot of resistance there to the new idea. But yeah, totally in agreement with your point. These are all the examples of like why I would probably myself not want to jump into hybrids right away if I was big tech or if I had a big existing presence in dense transformers. But yeah, there's a reason that they're moving slowly.

Yeah, that makes sense in that way. Do you see— so the lock-in effects are a way in which you commit, but then obviously the longer you commit, you're beginning to build out not just code but best practices where if you could imagine the possibility of Zyphra developing this architecture over the next 5+ years or longer? What kinds of— so I guess the positive side of lock-in effects is that you get really good at managing this architecture that maybe not that many others are managing. Do you have any insights as to what that could look like for you and Zyphra and the hybrid world?

There's some first mover advantage here where we can kind of steer the story. Like, we have like internal inferencing frameworks that we've already ported that if you want a Zamba model, we can run it for you and with you the fastest compared to everyone else. Um, if you want like community engagement, this is a negative. I think we're nimble enough to avoid the pitfalls of being locked in. So we have some very non-Zamba architectures that are cooking right now. As long as they are efficient on device and punch way above their weight. We're going to look at whatever that architecture is and we're going to try and find creative ways to make it train well. In terms of avoiding lock-in with that, it really just comes down to effort of you need to figure out fine-tuning again. You need to figure out sequence parallelism for every new architecture you make and this kind of thing. And putting the effort in is our way forward. Like, we don't have any super crazy intuition behind it. You just have to grind it out.

Nice. Yeah, that makes sense. I guess related to that is, and this might be part speculation, so I apologize for that, but it seems like you do have some real-world data given that Cypher is developing not just one flagship Mamba hybrid model, but various parameter counts, various sizes. Any sense in which you're getting sense of the scaling properties of this kind of model architecture, how those scaling properties are, you know, changing? You have Chinchilla scaling laws, like we know what it looks like for just base Transformers. Do you have any sense of what the scaling laws could look like for the Mamba hybrid models? Are they even that much different than what we've seen in Transformers?

I think they're pretty similar. We just have a slightly better slope for a lot of these. I would say I'm curious to see how shared blocks affects Chinchilla because we reach Chinchilla optimality way faster on the shared blocks. So if we were— if we train for 15 trillion tokens, at what point does the attention block saturate? It's kind of a similar thing to like, no one I don't think has really cracked scaling laws for MOEs. A similar— you can think of MOEs as a similar case where some blocks are getting trained more than others. I really don't see any reason why a Zamba architecture wouldn't scale well to 20B, 40B, 70B, and I think we're— we at some point will try this ourselves to get distillation in-house. But those are my only hot takes. I think they'll look pretty similar, just a little bit better.

That's awesome. This is maybe a big picture question. What does it take to do foundational ML research? And in particular, what is the kind of approach you and your team have? Because Zephyr seems fairly interesting. It's a company, but it also seems to be a lab doing what I would consider pretty clearly research. So it's not necessarily product-focused solely, but there's some end-commercial goals. Like, what is the process? Is it a lot of open-ended search in the team, or are there always specific objectives with the commercial interests in mind?

We have research guided by specific objectives. So those objectives are efficient inference is one. So You get a big first mover's advantage if you come up with a novel architecture that actually runs efficiently on device, because now everyone else either has to use your model or figure out all of the arcane knobs that you had to turn to get your model to work. We're sort of an applied research lab. We're not like making Mamba 3, right? We see all of the available blocks and training processes and LoRAs. Like every element of this Zamba model is built by someone else. But in order to recognize what all of those knobs do and then to turn them in a way that you get the same modeling performance as a dense transformer, but with a much more amenable on-device performance and even cloud performance, is the research that we're doing. So how do you take everyone else's work and how do you put it together in a way that's perfect at inference time? Yeah. And then this ties into the product story of if you have that inference-efficient, in-training-efficient model, It's way easier. It is now achievable to make it personalized to everybody. It's achievable to put it on a phone. It's achievable to personalize it to enterprise applications and these sort of things. It's definitely guided by product.

Yeah, that makes sense. And this is just a side question. The, you know, when I think about the properties that LLaMA has, but in particular, these hybrid models and I've had increasing interest in robotics. I feel like that could be the next big market. It's not really a market right now.

Yeah.

And what you need to run models locally, that just ties in almost directly with what your lab is focusing on. Any sense— obviously robotics usually involves completely different stack of data, so it's multimodal. You have things like the ACT Transformer. It'd be interesting to know what that would look like for Mamba. If you were to apply your work, which is primarily focused on NLP, how big of a jump is it to now be looking at like the robotics technical use cases? Or is that a bigger jump than what I'm thinking in my head? Do y'all ever think about things like that?

Oh, definitely. So robotics, gaming, there's all kinds of markets where it's really nice to have flexible models on device. For robotics specifically, to make— if it's just Like chat with humans and maybe ability to do visual question answering about like how to recognize and respond what the robot is seeing and speaking. That's pretty easy. That's pretty achievable. And I think we probably have the best offering for that right now. Once we get voice-to-voice done, then yeah, a robot walking around that can talk to you and is totally on devices is like nothing for us. What else would be required? It would all be very low power as well, so you could get it running on a Raspberry Pi or a Jetson. I think there's a lot of cases with like smart cities, with— you could put this on a Tesla and you can put in a self-driving car and just have a chat interface. Uh, you run the windshields, speed them up, run the windshields at the speed that I like. All these sort of things become possible when you have a model architecture like Zamba that's multimodal and voice-to-voice. And that's why we're pushing for it so hard.

Yeah, that's awesome. I guess just one more practical question is oftentimes, really, if you look at the adoption of models, it often comes down to just like how easy are they to use? And that often involves like how compatible are they with the Transformers library from Hugging Face? But I guess if you could just talk with Talk about that. If I wanted to do some really interesting experiments with Zamba2 right now, or I wanted to— whether that was fine-tunes, whether that was like trying to do some extended pre-training stuff, or whether, you know, how easy it is to use— what do I actually do? What libraries do I need to get into to use Zamba?

So right now you're restricted to Hugging Face Transformers. So we have Zombo1 has been upstreamed now. It's in Transformers. You can use it with all Accelerate and PEFT and all of the Hugging Face, like, class of frameworks. For Zombo2, we have our own fork of Transformers. It's currently getting upstreamed, but there's always just a lot of compatibility work with this kind of stuff. You can get, like, basic functionality right away, which is why we have our own fork, but getting it really solidly ingrained in the entire Hugging Face ecosystem Is the hard leg of the journey that we're still working on. We also have a pure PyTorch implementation for both Zamba 1 and Zamba 2. So if people have their own custom pipelines, they want to port it to their own internal inference framework, we tend to find people like that a lot. It's also more performant than putting it in like the Hugging Face ecosystem to just have pure PyTorch. You can compile it with torch.compile. You can create custom Triton kernels very easily. Going forward, I think we're about to finally like become more ingrained because we're about ready to crack GGML. And Llama CPP and Ollama and that entire ecosystem. We've pretty much finished porting Zamba2. I don't know, I struggle to put a date on it. I don't want to miss a date, but like in the next couple weeks, I think people will be actually— will be able to use this seamlessly with VLM, Ollama, Llama CPP, GGML, the entire ecosystem. I think we're right about to jump onto.

Very interesting. So the choice to fork Transformers, simply we'd like people to be able to use it immediately, and there's still some Kinks to work out. Yeah, to get it integrated. That's essentially it, right?

Yeah. And yeah, we understand that like people will balk, uh, balk when they see, oh, it's a fork, done, I'm not jumping into it. And even like, we've produced the best model at the weight class, so there's definitely incentive for those who want to use it. You got to use the fork. But for all of those who are more focused on what is the fastest way to get this running on my laptop or something, You will very soon be able to actually do that. Very soon. But currently, Hugging Face fork or pure PyTorch are your options.

And you said, like, what would be the most performant, both for inference and for, let's say, I wanted to, like, do fine-tuning or training? If I use one of those training frameworks, is it, like, fully optimized, or am I— are there some trade-offs with having to make sacrifices with getting compatibility? Or should I go straight to PyTorch, or Could I use those frameworks? It's not going to come with costs.

Yeah, so for small scale, like by small scale I mean like fine-tunes. So if you're at the pre-training scale, you've got to have your own thing. You've got to have your own GPT Neo X or Megatron DeepSpeed or MegatronLM. You've got to have a pre-training framework, and those are not super Mamba-friendly right now. We have our stuff is in-house unfortunately right now. It's just the name of the game. How people watch your PRs otherwise. In terms of fine-tuning, at the fine-tuning scale, Hugging Face is totally fine. We even use Hugging Face internally to do some of our fine-tunes, some of our context-length extensions. When you're working on sub, you know, a couple billion tokens to 10 billion and below, 50 billion below, Hugging Face is totally fine, and we recommend people use that. So yeah. If I had like a mid-range kind of thing, maybe I would try to wrap the pure PyTorch implementation of Zamba in some like simple data loader and just try and get that going if I have maybe 50 billion or something, 50 to 100 billion range. But I think most people fall into the fine-tuning range where Hugging Face suffices.

Very good, makes sense. Excellent. And if people today, like right now, wanted to start experimenting fine-tuning, I believe there's a GitHub linked to the blog. There I can access the fork of the Transformers library. If someone was motivated, they could spin something up today. Is that correct?

Oh yeah, definitely. Yeah, it's not hard to fork and to just build our fork and load our checkpoint. Yeah. And then if you just want to talk to the model, then we had a day zero release with NVIDIA as an inference microservice. We self-hosted it on like the Maya endpoint on our website. So if you just want to talk to the model, use one of those. If you want to spin it up, you can do it today with our fork. Excellent.

So I guess we've been talking a lot about LLaMA and the Zamba models, but Cypher also published a very interesting paper recently, the Tree Attention paper, and we've barely talked about it. I've not been able to dedicate myself to a complete deep dive, but my high-level view is it's very interesting in the sense that it seems to be in some aspects improving upon Ring Attention, which was the huge thing earlier this year where it seemed like everyone was training million-plus context length LLaMA models. If you could just talk about tree attention, how is it different from ring attention? And then what impact can it have? What is it really pushing on?

Sure. So ring attention is good when you have a network topology that's a bit more flat. So like a mesh topology, like with TPUs and stuff, I think was what it was designed more for. And that's because you have a ring with point-to-point across each GPU in the ring. If you're in GPU land, you have 2-level topologies where on a single node you have NVLink that connects all of your GPUs, is really high bandwidth. But across nodes you have a much slower InfiniBand, or maybe you have a RoCE, like a converged Ethernet thing. That's an order of magnitude lower bandwidth than the intra-node interconnect. So there's a reason why they only did up to 2 nodes in the original ring attention paper. And it's because you're bottlenecked by this cross-node link. So the boundary GPUs at the cross-node are going to be dictating how quick your inference speed is, and it's not great. As you scale up to more and more nodes, this becomes worse and worse. If you have smaller PV states, then you're going to be— the compute is not enough to overlap that expensive communication. So it's not scalable. The whole point of tree attention is to recognize this fact that 2-level topologies are not good for rings. And to reformulate the distributed attention operation as an energy function, which allows you to use all-reduce operations, communication operations, instead of point-to-point. All-reduce operations in existing communication libraries— NVIDIA's communication library called NCCL, NVIDIA Collective Communication Library, AMD has something like RCL, the Rockham Communication Library— all-reduces themselves are topology-aware. For GPU clusters. So they account for this 2-level topology, and they also add some of the computation is like in-network. So with an all-reduce operation, you can do some of the sum in the network card, and you can overlap that sum computation, that kernel, be it on the GPU or on the card if you're doing SHARP, with the communication of those qubit states. So basically it can scale. In short, it can scale on 2-level topologies because you're reformulating to an all-reduce Which is inherently topology aware and adds some of the compute, is easier to overlap versus point-to-point operations, which are restricted and sort of assume a flat topology of network.

So I guess if you could help me, maybe let's begin with the end in mind. Like, what will that enable me to do that ring attention would struggle with? What, in what cases will tree attention help me?

So. This all boils down to parallelism. Whenever you have parallelism, it boils down to memory. GPUs have a limited memory. If you have long contexts, so say if I want to do a million contexts on a Llama 8B, or if I need a lot of parameters that I need to store in memory at inference time, then I need more GPUs. For both of those cases I mentioned, with super long contexts, with super large models, you need lots of GPUs to fit, more GPUs than are on a node. So you're going to 2 nodes, 4 nodes, et cetera. When you're going across nodes, you have a 2-level topology, and all-reduces work better. So if you wanted a 16-node context extension training run with 2 million contexts, tree attention will work much better. If you have an inference run of LLaMA 4.0 5B on older GPUs, you need 4 nodes or 2 nodes. Tree attention will scale much better on that. Anytime you're scaling across nodes for GPU clusters or any imbalance of communication costs, tree attention will win.

Interesting. Okay. So this is not necessarily something where we had ring attention and all of a sudden everyone's just going to implement tree attention. It's on a case-by-case basis of which one you might use if you were to try to do multi-node training and, you know, distribute the memory.

Yeah, so there are cases where like if you perfectly have enough compute and your communication costs, like those 2 perfectly overlap one another and it's effectively free and you're on a single node, ring attention is pretty good. Even though all reduces are really closely optimized because they're used in training so much, you can't really do much better than that in free communication. It's totally overlapped. Yeah, not everyone will benefit right away, but Those cases where you have a smaller model, longer context, which translates like smaller KV states, that means you're going to be communication bottlenecked. There's not much compute to hide the communication. So those people will benefit from us, and then the people with really large models that are inherently across nodes will benefit from us. Yeah, but yeah, you're right, it's not—

Not everyone's case. Trying to— for like a very, very long training on very, very long sequences, it would help for that if that was a Only if it was a multi-node training setup, or it doesn't come into play?

It would help more. It would help more. Even on a single node, all-reduces are a little bit more effective than point-to-point, just because tuning libraries, or collective libraries today are really closely tuned for that specific operation. We're kind of exploiting the fact that all-reduce is really good than point-to-point, because you don't really use point-to-point much in training unless you're doing pipeline parallelism, which is less common than tensor parallelism or data parallelism, which are both all-reduce. So yeah, again, it'll help you a lot more at scale. It might help you a little bit at smaller scales, but I wouldn't promise big speedups.

Okay. No, that, that's very interesting. I guess just because we hit that topic, I have a— I thought of another question. The Zamba-1, I believe, was natively trained with, I think if I'm reading this correctly, 4,096, uh, token context length. Did Zamba 2 train at the same native token context length, or is that higher?

Yeah, yeah. And then we were able to extend to the 16K range for free just by interpolating, uh, rotary positional embeddings. Getting beyond that is an open question that we're still working on. So I think we have a path now for 64K and 100K, which should be coming out soon, especially for the smaller models, because remember I mentioned that it's a memory overhead thing at the end of the day, right? So like Zamba 2. 7B, we can go to way longer contexts than this, but just the engineering challenges to get us to a million contexts, still working on that. But yeah, trained on 4K.

Yeah. So I guess just because it's something I've experimented with, Mamba 1, like really basically the idea of, okay, we're natively trained at even 2,000 token context length, but I want something capable of much more. So I'm just going to do continuous pre-training and now all of a sudden my data distribution Mm-hmm. If you want to do it smart, maybe there's a curriculum learning approach where you don't just throw that much more context length at the training at once. So you want to do it smart, but it seems there are all these papers and my own experimentation is it works pretty well to do continuous pre-training to expand context length. To what extent is that a solution if somebody wanted to grab, let's say, the Zamba model, or we even see this with some of the just-a-transformer LLaMA models? Do you have any thoughts on continuous pre-training to extend context length? Is that sort of like a hack that you're just not going to get the quality that a natively trained model would have?

No, you can get the quality. That's what we're doing ourselves. So yeah, definitely continue to pre-train. You definitely need to do a curriculum where you're warming up the sequence length and batch size across your continual pre-training run, but eventually you're going to hit a memory wall. And then that's really it. That's all I'm trying to say here is that you gotta finish up sequence parallelism for the MambaBlocks as well if you want to get past that memory wall, but it totally works.

Yeah, absolutely. So playing around with the 2.8 billion parameter model, I could never figure out sequence parallelism. I, I use the DeepSpeed library because it just makes it easy to do all the 0-3 stuff. Of course, you have to pay tremendous costs with training efficiency when you start offloading everything to CPU memory or Yeah, so sequence parallelism is the big unlock. That's going to be exciting. Okay, cool. Does Xiphra have plans to open source that once that gets cracked?

Sequence parallelism for Mamba will definitely be open sourced. The, the Xiphra-specific training stack that we're using to produce Mamba will probably not be open source, but the general process and like what we did with tree attention, the kernels that we use, the scripts that we use, those will all be open source because we want the community to like move towards Mamba and see that there's something there.

Yeah, awesome. Great to hear. Anything else that we didn't cover about Zamba, the timing as in around the Zamba 2 release? Anything else you think we should cover before we wrap?

No, I think we've covered pretty much everything.

Appreciate your time, Quentin. It's absolutely great speaking with you.

Likewise.

So let's do this zoomed out one. I think this has been fantastic. I think that this is for the real ones out there who are, you know, very interested in learning more about what it actually takes to make these models work and just how many little nitty-gritty details go into that. So I've learned from this. I think a lot of people will appreciate all of the lessons that you've shared from many long days and probably some long nights working on this stuff. I think it's been great. Let's do the zoomed out thing with all that in mind.

Mm-hmm.

Where is Zypher going as a company? What role do you guys want to play in the lives of users? You talked a little bit about how you see yourselves competitively against the big guys, but I'd love to just hear on whatever timescale, you know, what you think my AI-assisted life is going to look like and what role Zyphra and Zyphra models will play in that.

Yeah, definitely. We've focused really hard on producing the best models for one modality, and now we're ready to expand out. So this includes other modalities. So we want like visual question answering. I want to be able to live edit my pictures. I want to be able to talk voice to voice, like I talk to you voice to voice. I want to talk to my AI voice to voice as well. This includes like actually deploying personalizability to people. This includes like actually launching like a Maya both enterprise and on device for consumers. This includes the broader ecosystem, for example, Ollama and Llama CPP and stuff. Those are all transformer based and getting the hybrid architectures into those is One of our next priorities so that everyone can actually deploy the model really quickly instead of using our fork of Hugging Face and all this sort of thing, and having a much more like a better user experience for hackers and for consumers. And then there's some also just some higher level things we think like memory. So retrieval is really important for personalizability that we're going to like, how does retrieval interplay with long context reinforcement learning so that like we can really tie down, it's not just continual pre-training when your phone, like, we definitely need some sort of approximation at least of RLHF or RL-AIF for users so that we can really extract as much as possible from the little data that we're going to get per user. There's a lot of ways to expand, but those are some of the main ones. Cool.

I love it. For now, I'll say Jason Mo and Quintin Anthony from Zyphra, thank you both for being part of the Cognitive Revolution.

Thanks, Hasan.

Thanks.

It is both energizing and enlightening to hear why people listen and learn what they value about the show. So please don't hesitate to reach out via email at tcr@turpentine.co, or you can DM me on the social media platform of your choice.