Upload 2 files

load_SlimPajama_JiRackTernary_405b.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 - 405b 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_405b import JiRackTernaryMoE405B , JiRackMoEConfig
+
+# --- КОНФИГУРАЦИЯ 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 = JiRackMoEConfig(num_hidden_layers=160)
+    model = JiRackTernaryMoE405B(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[!] Остановка по требованию пользователя. Прогресс сохранен.")

train_405b_heavy_mixed_val_data.py

@@ -0,0 +1,139 @@
+# ==============================================================================
+# 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
+# ==============================================================================
+# COPYRIGHT (C) 2025 KONSTANTIN VLADIMIROVICH GRABKO. ALL RIGHTS RESERVED.
+# PATENT PENDING | CMS MANHATTAN JIRACK TECHNOLOGY | VERSION 405B MoE
+# Optimized for Multi-GPU ROCm/CUDA Clusters (Tesla M10 Support)
+# ==============================================================================
+
+import torch
+import torch.nn as nn
+import os
+import random
+import json
+from torch.utils.data import DataLoader, IterableDataset
+from transformers import AutoTokenizer, get_cosine_schedule_with_warmup
+from datasets import load_dataset
+from accelerate import Accelerator # Заменяем DataParallel на более мощный инструмент
+import sys
+
+# Импорт вашей архитектуры 405B MoE
+from JiRackTernaryPyTorch_405b import JiRackTernaryMoE405B, JiRackMoEConfig
+
+# --- КОНФИГУРАЦИЯ CMS MANHATTAN ---
+MODEL_ID = "./models/jirack_405b_init"
+CULTURAL_DATA = "cultural_finetune.jsonl"
+SLIM_PAJAMA = "cerebras/SlimPajama-627B" # Основной массив общих знаний
+CHECKPOINT_DIR = "checkpoints_jirack_405b"
+
+MIX_RATIO = 0.40         # 40% Культурный код / 60% SlimPajama
+BATCH_SIZE = 1           # Только 1 при таком масштабе
+GRAD_ACCUM_STEPS = 64    # Огромное накопление для стабильности 405B
+LEARNING_RATE = 2.0e-6   # Ультра-низкий LR для MoE
+BLOCK_SIZE = 1024        # Для M10 лучше держать 1k-2k
+
+# --- УМНЫЙ МИКСЕР ДАННЫХ (SlimPajama + Client Code) ---
+class CMSMoEMixer(IterableDataset):
+    def __init__(self, tokenizer, client_file, pj_link, mix_ratio=0.4):
+        self.tokenizer = tokenizer
+        self.mix_ratio = mix_ratio
+        print(f">>> [CMS] Streaming SlimPajama-627B...")
+        self.pj_stream = load_dataset(pj_link, split="train", streaming=True)
+        
+        self.cultural_data = []
+        if os.path.exists(client_file):
+            with open(client_file, 'r', encoding='utf-8') as f:
+                for line in f:
+                    self.cultural_data.append(json.loads(line))
+            print(f">>> [CMS] Loaded {len(self.cultural_data)} cultural samples.")
+
+    def __iter__(self):
+        pj_iter = iter(self.pj_stream)
+        while True:
+            if random.random() < self.mix_ratio and self.cultural_data:
+                sample = random.choice(self.cultural_data)
+                text = f"Question: {sample['question']}\nAnswer: {sample['answer']}"
+            else:
+                try:
+                    sample = next(pj_iter)
+                    text = sample['text']
+                except StopIteration:
+                    pj_iter = iter(self.pj_stream)
+                    continue
+
+            tokens = self.tokenizer(
+                text, truncation=True, max_length=BLOCK_SIZE, padding="max_length", return_tensors="pt"
+            )
+            yield {"input_ids": tokens["input_ids"].squeeze(0), "labels": tokens["input_ids"].squeeze(0)}
+
+def train():
+    # Инициализация Accelerator (автоматически распределит 405B по всем GPU)
+    accelerator = Accelerator(gradient_accumulation_steps=GRAD_ACCUM_STEPS)
+    device = accelerator.device
+    
+    # 1. Токенайзер
+    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B")
+    if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token
+
+    # 2. Модель 405B MoE
+    # Для M10 важно: модель должна загружаться в bfloat16 или float16
+    config = JiRackMoEConfig()
+    model = JiRackTernaryMoE405B(config)
+    
+    # Включаем Gradient Checkpointing (жизненно важно для 405B)
+    model.gradient_checkpointing_enable()
+
+    # 3. Данные
+    dataset = CMSMoEMixer(tokenizer, CULTURAL_DATA, SLIM_PAJAMA, mix_ratio=MIX_RATIO)
+    loader = DataLoader(dataset, batch_size=BATCH_SIZE)
+
+    # 4. Оптимизатор
+    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=0.1)
+    
+    # Подготовка через Accelerator
+    model, optimizer, loader = accelerator.prepare(model, optimizer, loader)
+
+    print(f"\n--- [CMS MANHATTAN] 405B MoE ENGINE START ---")
+    print(f"GPUs: {accelerator.num_processes} | Strategy: Mixed MoE (SlimPajama + Cultural)")
+
+    model.train()
+    for step, batch in enumerate(loader):
+        with accelerator.accumulate(model):
+            outputs = model(**batch)
+            loss = outputs.loss
+            accelerator.backward(loss)
+            
+            # Клиппинг градиентов (защита экспертов от взрыва)
+            if accelerator.sync_gradients:
+                accelerator.clip_grad_norm_(model.parameters(), 1.0)
+            
+            optimizer.step()
+            optimizer.zero_grad()
+
+        if step % 10 == 0 and accelerator.is_main_process:
+            print(f"CMS 405B | Step: {step} | Loss: {loss.item():.4f} | LR: {LEARNING_RATE}")
+
+        # Сохранение чекпоинта
+        if step > 0 and step % 250 == 0 and accelerator.is_main_process:
+            save_path = os.path.join(CHECKPOINT_DIR, f"step_{step}")
+            accelerator.save_state(save_path)
+            print(f">>> [CMS] Checkpoint saved: {save_path}")
+            torch.cuda.empty_cache()
+
+if __name__ == "__main__":
+    # Оптимизация аллокатора CUDA для Tesla M10
+    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+    try:
+        train()
+    except Exception as e:
+        print(f"FATAL: {e}")
+        sys.exit(1)