OpenArch – PyTorch implementations of modern LLM architectures
github.com
github.com
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.
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?
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/
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.
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).
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.
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.
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
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.
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!
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?