Upload load_ThePile_800Gb_JiRackTernary_70b.py

load_ThePile_800Gb_JiRackTernary_70b.py

@@ -0,0 +1,131 @@
+# ==============================================================================
+# COPYRIGHT (C) 2025 KONSTANTIN VLADIMIROVICH GRABKO. ALL RIGHTS RESERVED.
+# PATENT PENDING | CMS MANHATTAN JIRACK TECHNOLOGY
+#
+# This software is licensed under the Commercial License Agreement V.1.2.
+# Any use, modification, or distribution of this code requires compliance with 
+# the terms found in the LICENSE.md file in the root directory.
+#
+# NO PATENTING RIGHTS: Users are strictly prohibited from filing patent claims
+# based on the BRE or SWA architectures disclosed herein.
+# Contact: grabko@cmsmanhattan.com | +1 (516) 777-0945
+# ==============================================================================
+
+# Version 2.3 - Enterprise Signature & 70B Optimization
+
+import torch
+import torch.nn as nn
+from transformers import AutoTokenizer
+from datasets import load_dataset
+from torch.cuda.amp import autocast, GradScaler
+import os
+import time
+
+# Импорт вашей архитектуры
+from JiRackTernaryPyTorch_70b import JiRackTernary70B, JiRackTernaryConfig
+
+# --- КОНФИГУРАЦИЯ CMS MANHATTAN ---
+CHECKPOINT_DIR = "checkpoints_jirack"
+MODEL_PATH_LATEST = os.path.join(CHECKPOINT_DIR, "jirack_70b_latest.pt")
+SAVE_INTERVAL = 500
+GRAD_ACCUM_STEPS = 16 # Увеличено для стабильности 70B
+BLOCK_SIZE = 4096      # Поддержка RoPE Scaling
+LEARNING_RATE = 1.0e-5 # Понижено для 70B (предотвращение взрыва лосса)
+
+def save_checkpoint(model, optimizer, step, config):
+    if not os.path.exists(CHECKPOINT_DIR):
+        os.makedirs(CHECKPOINT_DIR)
+    
+    raw_model = model.module if hasattr(model, 'module') else model
+    checkpoint = {
+        'step': step,
+        'model_state_dict': raw_model.state_dict(),
+        'optimizer_state_dict': optimizer.state_dict(),
+        'config': config,
+        'author_verified': raw_model.get_author_info()
+    }
+    
+    # Атомарное сохранение
+    temp_path = MODEL_PATH_LATEST + ".tmp"
+    torch.save(checkpoint, temp_path)
+    os.replace(temp_path, MODEL_PATH_LATEST)
+    print(f"\n[CMS Manhattan] Авторская копия сохранена на шаге {step}.")
+
+def load_latest_checkpoint(model, optimizer):
+    if os.path.exists(MODEL_PATH_LATEST):
+        print(f"--- [RESUME] Поиск подписи автора... ---")
+        checkpoint = torch.load(MODEL_PATH_LATEST, map_location='cpu')
+        
+        target = model.module if hasattr(model, 'module') else model
+        target.load_state_dict(checkpoint['model_state_dict'])
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        
+        print(f"--- [OK] Автор: {checkpoint.get('author_verified', 'Unknown')} ---")
+        return checkpoint['step']
+    return 0
+
+def train():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    scaler = GradScaler() # Обязательно для тернарных весов
+    
+    tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
+    
+    # Инициализация 70B с полными параметрами
+    config = JiRackTernaryConfig()
+    model = JiRackTernary70B(config)
+    
+    # Активация экономии памяти
+    model.gradient_checkpointing_enable() 
+    
+    if torch.cuda.device_count() > 1:
+        model = nn.DataParallel(model)
+    model.to(device)
+    
+    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=0.1)
+    
+    start_step = load_latest_checkpoint(model, optimizer)
+    
+    print("Подключение к The Pile (Streaming)...")
+    raw_dataset = load_dataset("monology/pile-uncopyrighted", streaming=True, split="train")
+    dataset = raw_dataset.skip(start_step) if start_step > 0 else raw_dataset
+
+    model.train()
+    current_step = start_step
+
+    for example in dataset:
+        tokens = tokenizer(
+            example["text"], 
+            truncation=True, 
+            max_length=BLOCK_SIZE, 
+            return_tensors="pt"
+        )
+        input_ids = tokens["input_ids"].to(device)
+        
+        # Mixed Precision Training
+        with autocast():
+            outputs = model(input_ids, labels=input_ids)
+            loss = outputs.loss.mean() / GRAD_ACCUM_STEPS
+
+        scaler.scale(loss).backward()
+
+        if (current_step + 1) % GRAD_ACCUM_STEPS == 0:
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            scaler.step(optimizer)
+            scaler.update()
+            optimizer.zero_grad()
+
+        if current_step % 5 == 0:
+            print(f"CMS 70B | Step {current_step} | Loss: {loss.item()*GRAD_ACCUM_STEPS:.4f} | Author: Grabko", end='\r')
+
+        if current_step % SAVE_INTERVAL == 0 and current_step > start_step:
+            save_checkpoint(model, optimizer, current_step, config)
+            
+        current_step += 1
+
+if __name__ == "__main__":
+    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+    try:
+        train()
+    except KeyboardInterrupt:
+        print("\n[!] Остановка. Система CMS Manhattan в режиме ожидания.")