DreamBooth DeepSpeed support for under 8 GB VRAM training #735
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
The documentation is not available anymore as the PR was closed or merged. |
|
Just tried this version, it works pretty well! DreamBooth finally fits onto a 3060, although it used 10743 MiB for me. Does DeepSpeed adjust the training to fit into VRAM or was it something else? Here's time comparison for different cases (800 steps in each case, not including time to generate class images):
*ShivamShrirao's fork caches the latents, which speeds up the training, but takes some time itself. For me, it was 02:34 for latent caching (for 103 instance images) and 14:36 for the training itself, resulting in 17:10. **I've noticed that with DeepSpeed enabled my CPU was under pretty heavy load, so including it into table too. By the way, if you don't want to run accelerate launch --use_deepspeed --zero_stage=2 --gradient_accumulation_steps=1 --offload_param_device=cpu --offload_optimizer_device=cpu train_dreambooth.py *training_arguments*Full command I've used for training: accelerate launch --use_deepspeed --zero_stage=2 --gradient_accumulation_steps=1 --offload_param_device=cpu --offload_optimizer_device=cpu train_dreambooth.py \
--pretrained_model_name_or_path=$MODEL_NAME --use_auth_token \
--instance_data_dir=$INSTANCE_DIR \
--class_data_dir=$CLASS_DIR \
--output_dir=$OUTPUT_DIR \
--with_prior_preservation --prior_loss_weight=1.0 \
--instance_prompt="a photo of sks dog" \
--class_prompt="a photo of dog" \
--resolution=512 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--gradient_checkpointing \
--learning_rate=5e-6 \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--num_class_images=200 \
--max_train_steps=800 \
--sample_batch_size=2 \
--mixed_precision=fp16 |
patrickvonplaten
approved these changes
Oct 7, 2022
|
Unfortunately, this did not work for me, tested on a 3080 10GB and 64GB of RAM. |
Due to recent commits some casts to half precision are not necessary anymore. Mention that DeepSpeed's version of Adam is about 2x faster.
|
Enabling |
There is a patch in the deepspeed repo to allow 8bit_adam to work with deepspeed, I haven't tested it so not sure if it works since there hasn't been any progress/comments on it since December. |
patil-suraj
approved these changes
Oct 10, 2022
|
|
||
| # Add the prior loss to the instance loss. | ||
| loss = loss + args.prior_loss_weight * prior_loss | ||
| else: | ||
| loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean() | ||
| loss = F.mse_loss(noise_pred.float(), noise.float(), reduction="mean") |
There was a problem hiding this comment.
why cast to float32 here ? Do we always want to compute loss in full precision ?
cc @patrickvonplaten
There was a problem hiding this comment.
It's mixed precision best practice to calculate large reduction in higher precision. This calculates mean of batch_size * 4 * 64 * 64 halfs and mse_loss is one of the operations that would be automatically casted to fp32 in fp16 with autocast. I'm not sure it it's necessary at low batch size as it does seem to work without it, but it doesn't really affect memory consumption since it's only one operation and should give some safety.
|
Good to merge for me if @patil-suraj is happy with it! |
|
Thanks a lot for the amazing contribution ! |
|
It appears that mixed precision is broken now |
prathikr
pushed a commit
to prathikr/diffusers
that referenced
this pull request
Oct 26, 2022
…e#735) * Support deepspeed * Dreambooth DeepSpeed documentation * Remove unnecessary casts, documentation Due to recent commits some casts to half precision are not necessary anymore. Mention that DeepSpeed's version of Adam is about 2x faster. * Review comments
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Add instructions on how to enable DeepSpeed in DreamBooth example to allow training on under 8 GB VRAM. I was able to train a working network using this on 8 GB VRAM GPU.
It did not work out of the box with fp16 mixed precision and I had to add some explicit casts to make it run. Without them it raises Exception:
RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.cuda.HalfTensor) should be the same.