OpenArch – PyTorch implementations of modern LLM architectures
github.com97 points by anuj0456 7 hours ago
97 points by anuj0456 7 hours ago
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.
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".
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).
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.
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.
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.
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?
One thing i am hoping can make easier understanding how architectural changes trasnlate into actual implementation and performance
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.
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.
[flagged]