gtcode
Porting Sakana AI's TRINITY Qwen-based model to Elixir/Bumblebee/Nx/Axon
Aloha gang,
I’m working on a port of Sakana AI’s TRINITY, an evolved LLM coordinator:
TRINITY Paper
OpenReview
Downloadable Assets
I started by attempting to reconstruct the work itself, but that isn’t realistic for me, given skill/resource constraints. So I’ve instead pivoted to porting their Python mechanism that uses a base Qwen model to build their coordinator: Trinity Coordinator. (Right now it’s been deconstructed to use a local path-dep in mix.exs related to an inference library I’m building to generalize abstracting LLM providers, thus not “clone friendly” yet.)
Just seeing if this is of interest to anyone. Certainly open to input/feedback/ideas/critiques on approach. Please respond here or open an issue with your candid feedback. There must be someone out there with more knowledge/experience on such matters who can provide guidance?
I’ve created the safetensors file from the original python scripts, so nx can talk numbers properly. I’ve been working on a staged verification process so that the resulting coordinator based on Qwen will behave the same as the generated .pt file from their Python system. There are some nuances related to numpy -> nx and others that might prevent perfect alignment but I’m aiming for behavioral and functional parity.
One thing I see often these days is people creating amazing work and ideas in Elixir, but often hard coupled to providers and the like. One goal for trinity_coordinator is to have a working standalone system with built in routing to LLM’s, but also pluggable/modular for integration into any other codebase/framework/system.
ps: I wasn’t sure if this is the right forum category, but there was a note that said to use the nx forum if it’s nx related.
Most Liked
polvalente
@gtcode I just shipped some improvements to EMLX in the past few days focusing on performance.
I took inspiration from what @ausimian did on his backend with the lowering compiler, and the final PR of the series has a benchmark in the description: feat: add fused kv_cache+sdpa by polvalente · Pull Request #124 · elixir-nx/emlx · GitHub
Our custom Qwen3 implementation now reaches 300 tok/s on 0.6B 4bit and the emlx_axon bumblebee rewrites can reach 120 tok/s on the non-quatized 0.6B! Tomorrow I’ll release a new emlx/emlx_axon version with these improvements. Improvements are available in emlx/emlx_axon 0.4!
ausimian
Just for fun, I ported this to work on my 24GB M4 MBP using the Emily backend. In doing so, during the export, I ran into a limitation of the current native mlx libraries - their svd functions have no ‘thin’ mode and always materialise the full matrix. For the Qwen embedder that turned out to be ~92GB of memory.
I updated Emily to support this mode (in specific cases) directly via the Gram matrix, could do the one-time export in ~2s and was able to run the qwen router example.
If you are interested, the changes I made are here
Popular in Questions
Other popular topics
Categories:
Sub Categories:
Forums
Popular Tags
- #ecto
- #liveview
- #troubleshooting
- #learning-elixir
- #deployment
- #library
- #erlang
- #testing
- #genserver
- #mix
- #absinthe
- #remote-other
- #otp
- #plug
- #how-to-question
- #macros
- #postgres
- #channels
- #elixirconf
- #exunit
- #discussion
- #code-sync
- #javascript
- #podcasts
- #onsite
- #dialyzer
- #docker
- #authentication
- #umbrella
- #full-time-contract
- #podcasts-by-brainlid
- #ecto-query
- #elixir-ls
- #phoenix_html
- #iex
- #blog-post
- #graphql
- #genstage
- #ai
- #websockets
- #supervisor
- #advent-of-code
- #elixirconf-us
- #distillery
- #processes
- #forms
- #api
- #metaprogramming
- #security
- #performance










