gtcode

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

polvalente

Nx Core Team

@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

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

polvalente

polvalente

Nx Core Team

@gtcode Nx 0.12.1 and EMLX 0.3.1 have been released with the changes I had to apply to get your lib working without OOMs

Where Next?

Popular in Questions Top

hariharasudhan94
lets say i have a sample like a = 20; b = 10; if (a > b) do {:ok, "a"} end if (a < b) do {:ok, b} end if (a == b) do {:ok, "equa...
New
skosch
To my knowledge, put_in, Map.update etc. all have the one limitation of not automatically creating intermediate keys when needed (for exa...
New
lessless
I believe there are people here who are dealing with CSV files import on the daily basis, and since Excel is a really popular tool there ...
New
nsuchy
Hi. I’ve noticed that Windows Powershell has it’s own IEX command and you cannot access Elixir’s IEX due to the conflict. This isn’t a cr...
New
gshaw
What is the idiomatic way of matching for not nil in Elixir? E.g., First way: defp halt_if_not_signed_in(conn, signed_in_account) when...
New
sergio_101
I am VERY much an elixir newbie. I have taken one elixir course and one phoenix course on Udemy. During that course, I saw the instructor...
New
JorisKok
I have a server on AWS, and was running a load test using artillery. When looking at the Phoenix dashboard I see the Ports going to 100% ...
New

Other popular topics Top

Qqwy
Original source of discussion: This topic on the Pragmatic Programmers’ Functional Web Development with Elixir, OTP, and Phoenix forum. ...
New
sen
Hi All, I set a environment variables in dev.exs , like below code. when i start server, how can i set the ${enable} value? thanks. d...
New
msaraiva
Surface is an experimental library built on top of Phoenix LiveView and its new LiveComponent API that aims to provide a more declarative...
564 44139 214
New
TunkShif
This post is an instruction guide to help you setup your Neovim for Elixir development from scratch. It includes general information on h...
274 42533 114
New
Harrisonl
We have an ECS cluster with 4 services, where each task joins a single cluster, via discovery ECS discovery service. Currently when I de...
New
jason.o
In the code below, if the create action is not set to accept “extra_key” as an input, it errors out with a message shown above. Is there ...
New

We're in Beta

About us Mission Statement