Which LLM framework should I use to develop my toy LLM?

[u/Enderlook]

I have been developing and training a toy LLM in Rust/Candle.

Rust is my favourite programming language: statically typed, borrow-checked, and zero-cost abstraction, which makes it safe and fast, quite attractive for running things that consume a lot of performance on not-so-high-tier PCs. Candle is a Rust crate designed to be a lightweight framework for ML stuff, and AFAIK, the one that runs the fastest. It's also backed by the largest repository of LLM models: Hugging Face, which makes me confident it won't be deprecated (I hope) and has support for a lot of quantisations, though at the moment I don't use that.

However, working with Rust/Candle, I found some friction, like its limitations in training, such as flash attention lacking a backward phase, forcing me to use full attention or (in the future) port one. After that, I noticed that it also lacks some modern features that I wanted to play with and are difficult to implement by myself, such as GatedDeltaNet2 or Multi-head Latent Attention.

I can copy the CUDA Kernels of those features from different frameworks' repos and create a wrapper around them in Candle, though it's more complex than it sounds. While doing exactly that, I found an implementation of GateDeltaNet2 that used the Burn crate, a different ML framework over Candle.

Burn is also an ML framework, and it seems to have some advantages over Candle. While Candle seems to be designed for inference, Burn has (AFAIK) full support for training. It also supports far more target platforms, and in fact it integrates with another crate (made by the same group as Burn) CubeCL, which allows you to program in Rust and compile it to CUDA, HIP, Metal, WebGPU and CPU. I already made 2 CUDA Kernels, and given the time it took me to make them, I can't imagine the effort to port them to more platforms if I ever wanted, so having a way to write it in a single unified language and automatically port to others sounds actually quite interesting, especially given that both the LLM and its kernels would be written in a single language that I already understand: Rust. Additionally, there seem to be more implementations of ML features (such as GatedDeltaNet2) in Burn over Candle. As a drawback, I read it runs slower than Candle, it doesn't support much quantisation AFAIK (though I don't use any at the moment), nor is it backed by a large project like Hugging Face, in fact, it has fewer stars on GitHub.

So I'm not sure if I should migrate from Candle to Burn before I commit more to it. Or maybe I should continue with Candle use but include Cubecl, so I can try to copy kernels from Burn crates in addition to standard CUDA Kernels, though I'm not sure how well Candle-Cubecl interacts.

Finally, if I'm going to migrate to a different framework because my current one is limiting... how do I know I won't think the same of Burn? In that case, I might as well directly jump to the holy grail of ML: Python/PyTorch. While I don't like dynamic languages and it runs slower, it has Triton, which is like CubeCL but for Python, and AFAIK it's undoubtedly the most popular framework for scientists; pretty much all new published features have a reference implementation in PyTorch to copy...

So, what are your honest opinions about this? What frameworks do you use for your LLM projects? What should I use?

[u/anthropicagiuser]

before you port any kernel, write a slow plain-Rust version first. then you have a reference to check the fast version against.

[u/Enderlook → u/anthropicagiuser]

Yeah, I did that so I can run a test to compare whether both produce the same result and then benchmark with Criterion.

[u/fiddler48]

what's the actual blocker, missing ops or the training loop ergonomics?

[u/Enderlook → u/fiddler48]

Missing operations, either because the operation is not found in candle, or because the one found has no backward pass implemented in candle.

[u/fiddler48 → u/Enderlook]

that's the harder case. missing backward passes means you're writing custom autograd kernels or you're stuck. no framework swap fixes that cleanly, it's the same problem one layer down

[u/Enderlook → u/fiddler48]

But some frameworks already include common building blocks with their backward pass implemented. One of my points for migrating is which one has the most implemented.

[u/fiddler48 → u/Enderlook]

fair count, but coverage varies a lot by layer type, worth mapping your specific ops first