Upload folder using huggingface_hub

.gitignore

@@ -17,3 +17,4 @@ metadata/
 store/
 preprocessed_configs/
 .early.coverage
+.ipynb_checkpoints/

data/ohlc_stats.npz

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f6a84b63ec605e83a655f404bc89d825aa8ffbb5ac3ea24c7d2197324646d016
+oid sha256:fc81a5fbd342767ca491eceeca805f54dd9ffe4ff0bb723cdda26e52f54f914d
 size 1660

log.log

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:21f6421b07eb49c1e0a5518a628403ce0ae7149fb81a600aebad2dfcaf0313c9
+oid sha256:e4ceaef802908dd650ce1ade210a0e827ec433b904adc3bf17c3d8a877e59ae6
 size 2854

train.py

@@ -38,7 +38,7 @@ from torch.optim import AdamW
 from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import ProjectConfiguration, set_seed
-from transformers import get_linear_schedule_with_warmup
+from transformers import get_cosine_schedule_with_warmup
 
 # Logging
 from tqdm.auto import tqdm
@@ -119,7 +119,7 @@ def quantile_pinball_loss(preds: torch.Tensor,
     return sum(losses) / mask.sum().clamp_min(1.0)
 
 
-def create_balanced_split(dataset, val_ratio: float = 0.1, seed: int = 42):
+def create_balanced_split(dataset, n_val_per_class: int = 1, seed: int = 42):
     """
     Create train/val split with balanced classes in validation set.
     Uses dataset's internal file_class_map for speed (no file loading).
@@ -153,10 +153,10 @@ def create_balanced_split(dataset, val_ratio: float = 0.1, seed: int = 42):
     train_indices = []
     val_indices = []
 
-    # For each class, take val_ratio samples for validation
+    # For each class, take n_val_per_class samples for validation
     for class_id, indices in class_to_indices.items():
         random.shuffle(indices)
-        n_val = max(1, int(len(indices) * val_ratio))  # At least 1 sample per class
+        n_val = min(len(indices), n_val_per_class) # Ensure we don't take more than we have
         val_indices.extend(indices[:n_val])
         train_indices.extend(indices[n_val:])
 
@@ -329,7 +329,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--pin_memory", dest="pin_memory", action="store_true", default=True)
     parser.add_argument("--no-pin_memory", dest="pin_memory", action="store_false")
     parser.add_argument("--resume_from_checkpoint", type=str, default=None, help="Path to checkpoint or 'latest'")
-    parser.add_argument("--val_split", type=float, default=0.1, help="Fraction of data for validation (default 0.1)")
+    parser.add_argument("--val_samples_per_class", type=int, default=1, help="Number of validation samples per class (default 1)")
     parser.add_argument("--val_every", type=int, default=1000, help="Run validation every N steps (default 1000)")
     return parser.parse_args()
 
@@ -374,7 +374,10 @@ def main() -> None:
         logger.info("Initialized with CLI arguments.")
         tensorboard_dir.mkdir(parents=True, exist_ok=True)
         checkpoint_dir.mkdir(parents=True, exist_ok=True)
-        accelerator.init_trackers("oracle_training")
+        from datetime import datetime
+        run_name = f"oracle_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+        accelerator.init_trackers(run_name)
+        logger.info(f"📊 TensorBoard run: {run_name}")
 
     device = accelerator.device
     
@@ -413,7 +416,16 @@ def main() -> None:
     logger.info(f"Initializing Encoders with dtype={init_dtype}...")
     
     # Encoders
-    multi_modal_encoder = MultiModalEncoder(dtype=init_dtype)
+    logger.info("Initializing Shared MultiModalEncoder (SigLIP) on GPU...")
+    # Initialize ONCE on GPU for both WalletEncoder (dims) and Collator (encoding)
+    multi_modal_encoder = MultiModalEncoder(dtype=init_dtype, device="cuda")
+    
+    # Use this shared instance for setting the worker encoder (num_workers=0 optimization)
+    # This avoids loading a second copy of SigLIP
+    logger.info("Setting shared encoder for collator...")
+    from models.multi_modal_processor import MultiModalEncoder as CollatorEncoder
+    _set_worker_encoder(multi_modal_encoder)
+
     time_encoder = ContextualTimeEncoder(dtype=init_dtype)
     token_encoder = TokenEncoder(multi_dim=multi_modal_encoder.embedding_dim, dtype=init_dtype)
     wallet_encoder = WalletEncoder(encoder=multi_modal_encoder, dtype=init_dtype)
@@ -430,18 +442,6 @@ def main() -> None:
         max_seq_len=max_seq_len
     )
 
-    # --- OPTIMIZATION: Pre-load SigLIP encoder to avoid lazy-loading on first batch ---
-    # This moves the ~8s model load from first batch to startup (where it's expected)
-    logger.info("Pre-loading SigLIP encoder for collator (avoids first-batch delay)...")
-    from models.multi_modal_processor import MultiModalEncoder as CollatorEncoder
-    collator_encoder = CollatorEncoder(
-        model_id=collator.model_id,
-        dtype=init_dtype,
-        device="cuda"  # Use GPU for encoding (requires num_workers=0)
-    )
-    _set_worker_encoder(collator_encoder)
-    logger.info("SigLIP encoder pre-loaded successfully.")
-
     # ==========================================================================
     # OFFLINE MODE: No DB connections during training for maximum GPU utilization
     # ==========================================================================
@@ -475,18 +475,18 @@ def main() -> None:
         raise RuntimeError("Dataset is empty.")
 
     # --- NEW: Create balanced train/val split ---
-    logger.info(f"Creating {1-args.val_split:.0%}/{args.val_split:.0%} train/val split with balanced classes...")
+    logger.info(f"Creating balanced split with {args.val_samples_per_class} validation samples per class...")
     train_indices, val_indices, class_distribution = create_balanced_split(
-        dataset, val_ratio=args.val_split, seed=seed
+        dataset, n_val_per_class=args.val_samples_per_class, seed=seed
     )
 
     # Log class distribution (use set for O(1) lookup)
     train_set = set(train_indices)
     logger.info(f"Total samples: {len(dataset)}, Train: {len(train_indices)}, Val: {len(val_indices)}")
     for class_id, indices in sorted(class_distribution.items()):
-        n_val = int(len(indices) * args.val_split)  # Approximate from split ratio
+        n_val = min(len(indices), args.val_samples_per_class)
         n_train = len(indices) - n_val
-        logger.info(f"  Class {class_id}: {len(indices)} total (~{n_train} train, ~{n_val} val)")
+        logger.info(f"  Class {class_id}: {len(indices)} total (~{n_train} train, {n_val} val)")
 
     # --- Compute class weights for loss weighting ---
     num_classes = max(class_distribution.keys()) + 1 if class_distribution else 7
@@ -544,7 +544,7 @@ def main() -> None:
     # Validation dataloader (no shuffle, no weighted sampling)
     val_dl_kwargs = dict(
         dataset=val_dataset,
-        batch_size=batch_size,
+        batch_size=1,  # Force batch size 1 for validation to prevent OOM
         shuffle=False,
         num_workers=int(args.num_workers),
         pin_memory=bool(args.pin_memory),
@@ -588,7 +588,7 @@ def main() -> None:
     max_train_steps = epochs * num_update_steps_per_epoch
     num_warmup_steps = int(max_train_steps * warmup_ratio)
 
-    scheduler = get_linear_schedule_with_warmup(
+    scheduler = get_cosine_schedule_with_warmup(
         optimizer,
         num_warmup_steps=num_warmup_steps,
         num_training_steps=max_train_steps
@@ -623,7 +623,9 @@ def main() -> None:
     if args.resume_from_checkpoint:
         if args.resume_from_checkpoint == "latest":
              if latest_checkpoint:
-                logger.info(f"Resuming from latest checkpoint: {latest_checkpoint}")
+                logger.info("=" * 60)
+                logger.info(f"🔄 RESUMING FROM LATEST CHECKPOINT: {latest_checkpoint}")
+                logger.info("=" * 60)
                 accelerator.load_state(str(latest_checkpoint))
              else:
                 logger.warning("Resume requested but no checkpoint found in dir. Starting fresh.")
@@ -631,12 +633,15 @@ def main() -> None:
              # Specific path
              custom_ckpt = Path(args.resume_from_checkpoint)
              if custom_ckpt.exists():
-                  logger.info(f"Resuming from specific checkpoint: {custom_ckpt}")
+                  logger.info("=" * 60)
+                  logger.info(f"🔄 RESUMING FROM CHECKPOINT: {custom_ckpt}")
+                  logger.info("=" * 60)
                   accelerator.load_state(str(custom_ckpt))
+                  logger.info("✅ Model, optimizer, scheduler, and dataloader states restored.")
              else:
                   raise FileNotFoundError(f"Checkpoint not found at {custom_ckpt}")
     else:
-        logger.info("No resume flag provided. Starting fresh.")
+        logger.info("🆕 No resume flag provided. Starting fresh training run.")
 
     # --- 7. Training Loop ---
     total_steps = 0
@@ -831,15 +836,14 @@ def main() -> None:
                 logger.info(f"Epoch {epoch+1} complete. Avg loss: {avg_loss:.6f}")
                 accelerator.log({"train/loss_epoch": avg_loss}, step=total_steps)
                 
-                # Save Checkpoint at end of epoch (REMOVED: saving every epoch is too much)
+                # Save Checkpoint at end of epoch
                 save_path = checkpoint_dir / f"epoch_{epoch+1}"
-                # accelerator.save_state(output_dir=str(save_path))
-                # hf_save_path = save_path / "hf_model"
-                # unwrapped_model = accelerator.unwrap_model(model)
-                # unwrapped_model.save_pretrained(str(hf_save_path))
+                accelerator.save_state(output_dir=str(save_path))
+                hf_save_path = save_path / "hf_model"
+                unwrapped_model = accelerator.unwrap_model(model)
+                unwrapped_model.save_pretrained(str(hf_save_path))
                 
-                # logger.info(f"Saved and HF-style model (EOF) to {save_path}")
-                pass
+                logger.info(f"Saved epoch checkpoint and HF-style model to {save_path}")
         else:
             if accelerator.is_main_process:
                 logger.warning(f"Epoch {epoch+1}: No valid batches processed.")

train.sh

@@ -1,22 +1,24 @@
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+
 accelerate launch train.py \
-  --epochs 1 \
-  --batch_size 8 \
+  --epochs 5 \
+  --batch_size 2 \
   --learning_rate 1e-4 \
   --warmup_ratio 0.1 \
-  --grad_accum_steps 2 \
+  --grad_accum_steps 8 \
   --max_grad_norm 1.0 \
   --seed 42 \
-  --log_every 3 \
-  --save_every 2000 \
+  --log_every 10 \
+  --save_every 1000 \
   --tensorboard_dir runs/oracle \
   --checkpoint_dir checkpoints \
   --mixed_precision bf16 \
   --max_seq_len 4096 \
-  --horizons_seconds 30 60 120 240 420 \
+  --horizons_seconds 300 900 1800 3600 7200 \
   --quantiles 0.1 0.5 0.9 \
   --ohlc_stats_path ./data/ohlc_stats.npz \
   --num_workers 0 \
-  --pin_memory \
-  --val_split 0.1 \
-  --val_every 50 \
+  --val_samples_per_class 2 \
+  --val_every 100 \
+  --resume_from_checkpoint checkpoints/checkpoint-3000 \
   "$@"