I have been studying modern LLM architectures and started implementing them from scratch in PyTorch to better understand the design choices behind each model.
OpenArch is a collection of these implementations, including Llama, Qwen, DeepSeek, Gemma, Kimi, GPT-OSS and others.
The goal is to keep the code readable and useful as a reference when going from the paper to an actual implementation.
Would be interested in feedback from people working on model architecture and training.
I was yesterday years old when I learned that those open weight models need custom code to run.
Somehow I expected inference engines are generic LLM runtimes that can execute any weight.
So, to get this right.
Someone trains a model.
They release the weights and a reference implementation of the model architecture.
Then a provider has to host this model either by running inference via the reference implementation, an open source implementation, or build their own.
Does this mean, providers don't just differ in quantisation and configuration, but also in inference engine implementation?
The implementation of a language model is usually good enough for running _a single forward pass_ through the model, but to host it via an inference engine you typically need to convert a few operations. For instance the MLP can be easily split across GPUs (tensor parallelism) and MoEs also have a way of parallelizing.
Most of it is pretty standard, since not that many different layers and primitives are used in LLM architectures, but once in a while something new comes along that needs more effort. MoEs are one example, they are sparse and allow for completely different inference patterns, which takes a while to figure out.
Last year I started a blogpost series about this topic (that I hope to update some time). I start from a minimal gpt implementation by Karpathy and build the engine around it, you might like it: https://pieter.ai/blog/2025/nanogpt-inference/
yes. this is just raw implementation of the model arch as described in papers. for complete model training with back propogation we need training pipeline with optmizer and loss calculation.
> Somehow I expected inference engines are generic LLM runtimes that can execute any weight.
In fact, this was close to be true until last year: almost every open model except DeepSeek had a very similar architecture that was pretty close to the GPT-2 one with very few variations on top (and sometimes an MoE architecture, which itself was a few year old at that point).
But a year ago there's been a cambrian explosion, first in attention mechanism but also in a bunch of other directions, mostly coming from China, and now there's a very massive diversity today's space.
Makes me feel better as a guy who tries to get O(1000) times smaller models working for scientific applications. Writing a backend that works for all the models people train is quite a chore, so the common refrain is "why don't we just do what industry does?" (he answer is that "industry" has billions of dollars). You also get "no one uses X backend any more" to which the reply is also "yes, but they have billions of dollars and a team of software engineers".
Go look at the github commits for llama.cpp and read the actual code they're adding to support new model/quant types... to me at least it's some serious black magic, and the sheer number of genius developers and activity level in this repo is absolutely wild.
Imagine what could happen if other open source projects had this level of engagement and expertise at hand and eager to contribute... instead of arguing over politics and making yet another fork of something.
Yes. https://inferencex.semianalysis.com provides some comparisons wrt. certain cost metrics. Some commonly used ones are vLLM, TensorRT-LLM, and SGLang. These three at least are open source, and all come with an Apache 2.0 license.
Some providers also have to implement their own engines, e.g., Cerebras has their own inference serving stack for their wafer-scale chips, as does Google for their TPUs (XLA compiler).
yes, most of them are similar. but implementation of GQA, MLA, mHC, Sliding Window changes the implementation drastically because of which the overall model effeciency changes.
Yes. You can see this affecting perf benchmarks as well. Usually the cheapest inference providers either use approximations like tanh instead of sigmoid, nvfp4 quantizarion, etc.
There was a post here the other day highlighting this by showing the benchmark perf of different I defence providers, it's a fantastic area to cheap out in, because you can never really tell if a model is 75% good or 83% good on some specific benchmark when you use it to build your own stuff
Absolutely. The worst part is that ML being probabilistic means the models still output something comprehensible meaning you can't determine if it's correct or not without a lot of scrutiny and engineering.
This is excellent for understanding. I'm having some trouble to get into understanding - pytorch is for me the RL which is used as gym/training. There I can chose ppo, dnq and other agents to perform some predefined actions in a predefined gym/world.
The repo you are showing - I really have problems to get it into RL understanding of mine. What's the gym? What are the agents. Can it be used to train that models with pytorch?
Sorry for the noob question. Papers are overwhelming my noob brain.
Your comment is a little unclear, so it‘s hard to parse your exact question. But it seems you are conflating 3 things, PyTorch, RL and Gym/Training (?).
– PyTorch is a framework which lets you define neural network models.
– RL is a collection of methods to train neural networks (change the network parameters to improve its performance).
– An RL-Gym is a framework to apply the neural networks to some problem. This lets you collect the data necessary to later use the methods of RL to train your model.
thx - yes, after reading myself what I have written, i need to ask for Sorry for writing it like it is :)
I think my problem is with the imagination and knowledge transfer:
As i remember, in the year ~2010 +/- 2-3y, machine learning became (again) POI for technology. I remember convolutional networks, deepQ, Genetic algorithms, etc being in the press. I remember Michael Schmidt, a biology student at that time (~roundabout~), being as "the influential data scientist of the year" at forbe's cover. MS did symbolic regression combined with genetic algorithm and founded nutopian that was developing this propriatary "eureqa" algorithm.
At that time, PyTorch was "created" by Facebook as an answer to google's tensorflow. I remember a guy trained on PyTorch a model that could play Nintendo's Mario. I wanted to understand how to do it, but it was toooo much for my brain. What i remember:
- one needs a gym that stands for the world the agent "lives in". The gym/world is data, like the input of a screenshot, or a gym/world for Stock Trading might have historical courses, prices etc.. also rewards, penalties are defined in the gym.
- the agent living in this gym/world have "actions" - left,right,buy, sell, ...
- RL training is then to let that agent "move through world/gym and perform some of its and agent's actions, where the outcome is rewarded or not. agent learns through rewards.."
Gym and RL training can be PyTorch or TensorFlow. In my understanding that gym is only a definition of what data is avaible, shape of data, etc. Then, the agent needs to be defined too - i remember Stable Baselines 3..
and then the RL training is just agent does x -> gym checks for outcome -> reward/penalty & adaption of weights -> inference -> repeat
This is cleary imaginable for me, even its not fully correct - there is an agent, there is a world, there is method to act&compare/data manipulation of weights - there is PyTorch that offers the tools to build a world, agent and to do the weight manipulations ..
OP have implementations of different Archs in PyTorch and my knowledge implodes. I cant imagine what is the "world/gym" in here? What is the data and actionables , whats the agent and whats its role/actions - and how would I use this implementations??
can I clone the repo, pick the one arch i want - and, when i would feed in data, after training, i will get the model that i can use for inference???
I definitely have a knot in my brain because of this!
Thank you for taking your time and trying to understand what my previous comment meant :) Are there some usefull sources for seeing how the learning is actually done (despite the tokenization) - and the skills training?? how are models trained to find the loophole that allows them to break out and reach their objectives "illegaly" - thats bugs my head since weeks!
This reminded me of lucidrains, perhaps one of the most prolific PyTorch users and paper implementers on the Internet. His implementations were crucial references to my own ability to boostrap my paper-implementing skills. Looks like he is still around and doing really awesome stuff. https://github.com/lucidrains
Had a question on the MoE: in kimmi-K2/model.py, the router does torch.topk(..., k=self.num_experts) while every other MoE uses k=self.top_k. The ctor's top_k=8 is never stored either.. Is that intentional dense routing or should that be self.top_k?
ur most welcome! if you're curious how I found it, I have made a code intelligence system. I literally fed it your repo and asked for bugs. You can try it yourself here: https://benzi.fly.dev/
This is really cool, great way to reinforce our understanding of model architectures! But how is the author confirming that these model architecture implementations are correct though? I don't see any details in the README.md.
hey, Building these from scratch in pure PyTorch is honestly the best way to deeply understand the paper details.
something better than simply implementing a traditional Transformer or GPT-2，As an individual maintainer, will be able to keep up with future model updates?
I think depending on your perspective it might be even better to eschew PyTorch and implement with pure NumPy. Sure it will be slow, but perhaps as a pedagogical exercise it is more worthwhile.
I have been studying modern LLM architectures and started implementing them from scratch in PyTorch to better understand the design choices behind each model.
OpenArch is a collection of these implementations, including Llama, Qwen, DeepSeek, Gemma, Kimi, GPT-OSS and others.
The goal is to keep the code readable and useful as a reference when going from the paper to an actual implementation.
Would be interested in feedback from people working on model architecture and training.
I was yesterday years old when I learned that those open weight models need custom code to run.
Somehow I expected inference engines are generic LLM runtimes that can execute any weight.
So, to get this right.
Someone trains a model.
They release the weights and a reference implementation of the model architecture.
Then a provider has to host this model either by running inference via the reference implementation, an open source implementation, or build their own.
Does this mean, providers don't just differ in quantisation and configuration, but also in inference engine implementation?
The implementation of a language model is usually good enough for running _a single forward pass_ through the model, but to host it via an inference engine you typically need to convert a few operations. For instance the MLP can be easily split across GPUs (tensor parallelism) and MoEs also have a way of parallelizing.
Most of it is pretty standard, since not that many different layers and primitives are used in LLM architectures, but once in a while something new comes along that needs more effort. MoEs are one example, they are sparse and allow for completely different inference patterns, which takes a while to figure out.
Last year I started a blogpost series about this topic (that I hope to update some time). I start from a minimal gpt implementation by Karpathy and build the engine around it, you might like it: https://pieter.ai/blog/2025/nanogpt-inference/
yes. this is just raw implementation of the model arch as described in papers. for complete model training with back propogation we need training pipeline with optmizer and loss calculation.
> Somehow I expected inference engines are generic LLM runtimes that can execute any weight.
In fact, this was close to be true until last year: almost every open model except DeepSeek had a very similar architecture that was pretty close to the GPT-2 one with very few variations on top (and sometimes an MoE architecture, which itself was a few year old at that point).
But a year ago there's been a cambrian explosion, first in attention mechanism but also in a bunch of other directions, mostly coming from China, and now there's a very massive diversity today's space.
yes, that is correct. after chinca came into picture the advancement in this field sky rockted
Makes me feel better as a guy who tries to get O(1000) times smaller models working for scientific applications. Writing a backend that works for all the models people train is quite a chore, so the common refrain is "why don't we just do what industry does?" (he answer is that "industry" has billions of dollars). You also get "no one uses X backend any more" to which the reply is also "yes, but they have billions of dollars and a team of software engineers".
Go look at the github commits for llama.cpp and read the actual code they're adding to support new model/quant types... to me at least it's some serious black magic, and the sheer number of genius developers and activity level in this repo is absolutely wild.
Imagine what could happen if other open source projects had this level of engagement and expertise at hand and eager to contribute... instead of arguing over politics and making yet another fork of something.
Yes. https://inferencex.semianalysis.com provides some comparisons wrt. certain cost metrics. Some commonly used ones are vLLM, TensorRT-LLM, and SGLang. These three at least are open source, and all come with an Apache 2.0 license.
Some providers also have to implement their own engines, e.g., Cerebras has their own inference serving stack for their wafer-scale chips, as does Google for their TPUs (XLA compiler).
Thanks for sharing
FYI Mistral at launch just dropped the weights without any model architecture mentioned.
Most of the OSS models follow the same architecture which is Llama +- a few things, so it wasn't too hard for people to make it work.
> Most of the OSS models follow the same architecture which is Llama +- a few things, so it wasn't too hard for people to make it work.
It used to be the case until last year, but now almost every Chinese model come with their own linear attention mechanism.
yes, most of them are similar. but implementation of GQA, MLA, mHC, Sliding Window changes the implementation drastically because of which the overall model effeciency changes.
Yes. You can see this affecting perf benchmarks as well. Usually the cheapest inference providers either use approximations like tanh instead of sigmoid, nvfp4 quantizarion, etc.
There was a post here the other day highlighting this by showing the benchmark perf of different I defence providers, it's a fantastic area to cheap out in, because you can never really tell if a model is 75% good or 83% good on some specific benchmark when you use it to build your own stuff
Absolutely. The worst part is that ML being probabilistic means the models still output something comprehensible meaning you can't determine if it's correct or not without a lot of scrutiny and engineering.
Hey anuj
This is excellent for understanding. I'm having some trouble to get into understanding - pytorch is for me the RL which is used as gym/training. There I can chose ppo, dnq and other agents to perform some predefined actions in a predefined gym/world.
The repo you are showing - I really have problems to get it into RL understanding of mine. What's the gym? What are the agents. Can it be used to train that models with pytorch?
Sorry for the noob question. Papers are overwhelming my noob brain.
Not the poster, but maybe I can help.
Your comment is a little unclear, so it‘s hard to parse your exact question. But it seems you are conflating 3 things, PyTorch, RL and Gym/Training (?).
– PyTorch is a framework which lets you define neural network models.
– RL is a collection of methods to train neural networks (change the network parameters to improve its performance).
– An RL-Gym is a framework to apply the neural networks to some problem. This lets you collect the data necessary to later use the methods of RL to train your model.
Thanks, Yes PyTorch is a framework largely used to create neural network models.
thx - yes, after reading myself what I have written, i need to ask for Sorry for writing it like it is :)
I think my problem is with the imagination and knowledge transfer:
As i remember, in the year ~2010 +/- 2-3y, machine learning became (again) POI for technology. I remember convolutional networks, deepQ, Genetic algorithms, etc being in the press. I remember Michael Schmidt, a biology student at that time (~roundabout~), being as "the influential data scientist of the year" at forbe's cover. MS did symbolic regression combined with genetic algorithm and founded nutopian that was developing this propriatary "eureqa" algorithm.
At that time, PyTorch was "created" by Facebook as an answer to google's tensorflow. I remember a guy trained on PyTorch a model that could play Nintendo's Mario. I wanted to understand how to do it, but it was toooo much for my brain. What i remember:
- one needs a gym that stands for the world the agent "lives in". The gym/world is data, like the input of a screenshot, or a gym/world for Stock Trading might have historical courses, prices etc.. also rewards, penalties are defined in the gym.
- the agent living in this gym/world have "actions" - left,right,buy, sell, ...
- RL training is then to let that agent "move through world/gym and perform some of its and agent's actions, where the outcome is rewarded or not. agent learns through rewards.."
Gym and RL training can be PyTorch or TensorFlow. In my understanding that gym is only a definition of what data is avaible, shape of data, etc. Then, the agent needs to be defined too - i remember Stable Baselines 3..
and then the RL training is just agent does x -> gym checks for outcome -> reward/penalty & adaption of weights -> inference -> repeat
This is cleary imaginable for me, even its not fully correct - there is an agent, there is a world, there is method to act&compare/data manipulation of weights - there is PyTorch that offers the tools to build a world, agent and to do the weight manipulations ..
OP have implementations of different Archs in PyTorch and my knowledge implodes. I cant imagine what is the "world/gym" in here? What is the data and actionables , whats the agent and whats its role/actions - and how would I use this implementations??
can I clone the repo, pick the one arch i want - and, when i would feed in data, after training, i will get the model that i can use for inference???
I definitely have a knot in my brain because of this!
Thank you for taking your time and trying to understand what my previous comment meant :) Are there some usefull sources for seeing how the learning is actually done (despite the tokenization) - and the skills training?? how are models trained to find the loophole that allows them to break out and reach their objectives "illegaly" - thats bugs my head since weeks!
This reminded me of lucidrains, perhaps one of the most prolific PyTorch users and paper implementers on the Internet. His implementations were crucial references to my own ability to boostrap my paper-implementing skills. Looks like he is still around and doing really awesome stuff. https://github.com/lucidrains
This is really cool!
Had a question on the MoE: in kimmi-K2/model.py, the router does torch.topk(..., k=self.num_experts) while every other MoE uses k=self.top_k. The ctor's top_k=8 is never stored either.. Is that intentional dense routing or should that be self.top_k?
thanks for pointing out. its a bug
ur most welcome! if you're curious how I found it, I have made a code intelligence system. I literally fed it your repo and asked for bugs. You can try it yourself here: https://benzi.fly.dev/
fixed
This is really cool, great way to reinforce our understanding of model architectures! But how is the author confirming that these model architecture implementations are correct though? I don't see any details in the README.md.
the only way i am aware of is to cross check with what deployed on huggingface/transformers repo
Really cool. I know the sauce is in the layers and weights, but its kind of mind blowing how short / simple the implementations are
One thing i am hoping can make easier understanding how architectural changes trasnlate into actual implementation and performance
i will see what i can add here. contributions are welcome
hey, Building these from scratch in pure PyTorch is honestly the best way to deeply understand the paper details. something better than simply implementing a traditional Transformer or GPT-2，As an individual maintainer, will be able to keep up with future model updates?
I think depending on your perspective it might be even better to eschew PyTorch and implement with pure NumPy. Sure it will be slow, but perhaps as a pedagogical exercise it is more worthwhile.
will try my best, but open for contribution.
[flagged]