Upload folder using huggingface_hub

data/ohlc_stats.npz

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

inference_utils.py

@@ -0,0 +1,34 @@
+import torch
+
+def transform_targets(targets):
+    """
+    Applies the log-transform used during training:
+    y_trans = sign(y) * log(1 + |y|)
+    
+    Args:
+        targets: torch.Tensor or float, raw returns (e.g. 1.5 for 150%)
+    Returns:
+        transformed targets in the same shape/type
+    """
+    if isinstance(targets, torch.Tensor):
+        return torch.sign(targets) * torch.log1p(torch.abs(targets))
+    else:
+        # Handle float/numpy
+        import numpy as np
+        return np.sign(targets) * np.log1p(np.abs(targets))
+
+def inverse_transform_targets(transformed_targets):
+    """
+    Inverts the log-transform to get back raw returns:
+    y = sign(y_trans) * (exp(|y_trans|) - 1)
+    
+    Args:
+        transformed_targets: torch.Tensor, model outputs (logits)
+    Returns:
+        raw returns (e.g. 1.5 for 150%)
+    """
+    if isinstance(transformed_targets, torch.Tensor):
+        return torch.sign(transformed_targets) * (torch.exp(torch.abs(transformed_targets)) - 1)
+    else:
+        import numpy as np
+        return np.sign(transformed_targets) * (np.exp(np.abs(transformed_targets)) - 1)

ingest.sh

@@ -20,7 +20,7 @@ error() { echo -e "${RED}[ERROR]${NC} $1"; exit 1; }
 #===============================================================================
 header "Step 5-6/7: Processing Epochs (Download → Ingest → Delete)"
 
-EPOCHS=(844 845 846)
+EPOCHS=(844 845 846 847 848 849)
 
 
 log "Processing epochs one at a time to minimize disk usage..."

train.py

@@ -1,4 +1,5 @@
 import os
+import json
 import argparse
 import math
 import logging
@@ -202,6 +203,8 @@ def run_validation(model, val_dataloader, accelerator, quantiles, quality_loss_f
             if labels_mask.sum() == 0:
                 return_loss = torch.tensor(0.0, device=accelerator.device)
             else:
+                # Log-transform targets for validation too (so val loss matches train loss scale)
+                labels = torch.sign(labels) * torch.log1p(torch.abs(labels))
                 return_loss = quantile_pinball_loss(preds, labels, labels_mask, quantiles)
 
             quality_loss = quality_loss_fn(quality_preds, quality_targets)
@@ -608,43 +611,102 @@ def main() -> None:
     # Load checkpoint if it exists
     starting_epoch = 0
     resume_step = 0
+    total_steps = 0
+    latest_checkpoint = None
     
     # Check for existing checkpoints
     if checkpoint_dir.exists():
-        # Look for subfolders named 'checkpoint-X' or 'epoch_X'
-        # Accelerate saves to folders.
         dirs = [d for d in checkpoint_dir.iterdir() if d.is_dir()]
         if dirs:
-            # Sort by modification time or name to find latest
-    # Sort by modification time or name to find latest
             dirs.sort(key=lambda x: x.stat().st_mtime)
             latest_checkpoint = dirs[-1]
     
     if args.resume_from_checkpoint:
+        resolved_ckpt = None
         if args.resume_from_checkpoint == "latest":
              if latest_checkpoint:
-                logger.info("=" * 60)
-                logger.info(f"🔄 RESUMING FROM LATEST CHECKPOINT: {latest_checkpoint}")
-                logger.info("=" * 60)
-                accelerator.load_state(str(latest_checkpoint))
+                resolved_ckpt = latest_checkpoint
              else:
                 logger.warning("Resume requested but no checkpoint found in dir. Starting fresh.")
         else:
-             # Specific path
              custom_ckpt = Path(args.resume_from_checkpoint)
              if custom_ckpt.exists():
-                  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.")
+                  resolved_ckpt = custom_ckpt
              else:
                   raise FileNotFoundError(f"Checkpoint not found at {custom_ckpt}")
+
+        if resolved_ckpt is not None:
+            logger.info("=" * 60)
+            logger.info(f"🔄 RESUMING FROM CHECKPOINT: {resolved_ckpt}")
+            logger.info("=" * 60)
+            accelerator.load_state(str(resolved_ckpt))
+
+            # Restore epoch / step counters from training_state.json
+            state_file = resolved_ckpt / "training_state.json"
+            if state_file.exists():
+                with open(state_file, "r") as f:
+                    training_state = json.load(f)
+                starting_epoch = training_state.get("epoch", 0)
+                resume_step = training_state.get("step_in_epoch", 0)
+                total_steps = training_state.get("global_step", 0)
+                logger.info(
+                    f"✅ Resumed: epoch={starting_epoch}, step_in_epoch={resume_step}, "
+                    f"global_step={total_steps}"
+                )
+            else:
+                # Try to infer step count from the restored scheduler state
+                inner_sched = scheduler.scheduler if hasattr(scheduler, 'scheduler') else scheduler
+                restored_step = getattr(inner_sched, 'last_epoch', 0)
+                if restored_step > 0:
+                    total_steps = restored_step
+                    logger.warning(
+                        f"⚠️  training_state.json not found. "
+                        f"Inferred global_step={total_steps} from scheduler state."
+                    )
+                else:
+                    logger.warning(
+                        "⚠️  training_state.json not found and scheduler step is 0. "
+                        "Epoch/step counters start from 0."
+                    )
+            logger.info("✅ Model, optimizer, scheduler, and dataloader states restored.")
+
+            # --- FIX: Rebuild scheduler to extend over new epochs ---
+            # Without this, the cosine schedule wraps/oscillates because the
+            # internal step counter (from past runs) exceeds the original
+            # max_train_steps (computed for just --epochs in this run).
+            # We extend the schedule: completed_steps → completed_steps + this_run_steps
+            # with NO warmup (model is already trained).
+            extended_total = total_steps + max_train_steps
+            logger.info(
+                f"♻️  Rebuilding LR schedule: "
+                f"completed={total_steps}, this_run={max_train_steps}, "
+                f"extended_total={extended_total} (no warmup)"
+            )
+
+            # Get the base optimizer (unwrap Accelerate wrapper)
+            base_opt = optimizer.optimizer if hasattr(optimizer, 'optimizer') else optimizer
+            new_sched = get_cosine_schedule_with_warmup(
+                base_opt,
+                num_warmup_steps=0,          # no re-warmup
+                num_training_steps=extended_total
+            )
+            # Fast-forward to current position
+            new_sched.last_epoch = total_steps
+            new_sched._step_count = total_steps + 1
+
+            # Swap into the Accelerate wrapper
+            if hasattr(scheduler, 'scheduler'):
+                scheduler.scheduler = new_sched
+            else:
+                scheduler = new_sched
+
+            # Log the resulting LR for verification
+            current_lr = base_opt.param_groups[0]['lr']
+            logger.info(f"♻️  LR after rebuild: {current_lr:.6e}")
     else:
         logger.info("🆕 No resume flag provided. Starting fresh training run.")
 
     # --- 7. Training Loop ---
-    total_steps = 0
     quality_loss_fn = nn.MSELoss()
     
     logger.info("***** Running training *****")
@@ -660,18 +722,21 @@ def main() -> None:
         epoch_loss = 0.0
         valid_batches = 0
         
+        # Skip already-processed batches when resuming mid-epoch
+        active_dataloader = dataloader
+        if epoch == starting_epoch and resume_step > 0:
+            logger.info(f"⏩ Skipping {resume_step} batches in epoch {epoch+1} (already processed)")
+            active_dataloader = accelerator.skip_first_batches(dataloader, num_batches=resume_step)
+
         # Tqdm only on main process
         progress_bar = tqdm(
-            dataloader, 
+            active_dataloader, 
             desc=f"Epoch {epoch+1}/{epochs}", 
             disable=not accelerator.is_local_main_process,
-            initial=resume_step # If you calculate resume_step from checkpoint
+            initial=resume_step if epoch == starting_epoch else 0
         )
 
         for step, batch in enumerate(progress_bar):
-            # Skip steps if resuming (Accelerate dataloader might handle this automatically if configured, 
-            # but 'skip_first_batches' is often manual. 
-            # For simplicity here, we assume load_state restored the dataloader iterator.)
             
             if batch is None:
                 continue
@@ -685,7 +750,8 @@ def main() -> None:
 
             grad_stats: Optional[Dict[str, float]] = None
             module_grad_stats: Dict[str, float] = {}
-            with accelerator.accumulate(model):
+            try:
+              with accelerator.accumulate(model):
                 outputs = model(batch)
                 
                 preds = outputs["quantile_logits"]
@@ -706,6 +772,11 @@ def main() -> None:
                         t_cutoffs[0] if t_cutoffs else "unknown",
                     )
 
+                # Log-transform targets to stabilize Class 5 gradients
+                # y_trans = sign(y) * log(1 + |y|)
+                # Compresses 10000x return (label ~15) to ~2.7
+                labels = torch.sign(labels) * torch.log1p(torch.abs(labels))
+
                 if labels_mask.sum() == 0:
                     return_loss = torch.tensor(0.0, requires_grad=True, device=accelerator.device)
                 else:
@@ -747,6 +818,26 @@ def main() -> None:
                 scheduler.step()
                 optimizer.zero_grad()
 
+            except torch.cuda.OutOfMemoryError:
+                # Log the offending batch for debugging
+                token_addresses = batch.get('token_addresses', [])
+                sample_indices = batch.get('sample_indices', [])
+                seq_len = batch.get('event_type_ids', torch.empty(0)).shape[-1] if 'event_type_ids' in batch else '?'
+                logger.warning(
+                    "⚠️ CUDA OOM — skipping batch! "
+                    "seq_len=%s sample_idx=%s token=%s | "
+                    "Clearing cache and continuing.",
+                    seq_len,
+                    sample_indices[0] if sample_indices else "unknown",
+                    token_addresses[0] if token_addresses else "unknown",
+                )
+                # Clean up to recover
+                optimizer.zero_grad(set_to_none=True)
+                torch.cuda.empty_cache()
+                if hasattr(torch.cuda, 'reset_peak_memory_stats'):
+                    torch.cuda.reset_peak_memory_stats()
+                continue
+
             # Logging
             if accelerator.sync_gradients:
                 total_steps += 1
@@ -794,7 +885,16 @@ def main() -> None:
                         save_path = checkpoint_dir / f"checkpoint-{total_steps}"
                         accelerator.save_state(output_dir=str(save_path))
 
-                        # NEW: Save in standard HF-loadable way
+                        # Save resume metadata
+                        state_file = save_path / "training_state.json"
+                        with open(state_file, "w") as f:
+                            json.dump({
+                                "epoch": epoch,
+                                "step_in_epoch": step + (resume_step if epoch == starting_epoch else 0) + 1,
+                                "global_step": total_steps,
+                            }, f, indent=2)
+
+                        # Save in standard HF-loadable way
                         hf_save_path = save_path / "hf_model"
                         unwrapped_model = accelerator.unwrap_model(model)
                         unwrapped_model.save_pretrained(str(hf_save_path))
@@ -839,6 +939,16 @@ def main() -> None:
                 # Save Checkpoint at end of epoch
                 save_path = checkpoint_dir / f"epoch_{epoch+1}"
                 accelerator.save_state(output_dir=str(save_path))
+
+                # Save resume metadata (epoch completed, so next epoch starts fresh)
+                state_file = save_path / "training_state.json"
+                with open(state_file, "w") as f:
+                    json.dump({
+                        "epoch": epoch + 1,
+                        "step_in_epoch": 0,
+                        "global_step": total_steps,
+                    }, f, indent=2)
+
                 hf_save_path = save_path / "hf_model"
                 unwrapped_model = accelerator.unwrap_model(model)
                 unwrapped_model.save_pretrained(str(hf_save_path))

train.sh

@@ -1,7 +1,7 @@
 export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
 
 accelerate launch train.py \
-  --epochs 5 \
+  --epochs 10 \
   --batch_size 2 \
   --learning_rate 1e-4 \
   --warmup_ratio 0.1 \
@@ -20,5 +20,5 @@ accelerate launch train.py \
   --num_workers 0 \
   --val_samples_per_class 2 \
   --val_every 100 \
-  --resume_from_checkpoint checkpoints/checkpoint-3000 \
+  --resume_from_checkpoint checkpoints/apollo-1 \
   "$@"