Update train_dreambooth_colossalai.py

accelerator.num_processes -> gpc.get_world_size(ParallelMode.DATA)

Update train_dreambooth_colossalai.py
accelerator.num_processes -> gpc.get_world_size(ParallelMode.DATA)
9edd0aa7 · Haofan Wang · GitHub · f1bc2418 · 9edd0aa7
Unverified Commit 9edd0aa7 authored Jan 05, 2023 by Haofan Wang Committed by GitHub Jan 05, 2023
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

examples/images/dreambooth/train_dreambooth_colossalai.py examples/images/dreambooth/train_dreambooth_colossalai.py +1 -1

No files found.
--- a/examples/images/dreambooth/train_dreambooth_colossalai.py
+++ b/examples/images/dreambooth/train_dreambooth_colossalai.py
@@ -484,7 +484,7 @@ def main(args):
        unet.enable_gradient_checkpointing()

    if args.scale_lr:
-        args.learning_rate = args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * 2
+        args.learning_rate = args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * gpc.get_world_size(ParallelMode.DATA)

    unet = gemini_zero_dpp(unet, pg, args.placement)