Create merge_70b_shards_v2.py

prepared_sft_data/merge_70b_shards_v2.py

@@ -0,0 +1,61 @@
+# ==============================================================================
+# COPYRIGHT (C) 2025-2026 KONSTANTIN VLADIMIROVICH GRABKO. ALL RIGHTS RESERVED.
+# PATENT PENDING | CMS MANHATTAN JIRACK TECHNOLOGY
+# ==============================================================================
+import os
+import torch
+import glob
+from safetensors.torch import load_file, save_file
+from tqdm import tqdm
+
+# --- ПУТИ ---
+# Твой обученный чекпоинт (результат Full SFT)
+SFT_CHECKPOINT_PATH = "/content/full_checkpoints_70b/jirack_70b_full_step_200.safetensors"
+# Папка с твоими оригинальными 30 шардами
+ORIGINAL_SHARDS_DIR = "/content/JiRack_BitNet_70B_Packed/checkpoints/checkpoint-220000"
+# Куда сохранить результат
+OUTPUT_DIR = "/content/JiRack_70B_SFT_Merged"
+
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+def merge_shards():
+    print(f"🚀 Загрузка SFT чекпоинта: {SFT_CHECKPOINT_PATH}")
+    sft_weights = load_file(SFT_CHECKPOINT_PATH, device="cpu")
+    
+    # Получаем список всех шардов оригинальной модели
+    shard_files = sorted(glob.glob(f"{ORIGINAL_SHARDS_DIR}/*.safetensors"))
+    
+    print(f"📦 Найдено шардов для обработки: {len(shard_files)}")
+    
+    for shard_path in tqdm(shard_files, desc="Merging Shards"):
+        shard_name = os.path.basename(shard_path)
+        # Загружаем оригинальный шард
+        current_shard = load_file(shard_path, device="cpu")
+        updated_shard = {}
+        merge_count = 0
+        
+        for key, weight in current_shard.items():
+            # Убираем префикс 'model.', если он есть в ключах чекпоинта, но нет в шардах (или наоборот)
+            # Мы ищем точное совпадение ключа в sft_weights
+            
+            # Проверяем ключ как есть
+            if key in sft_weights:
+                updated_shard[key] = sft_weights[key]
+                merge_count += 1
+            # Проверяем с учетом возможной разницы в префиксах (model.layers... vs layers...)
+            elif key.replace("model.", "") in sft_weights:
+                updated_shard[key] = sft_weights[key.replace("model.", "")]
+                merge_count += 1
+            else:
+                # Если веса не обучались (не попали в SFT чекпоинт), оставляем оригинал
+                updated_shard[key] = weight
+        
+        # Сохраняем обновленный шард
+        save_path = os.path.join(OUTPUT_DIR, shard_name)
+        save_file(updated_shard, save_path)
+        # print(f"✅ {shard_name}: обновлено {merge_count} тензоров")
+
+    print(f"\n✨ Мердж завершен! Готовая модель здесь: {OUTPUT_DIR}")
+
+if __name__ == "__main__":
+    merge_shards()