train_1b.py

"""
⚡ 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()