nindanaoto / bitnet158b Goto Github PK

論文に書かれていたBitNet1.58bのBitLinear部の実装です。
https://arxiv.org/abs/2310.11453
https://arxiv.org/abs/2402.17764

論文を読み自力で書いたところもありますが、以下の方々のコードが非常に参考になりました。
https://github.com/Beomi/BitNet-Transformers
https://github.com/frodo821/BitNet-Transformers
https://github.com/kyegomez/BitNet

参考にさせてもらったコードから大きな変更点としては、activationの順番とdeqauntizeの部分を直しています。
また、Beomiさんのコードのllamaに追加してみたところ、epoch60000でlossは3.7付近でした。


実験1:int8への変換 少し試してみましたが、matmal時に"addmm_cuda" not implemented for 'Char'となってしまうため実装できません。CPUで動かしてみましたが、あまりの遅さで実験できませんでした。仕方ないのでroundをして四捨五入してみましたが、メモリ使用量は変化ありませんでした。 補足1:torch.compile
自作したGPTではtorch.compileはmatmal後に2の累乗以外をかけるとlossがnanになりました。ところがllamaだと問題ないようです。

補足2:モデルサイズと学習率について
モデルサイズを大きくするとともに、論文通りに、あるいはそれよりさらに学習率を下げたほうが良さそうです。

補足3:上記のllamaのweightの保存時でエラーが出る
epoch67892でBeomiさんのコードとの相性で止まってしまいますが、bitnetの感覚はつかめると思います。
line 893あたりの保存設定を変えればなんとかなりそうですが、llamaに詳しくないため現在対処が難しいです。
時間があったらその部分にも挑戦します。