Update train_jirack_accelerate.py

train_jirack_accelerate.py

@@ -1,3 +1,5 @@
+#%%writefile  train_jirack_accelerate_v3.py
+
 # =============================================================================
 # COPYRIGHT © 2025-2026  Konstantin Vladimirovich Grabko. ALL RIGHTS RESERVED.
 # CMS Manhattan JiRack Technology — PATENT PENDING
@@ -10,11 +12,11 @@
 # Unauthorized commercial use is strictly prohibited.
 # Contact: grabko@cmsmanhattan.com
 # =============================================================================
-
 import os
 # Включаем оптимизацию памяти для ROCm/HIP ДО импорта torch!
 os.environ["PYTORCH_HIP_ALLOC_CONF"] = "expandable_segments:True"
 
+import re
 import glob
 import math
 import torch
@@ -36,6 +38,14 @@ class SingleShardDataset(Dataset):
     def __getitem__(self, idx):
         return self.data[idx].long()
 
+# --- Helpers for Natural Sorting ---
+def natural_sort_key(s):
+    return [int(text) if text.isdigit() else text.lower() for text in re.split(r'(\d+)', s)]
+
+def extract_shard_number(filename):
+    match = re.search(r'model_weights_shard_(\d+)\.pt', filename)
+    return int(match.group(1)) if match else 0
+
 # --- 2. Main Training Function ---
 def train():
     grad_accumulation_steps = 24
@@ -47,9 +57,12 @@ def train():
         gradient_accumulation_steps=grad_accumulation_steps
     )
 
-    pt_chunks_mask = "pretraindata/jirack_pretrain_chunk_*.pt"
-    checkpoint_dir = "checkpoints"
-    pt_files = sorted(glob.glob(pt_chunks_mask))
+    # Определяем абсолютные пути относительно расположения самого скрипта
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    pt_chunks_mask = os.path.join(script_dir, "pretraindata/jirack_pretrain_chunk_*.pt")
+    checkpoint_dir = os.path.join(script_dir, "checkpoints")
+    
+    pt_files = sorted(glob.glob(pt_chunks_mask), key=natural_sort_key)
 
     if not pt_files:
         raise FileNotFoundError(f"No chunk files found matching the mask: {pt_chunks_mask}")
@@ -64,13 +77,29 @@ def train():
     config = TernaryConfig()
     model = TernaryTransformer3B(config)
 
-    # === PINPOINT CHECKPOINT LOADING ===
-    checkpoint_load_path = "model_weights.pt"
+    # === DYNAMIC CHECKPOINT DISCOVERY (ABSOLUTE PATHS) ===
+    checkpoint_load_path = os.path.join(script_dir, "model_weights.pt")
+    start_shard_idx = 0
+
+    # Сканируем папку checkpoints по абсолютному пути
+    shard_checkpoints = glob.glob(os.path.join(checkpoint_dir, "model_weights_shard_*.pt"))
+    
+    if shard_checkpoints:
+        shard_checkpoints = sorted(shard_checkpoints, key=natural_sort_key)
+        latest_shard_checkpoint = shard_checkpoints[-1]
+        completed_shards = extract_shard_number(latest_shard_checkpoint)
+        
+        checkpoint_load_path = latest_shard_checkpoint
+        start_shard_idx = completed_shards  # Если закончили шард 2, индекс следующего равен 2 (3-й шард)
+
+    # Загружаем найденные веса
     if os.path.exists(checkpoint_load_path):
         if accelerator.is_local_main_process:
             print(f"-> Loading saved weights from: {checkpoint_load_path}")
         state_dict = torch.load(checkpoint_load_path, map_location="cpu", weights_only=True)
         model.load_state_dict(state_dict)
+        if accelerator.is_local_main_process and start_shard_idx > 0:
+            print(f"-> Resuming training from Shard {start_shard_idx + 1} (Skipping first {start_shard_idx} shards)")
     else:
         if accelerator.is_local_main_process:
             print(f"-> Checkpoint not found at {checkpoint_load_path}, training will start from scratch.")
@@ -81,14 +110,12 @@ def train():
         print("-> Gradient Checkpointing enabled.")
 
     criterion = nn.CrossEntropyLoss()
-
-    # Переносим модель на устройство через accelerator
     model = accelerator.prepare(model)
 
     # === ADAFACTOR CONFIGURATION FOR GPU ===
     optimizer = Adafactor(
         model.parameters(),
-        lr=2e-4,  # Peak LR
+        lr=2e-4,
         weight_decay=0.01,
         relative_step=False,
         scale_parameter=False,
@@ -98,7 +125,6 @@ def train():
     # === CALCULATING AND INITIALIZING SCHEDULER WITH WARMUP ===
     scheduler = None
     if USE_COSINE_SCHEDULER:
-        # Примерный расчет для сквозного или локального графика
         steps_per_shard = math.ceil(2000 / (batch_size * grad_accumulation_steps))
         total_steps = len(pt_files) * steps_per_shard
         num_warmup_steps = int(0.05 * total_steps)
@@ -110,25 +136,24 @@ def train():
         )
 
         if accelerator.is_local_main_process:
-            print(f"-> Scheduler ACTIVATED.")
-            print(f"   Total training steps: {total_steps}")
-            print(f"   Warmup steps: {num_warmup_steps}")
+            print(f"-> Scheduler ACTIVATED. Total steps: {total_steps}")
 
-    # Подготавливаем оптимизатор и планировщик
     if scheduler is not None:
         optimizer, scheduler = accelerator.prepare(optimizer, scheduler)
     else:
         optimizer = accelerator.prepare(optimizer)
 
     model.train()
-    shard_counter = 0
-
     if accelerator.is_local_main_process:
         print("Starting training...")
         os.makedirs(checkpoint_dir, exist_ok=True)
 
     # === ОСНОВНОЙ ЦИКЛ ПО ШАРДАМ ===
-    for shard_path in pt_files:
+    for shard_counter, shard_path in enumerate(pt_files):
+        # Пропускаем уже обработанные шарды
+        if shard_counter < start_shard_idx:
+            continue
+
         shard_name = os.path.basename(shard_path)
         if accelerator.is_local_main_process:
             print(f"\n[Shard {shard_counter + 1}/{len(pt_files)}] {shard_name}")
@@ -138,16 +163,14 @@ def train():
         train_loader = DataLoader(
             shard_dataset,
             batch_size=batch_size,
-            shuffle=True,
-            num_workers=2,
-            pin_memory=True
+            shuffle=True
         )
 
         train_loader = accelerator.prepare(train_loader)
 
         progress_bar = tqdm(
             train_loader,
-            desc=f"Processing {shard_name}",
+            desc=f"Training {shard_name}",
             disable=not accelerator.is_local_main_process
         )
 
@@ -162,7 +185,6 @@ def train():
                 loss = criterion(logits.reshape(-1, logits.size(-1)), targets.reshape(-1))
 
                 accelerator.backward(loss)
-
                 optimizer.step()
 
                 if scheduler is not None and accelerator.sync_gradients:
@@ -184,14 +206,12 @@ def train():
                         "lr": f"{current_lr:.2e}"
                     })
 
-        shard_counter += 1
-
-        # Сохранение весов прямо в общую папку с указанием номера шарда в имени файла
+        # Сохраняем веса с абсолютным путем
         if accelerator.is_local_main_process:
-            shard_checkpoint_path = os.path.join(checkpoint_dir, f"model_weights_shard_{shard_counter}.pt")
+            shard_checkpoint_path = os.path.join(checkpoint_dir, f"model_weights_shard_{shard_counter + 1}.pt")
             unwrapped_model = accelerator.unwrap_model(model)
             torch.save(unwrapped_model.state_dict(), shard_checkpoint_path)
-            print(f"✅  Saved after shard {shard_counter}: {shard_checkpoint_path}")
+            print(f"✅   Saved after shard {shard_counter + 1}: {shard_checkpoint_path}")
 
     # Финальное сохранение
     accelerator.wait_for_everyone()