はじめに
前回の続きです。touch-sp.hatenablog.com
自分の環境では「xFormers」がうまく使えなかったので「DeepSpeed」を用いてVRAM消費の削減を試みました。
チュートリアルには「bitsandbytes」の「--use_8bit_adam」との併用が現時点ではできないと書いてあります。そもそもoptimizerをCPUにoffloadしているため「--use_8bit_adam」を使用するメリットは少ないのではないかと思います。
（注意）おそらくCPU側のRAMが16GBでは動かないと思います。
環境構築
「DeepSpeed」はMicrosoftが開発しているライブラリです。でもWindowsでのインストールに失敗しました。そのためUbuntu 20.04 on WSL2を使用しています。Ubuntu 20.04 on WSL2 (Windows 11) CUDA 11.6.2 Python 3.8.10
pip install torch==1.13.1+cu116 torchvision==0.14.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116 pip install git+https://github.com/huggingface/diffusers.git pip install transformers==4.26.0 pip install accelerate==0.15.0 scipy==1.10.0 datasets==2.8.0 ftfy==6.1.1 tensorboard==2.11.2 pip install deepspeed==0.8.0
accelerate config
------------------------------------------------------------------------------------------------------------------------ In which compute environment are you running? This machine ------------------------------------------------------------------------------------------------------------------------ Which type of machine are you using? No distributed training Do you want to run your training on CPU only (even if a GPU is available)? [yes/NO]:NO Do you wish to optimize your script with torch dynamo?[yes/NO]:NO Do you want to use DeepSpeed? [yes/NO]: yes Do you want to specify a json file to a DeepSpeed config? [yes/NO]: NO ------------------------------------------------------------------------------------------------------------------------ What should be your DeepSpeed's ZeRO optimization stage? 2 ------------------------------------------------------------------------------------------------------------------------ Where to offload optimizer states? cpu ------------------------------------------------------------------------------------------------------------------------ Where to offload parameters? cpu How many gradient accumulation steps you're passing in your script? [1]: 1 Do you want to use gradient clipping? [yes/NO]: yes What is the gradient clipping value? [1.0]: 1.0 Do you want to enable `deepspeed.zero.Init` when using ZeRO Stage-3 for constructing massive models? [yes/NO]: yes How many GPU(s) should be used for distributed training? [1]:1 ------------------------------------------------------------------------------------------------------------------------ Do you wish to use FP16 or BF16 (mixed precision)? fp16
実験と結果
今回はすべて「prior-preservation loss」ありで実験しています。text encoderのファインチューニングなし
no gradient_checkpointing, no set_grads_to_none
accelerate launch train_dreambooth.py \ --pretrained_model_name_or_path="stable-diffusion-v1-4" \ --instance_data_dir="RoboData" \ --output_dir="dreambooth_robo" \ --instance_prompt="a photo of sks robo" \ --with_prior_preservation --prior_loss_weight=1.0 \ --class_data_dir="robo-class-images" \ --class_prompt="a photo of robo" \ --num_class_images=400 \ --resolution=512 \ --train_batch_size=1 \ --learning_rate=5e-6 \ --lr_scheduler="constant" \ --lr_warmup_steps=0 \ --checkpointing_steps=400 \ --max_train_steps=1200 \
with gradient_checkpointing, no set_grads_to_none
accelerate launch train_dreambooth.py \ --pretrained_model_name_or_path="stable-diffusion-v1-4" \ --instance_data_dir="RoboData" \ --output_dir="dreambooth_robo" \ --instance_prompt="a photo of sks robo" \ --with_prior_preservation --prior_loss_weight=1.0 \ --class_data_dir="robo-class-images" \ --class_prompt="a photo of robo" \ --num_class_images=400 \ --resolution=512 \ --train_batch_size=1 \ --learning_rate=5e-6 \ --lr_scheduler="constant" \ --lr_warmup_steps=0 \ --checkpointing_steps=400 \ --max_train_steps=1200 \ --gradient_checkpointing
no gradient_checkpointing, with set_grads_to_none
accelerate launch train_dreambooth.py \ --pretrained_model_name_or_path="stable-diffusion-v1-4" \ --instance_data_dir="RoboData" \ --output_dir="dreambooth_robo" \ --instance_prompt="a photo of sks robo" \ --with_prior_preservation --prior_loss_weight=1.0 \ --class_data_dir="robo-class-images" \ --class_prompt="a photo of robo" \ --num_class_images=400 \ --resolution=512 \ --train_batch_size=1 \ --learning_rate=5e-6 \ --lr_scheduler="constant" \ --lr_warmup_steps=0 \ --checkpointing_steps=400 \ --max_train_steps=1200 \ --set_grads_to_none
with gradient_checkpointing, with set_grads_to_none
text encoderのファインチューニングあり
現時点でDeepSpeedを使ってtext encoderのファインチューニングはできないようです。github.com
現時点での結論
- 「prior-preservation loss」ありでもDeepSpeedを使えばVRAM 10GBで学習可能。
- DeepSpeedを使ってtext encoderのファインチューニングはできない。
- DeepSpeedを今回の設定で使用した場合、「--set_grads_to_none」はVRAM削減にあまり貢献しない