From 3af040e950b374b345614af7ce94f6b1723cbebc Mon Sep 17 00:00:00 2001 From: Oleksandr Shchur Date: Wed, 19 Aug 2026 14:21:36 +0000 Subject: [PATCH] toto-2.0: load model once; tirex-2: unpin datasets toto-2.0: the wrapper reloaded the full model from scratch on every task (fit_predict is called per task), so the previous task's model was still resident when the next was moved to the GPU -> transient 2x model memory and CUDA OOM on multi-task runs. Load the model once and reuse it across tasks. tirex-2: drop the datasets==3.6.0 pin. 3.6.0 can't read parquet whose schema metadata uses the newer 'List' feature type (e.g. the TIME-bench datasets), failing with "Feature type 'List' not found". Unpinning inherits fev's datasets>=2.15,<5.0 and resolves to 4.x, matching the other wrappers. --- models/tirex-2/requirements.txt | 1 - models/toto-2.0/model.py | 8 ++++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/models/tirex-2/requirements.txt b/models/tirex-2/requirements.txt index 617d714..401a2e9 100644 --- a/models/tirex-2/requirements.txt +++ b/models/tirex-2/requirements.txt @@ -1,4 +1,3 @@ tirex-2==0.1.0 fev -datasets==3.6.0 pyarrow>=21.0.0 diff --git a/models/toto-2.0/model.py b/models/toto-2.0/model.py index 02fdd9b..b34cfe6 100644 --- a/models/toto-2.0/model.py +++ b/models/toto-2.0/model.py @@ -29,6 +29,7 @@ def __init__( self.decode_block_size = decode_block_size self.as_univariate = as_univariate self.device = device + self._model = None def _fit_predict(self, task: fev.Task) -> list[datasets.DatasetDict]: import torch @@ -40,7 +41,10 @@ def _fit_predict(self, task: fev.Task) -> list[datasets.DatasetDict]: target_columns = ["target"] if self.as_univariate else task.target_columns - model = PretrainedToto2.from_pretrained(fev.utils.maybe_cache_from_s3(self.model_path)) + if self._model is None: # load the large model once and reuse it across tasks + self._model = PretrainedToto2.from_pretrained(fev.utils.maybe_cache_from_s3(self.model_path)) + self._model = self._model.to(self.device).eval() + config = Toto2GluonTSModelConfig( prediction_length=task.horizon, context_length=self.context_length, @@ -48,7 +52,7 @@ def _fit_predict(self, task: fev.Task) -> list[datasets.DatasetDict]: decode_block_size=self.decode_block_size, quantiles=task.quantile_levels, ) - gts_model = Toto2GluonTSModel(model.to(self.device).eval(), config) + gts_model = Toto2GluonTSModel(self._model, config) predictor = gts_model.create_predictor(batch_size=self.batch_size, device=self.device) logging.getLogger("gluonts").setLevel(100)