Merge pull request #1064 from KohakuBlueleaf/fix-grad-sync

Avoid grad sync on each step even when doing accumulation
kohya-ss · Jan 23, 2024 · 7a20df5 · 7a20df5
2 parents bea4362 + 711b40c
commit 7a20df5
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/train_network.py b/train_network.py
@@ -842,10 +842,11 @@ def remove_model(old_ckpt_name):
                     loss = loss.mean()  # 平均なのでbatch_sizeで割る必要なし
 
                     accelerator.backward(loss)
-                    self.all_reduce_network(accelerator, network)  # sync DDP grad manually
-                    if accelerator.sync_gradients and args.max_grad_norm != 0.0:
-                        params_to_clip = accelerator.unwrap_model(network).get_trainable_params()
-                        accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
+                    if accelerator.sync_gradients:
+                        self.all_reduce_network(accelerator, network)  # sync DDP grad manually
+                        if args.max_grad_norm != 0.0:
+                            params_to_clip = accelerator.unwrap_model(network).get_trainable_params()
+                            accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
 
                     optimizer.step()
                     lr_scheduler.step()