Train a Time Series Foundation Model

Summary

Trained Moirai-Small (14M parameters) from scratch on a curated subset of GiftEvalPretrain, evaluated zero-shot on the GIFT-Eval benchmark.
Final result for best model : rank #72 (vs Moirai-Small original at #68), trained from scratch in 4h11 on 8× A100 (40Gb of VRAM) for 100K steps (1000 epochs of num_batches_per_epoch = 100).

Also trained a morai-base (91M parameters) for 8h20 with 100K steps (1000 epochs of num_batches_per_epoch = 100), reached rank #75.

Full details can be found in the report.

Dataset

Reduced GiftEvalPretrain from 149 datasets to 67 with 3 things in mind, removing :

• very big datasets (buildings_900k, era5_, cmip6_)
• most of the multivariate
• some other big datasets because of filesystem needs

I meant to only train on univariate data (to reduce complexity) but ended up doing multivariate training anyway. 7 datasets in the training subset contained multiple channels despite initial filtering.

Full selection rationale in the report.

Model

Initially launched a training of morai-base model (94M) only to realise its bad results were due to the low amount of training steps (1000 epochs * 100 batch_size = 100K steps (10% of total training stated in the paper)).

Then trained a morai-small on 100K steps and reached #72 global benchmark (original morai-small is #68).

Results

See full analysis in the report.

Reproducibility

Install uv, env and dependencies

curl -LsSf https://astral.sh/uv/install.sh | sh
uv sync

cp .env.template .env # and add var
set -a && source ./uni2ts/.env && set +a

Train the model (auto detect multi GPU and launches DDP training) :

uv run python -m cli.train \                                                                                           
    -cp conf/pretrain \                                                                                                  
    run_name=my_moirai_small \                                                                                            
    model=moirai_small_gift_eval \                                                                                        
    data=gift_eval_pretrain

Evaluate the model :

# Download eval data (one-time)
huggingface-cli download Salesforce/GiftEval --repo-type=dataset --local-dir /home/sagemaker-user/gift_eval_data

# Run evaluation
./project/moirai-1/eval/gift_eval_small.sh \
  outputs/pretrain/moirai_base_gift_eval/gift_eval_pretrain/my_moirai_small/checkpoints/last.ckpt

Changes to original uni2ts

Dependencies

Added uv as python dependency manager

Bug Fixes

New Code

---

Time Spent