Migrate burn-models to burn main branch with burn-store

[antimora]

Summary

Comprehensive migration of burn-models from burn 0.19.1 to burn main branch, replacing burn-import with the new burn-store crate for model weight loading and storage.

Changes by Model

bert-burn

• Update burn dependency to git main branch
• Migrate loader from burn-import to burn-store API

llama-burn

• Update burn/burn-store dependencies to git main branch
• Replace burn-import with burn-store (PytorchStore, SafetensorsStore, BurnpackStore)
• Add SafeTensors format support with PyTorchToBurnAdapter for HuggingFace models
• Fix TinyLlama weight permutation using burn-store's direct tensor access API
• Refactor key mappings using KeyRemapper::from_patterns() for cleaner code
• Add ndarray backend feature
• Add test_tiny.rs example for testing TinyLlama import

mobilenetv2-burn

• Update burn dependency to git main branch
• Migrate from NamedMpkFileRecorder to BurnpackStore

resnet-burn

• Update burn dependency to git main branch
• Migrate from NamedMpkFileRecorder to BurnpackStore
• Update inference example to use burn-store API

squeezenet-burn

• Update burn dependency to git main branch
• Update build script for burn-store compatibility

yolox-burn

• Update burn dependency to git main branch
• Migrate from NamedMpkFileRecorder to BurnpackStore

Key Migration Patterns

Old API	New API
burn-import	burn-store
PyTorchFileRecorder	PytorchStore
NamedMpkFileRecorder	BurnpackStore
SafetensorsFileRecorder	SafetensorsStore (new)
model.load_file(path, recorder)	model.load_from(&mut store)
model.save_file(path, recorder)	model.save_into(&mut store)

Test Plan

• bert-burn compiles
• llama-burn compiles and TinyLlama generates correct output
• mobilenetv2-burn inference example runs correctly (predicts "Labrador retriever")
• resnet-burn inference example runs correctly (predicts "Labrador retriever")
• squeezenet-burn compiles
• yolox-burn inference example runs correctly (detects person, bicycle, dog)

NOTES

1. We still need to update Cargo.toml to point to the published Burn crates
2. We need to update to the burnpack version (@laggui will do it )

[laggui]

Tried to execute the very first example in the repo and it failed to compile:

cargo run --example infer-embedding --release --features wgpu,fusion,safetensors

   Compiling burn-cubecl-fusion v0.20.0-pre.5 (https://github.com/tracel-ai/burn?branch=main#6c9ba5cf)
error[E0412]: cannot find type `Simple` in this scope
   --> /home/laggui/.cargo/git/checkouts/burn-6c277d792b0d5d7a/6c9ba5c/crates/burn-cubecl-fusion/src/optim/matmul/optimization.rs:172:36
    |
172 |         if arg.execute_fused::<BT, Simple>(context).is_err() {
    |                                    ^^^^^^ not found in this scope
    |
help: there is an enum variant `crate::optim::matmul::FusedMatmulSelector::Simple`; try using the variant's enum
    |
172 -         if arg.execute_fused::<BT, Simple>(context).is_err() {
172 +         if arg.execute_fused::<BT, crate::optim::matmul::FusedMatmulSelector>(context).is_err() {
    |
help: you might be missing a type parameter
    |
133 | impl<R: Runtime, Simple> MatmulOptimization<R> {

That's an issue with burn, but just reinforces the point that we should only have models point to a released version.

/edit: fixed this issue in tracel-ai/burn#4193

[laggui]

That feels a lot more verbose and possibly error prone (especially since we have to check for the "path" names) compared to the previous usage 😅

That's a good signal to improve the UX though.

[antimora]

Agreed! I've refactored to extract helper functions (permute_rotary_weights, permute_attention_weight) which makes the call site cleaner:

permute_rotary_weights(&mut llama.model, n_heads, n_kv_heads, d_model, device);

The verbosity stems from burn-store's ModuleAdapter not having access to:

• Model config (n_heads, d_model)
• Device (for tensor operations)

A potential future improvement could be a "post-load transformation hook" in burn-store that receives both the snapshot and model context. Would you like me to open an issue for this?

[laggui]

The call site is cleaner since it has been refactored into a function, but it still requires a lot more manipulations then previously. So I think there's still room for improvement for the API!

[laggui]

Reverted the llama load changes since the pre-trained weights are still uploaded online in the mpk format.

All models are working with the current PR.

I don't remember what your benchmarks initially showed on loading speed, but the PyTorchFileRecorder (while definitely wasn't fast) appears to load faster than the burn store equivalent. In my tests, it always loaded the llama 3.2 1b pth model in 7s vs 11s for the store.