Upload load_JiRack5_RedPajama_140b.py

load_JiRack5_RedPajama_140b.py

@@ -0,0 +1,87 @@
+# ==============================================================================
+# COPYRIGHT (C) 2025 KONSTANTIN VLADIMIROVICH GRABKO. ALL RIGHTS RESERVED.
+# PATENT PENDING | CMS MANHATTAN JIRACK TECHNOLOGY
+# ==============================================================================
+# Version 4.1 - 140B Dense | RedPajama-Data-1T Integration
+# Optimized for 160 Layers & SwiGLU-Attention (SWA) Fusion
+
+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 the Dense Architecture
+from JiRackPyTorch_GPT5_class_140b import JiRackPyTorch 
+
+# --- CMS MANHATTAN CONFIGURATION ---
+CHECKPOINT_DIR = "checkpoints_jirack_140b_dense"
+SAVE_INTERVAL = 500    
+GRAD_ACCUM_STEPS = 64  # High accumulation to stabilize the 160-layer gradient
+BLOCK_SIZE = 2048      
+LEARNING_RATE = 4.0e-6 
+
+def train():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    scaler = GradScaler()
+    
+    # --- REDPAJAMA INTEGRATION ---
+    # Using the INCITE-Base-3B tokenizer for its high-efficiency vocabulary
+    tokenizer = AutoTokenizer.from_pretrained("togethercomputer/RedPajama-INCITE-Base-3B")
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # Load RedPajama-Data-1T in Streaming Mode to save Disk I/O
+    print("Connecting to RedPajama-Data-1T (Streaming Mode)...")
+    dataset = load_dataset("togethercomputer/RedPajama-Data-1T", split="train", streaming=True)
+
+    # Initialize 140B Dense Flagship
+    model = JiRackPyTorch() 
+    model.gradient_checkpointing_enable() # BRE Strategy: Trade compute for VRAM
+    
+    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.01)
+    
+    model.train()
+    print("--- Training Started: JiRack 140B Dense ---")
+    
+    for current_step, example in enumerate(dataset):
+        # RedPajama uses the "text" key for content
+        tokens = tokenizer(
+            example["text"], 
+            truncation=True, 
+            max_length=BLOCK_SIZE, 
+            padding="max_length", 
+            return_tensors="pt"
+        )
+        
+        input_ids = tokens["input_ids"].to(device)
+        
+        # SWA Fusion Forward Pass (Mixed Precision)
+        with autocast(dtype=torch.bfloat16):
+            logits, loss, _ = model(input_ids, targets=input_ids)
+            loss = loss.mean() / GRAD_ACCUM_STEPS
+
+        scaler.scale(loss).backward()
+
+        # Step Optimization
+        if (current_step + 1) % GRAD_ACCUM_STEPS == 0:
+            scaler.unscale_(optimizer)
+            # Tight clipping for deep 140B networks
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5) 
+            scaler.step(optimizer)
+            scaler.update()
+            optimizer.zero_grad()
+
+        if current_step % 5 == 0:
+            vram = torch.cuda.memory_reserved() / 1e9
+            print(f"CMS 140B | Step {current_step} | Loss: {loss.item()*GRAD_ACCUM_STEPS:.4f} | VRAM: {vram:.1f}GB", end='\r')
+
+if __name__ == "__main__":
+    # Allocator tuning for Tesla M10 32GB
+    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True,max_split_size_mb:64"
+    train()