""" ⚡ ULTRON Systems // High-Efficiency Autonomous Colab GPU Training Loop Orchestrates: - Month-by-month streaming dataset ingestion - Mixed Precision (FP16) with GradScaler - Gradient Accumulation for effective batch size - Auto-resume from Google Drive / local checkpoints - Real-time VRAM telemetry and directional accuracy metrics """ import os import sys import gc import time import signal import argparse import torch import torch.nn as nn from torch.utils.data import DataLoader from typing import Optional from config import DatasetConfig, ModelConfig, TrainingConfig from dataset import StreamingShardManager from model import OHLCVTransformer from loss import MultiTaskFinancialLoss class Trainer: def __init__(self, data_cfg: DatasetConfig, model_cfg: ModelConfig, train_cfg: TrainingConfig): self.data_cfg = data_cfg self.model_cfg = model_cfg self.train_cfg = train_cfg # Setup device self.device = torch.device(train_cfg.device) print(f"⚡ [ULTRON] Initializing training on compute device: {self.device}") if self.device.type == "cuda": print(f"⚡ [ULTRON] GPU: {torch.cuda.get_device_name(0)} | VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB") # Determine active checkpoint directory if os.path.exists("/content/drive/MyDrive"): self.checkpoint_dir = train_cfg.checkpoint_dir print(f"⚡ [ULTRON] Google Drive detected. Checkpoints will persist to: {self.checkpoint_dir}") else: self.checkpoint_dir = train_cfg.local_checkpoint_dir print(f"⚡ [ULTRON] Local mode active. Checkpoints will save to: {self.checkpoint_dir}") os.makedirs(self.checkpoint_dir, exist_ok=True) # Instantiate Model, Loss, Optimizer, Scaler self.model = OHLCVTransformer(model_cfg).to(self.device) self.criterion = MultiTaskFinancialLoss().to(self.device) self.optimizer = torch.optim.AdamW( self.model.parameters(), lr=train_cfg.learning_rate, weight_decay=train_cfg.weight_decay, betas=(0.9, 0.98) ) self.scaler = torch.cuda.amp.GradScaler(enabled=train_cfg.fp16 and self.device.type == "cuda") self.global_step = 0 self.current_epoch = 0 self.best_dir_acc = 0.0 # Register graceful termination handlers signal.signal(signal.SIGINT, self._handle_exit) signal.signal(signal.SIGTERM, self._handle_exit) # Attempt auto-resume self._load_checkpoint() def _get_lr(self, step: int) -> float: """Linear warmup followed by cosine decay.""" if step < self.train_cfg.warmup_steps: return self.train_cfg.learning_rate * float(step) / float(max(1, self.train_cfg.warmup_steps)) import math progress = float(step - self.train_cfg.warmup_steps) / float(max(1, 10000 - self.train_cfg.warmup_steps)) cosine_decay = 0.5 * (1.0 + math.cos(math.pi * min(1.0, progress))) return max(self.train_cfg.min_lr, self.train_cfg.learning_rate * cosine_decay) def _save_checkpoint(self, is_best: bool = False): state = { "global_step": self.global_step, "epoch": self.current_epoch, "best_dir_acc": self.best_dir_acc, "model_state": self.model.state_dict(), "optimizer_state": self.optimizer.state_dict(), "scaler_state": self.scaler.state_dict(), "model_cfg": self.model_cfg, "train_cfg": self.train_cfg } latest_path = os.path.join(self.checkpoint_dir, "checkpoint_latest.pt") torch.save(state, latest_path) if is_best: best_path = os.path.join(self.checkpoint_dir, "checkpoint_best.pt") torch.save(state, best_path) print(f"🌟 [ULTRON] Saved new BEST model checkpoint to {best_path}") print(f"💾 [ULTRON] Checkpoint synced (Step {self.global_step}).") def _load_checkpoint(self): latest_path = os.path.join(self.checkpoint_dir, "checkpoint_latest.pt") if os.path.exists(latest_path): print(f"⚡ [ULTRON] Found existing checkpoint at {latest_path}. Resuming state...") try: ckpt = torch.load(latest_path, map_location=self.device) self.model.load_state_dict(ckpt["model_state"]) self.optimizer.load_state_dict(ckpt["optimizer_state"]) if "scaler_state" in ckpt and self.scaler.is_enabled(): self.scaler.load_state_dict(ckpt["scaler_state"]) self.global_step = ckpt.get("global_step", 0) self.current_epoch = ckpt.get("epoch", 0) self.best_dir_acc = ckpt.get("best_dir_acc", 0.0) print(f"✅ [ULTRON] Resumed successfully from Step {self.global_step} (Epoch {self.current_epoch}).") except Exception as e: print(f"⚠️ [ULTRON] Checkpoint load failed: {e}. Starting fresh.") else: print("⚡ [ULTRON] No previous checkpoint found. Initializing fresh weights.") def _handle_exit(self, signum, frame): print(f"\n⚠️ [ULTRON] Interrupt signal received ({signum}). Securing weights...") self._save_checkpoint() sys.exit(0) def train(self): print("\n" + "="*70) print("⚡ ULTRON // CONVOLUTIONAL TRANSFORMER AUTONOMOUS TRAINING ONLINE") print("="*70) manager = StreamingShardManager(self.data_cfg) self.model.train() for epoch in range(self.current_epoch, self.train_cfg.max_epochs): self.current_epoch = epoch print(f"\n🚀 [ULTRON] Beginning Epoch {epoch + 1}/{self.train_cfg.max_epochs}") # Stream monthly shards one by one for shard_idx, shard_dataset in enumerate(manager.stream_monthly_shards()): loader = DataLoader( shard_dataset, batch_size=self.train_cfg.batch_size, shuffle=True, num_workers=2, pin_memory=True if self.device.type == "cuda" else False, drop_last=True ) shard_loss = 0.0 shard_acc = 0.0 steps_in_shard = 0 t0 = time.time() self.optimizer.zero_grad() for batch_idx, (batch_x, batch_y) in enumerate(loader): batch_x = batch_x.to(self.device, non_blocking=True) batch_y = batch_y.to(self.device, non_blocking=True) # Dynamic learning rate schedule lr = self._get_lr(self.global_step) for param_group in self.optimizer.param_groups: param_group["lr"] = lr # Mixed precision forward pass with torch.cuda.amp.autocast(enabled=self.train_cfg.fp16 and self.device.type == "cuda"): pred_returns, pred_logits = self.model(batch_x) loss, metrics = self.criterion(pred_returns, pred_logits, batch_y) # Scale loss for gradient accumulation loss_scaled = loss / self.train_cfg.grad_accum_steps # Backward pass self.scaler.scale(loss_scaled).backward() shard_loss += metrics["total_loss"] shard_acc += metrics["dir_acc_bce"] steps_in_shard += 1 # Optimizer step on accumulation boundaries if (batch_idx + 1) % self.train_cfg.grad_accum_steps == 0: self.scaler.unscale_(self.optimizer) torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) self.scaler.step(self.optimizer) self.scaler.update() self.optimizer.zero_grad() self.global_step += 1 # Logging telemetry if self.global_step % self.train_cfg.log_interval == 0: vram_gb = torch.cuda.memory_allocated() / 1e9 if self.device.type == "cuda" else 0.0 rate = (self.train_cfg.batch_size * self.train_cfg.grad_accum_steps) / max(1e-4, (time.time() - t0)) print( f"⚡ [Step {self.global_step:05d}] " f"Loss: {metrics['total_loss']:.4f} " f"(Reg: {metrics['reg_loss']:.4f}, Dir: {metrics['dir_loss']:.4f}) | " f"Dir Acc: {metrics['dir_acc_bce']:.1f}% | " f"LR: {lr:.2e} | " f"VRAM: {vram_gb:.2f}GB | " f"{rate:.1f} win/s" ) t0 = time.time() # Periodic checkpointing if self.global_step % self.train_cfg.save_every_steps == 0: avg_acc = shard_acc / max(1, steps_in_shard) is_best = avg_acc > self.best_dir_acc if is_best: self.best_dir_acc = avg_acc self._save_checkpoint(is_best=is_best) # Cap steps per shard if configured to cycle through years faster if steps_in_shard >= self.train_cfg.max_steps_per_shard: break # Clean up shard memory explicitly del loader, shard_dataset if self.device.type == "cuda": torch.cuda.empty_cache() gc.collect() # End of epoch checkpoint self._save_checkpoint() print("\n✅ [ULTRON] Training loop complete. Final weights secured.") def parse_args(): parser = argparse.ArgumentParser(description="ULTRON OHLCV Transformer Trainer") parser.add_argument("--start_year", type=int, default=2021, help="Start year of shards") parser.add_argument("--end_year", type=int, default=2025, help="End year of shards") parser.add_argument("--batch_size", type=int, default=64, help="Batch size per GPU step") parser.add_argument("--grad_accum", type=int, default=4, help="Gradient accumulation steps") parser.add_argument("--lr", type=float, default=3e-4, help="Learning rate") parser.add_argument("--epochs", type=int, default=5, help="Total training epochs") return parser.parse_args() if __name__ == "__main__": args = parse_args() data_cfg = DatasetConfig(start_year=args.start_year, end_year=args.end_year) model_cfg = ModelConfig() train_cfg = TrainingConfig( batch_size=args.batch_size, grad_accum_steps=args.grad_accum, learning_rate=args.lr, max_epochs=args.epochs ) trainer = Trainer(data_cfg, model_cfg, train_cfg) trainer.train()