Add LoRA-GGPO for Flux

2026-04-08 22:35:09 +00:00 · 2025-03-06 00:00:38 -05:00
parent 3d79239be4
commit ea53290f62
2 changed files with 137 additions and 1 deletions
--- a/train_network.py
+++ b/train_network.py
@@ -1400,6 +1400,10 @@ class NetworkTrainer:
                            params_to_clip = accelerator.unwrap_model(network).get_trainable_params()
                            accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)

+                    if global_step % 5 == 0:
+                        if hasattr(network, "update_norms"):
+                            network.update_norms()
+
                    optimizer.step()
                    lr_scheduler.step()
                    optimizer.zero_grad(set_to_none=True)