Create merge_sft_sandwich.py

merge_sft_sandwich.py

@@ -0,0 +1,57 @@
+#%%writefile merge_sft_v30.py
+import torch, glob, os
+from safetensors.torch import load_file, save_file
+from tqdm import tqdm
+
+# Настройки
+BASE_MODEL_DIR = "/content/JiRack_BitNet_405B_Packed"
+SFT_CHECKPOINT = "/content/drive/MyDrive/405b_sft_results/sft_405b_v30_step_250.safetensors"
+OUTPUT_DIR = "/content/JiRack_405B_SFT_Merged"
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+# Маппинг слоев (тот же, что был в v30)
+layer_map = {0:0, 1:1, 2:124, 3:125}
+
+def merge():
+    print(f"💉 Загружаем SFT веса из {SFT_CHECKPOINT}...")
+    sft_sd = load_file(SFT_CHECKPOINT)
+    
+    shard_files = sorted(glob.glob(f"{BASE_MODEL_DIR}/**/*.safetensors", recursive=True))
+    
+    for shard in tqdm(shard_files, desc="Merging shards"):
+        sd = load_file(shard)
+        new_sd = {}
+        changed = False
+        
+        for name, param in sd.items():
+            # Проверяем, есть ли этот параметр в наших обученных весах
+            if name in sft_sd:
+                new_sd[name] = sft_sd[name].to(param.dtype)
+                changed = True
+                continue
+            
+            # Проверяем маппинг слоев (для слоев 0, 1, 124, 125)
+            # В SFT весах они лежат как layers.0, layers.1, layers.2, layers.3
+            if "layers." in name:
+                parts = name.split('.')
+                old_idx = int(parts[1])
+                # Находим, какой индекс в SFT соответствует этому слою в базе
+                sft_idx = next((k for k, v in layer_map.items() if v == old_idx), None)
+                
+                if sft_idx is not None:
+                    sft_name = name.replace(f"layers.{old_idx}", f"layers.{sft_idx}")
+                    if sft_name in sft_sd:
+                        new_sd[name] = sft_sd[sft_name].to(param.dtype)
+                        changed = True
+                        continue
+            
+            new_sd[name] = param
+        
+        # Сохраняем обновленный шард
+        output_shard = os.path.join(OUTPUT_DIR, os.path.basename(shard))
+        save_file(new_sd, output_shard)
+        
+    print(f"✅ Мердж завершен! Модель сохранена в {OUTPUT_DIR}")
+
+if __name__ == "__main__":
+    merge()