Create src/training/trainer.py

src/training/trainer.py

@@ -0,0 +1,75 @@
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from torch.nn import functional as F
+from torch.cuda.amp import GradScaler, autocast
+import time
+
+# --- Sovereign Training Utilities ---
+
+def get_batch(data, block_size, batch_size, device):
+    """Generates a small batch of data of inputs x and targets y."""
+    ix = torch.randint(len(data) - block_size, (batch_size,))
+    x = torch.stack([data[i:i+block_size] for i in ix])
+    y = torch.stack([data[i+1:i+block_size+1] for i in ix])
+    x, y = x.to(device), y.to(device)
+    return x, y
+
+class SovereignTrainer:
+    def __init__(self, model, optimizer, config, device):
+        self.model = model.to(device)
+        self.optimizer = optimizer
+        self.config = config
+        self.device = device
+        self.scaler = GradScaler() # For Mixed-Precision Training
+        self.block_size = config['model_params']['n_positions']
+
+    def train_step(self, x, y):
+        self.optimizer.zero_grad(set_to_none=True)
+        
+        # 1. Mixed Precision Forward Pass (Speeds up training on modern GPUs)
+        with autocast():
+            logits, loss = self.model(x, y)
+        
+        # 2. Backpropagation with Scaling
+        self.scaler.scale(loss).backward()
+        
+        # 3. Gradient Clipping (Prevents 'Exploding Gradients' in scratch builds)
+        self.scaler.unscale_(self.optimizer)
+        torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
+        
+        # 4. Optimizer Step
+        self.scaler.step(self.optimizer)
+        self.scaler.update()
+        
+        return loss.item()
+
+    def run_pretraining(self, train_data, val_data, max_iters=10000):
+        """The core pre-training loop for ARAVALLI-1."""
+        print(f"Sovereign Pre-training Initiated on {self.device}...")
+        self.model.train()
+        
+        start_time = time.time()
+        for iter in range(max_iters):
+            # Fetch batch
+            xb, yb = get_batch(train_data, self.block_size, 32, self.device)
+            
+            # Execute step
+            loss = self.train_step(xb, yb)
+            
+            # Logging and Checkpointing
+            if iter % 100 == 0 or iter == max_iters - 1:
+                dt = time.time() - start_time
+                print(f"Iter {iter}: Loss {loss:.4f} | Time: {dt:.2f}s")
+                # Trigger Sovereign Checkpoint (to be signed by pyHanko)
+                self.save_checkpoint(iter)
+                start_time = time.time()
+
+    def save_checkpoint(self, iter):
+        checkpoint = {
+            'model': self.model.state_dict(),
+            'optimizer': self.optimizer.state_dict(),
+            'config': self.config,
+            'iter': iter,
+        }
+        torch.save(checkpoint, f"data/processed/ckpt_iter_{iter}.pt")