You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
[AINode] Add TiRex-2 as a built-in forecasting model - #19273
This PR adds TiRex-2 (NXAI) to AINode's built-in forecasting models, next to timer_xl, sundial, chronos2, moirai2, toto and moment. TiRex-2 is an xLSTM time series foundation model with about 83M parameters. It forecasts nine quantiles (0.1 to 0.9), up to 320 steps ahead, for univariate and multivariate series, and it can use past-only and future-known covariates. The model is registered as tirex2. AINode downloads its weights from the Hugging Face repository NX-AI/TiRex-2 at startup, as it does for the other built-in models.
Usage
SELECT*FROM FORECAST(
model_id =>'tirex2',
targets => (SELECTtime, s FROMdb.tWHERE device ='d1'ORDER BYtimeDESCLIMIT1024) ORDER BYtime,
output_length =>24)
Once the weights are downloaded, SHOW MODELS lists tirex2 with model type tirex2, category builtin and state active.
Example: the input is 1024 points of 50 + 10·sin(2πt/24) + 0.01·t. The request is a TForecastReq with outputLength = 24, sent to AINode's inference manager on CPU with the published checkpoint. The response is 24 values (the median forecast). The first six:
step
forecast
true value
1
51.636
51.580
2
50.628
50.591
3
50.294
50.260
4
50.641
50.611
5
51.603
51.620
6
53.277
53.219
The mean absolute error is 0.112 over 24 steps and 0.126 over 96 steps. If output_length is above the model's horizon, the request fails with a clear message: 1599: TiRex-2 forecasts at most 320 steps, but output_length is 400.
Model code
iotdb/ainode/core/model/tirex2/ contains the inference model from NX-AI/tirex-2 0.5.0 (Apache License 2.0), in the same way as chronos2/ contains the code from chronos-forecasting. The model code is unchanged except for these points:
Imports use the iotdb.ainode.core.model.tirex2 package, and files carry the ASF header followed by NXAI's copyright notice.
The recurrent kernels follow the device of their inputs. Upstream picks the kernels when the model is built. AINode builds the model on CPU and then moves it to the device it is loaded on, so the choice now happens at run time. CUDA inputs use the Triton sLSTM and mLSTM kernels when Triton is available and the GPU has compute capability 8.0 (Ampere) or newer, since the kernels multiply bfloat16 operands. Other inputs, and older GPUs such as Volta and Turing, use the PyTorch kernels.
Triton counts as available only if its kernel modules import. Triton compiles a kernel from its Python source. In a packaged AINode binary, Triton may be bundled without that source; the model then falls back to the PyTorch kernels instead of failing.
The activation named in config.json must be a torch.nn module, such as torch.nn.SiLU. Upstream imports any dotted path from the config.
A checkpoint in the per-gate layout must contain the mLSTM input and forget gate biases. Upstream fills a missing bias with zeros.
The PyTorch sLSTM kernel converts the recurrent weight once per sequence instead of at every step.
Code that AINode does not need is left out: the high-level ForecastModel API, dataframe/GluonTS/fev adapters, plotting and the demo. The narwhals dependency goes with them.
TiRex-2's own loader replays the forward pass from CUDA graphs. This PR leaves CUDA graphs off (use_cuda_graphs = False), so every forecast runs eagerly.
Upstream compiles its residual blocks with torch.compile. This PR runs them eagerly. Inductor needs a C++ compiler and the Python headers at run time, and the packaged AINode executable does not ship the headers. Eager and compiled forecasts are bit-identical on CPU. FlexAttention, which upstream also compiles on CUDA, is off in the published configuration.
AINode integration
modeling_tirex2.TiRex2ForPrediction is the model, loaded through PyTorchModelHubMixin (hub_mixin_cls in ModelInfo). AINode already supports this loading path. The transformers PreTrainedModel path is not used, because the published model.safetensors stores one weight per gate. TiRex-2's load_state_dict hooks fuse those weights into the layout of its recurrent layers, and transformers' loader bypasses the hooks. The weights load with strict=True, so a checkpoint with a missing or extra weight is rejected.
pipeline_tirex2.TiRex2Pipeline turns each input into one TiRex-2 series. The targets of an input, shape (target_count, input_length), are forecast jointly. A covariate that appears only in past_covariates becomes a past covariate. A covariate that also appears in future_covariates becomes a known covariate, built from its past and future values. All inputs of a request run in one batch. The pipeline returns the 0.5 quantile.
output_length must be between 1 and 320. TiRex-2 has no autoregressive mode, so longer horizons are rejected instead of truncated. ModelInfo gets an optional max_output_length, set to 320 for tirex2. InferenceManager checks output_length against it, together with AINode's own limit, before a request is dispatched to an inference pool. A rejected request therefore never reaches the pool, and the pool keeps serving. Models without the field keep AINode's limit.
Dependencies
There are no new dependencies. The model needs PyTorch, NumPy, safetensors and huggingface_hub, which AINode already pins. On Linux, the CUDA wheels of PyTorch include Triton. TiRex-2 0.5.0 replaced its earlier flashrnn and mlstm_kernels dependencies, and their CUDA toolkit requirement, with these Triton kernels. pyproject.toml, ainode.spec and ainode.xml are unchanged. ainode.spec already collects every module under iotdb.ainode.core.
Licences
TiRex-2's code and its published weights are both under the Apache License 2.0 (code licence, weights licence). The weights are downloaded at run time and are not part of the distribution. This PR adds TiRex-2 to LICENSE, NOTICE and AINode's NOTICE-binary, with the copyright line from TiRex-2's NOTICE file.
NXAI's intended-use notice says the release is not intended for high-risk uses under the EU AI Act. The notice states that it does not change the licence.
Tests
tests/test_tirex2.py (unittest, CPU only, no downloads) uses a two-block TiRex-2 with random weights. The tests:
save it in the published checkpoint's per-gate layout and check that it reloads with identical weights;
check that a checkpoint with a missing weight, or a missing input or forget gate bias, is rejected;
load it through AINode's built-in loader and the TiRex2Pipeline, and check the output shapes for 1-D and 2-D targets and that the result is the model's median;
check how covariates are assembled, that output_length above the horizon is rejected, and that the activation lookup is restricted;
check that CPU inputs and GPUs below compute capability 8.0 use the PyTorch kernels, and that a forecast runs without torch.compile.
tests/test_builtin_model_limits.py (unittest) checks that the model limit lowers AINode's limit. With tirex2 loaded in a pool, a 400-step forecast is rejected without being dispatched to the pool, and the next 24-step forecast is served.
The tirex2 entry is added to AINodeTestUtils.BUILTIN_LTSM_MAP, so AINodeSharedClusterIT lists, loads and forecasts with it on CPU and GPU, like the other built-in foundation models. forecastTableFunctionErrorTest checks each model's own output_length limit: 320 for tirex2 and 2880 for the others.
With the published checkpoint on CPU, the forecasts are bit-identical to the unmodified TiRex-2 0.5.0 code for a univariate series, a multivariate series with past and known covariates, and a batch of series of different lengths.
This PR has:
been self-reviewed.
added or updated version, license, or notice information
added comments explaining the "why" and the intent of the code wherever would not be obvious
for an unfamiliar reader.
added unit tests or modified existing tests to cover new code paths, ensuring the threshold
for code coverage.
added integration tests.
Key changed/added classes (or packages if there are too many classes) in this PR
iotdb/ainode/core/model/tirex2/ (new): the TiRex-2 model, TiRex2ForPrediction, TiRex2Pipeline
iotdb/ainode/core/model/model_info.py: ModelInfo.max_output_length; tirex2 entry in BUILTIN_HF_TRANSFORMERS_MODEL_MAP
iotdb/ainode/core/manager/inference_manager.py: output_length checked against the model's limit before pool dispatch
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
An unsupported forecast horizon can terminate the inference worker and leave callers waiting indefinitely, while Volta and Turing GPUs are incorrectly routed to incompatible Triton kernels.
Adds TiRex-2 to AINode’s built-in forecasting models, supporting joint multivariate forecasts and optional covariates with median predictions up to 320 steps ahead.
Changes:
Registers tirex2 and integrates checkpoint loading and forecasting pipelines.
Vendors the model components and device-aware PyTorch/Triton kernels.
Adds unit and cluster-test coverage, plus source and binary licensing notices.
File
Description
NOTICE
Adds TiRex-2 attribution.
LICENSE
Records licensing for vendored code.
iotdb-core/ainode/tests/test_tirex2.py
Tests checkpoint loading, forecasts, covariates, and validation.
The reason will be displayed to describe this comment to others. Learn more.
Confirmed with a bounded reproduction using the actual pool execution loop, batcher, and TiRex-2 pipeline validation (only the model was stubbed with future_len=320): a 400-step request terminates the execution thread with ValueError, leaves the finished queue empty, and leaves a subsequent valid request queued. Please add coverage through the loaded-model pool path, including a valid request after the rejection; the current test only verifies that the direct pipeline call raises.
The reason will be displayed to describe this comment to others. Learn more.
Reviewed 87af3d4. There is a blocking packaged-runtime failure below, confirmed by the current AINode CI logs. I also independently reproduced the invalid-horizon worker failure and acceptance of missing gate biases reported in the existing threads; I will add the reproduction details there.
The seven new TiRex-2 unit tests pass in a local Python 3.11 / PyTorch 2.11.0 environment. This is a source-level sanity check, not validation of the pinned PyTorch 2.14.1, GPU kernels, or packaged executable. The current AINode integration job fails, and the TODO check also flags the new TODO in slstm_torch.py.
The reason will be displayed to describe this comment to others. Learn more.
[P1] Make compilation optional for the packaged AINode runtime
This unconditional torch.compile makes ordinary CPU inference fail in the shipped PyInstaller executable. In the current AINode job, forecastTableFunctionTest fails with CppCompileError: fatal error: Python.h: No such file or directory: Inductor looks under the frozen executable's _internal/include/python3.11, which is not populated by the current packaging. The uploaded AINode log also shows the same error escaping _requests_execute_loop during the loaded-model test, killing the execution thread and leaving the query to time out. This is an actual runtime failure, beyond first-request compilation overhead. Please make compilation opt-in or provide a supported packaged compilation path with an eager fallback, and verify both direct and pooled forecasts using the built executable.
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
The Triton sLSTM and mLSTM kernels lack numerical parity tests against their reference implementations, leaving recurrence, state propagation, and padded-tail correctness unverified by CPU forecasts and shape-only GPU checks.
The new Triton mLSTM implementation is not checked numerically against mlstm_chunkwise_native. All model unit forecasts use CPU, and the GPU cluster smoke test verifies only shape/count, so errors in chunk-state propagation or padded-tail handling can pass the existing assertions. Please add CUDA-gated, download-free parity cases spanning a chunk boundary and a non-multiple-of-64 sequence length, with tolerances appropriate to the kernel's dtype casts.
Add numerical parity tests for the sLSTM Triton recurrence
The sLSTM Triton recurrence has no numerical regression coverage: test_tirex2.py runs forecasts only on CPU and mocks CUDA dispatch, while the cluster smoke test checks result shape/count rather than values. In particular, the recurrent-weight transpose and gate/state layout are not checked against the PyTorch implementation. Please add a CUDA-gated, download-free comparison of outputs and final states using nonzero recurrent weights, nonzero initial states, and multiple time steps.
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
TiRex-2 causes a forecast-limit test assertion to fail, while zero-length requests can terminate the inference loop and leave callers waiting indefinitely.
The new pre-dispatch validation still allows TForecastReq(outputLength=0). With tirex2 loaded, this reaches TiRex2._prepare_input(), which raises ValueError; the uncaught exception exits _requests_execute_loop(), leaving this caller waiting indefinitely and subsequent requests unprocessed. Reject values below 1 here as well, before dispatch, and cover a zero-length rejection followed by a valid pooled request.
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
Adding TiRex-2 to shared integration coverage causes the horizon-rejection test to fail because it expects an upper bound of 2880 while the model-specific validation returns 320.
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
Adding TiRex-2 breaks the shared integration test because it expects a 2880-step limit while the model correctly reports 320; the range-error assertions need to be model-aware.
The reason will be displayed to describe this comment to others. Learn more.
Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.
The reason will be displayed to describe this comment to others. Learn more.
Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.
The reason will be displayed to describe this comment to others. Learn more.
Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.
The reason will be displayed to describe this comment to others. Learn more.
Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.
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
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.
CRZbulabula
left a comment
Description
This PR adds TiRex-2 (NXAI) to AINode's built-in forecasting models, next to
timer_xl,sundial,chronos2,moirai2,totoandmoment. TiRex-2 is an xLSTM time series foundation model with about 83M parameters. It forecasts nine quantiles (0.1 to 0.9), up to 320 steps ahead, for univariate and multivariate series, and it can use past-only and future-known covariates. The model is registered astirex2. AINode downloads its weights from the Hugging Face repositoryNX-AI/TiRex-2at startup, as it does for the other built-in models.Usage
Once the weights are downloaded,
SHOW MODELSliststirex2with model typetirex2, categorybuiltinand stateactive.Example: the input is 1024 points of
50 + 10·sin(2πt/24) + 0.01·t. The request is aTForecastReqwithoutputLength = 24, sent to AINode's inference manager on CPU with the published checkpoint. The response is 24 values (the median forecast). The first six:The mean absolute error is 0.112 over 24 steps and 0.126 over 96 steps. If
output_lengthis above the model's horizon, the request fails with a clear message:1599: TiRex-2 forecasts at most 320 steps, but output_length is 400.Model code
iotdb/ainode/core/model/tirex2/contains the inference model from NX-AI/tirex-2 0.5.0 (Apache License 2.0), in the same way aschronos2/contains the code from chronos-forecasting. The model code is unchanged except for these points:iotdb.ainode.core.model.tirex2package, and files carry the ASF header followed by NXAI's copyright notice.config.jsonmust be atorch.nnmodule, such astorch.nn.SiLU. Upstream imports any dotted path from the config.ForecastModelAPI, dataframe/GluonTS/fev adapters, plotting and the demo. Thenarwhalsdependency goes with them.TiRex-2's own loader replays the forward pass from CUDA graphs. This PR leaves CUDA graphs off (
use_cuda_graphs = False), so every forecast runs eagerly.Upstream compiles its residual blocks with
torch.compile. This PR runs them eagerly. Inductor needs a C++ compiler and the Python headers at run time, and the packaged AINode executable does not ship the headers. Eager and compiled forecasts are bit-identical on CPU. FlexAttention, which upstream also compiles on CUDA, is off in the published configuration.AINode integration
modeling_tirex2.TiRex2ForPredictionis the model, loaded throughPyTorchModelHubMixin(hub_mixin_clsinModelInfo). AINode already supports this loading path. The transformersPreTrainedModelpath is not used, because the publishedmodel.safetensorsstores one weight per gate. TiRex-2'sload_state_dicthooks fuse those weights into the layout of its recurrent layers, and transformers' loader bypasses the hooks. The weights load withstrict=True, so a checkpoint with a missing or extra weight is rejected.pipeline_tirex2.TiRex2Pipelineturns each input into one TiRex-2 series. The targets of an input, shape(target_count, input_length), are forecast jointly. A covariate that appears only inpast_covariatesbecomes a past covariate. A covariate that also appears infuture_covariatesbecomes a known covariate, built from its past and future values. All inputs of a request run in one batch. The pipeline returns the 0.5 quantile.output_lengthmust be between 1 and 320. TiRex-2 has no autoregressive mode, so longer horizons are rejected instead of truncated.ModelInfogets an optionalmax_output_length, set to 320 fortirex2.InferenceManagerchecksoutput_lengthagainst it, together with AINode's own limit, before a request is dispatched to an inference pool. A rejected request therefore never reaches the pool, and the pool keeps serving. Models without the field keep AINode's limit.Dependencies
There are no new dependencies. The model needs PyTorch, NumPy, safetensors and huggingface_hub, which AINode already pins. On Linux, the CUDA wheels of PyTorch include Triton. TiRex-2 0.5.0 replaced its earlier
flashrnnandmlstm_kernelsdependencies, and their CUDA toolkit requirement, with these Triton kernels.pyproject.toml,ainode.specandainode.xmlare unchanged.ainode.specalready collects every module underiotdb.ainode.core.Licences
TiRex-2's code and its published weights are both under the Apache License 2.0 (code licence, weights licence). The weights are downloaded at run time and are not part of the distribution. This PR adds TiRex-2 to
LICENSE,NOTICEand AINode'sNOTICE-binary, with the copyright line from TiRex-2's NOTICE file.NXAI's intended-use notice says the release is not intended for high-risk uses under the EU AI Act. The notice states that it does not change the licence.
Tests
tests/test_tirex2.py(unittest, CPU only, no downloads) uses a two-block TiRex-2 with random weights. The tests:TiRex2Pipeline, and check the output shapes for 1-D and 2-D targets and that the result is the model's median;output_lengthabove the horizon is rejected, and that the activation lookup is restricted;torch.compile.tests/test_builtin_model_limits.py(unittest) checks that the model limit lowers AINode's limit. Withtirex2loaded in a pool, a 400-step forecast is rejected without being dispatched to the pool, and the next 24-step forecast is served.tirex2entry is added toAINodeTestUtils.BUILTIN_LTSM_MAP, soAINodeSharedClusterITlists, loads and forecasts with it on CPU and GPU, like the other built-in foundation models.forecastTableFunctionErrorTestchecks each model's ownoutput_lengthlimit: 320 fortirex2and 2880 for the others.This PR has:
for an unfamiliar reader.
for code coverage.
Key changed/added classes (or packages if there are too many classes) in this PR
iotdb/ainode/core/model/tirex2/(new): the TiRex-2 model,TiRex2ForPrediction,TiRex2Pipelineiotdb/ainode/core/model/model_info.py:ModelInfo.max_output_length;tirex2entry inBUILTIN_HF_TRANSFORMERS_MODEL_MAPiotdb/ainode/core/manager/inference_manager.py:output_lengthchecked against the model's limit before pool dispatchiotdb-core/ainode/tests/test_tirex2.py,tests/test_builtin_model_limits.py(new)integration-test/.../ainode/utils/AINodeTestUtils.java:tirex2inBUILTIN_LTSM_MAP,getMaxOutputLengthintegration-test/.../ainode/it/AINodeSharedClusterIT.java: per-model limit inforecastTableFunctionErrorTestLICENSE,NOTICE,iotdb-core/ainode/src/assembly/resources/NOTICE-binary