LFX Workspace: Integrate burn.rs as a new WASI-NN backend
derekwin opened this issue · comments
Summary
Motivation
Burn.rs is a new comprehensive dynamic Deep Learning Framework built using Rust with extreme flexibility, compute efficiency and portability as its primary goals. The task of this project is to integrate burn.rs as a WASI-NN backend in the form of a plugin within wasmedge.
Details
Milestones:
-
The basics of Burn.rs: (week2)
- Dive into the wasi-nn wit (ref https://github.com/WebAssembly/wasi-nn).
- Design the structure of plugin project (ref https://github.com/derekwin/WasmEdge/tree/wasm-rustrs/plugins/wasmedge_rustls).
- Determining the burn.rs's core functions and data structures that need to be exposed. (https://github.com/tracel-ai/burn/tree/main/crates)
-
The basics of Burn.rs: (week3)
- Dive into the burn.rs.
- Dive into the wasmedge-wasi-nn (ref https://github.com/second-state/wasmedge-wasi-nn).
- Complete the structure of plugin project.
- Determining the burn.rs's backend mods that need to be exposed. (https://github.com/tracel-ai/burn/tree/main/crates)
-
Implement the detail of burn.rs plugin: (week4-8)
- Improve the structure of the project.
- Error types. (https://github.com/WebAssembly/wasi-nn/blob/main/wit/wasi-nn.wit)(https://github.com/tracel-ai/burn/blob/main/crates/burn-core/src/record/serde/error.rs)(https://github.com/tracel-ai/burn/blob/main/crates/burn-import/src/pytorch/error.rs)
- Implement the functions and data structures need be exposed.
- First, implement the core crate of burn.rs.
- Integrate different backends for burn.rs across several different time periods.
- Unit test
-
The Unit test and examples code: (weeks9-10)
- complete all test.
- Design and implement some usages examples.
-
Documents and Tutorials: (week11)
- Finish the documents and tutorials of this plugin.
-
1 week as buffer for unexpected delay.
Appendix
Please notice that we are using wasmedge-wasi-nn, which is forked from wasi-nn with some LLM-required extensions
You can focus on the original wasi-nn spec first because our modification extends the original spec. Once finishing the basic features, you can move forward to support the whole wasmedge-wasi-nn extensions.
ok, i have added the wasmedge-wasi-nn to the plan.
week1 update:
After discussing with mentor, the following viewpoints were obtained:
- Plug-in development will only focus on inference.
- Plug-in development don't need no_std, only support native platform.
- Plug-in development should focus on how to encapsulate the contents of burn.rs into the form required by wasmedge-wasi-nn
- burn.rs may have provided some similar data structures like wasi-nn.
week2 update:
- Conducted an in-depth study on the differences between wasi-nn and wasmedge-wasi-nn.
- Essentially determined the division of mods.
- Basically confirmed the structure of plugin implementation.
- Sorted out the code logic of the relevant parts in burn.rs.
More details: https://zrsnlqw1xd.feishu.cn/docx/DOYodYbeSoMRuQx82YacWuUpnzc?from=from_copylink
week3&4 update:
- Constructed the initial framework for the plugin.
- At Wednesday's meeting, with the assistance of Hydai, clarified the use of 'loadbynamewithconfig' to address the multi-backend adaptation issue in Burn.
- The device will be also determined by metadata and auto.