8 bit AdamW optimizer

[Redhawk18]

Pull Request Template

Checklist

• Confirmed that cargo run-checks command has been executed.
• Made sure the book is up to date with changes in this PR.

Related Issues/PRs

Provide links to relevant issues and dependent PRs.
Decreases memory usage, a major problem in Burn.

Changes

Summarize the problem being addressed and your solution.
Adds two impls of 8 bit optimizers based on Tim Dettmers, the tensor operation one is mostly just an artifact of in progress work. The 8 bit kernel is the one that improves results overall.

Testing

Describe how these changes have been tested.
Unit testing everywhere, then the optimizers were compared directly with this integration benchmark that also produces figures.

Loss for the synthetic benchmark

Wikitext 2 benchmark

Table: Comparison of optimizer implementations at learning rate 1e-4: average step time (with standard deviation), optimizer state size, and average process memory.

Optimizer	Step time (ms ± µs)	State size (GB)	Process video memory (GB)
AdamW (fp32)	141 ± 675	11.988	38.80
8-bit AdamW (tensor-op)	1992 ± 8000	6.041	35.31
8-bit AdamW (kernel)	66 ± 650	3.044	20.88