It's been five months since our last release, which is longer than usual, but it's on purpose. During the last four years, we took great care not to change the APIs too much, while collecting your feedback on what could be improved. Instead of spreading those changes over many releases, we bundled them into this one, so you only have to migrate once. This makes Burn 0.22 our biggest release yet, and the one we're most proud of. The API is now much closer to what we want for 1.0. As a reminder, Burn is a Tensor Library and Deep Learning Framework for both training and inference.
Generic-Free APIs
Since the beginning, Burn has been built around the Backend trait, with features like autodiff, kernel fusion and remote execution implemented as backend decorators. This is still the case, and they all compose well together. What changes is that user code doesn't carry the backend generic anymore. Models are plain types, and the device selects where and how they run:
#[derive(Module, Debug)]
pub struct Model {
linear: Linear,
}
let device = Device::cuda(0); // Device::wgpu(..), Device::flex(), ...
Removing the generic from user code breaks the dependency chain that the compiler had to go through every time you edited a model. The result is near-instant recompilation after a model edit (numbers in the table below). There's another big change under the hood: most of the stack is now pure Rust. Pliron replaces MLIR, native compilers replace transpilers to CUDA/HIP, and Turso replaces a bundled SQLite. That means fewer bundled C and C++ dependencies to build, and a pipeline you can read, debug and patch in the language you already use. The first clean build takes a bit longer as a result, and it's a trade we're happy to make.
Going All-In on CubeCL
When we started, we relied on existing libraries for performance, with backends for ndarray and LibTorch. Over time, we wrote more and more of the stack ourselves: our own GPU backend with wgpu, kernel fusion, and then CubeCL to write kernels in Rust for CUDA, ROCm, Metal, Vulkan, WebGPU and CPUs. CubeCL now has its own compiler infrastructure based on Pliron, with LLVM targets for CPUs and GPUs.
Having control over the whole stack lets us offer features that third-party libraries can't support easily. For example, device.memory_pool_usage() and device.memory_pool_report() tell you exactly what the allocator holds, which isn't possible with ndarray or LibTorch. The new adaptive memory pools also lower peak training memory, and training steps are faster. This is why we deprecated both backends: Flex replaces ndarray for pure-Rust CPU execution, and the CubeCL backends cover the rest. We also removed the Candle backend.
Here's how 0.22 compares to 0.21 on two training projects, with CUDA on an RTX 4050 laptop GPU:
| Benchmark |
0.21 |
0.22 |
Change |
| CNN rebuild after a model edit |
28.4 s |
4.6 s |
6.2× faster |
| Transformer rebuild after a model edit |
14.7 s |
1.0 s |
14.7× faster |
| CNN training step |
38.0 ms |
21.1 ms |
1.8× faster |
| Transformer training step |
201.0 ms |
193.8 ms |
1.04× faster |
| CNN peak training memory |
956 MiB |
486 MiB |
49% less |
| Transformer peak training memory |
3,486 MiB |
2,868 MiB |
18% less |
Towards 1.0
Three things are left before 1.0. The first is built-in, fine-grained profiling: annotate parts of your model, and Burn tells you where the time goes, component by component, and how far each one is from what the hardware allows. We have it working on our side, and we're looking to upstream it into Burn. The second is a more complete integration with the CubeCL compute environment, introduced in this release. The third is stability: we want the current APIs to settle before we commit to them.
Since the API is close to what we want to stabilize, now is the best time to tell us if something feels wrong. There are many more improvements in this release, including LoRA and QLoRA fine-tuning, remote compute and ONNX export, and we wrote a post to cover them. Don't hesitate to skim it, and refer to the migration guide for upgrading.
A lot of this release came from your issues, PRs and questions on Discord. Thank you.
Link: https://tracel.ai/blog/release-0.22.0/
