kyegomez/BitNet

解決する課題

BitNetは、大規模言語モデル（LLM）のメモリおよび計算要件を大幅に削減するように設計された、1-bit TransformerのPyTorch実装です。標準的な線形投影を二値化レイヤーに置き換えることで、モデルの重みの超圧縮を可能にします。

仕組み

このプロジェクトは、Transformer内の標準的な nn.Linear を置き換える BitLinear モジュールを実装しています。プロセスには、レイヤー正規化、二値化、abs-max量子化、それに続く逆量子化のシーケンスが含まれます。このリポジトリでは、replace_linears_in_pytorch_model および replace_linears_in_hf を使用して、既存のPyTorchまたはHugging Faceモデルの線形層を入れ替えるツールを提供しています。

対象者

1-bit LLMを実装したい、あるいはテキスト、画像、音声処理の超圧縮を探索したいAI研究者や開発者向けです。これらのレイヤーを使用するモデルは、正しく機能させるために、最初からトレーニングまたはファインチューニングする必要があることに注意してください。

ハイライト

• 幅広いアーキテクチャのサポート: BitNetTransformer、BitAttention（Multi-Grouped Query Attentionを使用）、BitFeedForward、BitMoE（Mixture of Experts）、および BitMamba の実装が含まれています。
• 統合ツール: 標準的なPyTorchおよびHugging Faceモデルの線形層を簡単に置き換えるためのユーティリティ関数を提供します。
• Visionサポート: Vision Transformerに1-bit圧縮を適用するための OneBitViT を含んでいます。
• パフォーマンスの最適化: 低ビットGEMM操作用の最適化されたCUDAカーネルを搭載しています。