はじめに
この記事は単にSDXLのfloat16型のモデル重みを一度int型に変換してから再度float16型に変換して代入している。このため推論自体はfloat16のままである。そのため実行速度が向上したり、学習メモリが低減したりはしない。
ただ保存モデルの容量が低減できる可能性を示せるだけである。
QLoRAの論文中の量子化技術がいくつか書いてあるのだが、これの効果が気になってやってみている。
また、以下の記事も参考にした。
1. 全結合のみInt8型
最初に範囲が$-127～127$のInt8型に変換する場合を考える。
まずif len(param.size())==2 and not('emb' in name):としているのは行列が二次元の全結合重み(nn.linear)のみ指している。len(param.size())==1は正規化レイヤーやバイアスのレイヤーで、len(param.size())==4はConv2dの重みであるが、全結合重みにしぼった。
全結合レイヤーのみに絞った理由はLoRA自体が二次元重みにしか適用できないことによる。
embの名の付くレイヤーも除いているが、これはembeddingの入力がindexで出力が一次元なので、実質一次元データの集合になっていると見なして除外した。
さて、int8への変換については各レイヤーの重みの絶対値最大値を探してこれが127になるような係数$c$をかけ、四捨五入関数(round)を掛けて整数にして、これをint8に変換する。四捨五入関数を掛けないとゼロ点の頻度が何故か倍になる。あとは$-127～127$のInt8型を係数$c$で割って元に戻す。
また、各モデル重みの$c$を保持する必要があるが、1パラメータにすぎないので重みの全体量に対してほとんどない。
from diffusers import DiffusionPipeline
import torch
model_id = './stable-diffusion-xl-base-1.0/'
output_path = './quant/'
pipe = DiffusionPipeline.from_pretrained(model_id, use_safetensors=True, torch_dtype=torch.float16, variant="fp16").to("cuda")
pipe.enable_model_cpu_offload()
prompt = "a photo of an astronaut riding a horse on mars"
seed = 42
generator = torch.Generator(device="cuda")
generator = generator.manual_seed(seed)
image = pipe(prompt=prompt, generator=generator).images[0]
image.save(output_path + "img_float16.png")
for name, param in pipe.unet.named_parameters():
#if len(param.size())==2 and not('emb' in name) and not('proj' in name) and not('ff' in name):
if len(param.size())==2 and not('emb' in name):
print(name, param.size(), type(param))
M = param.data.clone().to(torch.float32)
c = 127.0 / torch.max(torch.abs(M))
M2 = (torch.round(M * c)).to(torch.int8)
M2 = M2.to(torch.float16) / c
param.data = M2
# param = torch.nn.Parameter(M2)
generator = generator.manual_seed(seed)
image = pipe(prompt=prompt, generator=generator).images[0]
image.save(output_path + "img_int8.png")
この時、parameter数を数えると以下の様になった。
上述の全結合重みは太枠の範囲でモデル重み全体の86.0%を占める。
2. 全結合のみInt7~2型
Int7～2型を考える。
bit-1となるのは最初の1bitは符号に使われるため。
変換はtorch.int8型のままなので保持する容量自体には変化はない。
容量を減らすには7bit以下の変数を上手く保持するトリックがいる。
for bit in [7,6,5,4,3,2]:
pipe = DiffusionPipeline.from_pretrained(model_id, use_safetensors=True, torch_dtype=torch.float16, variant="fp16").to("cuda")
pipe.enable_model_cpu_offload()
for name, param in pipe.unet.named_parameters():
if len(param.size())==2 and not('emb' in name):
print(name, param.size(), type(param))
M = param.data.clone().to(torch.float32)
c = (2.0**(bit-1)-1.0) / torch.max(torch.abs(M))
M2 = (torch.round(M * c)).to(torch.int8)
M2 = M2.to(torch.float16) / c
param.data = M2
generator = generator.manual_seed(seed)
image = pipe(prompt=prompt, generator=generator).images[0]
image.save(output_path + "img_int%d.png" % (bit))
3. 全ての重みInt8~2型
前項で全結合重みのみint8型に変換する事を考えていたが、バイアスやら畳み込み重みも量子化してみる。
また全結合重みとConv2D重みのみ量子化するパターンも考える。
for bit in [8,7,6,5,4,3,2]:
pipe = DiffusionPipeline.from_pretrained(model_id, use_safetensors=True, torch_dtype=torch.float16, variant="fp16").to("cuda")
pipe.enable_model_cpu_offload()
for name, param in pipe.unet.named_parameters():
print(name, param.size(), type(param))
M = param.data.clone().to(torch.float32)
c = (2.0**(bit-1)-1.0) / torch.max(torch.abs(M))
M2 = (torch.round(M * c)).to(torch.int8)
M2 = M2.to(torch.float16) / c
param.data = M2
generator = generator.manual_seed(seed)
image = pipe(prompt=prompt, generator=generator).images[0]
image.save(output_path + "img_int%d_all.png" % (bit))
for bit in [8,7,6,5,4,3,2]:
pipe = DiffusionPipeline.from_pretrained(model_id, use_safetensors=True, torch_dtype=torch.float16, variant="fp16").to("cuda")
pipe.enable_model_cpu_offload()
for name, param in pipe.unet.named_parameters():
if len(param.size())==4 or (len(param.size())==2 and not('emb' in name)):
print(name, param.size(), type(param))
M = param.data.clone().to(torch.float32)
c = (2.0**(bit-1)-1.0) / torch.max(torch.abs(M))
M2 = (torch.round(M * c)).to(torch.int8)
M2 = M2.to(torch.float16) / c
param.data = M2
generator = generator.manual_seed(seed)
image = pipe(prompt=prompt, generator=generator).images[0]
image.save(output_path + "img_int%d_conv.png" % (bit))
4. block_size=64, 全結合のみInt8~2型
前述の量子化は外れ値に非常に弱い。例えば、$M=[-10.0,-1.0,-0.9,-0.8,..,0.8,0.9,1.0]$となるようなほとんどの値は$-1.0～1.0$にあり、一点$-10.0$だけ外れている。
この場合、前述の量子化では$-10～10$を等間隔に刻むがこれでは多くのbitを有効に使えない。
したがってこの場合、先ほどのような重み毎に決めた$c$をblock毎に決める必要がある。
block_size=64では変数64個毎に32bitの$c$を一個保存する。
この時、1パラメータに占める$c$の割合は$32bit/64=0.5bit$となる。
LLMの量子化ではblock_size=128(group_size=128)で$scale, zero\ point$の二個の変数を持って$32bit×2/128=0.5bit$としているのを見かけたが、QLoRAの論文では$zero\ point$は持たず、$scale$のみに見える。
これは後述のNF4の理由からそうなっているのだと思われる。
パラメータがblock_sizeでちょうど割り切れるなら簡単なのだが、割り切れない時重み行列を伸ばしてやらないといけなくちょっと手間である。
for bit in [8,7,6,5,4,3,2]:
pipe = DiffusionPipeline.from_pretrained(model_id, use_safetensors=True, torch_dtype=torch.float16, variant="fp16").to("cuda")
pipe.enable_model_cpu_offload()
for name, param in pipe.unet.named_parameters():
if len(param.size())==2 and not('emb' in name):
print(name, param.size(), type(param))
M = param.data.clone().to(torch.float32)
block_size = 64
orig_shape = M.shape
orig_length = len(M.flatten())
if orig_length%block_size==0:
extend_M = M.flatten()
else:
extend_M = torch.cat((M.flatten(), torch.zeros(block_size - orig_length%block_size)))
extend_length = len(extend_M)
Mmax = []
for i in range(extend_length//block_size):
Mmax.append(torch.max(torch.abs(extend_M[block_size*i:block_size*(i+1)])))
Mmax = torch.stack(Mmax, dim=0)
Mmax = Mmax.repeat_interleave(block_size)
c = (2.0**(bit-1)-1.0) / Mmax
M2 = (torch.round(extend_M * c)).to(torch.int8)
M2 = M2.to(torch.float16) / c
M2 = M2[:orig_length].reshape(orig_shape).to(torch.float16)
param.data = M2
generator = generator.manual_seed(seed)
image = pipe(prompt=prompt, generator=generator).images[0]
image.save(output_path + "img_int%d_block64.png" % (bit))
5. 全結合のみNF4型
さて、モデル重みが正規分布（ガウス分布）になっていると仮定すると、そもそも範囲を等間隔で刻むという事自体が効率的ではないかもしれない。
例えば正規分布では0付近が最も頻度が高く、最大値付近にはほとんど値がない。
従って0付近を小さく刻み、最大値付近では大きく刻みたい。
理論上は一様分布に対して正規分布の累積分布関数の逆関数torch.erfinvを掛ければいいが、QLoRAの論文ではNF4の量子化に$q=[-1.0, -0.6961928009986877, -0.5250730514526367,
-0.39491748809814453, -0.28444138169288635, -0.18477343022823334,
-0.09105003625154495, 0.0, 0.07958029955625534, 0.16093020141124725,
0.24611230194568634, 0.33791524171829224, 0.44070982933044434,
0.5626170039176941, 0.7229568362236023, 1.0]$
という値があるのでこれを使わしてもらう。
この関数自体は正規分布の累積分布関数の逆関数の一様刻みの中点から求まるっぽい。
torch.argmin(torch.abs(q-M2))のように書きたかったが、上手い実装がよく分からないのでtorch.where()を使って書けば以下の様になる。
このときM2は$0～15+100$のindex値になる。M3は$-1.0～1.0$の値に変数64個毎の最大値を掛けて戻す。このときの1パラメータあたり4.5bitである。
q = [-1.0, -0.6961928009986877, -0.5250730514526367,
-0.39491748809814453, -0.28444138169288635, -0.18477343022823334,
-0.09105003625154495, 0.0, 0.07958029955625534, 0.16093020141124725,
0.24611230194568634, 0.33791524171829224, 0.44070982933044434,
0.5626170039176941, 0.7229568362236023, 1.0]
q.append(1.01)
bit = 4
pipe = DiffusionPipeline.from_pretrained(model_id, use_safetensors=True, torch_dtype=torch.float16, variant="fp16").to("cuda")
pipe.enable_model_cpu_offload()
for name, param in pipe.unet.named_parameters():
if len(param.size())==2 and not('emb' in name):
print(name, param.size(), type(param))
M = param.data.clone().to(torch.float32)
block_size = 64
orig_shape = M.shape
orig_length = len(M.flatten())
if orig_length%block_size==0:
extend_M = M.flatten()
else:
extend_M = torch.cat((M.flatten(), torch.zeros(block_size - orig_length%block_size)))
extend_length = len(extend_M)
Mmax = []
for i in range(extend_length//block_size):
Mmax.append(torch.max(torch.abs(extend_M[block_size*i:block_size*(i+1)])))
Mmax = torch.stack(Mmax, dim=0)
Mmax = Mmax.repeat_interleave(block_size)
M2 = extend_M / Mmax
for i in range(16):
M2 = torch.where(M2 <= (q[i]+q[i+1])/2, 100+i, M2)
M2 = torch.round(M2)
M2 = M2.to(torch.int8)
M3 = M2.to(torch.float16)
for i in range(16):
M3 = torch.where(i+100==M2, q[i], M3)
M3 = M3 * Mmax
M3 = M3[:orig_length].reshape(orig_shape).to(torch.float16)
param.data = M3
generator = generator.manual_seed(seed)
image = pipe(prompt=prompt, generator=generator).images[0]
image.save(output_path + "img_NF4.png")
6. NF4型のqについて
QLoRAの論文を読みつつqの値について考える。
$-8～8$の17個($2^k+1$)の整数値を入れた場合の$Qx$を定義し、これの隣接点との平均を16点を$q$とする。係数として適当な$0.96$を入れている理由は、これが$1.0$だとtorch.erfinvに1や-1が入った時、∞や-∞に発散してしまうため。
この時、大体同じ分布を再現出来たのだが微妙に一致しない。論文の$q$が非対称である。
import torch
import numpy as np
import matplotlib.pyplot as plt
x = np.arange(16)
Qx = []
for i in range(-8, 9):
Qx.append(torch.erfinv(torch.tensor(i/8.0*0.96, dtype=torch.float32)))
Qx = torch.stack(Qx)
Qx = Qx.to('cpu').detach().numpy().copy()
q = np.zeros(16)
for i in range(16):
q[i] = (Qx[i]+Qx[i+1])/2
q = q/np.max(q)
print(len(Qx), Qx)
print(len(q), q)
q2 = np.array([-1.0, -0.6961928009986877, -0.5250730514526367,
-0.39491748809814453, -0.28444138169288635, -0.18477343022823334,
-0.09105003625154495, 0.0, 0.07958029955625534, 0.16093020141124725,
0.24611230194568634, 0.33791524171829224, 0.44070982933044434,
0.5626170039176941, 0.7229568362236023, 1.0])
print(len(q2), q2)
plt.plot(x, q)
plt.plot(x, q2)
plt.show()
「非対称量子化(Asymmetric Quantization)」あたりで検索するとそういうのがあるが、内容はよく分かってない。例えば以下の様にx座標を操作すればplotの一致は高くなる。
x2 = np.abs(np.linspace(-0.5, 0.5, 16)) - 0.5
...
plt.plot(x+x2, q)
plt.plot(x, q2)
また、こういう理論分布から求めるのではなく実際のモデル重みの分布を解析して、現実の累積分布関数の逆関数から求める手法もあるが、解析に時間が掛かるのと各重み毎に異なる長さ16のqテーブルを保持する必要がある。
これは以下のPalettizationがそれに該当すると思われる。
また、float8などを考えることもゼロ点近傍では細かく刻む、端では大きく刻むデータ型を考えるという観点からすると方針は近いかもしれない。
ところで、個人的には前述のブロックごとに分割するblockwiseとこのqテーブルを使うNF4の相性について気にかかった。結局、ブロックごとに倍率がことなるのでこの$q$の中央付近の傾きは各倍率ごとにバラバラになって実際の分布と一致しないのではないかと思った。
例えばブロック内に外れ値がない場合、最大値付近の値の数は多いのに大きく刻んでしまう。それでも等間隔に刻むよりはマシなのだろうか？
ブロックごとの正確なqテーブルを作るにはスケール以外に分布の尖度（一様分布と正規分布を見分ける指標？）が必要なのではないだろうかと思った。
仮にブロック内の分布が一様分布なら等間隔に刻むのが最も良い。
7.1 二重量子化(Double Quantization)
QLoRAの論文に載っている手法。
前述のblockwiseでは変数64個毎に32bitの変数を1個持って$32bit/64=0.5bit$であったが、二重量子化では変数64個毎に8bitの変数を1個持ち、かつ変数64・256個毎に32bitの変数を1個持つ。
この時、$8bit/64+32bit/(64\cdot 256)= 0.127 bit$で元のblockwiseより1パラメータ当たり$0.373bit$削れる。
しかし、現時点ではこれを検討の必要を感じるほどSDXLの量子化の知見は得ていない。
7.2 Mixed-Bit
参考記事に出てくる手法。
レイヤーごとに量子化を試してみて性能(PSNR)低下が激しいレイヤーは量子化bitを多めにして、性能低下がそれほどでもないレイヤーは量子化bitを少なめに設定する。
これによってレイヤーごとに異なる量子化bitを採用して削減を目指す。
参考の記事中に4.5bitの量子化があったので最初てっきり4bit+0.5bitのblockwiseを指すのかと思ったがどうやらそうではないらしい。
自分がlinearレイヤーに絞って量子化するのも意味的には近い。
8. Backwardの精度
仮にforwardが4bitで済んだとしても学習時にbackwardの精度が4bitでは恐らく上手く学習出来ないと思われる。学習時のbackwardには推論時より高い精度が必要なことが知られている。
これはQLoRAにおいてもそうでLLMの4bitのモデル重み自体は学習対象からは外れており、全結合重みに並列する16bitのLoRA重みのみを学習するため、4bitのbackwardは考えなくてよく16bitのbackwardのみで学習出来る。
従って仮にSDXLの推論に4bitモデルを構築できたとしても、モデル全体を学習するDreamboothやfinetuningの勾配の精度は不足しているため学習は上手く行かないと思われる。
QLoRAのように4bitのSDXLと16bitのLoRAを構築すればLoRAの学習は上手く行くのではと思われる。
9. 結果
9-1. 量子化するレイヤー
全てのレイヤーを量子化するより全結合重み(nn.linear)に絞って量子化する方がよく見える。
また、畳み込み層(Conv2D)を量子化した場合にも性能低下がある。
何故か分からないがbit7でbit6より性能低下がある。
linearは重みの86%、Conv2Dは重みの13%からモデルの削減率を求めると以下の様になる
9-2. 量子化する方法
Blockwiseは変数64個毎の係数$c$に1パラメータ当たり0.5bitとるのでInt型より0.5bitだけ大きい。
NF4 blockwiseは同じ4bit Int blockwiseより馬の脚の形でよく、5bit Intより背景の描写で元の形に近い。なのでこの中ではNF4 blockwiseがよく見える。
まとめ
SDXLの量子化について検討した。
量子化の改善手法は大きく分けてBlockごとに最大値を保存するBlockwiseと、範囲の分割を非等間隔に分割するPalettizationに分かれる。
前者は外れ値の影響を低減する効果があるが、1パラメータあたり0.5bit増加する。
後者は高い頻度の場所ほど細かく刻めるが、qテーブルなどを求めないといけない。
両方の組み合わせは個人的に相性は悪い気がしたが、結果を見る限りNF4 Blockwiseは悪くはない。
出力を見た感じ4割程度までは劣化は少なく削減できそうだが、それ以降は劣化が目につく。
Mixed-Bitのアイデアを採用してlinearレイヤーを4.5bit、Conv2Dレイヤーを8bitに削れば3割程度まで削れるかもしれない。その場合、約5.1bit相当である。