Add TimesFM 3 multivariate forecasting support - #718
Open
justinchuby wants to merge 4 commits into
Open
Conversation
Implement the TimesFM 3 patched-core forward graph with running RevIN, temporal and variate attention, iterative CPM refinement, and the forecasting task contract. Detect the official checkpoint config despite its missing Transformers model_type and cover FP32/FP16 execution and reference behavior. Co-authored-by: Copilot App <[email protected]> Copilot-Session: 3b9976ac-d466-4ff5-a901-3503f99adfc7 Signed-off-by: Justin Chu <[email protected]>
Performance Comparison
|
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
There are a couple of concrete correctness/robustness issues (config validation and task config typing/guarding) that can lead to invalid graphs or CI type-check failures.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds a native declarative ONNX implementation of TimesFM 3’s patched-core multivariate forecasting graph, including a new time-series-forecasting task and config detection for checkpoints that omit a Transformers model_type.
Changes:
- Introduces
TimesFM3Model/TimesFM3Configplus a new forecasting task (TimeSeriesForecastingTask) and registers them in the model/task registries. - Extends Transformers config resolution to infer
timesfm3from a distinctive config schema whenmodel_typeis missing. - Adds co-located unit/runtime tests and hooks the model into the repo’s build-graph/model-coverage harnesses.
File summaries
| File | Description |
|---|---|
| tests/model_coverage_test.py | Marks timesfm3 as covered by co-located tests due to non-redistributable official weights. |
| tests/build_graph/_support.py | Includes timesfm3 in the build-graph coverage list. |
| src/mobius/tasks/_time_series_forecasting.py | New task defining the patched-input forecasting I/O contract and exporting the model graph. |
| src/mobius/tasks/init.py | Exposes/registers the new task under time-series-forecasting. |
| src/mobius/models/timesfm3.py | Implements TimesFM 3 patched-core forward pass and config parsing. |
| src/mobius/models/timesfm3_test.py | Adds tests for config parsing, graph contract/weight names, runtime execution, and registry/task wiring. |
| src/mobius/models/init.py | Exports the new TimesFM3 model/config symbols. |
| src/mobius/integrations/transformers/_config_resolver.py | Adds schema-signature detection to infer timesfm3 when model_type is absent. |
| src/mobius/_registry.py | Registers timesfm3 with the new task/config class and adds fallback registration metadata. |
Review details
- Files reviewed: 9/9 changed files
- Comments generated: 2
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Separate raw-series preparation, patched preprocessing, learned model execution, CPM postprocessing, and forecast stitching into tensor-only ONNX components. Keep all learned layers in the control-flow-free model component for execution-provider capture while supporting padded batches with heterogeneous context and horizon lengths. Add ONNX interpolation, detrending, stitching, quantile policies, and regressions for padding invariance and upstream numerical behavior. Co-authored-by: Copilot App <[email protected]> Copilot-Session: 3b9976ac-d466-4ff5-a901-3503f99adfc7 Signed-off-by: Justin Chu <[email protected]>
Demonstrate the full five-component ONNX pipeline with reusable sessions, CUDA I/O binding, fixed-shape CUDA graph capture, latency reporting, and an official PyTorch comparison. Co-authored-by: Copilot App <[email protected]> Copilot-Session: 3b9976ac-d466-4ff5-a901-3503f99adfc7 Signed-off-by: Justin Chu <[email protected]>
Reject invalid patch lengths before graph construction and guard the forecasting task against incompatible base configurations. Co-authored-by: Copilot App <[email protected]> Copilot-Session: 3b9976ac-d466-4ff5-a901-3503f99adfc7 Signed-off-by: Justin Chu <[email protected]>
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
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.
Summary
raw_preprocessor: padded-batch interpolation, dynamic padding, detrending, role/mask construction, and patchingpreprocessor: running RevIN and patched feature constructionmodel: every learned layer (residual projection, temporal/variate transformer, output head), with noScan,Loop, orIf, so it can be captured independentlypostprocessor: iterative CPM RevIN refinement and reverse normalizationstitcher: overlap stitching, trend reversal, quantile sorting, optional positivity, and padded validity masksgoogle/timesfm-3.0-pytorchconfig, which does not declare a Transformersmodel_typeexamples/timesfm3_forecasting.py, a reusable five-session driver with CUDA I/O binding, fixed-shape CUDA graph capture, latency reporting, CSV output, and official PyTorch parity/benchmark comparisonThe implementation follows
google-research/timesfmcommit331c6d33cb1ac2611de3056d0ac7164aab6301eband checkpoint revision900fcab43d1bfe71733a33b3fec61a41fce28a27.Input contract
raw_preprocessoraccepts right-aligned context tensors[B,V,C], left-aligned future-covariate tensors[B,V,H], observed-value masks, per-row context/horizon lengths, and roles0=target,1=past-only,2=past-future. Outputs remain padded and include explicit validity/target masks; no ragged tensor convention is required.Symmetric averaging and outer z-normalization remain opt-in host orchestration policies because they duplicate/remap whole requests rather than form part of the checkpoint forward pass.
Validation
1.1e-69.2e-536.1versus seasonal-naive MAE71.3modelcomponentCUDA benchmark
Measured on an NVIDIA RTX 4060 Laptop GPU with ONNX Runtime 1.29.0 and PyTorch 2.11.0+cu128, using the official 330M checkpoint, batch 1, one variate, 512 context points, 64 forecast points, and 50 steady-state learned-core iterations:
CUDA graph capture reduced ORT learned-core median latency by about 21% in the controlled same-process comparison. FP16 ONNX versus FP32 PyTorch is not a like-for-like precision comparison; the example reports this fallback explicitly and still checks outputs (learned-core max absolute difference
0.00843, final forecast0.00926).Licensing and waivers
The upstream implementation is Apache-2.0. The official TimesFM 3 weights use Google's non-commercial license, so this PR does not include, redistribute, cache, or download those weights in CI.
Waivers: