2507

2b/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:42851d6484909485ba508da8b9605eb46fc3ccb39e73d6b9f00e84bc0c7980a2
+oid sha256:33928b6c0aaa936c3f2f72ae338c4c3712c4b6fc97d3bf7ba2ba9b8ac116055e
 size 7993399544

train.py

@@ -27,10 +27,10 @@ import torch.nn.functional as F
 # --------------------------- Параметры ---------------------------
 ds_path = "datasets/384"
 project = "2b"
-batch_size = 32
+batch_size = 64
 base_learning_rate = 8e-5
 min_learning_rate = 4e-5
-num_epochs = 27
+num_epochs = 25
 # samples/save per epoch
 sample_interval_share = 5
 use_wandb = True
@@ -43,7 +43,7 @@ unet_gradient = True
 clip_sample = False #Scheduler
 fixed_seed = False
 shuffle = True
-dispersive_loss_enabled = False
+dispersive_loss_enabled = True
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.allow_tf32 = True
 torch.backends.cuda.enable_mem_efficient_sdp(False)
@@ -59,7 +59,7 @@ steps_offset = 1 # Scheduler
 limit = 0
 checkpoints_folder = ""
 mixed_precision = "bf16" #"fp16"
-gradient_accumulation_steps = 2
+gradient_accumulation_steps = 1
 accelerator = Accelerator(
     mixed_precision=mixed_precision,
     gradient_accumulation_steps=gradient_accumulation_steps
@@ -522,6 +522,15 @@ else:
                (1 + math.cos(math.pi * decay_ratio))
     
     lr_scheduler = LambdaLR(optimizer, lambda step: lr_schedule(step) / base_learning_rate)
+
+    num_params = sum(p.numel() for p in unet.parameters())
+    print(f"[rank {accelerator.process_index}] total params: {num_params}")
+    # Проверка на NaN/Inf
+    for name, param in unet.named_parameters():
+        if torch.isnan(param).any() or torch.isinf(param).any():
+            print(f"[rank {accelerator.process_index}] NaN/Inf in {name}")
+            # Опционально: заменить на нормальные значения
+            #param.data = torch.randn_like(param) * 0.01
     unet, optimizer, lr_scheduler = accelerator.prepare(unet, optimizer, lr_scheduler)
 
 # Регистрация хуков ПОСЛЕ prepare