Merge pull request #693 from kohya-ss/dev

Support for bitsandbytes 0.39.1 with Paged Optimizer
bmaltais · Jul 29, 2023 · 3a7326a · 3a7326a
2 parents 7b31c08 + 38b59a9
commit 3a7326a
Show file tree

Hide file tree

Showing 5 changed files with 54 additions and 12 deletions.
diff --git a/README-ja.md b/README-ja.md
@@ -125,6 +125,16 @@ pip install https://github.com/jllllll/bitsandbytes-windows-webui/raw/main/bitsa
 
 アップグレード時には`pip install .`でこのリポジトリを更新し、必要に応じて他のパッケージもアップグレードしてください。
 
+### オプション：PagedAdamW8bitとPagedLion8bitを使う
+
+PagedAdamW8bitとPagedLion8bitを使う場合には`bitsandbytes`を0.39.0以降にアップグレードする必要があります。`bitsandbytes`をアンインストールし、Windows環境では例えば[こちら](https://github.com/jllllll/bitsandbytes-windows-webui)などからWindows版のwhlファイルをインストールしてください。たとえば以下のような手順になります。
+
+```powershell
+pip install https://github.com/jllllll/bitsandbytes-windows-webui/releases/download/wheels/bitsandbytes-0.39.1-py3-none-win_amd64.whl
+```
+
+アップグレード時には`pip install .`でこのリポジトリを更新し、必要に応じて他のパッケージもアップグレードしてください。
+
 ## アップグレード
 
 新しいリリースがあった場合、以下のコマンドで更新できます。

diff --git a/README.md b/README.md
@@ -109,6 +109,16 @@ pip install https://github.com/jllllll/bitsandbytes-windows-webui/raw/main/bitsa
 
 For upgrading, upgrade this repo with `pip install .`, and upgrade necessary packages manually.
 
+### Optional: Use PagedAdamW8bit and PagedLion8bit
+
+For PagedAdamW8bit and PagedLion8bit, you need to upgrade `bitsandbytes` to 0.39.0 or later. Uninstall `bitsandbytes`, and for Windows, install the Windows version whl file from [here](https://github.com/jllllll/bitsandbytes-windows-webui) or other sources, like:
+
+```powershell
+pip install https://github.com/jllllll/bitsandbytes-windows-webui/releases/download/wheels/bitsandbytes-0.39.1-py3-none-win_amd64.whl
+```
+
+For upgrading, upgrade this repo with `pip install .`, and upgrade necessary packages manually.
+
 ## Upgrade
 
 When a new release comes out you can upgrade your repo with the following command:

diff --git a/docs/train_README-ja.md b/docs/train_README-ja.md
@@ -609,10 +609,12 @@ masterpiece, best quality, 1boy, in business suit, standing at street, looking b
     - AdamW : [torch.optim.AdamW](https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html)
     - 過去のバージョンのオプション未指定時と同じ
     - AdamW8bit : 引数は同上
+    - PagedAdamW8bit : 引数は同上
     - 過去のバージョンの--use_8bit_adam指定時と同じ
     - Lion : https://github.com/lucidrains/lion-pytorch
     - 過去のバージョンの--use_lion_optimizer指定時と同じ
     - Lion8bit : 引数は同上
+    - PagedLion8bit : 引数は同上
     - SGDNesterov : [torch.optim.SGD](https://pytorch.org/docs/stable/generated/torch.optim.SGD.html), nesterov=True
     - SGDNesterov8bit : 引数は同上
     - DAdaptation(DAdaptAdamPreprint) : https://github.com/facebookresearch/dadaptation

diff --git a/docs/train_README-zh.md b/docs/train_README-zh.md
@@ -546,9 +546,12 @@ masterpiece, best quality, 1boy, in business suit, standing at street, looking b
     -- 指定优化器类型。您可以指定
     - AdamW : [torch.optim.AdamW](https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html)
     - 与过去版本中未指定选项时相同
-    - AdamW8bit : 同上
+    - AdamW8bit : 参数同上
+    - PagedAdamW8bit : 参数同上
     - 与过去版本中指定的 --use_8bit_adam 相同
     - Lion : https://github.com/lucidrains/lion-pytorch
+    - Lion8bit : 参数同上
+    - PagedLion8bit : 参数同上
     - 与过去版本中指定的 --use_lion_optimizer 相同
     - SGDNesterov : [torch.optim.SGD](https://pytorch.org/docs/stable/generated/torch.optim.SGD.html), nesterov=True
     - SGDNesterov8bit : 参数同上

diff --git a/library/train_util.py b/library/train_util.py
@@ -2033,7 +2033,7 @@ def add_optimizer_arguments(parser: argparse.ArgumentParser):
         "--optimizer_type",
         type=str,
         default="",
-        help="Optimizer to use / オプティマイザの種類: AdamW (default), AdamW8bit, Lion8bit, Lion, SGDNesterov, SGDNesterov8bit, DAdaptation(DAdaptAdamPreprint), DAdaptAdaGrad, DAdaptAdam, DAdaptAdan, DAdaptAdanIP, DAdaptLion, DAdaptSGD, AdaFactor",
+        help="Optimizer to use / オプティマイザの種類: AdamW (default), AdamW8bit, PagedAdamW8bit, Lion8bit, PagedLion8bit, Lion, SGDNesterov, SGDNesterov8bit, DAdaptation(DAdaptAdamPreprint), DAdaptAdaGrad, DAdaptAdam, DAdaptAdan, DAdaptAdanIP, DAdaptLion, DAdaptSGD, AdaFactor",
     )
 
     # backward compatibility
@@ -2649,7 +2649,7 @@ def task():
 
 
 def get_optimizer(args, trainable_params):
-    # "Optimizer to use: AdamW, AdamW8bit, Lion, SGDNesterov, SGDNesterov8bit, Lion8bit, DAdaptation(DAdaptAdamPreprint), DAdaptAdaGrad, DAdaptAdam, DAdaptAdan, DAdaptAdanIP, DAdaptLion, DAdaptSGD, Adafactor"
+    # "Optimizer to use: AdamW, AdamW8bit, Lion, SGDNesterov, SGDNesterov8bit, PagedAdamW8bit, Lion8bit, PagedLion8bit, DAdaptation(DAdaptAdamPreprint), DAdaptAdaGrad, DAdaptAdam, DAdaptAdan, DAdaptAdanIP, DAdaptLion, DAdaptSGD, Adafactor"
 
     optimizer_type = args.optimizer_type
     if args.use_8bit_adam:
@@ -2726,20 +2726,37 @@ def get_optimizer(args, trainable_params):
         print(f"use Lion optimizer | {optimizer_kwargs}")
         optimizer_class = lion_pytorch.Lion
         optimizer = optimizer_class(trainable_params, lr=lr, **optimizer_kwargs)
-
-    elif optimizer_type == "Lion8bit".lower():
+        
+    elif optimizer_type.endswith("8bit".lower()):
         try:
             import bitsandbytes as bnb
         except ImportError:
             raise ImportError("No bitsandbytes / bitsandbytesがインストールされていないようです")
 
-        print(f"use 8-bit Lion optimizer | {optimizer_kwargs}")
-        try:
-            optimizer_class = bnb.optim.Lion8bit
-        except AttributeError:
-            raise AttributeError(
-                "No Lion8bit. The version of bitsandbytes installed seems to be old. Please install 0.38.0 or later. / Lion8bitが定義されていません。インストールされているbitsandbytesのバージョンが古いようです。0.38.0以上をインストールしてください"
-            )
+        if optimizer_type == "Lion8bit".lower():
+                print(f"use 8-bit Lion optimizer | {optimizer_kwargs}")
+                try:
+                    optimizer_class = bnb.optim.Lion8bit
+                except AttributeError:
+                    raise AttributeError(
+                        "No Lion8bit. The version of bitsandbytes installed seems to be old. Please install 0.38.0 or later. / Lion8bitが定義されていません。インストールされているbitsandbytesのバージョンが古いようです。0.38.0以上をインストールしてください"
+                    )
+        elif optimizer_type == "PagedAdamW8bit".lower():
+                print(f"use 8-bit PagedAdamW optimizer | {optimizer_kwargs}")
+                try:
+                    optimizer_class = bnb.optim.PagedAdamW8bit
+                except AttributeError:
+                    raise AttributeError(
+                        "No PagedAdamW8bit. The version of bitsandbytes installed seems to be old. Please install 0.39.0 or later. / PagedAdamW8bitが定義されていません。インストールされているbitsandbytesのバージョンが古いようです。0.39.0以上をインストールしてください"
+                    )
+        elif optimizer_type == "PagedLion8bit".lower():
+                print(f"use 8-bit Paged Lion optimizer | {optimizer_kwargs}")
+                try:
+                    optimizer_class = bnb.optim.PagedLion8bit
+                except AttributeError:
+                    raise AttributeError(
+                        "No PagedLion8bit. The version of bitsandbytes installed seems to be old. Please install 0.39.0 or later. / PagedLion8bitが定義されていません。インストールされているbitsandbytesのバージョンが古いようです。0.39.0以上をインストールしてください"
+                    )
 
         optimizer = optimizer_class(trainable_params, lr=lr, **optimizer_kwargs)