2512

snooc.py

@@ -0,0 +1,89 @@
+import torch
+from torch.optim import Optimizer
+
+class SnooC(Optimizer):
+    """
+    @DominikKallusky, @vishal9-team, @vinaysrao
+
+    Sparse Nesterov Outer Optimizer (Snoo) is a momentum-based wrapper to any optimizer that can
+    improve the stability and smoothness of the optimization process and thus the quality
+    of large language models (LLM) and other models. Snoo implicitly adds temporal regularization
+    to the parameters, thus smoothing the training trajectory and instilling a bias towards flatter
+    minima and lower parameter norms. Snoo is computationally efficient, incurring minimal overhead
+    in compute and moderate memory usage.
+    """
+
+    @torch.no_grad()
+    def __init__(self, optimizer, lr: float = 0.67, momentum: float = 0.67, k: int = 20) -> None:
+        self.optimizer = optimizer
+        self.lr = lr
+        self.momentum = momentum
+        self.k = k
+        self.current_step = 0
+        self.model_params = None
+        self.outer_buf = None
+        self.outer_optimizer = None
+
+        # Check if the optimizer already has parameters
+        if self.optimizer.param_groups:
+            self.param_groups = self.optimizer.param_groups
+    
+    @torch.no_grad()
+    def _initialize_outer_optimizer(self):
+        params = []
+        for pg in self.optimizer.param_groups:
+            if len(pg['params']) > 1:
+                for param in pg['params']:
+                    if isinstance(param, torch.Tensor):
+                        params.append(param)
+            else:
+                params = pg['params']
+        
+        if not params:
+            return
+
+        self.model_params = list(params)
+        self.outer_buf = [p.clone() for p in self.model_params]
+        self.outer_optimizer = torch.optim.SGD(
+            self.model_params,
+            lr=self.lr,
+            momentum=self.momentum,
+            nesterov=True,
+            fused=True,
+        )
+        self.param_groups = self.optimizer.param_groups
+        del params
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        if self.outer_optimizer is None or self.current_step == 0:
+            # If the optimizer has been updated with parameters, initialize.
+            if self.optimizer.param_groups:
+                self._initialize_outer_optimizer()
+            else:
+                # If there are still no parameters, we cannot perform a step.
+                # Depending on the use case, you might want to raise an error
+                # or simply return without doing anything.
+                return self.optimizer.step(closure)
+
+        loss = self.optimizer.step(closure)
+        if self.current_step % self.k == 0:
+            for p_new, p_old in zip(self.model_params, self.outer_buf):
+                p_new.grad = p_old.data - p_new.data
+                p_new.copy_(p_old, non_blocking=True)
+
+            self.outer_optimizer.step()
+
+            for p_new, p_old in zip(self.model_params, self.outer_buf):
+                p_old.copy_(p_new, non_blocking=True)
+        self.current_step += 1
+        return loss
+    
+    def zero_grad(self, set_to_none: bool = False):
+        self.optimizer.zero_grad(set_to_none=set_to_none)
+
+    def state_dict(self):
+        return self.optimizer.state_dict()
+
+    def load_state_dict(self, state_dict):
+        self.optimizer.load_state_dict(state_dict)

train.py

@@ -29,14 +29,14 @@ from transformers import AutoTokenizer, AutoModel
 # --------------------------- Параметры ---------------------------
 ds_path = "/workspace/sdxs/datasets/768"
 project = "unet"
-batch_size = 40
+batch_size = 36
 base_learning_rate = 4e-5
 min_learning_rate = 2.7e-5
 num_epochs = 100
 sample_interval_share = 5
 max_length = 192
-use_wandb = True
-use_comet_ml = False
+use_wandb = False
+use_comet_ml = True
 save_model = True
 use_decay = True
 fbp = False
@@ -53,8 +53,8 @@ torch.backends.cudnn.allow_tf32 = True
 dtype = torch.float32
 save_barrier = 1.006
 warmup_percent = 0.01
-percentile_clipping = 98
-betta2 = 0.999 #0.999?
+percentile_clipping = 97
+betta2 = 0.999
 eps = 1e-7
 clip_grad_norm = 1.0
 limit = 0
@@ -377,68 +377,21 @@ def create_optimizer(name, params):
         return torch.optim.AdamW(
             params, lr=base_learning_rate, betas=(0.9, betta2), eps=1e-8, weight_decay=0.01
         )
-    elif name == "adamuon":
-        from pytorch_optimizer import AdaMuon
-        # 1. РАЗДЕЛЕНИЕ ПАРАМЕТРОВ
-        # Разделяем параметры на те, которые Muon-совместимы (ndim >= 2), 
-        # и те, которые лучше оставить для AdamW8bit (ndim < 2).
-        
-        muon_params = []
-        adam8bit_params = []
-
-        if isinstance(params, (list, tuple)):
-            all_params = params
-        else:
-            # Если params - это генератор (unet.parameters()), преобразуем его в list
-            all_params = list(params)
-        
-        for p in all_params:
-            if p.requires_grad:
-                if p.ndim >= 2:
-                    muon_params.append(p)
-                else:
-                    adam8bit_params.append(p)
-        
-        # 2. СОЗДАНИЕ ГРУПП ПАРАМЕТРОВ
-        param_groups = []
-        
-        # --- Группа 1: AdaMuon (для матриц/сверток) ---
-        if muon_params:
-            print(f"Используем AdaMuon для {len(muon_params)} параметров (ndim >= 2)")
-            param_groups.append({
-                'params': muon_params,
-                'optimizer': 'AdaMuon',              # Явно указываем оптимизатор
-                'use_muon': True,                    # Обязательно для AdaMuon
-                'lr': base_learning_rate,            # Использование обычных пробелов
-                'betas': (0.9, betta2),
-                'eps': eps,
-                'weight_decay': 0.01,
-                'nesterov': True,
-                'decouple_weight_decay': True
-            })
-            
-        # --- Группа 2: AdamW8bit (для смещений/нормализации) ---
-        if adam8bit_params:
-            print(f"Используем AdamW8bit для {len(adam8bit_params)} параметров (ndim < 2)")
-            param_groups.append({
-                'params': adam8bit_params,
-                'optimizer': 'AdamW8bit',            # Явно указываем оптимизатор
-                'use_muon': False,
-                # AdamW8bit не требует use_muon
-                'lr': base_learning_rate,            # Использование обычных пробелов
-                'betas': (0.9, betta2),
-                'eps': eps,
-                'weight_decay': 0.01,
-                'percentile_clipping': percentile_clipping
-            })
-
-        if not param_groups:
-            raise ValueError("Нет параметров для оптимизации.")
-            
-        # 3. ИНИЦИАЛИЗАЦИЯ (смешанный оптимизатор)
-        # При передаче списка групп оптимизатор AdaMuon автоматически 
-        # распознает ключи 'optimizer' и инициализирует внутренние оптимизаторы.
-        return AdaMuon(param_groups)
+    elif name == "muon":
+        from muon import MuonWithAuxAdam
+        trainable_params = [p for p in params if p.requires_grad]
+        hidden_weights = [p for p in trainable_params if p.ndim >= 2]
+        hidden_gains_biases = [p for p in trainable_params if p.ndim < 2]
+    
+        param_groups = [ 
+            dict(params=hidden_weights, use_muon=True,
+                lr=1e-3, weight_decay=1e-4),
+            dict(params=hidden_gains_biases, use_muon=False,
+                lr=1e-4, betas=(0.9, 0.95), weight_decay=1e-4),
+        ]
+        optimizer = MuonWithAuxAdam(param_groups)
+        from snooc import SnooC
+        return SnooC(optimizer)
     else:
         raise ValueError(f"Unknown optimizer: {name}")
 
@@ -723,12 +676,20 @@ for epoch in range(start_epoch, start_epoch + num_epochs):
                         # Передаем tuple (emb, mask) для негатива
                         generate_and_save_samples(fixed_samples, (uncond_emb, uncond_mask), global_step)
                         last_n = sample_interval
-    
+                        
                         if save_model:
-                            avg_sample_loss = np.mean(batch_losses[-sample_interval:]) if len(batch_losses) > 0 else 0.0
-                            print("saving:", avg_sample_loss < min_loss * save_barrier, "Avg:", avg_sample_loss)
-                            if avg_sample_loss is not None and avg_sample_loss < min_loss * save_barrier:
-                                min_loss = avg_sample_loss
+                            has_losses = len(batch_losses) > 0
+                            avg_sample_loss = np.mean(batch_losses[-sample_interval:]) if has_losses else 0.0
+                            last_loss = batch_losses[-1] if has_losses else 0.0
+                            max_loss = max(avg_sample_loss, last_loss)    
+                            should_save = max_loss < min_loss * save_barrier
+                            print(
+                                f"Saving: {should_save} | Max: {max_loss:.4f} | "
+                                f"Last: {last_loss:.4f} | Avg: {avg_sample_loss:.4f}"
+                            )
+                            # 6. Сохранение и обновление
+                            if should_save:
+                                min_loss = max_loss
                                 save_checkpoint(unet)
 
     if accelerator.is_main_process:

unet/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d40e06c2a67bb2729fd4dae0a554773051648c61bc3ffa00fedf9412bf40ec74
+oid sha256:e910ef5f2cb2a67f6f4259d48422d509e2b77e51021e41b611e7329aecd84aa0
 size 7444321360