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

chokchit
** (DBConnection.ConnectionError) connection not available and request was dropped from queue after 2733ms. You can configure how long re...
New
aadeshere1
I have a another noob question about loop. Since elixir is immutable, while loop is not directly possible. total = 10 while total != 0 ...
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
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
minhajuddin
I have seen a lot of code which picks the first element from a list using Enum.at(0) instead of List.first. Is there a reason why people ...
New
vac
Hi, I’m quite new in Elixir and I’m trying to format a string to a PEM format. I have the certificate value like MIIDBTCCAe2...... and I...
New
ycv005
I have followed this StackOverflow post to install the specific version of Erlang. And When I am running mix ecto.setup then getting fol...
New
aalberti333
As the title describes, I’m trying to run Enum.map() over a list of key/value pairs, where the value is a map. My data looks like this: ...
New
lucidguppy
I have a super simple question about elixir - how would I take a file like this foo bar baz and output a new file that enumerates th...
New
jononomo
For some reason my phoenix channels are working for me in my local dev environment, but as soon as I deploy via Docker, I get a 403 error...
New

Other popular topics Top

9mm
I am constructing a JSON object (map) and I need to conditionally set a field. I’m trying to write proper elixir-way code… and I’m at a l...
New
shahryarjb
Hello, I have map which I want to convert it to string like this: the map: %{last_name: "tavakkoli", name: "shahryar"} the string I ne...
New
jononomo
I am trying to figure out how Mix knows whether the environment is test, dev, or prod – where is this set? Thanks.
New
minhajuddin
I have seen a lot of code which picks the first element from a list using Enum.at(0) instead of List.first. Is there a reason why people ...
New
freewebwithme
Using vs code and installed ElixirLS: support and debugger. And I got an error popped up on start up says Failed to run ‘elixir’ comma...
New
vonH
When I run the Plug and I recompile I wind up having to use Ctrl C to quit iex and start again. Witht the help of rlwrap I can use the cu...
New
ashish173
I am using Ecto timestamps with postgres, I can see the timestamps() use the :naive_dateime but for my use case I wanted to store the ti...
New
KronicDeth
Elixir plugin for JetBrain’s IntelliJ Platform (including Rubymine) This is a plugin that adds support for Elixir to JetBrains IntelliJ...
289 36352 110
New
klo
Got a question about when to concat vs. prepending items to list then reversing to achieve appending. So i know lists boil down to [1 | ...
New
dokuzbir
I want to highlight html closing tags when i click a html tag. That works in .html files but doesnt work for html.eex templates. How can...
New

Latest on Elixir Forum

Elixir Forum

We're in Beta

About us Mission Statement