diff --git a/rdagent/app/qlib_rd_loop/quant.py b/rdagent/app/qlib_rd_loop/quant.py index 99506adf..64b6e66c 100644 --- a/rdagent/app/qlib_rd_loop/quant.py +++ b/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, diff --git a/rdagent/components/coder/kronos_adapter.py b/rdagent/components/coder/kronos_adapter.py index ef90d85d..5f51ca91 100644 --- a/rdagent/components/coder/kronos_adapter.py +++ b/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)