Support training in no-std

[ivila]

Support training in no-std

Checklist

• [*] Confirmed that run-checks all script has been executed.
• [*] Made sure the book is up to date with changes in this PR.

Related Issues/PRs

#2829

Changes

1. burn-core

Make optim module support no-std:

1. Manually import alloc::vec::Vec and alloc::boxed::Box when no-std
2. Use num_traits::float::Float to support sqrt function in no-std
3. Mark all of the tests in optim std only as they require TestAutodiffBackend.

2. burn-autodiff

Fix no-std feature and provide a facade crate::libs for clean codes:

1. Changing some imports from std to core or alloc.
2. Replacing std::collections::{HashMap, HashSet} with hashbrown::{HashMap, HashSet} when no-std
3. Use num_traits::float::Float to support some math functions in no-std
4. Mark burn-autodiff NO_STD in xtask

3. burn

Fix dependency of burn-autodiff: add default-features = false on dependency of burn-autodiff

Testing

Successfully compile a TA(trusted application) and run with it.

PS: I tried run-checks.sh all but it always failed by network reason(cannot reach huggingface.co).

[codecov]

Codecov Report

Attention: Patch coverage is 92.78351% with 7 lines in your changes missing coverage. Please review.

Project coverage is 82.23%. Comparing base (d67cade) to head (9f77209).
Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/burn-autodiff/src/ops/int_tensor.rs	0.00%	3 Missing ⚠️
crates/burn-autodiff/src/ops/bool_tensor.rs	0.00%	2 Missing ⚠️
crates/burn-autodiff/src/ops/transaction.rs	0.00%	1 Missing ⚠️
crates/burn-autodiff/src/runtime/mutex.rs	0.00%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff            @@
##             main    #2830    +/-   ##
========================================
  Coverage   82.23%   82.23%            
========================================
  Files         853      853            
  Lines      113704   113888   +184     
========================================
+ Hits        93499    93654   +155     
- Misses      20205    20234    +29     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[laggui]

Will take a look at this tomorrow! But just a preliminary comment: I think we can just replace most types with their alloc types in all cases. That would probably just leave the collections stuff which is only used in a couple modules, so maybe we don't need to have a crate-level module to re-export compatible types.

[ivila]

Updates:

1. Directly import from alloc and core, but retain a crate-level module named collections that contains HashMap and HashSet to avoid duplicate feature-checking code in 10 modules.
2. Add the no-std target aarch64-unknown-none and test the no-std build of burn-autodiff specifically for that target.

[laggui]

Thanks for extending the portability!

Looks good so far, just a some comments below.

Mark all of the tests in optim std only as they require TestAutodiffBackend.

Now that you've made burn-autodiff no-std compatible this should be good to go (on 64-bit platforms, that is), no? Some tests save or load a state from a record file so these will need to be marked for std, but I think the rest should be fine 🤔

Also, NO_STD_64BIT_ONLY_CRATES with burn-autodiff is only due to the atomics at his point, unless I am forgetting something else... If that's the case, we could use portable-atomic to add the support. But that can be left for the future, mostly thinking out loud here 😄

/edit: oh and also the issue in formatting CI

[ivila]

Updates:

1. In burn-autodiff: Replace core::sync::atomic::{AtomicU64, Ordering} with portable-atomic::{AtomicU64, Ordering} to support 32-bit platforms.
2. In burn-autodiff: Remove the dependency on cfg-block.
3. In xtask: Revert changes in xtask and add burn-autodiff to NO_STD_CRATES.
4. In burn-core: Enable tests for the optim module in no-std, replace BinFileRecorder with BinBytesRecorder when no-std, and remove the dev-dependency on tempfile.

[ivila]

checking on the failed build of target thumbv6m-none-eabi

[ivila]

I tried using portable_atomic::Arc to replace alloc::sync::Arc for the target thumbv6m-none-eabi. It works, but I encountered another error:

The method clone_if_require_grad in burn_diff::graph::Node requires the arbitrary_self_types feature under the thumbv6m-none-eabi target.

After consideration, I believe we should leave this as is and skip building burn-autodiff for the thumbv6m-none-eabi target.

[laggui]

Awesome! Thanks for the changes

After consideration, I believe we should leave this as is and skip building burn-autodiff for the thumbv6m-none-eabi target.

Makes sense, this PR already makes it much closer to full no-std.

Will come back to this on Monday but should be good to go 🙂

[ivila]

Updates: rebase to latest mainline

[laggui]

One minor change left, missed it last time 😅

After that LGTM!

[laggui]

LGTM! Thanks for addressing the changes