feat!: Device::metal() is cubecl’s Metal; sliding-window attention (#6333)
Co-authored-by: Thierry Cantin-Demers 48504765+ThierryCantin-Demers@users.noreply.github.com Co-authored-by: Guillaume Lagrange lagrange.guillaume.1@gmail.com
版权所有:中国计算机学会技术支持:开源发展技术委员会
京ICP备13000930号-9
京公网安备 11010802047560号
A tensor library and deep learning framework for numerical computing, training and inference.
Training and inference usually live in separate worlds. Models are typically trained in Python then exported to an open format like ONNX or optimized for production engines like vLLM, ONNX Runtime, or TensorRT. This export step is often brittle and lossy, ruling out complex architectures and advanced deployment use cases.
Burn unifies the two. By executing multi-platform tensor operations via a single, unified API, the exact code used for training is the exact code that runs in production. This makes workloads like on-device personalization and federated learning straightforward, while enabling teams to go from prototype to deployment in a single codebase.
Burn preserves the intuitive ergonomics of PyTorch, with dynamic shapes and graphs, but JIT-compiles streams of tensor operations, performing automatic kernel fusion. You get the flexibility of dynamic graphs without the performance drop.
Rust for Research?
Rust used to be a tough sell for research: long compilation times disrupted the fast edit-compile-run loop that draws researchers to Python. Burn changes this paradigm. Designed around incremental compilation, modifying model code recompiles in under 5 seconds, even in release mode. This delivers a Python-like feedback loop with the speed and safety of Rust.
Ecosystem
Burn is the core of a growing, fully open-source Rust AI ecosystem. You are not adopting a single library, you are joining a stack that spans GPU compute, model interop and domain toolkits, with plenty of room to help shape what comes next.
burn-storeburn-visionburn-rlburn-datasetBurn’s CubeCL backends (CUDA, ROCm, Metal, Vulkan, WebGPU, CPU) compose with autodiff, fusion and remote-execution decorators, while the pure-Rust Flex CPU/
no_stdbackend composes with autodiff only. See Supported Backends below for the full matrix.Every project here is open-source and actively developed. Want to help build the Rust AI ecosystem? The good first issues are a great place to start, and the Contributing guide will get you set up.
Community crates 🌱
These crates are not maintained by Tracel, but they are part of the same Rust AI story. Anything that helps you load data, build environments, or ship models belongs here. Built something that fits? Open a PR to add it!
Backend
Burn strives to be as fast as possible on as many hardwares as possible, with robust implementations. We believe this flexibility is crucial for modern needs where you may train your models in the cloud, then deploy on customer hardwares, which vary from user to user.
Supported Backends
Most backends support all operating systems, so we don’t mention them in the tables below.
GPU Backends:
CPU Backends:
The two native CPU backends are independent. Cpu (
cpufeature) is the CubeCL runtime for the CPU: it JIT-compiles the same kernels as the GPU backends through LLVM and supports fusion. Flex (flexfeature) is a pure-Rust eager backend with no native dependencies that also runs on Wasm andno_std.Burn’s backend architecture lets you swap backends while keeping the same model code. You can enable multiple backends in the same application and choose the device for your tensors and modules at runtime through
Device. This gives you the freedom to use different backends side by side and select the hardware best suited to each workload.Autodifferentiation and automatic kernel fusion integrate with the same tensor and module APIs, so models benefit from these capabilities on supported backends without changing their implementation.
Autodiff: Bringing backpropagation to any backend 🔄
In application code, autodiff is runtime context carried by tensors. Devices provide the default context for newly created tensors, and each tensor can later enable or remove autodiff independently. Enabling autodiff permits graph recording; it does not by itself make a tensor retain gradients.
Internally, Burn implements this by decorating a concrete backend, so autodiff cannot execute by itself. Enable it on a device before creating tensors or initializing a model. With the
autodiffandwgpufeatures enabled:backward()checks graph participation at runtime. Enable autodiff before the forward pass and callrequire_grad()on source leaves whose gradients you need.is_autodiff(),is_tracked(), andis_require_grad()inspect autodiff association, graph participation, and gradient retention respectively. See the autodiff guide.See the Autodiff Backend README for more details.
Fusion: Backend decorator that brings kernel fusion to supported backends
This backend decorator enhances a backend with kernel fusion, provided that the inner backend supports it. Note that you can compose this backend with other backend decorators such as Autodiff. All first-party accelerated backends (like WGPU and CUDA) use Fusion by default (
burn/fusionfeature flag), so you typically don’t need to apply it manually.Device::autodiff().gradient_checkpointing()enables the balanced gradient-checkpointing strategy, which trades recomputation for reduced activation storage during training.See the Fusion Backend README for more details.
Remote (Beta): Backend decorator for remote backend execution, useful for distributed computations
Remote execution has a client and a server. The server’s devices select the compute backend; clients use a remote
Devicewith the same tensor API. Iroh, the default transport, reaches a server across any network, authenticated and encrypted; see the server example and the distributed computing guide. On a trusted network, WebSocket is the simplest setup: enableremote-server,remote-websocket, andcudaon the server, andremote-websocketplusautodiffon the client:Training & Inference
The whole deep learning workflow is made easy with Burn, as you can monitor your training progress with an ergonomic dashboard, and run inference everywhere from embedded devices to large GPU clusters.
Burn was built from the ground up with training and inference in mind. It’s also worth noting how Burn, in comparison to frameworks like PyTorch, simplifies the transition from training to deployment, eliminating the need for code changes.
Click on the following sections to expand 👇
Training Dashboard 📈
As you can see in the previous video (click on the picture!), a new terminal UI dashboard based on the Ratatui crate allows users to follow their training with ease without having to connect to any external application.
You can visualize your training and validation metrics updating in real-time and analyze the lifelong progression or recent history of any registered metrics using only the arrow keys. Break from the training loop without crashing, allowing potential checkpoints to be fully written or important pieces of code to complete without interruption 🛡
ONNX Support 🐫
Burn supports importing ONNX (Open Neural Network Exchange) models through the burn-onnx crate, allowing you to easily port models from TensorFlow or PyTorch to Burn. The ONNX model is converted into Rust code that uses Burn’s native APIs, enabling the imported model to run on any Burn backend (CPU, GPU, WebAssembly) and benefit from all of Burn’s optimizations like automatic kernel fusion.
Our ONNX support is further described in this section of the Burn Book 🔥.
Importing PyTorch or Safetensors Models 🚚
You can load weights from PyTorch or Safetensors formats directly into your Burn-defined models. This makes it easy to reuse existing models while benefiting from Burn’s performance and deployment features.
Learn more in the Saving & Loading Models section of the Burn Book.
Inference in the Browser 🌐
Several of our backends can run in WebAssembly environments: Flex for CPU execution, and WGPU for GPU acceleration via WebGPU. This means that you can run inference directly within a browser. We provide several examples of this:
Embedded: no_std support ⚙️
Burn’s core components support no_std. This means it can run in bare metal environment such as embedded devices without an operating system.
Benchmarks
To evaluate performance across different backends and track improvements over time, we provide a dedicated benchmarking suite.
Run and compare benchmarks using burn-bench.
Getting Started
Just heard of Burn? You are at the right place! Just continue reading this section and we hope you can get on board really quickly.
The Burn Book 🔥
To begin working effectively with Burn, it is crucial to understand its key components and philosophy. This is why we highly recommend new users to read the first sections of The Burn Book 🔥. It provides detailed examples and explanations covering every facet of the framework, including building blocks like tensors, modules, and optimizers, all the way to advanced usage, like coding your own GPU kernels.
Examples 🙏
Let’s start with a code snippet that shows how intuitive the framework is to use! In the following, we declare a neural network module with some parameters along with its forward pass.
We have a somewhat large amount of examples in the repository that shows how to use the framework in different scenarios.
Following the book:
Moduleto train on the MNIST dataset and use for inference.Learner.Additional examples:
Learnerprogress.Module(MLP) with theLearnerconfigured to log metrics and keep training checkpoints.For more practical insights, you can clone the repository and run any of them directly on your computer!
Pre-trained Models 🤖
We keep an updated and curated list of models and examples built with Burn, see the tracel-ai/models repository for more details.
Don’t see the model you want? Don’t hesitate to open an issue, and we may prioritize it. Built a model using Burn and want to share it? You can also open a Pull Request and add your model under the community section!
Why use Rust for AI? 🦀
Deep Learning is a special form of software where you need very high level abstractions as well as extremely fast execution time. Rust is the perfect candidate for that use case since it provides zero-cost abstractions to easily create neural network modules, and fine-grained control over memory to optimize every detail. To this day, the mainstream solution has been to offer APIs in Python but rely on bindings to low-level languages such as C/C++. This reduces portability, increases complexity and creates friction between researchers and engineers. Rust’s approach to abstractions is versatile enough to tackle this two-language dichotomy, and Cargo makes it easy to build, test and deploy from any environment, which is usually a pain in Python.
Rust’s AI ecosystem is young, but it is real and growing quickly. Foundational pieces are already here: Burn and CubeCL for training and compute, candle for inference, Hugging Face’s
tokenizersandsafetensors, andpolarsandndarrayfor data. Betting on Rust today means betting on a stack that is growing, and one where contributors still shape the direction. The pieces that don’t exist yet are opportunities rather than dead-ends (see Contributing).Rust is also what makes one-stack-everywhere possible: a single self-contained binary with no Python runtime to ship, running from servers down to
no_stdembedded targets.Loading Model Records From Previous Versions ⚠️
Burn 0.22 uses burnpack for native records and cannot directly read legacy
Recorderformats such as.mpk,.bin, or JSON. Load the checkpoint in a compatible older Burn project, export the model weights throughburn-store, and import them into your 0.22 model. See Migrating checkpoints for an example and the distinction between transferring weights and resuming training state.For records saved before
0.14.0, an earlier migration step may also be needed: use Burn0.14,0.15, or0.16with therecord-backward-compatfeature to load and re-save the record with the newer tensor-data representation. For legacy binary records, first use the version that wrote them to export a self-describing format such asNamedMpkFileRecorder. This historical migration does not convert records to the format required by 0.22; follow the checkpoint migration guide afterward.Community
If you are excited about the project, don’t hesitate to join our Discord! We try to be as welcoming as possible to everybody from any background. You can ask your questions and share what you built with the community!
Contributing
Before contributing, please read the Contributing Guidelines and our Code of Conduct. The Contributor Book covers architecture, environment setup, and guides for common tasks.
Status
Burn is currently in active development, and there will be breaking changes. While any resulting issues are likely to be easy to fix, there are no guarantees at this stage.
License
Burn is distributed under the terms of either the MIT License or the Apache License, Version 2.0. Unless otherwise noted, this applies to all contents of this repository, including source code, examples, documentation, and code snippets in the Burn Book and Contributor Book. See the licenses for details.
Third-party materials and externally downloaded datasets or models remain subject to their respective licenses.
Opening a pull request is assumed to signal agreement with these licensing terms.