diff --git a/.github/workflows/cluster-it-1c1d1a.yml b/.github/workflows/cluster-it-1c1d1a.yml index 111449fb335d6..aa69554225f96 100644 --- a/.github/workflows/cluster-it-1c1d1a.yml +++ b/.github/workflows/cluster-it-1c1d1a.yml @@ -81,6 +81,9 @@ jobs: PY "$AINODE_VENV_PY" -m unittest discover -v -s tests -p 'test_dependency_security.py' "$AINODE_VENV_PY" -m unittest discover -v -s tests -p 'test_sundial_flow_loss.py' + "$AINODE_VENV_PY" -m unittest discover -v -s tests -p 'test_timestral.py' + "$AINODE_VENV_PY" -m unittest discover -v -s tests -p 'test_builtin_model_download.py' + "$AINODE_VENV_PY" -m unittest discover -v -s tests -p 'test_builtin_model_load.py' - name: Upload Artifact if: failure() uses: actions/upload-artifact@v6 diff --git a/LICENSE b/LICENSE index 222444ed3fbc9..5338155329911 100644 --- a/LICENSE +++ b/LICENSE @@ -426,3 +426,14 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -------------------------------------------------------------------------------- + +The following files include code modified from Timestral project. + +./iotdb-core/ainode/iotdb/ainode/core/model/timestral/* + +The Timestral is open source software licensed under the Apache License 2.0 +Project page: https://github.com/mistralai/timestral +License: https://github.com/mistralai/timestral/blob/main/LICENSE +Copyright 2026 Mistral AI + +-------------------------------------------------------------------------------- diff --git a/NOTICE b/NOTICE index b457766a81ba0..2b850e4469d96 100644 --- a/NOTICE +++ b/NOTICE @@ -34,6 +34,15 @@ This product includes source code derived from the chronos-forecasting project: ============================================================================ +This product includes source code derived from the Timestral project: + + Timestral + Copyright 2026 Mistral AI + Licensed under the Apache License, Version 2.0 + https://github.com/mistralai/timestral + +============================================================================ + Apache Commons Collections Copyright 2001-2019 The Apache Software Foundation diff --git a/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeSharedClusterIT.java b/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeSharedClusterIT.java index cbd0f62f16ef2..4f3aa77976978 100644 --- a/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeSharedClusterIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeSharedClusterIT.java @@ -52,6 +52,7 @@ import static org.apache.iotdb.ainode.utils.AINodeTestUtils.BUILTIN_LTSM_MAP; import static org.apache.iotdb.ainode.utils.AINodeTestUtils.BUILTIN_MODEL_MAP; +import static org.apache.iotdb.ainode.utils.AINodeTestUtils.ON_REQUEST_BUILTIN_MODEL_MAP; import static org.apache.iotdb.ainode.utils.AINodeTestUtils.checkHeader; import static org.apache.iotdb.ainode.utils.AINodeTestUtils.checkModelNotOnSpecifiedDevice; import static org.apache.iotdb.ainode.utils.AINodeTestUtils.checkModelOnSpecifiedDevice; @@ -236,13 +237,20 @@ private void showBuiltInModelTest(Statement statement) throws SQLException { ResultSetMetaData resultSetMetaData = resultSet.getMetaData(); checkHeader(resultSetMetaData, "ModelId,ModelType,Category,State"); while (resultSet.next()) { - built_in_model_count++; FakeModelInfo modelInfo = new FakeModelInfo( resultSet.getString(1), resultSet.getString(2), resultSet.getString(3), resultSet.getString(4)); + FakeModelInfo onRequestModelInfo = ON_REQUEST_BUILTIN_MODEL_MAP.get(modelInfo.getModelId()); + if (onRequestModelInfo != null) { + // Inactive until another test loads it + assertEquals(onRequestModelInfo.getModelType(), modelInfo.getModelType()); + assertEquals(onRequestModelInfo.getCategory(), modelInfo.getCategory()); + continue; + } + built_in_model_count++; assertTrue(BUILTIN_MODEL_MAP.containsKey(modelInfo.getModelId())); assertEquals(BUILTIN_MODEL_MAP.get(modelInfo.getModelId()), modelInfo); } @@ -459,6 +467,21 @@ public void largeTimeSeriesModelLoadInferenceAndForecastTest() } } + @Test + public void onRequestBuiltInModelLoadInferenceAndForecastTest() + throws SQLException, InterruptedException { + try (Connection treeConnection = EnvFactory.getEnv().getConnection(BaseEnv.TREE_SQL_DIALECT); + Statement treeStatement = treeConnection.createStatement(); + Connection tableConnection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT); + Statement tableStatement = tableConnection.createStatement()) { + for (FakeModelInfo modelInfo : ON_REQUEST_BUILTIN_MODEL_MAP.values()) { + // LOAD MODEL downloads the weights before loading the model + loadRunAndUnloadModelOnDevices( + treeStatement, tableStatement, modelInfo.getModelId(), "cpu"); + } + } + } + private void loadRunAndUnloadModelOnDevices( Statement treeStatement, Statement tableStatement, String modelId, String devices) throws SQLException, InterruptedException { diff --git a/integration-test/src/test/java/org/apache/iotdb/ainode/utils/AINodeTestUtils.java b/integration-test/src/test/java/org/apache/iotdb/ainode/utils/AINodeTestUtils.java index 8f5567c7c9cf7..d92323d65cc8d 100644 --- a/integration-test/src/test/java/org/apache/iotdb/ainode/utils/AINodeTestUtils.java +++ b/integration-test/src/test/java/org/apache/iotdb/ainode/utils/AINodeTestUtils.java @@ -65,6 +65,16 @@ public class AINodeTestUtils { "moment", new FakeModelInfo("moment", "moment", "builtin", "active"))) .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + /** + * Built-in models whose weights are downloaded only when they are loaded. They are inactive until + * then. + */ + public static final Map ON_REQUEST_BUILTIN_MODEL_MAP = + Stream.of( + new AbstractMap.SimpleEntry<>( + "timestral", new FakeModelInfo("timestral", "timestral", "builtin", "inactive"))) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + public static final Map BUILTIN_MODEL_MAP; static { diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_controller.py b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_controller.py index 73bf01c3f7a63..517455e37458a 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_controller.py +++ b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_controller.py @@ -42,6 +42,7 @@ ) from iotdb.ainode.core.log import Logger from iotdb.ainode.core.manager.model_manager import ModelManager +from iotdb.ainode.core.model.model_constants import ModelStates from iotdb.ainode.core.util.atmoic_int import AtomicInt from iotdb.ainode.core.util.batch_executor import BatchExecutor from iotdb.ainode.core.util.decorator import synchronized @@ -66,6 +67,9 @@ def __init__(self, result_queue: mp.Queue): # for pool instances control self._task_queue = queue.Queue() + # Models whose weights are downloading for a LOAD MODEL, with the load and unload + # tasks for them that arrived during the download. Only the worker thread uses it. + self._tasks_waiting_for_download: Dict[str, list] = {} self._pool_control_worker_thread = threading.Thread( target=self._worker_loop, daemon=True ) @@ -192,15 +196,39 @@ def _worker_loop(self): "PoolController received task None, the worker loop is existed." ) break - task_fn, args, kwargs = task try: - task_fn(*args, **kwargs) + self._run_task(*task) except Exception as e: logger.error(f"Error executing task: {e}") finally: self._task_queue.task_done() + def _run_task(self, task_fn, args, kwargs): + model_id = args[0] + if ( + task_fn != self._finish_download_task + and model_id in self._tasks_waiting_for_download + ): + # Keep the order of a model's tasks: run them after its download ends + self._tasks_waiting_for_download[model_id].append((task_fn, args, kwargs)) + return + task_fn(*args, **kwargs) + def _load_one_model_task(self, model_id: str, device_id_list: list[torch.device]): + model_info = self._model_manager.get_model_info(model_id) + if not model_info.download_on_start and model_info.state != ModelStates.ACTIVE: + # Built-in models that are not downloaded at start are downloaded when loaded. + # The download runs in its own thread, so that tasks for other models and + # stop() do not wait for it. + self._tasks_waiting_for_download[model_id] = [] + threading.Thread( + target=self._download_model, + args=(model_id, device_id_list), + name=f"{ThreadName.INFERENCE_POOL_CONTROLLER.value}-download-{model_id}", + daemon=True, + ).start() + return + def _load_one_model_on_device_task(device_id: torch.device): if not self.has_pool_on_device(device_id): self._expand_pools_on_device(model_id, device_id, 1) @@ -216,6 +244,37 @@ def _load_one_model_on_device_task(device_id: torch.device): load_model_futures, return_when=concurrent.futures.ALL_COMPLETED ) + def _download_model(self, model_id: str, device_id_list: list[torch.device]): + try: + succeeded = self._model_manager.download_builtin_model(model_id) + except Exception as e: + logger.error(f"[Inference] Error downloading model {model_id}: {e}") + succeeded = False + if not self._stop_event.is_set(): + self._task_queue.put( + ( + self._finish_download_task, + (model_id, device_id_list, succeeded), + {}, + ) + ) + + def _finish_download_task( + self, model_id: str, device_id_list: list[torch.device], succeeded: bool + ): + waiting_tasks = self._tasks_waiting_for_download.pop(model_id) + if succeeded: + self._load_one_model_task(model_id, device_id_list) + else: + logger.error( + f"[Inference] Cannot load model {model_id}: downloading it failed." + ) + for task in waiting_tasks: + try: + self._run_task(*task) + except Exception as e: + logger.error(f"Error executing task: {e}") + def _unload_one_model_task(self, model_id: str, device_id_list: list[torch.device]): def _unload_one_model_on_device_task(device_id: torch.device): if self.has_request_pools(model_id, device_id): diff --git a/iotdb-core/ainode/iotdb/ainode/core/manager/model_manager.py b/iotdb-core/ainode/iotdb/ainode/core/manager/model_manager.py index e1d67873b5334..c6eab05f75bd8 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/manager/model_manager.py +++ b/iotdb-core/ainode/iotdb/ainode/core/manager/model_manager.py @@ -87,6 +87,9 @@ def get_model_info( ) -> Optional[ModelInfo]: return self._model_storage.get_model_info(model_id, category) + def download_builtin_model(self, model_id: str) -> bool: + return self._model_storage.download_builtin_model(model_id) + def _refresh(self): """Refresh the model list (re-scan the file system)""" self._model_storage.discover_all_models() diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/model_info.py b/iotdb-core/ainode/iotdb/ainode/core/model/model_info.py index 50b0aaf66bcc2..0ad4c10ebcfa8 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/model/model_info.py +++ b/iotdb-core/ainode/iotdb/ainode/core/model/model_info.py @@ -32,6 +32,7 @@ def __init__( repo_id: str = "", auto_map: Optional[Dict] = None, hub_mixin_cls: Optional[str] = None, + download_on_start: bool = True, ): self.model_id = model_id self.model_type = model_type @@ -41,6 +42,8 @@ def __init__( self.repo_id = repo_id self.auto_map = auto_map self.hub_mixin_cls = hub_mixin_cls + # False for a built-in model whose weights are downloaded only when it is loaded + self.download_on_start = download_on_start def __repr__(self): return ( @@ -178,4 +181,16 @@ def __repr__(self): "AutoModelForCausalLM": "modeling_moment.MomentForPrediction", }, ), + "timestral": ModelInfo( + model_id="timestral", + category=ModelCategory.BUILTIN, + state=ModelStates.INACTIVE, + model_type="timestral", + pipeline_cls="pipeline_timestral.TimestralPipeline", + repo_id="mistralai/Timestral-209M", + hub_mixin_cls="modeling_timestral.TimestralForPrediction", + # The weights are not under an open-source licence, so AINode downloads them only + # when the model is loaded (LOAD MODEL), not at start + download_on_start=False, + ), } diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py b/iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py index b824fe84fc168..47632d821070b 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py +++ b/iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py @@ -20,6 +20,7 @@ import json import os import shutil +import threading from typing import Dict, Optional from huggingface_hub import hf_hub_download @@ -74,6 +75,8 @@ def __init__(self): } # Async download executor self._executor = concurrent.futures.ThreadPoolExecutor(max_workers=2) + # Serializes the on-request downloads of built-in models + self._download_lock = threading.Lock() # Thread lock pool for protecting concurrent access to model information self._lock_pool = ModelLockPool() self._initialize_directories() @@ -132,55 +135,117 @@ def _process_builtin_model_directory(self, model_dir: str, model_id: str): if existing_model.state in (ModelStates.ACTIVATING, ModelStates.ACTIVE): return + model_info = BUILTIN_HF_TRANSFORMERS_MODEL_MAP[model_id] + self._models[ModelCategory.BUILTIN.value][model_id] = model_info + # A model whose weights are not downloaded at start stays INACTIVE until it is + # loaded, which downloads them (see download_builtin_model) + if not model_info.download_on_start and not self._builtin_files_exist( + model_dir + ): + model_info.state = ModelStates.INACTIVE + logger.info( + f"Model {model_id} is not downloaded at start; " + f"LOAD MODEL {model_id} downloads it from {model_info.repo_id}." + ) + return + # If model not exists or is INACTIVE, we'll try to update its info and download its weights - self._models[ModelCategory.BUILTIN.value][model_id] = ( - BUILTIN_HF_TRANSFORMERS_MODEL_MAP[model_id] - ) - self._models[ModelCategory.BUILTIN.value][ - model_id - ].state = ModelStates.ACTIVATING - - def _download_model_if_necessary() -> bool: - """Returns: True if the model is existed or downloaded successfully, False otherwise.""" - repo_id = BUILTIN_HF_TRANSFORMERS_MODEL_MAP[model_id].repo_id - weights_path = os.path.join(model_dir, MODEL_SAFETENSORS) - config_path = os.path.join(model_dir, CONFIG_JSON) - if not os.path.exists(weights_path): - try: - hf_hub_download( - repo_id=repo_id, - filename=MODEL_SAFETENSORS, - local_dir=model_dir, - ) - except Exception as e: - logger.error( - f"Failed to download model weights from HuggingFace: {e}" - ) - return False - if not os.path.exists(config_path): - try: - hf_hub_download( - repo_id=repo_id, - filename=CONFIG_JSON, - local_dir=model_dir, - ) - except Exception as e: - logger.error( - f"Failed to download model config from HuggingFace: {e}" - ) - return False - return True + model_info.state = ModelStates.ACTIVATING - future = self._executor.submit(_download_model_if_necessary) + future = self._executor.submit( + self._download_builtin_model_files, model_id, model_dir + ) future.add_done_callback( lambda f, mid=model_id: self._callback_model_download_result(f, mid) ) + @staticmethod + def _builtin_files_exist(model_dir: str) -> bool: + return os.path.exists( + os.path.join(model_dir, MODEL_SAFETENSORS) + ) and os.path.exists(os.path.join(model_dir, CONFIG_JSON)) + + def _download_builtin_model_files(self, model_id: str, model_dir: str) -> bool: + """Returns: True if the model is existed or downloaded successfully, False otherwise.""" + repo_id = BUILTIN_HF_TRANSFORMERS_MODEL_MAP[model_id].repo_id + weights_path = os.path.join(model_dir, MODEL_SAFETENSORS) + config_path = os.path.join(model_dir, CONFIG_JSON) + if not os.path.exists(weights_path): + try: + hf_hub_download( + repo_id=repo_id, + filename=MODEL_SAFETENSORS, + local_dir=model_dir, + ) + except Exception as e: + logger.error(f"Failed to download model weights from HuggingFace: {e}") + return False + if not os.path.exists(config_path): + try: + hf_hub_download( + repo_id=repo_id, + filename=CONFIG_JSON, + local_dir=model_dir, + ) + except Exception as e: + logger.error(f"Failed to download model config from HuggingFace: {e}") + return False + return True + + def download_builtin_model(self, model_id: str) -> bool: + """ + Download the weights of a built-in model that is not downloaded at start, and wait + until the download ends. + Args: + model_id (str): Unique identifier for the model. + Returns: + bool: True if the model is ready to use, False if the download failed. + Raises: + ModelNotExistException: If the model_id is not a built-in model. + """ + with self._download_lock: + with self._lock_pool.get_lock(model_id).write_lock(): + model_info = self._models[ModelCategory.BUILTIN.value].get(model_id) + if ( + model_info is None + or model_id not in BUILTIN_HF_TRANSFORMERS_MODEL_MAP + ): + raise ModelNotExistException(model_id) + if model_info.state == ModelStates.ACTIVE: + return True + model_info.state = ModelStates.ACTIVATING + logger.info( + f"Downloading model {model_id} from HuggingFace: {model_info.repo_id}" + ) + model_dir = os.path.join( + self._models_dir, ModelCategory.BUILTIN.value, model_id + ) + succeeded = False + try: + succeeded = self._download_builtin_model_files(model_id, model_dir) + finally: + succeeded = self._set_model_download_result(model_id, succeeded) + return succeeded + def _callback_model_download_result(self, future, model_id: str): """Callback function for handling model download results""" + try: + succeeded = future.result() + except Exception as e: + logger.error(f"Error in download callback for model {model_id}: {e}") + succeeded = False + self._set_model_download_result(model_id, succeeded) + + def _set_model_download_result(self, model_id: str, succeeded: bool) -> bool: + """ + Mark a built-in model ACTIVE after a successful download, INACTIVE otherwise. + Returns: + bool: True if the model is ACTIVE, False if the download or reading its + config.json failed. + """ with self._lock_pool.get_lock(model_id).write_lock(): try: - if future.result(): + if succeeded: model_info = self._models[ModelCategory.BUILTIN.value][model_id] model_info.state = ModelStates.ACTIVE config_path = os.path.join( @@ -211,6 +276,10 @@ def _callback_model_download_result(self, future, model_id: str): self._models[ModelCategory.BUILTIN.value][ model_id ].state = ModelStates.INACTIVE + return ( + self._models[ModelCategory.BUILTIN.value][model_id].state + == ModelStates.ACTIVE + ) def _process_user_defined_model_directory(self, model_dir: str, model_id: str): """Handling the discovery logic for a user-defined model directory.""" diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/timestral/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/model/timestral/__init__.py new file mode 100644 index 0000000000000..05a224663c0e0 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/timestral/__init__.py @@ -0,0 +1,19 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# Copyright 2026 Mistral AI +# SPDX-License-Identifier: Apache-2.0 diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/timestral/config.py b/iotdb-core/ainode/iotdb/ainode/core/model/timestral/config.py new file mode 100644 index 0000000000000..2fefbd3c57d26 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/timestral/config.py @@ -0,0 +1,167 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# Copyright 2026 Mistral AI +# SPDX-License-Identifier: Apache-2.0 + +"""Inference configuration for Timestral checkpoints. + +Only the fields that change the forward pass at inference are kept. The package +implements one family of Timestral models: causal modeling trained with Decoupled Horizon +Forecasting (DHF), with monotone quantiles, decoded in the advancing refine +frame. +""" + +import json +from dataclasses import asdict, dataclass, fields +from pathlib import Path +from typing import Any + +LAYER_KINDS = ("time", "variate", "time_variate", "variate_time") +ATTENTION_GATES = ("none", "sdpa_sigmoid") + +# Ordered attention axes each layer kind runs before its shared MLP; True = time axis. +LAYER_KIND_AXES: dict[str, tuple[bool, ...]] = { + "time": (True,), + "variate": (False,), + "time_variate": (True, False), + "variate_time": (False, True), +} + + +@dataclass(frozen=True) +class TimestralConfig: + """Architecture and decode settings of a Timestral checkpoint. + + Attributes: + patch_size: Time steps per patch. + d_model: Model width. + num_heads: Attention heads per attention sublayer. + mlp_hidden_dim: Hidden width of the SwiGLU MLP. + block_schedule: One entry per layer, each one of `LAYER_KINDS`. + quantile_knots: Quantile levels of the head, symmetric about 0.5. + block_size: Trained DHF block size K, the maximum patches decoded per block. + context_length: Trained context length, the default context cap. + quantile_width_factor: Increment scale of the monotone quantile head. + """ + + patch_size: int + d_model: int + num_heads: int + mlp_hidden_dim: int + block_schedule: tuple[str, ...] + quantile_knots: tuple[float, ...] + block_size: int + context_length: int + rope_theta: float = 10000.0 + rope_partial_factor: float = 0.5 + norm_eps: float = 1e-6 + attn_bias: bool = False + mlp_bias: bool = False + time_attn_gate: str = "sdpa_sigmoid" + variate_attn_gate: str = "sdpa_sigmoid" + minimum_scale: float = 1e-6 + use_asinh: bool = True + stabilize_with_global: bool = False + covariate_enabled: bool = True + split_context_horizon_attn: bool = True + attend_missing_patches: bool = True + quantile_width_factor: float = 2.0 + + def __post_init__(self) -> None: + assert ( + self.d_model % self.num_heads == 0 + ), f"d_model={self.d_model} not divisible by num_heads={self.num_heads}" + assert ( + self.rope_dim > 0 and self.rope_dim % 2 == 0 + ), f"rope_partial_factor={self.rope_partial_factor} gives an invalid rope_dim={self.rope_dim}" + for kind in self.block_schedule: + assert ( + kind in LAYER_KINDS + ), f"unknown layer kind {kind!r}; expected one of {LAYER_KINDS}" + for gate in (self.time_attn_gate, self.variate_attn_gate): + assert ( + gate in ATTENTION_GATES + ), f"unknown attention gate {gate!r}; expected one of {ATTENTION_GATES}" + knots = self.quantile_knots + assert 0.5 in knots, f"quantile_knots must contain 0.5, got {knots}" + assert all( + a < b for a, b in zip(knots, knots[1:]) + ), f"quantile_knots must be increasing, got {knots}" + if self.stabilize_with_global: + raise ValueError( + "the advancing refine frame needs stabilize_with_global=False" + ) + + @property + def head_dim(self) -> int: + return self.d_model // self.num_heads + + @property + def rope_dim(self) -> int: + """Head dimensions rotated by RoPE; the rest pass through unrotated.""" + return int(self.head_dim * self.rope_partial_factor) + + @property + def num_quantiles(self) -> int: + return len(self.quantile_knots) + + @property + def median_index(self) -> int: + return self.quantile_knots.index(0.5) + + @property + def supports_known_covariates(self) -> bool: + """Whether the layer schedule can route a known covariate's future to the targets. + + Some time op must be followed by at least one more time op and at least one + variate op; otherwise the future of a known covariate never reaches a target. + """ + if not self.covariate_enabled: + return False + ops = [ + is_time for kind in self.block_schedule for is_time in LAYER_KIND_AXES[kind] + ] + for i, is_time in enumerate(ops): + rest = ops[i + 1 :] + if is_time and any(rest) and not all(rest): + return True + return False + + def to_dict(self) -> dict[str, Any]: + d = asdict(self) + d["block_schedule"] = list(self.block_schedule) + d["quantile_knots"] = list(self.quantile_knots) + return d + + @classmethod + def from_dict(cls, d: dict[str, Any]) -> "TimestralConfig": + known = {f.name for f in fields(cls)} + unknown = set(d) - known + if unknown: + raise ValueError(f"unknown config keys: {sorted(unknown)}") + d = dict(d) + d["block_schedule"] = tuple(d["block_schedule"]) + d["quantile_knots"] = tuple(float(q) for q in d["quantile_knots"]) + return cls(**d) + + def save(self, path: str | Path) -> None: + Path(path).write_text(json.dumps(self.to_dict(), indent=2) + "\n") + + @classmethod + def load(cls, path: str | Path) -> "TimestralConfig": + return cls.from_dict(json.loads(Path(path).read_text())) diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/timestral/covariate_encoding.py b/iotdb-core/ainode/iotdb/ainode/core/model/timestral/covariate_encoding.py new file mode 100644 index 0000000000000..587a8f3a2ee60 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/timestral/covariate_encoding.py @@ -0,0 +1,106 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# Copyright 2026 Mistral AI +# SPDX-License-Identifier: Apache-2.0 + +"""Target encoding of categorical covariates. + +Timestral was trained on numeric, target-scaled covariates, so a categorical covariate +is turned into numbers before it reaches the model. The default is a +per-series smoothed mean-target encoding: in series `i`, category `c` becomes + + (smooth * mean_i + sum_{i,c}) / (smooth + count_{i,c}) + +over the series' context, with `smooth = 1`. Missing values (NaN) form their own +category. A category that never occurs in the context (for example a new value in the +horizon of a known covariate) resolves to the series mean. +""" + +import numpy as np + + +def target_encode( + target_context: np.ndarray, + past_codes: np.ndarray, + future_codes: np.ndarray | None = None, + smooth: float = 1.0, +) -> tuple[np.ndarray, np.ndarray | None]: + """Per-series smoothed mean-target encoding. + + Args: + target_context: `(T,)` target values. Non-finite entries are left out of the means. + past_codes: `(T,)` category codes aligned with `target_context`; NaN marks a + missing value, which is encoded as a category of its own. + future_codes: Optional `(H,)` codes over the horizon (known covariates). + smooth: Weight of the series mean in each category mean. + + Returns: + `(encoded_past, encoded_future)` float32 arrays; `encoded_future` is None when + `future_codes` is None. + """ + target = np.asarray(target_context, dtype=np.float64) + past_codes = np.asarray(past_codes, dtype=np.float64) + assert ( + target.shape == past_codes.shape + ), f"target {target.shape} and codes {past_codes.shape} must align" + observed = np.isfinite(target) + target = np.where(observed, target, 0.0) + item_count = float(observed.sum()) + item_mean = float(target.sum() / item_count) if item_count > 0 else 0.0 + + # Slot per distinct finite code, plus a final slot for NaN. + categories = np.unique(past_codes[np.isfinite(past_codes)]) + nan_slot = len(categories) + past_slots = _slots(past_codes, categories, nan_slot) + sums = np.bincount(past_slots, weights=target * observed, minlength=nan_slot + 1) + counts = np.bincount( + past_slots, weights=observed.astype(np.float64), minlength=nan_slot + 1 + ) + lookup = np.append( + (smooth * item_mean + sums) / (smooth + counts), item_mean + ) # last entry: unseen + + encoded_past = lookup[past_slots].astype(np.float32) + if future_codes is None: + return encoded_past, None + future_slots = _slots( + np.asarray(future_codes, dtype=np.float64), + categories, + nan_slot, + unseen=nan_slot + 1, + ) + return encoded_past, lookup[future_slots].astype(np.float32) + + +def _slots( + codes: np.ndarray, categories: np.ndarray, nan_slot: int, unseen: int | None = None +) -> np.ndarray: + """Slot of each code in `categories`; NaN goes to `nan_slot`, unknown codes to `unseen`.""" + idx = np.searchsorted(categories, np.nan_to_num(codes, nan=0.0)) + idx = np.minimum(idx, max(len(categories) - 1, 0)) + known = np.isfinite(codes) & (len(categories) > 0) + if len(categories) > 0: + known &= categories[idx] == codes + assert ( + unseen is not None or (known | np.isnan(codes)).all() + ), "every past code must be in categories" + return np.where( + np.isnan(codes), + nan_slot, + np.where(known, idx, unseen if unseen is not None else 0), + ).astype(np.intp) diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/timestral/forecaster.py b/iotdb-core/ainode/iotdb/ainode/core/model/timestral/forecaster.py new file mode 100644 index 0000000000000..8f5bace61e500 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/timestral/forecaster.py @@ -0,0 +1,628 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# Copyright 2026 Mistral AI +# SPDX-License-Identifier: Apache-2.0 + +"""Probabilistic forecasting with Timestral, including the clampsym test-time augmentation. + +`TimestralForecaster.forecast` takes a context `(B, V, T)` (series, variates, time; NaN +marks a missing value) and returns quantile forecasts `(B, V, H, Q)`. Horizons longer +than one decode block are forecast block by block, feeding each block's median back as +context. + +Clampsym (both options on by default): + +- Symmetric averaging forecasts the context and its negation, reflects the second + forecast back (`-Q(1 - p)`), and averages the two. Every covariate row is negated + with the context: numeric covariates, 0/1 indicators, target encodings and raw codes. +- Non-negative clipping clamps the forecast at 0 for every series whose observed + context is entirely non-negative. + +Categorical covariates are target-encoded against the series' forward-filled context +(see :mod:`timestral.covariate_encoding`) and then handled like numeric covariates. +""" + +import math +from dataclasses import dataclass + +import numpy as np +import torch +from torch import Tensor + +from iotdb.ainode.core.model.timestral.covariate_encoding import target_encode +from iotdb.ainode.core.model.timestral.model import ( + ROLE_KNOWN_COV, + ROLE_PAST_COV, + ROLE_TARGET, + TimestralModel, +) +from iotdb.ainode.core.model.timestral.scaler import ( + advancing_block_frames, + causal_patch_loc_scale, + clamp_nonfinite, +) + +ArrayLike = Tensor | np.ndarray + + +@dataclass(frozen=True) +class Forecast: + """Quantile forecast. + + Attributes: + quantiles: `(B, V, H, Q)` forecast values, sorted along the quantile axis. + quantile_levels: The `Q` probability levels. + """ + + quantiles: Tensor + quantile_levels: tuple[float, ...] + + @property + def median(self) -> Tensor: + """`(B, V, H)` median forecast.""" + return self.quantiles[..., self.quantile_levels.index(0.5)] + + def quantile(self, level: float) -> Tensor: + """`(B, V, H)` forecast at one of the model's quantile levels.""" + if level not in self.quantile_levels: + raise ValueError(f"level {level} is not one of {self.quantile_levels}") + return self.quantiles[..., self.quantile_levels.index(level)] + + +def ffill_with_mask(inputs: Tensor, mask: Tensor) -> Tensor: + """Forward-fill unobserved positions; leading unobserved positions become 0.""" + if mask.all(): + return inputs + idx = torch.arange(inputs.shape[-1], device=inputs.device).expand_as(inputs) + idx = torch.where(mask, idx, torch.zeros_like(idx)) + idx = idx.cummax(dim=-1).values + return torch.nan_to_num(torch.gather(inputs, dim=-1, index=idx), nan=0.0) + + +def pad_to_patch( + inputs: Tensor, mask: Tensor, patch_size: int +) -> tuple[Tensor, Tensor]: + """Left-pad to a multiple of `patch_size`; padding is 0 and unobserved.""" + ctx_len = inputs.shape[-1] + if ctx_len % patch_size != 0: + pad_len = patch_size - (ctx_len % patch_size) + inputs = torch.nn.functional.pad(inputs, (pad_len, 0), value=0.0) + mask = torch.nn.functional.pad(mask, (pad_len, 0), value=False) + return inputs, mask + + +def replace_extreme_values( + x: Tensor, threshold: float, replacement: float = 0.0 +) -> Tensor: + return torch.where( + torch.isfinite(x) & (x.abs() < threshold), x, torch.full_like(x, replacement) + ) + + +def nonnegative_context(inputs: Tensor, observation_mask: Tensor) -> Tensor: + """`(B, V)`: whether every observed context value is non-negative (and one exists).""" + observed = torch.isfinite(inputs) & observation_mask.bool() + return observed.any(dim=-1) & ((inputs >= 0) | ~observed).all(dim=-1) + + +class TimestralForecaster: + """Block-autoregressive quantile forecaster around a :class:`TimestralModel`. + + Args: + model: The network. + max_context: Context cap in time steps (most recent kept); defaults to the + trained context length. + per_shot_horizon_patches: Patches decoded per block, capped at the trained + block size. + quantile_real_cap_k: Clamp forecasts to `[min - k * s, max + k * s]` of the + observed context, `s` being the scale of the last context patch. 0 disables. + extreme_value_threshold: Fed-back medians at or above this magnitude become 0. + """ + + def __init__( + self, + model: TimestralModel, + max_context: int | None = None, + per_shot_horizon_patches: int = 32, + quantile_real_cap_k: float = 1e4, + extreme_value_threshold: float = 1e30, + ) -> None: + assert ( + per_shot_horizon_patches >= 1 + ), f"per_shot_horizon_patches must be >= 1, got {per_shot_horizon_patches}" + self.model = model + self.config = model.config + self.max_context = ( + max_context if max_context is not None else model.config.context_length + ) + self.per_shot_horizon_patches = min( + per_shot_horizon_patches, model.config.block_size + ) + self.quantile_real_cap_k = quantile_real_cap_k + self.extreme_value_threshold = extreme_value_threshold + self.quantile_levels = tuple(model.config.quantile_knots) + self.median_index = model.config.median_index + + @property + def patch_size(self) -> int: + return self.config.patch_size + + @property + def device(self) -> torch.device: + return next(self.model.parameters()).device + + @property + def supports_covariates(self) -> bool: + return self.config.covariate_enabled + + @property + def supports_known_covariates(self) -> bool: + return self.config.supports_known_covariates + + @torch.inference_mode() + def forecast( + self, + context: ArrayLike, + prediction_length: int, + observation_mask: ArrayLike | None = None, + past_covariates: ArrayLike | None = None, + known_covariates: ArrayLike | None = None, + categorical_past_covariates: ArrayLike | None = None, + categorical_known_covariates: ArrayLike | None = None, + symmetric_averaging: bool = True, + clip_nonnegative: bool = True, + ) -> Forecast: + """Forecast `prediction_length` steps after `context`. + + Args: + context: `(B, V, T)`, `(B, T)` (one variate) or `(T,)` (one series). NaN marks + a missing value. The variates of a series are forecast jointly. + prediction_length: Horizon in time steps. + observation_mask: Optional boolean mask shaped like `context`; False marks + additional missing values. + past_covariates: `(B, C_past, T)` covariates known over the context only. + known_covariates: `(B, C_known, T + prediction_length)` covariates known over + the context and the horizon. + categorical_past_covariates: `(B, C, T)` category codes (any numbers; NaN marks + a missing value) known over the context only. + categorical_known_covariates: `(B, C, T + prediction_length)` category codes + known over the context and the horizon. + Categorical covariates are target-encoded per series against the + forward-filled target context. A series with several target variates + has no single target to encode against, so its codes are used as is. + symmetric_averaging: Average with the reflected forecast of `-context`, with every + covariate row negated too. + clip_nonnegative: Clamp at 0 the series whose observed context is non-negative. + + Returns: + A :class:`Forecast` with `quantiles` `(B, V, prediction_length, Q)`. + """ + assert ( + prediction_length >= 1 + ), f"prediction_length must be >= 1, got {prediction_length}" + inputs, mask = self._as_context(context, observation_mask) + b = inputs.shape[0] + t = inputs.shape[-1] + past = self._as_covariates(past_covariates, b, t, "past_covariates") + known = self._as_covariates( + known_covariates, b, t + prediction_length, "known_covariates" + ) + cat_past = self._as_categorical( + categorical_past_covariates, b, t, "categorical_past_covariates" + ) + cat_known = self._as_categorical( + categorical_known_covariates, + b, + t + prediction_length, + "categorical_known_covariates", + ) + covariates = {"past_covariates": past, "known_covariates": known} + if cat_past is not None or cat_known is not None: + target = self._encoding_target(inputs, mask) + enc_past = self._encode_categorical(cat_past, target) + enc_known = self._encode_categorical(cat_known, target) + covariates = { + "past_covariates": _cat_rows(past, enc_past), + "known_covariates": _cat_rows(known, enc_known), + } + # The mirrored pass forecasts -context with every covariate row negated. Target encodings + # are linear in the target, so negating them gives the encodings against -context. + mirrored_covariates = {name: _negate(rows) for name, rows in covariates.items()} + if ( + covariates["known_covariates"] is not None + and not self.supports_known_covariates + ): + raise ValueError( + "this checkpoint cannot condition on known-future covariates" + ) + + quantiles = self._forecast_raw(inputs, mask, prediction_length, **covariates) + if symmetric_averaging: + levels = self.quantile_levels + assert all( + abs(lo + hi - 1.0) < 1e-9 for lo, hi in zip(levels, reversed(levels)) + ), f"symmetric averaging needs quantile levels symmetric about 0.5, got {levels}" + mirrored = self._forecast_raw( + -inputs, mask, prediction_length, **mirrored_covariates + ) + quantiles = 0.5 * (quantiles + (-mirrored).flip(-1)) + if clip_nonnegative: + eligible = nonnegative_context(inputs, mask) + quantiles = torch.where( + eligible[..., None, None], quantiles.clamp_min(0), quantiles + ) + return Forecast(quantiles=quantiles, quantile_levels=self.quantile_levels) + + def _as_context( + self, context: ArrayLike, observation_mask: ArrayLike | None + ) -> tuple[Tensor, Tensor]: + x = torch.as_tensor(context, dtype=torch.float32, device=self.device) + if x.dim() == 1: + x = x[None, None] + elif x.dim() == 2: + x = x[:, None] + assert ( + x.dim() == 3 + ), f"context must be (T,), (B, T) or (B, V, T), got {tuple(x.shape)}" + mask = torch.isfinite(x) + if observation_mask is not None: + user_mask = torch.as_tensor( + observation_mask, dtype=torch.bool, device=self.device + ).reshape(x.shape) + mask = mask & user_mask + return torch.where(mask, x, torch.zeros_like(x)), mask + + def _encoding_target(self, inputs: Tensor, mask: Tensor) -> Tensor | None: + """`(B, T)` target to encode against, or None when a series has several target variates. + + The forward-filled target from its first observation on (so left padding or a late + start does not enter the category means), over the last `max_context` steps, the + ones the model reads; NaN elsewhere. + """ + if inputs.shape[1] != 1: + return None + t = inputs.shape[-1] + started = mask[:, 0].int().cummax(dim=-1).values.bool() + in_window = torch.arange(t, device=inputs.device) >= t - self.max_context + return torch.where( + started & in_window, ffill_with_mask(inputs, mask)[:, 0], torch.nan + ) + + def _encode_categorical( + self, codes: Tensor | None, target_context: Tensor | None + ) -> Tensor | None: + """Target-encode `(B, C, L)` codes against `(B, T)` targets (NaN outside the encoding window); raw codes when None.""" + if codes is None or target_context is None: + return codes.to(self.device, torch.float32) if codes is not None else None + t = target_context.shape[-1] + ctx = target_context.double().cpu().numpy() + rows = codes.double().cpu().numpy() + out = np.empty(rows.shape, dtype=np.float32) + for i in range(rows.shape[0]): + for j in range(rows.shape[1]): + future = rows[i, j, t:] if rows.shape[-1] > t else None + past_enc, future_enc = target_encode(ctx[i], rows[i, j, :t], future) + # A missing code before the target starts is padding, not the missing category. + out[i, j, :t] = np.where( + np.isnan(ctx[i]) & np.isnan(rows[i, j, :t]), np.nan, past_enc + ) + if future_enc is not None: + out[i, j, t:] = future_enc + return torch.from_numpy(out).to(self.device) + + def _as_categorical( + self, values: ArrayLike | None, batch: int, length: int, name: str + ) -> Tensor | None: + """Category codes as float64 on CPU, so that distinct codes past 2**24 stay distinct.""" + if values is None: + return None + v = self._as_covariates( + torch.as_tensor(values, dtype=torch.float64), + batch, + length, + name, + dtype=torch.float64, + ) + if v is not None and torch.isinf(v).any(): + raise ValueError( + f"{name} must hold finite codes or NaN; got an infinite value" + ) + return v.cpu() if v is not None else None + + def _as_covariates( + self, + values: ArrayLike | None, + batch: int, + length: int, + name: str, + dtype: torch.dtype = torch.float32, + ) -> Tensor | None: + if values is None: + return None + if not self.supports_covariates: + raise ValueError( + f"this checkpoint was not trained with covariates; got {name}" + ) + v = torch.as_tensor( + values, dtype=dtype, device=self.device if dtype == torch.float32 else "cpu" + ) + if v.dim() == 2: + v = v[:, None] + assert ( + v.dim() == 3 and v.shape[0] == batch and v.shape[-1] == length + ), f"{name} must be (B={batch}, C, {length}), got {tuple(v.shape)}" + return v if v.shape[1] > 0 else None + + def _forecast_raw( + self, + inputs: Tensor, + observation_mask: Tensor, + prediction_length: int, + past_covariates: Tensor | None = None, + known_covariates: Tensor | None = None, + ) -> Tensor: + """One forecast pass: `(B, V, T)` inputs to `(B, V, H, Q)` raw-space quantiles.""" + self.model.eval() + rows, mask, role_ids, known_future, known_future_mask, n_tgt = ( + self._assemble_rows( + inputs, + observation_mask, + prediction_length, + past_covariates, + known_covariates, + ) + ) + quantiles = self._block_ar_decode( + rows, mask, prediction_length, role_ids, known_future, known_future_mask + ) + quantiles = quantiles[:, :n_tgt] + if self.quantile_real_cap_k > 0: + quantiles = self._cap_quantiles(quantiles, rows[:, :n_tgt], mask[:, :n_tgt]) + if torch.isnan(quantiles).any(): + raise ValueError( + f"forecast produced {int(torch.isnan(quantiles).sum())} NaN values; the input is degenerate" + ) + return quantiles + + def _cap_quantiles( + self, quantiles: Tensor, inputs: Tensor, observation_mask: Tensor + ) -> Tensor: + not_obs = ~observation_mask + ctx_max = torch.nan_to_num( + inputs.masked_fill(not_obs, float("-inf")).amax(-1, keepdim=True), + neginf=0.0, + ) + ctx_min = torch.nan_to_num( + inputs.masked_fill(not_obs, float("inf")).amin(-1, keepdim=True), posinf=0.0 + ) + c = self.config + _, scale = causal_patch_loc_scale( + inputs, + observation_mask, + c.patch_size, + c.minimum_scale, + c.stabilize_with_global, + ) + anchor_scale = scale[..., -1:] * self.quantile_real_cap_k + return quantiles.clamp( + (ctx_min - anchor_scale).unsqueeze(-1), + (ctx_max + anchor_scale).unsqueeze(-1), + ) + + def _assemble_rows( + self, + inputs: Tensor, + observation_mask: Tensor, + prediction_length: int, + past_covariates: Tensor | None, + known_covariates: Tensor | None, + ) -> tuple[Tensor, Tensor, Tensor | None, Tensor | None, Tensor | None, int]: + """Stack target and covariate rows; returns `(rows, mask, role_ids, known_future, known_future_mask, n_tgt)`.""" + ps = self.patch_size + device = inputs.device + b, n_tgt, ctx0 = inputs.shape + horizon = math.ceil(prediction_length / ps) * ps + + ctx_rows = [ffill_with_mask(inputs, observation_mask)] + ctx_masks = [observation_mask] + roles = [ROLE_TARGET] * n_tgt + + n_past = 0 + if past_covariates is not None: + assert ( + past_covariates.shape[-1] == ctx0 + ), "past covariates must cover exactly the context" + n_past = past_covariates.shape[1] + ctx_masks.append(~torch.isnan(past_covariates)) + ctx_rows.append(torch.nan_to_num(past_covariates, nan=0.0)) + roles += [ROLE_PAST_COV] * n_past + + n_known = 0 + known_future_vals: Tensor | None = None + known_future_obs: Tensor | None = None + if known_covariates is not None: + n_known = known_covariates.shape[1] + kc_ctx, kc_fut = ( + known_covariates[..., :-prediction_length], + known_covariates[..., -prediction_length:], + ) + assert ( + kc_ctx.shape[-1] == ctx0 + ), "known covariates must cover the context plus the horizon" + ctx_masks.append(~torch.isnan(kc_ctx)) + ctx_rows.append(torch.nan_to_num(kc_ctx, nan=0.0)) + roles += [ROLE_KNOWN_COV] * n_known + known_future_obs = ~torch.isnan(kc_fut) + known_future_vals = torch.nan_to_num(kc_fut, nan=0.0) + + rows = torch.cat(ctx_rows, dim=1) + mask = torch.cat(ctx_masks, dim=1) + if rows.shape[-1] > self.max_context: + rows = rows[:, :, -self.max_context :] + mask = mask[:, :, -self.max_context :] + rows, mask = pad_to_patch(rows, mask, ps) + v_all = rows.shape[1] + if n_past == 0 and n_known == 0: + return rows, mask, None, None, None, n_tgt + + role_ids = torch.tensor(roles, dtype=torch.long, device=device).repeat(b) + known_future: Tensor | None = None + known_future_mask: Tensor | None = None + if n_known > 0: + assert known_future_vals is not None and known_future_obs is not None + k0 = n_tgt + n_past + plen = known_future_vals.shape[-1] + kf = torch.zeros(b, v_all, horizon, device=device, dtype=rows.dtype) + kfm = torch.zeros(b, v_all, horizon, device=device, dtype=torch.bool) + kf[:, k0 : k0 + n_known, :plen] = known_future_vals + kfm[:, k0 : k0 + n_known, :plen] = known_future_obs + known_future = kf.reshape(b * v_all, horizon) + known_future_mask = kfm.reshape(b * v_all, horizon) + return rows, mask, role_ids, known_future, known_future_mask, n_tgt + + def _block_ar_decode( + self, + inputs: Tensor, + observation_mask: Tensor, + prediction_length: int, + role_ids: Tensor | None, + known_future: Tensor | None, + known_future_mask: Tensor | None, + ) -> Tensor: + """Decode block by block in the packed `(R, T)` layout; returns `(B, V, H, Q)`. + + After each block the context grows by the block's feedback: targets carry their + median, known-future covariates their real future, past-only covariates a blank. + Past `max_context` the oldest steps are dropped. + """ + ps = self.patch_size + b, v, ctx_len = inputs.shape + context = inputs.reshape(b * v, ctx_len) + mask = observation_mask.reshape(b * v, ctx_len) + group_ids = torch.arange(b, device=inputs.device).repeat_interleave(v) + + known_rows = past_rows = target_rows = None + if role_ids is not None: + known_rows = role_ids == ROLE_KNOWN_COV + past_rows = (role_ids != ROLE_TARGET) & ~known_rows + target_rows = role_ids == ROLE_TARGET + + blocks: list[Tensor] = [] + remaining = math.ceil(prediction_length / ps) + offset = 0 + while remaining > 0: + block = min(self.per_shot_horizon_patches, remaining) + blen = block * ps + kf_block = ( + known_future[:, offset : offset + blen] + if known_future is not None + else None + ) + kfm_block = ( + known_future_mask[:, offset : offset + blen] + if known_future_mask is not None + else None + ) + quantiles = self.model.predict_block( + context, mask, group_ids, block, role_ids, kf_block, kfm_block + ) + quantiles = self._refine_block(quantiles, context, mask, target_rows) + blocks.append(quantiles) + remaining -= block + offset += blen + if remaining <= 0: + break + feedback = replace_extreme_values( + quantiles[..., self.median_index], self.extreme_value_threshold + ) + feedback_mask = torch.ones_like(feedback, dtype=torch.bool) + if role_ids is not None: + assert known_rows is not None and past_rows is not None + if kf_block is not None: + assert kfm_block is not None + feedback = torch.where(known_rows[:, None], kf_block, feedback) + feedback_mask = torch.where( + known_rows[:, None], kfm_block, feedback_mask + ) + feedback = torch.where( + past_rows[:, None], torch.zeros_like(feedback), feedback + ) + feedback_mask = feedback_mask & ~past_rows[:, None] + context = torch.cat([context, feedback], dim=-1) + mask = torch.cat([mask, feedback_mask], dim=-1) + if context.shape[-1] > self.max_context: + keep = max(ps, (self.max_context // ps) * ps) + context = context[..., -keep:] + mask = mask[..., -keep:] + + packed = torch.cat(blocks, dim=-2)[..., :prediction_length, :] + return packed.reshape(b, v, prediction_length, -1) + + def _refine_block( + self, + quantiles: Tensor, + context: Tensor, + mask: Tensor, + target_rows: Tensor | None, + ) -> Tensor: + """Move one decoded block from the anchor frame to the advancing refine frame it was trained in. + + The model denormalizes the whole block `(R, H * patch_size, Q)` with the frame of the + last context patch. The first horizon patch keeps that frame; each later one is + re-denormalized with :func:`timestral.scaler.advancing_block_frames`. Covariate rows + keep their quantiles. + """ + ps = self.patch_size + n_pred = quantiles.shape[1] // ps + if n_pred <= 1: + return quantiles + _, loc, scale = self.model.scale(context, mask) + loc0, scale0 = loc[:, -1], scale[:, -1] + median = ( + quantiles[..., self.median_index] + .unflatten(-1, (n_pred, ps)) + .to(torch.float64) + ) + sinh_median = (median - loc0[:, None, None]) / scale0[:, None, None] + loc_j, scale_j = advancing_block_frames( + context, mask, sinh_median, self.config.minimum_scale + ) + first = torch.zeros(loc_j.shape, dtype=torch.bool, device=loc_j.device) + first[:, 0] = True + loc_j = torch.where( + first, loc0[:, None], loc_j.to(quantiles.dtype) + ).repeat_interleave(ps, dim=-1) + scale_j = torch.where( + first, scale0[:, None], scale_j.to(quantiles.dtype) + ).repeat_interleave(ps, dim=-1) + sinh_norm = (quantiles - loc0[:, None, None]) / scale0[:, None, None] + refined = clamp_nonfinite(sinh_norm * scale_j[..., None] + loc_j[..., None]) + if target_rows is not None: + refined = torch.where(target_rows[:, None, None], refined, quantiles) + return refined + + +def _negate(x: Tensor | None) -> Tensor | None: + return -x if x is not None else None + + +def _cat_rows(a: Tensor | None, b: Tensor | None) -> Tensor | None: + """Concatenate covariate rows along the channel axis, either side possibly None.""" + if a is None or b is None: + return a if b is None else b + return torch.cat([a, b], dim=1) diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/timestral/layers.py b/iotdb-core/ainode/iotdb/ainode/core/model/timestral/layers.py new file mode 100644 index 0000000000000..4da2a0469a5e8 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/timestral/layers.py @@ -0,0 +1,336 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# Copyright 2026 Mistral AI +# SPDX-License-Identifier: Apache-2.0 + +"""Transformer layers for Timestral inference. + +Parameter names match the released weights, so they load with +`strict=True`. All tensors use the packed layout: one row per variate of every series in +the batch, `(R, N, D)` with `N` patches. Time attention runs per row over patches; variate +attention runs per patch over the rows of the same series. +""" + +from dataclasses import dataclass + +import torch +import torch.nn.functional as F +from torch import Tensor, nn + +from iotdb.ainode.core.model.timestral.config import LAYER_KIND_AXES + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float) -> None: + super().__init__() + self.eps = eps + self.scale = nn.Parameter(torch.ones(dim)) + + def forward(self, x: Tensor) -> Tensor: + xf = x.float() + out = (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)).type_as(x) + return out * self.scale + + +class ResidualBlock(nn.Module): + """`output(silu(hidden(x))) + residual(x)`, used by the patch embedding and the head.""" + + def __init__( + self, input_dims: int, hidden_dims: int, output_dims: int, use_bias: bool + ) -> None: + super().__init__() + self.hidden_layer = nn.Linear(input_dims, hidden_dims, bias=use_bias) + self.output_layer = nn.Linear(hidden_dims, output_dims, bias=use_bias) + self.residual_layer = nn.Linear(input_dims, output_dims, bias=use_bias) + self.activation = nn.SiLU() + + def forward(self, x: Tensor) -> Tensor: + out = self.output_layer(self.activation(self.hidden_layer(x))) + return out + self.residual_layer(x) + + +class FusedSwiGLU(nn.Module): + """SwiGLU MLP with one fused input projection; computes `gate * silu(up)`.""" + + def __init__(self, embed_dim: int, hidden_dim: int, use_bias: bool) -> None: + super().__init__() + self.w12 = nn.Linear(embed_dim, 2 * hidden_dim, bias=use_bias) + self.w3 = nn.Linear(hidden_dim, embed_dim, bias=use_bias) + + def forward(self, x: Tensor) -> Tensor: + gate, up = self.w12(x).chunk(2, dim=-1) + return self.w3(gate * F.silu(up)) + + +class SinCosRoPE(nn.Module): + """Rotary position embedding (sin/cos concatenation), with partial rotation. + + Only the first `rotary_dim` dimensions of each head rotate. The inverse-frequency + table is zero-padded past `rotary_dim / 2`, so the remaining dimensions get `cos=1`, + `sin=0`, which is the identity rotation. + """ + + inv_freq: Tensor + + def __init__(self, dim: int, base: float, rotary_dim: int) -> None: + super().__init__() + assert ( + 0 < rotary_dim <= dim and rotary_dim % 2 == 0 + ), f"invalid rotary_dim={rotary_dim} for dim={dim}" + self.dim = dim + self.base = base + self.rotary_dim = rotary_dim + self.register_buffer("inv_freq", self._build_inv_freq(None), persistent=False) + + def _build_inv_freq(self, device: torch.device | None) -> Tensor: + inv_freq = torch.zeros(self.dim // 2, dtype=torch.float32, device=device) + idx = torch.arange(0, self.rotary_dim, 2, dtype=torch.float32, device=device) + inv_freq[: self.rotary_dim // 2] = 1.0 / (self.base ** (idx / self.rotary_dim)) + return inv_freq + + def rebuild_inv_freq(self) -> None: + """Recompute the table on the buffer's current device. + + A table computed on CPU and then moved could differ in the last bit. + """ + self.inv_freq = self._build_inv_freq(self.inv_freq.device) + + @torch.no_grad() + def forward(self, x: Tensor, position_ids: Tensor) -> tuple[Tensor, Tensor]: + """Return `(cos, sin)` tables, each `(R, N, head_dim)`, for integer `position_ids` `(R, N)`.""" + inv_freq_expanded = ( + self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + ) + position_ids_expanded = position_ids[:, None, :].float() + device_type = x.device.type if x.device.type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = (inv_freq_expanded @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +def _rotate_half(x: Tensor) -> Tensor: + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def _apply_rotary( + q: Tensor, k: Tensor, cos: Tensor, sin: Tensor +) -> tuple[Tensor, Tensor]: + """Rotate `q`, `k` shaped `(batch, seq, heads, head_dim)`; the tables broadcast over heads.""" + cos = cos.unsqueeze(2) + sin = sin.unsqueeze(2) + return (q * cos) + (_rotate_half(q) * sin), (k * cos) + (_rotate_half(k) * sin) + + +@dataclass(frozen=True) +class TwoStreamMask: + """Time mask over the `[Enc ; Dec]` sequence plus its split point. + + Attention with this mask runs two query passes: context queries `[:enc_len]` over + context keys only, and horizon queries `[enc_len:]` over all keys. This skips the + always-empty context-to-horizon quadrant. + """ + + mask: Tensor # (R, 1, S, S) bool, True = attend + enc_len: int + + +def _sdpa( + query: Tensor, + key: Tensor, + value: Tensor, + mask: Tensor | None, + is_causal: bool = False, +) -> Tensor: + """Scaled dot-product attention on `(batch, seq, heads, head_dim)`; returns `(batch, seq, heads * head_dim)`.""" + batch, seqlen, _, _ = query.shape + out = F.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + attn_mask=mask, + is_causal=is_causal, + dropout_p=0.0, + ) + return out.transpose(1, 2).contiguous().view(batch, seqlen, -1) + + +class Attention(nn.Module): + """Multi-head attention over the time axis (`is_time=True`) or the variate axis. + + Time attention applies RoPE; variate attention has no positional encoding. With the + `sdpa_sigmoid` gate, the attention output is multiplied by `sigmoid(x @ W_g)` before + the output projection. + """ + + def __init__( + self, embed_dim: int, num_heads: int, is_time: bool, use_bias: bool, gate: str + ) -> None: + super().__init__() + self.embed_dim = embed_dim + self.num_heads = num_heads + self.head_dim = embed_dim // num_heads + self.is_time = is_time + self.wQKV = nn.Linear(embed_dim, 3 * embed_dim, bias=use_bias) + self.wO = nn.Linear(embed_dim, embed_dim, bias=use_bias) + self._gate = ( + nn.Linear(embed_dim, embed_dim, bias=use_bias) + if gate == "sdpa_sigmoid" + else None + ) + + def _qkv(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]: + """Split the fused projection. The weight layout is `(qkv, head_dim, n_heads)`.""" + batch, seq, _ = x.shape + qkv = self.wQKV(x.contiguous()).reshape( + batch, seq, 3, self.head_dim, self.num_heads + ) + q, k, v = qkv.permute(2, 0, 1, 4, 3).unbind(dim=0) + return q.contiguous(), k.contiguous(), v.contiguous() + + def forward( + self, x: Tensor, mask: Tensor | TwoStreamMask | None, cos: Tensor, sin: Tensor + ) -> Tensor: + rows, num_patches, _ = x.shape + inputs = ( + x if self.is_time else x.permute(1, 0, 2) + ) # time: (R, N, D); variate: (N, R, D) + q, k, v = self._qkv(inputs) + if self.is_time: + q, k = _apply_rotary(q, k, cos, sin) + if isinstance(mask, TwoStreamMask): + assert ( + self.is_time + ), "the context/horizon split applies to time attention only" + n = mask.enc_len + out_enc = _sdpa(q[:, :n], k[:, :n], v[:, :n], mask.mask[:, :, :n, :n]) + out_dec = _sdpa(q[:, n:], k, v, mask.mask[:, :, n:, :]) + out = torch.cat([out_enc, out_dec], dim=1) + else: + out = _sdpa(q, k, v, mask, is_causal=self.is_time and mask is None) + if self._gate is not None: + out = out * self._gate(inputs).sigmoid() + if self.is_time: + out = out.reshape(rows, num_patches, self.embed_dim) + else: + out = out.reshape(num_patches, rows, self.embed_dim).permute(1, 0, 2) + return self.wO(out) + + +class TransformerBlock(nn.Module): + """The attention(s) of one layer kind, in order, each a pre-norm residual, then one MLP.""" + + def __init__( + self, + kind: str, + embed_dim: int, + num_heads: int, + mlp_hidden_dim: int, + norm_eps: float, + attn_bias: bool, + mlp_bias: bool, + time_attn_gate: str, + variate_attn_gate: str, + ) -> None: + super().__init__() + axes = LAYER_KIND_AXES[kind] + self.attn_norms = nn.ModuleList([RMSNorm(embed_dim, norm_eps) for _ in axes]) + self.attns = nn.ModuleList( + [ + Attention( + embed_dim, + num_heads, + is_time, + attn_bias, + time_attn_gate if is_time else variate_attn_gate, + ) + for is_time in axes + ] + ) + self.mlp_norm = RMSNorm(embed_dim, norm_eps) + self.mlp = FusedSwiGLU(embed_dim, mlp_hidden_dim, mlp_bias) + + def forward( + self, + x: Tensor, + time_mask: Tensor | TwoStreamMask | None, + variate_mask: Tensor | None, + cos: Tensor, + sin: Tensor, + ) -> Tensor: + for norm, attn in zip(self.attn_norms, self.attns): + x = x + attn(norm(x), time_mask if attn.is_time else variate_mask, cos, sin) + return x + self.mlp(self.mlp_norm(x)) + + +class Transformer(nn.Module): + def __init__( + self, + block_schedule: tuple[str, ...], + embed_dim: int, + num_heads: int, + mlp_hidden_dim: int, + rope_theta: float, + rope_dim: int, + norm_eps: float, + attn_bias: bool, + mlp_bias: bool, + time_attn_gate: str, + variate_attn_gate: str, + ) -> None: + super().__init__() + self.layers = nn.ModuleList( + [ + TransformerBlock( + kind=kind, + embed_dim=embed_dim, + num_heads=num_heads, + mlp_hidden_dim=mlp_hidden_dim, + norm_eps=norm_eps, + attn_bias=attn_bias, + mlp_bias=mlp_bias, + time_attn_gate=time_attn_gate, + variate_attn_gate=variate_attn_gate, + ) + for kind in block_schedule + ] + ) + self.rope = SinCosRoPE( + dim=embed_dim // num_heads, base=rope_theta, rotary_dim=rope_dim + ) + self.out_norm = RMSNorm(embed_dim, norm_eps) + + def forward( + self, + x: Tensor, + time_mask: Tensor | TwoStreamMask | None, + variate_mask: Tensor | None, + positions: Tensor, + ) -> Tensor: + """Run the stack on `x` `(R, N, D)` with integer time `positions` `(R, N)`.""" + assert ( + positions.shape == x.shape[:2] + ), f"positions {tuple(positions.shape)} do not match x {tuple(x.shape)}" + cos, sin = self.rope(x, positions) + for layer in self.layers: + x = layer(x, time_mask, variate_mask, cos, sin) + return self.out_norm(x) diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/timestral/masks.py b/iotdb-core/ainode/iotdb/ainode/core/model/timestral/masks.py new file mode 100644 index 0000000000000..8c50dbe900e70 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/timestral/masks.py @@ -0,0 +1,201 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# Copyright 2026 Mistral AI +# SPDX-License-Identifier: Apache-2.0 + +"""Attention masks and time positions for Decoupled Horizon Forecasting (DHF) decoding. + +The sequence is `[Enc ; Dec]`: encoder tokens embed the observed context patches; +decoder tokens are forecast queries for the horizon patches. All masks are boolean, +True meaning "may attend". +""" + +import torch +from torch import Tensor + + +def compute_enc_key_validity(obs_per_patch: Tensor, attend_missing: bool) -> Tensor: + """`(R, N)` encoder key validity. + + A patch is a valid key if it is observed, or if `attend_missing` is set and the + row's series has started at or before it. Leading all-unobserved patches are + left-padding and are never valid keys. + """ + if not attend_missing: + return obs_per_patch + started = obs_per_patch.cummax(dim=-1).values + return obs_per_patch | started + + +def enc_causal_mask(key_valid: Tensor) -> Tensor: + """`(R, L, L)` causal encoder mask over valid keys, with the diagonal always open.""" + _, ll = key_valid.shape + device = key_valid.device + causal = torch.tril(torch.ones(ll, ll, dtype=torch.bool, device=device)) + eye = torch.eye(ll, dtype=torch.bool, device=device) + return (causal[None] & key_valid[:, None, :]) | eye[None] + + +def build_forecast_time_mask( + context_len: int, horizon: int, key_valid: Tensor, device: torch.device +) -> Tensor: + """`(R, 1, L+h, L+h)` time mask for one block without covariates. + + The encoder is causal; every decoder token reads the whole valid context and its + own causal decoder chain; the context never reads the horizon. + """ + rows = key_valid.shape[0] + enc = enc_causal_mask(key_valid) + enc_dec = torch.zeros(rows, context_len, horizon, dtype=torch.bool, device=device) + dec_enc = key_valid[:, None, :].expand(rows, horizon, context_len) + dec_dec = torch.tril(torch.ones(horizon, horizon, dtype=torch.bool, device=device))[ + None + ].expand(rows, horizon, horizon) + top = torch.cat([enc, enc_dec], dim=2) + bottom = torch.cat([dec_enc, dec_dec], dim=2) + return torch.cat([top, bottom], dim=1).unsqueeze(1) + + +def build_forecast_positions( + context_len: int, horizon: int, first_obs: Tensor +) -> Tensor: + """`(R, L+h)` time positions `0..L+h-1`, anchored at each row's first observed patch.""" + p = torch.arange(context_len + horizon, device=first_obs.device) + return p[None, :] - first_obs[:, None] + + +def compute_block_layout( + min_context: Tensor, n_patches: int, block_size: int +) -> tuple[Tensor, Tensor]: + """Per decoder position: the anchor of its block and whether it is active (`p >= s`).""" + device = min_context.device + p = torch.arange(n_patches, device=device)[None, :] + s = min_context[:, None] + active = p >= s + offset = (p - s).clamp(min=0) + anchor = s + block_size * (offset // block_size) + anchor = torch.where(active, anchor, torch.zeros_like(anchor)) + return anchor, active + + +def build_two_stream_time_mask( + min_context: Tensor, + n_patches: int, + block_size: int, + key_valid: Tensor, + bidirectional_rows: Tensor | None, +) -> Tensor: + """`(R, 1, 2N, 2N)` time mask over a full `[Enc ; Dec]` window (the covariate path). + + - Enc to Enc: causal over valid keys; full (non-causal) for `bidirectional_rows`, + the known-future covariates. + - Enc to Dec: never. + - Dec to Enc: a decoder token in the block anchored at `b` reads valid `Enc[j]`, `j < b`. + - Dec to Dec: causal within the same block; every decoder token reads itself. + """ + anchor, active = compute_block_layout(min_context, n_patches, block_size) + rows, n = active.shape + device = active.device + causal = torch.tril(torch.ones(n, n, dtype=torch.bool, device=device)) + key_col = key_valid[:, None, :] + j = torch.arange(n, device=device).view(1, 1, n) + eye = torch.eye(n, dtype=torch.bool, device=device) + enc_enc = enc_causal_mask(key_valid) + if bidirectional_rows is not None: + enc_bi = key_valid[:, None, :].expand(rows, n, n) | eye[None] + enc_enc = torch.where(bidirectional_rows[:, None, None], enc_bi, enc_enc) + enc_dec = torch.zeros(rows, n, n, dtype=torch.bool, device=device) + dec_enc = active[:, :, None] & (j < anchor[:, :, None]) & key_col + same_block = anchor[:, :, None] == anchor[:, None, :] + dec_dec = same_block & causal[None] & active[:, :, None] & active[:, None, :] + dec_dec = dec_dec | eye[None] + top = torch.cat([enc_enc, enc_dec], dim=2) + bottom = torch.cat([dec_enc, dec_dec], dim=2) + return torch.cat([top, bottom], dim=1).unsqueeze(1) + + +def variate_firewall( + r: int, + device: torch.device, + target_rows: Tensor | None, + bidirectional_rows: Tensor | None, +) -> Tensor | None: + """`(R, R)` forbidden query-to-key row pairs, or None. + + - A covariate row may not read a target row (covariates are exogenous). + - A known-future (bidirectional) row may not read a causal row. Without this, a + covariate's view of the future could relay a target's own future back into it. + """ + if target_rows is None and bidirectional_rows is None: + return None + forbid = torch.zeros(r, r, dtype=torch.bool, device=device) + if target_rows is not None: + forbid = forbid | ((~target_rows[:, None]) & target_rows[None, :]) + if bidirectional_rows is not None: + forbid = forbid | (bidirectional_rows[:, None] & (~bidirectional_rows)[None, :]) + return forbid + + +def build_two_stream_variate_mask( + group_ids: Tensor, + enc_key_valid: Tensor, + n_dec: int, + target_rows: Tensor | None = None, + bidirectional_rows: Tensor | None = None, +) -> Tensor | None: + """`(N_enc + N_dec, 1, R, R)` variate mask, or None for a single row. + + A row reads another row of the same series; on encoder slots the key row must also + be a valid key. The firewalls apply on every slot; the diagonal is always open. + """ + r, _ = enc_key_valid.shape + if r == 1: + return None + device = enc_key_valid.device + row_eq = group_ids[:, None] == group_ids[None, :] + dec_valid = torch.ones(r, n_dec, dtype=torch.bool, device=device) + slot_valid = torch.cat([enc_key_valid, dec_valid], dim=1) + key_valid = slot_valid.permute(1, 0) + mask = row_eq[None, None] & key_valid[:, None, None, :] + forbid = variate_firewall(r, device, target_rows, bidirectional_rows) + if forbid is not None: + mask = mask & ~forbid[None, None] + eye = torch.eye(r, dtype=torch.bool, device=device) + return mask | eye[None, None] + + +def build_two_stream_positions(n_patches: int, first_obs: Tensor) -> Tensor: + """`(R, 2N)` positions: encoder and decoder token of the same patch share a time.""" + p = torch.arange(n_patches, device=first_obs.device) + row = p[None, :] - first_obs[:, None] + return torch.cat([row, row], dim=-1) + + +def group_first_observed_patch( + obs_per_patch: Tensor, inverse: Tensor, num_groups: int +) -> Tensor: + """`(num_groups,)` index of the first patch any row of the group observes (0 if none).""" + n = obs_per_patch.shape[-1] + idx = inverse[:, None].expand(-1, n) + group_obs = torch.zeros( + num_groups, n, dtype=torch.float32, device=obs_per_patch.device + ) + group_obs.scatter_reduce_( + 0, idx, obs_per_patch.float(), reduce="amax", include_self=True + ) + return group_obs.argmax(dim=-1) diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/timestral/model.py b/iotdb-core/ainode/iotdb/ainode/core/model/timestral/model.py new file mode 100644 index 0000000000000..44ce1b60e9770 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/timestral/model.py @@ -0,0 +1,392 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# Copyright 2026 Mistral AI +# SPDX-License-Identifier: Apache-2.0 + +"""The Timestral network and its single-block forecast primitive. + +The model normalizes each row causally per patch, embeds patches from `[values, +missing]` channels, runs a time/variate transformer over `[Enc ; Dec]` tokens (context +patches followed by learned forecast-query tokens), and reads quantiles for every +horizon time step from the decoder tokens. +""" + +import torch +from torch import Tensor, nn + +from iotdb.ainode.core.model.timestral.config import TimestralConfig +from iotdb.ainode.core.model.timestral.layers import ( + ResidualBlock, + Transformer, + TwoStreamMask, +) +from iotdb.ainode.core.model.timestral.masks import ( + build_forecast_positions, + build_forecast_time_mask, + build_two_stream_positions, + build_two_stream_time_mask, + build_two_stream_variate_mask, + compute_enc_key_validity, + group_first_observed_patch, +) +from iotdb.ainode.core.model.timestral.scaler import ( + apply_affine_asinh_masked, + causal_patch_loc_scale, + denormalize_quantiles, + instance_norm_stats, +) + +ROLE_TARGET = 0 +ROLE_PAST_COV = 1 +ROLE_KNOWN_COV = 2 + + +class PatchRepresentation(nn.Module): + """Patch embedding over per-step `[values, missing]` channels, plus the forecast-query token.""" + + def __init__(self, patch_size: int, d_model: int) -> None: + super().__init__() + self.patch_size = patch_size + self.patch_proj = ResidualBlock( + 2 * patch_size, 4 * d_model, d_model, use_bias=True + ) + self.forecast_query = nn.Parameter(torch.zeros(d_model)) + + def forward(self, scaled: Tensor, observed: Tensor) -> Tensor: + """`(R, T)` normalized values and observation mask to `(R, T / patch_size, d_model)`.""" + r, t = scaled.shape + n = t // self.patch_size + values = scaled.reshape(r, n, self.patch_size) + missing = (~observed).reshape(r, n, self.patch_size).to(scaled.dtype) + x = torch.cat([values, missing], dim=-1).to( + self.patch_proj.hidden_layer.weight.dtype + ) + return self.patch_proj(x) + + def forecast_tokens(self, rows: int, n_patches: int) -> Tensor: + return self.forecast_query.view(1, 1, -1).expand(rows, n_patches, -1) + + +def monotone_quantiles( + raw: Tensor, levels: tuple[float, ...], width_factor: float +) -> Tensor: + """Non-crossing quantiles from raw head outputs along the last axis (`monotone_cumsum`). + + The median channel is free. The others fan out from it as + `median -/+ cumsum(width_factor * dp * softplus(raw))`, `dp` being the probability gap + to the next level inward, so unequal knot spacing is respected. + """ + m = levels.index(0.5) + widths = [0.0] * len(levels) + for i in range(m + 1, len(levels)): + widths[i] = levels[i] - levels[i - 1] + for i in range(m): + widths[i] = levels[i + 1] - levels[i] + widths_t = torch.tensor(widths, device=raw.device, dtype=raw.dtype) + inc = width_factor * widths_t * nn.functional.softplus(raw) + median = raw[..., m : m + 1] + up = torch.cumsum(inc[..., m + 1 :], dim=-1) + down = torch.flip( + torch.cumsum(torch.flip(inc[..., :m], dims=[-1]), dim=-1), dims=[-1] + ) + big = torch.finfo(raw.dtype).max + return torch.cat([median - down, median, median + up], dim=-1).clamp( + min=-big, max=big + ) + + +class QuantileHead(nn.Module): + """Maps each decoder token to `patch_size` monotone quantile vectors in normalized space.""" + + def __init__(self, config: TimestralConfig) -> None: + super().__init__() + self.config = config + self.patch_size = config.patch_size + self.num_quantiles = config.num_quantiles + self.projection = ResidualBlock( + config.d_model, + 4 * config.d_model, + config.patch_size * config.num_quantiles, + use_bias=False, + ) + + def forward(self, x: Tensor) -> Tensor: + """`(R, N, D)` to normalized quantiles `(R, N * patch_size, Q)` in at least fp32.""" + raw = self.projection(x) + *leading, n_patches, _ = raw.shape + preds = raw.reshape(*leading, n_patches, self.patch_size, self.num_quantiles) + preds = preds.to(torch.promote_types(preds.dtype, torch.float32)) + preds = monotone_quantiles( + preds, self.config.quantile_knots, self.config.quantile_width_factor + ) + return preds.reshape(*leading, n_patches * self.patch_size, self.num_quantiles) + + +class TimestralModel(nn.Module): + def __init__(self, config: TimestralConfig) -> None: + super().__init__() + self.config = config + c = config + self.representation = PatchRepresentation(c.patch_size, c.d_model) + self.transformer = Transformer( + block_schedule=c.block_schedule, + embed_dim=c.d_model, + num_heads=c.num_heads, + mlp_hidden_dim=c.mlp_hidden_dim, + rope_theta=c.rope_theta, + rope_dim=c.rope_dim, + norm_eps=c.norm_eps, + attn_bias=c.attn_bias, + mlp_bias=c.mlp_bias, + time_attn_gate=c.time_attn_gate, + variate_attn_gate=c.variate_attn_gate, + ) + self.head = QuantileHead(c) + + def scale(self, inputs: Tensor, observed: Tensor) -> tuple[Tensor, Tensor, Tensor]: + """Causal per-patch normalization of packed `(R, T)` rows. + + Returns `(scaled, loc, scale)`, each `(R, T)`. + """ + c = self.config + loc, scale = causal_patch_loc_scale( + inputs.unsqueeze(1), + observed.unsqueeze(1), + c.patch_size, + c.minimum_scale, + c.stabilize_with_global, + ) + loc, scale = loc.squeeze(1), scale.squeeze(1) + scaled = apply_affine_asinh_masked( + inputs, loc, scale, observed, use_asinh=c.use_asinh + ) + return scaled, loc, scale + + def _known_cov_scaling( + self, + inputs: Tensor, + observed: Tensor, + scaled: Tensor, + loc: Tensor, + scale: Tensor, + known_rows: Tensor, + ) -> tuple[Tensor, Tensor, Tensor]: + """Normalize known-future covariate rows with non-causal whole-row statistics. + + Indicator rows (values in {0, 1}) skip the asinh squash. + """ + if not bool(known_rows.any()): + return scaled, loc, scale + loc_c, scale_c, is_binary = instance_norm_stats( + inputs, observed, minimum_scale=self.config.minimum_scale + ) + known_scaled = apply_affine_asinh_masked( + inputs, loc_c, scale_c, observed, use_asinh=False + ) + if self.config.use_asinh: + work = torch.promote_types(known_scaled.dtype, torch.float32) + squashed = known_scaled.to(work).asinh().to(known_scaled.dtype) + known_scaled = torch.where(observed & ~is_binary, squashed, known_scaled) + sel = known_rows[:, None] + return ( + torch.where(sel, known_scaled, scaled), + torch.where(sel, loc_c, loc), + torch.where(sel, scale_c, scale), + ) + + def _run_two_stream( + self, + x_enc: Tensor, + x_dec: Tensor, + time_mask: Tensor, + variate_mask: Tensor | None, + positions: Tensor, + ) -> Tensor: + """Transformer over `[Enc ; Dec]`, then the head on the decoder tokens (normalized space).""" + x = torch.cat([x_enc, x_dec], dim=1) + mask: Tensor | TwoStreamMask = time_mask + if self.config.split_context_horizon_attn: + mask = TwoStreamMask(mask=time_mask, enc_len=x_enc.shape[1]) + h = self.transformer(x, mask, variate_mask, positions) + return self.head(h[:, x_enc.shape[1] :, :]) + + def predict_block( + self, + inputs: Tensor, + observed: Tensor, + group_ids: Tensor, + horizon_patches: int, + role_ids: Tensor | None = None, + known_future: Tensor | None = None, + known_future_mask: Tensor | None = None, + ) -> Tensor: + """Forecast one block of `horizon_patches` patches after a packed context. + + Args: + inputs: `(R, T)` context values, `T` a multiple of the patch size. + observed: `(R, T)` observation mask. + group_ids: `(R,)` series id per row; rows of one series attend each other. + horizon_patches: Number of patches to forecast. + role_ids: `(R,)` row roles (`ROLE_TARGET`, `ROLE_PAST_COV`, `ROLE_KNOWN_COV`), + or None when every row is a target. + known_future: `(R, horizon_patches * patch_size)` known-future covariate + values over this block; zeros on other rows. + known_future_mask: Observation mask for `known_future`. + + Returns: + Raw-space quantiles `(R, horizon_patches * patch_size, Q)`, non-decreasing along Q. + """ + assert ( + inputs.dim() == 2 + ), f"expected packed (R, T) inputs, got {tuple(inputs.shape)}" + if ( + role_ids is None + or not self.config.covariate_enabled + or bool((role_ids == ROLE_TARGET).all()) + ): + return self._predict_context_only( + inputs, observed, group_ids, horizon_patches + ) + r = inputs.shape[0] + horizon_len = horizon_patches * self.config.patch_size + if known_future is None: + known_future = torch.zeros( + r, horizon_len, device=inputs.device, dtype=inputs.dtype + ) + known_future_mask = torch.zeros( + r, horizon_len, device=inputs.device, dtype=torch.bool + ) + assert ( + known_future_mask is not None + ), "known_future_mask is required alongside known_future" + assert known_future.shape == ( + r, + horizon_len, + ), f"known_future must be ({r}, {horizon_len}), got {tuple(known_future.shape)}" + full_inputs = torch.cat([inputs, known_future], dim=-1) + full_observed = torch.cat([observed, known_future_mask], dim=-1) + return self._predict_window( + full_inputs, full_observed, group_ids, horizon_patches, role_ids + ) + + def _predict_context_only( + self, inputs: Tensor, observed: Tensor, group_ids: Tensor, horizon_patches: int + ) -> Tensor: + """No covariates: encoder over the context, `h` forecast-query tokens after it.""" + c = self.config + ps = c.patch_size + rows, t = inputs.shape + assert t % ps == 0, f"context length {t} must be a multiple of patch_size {ps}" + assert group_ids.shape == ( + rows, + ), f"group_ids must be ({rows},), got {tuple(group_ids.shape)}" + ll, h = t // ps, horizon_patches + assert h >= 1, f"horizon_patches must be >= 1, got {h}" + + scaled, loc, scale = self.scale(inputs, observed) + x_enc = self.representation(scaled, observed) + x_dec = self.representation.forecast_tokens(rows, h) + + obs_per_patch = observed.unflatten(-1, (ll, ps)).any(dim=-1) + key_valid = compute_enc_key_validity(obs_per_patch, c.attend_missing_patches) + time_mask = build_forecast_time_mask(ll, h, key_valid, inputs.device) + variate_mask = build_two_stream_variate_mask(group_ids, key_valid, h) + unique_ids, inverse = torch.unique(group_ids, return_inverse=True) + first_obs = group_first_observed_patch( + obs_per_patch, inverse, unique_ids.numel() + )[inverse] + positions = build_forecast_positions(ll, h, first_obs) + + preds = self._run_two_stream(x_enc, x_dec, time_mask, variate_mask, positions) + anchor_loc = loc.unflatten(-1, (ll, ps))[..., 0][:, ll - 1 : ll] + anchor_scale = scale.unflatten(-1, (ll, ps))[..., 0][:, ll - 1 : ll] + return denormalize_quantiles( + preds, anchor_loc.unsqueeze(-1), anchor_scale.unsqueeze(-1), c.use_asinh + ) + + def _predict_window( + self, + full_inputs: Tensor, + full_observed: Tensor, + group_ids: Tensor, + horizon_patches: int, + role_ids: Tensor, + ) -> Tensor: + """With covariates: encoder over context plus horizon, so known-future values are keys. + + Target and past-only rows keep a blank horizon. The time and variate masks keep + target information out of covariate rows and keep any row's future out of the + targets. + """ + c = self.config + ps = c.patch_size + rows, w = full_inputs.shape + assert w % ps == 0, f"window length {w} must be a multiple of patch_size {ps}" + assert group_ids.shape == (rows,) and role_ids.shape == ( + rows, + ), "group_ids and role_ids must be (R,)" + h = horizon_patches + assert ( + 1 <= h <= c.block_size + ), f"decode block {h} must be in [1, {c.block_size}] with covariates" + n = w // ps + ll = n - h + assert ( + ll >= 1 + ), f"context must have at least one patch (window {n}, horizon {h})" + target_rows = role_ids == ROLE_TARGET + known_rows = role_ids == ROLE_KNOWN_COV + + scaled, loc, scale = self.scale(full_inputs, full_observed) + scaled, loc, scale = self._known_cov_scaling( + full_inputs, full_observed, scaled, loc, scale, known_rows + ) + x_enc = self.representation(scaled, full_observed) + x_dec = self.representation.forecast_tokens(rows, n) + + obs_per_patch = full_observed.unflatten(-1, (n, ps)).any(dim=-1) + key_valid = compute_enc_key_validity(obs_per_patch, c.attend_missing_patches) + s_rows = torch.full((rows,), ll, device=full_inputs.device, dtype=torch.long) + time_mask = build_two_stream_time_mask( + s_rows, n, c.block_size, key_valid, bidirectional_rows=known_rows + ) + variate_mask = build_two_stream_variate_mask( + group_ids, + key_valid, + n, + target_rows=target_rows, + bidirectional_rows=known_rows, + ) + unique_ids, inverse = torch.unique(group_ids, return_inverse=True) + first_obs = group_first_observed_patch( + obs_per_patch, inverse, unique_ids.numel() + )[inverse] + positions = build_two_stream_positions(n, first_obs) + + preds = self._run_two_stream( + x_enc, x_dec, time_mask, variate_mask, positions + ).unflatten(1, (n, ps)) + horizon_preds = preds[:, ll:n].reshape(rows, h * ps, preds.shape[-1]) + anchor_loc = loc.unflatten(-1, (n, ps))[..., 0][:, ll - 1 : ll] + anchor_scale = scale.unflatten(-1, (n, ps))[..., 0][:, ll - 1 : ll] + return denormalize_quantiles( + horizon_preds, + anchor_loc.unsqueeze(-1), + anchor_scale.unsqueeze(-1), + c.use_asinh, + ) diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/timestral/modeling_timestral.py b/iotdb-core/ainode/iotdb/ainode/core/model/timestral/modeling_timestral.py new file mode 100644 index 0000000000000..0737d92e9f51d --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/timestral/modeling_timestral.py @@ -0,0 +1,52 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +import torch +from huggingface_hub import PyTorchModelHubMixin +from safetensors.torch import load_file + +from iotdb.ainode.core.model.timestral.config import TimestralConfig +from iotdb.ainode.core.model.timestral.model import TimestralModel + + +class TimestralForPrediction(TimestralModel, PyTorchModelHubMixin): + """ + Timestral as an AINode model, loaded from a directory with the ``config.json`` and + ``model.safetensors`` of a Hugging Face checkpoint such as ``mistralai/Timestral-209M``. + + The keys of ``config.json`` are the fields of :class:`TimestralConfig`. + """ + + def __init__(self, **fields): + super().__init__(TimestralConfig.from_dict(fields)) + + def _apply(self, fn, *args, **kwargs): + module = super()._apply(fn, *args, **kwargs) + # Recompute the rotary frequencies on the new device rather than moving them, + # as Timestral's own loader does + self.transformer.rope.rebuild_inv_freq() + return module + + @classmethod + def _load_as_safetensor( + cls, model: "TimestralForPrediction", model_file: str, map_location: str, strict + ) -> "TimestralForPrediction": + state_dict = load_file(model_file, device=str(torch.device(map_location))) + model.load_state_dict(state_dict, strict=True) + model.transformer.rope.rebuild_inv_freq() + return model.eval() diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/timestral/pipeline_timestral.py b/iotdb-core/ainode/iotdb/ainode/core/model/timestral/pipeline_timestral.py new file mode 100644 index 0000000000000..bd47e4be4e14c --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/timestral/pipeline_timestral.py @@ -0,0 +1,157 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +import torch + +from iotdb.ainode.core.inference.pipeline.basic_pipeline import ForecastPipeline +from iotdb.ainode.core.model.model_info import ModelInfo +from iotdb.ainode.core.model.timestral.forecaster import TimestralForecaster + + +class TimestralPipeline(ForecastPipeline): + """ + Forecasting pipeline of Timestral. + + Each input is one series of one or more targets, forecast jointly, with optional past + covariates and future-known covariates. The model predicts 21 quantiles (0.01 to 0.99); + horizons longer than one decode block (1024 steps) are forecast block by block, feeding + each block's median back as context. The pipeline returns the median. + """ + + # Timestral forecasts a batch as one set of rows, one row per target, and its + # variate-attention mask grows with the square of the number of rows. Series of the same + # shape are batched together only up to this many rows, so that the memory of a batch of + # requests grows linearly with the number of requests. + max_batch_rows = 64 + + def __init__(self, model_info: ModelInfo, **model_kwargs): + super().__init__(model_info, **model_kwargs) + self.forecaster = TimestralForecaster(self.model) + + def _preprocess(self, inputs, **infer_kwargs) -> list[dict]: + """ + Convert the validated inputs to the arguments of the forecaster. + + Parameters + ---------- + inputs : list of dict + Each dictionary contains: + - 'targets': A tensor of shape (input_length,) or (target_count, input_length). + - 'past_covariates' (optional): A dictionary of tensors of shape (input_length,). + - 'future_covariates' (optional): A dictionary of tensors of shape (output_length,), + whose keys are a subset of 'past_covariates'. + + Returns + ------- + list of dict + One dictionary per input, with: + - 'context': The targets, a (1, target_count, input_length) tensor. + - 'past_covariates': The covariates known only in the past, a + (1, count, input_length) tensor, or None. + - 'known_covariates': The covariates also known in the future, a + (1, count, input_length + output_length) tensor of their past and future + values, or None. + """ + series = [] + for item in inputs: + targets = item["targets"] + if targets.ndim == 1: + targets = targets.unsqueeze(0) + past_covariates = item.get("past_covariates") or {} + future_covariates = item.get("future_covariates") or {} + past_only = [k for k in past_covariates if k not in future_covariates] + known = [k for k in past_covariates if k in future_covariates] + series.append( + { + "context": targets.float().unsqueeze(0), + "past_covariates": ( + torch.stack([past_covariates[k] for k in past_only]) + .float() + .unsqueeze(0) + if past_only + else None + ), + "known_covariates": ( + torch.stack( + [ + torch.cat([past_covariates[k], future_covariates[k]]) + for k in known + ] + ) + .float() + .unsqueeze(0) + if known + else None + ), + } + ) + return series + + def forecast(self, inputs: list[dict], **infer_kwargs) -> list[torch.Tensor]: + """ + Forecast the series, batching those of the same shape without covariates, up to + `max_batch_rows` targets per batch. + + Returns + ------- + list of torch.Tensor + One tensor per series, of shape (target_count, output_length, quantile_count). + """ + output_length = infer_kwargs.get("output_length", 96) + outputs: list[torch.Tensor | None] = [None] * len(inputs) + batches: list[list[int]] = [] + open_batches: dict[tuple, list[int]] = {} + for i, item in enumerate(inputs): + if ( + item["past_covariates"] is not None + or item["known_covariates"] is not None + ): + batches.append([i]) + continue + shape = tuple(item["context"].shape) + batch = open_batches.get(shape) + if batch is None or (len(batch) + 1) * shape[1] > self.max_batch_rows: + batch = [] + batches.append(batch) + open_batches[shape] = batch + batch.append(i) + for indices in batches: + items = [inputs[i] for i in indices] + forecast = self.forecaster.forecast( + torch.cat([item["context"] for item in items]), + output_length, + past_covariates=items[0]["past_covariates"], + known_covariates=items[0]["known_covariates"], + ) + for i, quantiles in zip(indices, forecast.quantiles): + outputs[i] = quantiles + return outputs + + def _postprocess( + self, outputs: list[torch.Tensor], **infer_kwargs + ) -> list[torch.Tensor]: + """ + Select the median of each forecast. + + Returns + ------- + list of torch.Tensor + One tensor per series, of shape (target_count, output_length). + """ + median = self.forecaster.median_index + return [output[..., median].float().cpu().clone() for output in outputs] diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/timestral/scaler.py b/iotdb-core/ainode/iotdb/ainode/core/model/timestral/scaler.py new file mode 100644 index 0000000000000..ac91a0fc3c9df --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/timestral/scaler.py @@ -0,0 +1,286 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# Copyright 2026 Mistral AI +# SPDX-License-Identifier: Apache-2.0 + +"""Normalization for Timestral: causal per-patch mean/std scaling with an asinh squash. + +Each patch is normalized with the mean and std of every observed value up to the end of +that patch, so the context never sees future statistics. Forecasts come back in the +normalized (asinh) space of the last context patch and are inverted with `sinh`. +""" + +import functools + +import torch +from torch import Tensor + + +def compute_causal_statistics( + data: Tensor, + observation_mask: Tensor, + minimum_scale: float, + stabilize_with_global: bool, + scale_factor_exponent: float = 10.0, +) -> tuple[Tensor, Tensor]: + """Causal running mean and std along the last axis (Welford form, in float64). + + Returns `(means, scales)`, both shaped like `data`; position `t` uses the observed + values at positions `0..t` only. Unobserved positions get zero weight. The variance + uses Bessel's correction and the std is clamped to `minimum_scale` after the sqrt. + """ + weights = torch.ones_like(data) * observation_mask + with torch.no_grad(): + x = data.to(torch.float64) + w = weights.to(torch.float64) + prev_deterministic = torch.are_deterministic_algorithms_enabled() + if prev_deterministic and data.device.type == "cuda": + torch.use_deterministic_algorithms(False) + try: + weighted = w * x + cum_w = torch.cumsum(w, dim=-1) + cum_x = torch.cumsum(weighted, dim=-1) + denominator = cum_w.clamp_min(1.0) + means = cum_x / denominator + shifted = torch.zeros_like(means) + shifted[..., 1:] = means[..., :-1] + increment = (x - shifted) * (x - means) * w + m2 = torch.cumsum(increment, dim=-1) + variance = m2 / torch.clamp(denominator - 1.0, min=1.0) + scales = torch.sqrt(variance).clamp(min=minimum_scale) + if stabilize_with_global: + lo, hi = 10.0 ** (-scale_factor_exponent), 10.0**scale_factor_exponent + global_den = ( + (weights * observation_mask).sum(-1, keepdim=True).clamp_min(1.0) + ) + global_means = torch.nan_to_num( + weighted.sum(-1, keepdim=True) / global_den + ) + global_var = ( + ((x - global_means) * weights * observation_mask) ** 2 + ).sum(-1, keepdim=True) / global_den + global_scale = torch.sqrt(global_var + minimum_scale).expand_as(scales) + floor = torch.max( + torch.tensor( + minimum_scale, device=scales.device, dtype=scales.dtype + ), + global_scale * lo, + ) + scales = torch.clamp(scales, min=floor, max=global_scale * hi) + means = means.to(data.dtype) + scales = scales.to(data.dtype) + finally: + if prev_deterministic and data.device.type == "cuda": + torch.use_deterministic_algorithms(True) + return means, scales + + +def causal_patch_loc_scale( + data: Tensor, + observation_mask: Tensor, + patch_size: int, + minimum_scale: float, + stabilize_with_global: bool, +) -> tuple[Tensor, Tensor]: + """Per-patch `(loc, scale)` broadcast back to every time step of the patch. + + Each patch uses the causal statistics at its last time step, i.e. everything + observed up to the end of that patch. + """ + assert ( + data.shape[-1] % patch_size == 0 + ), f"length {data.shape[-1]} is not a multiple of patch_size {patch_size}" + with torch.no_grad(): + means, scales = compute_causal_statistics( + data, observation_mask, minimum_scale, stabilize_with_global + ) + patch_means = means.unfold(-1, patch_size, patch_size)[..., -1] + patch_scales = scales.unfold(-1, patch_size, patch_size)[..., -1] + return patch_means.repeat_interleave( + patch_size, dim=-1 + ), patch_scales.repeat_interleave(patch_size, dim=-1) + + +def _fp32_io(*dtypes: torch.dtype) -> tuple[torch.dtype, torch.dtype]: + """`(out_dtype, work_dtype)`: compute in at least fp32, return the promoted input dtype.""" + out_dtype = functools.reduce(torch.promote_types, dtypes) + return out_dtype, torch.promote_types(out_dtype, torch.float32) + + +def apply_affine_asinh_masked( + data: Tensor, + loc: Tensor, + scale: Tensor, + observation_mask: Tensor, + use_asinh: bool, +) -> Tensor: + """`asinh((data - loc) / scale)` on observed positions; unobserved positions become 0.""" + out_dtype, work = _fp32_io(data.dtype, loc.dtype, scale.dtype) + scaled = (data.to(work) - loc.to(work)) / scale.to(work).clamp_min( + torch.finfo(work).eps + ) + scaled = scaled.masked_fill(~observation_mask, 0.0) + if use_asinh: + scaled = torch.where(observation_mask, scaled.asinh(), scaled) + return scaled.to(out_dtype) + + +def invert_affine_asinh( + x: Tensor, loc: Tensor, scale: Tensor, use_asinh: bool +) -> Tensor: + """Inverse of the forward transform: `sinh(x) * scale + loc`.""" + out_dtype, work = _fp32_io(x.dtype, loc.dtype, scale.dtype) + x = x.to(work) + if use_asinh: + x = x.sinh() + return (x * scale.to(work) + loc.to(work)).to(out_dtype) + + +def clamp_nonfinite(vals: Tensor) -> Tensor: + """Map +/-inf (a sinh overflow) to the nearest finite value along the last axis. + + A position whose whole last axis is non-finite becomes 0. NaN is left untouched so + that real instability still surfaces downstream. + """ + clamped = torch.where( + vals == float("inf"), + torch.where(vals.isfinite(), vals, -float("inf")).amax(dim=-1, keepdim=True), + torch.where( + vals == -float("inf"), + torch.where(vals.isfinite(), vals, float("inf")).amin(dim=-1, keepdim=True), + vals, + ), + ) + return torch.where(clamped.isinf(), torch.zeros_like(clamped), clamped) + + +def denormalize_quantiles( + preds: Tensor, loc: Tensor, scale: Tensor, use_asinh: bool +) -> Tensor: + """Invert the normalization and clamp infinities; monotone quantiles stay in order.""" + return clamp_nonfinite(invert_affine_asinh(preds, loc, scale, use_asinh)) + + +def welford_combine( + w_a: Tensor, mean_a: Tensor, m2_a: Tensor, w_b: Tensor, mean_b: Tensor, m2_b: Tensor +) -> tuple[Tensor, Tensor, Tensor]: + """Merge two Welford states `(weight, mean, m2)` element-wise; `(0, 0, 0)` is the identity.""" + w = w_a + w_b + w_safe = w.clamp(min=torch.finfo(w.dtype).tiny) + delta = mean_b - mean_a + mean = mean_a + delta * w_b / w_safe + m2 = m2_a + m2_b + delta.square() * w_a * w_b / w_safe + is_zero = w == 0 + mean = torch.where(is_zero, torch.zeros_like(mean), mean) + m2 = torch.where(is_zero, torch.zeros_like(m2), m2) + return w, mean, m2 + + +def cumulative_welford( + w: Tensor, mean: Tensor, m2: Tensor +) -> tuple[Tensor, Tensor, Tensor]: + """Inclusive prefix scan of Welford states along dim 1 (Hillis-Steele, `log2(N)` merges).""" + n = w.shape[1] + d = 1 + while d < n: + pad = (d, 0) + w_l = torch.nn.functional.pad(w[:, : n - d], pad) + mean_l = torch.nn.functional.pad(mean[:, : n - d], pad) + m2_l = torch.nn.functional.pad(m2[:, : n - d], pad) + w, mean, m2 = welford_combine(w_l, mean_l, m2_l, w, mean, m2) + d *= 2 + return w, mean, m2 + + +def advancing_block_frames( + context: Tensor, observation_mask: Tensor, sinh_median: Tensor, minimum_scale: float +) -> tuple[Tensor, Tensor]: + """Per-patch `(loc, scale)` of a decoded block in the advancing refine frame, each `(R, H)` float64. + + Horizon patch `j` is normalized with the statistics of the observed context plus the + block's medians of patches `0..j-1`, each median denormalized in its own frame before + it is folded in. With perfect medians this is causal per-patch normalization of the true + series. The statistics are float64 Welford states, accumulated exactly as in training. + + Args: + context: `(R, T)` block context, `T` a multiple of the patch size. + observation_mask: `(R, T)`. + sinh_median: `(R, H, patch_size)` the block's median in the last context patch's + frame with the asinh undone, `(median - loc) / scale`. + minimum_scale: Scale floor. + """ + r, n_pred, ps = sinh_median.shape + blank = sinh_median.new_zeros(r, n_pred, ps) + # Per-patch moments of the context, scanned together with the zero-weight horizon patches. + values = torch.cat( + [context.unflatten(-1, (-1, ps)).to(torch.float64), blank], dim=1 + ) + weights = torch.cat( + [observation_mask.unflatten(-1, (-1, ps)).to(torch.float64), blank], dim=1 + ) + cnt = weights.sum(-1) + mean = (values * weights).sum(-1) / cnt.clamp(min=1.0) + m2 = (weights * (values - mean.unsqueeze(-1)).square()).sum(-1) + is_zero = cnt == 0 + mean = torch.where(is_zero, torch.zeros_like(mean), mean) + m2 = torch.where(is_zero, torch.zeros_like(m2), m2) + cnt, mean, m2 = (s[:, -n_pred - 1] for s in cumulative_welford(cnt, mean, m2)) + + offset = mean # medians are clamped relative to the context mean, so a large level survives + clip = torch.finfo(torch.float32).max ** 0.5 / 1e3 + locs, scales = [], [] + for h in range(n_pred): + var = (m2 / (cnt - 1.0).clamp(min=1.0)).clamp(min=0.0) + scale = var.sqrt().clamp(min=minimum_scale) + locs.append(mean) + scales.append(scale) + folded = ( + sinh_median[:, h] * scale[:, None] + mean[:, None] - offset[:, None] + ).clamp(-clip, clip) + folded = folded + offset[:, None] + b_mean = folded.mean(-1) + b_m2 = (folded - b_mean[:, None]).square().sum(-1) + cnt, mean, m2 = welford_combine( + cnt, mean, m2, torch.full_like(cnt, ps), b_mean, b_m2 + ) + return torch.stack(locs, dim=1), torch.stack(scales, dim=1) + + +def instance_norm_stats( + values: Tensor, observation_mask: Tensor, minimum_scale: float +) -> tuple[Tensor, Tensor, Tensor]: + """Non-causal per-row mean/std over observed positions, for known-future covariates. + + Rows whose observed values are all 0 or 1 (indicator covariates) keep `loc=0`, + `scale=1`. Returns `(loc, scale, is_binary)`, each `(R, 1)`. + """ + orig_dtype = values.dtype + stat_dtype = torch.promote_types(orig_dtype, torch.float32) + x = values.to(stat_dtype) + mask_f = observation_mask.to(stat_dtype) + n = mask_f.sum(dim=-1, keepdim=True).clamp(min=1.0) + loc = (x * mask_f).sum(dim=-1, keepdim=True) / n + var = (((x - loc) * mask_f) ** 2).sum(dim=-1, keepdim=True) / n + scale = var.sqrt().clamp_min(minimum_scale) + is_zero_one = ((values == 0.0) | (values == 1.0) | ~observation_mask).all( + dim=-1, keepdim=True + ) + is_binary = is_zero_one & observation_mask.any(dim=-1, keepdim=True) + loc = torch.where(is_binary, torch.zeros_like(loc), loc) + scale = torch.where(is_binary, torch.ones_like(scale), scale) + return loc.to(orig_dtype), scale.to(orig_dtype), is_binary diff --git a/iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py b/iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py index e3b5b49e05760..b402e2100bf82 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py +++ b/iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py @@ -22,6 +22,7 @@ from iotdb.ainode.core.manager.device_manager import DeviceManager from iotdb.ainode.core.manager.inference_manager import InferenceManager from iotdb.ainode.core.manager.model_manager import ModelManager +from iotdb.ainode.core.model.model_constants import ModelStates from iotdb.ainode.core.rpc.status import get_status from iotdb.thrift.ainode import IAINodeRPCService from iotdb.thrift.ainode.ttypes import ( @@ -139,16 +140,29 @@ def _ensure_model_is_registered(self, model_id: str) -> TSStatus: ) return TSStatus(code=TSStatusCode.SUCCESS_STATUS.value) + def _ensure_model_is_downloaded(self, model_id: str) -> TSStatus: + model_info = self._model_manager.get_model_info(model_id) + if not model_info.download_on_start and model_info.state != ModelStates.ACTIVE: + return TSStatus( + code=TSStatusCode.MODEL_NOT_EXIST_ERROR.value, + message=f"Model [{model_id}] is not downloaded at start. Use 'LOAD MODEL {model_id} TO DEVICES ...' to download and load it; 'SHOW MODELS' shows it as active once it is downloaded.", + ) + return TSStatus(code=TSStatusCode.SUCCESS_STATUS.value) + # ==================== Inference ==================== def inference(self, req: TInferenceReq) -> TInferenceResp: status = self._ensure_model_is_registered(req.modelId) + if status.code == TSStatusCode.SUCCESS_STATUS.value: + status = self._ensure_model_is_downloaded(req.modelId) if status.code != TSStatusCode.SUCCESS_STATUS.value: return TInferenceResp(status, []) return self._inference_manager.inference(req) def forecast(self, req: TForecastReq) -> TForecastResp: status = self._ensure_model_is_registered(req.modelId) + if status.code == TSStatusCode.SUCCESS_STATUS.value: + status = self._ensure_model_is_downloaded(req.modelId) if status.code != TSStatusCode.SUCCESS_STATUS.value: return TForecastResp(status, []) return self._inference_manager.forecast(req) diff --git a/iotdb-core/ainode/src/assembly/resources/NOTICE-binary b/iotdb-core/ainode/src/assembly/resources/NOTICE-binary index e30b3ffd5b396..4d89656316b50 100644 --- a/iotdb-core/ainode/src/assembly/resources/NOTICE-binary +++ b/iotdb-core/ainode/src/assembly/resources/NOTICE-binary @@ -58,6 +58,15 @@ This product includes source code derived from the chronos-forecasting project: ============================================================================ +This product includes source code derived from the Timestral project: + + Timestral + Copyright 2026 Mistral AI + Licensed under the Apache License, Version 2.0 + https://github.com/mistralai/timestral + +============================================================================ + Apache Commons Collections Copyright 2001-2019 The Apache Software Foundation diff --git a/iotdb-core/ainode/tests/test_builtin_model_download.py b/iotdb-core/ainode/tests/test_builtin_model_download.py new file mode 100644 index 0000000000000..4a08d90a46eb9 --- /dev/null +++ b/iotdb-core/ainode/tests/test_builtin_model_download.py @@ -0,0 +1,146 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +import os +import shutil +import tempfile +import unittest +from unittest import mock + +from iotdb.ainode.core.config import AINodeDescriptor +from iotdb.ainode.core.constant import TSStatusCode +from iotdb.ainode.core.model.model_constants import ModelStates +from iotdb.ainode.core.model.model_info import BUILTIN_HF_TRANSFORMERS_MODEL_MAP +from iotdb.ainode.core.model.model_storage import ModelStorage +from iotdb.ainode.core.rpc.handler import AINodeRPCServiceHandler + + +class BuiltinModelDownloadTest(unittest.TestCase): + """Timestral's weights are downloaded when the model is loaded, not at start.""" + + def setUp(self): + self.cwd = os.getcwd() + self.work = tempfile.mkdtemp() + os.chdir(self.work) + self.model_dir = os.path.join( + self.work, + AINodeDescriptor().get_config().get_ain_models_dir(), + "builtin", + "timestral", + ) + self.downloads = [] + self.addCleanup(self.restore) + for model_info in BUILTIN_HF_TRANSFORMERS_MODEL_MAP.values(): + self.addCleanup(setattr, model_info, "state", model_info.state) + + def restore(self): + os.chdir(self.cwd) + shutil.rmtree(self.work) + + def fake_download(self, repo_id, filename, local_dir): + self.downloads.append(repo_id) + if repo_id != "mistralai/Timestral-209M": + raise OSError("offline") + with open(os.path.join(local_dir, filename), "w") as f: + f.write("{}") + return os.path.join(local_dir, filename) + + def discover(self): + with mock.patch( + "iotdb.ainode.core.model.model_storage.hf_hub_download", + side_effect=self.fake_download, + ): + storage = ModelStorage() + storage._executor.shutdown(wait=True) + return storage + + def test_not_downloaded_at_start(self): + storage = self.discover() + self.assertNotIn("mistralai/Timestral-209M", self.downloads) + self.assertIn("thuml/sundial-base-128m", self.downloads) + self.assertEqual( + storage.get_model_info("timestral").state, ModelStates.INACTIVE + ) + self.assertFalse(os.path.exists(os.path.join(self.model_dir, "config.json"))) + + def test_downloaded_on_request(self): + storage = self.discover() + with mock.patch( + "iotdb.ainode.core.model.model_storage.hf_hub_download", + side_effect=self.fake_download, + ): + self.assertTrue(storage.download_builtin_model("timestral")) + # a second request finds the model active and downloads nothing + count = len(self.downloads) + self.assertTrue(storage.download_builtin_model("timestral")) + self.assertEqual(len(self.downloads), count) + self.assertEqual(storage.get_model_info("timestral").state, ModelStates.ACTIVE) + self.assertTrue(os.path.exists(os.path.join(self.model_dir, "config.json"))) + self.assertTrue( + os.path.exists(os.path.join(self.model_dir, "model.safetensors")) + ) + + # once downloaded, the model is active from the next start on + storage = self.discover() + self.assertEqual(storage.get_model_info("timestral").state, ModelStates.ACTIVE) + + def test_failed_download_leaves_model_inactive(self): + storage = self.discover() + with mock.patch( + "iotdb.ainode.core.model.model_storage.hf_hub_download", + side_effect=OSError("offline"), + ): + self.assertFalse(storage.download_builtin_model("timestral")) + self.assertEqual( + storage.get_model_info("timestral").state, ModelStates.INACTIVE + ) + + def test_malformed_config_leaves_model_inactive(self): + storage = self.discover() + + def malformed_download(repo_id, filename, local_dir): + with open(os.path.join(local_dir, filename), "w") as f: + f.write("{not json") + return os.path.join(local_dir, filename) + + with mock.patch( + "iotdb.ainode.core.model.model_storage.hf_hub_download", + side_effect=malformed_download, + ): + self.assertFalse(storage.download_builtin_model("timestral")) + self.assertEqual( + storage.get_model_info("timestral").state, ModelStates.INACTIVE + ) + + def test_forecast_before_download_is_rejected(self): + storage = self.discover() + handler = mock.Mock(_model_manager=storage) + status = AINodeRPCServiceHandler._ensure_model_is_downloaded( + handler, "timestral" + ) + self.assertEqual(status.code, TSStatusCode.MODEL_NOT_EXIST_ERROR.value) + self.assertIn("LOAD MODEL timestral", status.message) + for model_id in ("sundial", "arima"): + status = AINodeRPCServiceHandler._ensure_model_is_downloaded( + handler, model_id + ) + self.assertEqual(status.code, TSStatusCode.SUCCESS_STATUS.value) + + +if __name__ == "__main__": + unittest.main() diff --git a/iotdb-core/ainode/tests/test_builtin_model_load.py b/iotdb-core/ainode/tests/test_builtin_model_load.py new file mode 100644 index 0000000000000..16e9b6e602db1 --- /dev/null +++ b/iotdb-core/ainode/tests/test_builtin_model_load.py @@ -0,0 +1,132 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +import queue +import threading +import time +import unittest +from unittest import mock + +import torch + +from iotdb.ainode.core.inference.pool_controller import PoolController +from iotdb.ainode.core.model.model_constants import ModelStates + + +class BuiltinModelLoadTest(unittest.TestCase): + """LOAD MODEL downloads on-request weights without blocking the pool controller.""" + + def setUp(self): + self.release_download = threading.Event() + self.download_succeeds = True + self.events = [] + self.models = { + "timestral": mock.Mock(download_on_start=False, state=ModelStates.INACTIVE), + "sundial": mock.Mock(download_on_start=True, state=ModelStates.ACTIVE), + } + self.controller = self._controller() + + def _controller(self): + controller = object.__new__(PoolController) + controller._stop_event = threading.Event() + controller._task_queue = queue.Queue() + controller._tasks_waiting_for_download = {} + controller._executor = mock.Mock() + controller._request_pool_map = {} + controller._model_manager = mock.Mock() + controller._model_manager.get_model_info.side_effect = self.models.get + controller._model_manager.download_builtin_model.side_effect = self._download + + def load_on_devices(model_ids, task): + self.events.append(("load", self._model_being_loaded)) + return [] + + def load(model_id, device_id_list): + self._model_being_loaded = model_id + PoolController._load_one_model_task(controller, model_id, device_id_list) + + controller._executor.submit_batch.side_effect = load_on_devices + controller._load_one_model_task = load + controller._unload_one_model_task = lambda model_id, _: self.events.append( + ("unload", model_id) + ) + controller._pool_control_worker_thread = threading.Thread( + target=controller._worker_loop, daemon=True + ) + controller._pool_control_worker_thread.start() + return controller + + def _download(self, model_id): + self.release_download.wait(10) + if self.download_succeeds: + self.models[model_id].state = ModelStates.ACTIVE + return self.download_succeeds + + def _wait_for(self, count): + deadline = time.monotonic() + 5 + while len(self.events) < count and time.monotonic() < deadline: + time.sleep(0.01) + + def _stop(self): + self.release_download.set() + self.controller.stop() + + def test_other_models_are_not_blocked_by_a_download(self): + self.addCleanup(self._stop) + cpu = [torch.device("cpu")] + self.controller.load_model("timestral", cpu) + self.controller.unload_model("timestral", cpu) + self.controller.load_model("sundial", cpu) + self.controller.unload_model("sundial", cpu) + self._wait_for(2) + # sundial's tasks run while timestral is still downloading + self.assertEqual(self.events, [("load", "sundial"), ("unload", "sundial")]) + + self.release_download.set() + self._wait_for(4) + # timestral's unload runs after its load, in the order they were requested + self.assertEqual( + self.events[2:], [("load", "timestral"), ("unload", "timestral")] + ) + + def test_failed_download_skips_the_load_only(self): + self.addCleanup(self._stop) + self.download_succeeds = False + cpu = [torch.device("cpu")] + self.controller.load_model("timestral", cpu) + self.controller.unload_model("timestral", cpu) + self.release_download.set() + self._wait_for(1) + self.assertEqual(self.events, [("unload", "timestral")]) + + def test_stop_does_not_wait_for_a_download(self): + self.controller.load_model("timestral", [torch.device("cpu")]) + deadline = time.monotonic() + 5 + while ( + not self.controller._model_manager.download_builtin_model.called + and time.monotonic() < deadline + ): + time.sleep(0.01) + started = time.monotonic() + self.controller.stop() + self.assertLess(time.monotonic() - started, 3) + self.release_download.set() + + +if __name__ == "__main__": + unittest.main() diff --git a/iotdb-core/ainode/tests/test_timestral.py b/iotdb-core/ainode/tests/test_timestral.py new file mode 100644 index 0000000000000..78774d0972a28 --- /dev/null +++ b/iotdb-core/ainode/tests/test_timestral.py @@ -0,0 +1,211 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +import copy +import json +import os +import tempfile +import unittest +from unittest import mock + +import torch +from safetensors.torch import save_file + +from iotdb.ainode.core.model.model_info import BUILTIN_HF_TRANSFORMERS_MODEL_MAP +from iotdb.ainode.core.model.timestral.forecaster import TimestralForecaster +from iotdb.ainode.core.model.timestral.modeling_timestral import ( + TimestralForPrediction, +) +from iotdb.ainode.core.model.timestral.pipeline_timestral import TimestralPipeline + +QUANTILES = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] + +# A small Timestral with the structure of the published checkpoints; one decode block is +# block_size * patch_size = 16 steps +TINY_CONFIG = { + "patch_size": 4, + "d_model": 16, + "num_heads": 2, + "mlp_hidden_dim": 32, + "block_schedule": ["time", "time_variate"], + "quantile_knots": QUANTILES, + "block_size": 4, + "context_length": 256, + "rope_theta": 10000.0, + "rope_partial_factor": 0.5, + "norm_eps": 1e-06, + "attn_bias": False, + "mlp_bias": False, + "time_attn_gate": "sdpa_sigmoid", + "variate_attn_gate": "sdpa_sigmoid", + "minimum_scale": 1e-06, + "use_asinh": True, + "stabilize_with_global": False, + "covariate_enabled": True, + "split_context_horizon_attn": True, + "attend_missing_patches": True, + "quantile_width_factor": 3.0, +} + + +class TimestralTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + torch.manual_seed(0) + cls.model = TimestralForPrediction(**copy.deepcopy(TINY_CONFIG)).eval() + # the zero-initialized forecast query would make every horizon step alike + with torch.no_grad(): + cls.model.representation.forecast_query.normal_() + cls.checkpoint = tempfile.TemporaryDirectory() + with open(os.path.join(cls.checkpoint.name, "config.json"), "w") as f: + json.dump(TINY_CONFIG, f) + save_file( + cls.model.state_dict(), + os.path.join(cls.checkpoint.name, "model.safetensors"), + ) + + @classmethod + def tearDownClass(cls): + cls.checkpoint.cleanup() + + def load_pipeline(self): + model_info = copy.copy(BUILTIN_HF_TRANSFORMERS_MODEL_MAP["timestral"]) + with mock.patch( + "iotdb.ainode.core.model.model_loader.get_model_path", + return_value=self.checkpoint.name, + ), mock.patch( + "iotdb.ainode.core.model.utils.get_model_path", + return_value=self.checkpoint.name, + ): + return TimestralPipeline(model_info, device=torch.device("cpu")) + + def run_pipeline(self, pipeline, inputs, output_length): + kwargs = {"output_length": output_length} + series = pipeline.preprocess(inputs, **kwargs) + return pipeline.postprocess(pipeline.forecast(series, **kwargs), **kwargs) + + def test_checkpoint_loads(self): + loaded = TimestralForPrediction.from_pretrained(self.checkpoint.name) + self.assertFalse(loaded.training) + expected = self.model.state_dict() + actual = loaded.state_dict() + self.assertEqual(set(expected), set(actual)) + for key in expected: + torch.testing.assert_close(actual[key], expected[key], rtol=0, atol=0) + + def test_checkpoint_with_missing_weight_is_rejected(self): + with tempfile.TemporaryDirectory() as directory: + with open(os.path.join(directory, "config.json"), "w") as f: + json.dump(TINY_CONFIG, f) + state_dict = dict(self.model.state_dict()) + state_dict.pop("head.projection.output_layer.weight") + save_file(state_dict, os.path.join(directory, "model.safetensors")) + with self.assertRaises(RuntimeError): + TimestralForPrediction.from_pretrained(directory) + + def test_unknown_config_key_is_rejected(self): + with self.assertRaises(ValueError): + TimestralForPrediction(**TINY_CONFIG, unknown_key=1) + + def test_rotary_frequencies_follow_the_model(self): + model = TimestralForPrediction(**copy.deepcopy(TINY_CONFIG)) + expected = model.transformer.rope.inv_freq.clone() + model.double() + self.assertEqual(model.transformer.rope.inv_freq.dtype, torch.float32) + torch.testing.assert_close(model.transformer.rope.inv_freq, expected) + + def test_builtin_pipeline_forecasts_median(self): + pipeline = self.load_pipeline() + self.assertIsInstance(pipeline.model, TimestralForPrediction) + torch.manual_seed(1) + inputs = [ + {"targets": torch.randn(100).cumsum(0)}, + {"targets": torch.randn(3, 40).cumsum(-1)}, + {"targets": torch.randn(100).cumsum(0)}, + ] + # longer than one decode block + outputs = self.run_pipeline(pipeline, copy.deepcopy(inputs), 40) + self.assertEqual([tuple(o.shape) for o in outputs], [(1, 40), (3, 40), (1, 40)]) + forecaster = TimestralForecaster(self.model) + for output, item in zip(outputs, inputs): + self.assertEqual(output.dtype, torch.float32) + self.assertFalse(output.is_inference()) + targets = item["targets"].reshape(-1, item["targets"].shape[-1]) + expected = forecaster.forecast(targets.unsqueeze(0), 40).median[0] + torch.testing.assert_close(output, expected) + + def test_batches_are_bounded_by_rows(self): + # Many requests of the same shape are split into batches of at most + # max_batch_rows targets, and each forecast equals the forecast of its own series. + pipeline = self.load_pipeline() + pipeline.max_batch_rows = 4 + torch.manual_seed(2) + inputs = [{"targets": torch.randn(2, 60).cumsum(-1)} for _ in range(5)] + inputs.append({"targets": torch.randn(6, 60).cumsum(-1)}) + rows = [] + forecast = pipeline.forecaster.forecast + + def record_rows(context, *args, **kwargs): + rows.append(context.shape[0] * context.shape[1]) + return forecast(context, *args, **kwargs) + + with mock.patch.object(pipeline.forecaster, "forecast", record_rows): + outputs = self.run_pipeline(pipeline, copy.deepcopy(inputs), 24) + self.assertEqual(rows, [4, 4, 2, 6]) + forecaster = TimestralForecaster(self.model) + for output, item in zip(outputs, inputs): + expected = forecaster.forecast(item["targets"].unsqueeze(0), 24).median[0] + torch.testing.assert_close(output, expected) + + def test_covariates(self): + pipeline = self.load_pipeline() + torch.manual_seed(2) + inputs = [ + { + "targets": torch.randn(2, 60), + "past_covariates": { + "past_only": torch.randn(60), + "known": torch.randn(60), + }, + "future_covariates": {"known": torch.randn(24)}, + } + ] + series = pipeline.preprocess(copy.deepcopy(inputs), output_length=24) + self.assertEqual(tuple(series[0]["context"].shape), (1, 2, 60)) + self.assertEqual(tuple(series[0]["past_covariates"].shape), (1, 1, 60)) + self.assertEqual(tuple(series[0]["known_covariates"].shape), (1, 1, 84)) + torch.testing.assert_close( + series[0]["known_covariates"][0, 0], + torch.cat( + [ + inputs[0]["past_covariates"]["known"], + inputs[0]["future_covariates"]["known"], + ] + ), + ) + outputs = self.run_pipeline(pipeline, copy.deepcopy(inputs), 24) + self.assertEqual(tuple(outputs[0].shape), (2, 24)) + self.assertTrue(torch.isfinite(outputs[0]).all()) + without = self.run_pipeline( + pipeline, [{"targets": inputs[0]["targets"].clone()}], 24 + ) + self.assertFalse(torch.equal(outputs[0], without[0])) + + +if __name__ == "__main__": + unittest.main()