Add support for TimesFM 3.0 foundation model (multivariate, past & future covariates) - #3199
Merged
Merged
Conversation
Port the TimesFM 3.0 PyTorch architecture (commit 9de33f6f487baf8adc26eb757f29b6a7a557b823) to darts/models/components/timesfm3_submodels.py as a first step towards TimesFM 3.0 support (issue unit8co#3196): - ResidualBlock, RMSNorm (torch>=2.0 compatible), PerDimScale, RotaryPositionalEmbedding, MultiHeadAttention (SDPA), MixingTransformer (sequence + variate attention) and stacked variant - Patching/stitching and CPM-RevIN refinement utilities - Config dataclasses buildable from the HuggingFace config.json Trims vs upstream (documented in the module docstring): no PyTorchModelHubMixin (weights are loaded via HuggingFaceConnector), no autoregressive KV-cache, no multi-segment support, SDPA-only attention path. Verified: state_dict keys and shapes match the google/timesfm-3.0-pytorch checkpoint exactly (445/445 tensors).
Implements the Darts integration of TimesFM 3.0 (issue unit8co#3196) on top of the ported submodels: - _TimesFM3Module: assembles the architecture from the HuggingFace config.json and ports the original preprocessing and single-pass decode (stitching, linear detrending, iterative CPM-RevIN refinement). Multivariate series are forecast natively with variate attention; past covariates are mapped to past-only channels and future covariates to past-and-future channels (output_chunk_shift gaps are masked out). - TimesFM3Model: FoundationModel wrapper with accept_license gate (TimesFM 3.0 weights are non-commercial), context (15,360) and prediction (1,024) caps, QuantileRegression likelihood restricted to the 9 pre-trained quantiles. Fine-tuning is not supported yet. Verified bit-exact parity (max abs diff 0.0) against the upstream timesfm3.TimesFM3Forecaster on univariate, multivariate, covariates, output_chunk_shift and batched scenarios.
- tiny_timesfm3: tiny random-weight checkpoint (2 layers, model_dims 64, 4 heads; structural values input/output patch 32/64 and 9 quantiles kept) generated with the upstream timesfm3 package, loaded by fast tests via local_dir (49 tensors, 716KB) - timesfm3_prediction/timesfm3.npz: reference forecasts (time, variables, quantiles) = (128, 2, 9) generated with the upstream timesfm3 TimesFM3Forecaster and the real google/timesfm-3.0-pytorch weights on the Zurich electricity dataset (last 1024 points of context, horizon 128) - update TimesFM3Model docstring examples with real AirPassengers outputs Fidelity verified: Darts predictions match the reference with rtol=1e-5, atol=1e-5 for both deterministic and probabilistic modes.
- test_timesfm3.py: fast tests (creation validation, deterministic, probabilistic, multivariate, past/future/both covariates, variate limit, multiple series) using the tiny local artefact, and slow fidelity tests comparing against the original implementation reference - register TimesFM3Model in the lazy public API and foundation model docs - add TimesFM3Model to the cross-model variable input chunk length tests in test_foundation.py
- CHANGELOG entry under Unreleased - model tables in README.md, docs/source/index.rst (multivariate, past and future covariates support, no static covariates), and docs/userguide/covariates.md - notebook 25: mention TimesFM 3.0's pre-trained quantiles and license acceptance, covariates support, and API reference link
Add TimesFM3Model to the remaining parametrized cross-model tests: - min_train_series_length for variable and fixed input chunk lengths - variable ICL save/load (save(), load_weights(); load_from_checkpoint is skipped for TimesFM3 as it requires fine-tuning, which is not supported yet - same as TiRex) - variable ICL predictions matching fixed ICL predictions (validates the masked context padding semantics) - loading weights into a model with different chunk parameters
Map ff_activation="swiglu" to plain SiLU, matching the released timesfm PyPI package (which implements no gated FFN in its PyTorch backend). This keeps configurations of checkpoints published with earlier timesfm releases loadable; the ported commit 9de33f6f had removed the alias. No effect on the released checkpoint, which uses relu.
…imesFM3Model - test_missing_values: NaNs in the target (interior and trailing) and in past / future covariates are handled via the masking logic of the ported decode() - predictions never contain NaNs - test_output_chunk_shift: predictions start after the shifted gap and auto-regressive prediction with output_chunk_shift raises - test_too_many_variates: also cover the boundary case of exactly 32 target components (the checkpoint maximum)
…M3Model - load the HuggingFace config.json at model construction (as done by Chronos2Model) and validate the QuantileRegression likelihood against the pre-trained quantiles read from it, instead of hardcoded constants - validate the total number of target components and covariates at fit time, early, using max_variates from the checkpoint configuration (the prediction-time check in the PL module is kept as a safety net) - document the provenance of the context / prediction length caps
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #3199 +/- ##
==========================================
- Coverage 97.14% 97.06% -0.09%
==========================================
Files 169 169
Lines 18826 18891 +65
==========================================
+ Hits 18288 18336 +48
- Misses 538 555 +17 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Collaborator
|
Thanks for the PR @JuanCruzC97, I'll have time to review soon (was a bit busy implementing the torch I/O refactor from #3204, will apply the required changes to your PR once I've merge that one). |
Contributor
Author
dennisbader
approved these changes
Sep 17, 2026
dennisbader
left a comment
dennisbader
left a comment
Collaborator
There was a problem hiding this comment.
Beautiful PR and implementation @JuanCruzC97, congrats 🚀 👏
It was in an excellent state already from the beginning, thanks a lot for that :)
I pushed a couple of updates that added the last couple of missing things:
- fine-tuning is now supported
- removed the upper limit for
output_chunk_length(if the original implementation doesn't have a limit, we should not set one either) - refactored
forward()a bit to leverage our newPLModuleInput(andPLModuleOutput) - some minor doc updates
We can merge once all tests have passed. Let me know if you would like to have a look at the changes before I merge.
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
2 participants
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.
Checklist before merging this PR:
Fixes #3196
Summary
This PR adds support for Google's TimesFM 3.0 time-series foundation model (
google/timesfm-3.0-pytorch, 330M parameters) asTimesFM3Model, following the discussion in #3196 and the license clarification in google-research/timesfm#501.Unlike previous versions, TimesFM 3.0 natively supports:
QuantileRegressionlikelihood, single-pass decoding with patch stitching).The implementation follows the established foundation-model pattern in Darts (
FoundationModelbase class,HuggingFaceConnectorweight loading, lazy model creation in_create_model()), with the architecture ported natively intodarts/models/components/timesfm3_submodels.pyfrom google-research/timesfm (Apache-2.0, commit pinned in the module docstring).Since the TimesFM 3.0 pre-trained weights are non-commercial (
timesfm-non-commercial-license-v1.0), model creation requiresaccept_license=True, following theTiRexModelpattern. There are no new runtime dependencies (huggingface-hubandsafetensorsare already part ofdarts[torch]).Validation: the ported architecture was verified to be bit-exact against the original implementation — the Darts one-shot forecast matches the upstream
timesfm3.TimesFM3Forecasteroutput with max abs diff0.0(univariate, multivariate, covariates,output_chunk_shift, and batched scenarios), and the slow fidelity test compares against reference forecasts generated with the original forecaster and the real weights (rtol=atol=1e-5).Design decisions for review
fit()time and at prediction time), instead of the non-deterministic random covariate subsampling performed by the upstreamTimesFM3Evaluator.max_variatesis read from the checkpoint'sconfig.json, not hardcoded.predict_batch()before callingdecode(); we use the masking logic ofdecode()itself (the model was trained with masked inputs). Behavioral deviation: with NaN-containing series, results do not match upstream 1:1. Easy to switch to interpolation if strict fidelity is preferred.input_chunk_length <= 15360with an error. Neither 15,360 nor the horizon cap are part ofconfig.json(they are upstreamTimesFM3Forecasterconstants); they remain documented class constants. What is read fromconfig.json(same spirit asChronos2Model): the pre-trained quantiles andmax_variates.n > output_chunk_length, Darts splits the forecast into autoregressive chunks, which can differ from upstream's single-pass stitched decoding. Documented in the docstring; setoutput_chunk_length >= nfor a single-pass forecast.Other Information
Fidelity & implementation notes
DecodeCache(autoregression is handled byTorchForecastingModel), nosegment_idssupport and SDPA-only attention (the path used by the released checkpoint), only theidentityinput transform,swiglumapped to plain SiLU (as upstream does — no gated FFN in the PyTorch backend).config.json(e.g.use_rope_var=False), not the inconsistent defaults of the upstream_make_torch_model().hub_model_revisionis pinned to the current commit ofgoogle/timesfm-3.0-pytorchfor reproducibility._get_running_statsloops over patches in Python (like upstream; ~480 iterations at maximum context). Not vectorized to avoid numerical divergence; can be a follow-up.Tests
test_timesfm3.py: creation validation (license, caps, likelihood, fine-tuning), deterministic / probabilistic / multivariate / past / future / both-covariates / multiple-series forecasts, missing values (interior & trailing target NaNs, past & future covariate NaNs),output_chunk_shiftalignment + autoregression lock, variate limit boundary (32 ok / 33 raises) — fast tests run against a committed tiny random-weight artefact (structural patch sizes 32/64 and 9 quantiles intact).@pytest.mark.slow) againsttimesfm3.npzreference forecasts generated with the original forecaster and the full checkpoint (Zurich electricity dataset, deterministic + probabilistic).test_foundation.py(variable input chunk length fit/predict/save/load/weights,min_train_series_length, fixed-vs-variable ICL equivalence); only excluded from theload_from_checkpointpart, which requires fine-tuning (as with TiRex).Docs & artefacts
README.md,docs/source/index.rst,docs/userguide/covariates.md;CHANGELOG.mdentry under Unreleased; notebook25-FoundationModel-examples.ipynbmentions the model (covariates support, quantiles, license note). NoINSTALL.mdchanges needed (no extra dependencies).