""" ⚡ ULTRON Systems // 1.189 Billion Parameter Foundation Model Training Pipeline Optimized for NVIDIA A100 (80GB) Tensor Core GPUs with Bfloat16 Precision. """ import os import sys import gc import time import math import signal import argparse import numpy as np import torch import torch.nn as nn from torch.utils.data import DataLoader from config import DatasetConfig from config_1b import Model1BConfig, Training1BConfig from dataset import StreamingShardManager from model_1b import FinancialConformer1B class Trainer1B: def __init__(self, model_cfg: Model1BConfig, train_cfg: Training1BConfig): self.model_cfg = model_cfg self.train_cfg = train_cfg self.device = torch.device(train_cfg.device) print(f"⚡ [ULTRON 1B] Initializing Foundation Model Training on: {self.device}") if self.device.type == "cuda": gpu_name = torch.cuda.get_device_name(0) vram_total = torch.cuda.get_device_properties(0).total_memory / 1e9 print(f"⚡ [ULTRON 1B] Hardware: {gpu_name} | VRAM: {vram_total:.2f} GB") # Enable TF32 for high throughput on Ampere torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True # Checkpoint directory if os.path.exists("/content/drive/MyDrive"): self.checkpoint_dir = train_cfg.checkpoint_dir else: self.checkpoint_dir = train_cfg.local_checkpoint_dir os.makedirs(self.checkpoint_dir, exist_ok=True) print(f"⚡ [ULTRON 1B] Checkpoints will be saved to: {self.checkpoint_dir}") # Instantiate Model print("⚡ [ULTRON 1B] Instantiating 1.189B Parameter Architecture...") self.model = FinancialConformer1B(model_cfg).to(self.device) total_params = self.model.count_parameters() print(f"✅ [ULTRON 1B] Model Loaded: {total_params:,} parameters ({total_params/1e9:.3f} Billion)") # Optimizer: 8-bit AdamW via bitsandbytes (drops static optimizer VRAM from 9.5GB to 2.4GB) try: import bitsandbytes as bnb print("⚡ [ULTRON 1B] Utilizing 8-bit AdamW (bitsandbytes) for ultra-low VRAM footprint.") self.optimizer = bnb.optim.AdamW8bit( self.model.parameters(), lr=train_cfg.learning_rate, weight_decay=train_cfg.weight_decay, betas=(0.9, 0.95), eps=1e-8 ) except Exception as e: print(f"⚠️ [ULTRON 1B] bitsandbytes unavailable ({e}). Falling back to standard AdamW.") self.optimizer = torch.optim.AdamW( self.model.parameters(), lr=train_cfg.learning_rate, weight_decay=train_cfg.weight_decay, betas=(0.9, 0.95), eps=1e-8 ) # Losses self.huber_loss = nn.HuberLoss(delta=0.01) self.bce_loss = nn.BCEWithLogitsLoss() self.global_step = 0 self.current_epoch = 0 self.current_shard_idx = 0 self.best_dir_acc = 0.0 self.hf_token = os.environ.get("HF_TOKEN", "") self.hf_repo_id = os.environ.get("HF_REPO_ID", "Ultron28/ohlcv-transformer") signal.signal(signal.SIGINT, self._handle_exit) signal.signal(signal.SIGTERM, self._handle_exit) self._load_checkpoint() import threading def _heartbeat_worker(): while True: time.sleep(30) try: with open("/content/heartbeat.txt", "w") as f: f.write(f"HEARTBEAT:{time.time()}:STEP:{self.global_step}:SHARD:{self.current_shard_idx}\n") except Exception: pass threading.Thread(target=_heartbeat_worker, daemon=True).start() def _get_lr(self, step: int) -> float: if step < self.train_cfg.warmup_steps: return self.train_cfg.learning_rate * float(step) / float(max(1, self.train_cfg.warmup_steps)) progress = float(step - self.train_cfg.warmup_steps) / float(max(1, 20000 - 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 _upload_to_hf(self, file_path: str, repo_filename: str): if not self.hf_token: return def _worker(): try: from huggingface_hub import HfApi api = HfApi(token=self.hf_token) api.upload_file( path_or_fileobj=file_path, path_in_repo=f"checkpoints_1b/{repo_filename}", repo_id=self.hf_repo_id, repo_type="model" ) print(f"☁️ [ULTRON 1B] Background Sync Complete: {repo_filename} -> Hugging Face ({self.hf_repo_id}).", flush=True) except Exception as e: print(f"⚠️ [ULTRON 1B] HF background upload notice: {e}", flush=True) import threading threading.Thread(target=_worker, daemon=True).start() def _log_telemetry(self, entry: dict): telemetry_file = "/content/telemetry_1b.jsonl" try: import json with open(telemetry_file, "a") as f: f.write(json.dumps(entry) + "\n") except Exception: pass def _sync_remote_log(self): if not self.hf_token: return def _worker(): try: from huggingface_hub import HfApi api = HfApi(token=self.hf_token) for fname in ["train_1b.log", "telemetry_1b.jsonl"]: fpath = f"/content/{fname}" if os.path.exists(fpath): api.upload_file( path_or_fileobj=fpath, path_in_repo=f"logs/{fname}", repo_id=self.hf_repo_id, repo_type="model" ) except Exception: pass import threading threading.Thread(target=_worker, daemon=True).start() def _save_checkpoint(self, is_best: bool = False, push_hf: bool = True): # Convert floating point weights to bfloat16 to drop file size from 7.2GB to 2.4GB (3x faster upload) model_state = {k: v.to(torch.bfloat16) if v.is_floating_point() else v for k, v in self.model.state_dict().items()} state = { "global_step": self.global_step, "epoch": self.current_epoch, "current_shard_idx": self.current_shard_idx, "best_dir_acc": self.best_dir_acc, "model_state": model_state, "optimizer_state": self.optimizer.state_dict(), "model_cfg": self.model_cfg, "train_cfg": self.train_cfg } latest_path = os.path.join(self.checkpoint_dir, "checkpoint_1b_latest.pt") torch.save(state, latest_path) print(f"💾 [ULTRON 1B] Checkpoint saved locally (Step {self.global_step}, Shard {self.current_shard_idx}).", flush=True) if is_best: best_path = os.path.join(self.checkpoint_dir, "checkpoint_1b_best.pt") torch.save(state, best_path) print(f"🌟 [ULTRON 1B] New BEST Foundation Model Checkpoint: {best_path}", flush=True) if push_hf: self._upload_to_hf(best_path, "checkpoint_1b_best.pt") if push_hf: self._upload_to_hf(latest_path, "checkpoint_1b_latest.pt") self._sync_remote_log() def _load_checkpoint(self): latest_path = os.path.join(self.checkpoint_dir, "checkpoint_1b_latest.pt") # If not present on local disk, pull from Hugging Face Model Hub if not os.path.exists(latest_path) and self.hf_token: try: from huggingface_hub import hf_hub_download print(f"☁️ [ULTRON 1B] Checking Hugging Face ({self.hf_repo_id}) for remote checkpoint...") downloaded = hf_hub_download( repo_id=self.hf_repo_id, filename="checkpoints_1b/checkpoint_1b_latest.pt", token=self.hf_token, local_dir=self.checkpoint_dir ) if os.path.exists(downloaded): latest_path = downloaded print(f"✅ [ULTRON 1B] Cloud checkpoint retrieved from Hugging Face: {latest_path}") except Exception as e: print(f"ℹ️ [ULTRON 1B] No existing remote checkpoint on Hugging Face: {e}") if os.path.exists(latest_path): print(f"⚡ [ULTRON 1B] Resuming state from checkpoint: {latest_path} ...") try: ckpt = torch.load(latest_path, map_location=self.device, weights_only=False) self.model.load_state_dict(ckpt["model_state"]) self.optimizer.load_state_dict(ckpt["optimizer_state"]) self.global_step = ckpt.get("global_step", 0) self.current_epoch = ckpt.get("epoch", 0) self.current_shard_idx = ckpt.get("current_shard_idx", 0) self.best_dir_acc = ckpt.get("best_dir_acc", 0.0) print(f"✅ [ULTRON 1B] State resumed at Step {self.global_step}, Shard {self.current_shard_idx}.") except Exception as e: print(f"⚠️ [ULTRON 1B] Checkpoint load failed: {e}. Starting fresh.") else: print("⚡ [ULTRON 1B] Initializing fresh weights for 1.189B model.") def _handle_exit(self, signum, frame): print(f"\n⚠️ [ULTRON 1B] Termination signal ({signum}). Securing weights to disk and cloud...") self._save_checkpoint(push_hf=True) sys.exit(0) def train(self): print("\n" + "="*75) print("⚡ ULTRON // 1.189B FINANCIAL CONFORMER TRAINING ONLINE (A100 BF16)") print("="*75) data_cfg = DatasetConfig( start_year=2010, end_year=2023, seq_len=self.train_cfg.seq_len, window_stride=self.train_cfg.window_stride, pred_horizons=self.train_cfg.pred_horizons ) manager = StreamingShardManager(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 1B] Beginning Foundation Epoch {epoch + 1}/{self.train_cfg.max_epochs}") for shard_idx, shard_dataset in enumerate(manager.stream_monthly_shards(start_shard_idx=self.current_shard_idx), start=self.current_shard_idx): shard_name = manager.parquet_files[shard_idx] if shard_idx < len(manager.parquet_files) else f"shard_{shard_idx}" print(f"\n⚡ [ULTRON 1B] Processing Shard {shard_idx + 1}/{len(manager.parquet_files)}: {shard_name}") 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 ) 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) lr = self._get_lr(self.global_step) for param_group in self.optimizer.param_groups: param_group["lr"] = lr # Native Bfloat16 Mixed Precision on A100 Tensor Cores with torch.amp.autocast('cuda', dtype=torch.bfloat16): pred_returns, pred_logits = self.model(batch_x) reg_loss = self.huber_loss(pred_returns, batch_y) target_dir = (batch_y > 0.0).float() dir_loss = self.bce_loss(pred_logits, target_dir) total_loss = (reg_loss + 0.5 * dir_loss) / self.train_cfg.grad_accum_steps total_loss.backward() if (batch_idx + 1) % self.train_cfg.grad_accum_steps == 0: torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) self.optimizer.step() self.optimizer.zero_grad() self.global_step += 1 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)) with torch.no_grad(): horizon_matches = ((pred_logits > 0.0) == target_dir).float() dir_acc = horizon_matches.mean().item() * 100.0 # Per-horizon directional accuracies: [1m, 3m, 5m, 15m, 30m, 60m, 390m] horizon_accs = [round(horizon_matches[:, h].mean().item() * 100.0, 1) for h in range(horizon_matches.shape[1])] log_entry = { "step": self.global_step, "epoch": epoch + 1, "shard_idx": shard_idx, "shard_file": shard_name, "total_loss": round(total_loss.item()*self.train_cfg.grad_accum_steps, 4), "reg_loss": round(reg_loss.item(), 4), "dir_loss": round(dir_loss.item(), 4), "dir_acc_total": round(dir_acc, 2), "horizon_accs": horizon_accs, "lr": float(f"{lr:.2e}"), "vram_gb": round(vram_gb, 2), "rate_win_s": round(rate, 1), "timestamp": time.time() } self._log_telemetry(log_entry) h_str = f"1m:{horizon_accs[0]}% 5m:{horizon_accs[2]}% 15m:{horizon_accs[3]}% 1h:{horizon_accs[5]}%" print( f"⚡ [1B Step {self.global_step:05d}] " f"Loss: {total_loss.item()*self.train_cfg.grad_accum_steps:.4f} " f"(Reg: {reg_loss.item():.4f}, Dir: {dir_loss.item():.4f}) | " f"Dir Acc: {dir_acc:.1f}% ({h_str}) | " f"LR: {lr:.2e} | " f"VRAM: {vram_gb:.2f}GB | " f"{rate:.1f} win/s", flush=True ) t0 = time.time() if self.global_step % self.train_cfg.save_every_steps == 0: self._save_checkpoint(push_hf=True) del loader, shard_dataset if self.device.type == "cuda": torch.cuda.empty_cache() gc.collect() self.current_shard_idx = shard_idx + 1 self._save_checkpoint(push_hf=True) print(f"✅ [ULTRON 1B] Shard {shard_idx + 1}/{len(manager.parquet_files)} finished and backed up to Hugging Face Hub.") # End of epoch reset self.current_shard_idx = 0 self._save_checkpoint(push_hf=True) print("\n✅ [ULTRON 1B] Foundation Model Training Complete. All weights secured.") if __name__ == "__main__": model_cfg = Model1BConfig() train_cfg = Training1BConfig() trainer = Trainer1B(model_cfg, train_cfg) trainer.train()