527

train.py

@@ -37,7 +37,7 @@ use_wandb = True
 save_model = True
 use_decay = True
 fbp = False # fused backward pass
-adam8bit = True
+optimizer_type = "adam8bit"
 torch_compile = False
 unet_gradient = True
 clip_sample = False #Scheduler
@@ -187,7 +187,7 @@ if use_wandb and accelerator.is_main_process:
         "base_learning_rate": base_learning_rate,
         "num_epochs": num_epochs,
         "fbp": fbp,
-        "adam8bit": adam8bit,
+        "optimizer_type": optimizer_type,
     })
 
 # Включение Flash Attention 2/SDPA
@@ -495,89 +495,63 @@ else:
     if fbp:
         trainable_params = list(unet.parameters())
 
-if fbp:
-    # [1] Создаем словарь оптимизаторов (fused backward)
-    if adam8bit:
-        optimizer_dict = {
-            p: bnb.optim.AdamW8bit(
-                [p],  # Каждый параметр получает свой оптимизатор
-                lr=base_learning_rate,
-                eps=1e-8
-            ) for p in trainable_params
-        }
+def create_optimizer(name, params):
+    if name == "adam8bit":
+        import bitsandbytes as bnb
+        return bnb.optim.AdamW8bit(
+            params, lr=base_learning_rate, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01
+        )
+    elif name == "adam":
+        return torch.optim.AdamW(
+            params, lr=base_learning_rate, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01
+        )
+    elif name == "lion8bit":
+        import bitsandbytes as bnb
+        return bnb.optim.Lion8bit(
+            params, lr=base_learning_rate, betas=(0.9, 0.97), weight_decay=0.01,
+            percentile_clipping=percentile_clipping
+        )
+    elif name == "adafactor":
+        from transformers import Adafactor
+        return Adafactor(
+            params, lr=base_learning_rate, scale_parameter=True, relative_step=False,
+            warmup_init=False, eps=(1e-30, 1e-3), clip_threshold=1.0,
+            beta1=0.9, weight_decay=0.01
+        )
     else:
-        optimizer_dict = {
-            p: bnb.optim.Lion8bit(
-                [p],  # Каждый параметр получает свой оптимизатор
-                lr=base_learning_rate,
-                betas=(0.9, 0.97),
-                weight_decay=0.01,
-                percentile_clipping=percentile_clipping,
-            ) for p in trainable_params
-        }
+        raise ValueError(f"Unknown optimizer: {name}")
+
+if fbp:
+    # Создаем отдельный оптимизатор для каждого параметра
+    optimizer_dict = {p: create_optimizer(optimizer_type, [p]) for p in trainable_params}
     
-    # [2] Определяем hook для применения оптимизатора сразу после накопления градиента
     def optimizer_hook(param):
         optimizer_dict[param].step()
         optimizer_dict[param].zero_grad(set_to_none=True)
     
-    # [3] Регистрируем hook для trainable параметров модели
     for param in trainable_params:
         param.register_post_accumulate_grad_hook(optimizer_hook)
     
-    # Подготовка через Accelerator
     unet, optimizer = accelerator.prepare(unet, optimizer_dict)
 else:
-    if adam8bit:
-        optimizer = bnb.optim.AdamW8bit(
-            params=unet.parameters(),
-            lr=base_learning_rate,
-            betas=(0.9, 0.999),
-            eps=1e-8,
-            weight_decay=0.01
-        )
-        #from torch.optim import AdamW
-        #optimizer = AdamW(
-        #    params=unet.parameters(),
-        #    lr=base_learning_rate,
-        #    betas=(0.9, 0.999),
-        #    eps=1e-8,
-        #    weight_decay=0.01
-        #)
-    else:
-        optimizer = bnb.optim.Lion8bit(
-            params=unet.parameters(),
-            lr=base_learning_rate,
-            betas=(0.9, 0.97),
-            weight_decay=0.01,
-            percentile_clipping=percentile_clipping,
-        )
-
-    def lr_schedule(step, max_steps, base_lr, min_lr, use_decay=True):
-        # Если не используем затухание, возвращаем базовый LR
+    optimizer = create_optimizer(optimizer_type, unet.parameters())
+    
+    def lr_schedule(step):
+        x = step / (total_training_steps * world_size)
+        warmup = 0.05
+        
         if not use_decay:
-            return base_lr
+            return base_learning_rate
+        if x < warmup:
+            return min_learning_rate + (base_learning_rate - min_learning_rate) * (x / warmup)
         
-        # Иначе используем линейный прогрев и косинусное затухание
-        x = step / max_steps
-        percent = 0.05
-        if x < percent:
-            # Линейный прогрев до percent% шагов
-            return min_lr + (base_lr - min_lr) * (x / percent)
-        else:
-            # Косинусное затухание
-            decay_ratio = (x - percent) / (1 - percent)
-            return min_lr + 0.5 * (base_lr - min_lr) * (1 + math.cos(math.pi * decay_ratio))
+        decay_ratio = (x - warmup) / (1 - warmup)
+        return min_learning_rate + 0.5 * (base_learning_rate - min_learning_rate) * \
+               (1 + math.cos(math.pi * decay_ratio))
     
-    
-    def custom_lr_lambda(step):
-        return lr_schedule(step, total_training_steps*world_size, 
-                         base_learning_rate, min_learning_rate, 
-                         use_decay) / base_learning_rate
-    
-    lr_scheduler = LambdaLR(optimizer, lr_lambda=custom_lr_lambda)
+    lr_scheduler = LambdaLR(optimizer, lambda step: lr_schedule(step) / base_learning_rate)
     unet, optimizer, lr_scheduler = accelerator.prepare(unet, optimizer, lr_scheduler)
-    
+
  # Регистрация хуков ПОСЛЕ prepare
 dispersive_hook.register_hooks(unet, "down_blocks.2")
 

unet/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e00ac22a38007a5fadae89d5d39b7286f8d1189c3ba4806653a29214dfd04e3f
+oid sha256:7a3a5c9271f789977d7738b696c288c509253bbb416e182ab58925101a0a9409
 size 7014306128