Bitnet 1.58 prework, POC, and staging

[CoffeeVampir3]

Bitnet 1.58 Groundwork

After some talks with Saroufim and the cuda mode team working on bitnet, we've outlined a strategy for implementing bitnet 1.58 method into torch. This issue lays the groundwork for 2-bit trinary tensor quantization and bitnet linear work for Bitnet 1.58

I've set up a staging repo Staging with a number of items:

• To the point minimal lib
• Training notebook for creating a full model, up to the point where we quantize and pack
• Cleaned up minimal training example for running as a script
• Example of the compiled kernel

This covers the initial groundwork for getting working trinary networks into torch.

• [x] Example Quantization Method
• [x] POC layer quantization
• [ ] Runnable example model with quantized layers (In progress Dtype and Runnable Model)
• [ ] AO dtype
• [ ] AO layer type (?) for bitnet linear
• [ ] Runnable example model with full dtype + bitnet linear layer, shippable

[msaroufim]

Very cool! Thank you for writing up such a clear plan. We can start merging the bit-packing logic and the layer quantization so feel free to send a PR whenever you're ready. This is very much following a similar to the playbook for fp6 @gau-nernst followed

Related work

• Subclass work: https://github.com/pytorch/ao/pull/282
• Generic bit packing https://github.com/pytorch/ao/issues/284

[CoffeeVampir3]

:+1: https://github.com/pytorch/ao/pull/285 I've kept this seperate from Andreas' commit for now, this encapsulates only the working bits.