toto-2.0: load model once across tasks; tirex-2: unpin datasets - #176
Merged
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Two small model-wrapper fixes surfaced while running multi-task evals.
toto-2.0: load the model once, reuse across tasks_fit_predictis called once per task, and it reloaded the full pretrained model from scratch every time. On GPU the previous task's model was often still resident when the next was moved onto the device, giving a transient 2× model footprint andtorch.OutOfMemoryErrorpartway through a multi-task run (observed on a 2.5B checkpoint: process pinned at ~21.5 GiB and failing to allocate tens of MiB, regardless ofbatch_size). Now the model is loaded once (cached onself._model) and reused; only the lightweight per-taskconfig/predictor is rebuilt.tirex-2: drop thedatasets==3.6.0pindatasets==3.6.0can't deserialize parquet whose embedded schema metadata uses the newerListfeature type (e.g. the TIME-bench datasets), failing withValueError: Feature type 'List' not found. The pin came fromtirex-2'sfevextra; we don't install that extra, so removing the explicit line lets the resolver honorfev'sdatasets>=2.15,<5.0and pick 4.x — matching every other wrapper. Verified onCommodity_Import/M/short(resolves todatasets 4.8.5, loads and produces metrics).Both verified locally.