Update train_jirack_accelerate.py

train_jirack_accelerate.py

@@ -12,7 +12,7 @@
 # =============================================================================
 
 import os
-# Enable memory optimization for ROCm/HIP BEFORE importing torch!
+# Включаем оптимизацию памяти для ROCm/HIP ДО импорта torch!
 os.environ["PYTORCH_HIP_ALLOC_CONF"] = "expandable_segments:True"
 
 import glob
@@ -25,7 +25,7 @@ from tqdm import tqdm
 from transformers import Adafactor, get_cosine_schedule_with_warmup
 
 # --- GLOBAL CONSTANTS ---
-USE_COSINE_SCHEDULER = True  # Set to False if you want to disable the scheduler (warmup + cosine)
+USE_COSINE_SCHEDULER = False  # Set to False if you want to disable the scheduler (warmup + cosine)
 
 # --- 1. Lightweight Dataset ---
 class SingleShardDataset(Dataset):
@@ -38,8 +38,8 @@ class SingleShardDataset(Dataset):
 
 # --- 2. Main Training Function ---
 def train():
-    grad_accumulation_steps = 1
-    batch_size = 1
+    grad_accumulation_steps = 24
+    batch_size = 8
 
     # Initialize Accelerator
     accelerator = Accelerator(
@@ -58,7 +58,7 @@ def train():
         print(f"Shards found: {len(pt_files)}")
         print("Initializing JiRack 3.3B...")
 
-    # Strictly import your classes from the local file
+    # Импортируем строго ваши классы из локального файла
     from JiRackNative_3b import TernaryTransformer3B, TernaryConfig
 
     config = TernaryConfig()
@@ -74,7 +74,6 @@ def train():
     else:
         if accelerator.is_local_main_process:
             print(f"-> Checkpoint not found at {checkpoint_load_path}, training will start from scratch.")
-    # ==================================
 
     model.gradient_checkpointing_enable()
 
@@ -83,7 +82,7 @@ def train():
 
     criterion = nn.CrossEntropyLoss()
 
-    # First, move the model to the ROCm/HIP device via accelerator
+    # Переносим модель на устройство через accelerator
     model = accelerator.prepare(model)
 
     # === ADAFACTOR CONFIGURATION FOR GPU ===
@@ -99,15 +98,11 @@ def train():
     # === CALCULATING AND INITIALIZING SCHEDULER WITH WARMUP ===
     scheduler = None
     if USE_COSINE_SCHEDULER:
-        # Calculate total number of actual weight updates (optimizer steps)
-        # 2000 rows per shard / batch_size 1 / grad_accumulation_steps 4 = 500 steps per shard.
+        # Примерный расчет для сквозного или локального графика
         steps_per_shard = math.ceil(2000 / (batch_size * grad_accumulation_steps))
         total_steps = len(pt_files) * steps_per_shard
-
-        # Set warmup (e.g., 5% of total training steps)
         num_warmup_steps = int(0.05 * total_steps)
 
-        # Combined scheduler: smoothly increases LR up to 2e-4, then decays following a cosine curve to 0
         scheduler = get_cosine_schedule_with_warmup(
             optimizer=optimizer,
             num_warmup_steps=num_warmup_steps,
@@ -119,7 +114,7 @@ def train():
             print(f"   Total training steps: {total_steps}")
             print(f"   Warmup steps: {num_warmup_steps}")
 
-    # Prepare optimizer and scheduler via accelerator
+    # Подготавливаем оптимизатор и планировщик
     if scheduler is not None:
         optimizer, scheduler = accelerator.prepare(optimizer, scheduler)
     else:
@@ -130,7 +125,9 @@ def train():
 
     if accelerator.is_local_main_process:
         print("Starting training...")
+        os.makedirs(checkpoint_dir, exist_ok=True)
 
+    # === ОСНОВНОЙ ЦИКЛ ПО ШАРДАМ ===
     for shard_path in pt_files:
         shard_name = os.path.basename(shard_path)
         if accelerator.is_local_main_process:
@@ -168,7 +165,6 @@ def train():
 
                 optimizer.step()
 
-                # Step the scheduler only during an actual gradient sync update
                 if scheduler is not None and accelerator.sync_gradients:
                     if not getattr(accelerator, "optimizer_step_was_skipped", False):
                         scheduler.step()
@@ -180,7 +176,6 @@ def train():
                 ppl = math.exp(avg_loss) if avg_loss < 20 else float('inf')
 
                 if accelerator.is_local_main_process:
-                    # Extract current LR to display on the tqdm panel
                     current_lr = scheduler.get_last_lr()[0] if scheduler is not None else optimizer.param_groups[0]['lr']
                     progress_bar.set_postfix({
                         "loss": f"{loss.item():.4f}",
@@ -191,22 +186,20 @@ def train():
 
         shard_counter += 1
 
-        # Save checkpoint after each shard
+        # Сохранение весов прямо в общую папку с указанием номера шарда в имени файла
         if accelerator.is_local_main_process:
-            current_checkpoint_path = os.path.join(checkpoint_dir, f"jirack_shard_{shard_counter}")
-            os.makedirs(current_checkpoint_path, exist_ok=True)
+            shard_checkpoint_path = os.path.join(checkpoint_dir, f"model_weights_shard_{shard_counter}.pt")
             unwrapped_model = accelerator.unwrap_model(model)
-            torch.save(unwrapped_model.state_dict(), os.path.join(current_checkpoint_path, "model_weights.pt"))
-            print(f"✅ Saved after shard {shard_counter}")
+            torch.save(unwrapped_model.state_dict(), shard_checkpoint_path)
+            print(f"✅  Saved after shard {shard_counter}: {shard_checkpoint_path}")
 
-    # Final saving step
+    # Финальное сохранение
     accelerator.wait_for_everyone()
     if accelerator.is_local_main_process:
-        final_dir = os.path.join(checkpoint_dir, "jirack_final_3b")
-        os.makedirs(final_dir, exist_ok=True)
+        final_path = os.path.join(checkpoint_dir, "model_final_weights.pt")
         unwrapped_model = accelerator.unwrap_model(model)
-        torch.save(unwrapped_model.state_dict(), os.path.join(final_dir, "model_final_weights.pt"))
-        print("Final save completed.")
+        torch.save(unwrapped_model.state_dict(), final_path)
+        print(f"Final save completed: {final_path}")
 
 if __name__ == "__main__":
-    train()
+    train()