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

vertexbuffer
Hello, can anybody help here..? I have a list of players and I what to delete an element, but every for loop the list is reverting to ori...
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
siddhant3030
Hi, I have to write a raw query for one of my project. But till now I have used ecto queries and don’t have much experience writing raw ...
New
WestKeys
Currently suffering from paralysis by [HTTP client] analysis. This is rather unusual in Elixirland as there tends to be consensus on the ...
New
JeremM34
Hello, how can I check the Phoenix version ? Thanks !
New
pmjoe
I have a relationship of love and hate with Elixir. Lots of things are just absolutely right, but there are some things that are kind of ...
New
jerry
Good day to you all. I have been struggling to get a query involving like and ilike to work. Can anyone assist me on this, please? pro...
New
Qqwy
Original source of discussion: This topic on the Pragmatic Programmers’ Functional Web Development with Elixir, OTP, and Phoenix forum. ...
New
srinivasu
How to handle excepions in elixir? Suppose i have A, B, C ,D, E modules. and each module has get() function. A.get() method will call t...
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

senggen
Erlang/OTP 25 [erts-13.2.2] [source] [64-bit] [smp:8:8] [ds:8:8:10] [async-threads:1] 15:22:35.803 [error] gen_event {lager_file_backend...
New
mcarvalho
What is the difference between System.get_env and Application.get_env? For example, what are best practices to use one versus another.
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
Nvim
Anybody knows a comprehensive comparison of Django and Phoenix, thanks for the help. Where are they similar? Where do they differ the m...
New
chrismccord
Phoenix 1.4.0 released Phoenix 1.4 is out! This release ships with exciting new features, most notably with HTTP2 support, improved deve...
688 31013 112
New
pmjoe
I have a relationship of love and hate with Elixir. Lots of things are just absolutely right, but there are some things that are kind of ...
New
fayddelight
I tried installing elixir 1.11.2 erlang 23.3.4 via asdf in my zsh shell. Enabled the versions locally and globally. When I list them ...
New
rms.mrcs
Hi, I need to transform a list of numbers into a map where the keys are the indexes and the values are the original values of the list. ...
New
JakeBecker
TL;DR: I’ve just released an implementation of Microsoft’s IDE-independent Language Server Protocol for Elixir. It adds language support ...
1144 54120 245
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

Latest on Elixir Forum

Elixir Forum

We're in Beta

About us Mission Statement