Upload load_SlimPajama_JiRackTernary_236b.py

load_SlimPajama_JiRackTernary_236b.py

@@ -0,0 +1,146 @@
+# ==============================================================================
+# 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 3.0 - 236B Extreme Scale & SlimPajama Integration
+# Optimized for 4x Tesla M10 (128GB VRAM Total)
+
+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
+
+# Импорт вашей новой архитектуры 140B
+from JiRackTernaryPyTorch_236b import JiRackTernary236B, JiRackTernaryConfig
+
+# --- КОНФИГУРАЦИЯ CMS MANHATTAN ---
+CHECKPOINT_DIR = "checkpoints_jirack_140b"
+MODEL_PATH_LATEST = os.path.join(CHECKPOINT_DIR, "jirack_140b_latest.pt")
+SAVE_INTERVAL = 250    # Чаще для такой огромной модели
+GRAD_ACCUM_STEPS = 32  # Увеличено для стабильности на 140B
+BLOCK_SIZE = 2048      # Уменьшено с 4096 до 2048 для экономии VRAM на M10
+LEARNING_RATE = 5.0e-6 # Еще ниже для 160 слоев (очень глубокая сеть)
+
+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] Авторская копия 140B сохранена на шаге {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] Модель 140B загружена. Автор: {checkpoint.get('author_verified')} ---")
+        return checkpoint['step']
+    return 0
+
+def train():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    scaler = GradScaler()
+    
+    # Используем токенайзер Llama-3 (он лучше подходит для больших словарей)
+    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B")
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # Инициализация 140B (160 слоев)
+    config = JiRackTernaryConfig(num_hidden_layers=160)
+    model = JiRackTernary236B(config)
+    
+    # Принудительно включаем Gradient Checkpointing для выживания на M10
+    model.gradient_checkpointing_enable() 
+    
+    if torch.cuda.device_count() > 1:
+        print(f"Обнаружено {torch.cuda.device_count()} GPU. Активация DataParallel...")
+        model = nn.DataParallel(model)
+    
+    model.to(device)
+    
+    # Weight decay важен для предотвращения переполнения тернарных весов
+    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=0.1)
+    
+    start_step = load_latest_checkpoint(model, optimizer)
+    
+    # ПОДКЛЮЧЕНИЕ SLIM PAJAMA (627B tokens)
+    print("Подключение к SlimPajama-627B (Streaming mode)...")
+    dataset = load_dataset(
+        "cerebras/SlimPajama-627B", 
+        streaming=True, 
+        split="train"
+    )
+    
+    if start_step > 0:
+        dataset = dataset.skip(start_step)
+
+    model.train()
+    current_step = start_step
+
+    for example in dataset:
+        # SlimPajama использует поле 'text'
+        tokens = tokenizer(
+            example["text"], 
+            truncation=True, 
+            max_length=BLOCK_SIZE, 
+            padding="max_length",
+            return_tensors="pt"
+        )
+        
+        input_ids = tokens["input_ids"].to(device)
+        
+        # Mixed Precision для ускорения на Tesla M10
+        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)
+            # Градиентный клиппинг обязателен для глубоких 160-слойных сетей
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            scaler.step(optimizer)
+            scaler.update()
+            optimizer.zero_grad()
+
+        if current_step % 1 == 0:
+            print(f"CMS 140B | Step {current_step} | Loss: {loss.item()*GRAD_ACCUM_STEPS:.4f} | VRAM: {torch.cuda.memory_reserved() / 1e9:.1f}GB", 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__":
+    # Настройка аллокатора для Tesla M10 (предотвращает фрагментацию)
+    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True,max_split_size_mb:128"
+    try:
+        train()
+    except KeyboardInterrupt:
+        print("\n[!] Остановка по требованию пользователя. Прогресс сохранен.")