feat: support Kronos-small and Kronos-base models, auto-select GPU/CPU

rdagent/app/qlib_rd_loop/quant.py

@@ -100,7 +100,7 @@ class QuantRDLoop(RDLoop):
             if json_path.exists() and parquet_path.exists():
                 try:
                     existing = _json.loads(json_path.read_text())
-                    if existing.get("ic") is not None:
+                    if existing.get("ic") is not None and existing.get("model_size") == "small":
                         logger.info(f"Kronos: {factor_name} exists (IC={existing['ic']:.4f}), skip")
                         continue
                 except Exception:
@@ -109,14 +109,23 @@ class QuantRDLoop(RDLoop):
             try:
                 from rdagent.components.coder.kronos_adapter import build_kronos_factor, evaluate_kronos_model
 
-                logger.info(f"Kronos: generating {factor_name} (pred={pred_bars}, stride=500, CPU)...")
+                has_cuda = False
+                try:
+                    import torch
+                    has_cuda = torch.cuda.is_available()
+                except Exception:
+                    pass
+                device = "cuda" if has_cuda else "cpu"
+
+                logger.info(f"Kronos-small: generating {factor_name} (pred={pred_bars}, stride=500, {device})...")
                 factor_df = build_kronos_factor(
                     hdf5_path=data_path,
                     context_bars=100,
                     pred_bars=pred_bars,
                     stride_bars=500,
-                    device="cpu",
-                    batch_size=16,
+                    device=device,
+                    batch_size=32,
+                    model_size="small",
                 )
 
                 values_dir.mkdir(parents=True, exist_ok=True)
@@ -128,8 +137,9 @@ class QuantRDLoop(RDLoop):
                     context_bars=100,
                     pred_bars=pred_bars,
                     stride_bars=2000,
-                    device="cpu",
-                    batch_size=16,
+                    device=device,
+                    batch_size=32,
+                    model_size="small",
                 )
                 ic = metrics.get("IC_mean", 0.0) or 0.0
 
@@ -138,6 +148,7 @@ class QuantRDLoop(RDLoop):
                     "factor_name": factor_name,
                     "status": "success",
                     "ic": ic,
+                    "model_size": "small",
                     "model": "NeoQuasar/Kronos-mini",
                     "context_bars": 100,
                     "pred_bars": pred_bars,

rdagent/components/coder/kronos_adapter.py

@@ -90,9 +90,19 @@ class KronosAdapter:
     MODEL_ID = "NeoQuasar/Kronos-mini"
     TOKENIZER_ID = "NeoQuasar/Kronos-Tokenizer-2k"
 
-    def __init__(self, device: Optional[str] = None, max_context: int = 512):
+    # Mapping for larger Kronos variants
+    _MODEL_MAP = {
+        "mini": ("NeoQuasar/Kronos-mini", "NeoQuasar/Kronos-Tokenizer-2k"),
+        "small": ("NeoQuasar/Kronos-small", "NeoQuasar/Kronos-Tokenizer-base"),
+        "base": ("NeoQuasar/Kronos-base", "NeoQuasar/Kronos-Tokenizer-base"),
+    }
+
+    def __init__(self, device: Optional[str] = None, max_context: int = 512, model_size: str = "mini"):
         self.device = device or "cpu"
         self.max_context = max_context
+        self.model_size = model_size
+        if model_size in self._MODEL_MAP:
+            self.MODEL_ID, self.TOKENIZER_ID = self._MODEL_MAP[model_size]
         self._predictor = None
 
     def load(self) -> "KronosAdapter":
@@ -102,11 +112,10 @@ class KronosAdapter:
             raise RuntimeError("Kronos not available — see warning above.")
         from model import Kronos, KronosTokenizer, KronosPredictor  # type: ignore
 
-        logger.info(f"Loading Kronos-mini from HuggingFace ({self.MODEL_ID})...")
+        logger.info(f"Loading Kronos-{self.model_size} from HuggingFace ({self.MODEL_ID})...")
         tokenizer = KronosTokenizer.from_pretrained(self.TOKENIZER_ID)
         model = Kronos.from_pretrained(self.MODEL_ID)
-        self._predictor = KronosPredictor(model, tokenizer, device=self.device, max_context=self.max_context)
-        logger.info("Kronos-mini loaded.")
+        logger.info(f"Kronos-{self.model_size} loaded.")
         return self
 
     def predict_next_bars(
@@ -223,6 +232,7 @@ def build_kronos_factor(
     stride_bars: int = 96,
     device: Optional[str] = None,
     batch_size: int = 32,
+    model_size: str = "mini",
 ) -> pd.DataFrame:
     """
     Generate the Kronos predicted-return factor for all EUR/USD 1-min bars.
@@ -244,7 +254,7 @@ def build_kronos_factor(
     df = raw.xs(instrument, level="instrument")
     ohlcv = _ohlcv_from_predix(df)
 
-    adapter = KronosAdapter(device=device, max_context=min(context_bars, 512))
+    adapter = KronosAdapter(device=device, max_context=min(context_bars, 512), model_size=model_size)
     adapter.load()
 
     bar_indices = list(range(context_bars, len(ohlcv), stride_bars))
@@ -302,6 +312,7 @@ def evaluate_kronos_model(
     stride_bars: int = 30,
     device: Optional[str] = None,
     batch_size: int = 32,
+    model_size: str = "mini",
 ) -> dict:
     """
     Evaluate Kronos as a standalone model (Option B, alongside LightGBM).
@@ -318,7 +329,7 @@ def evaluate_kronos_model(
     df = raw.xs(instrument, level="instrument")
     ohlcv = _ohlcv_from_predix(df)
 
-    adapter = KronosAdapter(device=device, max_context=min(context_bars, 512))
+    adapter = KronosAdapter(device=device, max_context=min(context_bars, 512), model_size=model_size)
     adapter.load()
 
     n = len(ohlcv)