diff --git a/LICENSE b/LICENSE index 222444ed3fbc9..e313d7a2e182c 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 TiRex-2 project. + +./iotdb-core/ainode/iotdb/ainode/core/model/tirex2/* + +The TiRex-2 is open source software licensed under the Apache License 2.0 +Project page: https://github.com/NX-AI/tirex-2 +License: https://github.com/NX-AI/tirex-2/blob/main/LICENSE +Copyright 2026 NXAI GmbH + +-------------------------------------------------------------------------------- diff --git a/NOTICE b/NOTICE index b457766a81ba0..a583df63ed022 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 TiRex-2 project: + + TiRex-2 + Copyright 2026 NXAI GmbH + Licensed under the Apache License, Version 2.0 + https://github.com/NX-AI/tirex-2 + +============================================================================ + 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..fc4cf1f404a13 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 @@ -397,6 +397,7 @@ public static void forecastTableFunctionErrorTest( errorTest( statement, invalidOutputLengthSQLWithZero, "701: OUTPUT_LENGTH should be greater than 0"); + int maxOutputLength = AINodeTestUtils.getMaxOutputLength(modelInfo.getModelId()); String invalidOutputLengthSQLWithOutOfRange = String.format( FORECAST_TABLE_FUNCTION_SQL_TEMPLATE, @@ -405,13 +406,15 @@ public static void forecastTableFunctionErrorTest( 5760, 2880, 5760, - 2881, + maxOutputLength + 1, 1, "time"); errorTest( statement, invalidOutputLengthSQLWithOutOfRange, - "1599: Error occurred while executing forecast:[Attribute output_length expect value between 1 and 2880, got 2881 instead.]"); + String.format( + "1599: Error occurred while executing forecast:[Attribute output_length expect value between 1 and %d, got %d instead.]", + maxOutputLength, maxOutputLength + 1)); String invalidOutputIntervalSQL = String.format( 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..79c242aff400f 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 @@ -62,7 +62,9 @@ public class AINodeTestUtils { new AbstractMap.SimpleEntry<>( "toto", new FakeModelInfo("toto", "toto", "builtin", "active")), new AbstractMap.SimpleEntry<>( - "moment", new FakeModelInfo("moment", "moment", "builtin", "active"))) + "moment", new FakeModelInfo("moment", "moment", "builtin", "active")), + new AbstractMap.SimpleEntry<>( + "tirex2", new FakeModelInfo("tirex2", "tirex2", "builtin", "active"))) .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); public static final Map BUILTIN_MODEL_MAP; @@ -95,6 +97,17 @@ public class AINodeTestUtils { BUILTIN_MODEL_MAP = Collections.unmodifiableMap(tmp); } + /** AINode's default limit on output_length (ain_inference_max_output_length). */ + public static final int DEFAULT_MAX_OUTPUT_LENGTH = 2880; + + /** Built-in models whose own forecast horizon is below AINode's limit. */ + private static final Map MODEL_MAX_OUTPUT_LENGTH = + Collections.singletonMap("tirex2", 320); + + public static int getMaxOutputLength(String modelId) { + return MODEL_MAX_OUTPUT_LENGTH.getOrDefault(modelId, DEFAULT_MAX_OUTPUT_LENGTH); + } + private static final Logger LOGGER = LoggerFactory.getLogger(AINodeTestUtils.class); public static void checkHeader(ResultSetMetaData resultSetMetaData, String title) diff --git a/iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py b/iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py index 1f1bd083f26fb..426c5a2094381 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py +++ b/iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py @@ -217,6 +217,21 @@ def _do_inference_and_construct_resp( resp_list.append(resp) return resp_list + def _max_output_length(self, model_id: str) -> int: + """AINode's output_length limit, lowered to the model's own limit if it has one. + + The check runs before a request reaches an inference pool, where a failing + request would stop the pool. + """ + max_output_length = ( + AINodeDescriptor().get_config().get_ain_inference_max_output_length() + ) + if self._model_manager.is_model_registered(model_id): + model_limit = self._model_manager.get_model_info(model_id).max_output_length + if model_limit is not None: + max_output_length = min(max_output_length, model_limit) + return max_output_length + def _run_forecast( self, req, @@ -232,17 +247,10 @@ def _run_forecast( inference_attrs = extract_attrs(req) output_length = int(inference_attrs.get("output_length", 96)) - if ( - output_length - > AINodeDescriptor().get_config().get_ain_inference_max_output_length() - ): + max_output_length = self._max_output_length(model_id) + if not 1 <= output_length <= max_output_length: raise NumericalRangeException( - "output_length", - output_length, - 1, - AINodeDescriptor() - .get_config() - .get_ain_inference_max_output_length(), + "output_length", output_length, 1, max_output_length ) model_inputs: list[dict[str, torch.Tensor | dict[str, torch.Tensor]]] = [ 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..ec9e125322c5b 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, + max_output_length: Optional[int] = None, ): self.model_id = model_id self.model_type = model_type @@ -41,6 +42,9 @@ def __init__( self.repo_id = repo_id self.auto_map = auto_map self.hub_mixin_cls = hub_mixin_cls + # The longest forecast the model can produce, if it is shorter than AINode's own limit. + # Requests are checked against it before they reach an inference pool. + self.max_output_length = max_output_length def __repr__(self): return ( @@ -178,4 +182,14 @@ def __repr__(self): "AutoModelForCausalLM": "modeling_moment.MomentForPrediction", }, ), + "tirex2": ModelInfo( + model_id="tirex2", + category=ModelCategory.BUILTIN, + state=ModelStates.INACTIVE, + model_type="tirex2", + pipeline_cls="pipeline_tirex2.TiRex2Pipeline", + repo_id="NX-AI/TiRex-2", + hub_mixin_cls="modeling_tirex2.TiRex2ForPrediction", + max_output_length=320, + ), } diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/__init__.py new file mode 100644 index 0000000000000..188f0cc57cb69 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/__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 (c) NXAI GmbH. +# SPDX-License-Identifier: Apache-2.0 diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/__init__.py new file mode 100644 index 0000000000000..188f0cc57cb69 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/__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 (c) NXAI GmbH. +# SPDX-License-Identifier: Apache-2.0 diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/attention_block.py b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/attention_block.py new file mode 100644 index 0000000000000..717955771d70d --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/attention_block.py @@ -0,0 +1,438 @@ +# 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 (c) NXAI GmbH. +# SPDX-License-Identifier: Apache-2.0 + +import math +import warnings +from typing import Any + +import torch +from torch.nn import functional as F + +from iotdb.ainode.core.model.tirex2.component.feedforward import MLP + +FLEX_MASK_BLOCK_SIZE = 128 +FLEX_KERNEL_BLOCK_SIZE = 32 + +create_block_mask = None +_flex_attention = None +flex_attention = None +_flex_attention_import_attempted = False + + +def _load_flex_attention() -> bool: + """Load FlexAttention on demand, only for layers that opt into it.""" + global create_block_mask + global _flex_attention + global flex_attention + global _flex_attention_import_attempted + + if flex_attention is not None: + return True + if _flex_attention_import_attempted: + return False + + _flex_attention_import_attempted = True + try: + from torch.nn.attention.flex_attention import ( + create_block_mask as imported_create_block_mask, + ) + from torch.nn.attention.flex_attention import ( + flex_attention as imported_flex_attention, + ) + except ImportError: # pragma: no cover - depends on installed PyTorch version + return False + + create_block_mask = imported_create_block_mask + _flex_attention = imported_flex_attention + # FlexAttention is intended to run through the compiled fused kernel. + # dynamic=True avoids recompiling for every distinct sequence length. + flex_attention = torch.compile(imported_flex_attention, dynamic=True) + return True + + +def is_flex_attention_available() -> bool: + """Return whether FlexAttention can be loaded in the current PyTorch install.""" + return _load_flex_attention() + + +def _require_flex_attention() -> None: + if not _load_flex_attention(): + raise RuntimeError( + "FlexAttention was requested, but torch.nn.attention.flex_attention " + "is not available in this PyTorch installation." + ) + + +def norm(x: torch.Tensor) -> torch.Tensor: + """Apply RMS normalization over last dimension.""" + return F.rms_norm(x, (x.size(-1),)) + + +class AttentionLayer(torch.nn.Module): + """ + Multi-head attention with an optional group-wise mask. + + Input is expected as `[B, L, D]`, where attention is computed across the length dimension `L` (this may + represent time or variates depending on the caller). When a group mask is provided, attention scores + between tokens with different group ids are masked out to form a block-diagonal attention pattern. + Key/query/value vectors are assumed to be pre-normalized and any residual connections or output + normalization should be handled by the caller. + + Parameters + ---------- + input_dim: int + Dimension of the incoming key/query/value vectors. + kv_proj_dim: int + Dimension of each attention head after projection. + n_heads: int + Number of attention heads. + dropout: float + Dropout probability used inside scaled dot-product attention during training. + use_qk_norm: bool + Whether to apply RMS normalization to the query and key vectors before + computing attention scores (default: True). + use_flex_attention: bool + Whether to use FlexAttention instead of scaled dot product attention + (default: False). + """ + + def __init__( + self, + input_dim: int, + kv_proj_dim: int, + n_heads: int = 4, + dropout: float = 0, + disable_singleton_attention: bool = False, + use_rope: bool = False, + no_qk_scale: bool = False, + use_qk_norm: bool = True, + use_flex_attention: bool = False, + ): + super().__init__() + self.input_dim: int = input_dim + self.kv_proj_dim: int = kv_proj_dim + self.n_heads: int = n_heads + self.dropout: float = dropout + self.disable_singleton_attention: bool = disable_singleton_attention + self.use_rope: bool = use_rope + self.use_qk_norm: bool = use_qk_norm + self.use_flex_attention: bool = use_flex_attention + if self.use_flex_attention: + _require_flex_attention() + + self.embedding_dim = kv_proj_dim * n_heads + + self.WK = self.create_weights(self.input_dim, self.embedding_dim, bias=False) + self.WQ = self.create_weights(self.input_dim, self.embedding_dim, bias=False) + self.WV = self.create_weights(self.input_dim, self.embedding_dim, bias=False) + self.WO = self.create_weights(self.embedding_dim, self.input_dim, bias=False) + + self.scale = 1 / math.sqrt(self.kv_proj_dim) if not no_qk_scale else 1 + + def create_weights(self, input_dim, embedding_dim, bias: bool) -> torch.nn.Linear: + """Create and initialize a linear projection layer.""" + linear = torch.nn.Linear( + in_features=input_dim, out_features=embedding_dim, bias=bias + ) + return linear + + def _build_group_block_mask( + self, + group_vector: torch.Tensor, + target_mask: torch.Tensor, + seq_len: int, + ): + """Build the block-sparse FlexAttention mask matching the dense group mask.""" + _require_flex_attention() + assert create_block_mask is not None + + group_vector = ( + group_vector.squeeze(-1) if group_vector.ndim == 2 else group_vector + ) + target_mask = target_mask.squeeze(-1) if target_mask.ndim == 2 else target_mask + target_mask = target_mask.to(torch.bool) + + def mask_mod(b, h, q_idx, kv_idx): + same_group = group_vector[q_idx] == group_vector[kv_idx] + covariate_to_target = (~target_mask[q_idx]) & target_mask[kv_idx] + return same_group & ~covariate_to_target + + return create_block_mask( + mask_mod, + B=None, + H=None, + Q_LEN=seq_len, + KV_LEN=seq_len, + device=group_vector.device, + BLOCK_SIZE=FLEX_MASK_BLOCK_SIZE, + ) + + def forward( + self, + x: torch.Tensor, + group_vector: torch.Tensor | None = None, + target_mask: torch.Tensor | None = None, + state: Any | None = None, + ) -> torch.Tensor: + """ + Compute attention with an optional group mask. + + Inputs are expected to be pre-normalized. The output contains only the attention transformation; any + residual connection must be applied by the caller if needed. + + Parameters + ---------- + x: torch.Tensor + Pre-normalized input tensor of shape [B, L, D]. + group_vector: torch.Tensor | None + Optional group ids with shape `[L]` or `[L, 1]` aligned to the sequence dimension. Positions + with different ids are masked from attending to each other. If `None`, attention is computed + across all positions. + state: Any | None + Unused placeholder for compatibility with other mixer interfaces. + + Returns + ------- + torch.Tensor + The attention-transformed tokens of shape `[B, L, D]`. + """ + _, L, _ = x.shape + + def unpack(x): + """Reshape from (B, L, D) to multi-head format (B, H, L, D_head).""" + return x.reshape(x.shape[0], L, self.n_heads, self.kv_proj_dim).permute( + 0, 2, 1, 3 + ) + + def pack(x): + """Reshape from multi-head format (B, H, L, D_head) to (B, L, D).""" + return x.permute(0, 2, 1, 3).reshape( + x.shape[0], L, self.n_heads * self.kv_proj_dim + ) + + qk_norm = norm if self.use_qk_norm else (lambda t: t) + Q = qk_norm( + unpack(self.WQ(x)) + ) # [B, L, D_in] -> [B, L, D_embed] -> [B, H, L, D_head] + K = qk_norm( + unpack(self.WK(x)) + ) # [B, L, D_in] -> [B, L, D_embed] -> [B, H, L, D_head] + V = unpack(self.WV(x)) # [B, L, D_in] -> [B, L, D_embed] -> [B, H, L, D_head] + + mask = None + block_mask = None + # create mask, allow attention only between elements with the same group id + if group_vector is not None: + if group_vector.shape[0] != L: + raise RuntimeError( + f"Group mask length has to be the same as first dimension of x, " + f"but found {group_vector.shape[0]} and {L}, respectively" + ) + + group_vector = ( + group_vector.squeeze(-1) if group_vector.ndim == 2 else group_vector + ) + + if target_mask is None: + raise RuntimeError( + f"Attention shall restrict information flow between target and covariates but no target_mask provided" + ) + + if self.use_flex_attention: + block_mask = self._build_group_block_mask(group_vector, target_mask, L) + else: + mask = group_vector.unsqueeze(0) == group_vector.unsqueeze(1) # [L, L] + target_mask = ( + target_mask.squeeze(-1) if target_mask.ndim == 2 else target_mask + ) + # Block covariates (rows where ~target_mask) from attending to targets (cols where target_mask) + covariate_to_target = ~target_mask.unsqueeze(1) & target_mask.unsqueeze( + 0 + ) # [L, L] + mask = mask & ~covariate_to_target + mask = mask[ + None, None, ... + ] # mask has to be broadcastable to query dimensions + + # Compute attention + if self.use_flex_attention: + _require_flex_attention() + if Q.is_cuda: + assert flex_attention is not None + attention_out = flex_attention( # [B, H, L, D_head] + Q, + K, + V, + block_mask=block_mask, + scale=self.scale, + kernel_options={ + "BLOCK_M": FLEX_KERNEL_BLOCK_SIZE, + "BLOCK_N": FLEX_KERNEL_BLOCK_SIZE, + }, + ) + else: + assert _flex_attention is not None + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message="flex_attention called without torch.compile", + category=UserWarning, + ) + attention_out = _flex_attention( # [B, H, L, D_head] + Q, + K, + V, + block_mask=block_mask, + scale=self.scale, + ) + else: + attention_out = ( + torch.nn.functional.scaled_dot_product_attention( # [B, H, L, D_head] + Q, + K, + V, + attn_mask=mask, + dropout_p=self.dropout if self.training else 0.0, + scale=self.scale, + ) + ) + + packed_attention = pack(attention_out) # [B, H, L, D_head] -> [B, L, D_embed] + x_trans = self.WO(packed_attention) # [B, L, D_embed] -> [B, L, D_in] + x_trans = torch.nn.functional.dropout( + x_trans, self.dropout, training=self.training + ) + + return x_trans + + +class AttentionBlock(torch.nn.Module): + """Transformer attention block with RMS-normalized attention and FFN residuals. + + The block combines multi-head attention with a feed-forward network. Optional + group-wise attention masks are applied when group IDs are supplied at forward time. + """ + + def __init__( + self, + input_dim: int, + n_heads: int, + dropout: float, + act_fn: torch.nn.Module, + use_qk_norm: bool = True, + use_flex_attention: bool = False, + ): + super().__init__() + self.input_dim = input_dim + self.n_heads = n_heads + self.dropout = dropout + self.use_qk_norm = use_qk_norm + self.use_flex_attention = use_flex_attention + self._collect_group_stats_this_step = False + self.use_post_norm = False + + assert ( + input_dim % n_heads == 0 + ), f"Input dimension {input_dim} has do be divisible by the number of heads {n_heads}" + self.kv_proj_dim = input_dim // n_heads + + self.norm_attn = torch.nn.RMSNorm(self.input_dim, eps=1e-6) + self.attn = AttentionLayer( + input_dim=input_dim, + kv_proj_dim=self.kv_proj_dim, + n_heads=self.n_heads, + dropout=self.dropout, + use_qk_norm=self.use_qk_norm, + use_flex_attention=self.use_flex_attention, + ) + + self.norm_ffn = torch.nn.RMSNorm(self.input_dim, eps=1e-6) + self.ffn = MLP(self.input_dim, 4 * self.input_dim, act_fn, dropout=self.dropout) + + def forward( + self, + x: torch.Tensor, + *args, + group_vector: torch.Tensor | None = None, + target_mask: torch.Tensor | None = None, + attention_idx: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + """Forward pass through attention and FFN with residual connections. + + Parameters + ---------- + x : torch.Tensor + Input tensor of shape [B, L, D] + group_vector : torch.Tensor or None, optional + Optional group ids for group-wise attention masking + attention_idx : torch.Tensor or None, optional + Variates in non-singleton groups (see ``variate_indices``), derived when omitted. + *args + Unused positional arguments for interface compatibility + **kwargs + Unused keyword arguments for interface compatibility + + Returns + ------- + torch.Tensor + Output tensor of shape [B, L, D]. + """ + x_in = x + + if group_vector is not None: + if attention_idx is None: + gv = ( + group_vector.squeeze(-1) if group_vector.ndim == 2 else group_vector + ) + _, inverse, counts = torch.unique( + gv, return_inverse=True, return_counts=True + ) + attention_idx = torch.where(counts[inverse] > 1)[0] + if attention_idx.numel() == 0: + return x_in # batch is entirely singleton groups -> block is a no-op + x = x[:, attention_idx, :] + group_vector = group_vector[attention_idx] + if target_mask is not None: + target_mask = target_mask[attention_idx] + + assert ( + x.shape[1] == group_vector.shape[0] + ), f"gather shape mismatch: x={x.shape[1]}, gv={group_vector.shape[0]}" + if target_mask is not None: + assert ( + target_mask.shape[0] == x.shape[1] + ), f"target_mask not sliced consistently: tm={target_mask.shape[0]}, x={x.shape[1]}" + + x_attn = self.attn( + self.norm_attn(x), group_vector=group_vector, target_mask=target_mask + ) + x = x + x_attn + + x_ffn = self.ffn(self.norm_ffn(x)) + x = x + x_ffn + # Scatter non-singleton outputs back into the full-L tensor + if group_vector is not None: + out = x_in.clone() + out[:, attention_idx, :] = x + assert out.shape == x_in.shape + x = out + + return x diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/bi_xlstm.py b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/bi_xlstm.py new file mode 100644 index 0000000000000..4c6c8b9c80092 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/bi_xlstm.py @@ -0,0 +1,184 @@ +# 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 (c) NXAI GmbH. +# SPDX-License-Identifier: Apache-2.0 + +"""This module exposes xLSTM blocks combining recurrent kernels with feed-forward adapters.""" + +from typing import Literal + +import torch +from torch import nn + +from iotdb.ainode.core.model.tirex2.component.feedforward import FeedForward +from iotdb.ainode.core.model.tirex2.component.mlstm_block import ( + mLSTMLayer, + mLSTMLayerConfig, +) +from iotdb.ainode.core.model.tirex2.component.norm import RMSNorm +from iotdb.ainode.core.model.tirex2.component.slstm_block import ( + sLSTMLayer, + sLSTMLayerConfig, +) + + +class BiXLSTM(nn.Module): + """Bi-directional xLSTM block that splits variates into forward and reverse subsets, + runs a separate xLSTM cell over each, and recombines the results. + + The cell type is controlled by ``cell_type``: + - ``"slstm"``: Uses sLSTM cells (original behaviour, equivalent to the former BiSLSTM). + - ``"mlstm"``: Uses mLSTM cells. + + The split is controlled by ``split_method``: + - ``"reverse_wo_target"``: The forward cell sees all variates, while the + reverse cell sees only the non-target covariates. The non-target variates + thus appear in both directions and are shared. + - ``"forward_only_target"``: The forward cell sees only the target variates + and the reverse cell sees only the non-target covariates. No variates are + shared between the two directions. + - ``"reverse_known_only"``: The forward cell sees all variates, while the + reverse cell sees only the known covariates (those with non-NaN values in + the future window). Past-only covariates are processed forward-only. + + Variates that appear in both directions have their forward and reverse + representations concatenated along the embedding dimension (2*D) and + downprojected back to D via a linear layer with dropout. Variates that appear + in only one direction keep their single-direction output unchanged. Residual + connections wrap both the recurrent stage and the subsequent feed-forward block. + + Parameters + ---------- + embedding_dim : int + Model dimension. + num_heads : int + Number of heads of the recurrent cell. + block_idx : int + Index of this block in the overall stack. + num_blocks : int + Total number of blocks in the stack. + dropout : float + Dropout probability applied in the combination projection. + split_method : {"forward_only_target", "reverse_wo_target", "reverse_known_only"} + Strategy for splitting variates between forward and reverse cells. + share_weights : bool + When True, the forward and reverse cells share the same parameters. + cell_type : {"slstm", "mlstm"} + Which recurrent cell to instantiate. + """ + + def __init__( + self, + embedding_dim: int, + num_heads: int, + block_idx: int, + num_blocks: int, + dropout: float = 0.0, + share_weights: bool = True, + cell_type: Literal["slstm", "mlstm"] = "slstm", + ): + super().__init__() + self.cell_type = cell_type + self.norm_lstm = RMSNorm(embedding_dim, eps=1e-6) + self._split = self._split_variates + self._recombine = self._recombine_variates + + match cell_type: + case "slstm": + slstm_config = sLSTMLayerConfig( + embedding_dim=embedding_dim, + num_heads=num_heads, + recurrent_weight_init="zeros", + bias_init="powerlaw_blockdependent", + ) + self.fwd_cell = sLSTMLayer( + slstm_config, block_idx=block_idx, num_blocks=num_blocks + ) + case "mlstm": + mlstm_config = mLSTMLayerConfig( + embedding_dim=embedding_dim, + num_heads=num_heads, + ) + self.fwd_cell = mLSTMLayer(mlstm_config) + case _: + raise ValueError( + f"Unknown cell type '({cell_type})'. Allowed 'mlstm' or 'slstm'." + ) + + self.combination = nn.Linear(2 * embedding_dim, embedding_dim, bias=False) + self.combination_dropout = nn.Dropout(dropout) + + self.norm_ffn = RMSNorm(embedding_dim, eps=1e-6) + self.ffn = FeedForward(embedding_dim) + + @property + def rev_cell(self) -> nn.Module: + """The reverse direction shares the forward cell; it is registered (and stored) once.""" + return self.fwd_cell + + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + # Older checkpoints also store the shared cell under ``rev_cell``. + for key in [key for key in state_dict if key.startswith(f"{prefix}rev_cell.")]: + del state_dict[key] + super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + + def forward(self, x: torch.Tensor, reverse_idx: torch.Tensor): + """Run the bidirectional xLSTM block; ``reverse_idx`` are the variates of the reverse cell.""" + x_norm = self.norm_lstm(x) + + x_fwd, x_rev, fwd_idc, rev_idc = self._split(x_norm, reverse_idx) + + x_fwd = self.fwd_cell(x_fwd) if x_fwd.shape[0] > 0 else x_fwd + if x_rev.shape[0] > 0: + x_rev_flipped = torch.flip(x_rev, dims=(1,)).contiguous() + x_rev_out = self.rev_cell(x_rev_flipped) + x_rev = torch.flip(x_rev_out, dims=(1,)) + + out = self._recombine(x_fwd, x_rev, fwd_idc, rev_idc, x) + + # residual connection around lstm + combination + out = x + out + + # ffn with residual connection + out = out + self.ffn(self.norm_ffn(out)) + + return out + + def _split_variates( + self, x: torch.Tensor, known_idc: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Forward cell sees all variates, reverse sees only known covariates (non-NaN future).""" + return x, x[known_idc], torch.arange(x.shape[0], device=x.device), known_idc + + def _recombine_variates( + self, + x_fwd: torch.Tensor, + x_rev: torch.Tensor, + fwd_idc: torch.Tensor, + rev_idc: torch.Tensor, + ref: torch.Tensor, + ) -> torch.Tensor: + """Covariates appear in both directions and are combined; targets are fwd-only.""" + # combine covariate representations from both directions + x_shared = torch.cat([x_fwd[rev_idc], x_rev], dim=-1) + x_shared = self.combination_dropout(self.combination(x_shared)) + + # start from fwd output (all variates), overwrite covariates with combined + out = x_fwd.to(ref.dtype) + out[rev_idc] = x_shared.to(ref.dtype) + return out diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/feedforward.py b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/feedforward.py new file mode 100644 index 0000000000000..207ca01464ad3 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/feedforward.py @@ -0,0 +1,58 @@ +# 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 (c) NXAI GmbH. +# SPDX-License-Identifier: Apache-2.0 + +"""Feed-forward networks: a plain MLP and the gated FeedForward of xLSTM Large.""" + +import math + +import torch +from torch import nn + + +class MLP(nn.Module): + """Feed-forward network ``wo(dropout(act(wi(x))))`` of the attention block.""" + + def __init__( + self, dim: int, hidden_dim: int, act_fn: nn.Module, dropout: float = 0.0 + ): + super().__init__() + self.wi = nn.Linear(dim, hidden_dim, bias=False) + self.wo = nn.Linear(hidden_dim, dim, bias=False) + self.act_fn = act_fn + self.dropout = nn.Dropout(dropout) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.wo(self.dropout(self.act_fn(self.wi(x)))) + + +class FeedForward(nn.Module): + """SwiGLU feed-forward network of xLSTM Large, 8/3 of the embedding rounded up to a multiple of 64.""" + + def __init__(self, dim: int): + super().__init__() + hidden_dim = math.ceil(dim * 2.6667 / 64) * 64 + self.proj_up_gate = nn.Linear(dim, hidden_dim, bias=False) + self.proj_up = nn.Linear(dim, hidden_dim, bias=False) + self.proj_down = nn.Linear(hidden_dim, dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.proj_down( + nn.functional.silu(self.proj_up_gate(x)) * self.proj_up(x) + ) diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/kernel/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/kernel/__init__.py new file mode 100644 index 0000000000000..1bb70e0f5c43f --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/kernel/__init__.py @@ -0,0 +1,67 @@ +# 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 (c) NXAI GmbH. +# SPDX-License-Identifier: Apache-2.0 + +"""sLSTM and mLSTM recurrence kernels. Import the modules directly; the triton ones need triton.""" + +from functools import cache + +import torch + + +@cache +def triton_available() -> bool: + """Whether the Triton kernels can be used; PyTorch's Windows wheels, for one, ship without Triton. + + Triton compiles a kernel from its Python source, so the kernel modules are imported here as well: + they fail to import where Triton is installed but the source is not, as in a frozen executable. + """ + try: + import triton # noqa: F401 + + from iotdb.ainode.core.model.tirex2.component.kernel import ( # noqa: F401 + mlstm_triton, + slstm_triton, + ) + except Exception: + return False + return True + + +# The Triton kernels multiply bfloat16 operands with tl.dot, which needs compute capability 8.0 +# (Ampere) or newer. Older GPUs, such as Volta and Turing, use the PyTorch kernels. +_MIN_TRITON_CAPABILITY = (8, 0) + + +@cache +def _supports_triton_kernels(device_index: int) -> bool: + return torch.cuda.get_device_capability(device_index) >= _MIN_TRITON_CAPABILITY + + +def use_triton(x: torch.Tensor) -> bool: + """Whether to run the recurrence of ``x`` on the Triton kernels. + + That is the case for CUDA inputs on a GPU of compute capability 8.0 or newer, when Triton is available. + """ + if not (x.is_cuda and triton_available()): + return False + index = x.device.index + return _supports_triton_kernels( + torch.cuda.current_device() if index is None else index + ) diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/kernel/mlstm_native.py b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/kernel/mlstm_native.py new file mode 100644 index 0000000000000..34408a32aab51 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/kernel/mlstm_native.py @@ -0,0 +1,117 @@ +# 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 (c) NXAI GmbH. +# SPDX-License-Identifier: Apache-2.0 + +"""mLSTM chunkwise forward pass in plain PyTorch, used on CPU/MPS and when gradients are needed. + +Vendored from https://github.com/NX-AI/mlstm_kernels (``mlstm_kernels==2.0.6``), reduced to what +TiRex-2 uses: ``mLSTMBackend`` in ``train_with_padding`` mode without initial or last states, with the +``chunkwise--native_autograd`` kernel. The computations and their order are kept as upstream, so the +results are bitwise identical. See ``mlstm_triton.py`` for the CUDA kernel. +Upstream's padding wrapper ignores the backend's chunk_size and eps, so they are always 64 and 1e-6. + +Notation: B batch, NH heads, S sequence length, DHQK / DHHV head dims of q, k / v, NC number of chunks, +L chunk size. vecB is the cumulative log forget gate within a chunk, vecA = vecB[-1] - vecB + vecI the +gate contribution from each step to the next chunk state, scaG = vecB[-1] the decay of a whole chunk. +""" + +import torch +import torch.nn.functional as F + + +def pad_to_chunks(q, k, v, i, f, chunk_size): + """Pad S to a multiple of the chunk size with steps that keep the state: input gate 0, forget gate 1.""" + S = q.shape[2] + pad = ((S - 1) // chunk_size + 1) * chunk_size - S + q, k, v = (F.pad(x, (0, 0, 0, pad), value=0.0) for x in (q, k, v)) + return q, k, v, F.pad(i, (0, pad), value=-25.0), F.pad(f, (0, pad), value=25.0) + + +def mlstm_chunkwise_native( + q: torch.Tensor, # (B, NH, S, DHQK) + k: torch.Tensor, # (B, NH, S, DHQK) + v: torch.Tensor, # (B, NH, S, DHHV) + i: torch.Tensor, # (B, NH, S) input gate preactivation + f: torch.Tensor, # (B, NH, S) forget gate preactivation + chunk_size: int = 64, + eps: float = 1e-6, +) -> torch.Tensor: # (B, NH, S, DHHV) + """mLSTM over a full sequence from zero states, in plain PyTorch (any device, differentiable).""" + B, NH, S, DHQK = q.shape + DHHV = v.shape[-1] + q, k, v, i, f = pad_to_chunks(q, k, v, i, f, chunk_size) + NC, L = q.shape[2] // chunk_size, chunk_size + qk_scale = DHQK**-0.5 + + vecI = i.reshape(B, NH, NC, L) + vecB = F.logsigmoid(f.reshape(B, NH, NC, L)).cumsum(-1) + matQ = q.reshape(B, NH, NC, L, DHQK) + matK = k.reshape(B, NH, NC, L, DHQK) + matV = v.reshape(B, NH, NC, L, DHHV) + + # Recurrently compute the states (C, n, m) entering each chunk, stored in the dtype of k. + vecA = vecB[..., -1, None] - vecB + vecI + scaG = vecB[..., -1] + scaA_max = vecA.max(-1).values + matC_k = torch.zeros((B, NH, DHQK, DHHV), dtype=k.dtype, device=k.device) + vecN_k = torch.zeros((B, NH, DHQK), dtype=k.dtype, device=k.device) + scaM_k = torch.zeros((B, NH), dtype=k.dtype, device=k.device) + matC_states, vecN_states, scaM_states = [], [], [] + for c in range(NC): + matC_states.append(matC_k.to(k.dtype)) + vecN_states.append(vecN_k.to(k.dtype)) + scaM_states.append(scaM_k.to(k.dtype)) + if c == NC - 1: + break + scaM_k_next = torch.max(scaG[:, :, c] + scaM_k, scaA_max[:, :, c]) + matK_gated = ( + matK[:, :, c] + * torch.exp(vecA[:, :, c, :] - scaM_k_next[..., None])[:, :, :, None] + ) + scaGbar_k = torch.exp(scaG[:, :, c] + scaM_k - scaM_k_next)[:, :, None] + matC_k = ( + scaGbar_k[..., None] * matC_k + matK_gated.transpose(-2, -1) @ matV[:, :, c] + ) + vecN_k = scaGbar_k * vecN_k + matK_gated.transpose(-2, -1).sum(-1) + scaM_k = scaM_k_next + matC_states = torch.stack(matC_states, dim=2) # (B, NH, NC, DHQK, DHHV) + vecN_states = torch.stack(vecN_states, dim=2) # (B, NH, NC, DHQK) + scaM_states = torch.stack(scaM_states, dim=2) # (B, NH, NC) + + # Compute the outputs within each chunk in parallel. + ltr = torch.tril(torch.ones((L, L), dtype=torch.bool, device=q.device)) + matLogD = ( + torch.where(ltr, vecB[:, :, :, :, None] - vecB[:, :, :, None, :], -float("inf")) + + vecI[:, :, :, None, :] + ) + vecM_inter = vecB + scaM_states[:, :, :, None] + vecM_combine = torch.maximum(vecM_inter, matLogD.max(dim=-1).values)[ + ..., None + ] # (B, NH, NC, L, 1) + vecM_inter = vecM_inter[..., None] + + matD = torch.exp(matLogD - vecM_combine) + matM = (matQ @ matK.transpose(-2, -1)) * qk_scale * matD + matQ_gated = matQ * torch.exp(vecM_inter - vecM_combine) * qk_scale + + matNumerator = matQ_gated @ matC_states + matM @ matV + vecDenom = matQ_gated @ vecN_states.unsqueeze(-1) + matM.sum(dim=-1, keepdim=True) + vecDenom_max = torch.maximum(torch.abs(vecDenom), torch.exp(-vecM_combine)) + matH = matNumerator / (vecDenom_max + eps) + return matH.reshape(B, NH, NC * L, DHHV)[:, :, :S, :] diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/kernel/mlstm_triton.py b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/kernel/mlstm_triton.py new file mode 100644 index 0000000000000..0364945a8c0d1 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/kernel/mlstm_triton.py @@ -0,0 +1,282 @@ +# 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 (c) NXAI GmbH. +# SPDX-License-Identifier: Apache-2.0 + +"""mLSTM chunkwise forward pass with Triton kernels, used on CUDA. + +Vendored from https://github.com/NX-AI/mlstm_kernels (``mlstm_kernels==2.0.6``), reduced to the forward +pass of the ``chunkwise--triton_limit_chunk`` kernel used by ``mLSTMBackend`` in ``train_with_padding`` +mode without initial or last states. Arithmetic and dtype casts follow upstream; memory accesses +use explicit offsets and masks. See ``mlstm_native.py`` for the notation. +""" + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from iotdb.ainode.core.model.tirex2.component.kernel.mlstm_native import ( + mlstm_chunkwise_native, + pad_to_chunks, +) + + +@triton.jit +def _recurrent_fw_C_kernel( + matK, # (B, NH, S, DHQK) + matV, # (B, NH, S, DHHV) + vecB, # (B, NH, NC, L) + vecI, # (B, NH, NC, L) + matC_states, # (B, NH, NC * DHQK, DHHV) + vecN_states, # (B, NH, NC * DHQK) + scaM_states, # (B, NH, NC) + S: tl.constexpr, + DHQK: tl.constexpr, + DHHV: tl.constexpr, + NC: tl.constexpr, + L: tl.constexpr, + siz_b_DHQK: tl.constexpr, + siz_b_DHHV: tl.constexpr, + DTYPE: tl.constexpr, +): + """Computes the states (C, n, m) entering each chunk.""" + idx_b_DHQK, idx_b_DHHV = tl.program_id(0), tl.program_id(1) + idx_b_BNH = tl.program_id(2).to(tl.int64) + offs_qk = idx_b_DHQK * siz_b_DHQK + tl.arange(0, siz_b_DHQK) + offs_hv = idx_b_DHHV * siz_b_DHHV + tl.arange(0, siz_b_DHHV) + mask_c = (offs_qk[:, None] < DHQK) & (offs_hv[None, :] < DHHV) + + matC_k_val = tl.zeros((siz_b_DHQK, siz_b_DHHV), dtype=tl.float32) + vecN_k_val = tl.zeros((siz_b_DHQK,), dtype=tl.float32) + scaM_k_val = 0.0 + + for k in range(NC): + offs_t = k * L + tl.arange(0, L) + matC_k_ptr = ( + matC_states + + (idx_b_BNH * NC + k) * DHQK * DHHV + + offs_qk[:, None] * DHHV + + offs_hv[None, :] + ) + tl.store(matC_k_ptr, matC_k_val, mask=mask_c) + if idx_b_DHHV == 0: + vecN_k_ptr = vecN_states + (idx_b_BNH * NC + k) * DHQK + offs_qk + tl.store(vecN_k_ptr, vecN_k_val, mask=offs_qk < DHQK) + if (idx_b_DHQK == 0) and (idx_b_DHHV == 0): + tl.store(scaM_states + idx_b_BNH * NC + k, scaM_k_val.to(tl.float32)) + + vecBI_offset = idx_b_BNH * S + k * L + vecB_last_k_val = tl.load(vecB + vecBI_offset + (L - 1)).to(tl.float32) + vecB_k_val = tl.load(vecB + vecBI_offset + tl.arange(0, L)).to(tl.float32) + vecI_k_val = tl.load(vecI + vecBI_offset + tl.arange(0, L)).to(tl.float32) + + vecA_k_val = (vecB_last_k_val - vecB_k_val) + vecI_k_val + scaG_k_val = vecB_last_k_val + scaM_next_val = tl.maximum(scaG_k_val + scaM_k_val, tl.max(vecA_k_val)) + + # load matK transposed (siz_b_DHQK, L) and matV (L, siz_b_DHHV) + matK_k_ptr = ( + matK + idx_b_BNH * S * DHQK + offs_qk[:, None] + offs_t[None, :] * DHQK + ) + matV_k_ptr = ( + matV + idx_b_BNH * S * DHHV + offs_t[:, None] * DHHV + offs_hv[None, :] + ) + matK_k_val = tl.load( + matK_k_ptr, mask=(offs_qk[:, None] < DHQK) & (offs_t[None, :] < S), other=0 + ).to(tl.float32) + matV_k_val = tl.load( + matV_k_ptr, mask=(offs_t[:, None] < S) & (offs_hv[None, :] < DHHV), other=0 + ).to(DTYPE) + + vecAbar_k_val = tl.exp(vecA_k_val - scaM_next_val) + scaGbar_k_val = tl.exp(scaG_k_val + scaM_k_val - scaM_next_val) + matKbar_k_val = (matK_k_val * vecAbar_k_val[None, :]).to(DTYPE) + + matC_k_val = scaGbar_k_val * matC_k_val + tl.dot(matKbar_k_val, matV_k_val) + vecN_k_val = scaGbar_k_val * vecN_k_val + tl.sum( + matKbar_k_val.to(tl.float32), axis=1 + ) + scaM_k_val = scaM_next_val + + +@triton.jit +def _parallel_fw_H_kernel( + matQ, # (B, NH, S, DHQK) + matK, # (B, NH, S, DHQK) + matV, # (B, NH, S, DHHV) + matC_states, # (B, NH, NC * DHQK, DHHV) + vecN_states, # (B, NH, NC * DHQK) + scaM_states, # (B, NH, NC) + vecI, # (B, NH, NC, L) + vecB, # (B, NH, NC, L) + matH, # (B, NH, S, DHHV) + qk_scale, + S: tl.constexpr, + DHQK: tl.constexpr, + DHHV: tl.constexpr, + NC: tl.constexpr, + L: tl.constexpr, + siz_b_DHQK: tl.constexpr, + siz_b_DHHV: tl.constexpr, + DTYPE: tl.constexpr, + EPS: tl.constexpr, +): + """Computes the outputs H of one chunk from the chunk's inputs and the state entering it.""" + idx_b_DHHV, idx_b_NC = tl.program_id(0), tl.program_id(1) + idx_b_BNH = tl.program_id(2).to(tl.int64) + offs_t = idx_b_NC * L + tl.arange(0, L) + offs_hv = idx_b_DHHV * siz_b_DHHV + tl.arange(0, siz_b_DHHV) + mask_hv = (offs_t[:, None] < S) & (offs_hv[None, :] < DHHV) + + vecBI_offset = idx_b_BNH * S + idx_b_NC * L + vecB_val = tl.load(vecB + vecBI_offset + tl.arange(0, L)).to(tl.float32) + vecI_val = tl.load(vecI + vecBI_offset + tl.arange(0, L)).to(tl.float32) + scaM_km1_val = tl.load(scaM_states + idx_b_BNH * NC + idx_b_NC).to(tl.float32) + + # gate matrix matDbar (L, L) and stabilizer vecM_combine (L,) + idx_mask = tl.arange(0, L) + mask = idx_mask[:, None] >= idx_mask[None, :] + matD_val = tl.where( + mask, vecB_val[:, None] - vecB_val[None, :] + vecI_val[None, :], -float("inf") + ) + vecM_combine_val = tl.maximum(vecB_val + scaM_km1_val, tl.max(matD_val, axis=1)) + matDbar_val = tl.exp(matD_val - vecM_combine_val[:, None]) + vecBbar_val = tl.exp(vecB_val + scaM_km1_val - vecM_combine_val) + + matS_val = tl.zeros((L, L), dtype=tl.float32) + matH_inter_val = tl.zeros((L, siz_b_DHHV), dtype=tl.float32) + vecH_inter_denom_val = tl.zeros((L,), dtype=tl.float32) + for idx_b_DHQK in range(tl.cdiv(DHQK, siz_b_DHQK)): + offs_qk = idx_b_DHQK * siz_b_DHQK + tl.arange(0, siz_b_DHQK) + mask_q = (offs_t[:, None] < S) & (offs_qk[None, :] < DHQK) + mask_k = (offs_qk[:, None] < DHQK) & (offs_t[None, :] < S) + mask_c = (offs_qk[:, None] < DHQK) & (offs_hv[None, :] < DHHV) + matQ_ptr = ( + matQ + idx_b_BNH * S * DHQK + offs_t[:, None] * DHQK + offs_qk[None, :] + ) + matK_ptr = ( + matK + idx_b_BNH * S * DHQK + offs_qk[:, None] + offs_t[None, :] * DHQK + ) + matC_km1_ptr = ( + matC_states + + (idx_b_BNH * NC + idx_b_NC) * DHQK * DHHV + + offs_qk[:, None] * DHHV + + offs_hv[None, :] + ) + vecN_km1_ptr = vecN_states + (idx_b_BNH * NC + idx_b_NC) * DHQK + offs_qk + + matQ_val = tl.load(matQ_ptr, mask=mask_q, other=0).to(DTYPE) + matK_val = tl.load(matK_ptr, mask=mask_k, other=0).to(DTYPE) + matS_val += (tl.dot(matQ_val, matK_val) * qk_scale).to(DTYPE) + + matQbar_val = (matQ_val * vecBbar_val[:, None] * qk_scale).to(DTYPE) + matC_km1_val = tl.load(matC_km1_ptr, mask=mask_c, other=0).to(DTYPE) + matH_inter_val += tl.dot(matQbar_val, matC_km1_val) + vecN_km1_val = tl.load(vecN_km1_ptr, mask=offs_qk < DHQK, other=0).to(DTYPE) + vecH_inter_denom_val += tl.sum( + (matQbar_val * vecN_km1_val[None, :]).to(tl.float32), axis=1 + ) + + matSbar_val = (matS_val * matDbar_val).to(DTYPE) + matV_ptr = matV + idx_b_BNH * S * DHHV + offs_t[:, None] * DHHV + offs_hv[None, :] + matV_val = tl.load(matV_ptr, mask=mask_hv, other=0).to(DTYPE) + matH_intra_val = tl.dot(matSbar_val, matV_val) + vecH_intra_denom_val = tl.sum(matSbar_val.to(tl.float32), axis=1) + + matH_num_val = matH_inter_val + matH_intra_val + vecH_denom_val = tl.maximum( + tl.abs(vecH_inter_denom_val + vecH_intra_denom_val), tl.exp(-vecM_combine_val) + ) + matH_val = matH_num_val / (vecH_denom_val[:, None] + EPS) + + matH_ptr = matH + idx_b_BNH * S * DHHV + offs_t[:, None] * DHHV + offs_hv[None, :] + tl.store(matH_ptr, matH_val.to(DTYPE), mask=mask_hv) + + +_TRITON_DTYPES = { + torch.float32: tl.float32, + torch.float16: tl.float16, + torch.bfloat16: tl.bfloat16, +} + + +def mlstm_chunkwise_triton( + q: torch.Tensor, # (B, NH, S, DHQK) + k: torch.Tensor, # (B, NH, S, DHQK) + v: torch.Tensor, # (B, NH, S, DHHV) + i: torch.Tensor, # (B, NH, S) input gate preactivation + f: torch.Tensor, # (B, NH, S) forget gate preactivation + chunk_size: int = 64, + eps: float = 1e-6, + autocast_kernel_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: # (B, NH, S, DHHV) + """mLSTM over a full sequence from zero states, with forward-only Triton kernels (CUDA).""" + # The triton kernels have no backward, so fall back to the native implementation when gradients are needed. + if torch.is_grad_enabled() and any(x.requires_grad for x in (q, k, v, i, f)): + return mlstm_chunkwise_native(q, k, v, i, f, chunk_size=chunk_size, eps=eps) + + S = q.shape[2] + q, k, v, i, f = pad_to_chunks(q, k, v, i, f, chunk_size) + # Under autocast, the inputs are cast and the kernel runs without autocast (upstream's custom_fwd). + if torch.is_autocast_enabled("cuda"): + q, k, v, i, f = (x.to(autocast_kernel_dtype) for x in (q, k, v, i, f)) + with torch.autocast("cuda", enabled=False): + q, k, v, i, f = (x.contiguous() for x in (q, k, v, i, f)) + B, NH, S_pad, DHQK = q.shape + DHHV = v.shape[-1] + NC, L = S_pad // chunk_size, chunk_size + assert L & (L - 1) == 0, "Chunk size must be a power of 2." + + vecI = i.reshape(B, NH, NC, L) + vecB = F.logsigmoid(f.reshape(B, NH, NC, L).to(torch.float32)).cumsum(-1) + + siz_b_DHQK = min(64, triton.next_power_of_2(DHQK)) + siz_b_DHHV = min(64, triton.next_power_of_2(DHHV)) + num_warps = 4 if siz_b_DHQK == 64 else 2 + dims = dict( + S=S_pad, + DHQK=DHQK, + DHHV=DHHV, + NC=NC, + L=L, + siz_b_DHQK=siz_b_DHQK, + siz_b_DHHV=siz_b_DHHV, + ) + DTYPE = _TRITON_DTYPES[q.dtype] + + matC_states = torch.empty( + B, NH, NC * DHQK, DHHV, device=q.device, dtype=torch.float32 + ) + vecN_states = torch.empty( + B, NH, NC * DHQK, device=q.device, dtype=torch.float32 + ) + scaM_states = torch.empty(B, NH, NC, device=q.device, dtype=torch.float32) + grid = (triton.cdiv(DHQK, siz_b_DHQK), triton.cdiv(DHHV, siz_b_DHHV), B * NH) + _recurrent_fw_C_kernel[grid]( + k, v, vecB, vecI, matC_states, vecN_states, scaM_states, + **dims, DTYPE=DTYPE, num_stages=1, num_warps=num_warps, + ) # fmt: skip + + matH = torch.empty(B, NH, S_pad, DHHV, device=q.device, dtype=q.dtype) + grid = (triton.cdiv(DHHV, siz_b_DHHV), NC, B * NH) + _parallel_fw_H_kernel[grid]( + q, k, v, matC_states, vecN_states, scaM_states, vecI, vecB, matH, DHQK**-0.5, + **dims, DTYPE=DTYPE, EPS=eps, num_stages=1, num_warps=num_warps, + ) # fmt: skip + return matH[:, :, :S, :] diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/kernel/slstm_torch.py b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/kernel/slstm_torch.py new file mode 100644 index 0000000000000..17acbdebe2718 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/kernel/slstm_torch.py @@ -0,0 +1,133 @@ +# 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 (c) NXAI GmbH. +# SPDX-License-Identifier: Apache-2.0 + +"""PyTorch sLSTM forward, used for the recurrence on CPU and MPS. + +Vendored from ``slstm_kernels/torch_kernel/slstm_torch.py`` in +https://github.com/NX-AI/slstm-kernels at commit ``b2bf942``. Keep the two in +sync; changes belong upstream first. Like the Triton kernel it keeps the states +in float32 and rounds only the hidden state to bfloat16 for the recurrent matmul. +""" + +import torch +import torch.nn.functional as F + +NG_GATES = 4 + + +def check_tensors(x, states, weight, bias, cu_seqlens): + Bx, S, NH, NG, HD = x.shape + B = states.shape[0] + + assert states.shape == ( + B, + NH, + NG, + HD, + ), f"states must have shape [B, NH, NG, HD] = {(B, NH, NG, HD)}, got {tuple(states.shape)}!" + assert weight.shape == ( + NH, + NG, + HD, + HD, + ), f"weight must have shape [NH, NG, HD, HD] = {(NH, NG, HD, HD)}, got {tuple(weight.shape)}!" + assert bias.shape == ( + NH, + NG, + HD, + ), f"bias must have shape [NH, NG, HD] = {(NH, NG, HD)}, got {tuple(bias.shape)}!" + + assert NG == NG_GATES, f"slstm requires NG={NG_GATES} gates (i, f, z, o), got {NG}!" + + assert x.dtype in (torch.bfloat16, torch.float32), "x must be bfloat16!" + assert states.dtype == torch.float32, "states must be float32!" + assert weight.dtype in (torch.bfloat16, torch.float32), "weight must be bfloat16!" + assert bias.dtype == torch.float32, "bias must be float32!" + + assert x.is_contiguous(), "x must be contiguous!" + assert states.is_contiguous(), "states must be contiguous!" + assert weight.is_contiguous(), "weight must be contiguous!" + assert bias.is_contiguous(), "bias must be contiguous!" + + assert states.device == x.device, "states must be on the same device as x!" + assert weight.device == x.device, "weight must be on the same device as x!" + assert bias.device == x.device, "bias must be on the same device as x!" + + if cu_seqlens is None: + assert Bx == B, "x and states must have the same batch size!" + else: + assert cu_seqlens.shape == (B + 1,), "cu_seqlens must have shape [B+1]!" + assert cu_seqlens.dtype == torch.int64, "cu_seqlens must be int64!" + assert ( + cu_seqlens.device == x.device + ), "cu_seqlens must be on the same device as x!" + assert cu_seqlens.is_contiguous(), "cu_seqlens must be contiguous!" + + +def slstm_torch( + x: torch.Tensor, # [B, S, NH, NG, HD] + states: torch.Tensor, # [B, NH, NG, HD] + weight: torch.Tensor, # [NH, NG, HD, HD] + bias: torch.Tensor, # [NH, NG, HD] + cu_seqlens: torch.Tensor = None, # [B+1] +) -> tuple[torch.Tensor, torch.Tensor]: # [B, S, NH, HD], [B, NH, NG, HD] + check_tensors(x, states, weight, bias, cu_seqlens) + assert cu_seqlens is None, "cu_seqlens isn't implementent for torch backend!" + B, S, NH, NG, HD = x.shape + assert NG == NG_GATES, f"torch backend requires NG={NG_GATES} gates, got {NG}!" + assert states.shape == (B, NH, NG, HD) + states = states.unbind(dim=-2) + + w_t = weight.view(NH, NG * HD, HD).transpose(-1, -2).float() + output = [] + for i in range(S): + y_bf16 = states[0].bfloat16().float() + Ry = torch.einsum("BHE,HEO->BHO", y_bf16, w_t).view(B, NH, NG, HD) + states = slstm_torch_pointwise( + x[:, i].float(), Ry.float(), bias.unsqueeze(0), states + ) + output.append(states[0].to(x.dtype)) + + return torch.stack(output, dim=1), torch.stack(states, dim=-2) + + +def slstm_torch_pointwise( + Wx: torch.Tensor, # [B, NH, NG, HD] + Ry: torch.Tensor, # [B, NH, NG, HD] + b: torch.Tensor, # [1, NH, NG, HD] + states: list[torch.Tensor], # NG x [B, NH, HD] +) -> list[torch.Tensor]: + raw = Wx + Ry + b + iraw, fraw, zraw, oraw = raw.unbind(dim=-2) + y, c, n, m = states + + # Equations reference the xlstm paper on page 4: https://arxiv.org/pdf/2405.04517 + logfplusm = m + F.logsigmoid(fraw.clamp(max=15)) # eq 15 # Clamp to avoid subnomals + mnew = torch.where(n == 0.0, iraw, torch.max(iraw, logfplusm)) # eq 15 + ogate = torch.sigmoid(oraw) # eq 14 + igate = torch.exp(iraw - mnew).clamp(max=1) # eq 16 + fgate = torch.exp(logfplusm - mnew).clamp(max=1) # eq 17 + zgate = torch.tanh(zraw) # eq 11 + cnew = fgate * c + igate * zgate # eq 8 + # TODO: Why do wee need the maximum here? + nnew = (fgate * n + igate).clamp(min=1) # eq 9 + hnew = ogate * cnew / nnew # eq 10 + + return [hnew, cnew, nnew, mnew] # NG x [B, NH, HD] diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/kernel/slstm_triton.py b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/kernel/slstm_triton.py new file mode 100644 index 0000000000000..d3fc23fec0788 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/kernel/slstm_triton.py @@ -0,0 +1,297 @@ +# 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 (c) NXAI GmbH. +# SPDX-License-Identifier: Apache-2.0 + +"""Triton sLSTM forward kernel, used for the recurrence on CUDA. + +Vendored from ``slstm_kernels/triton_kernel/slstm_triton.py`` in +https://github.com/NX-AI/slstm-kernels at commit ``b2bf942``, with int64 batch +offsets. Keep the two in sync; changes to the kernel belong upstream first. + +Unlike FlashRNN's CUDA kernel, it keeps the cell, normalizer and stabilizer +states in float32 across time steps and launches on the current stream without +synchronizing, so the recurrence can be captured in a CUDA graph. Only the +recurrent matmul runs on bfloat16 operands. +""" + +from functools import cache + +import torch +import triton +import triton.language as tl +from triton.language.extra import libdevice + +from iotdb.ainode.core.model.tirex2.component.kernel.slstm_torch import slstm_torch + +NG_GATES = 4 + +LOG2E = tl.constexpr(1.4426950408889634) + +SMEM_HEADROOM = 4096 + + +@cache +def _max_shared_mem(device_index: int) -> int: + return triton.runtime.driver.active.utils.get_device_properties(device_index)[ + "max_shared_mem" + ] + + +def check_tensors(x, states, weight, bias, cu_seqlens): + Bx, S, NH, NG, HD = x.shape + B = states.shape[0] + + assert states.shape == ( + B, + NH, + NG, + HD, + ), f"states must have shape [B, NH, NG, HD] = {(B, NH, NG, HD)}, got {tuple(states.shape)}!" + assert weight.shape == ( + NH, + NG, + HD, + HD, + ), f"weight must have shape [NH, NG, HD, HD] = {(NH, NG, HD, HD)}, got {tuple(weight.shape)}!" + assert bias.shape == ( + NH, + NG, + HD, + ), f"bias must have shape [NH, NG, HD] = {(NH, NG, HD)}, got {tuple(bias.shape)}!" + + assert NG == NG_GATES, f"slstm requires NG={NG_GATES} gates (i, f, z, o), got {NG}!" + + assert x.dtype in (torch.bfloat16, torch.float32), "x must be bfloat16!" + assert states.dtype == torch.float32, "states must be float32!" + assert weight.dtype in (torch.bfloat16, torch.float32), "weight must be bfloat16!" + assert bias.dtype == torch.float32, "bias must be float32!" + + assert x.is_contiguous(), "x must be contiguous!" + assert states.is_contiguous(), "states must be contiguous!" + assert weight.is_contiguous(), "weight must be contiguous!" + assert bias.is_contiguous(), "bias must be contiguous!" + + assert states.device == x.device, "states must be on the same device as x!" + assert weight.device == x.device, "weight must be on the same device as x!" + assert bias.device == x.device, "bias must be on the same device as x!" + + if cu_seqlens is None: + assert Bx == B, "x and states must have the same batch size!" + else: + assert cu_seqlens.shape == (B + 1,), "cu_seqlens must have shape [B+1]!" + assert cu_seqlens.dtype == torch.int64, "cu_seqlens must be int64!" + assert ( + cu_seqlens.device == x.device + ), "cu_seqlens must be on the same device as x!" + assert cu_seqlens.is_contiguous(), "cu_seqlens must be contiguous!" + + +@triton.jit +def slstm_pointwise(iraw, fraw, zraw, oraw, c, n, m): + logfplusm = m - tl.log(1.0 + tl.exp(-tl.minimum(fraw, 15.0))) + mnew = tl.where(n == 0.0, iraw, tl.maximum(iraw, logfplusm)) + igate = tl.minimum(tl.exp2(LOG2E * (iraw - mnew)), 1.0) + fgate = tl.minimum(tl.exp2(LOG2E * (logfplusm - mnew)), 1.0) + zgate = libdevice.tanh(zraw) + cnew = fgate * c + igate * zgate + nnew = tl.maximum((fgate * n + igate), 1.0) + hnew = cnew / (nnew * (1.0 + tl.exp2(-LOG2E * oraw))) + return hnew, cnew, nnew, mnew + + +@triton.autotune( + configs=[ + triton.Config({"B_BLOCK_SIZE": 2**i}, num_warps=w, num_stages=1) + for i in range(6) + for w in (4, 8, 16) + ], + key=["B_POW2", "NH", "NG", "HD", "RESIDENT_WEIGHTS"], + prune_configs_by={ + "early_config_prune": lambda configs, nargs, **kwargs: [ + c for c in configs if c.kwargs["B_BLOCK_SIZE"] <= nargs["B_POW2"] + ] + }, +) +@triton.jit +def slstm_kernel( + x_ptr, + states_ptr, + weight_ptr, + bias_ptr, + cu_seqlens_ptr, + out, + out_states, + B, + S, + B_POW2, + NH: tl.constexpr, + NG: tl.constexpr, + HD: tl.constexpr, + IS_VARLEN: tl.constexpr, + RESIDENT_WEIGHTS: tl.constexpr, + B_BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + r = tl.arange(0, HD)[None, :] + head_id = pid % NH + # int64, as the input/output offsets derived from it exceed int32 beyond ~14k rows of 74 steps + batch_id = ((pid // NH) * B_BLOCK_SIZE + tl.arange(0, B_BLOCK_SIZE)[:, None]).to( + tl.int64 + ) + batch_mask = batch_id < B + + if IS_VARLEN: + seq_start = tl.load(cu_seqlens_ptr + batch_id, mask=batch_mask, other=0) + seq_end = tl.load(cu_seqlens_ptr + batch_id + 1, mask=batch_mask, other=0) + seq_len = (seq_end - seq_start).to(tl.int32) + # Rows of the block run in lockstep, so step through the longest one and mask + num_steps = tl.max(tl.where(batch_mask, seq_len, 0)) + x_mask = batch_mask & (seq_len > 0) + else: + seq_start = batch_id * S + seq_len = S + num_steps = S + x_mask = batch_mask + + state_offs = batch_id * NH * NG * HD + head_id * NG * HD + r + y = tl.load(states_ptr + state_offs + 0 * HD, mask=batch_mask, other=0.0) + c = tl.load(states_ptr + state_offs + 1 * HD, mask=batch_mask, other=0.0) + n = tl.load(states_ptr + state_offs + 2 * HD, mask=batch_mask, other=0.0) + m = tl.load(states_ptr + state_offs + 3 * HD, mask=batch_mask, other=0.0) + + # weight arrives transposed, [NH, NG, HD_in, HD_out], so each gate is a row-major tile + w_ptrs = weight_ptr + head_id * NG * HD * HD + tl.arange(0, HD)[:, None] * HD + r + if RESIDENT_WEIGHTS: + w_i = tl.load(w_ptrs + 0 * HD * HD) + w_f = tl.load(w_ptrs + 1 * HD * HD) + w_z = tl.load(w_ptrs + 2 * HD * HD) + w_o = tl.load(w_ptrs + 3 * HD * HD) + + bias_offs = head_id * NG * HD + b_i = tl.load(bias_ptr + bias_offs + 0 * HD + r) + b_f = tl.load(bias_ptr + bias_offs + 1 * HD + r) + b_z = tl.load(bias_ptr + bias_offs + 2 * HD + r) + b_o = tl.load(bias_ptr + bias_offs + 3 * HD + r) + + x_base = x_ptr + seq_start * NH * NG * HD + head_id * NG * HD + r + out_base = out + seq_start * NH * HD + head_id * HD + r + + p_i = tl.load(x_base + 0 * HD, mask=x_mask, other=0.0) + p_f = tl.load(x_base + 1 * HD, mask=x_mask, other=0.0) + p_z = tl.load(x_base + 2 * HD, mask=x_mask, other=0.0) + p_o = tl.load(x_base + 3 * HD, mask=x_mask, other=0.0) + + for step_id in range(num_steps): + x_i, x_f, x_z, x_o = p_i, p_f, p_z, p_o # double buffer + + # clamped to the row's own last step in cu_seqlen mode + next_offs = x_base + tl.minimum(step_id + 1, seq_len - 1) * NH * NG * HD + p_i = tl.load(next_offs + 0 * HD, mask=x_mask, other=0.0) + p_f = tl.load(next_offs + 1 * HD, mask=x_mask, other=0.0) + p_z = tl.load(next_offs + 2 * HD, mask=x_mask, other=0.0) + p_o = tl.load(next_offs + 3 * HD, mask=x_mask, other=0.0) + + y_bf16 = y.to(tl.bfloat16) + if RESIDENT_WEIGHTS: + iraw = tl.dot(y_bf16, w_i, acc=x_i.to(tl.float32) + b_i) + fraw = tl.dot(y_bf16, w_f, acc=x_f.to(tl.float32) + b_f) + zraw = tl.dot(y_bf16, w_z, acc=x_z.to(tl.float32) + b_z) + oraw = tl.dot(y_bf16, w_o, acc=x_o.to(tl.float32) + b_o) + else: + # Re-read from L2 every step + iraw = tl.dot( + y_bf16, tl.load(w_ptrs + 0 * HD * HD), acc=x_i.to(tl.float32) + b_i + ) + fraw = tl.dot( + y_bf16, tl.load(w_ptrs + 1 * HD * HD), acc=x_f.to(tl.float32) + b_f + ) + zraw = tl.dot( + y_bf16, tl.load(w_ptrs + 2 * HD * HD), acc=x_z.to(tl.float32) + b_z + ) + oraw = tl.dot( + y_bf16, tl.load(w_ptrs + 3 * HD * HD), acc=x_o.to(tl.float32) + b_o + ) + + ynew, cnew, nnew, mnew = slstm_pointwise(iraw, fraw, zraw, oraw, c, n, m) + + if IS_VARLEN: + # Freeze a row on its final state + step_mask = batch_mask & (step_id < seq_len) + y = tl.where(step_mask, ynew, y) + c = tl.where(step_mask, cnew, c) + n = tl.where(step_mask, nnew, n) + m = tl.where(step_mask, mnew, m) + else: + step_mask = batch_mask + y, c, n, m = ynew, cnew, nnew, mnew + + tl.store(out_base + step_id * NH * HD, y, mask=step_mask) + + tl.store(out_states + state_offs + 0 * HD, y, mask=batch_mask) + tl.store(out_states + state_offs + 1 * HD, c, mask=batch_mask) + tl.store(out_states + state_offs + 2 * HD, n, mask=batch_mask) + tl.store(out_states + state_offs + 3 * HD, m, mask=batch_mask) + + +def slstm_triton_fw( + x: torch.Tensor, # [B, S, NH, NG, HD] + states: torch.Tensor, # [B, NH, NG, HD] + weight: torch.Tensor, # [NH, NG, HD, HD] + bias: torch.Tensor, # [NH, NG, HD] + cu_seqlens: torch.Tensor = None, # [B+1] +) -> tuple[torch.Tensor, torch.Tensor]: # [B, S, NH, HD], [B, NH, NG, HD] + if torch.is_grad_enabled() and any( + t.requires_grad for t in (x, states, weight, bias) + ): + return slstm_torch(x, states, weight, bias, cu_seqlens) + + check_tensors(x, states, weight, bias, cu_seqlens) + Bx, S, NH, NG, HD = x.shape + B = states.shape[0] + assert NG == NG_GATES, f"triton backend requires NG={NG_GATES} gates, got {NG}!" + # and transposed so the kernel loads each gate as the [HD_in, HD_out] dot operand + weight = weight.to(torch.bfloat16).transpose(-1, -2).contiguous() + resident_weights = ( + NG * HD * HD * weight.element_size() + SMEM_HEADROOM + <= _max_shared_mem(x.device.index) + ) + + out = torch.empty((Bx, S, NH, HD), dtype=x.dtype, device=x.device) + out_states = torch.empty_like(states) + + grid = lambda meta: (triton.cdiv(B, meta["B_BLOCK_SIZE"]) * NH,) + slstm_kernel[grid]( + x, + states, + weight, + bias, + cu_seqlens, + out, + out_states, + B, + S, + triton.next_power_of_2(B), + NH, + NG, + HD, + cu_seqlens is not None, + resident_weights, + ) + + return out, out_states diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/mlstm_block.py b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/mlstm_block.py new file mode 100644 index 0000000000000..88df4dfe48535 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/mlstm_block.py @@ -0,0 +1,222 @@ +# 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 (c) NXAI GmbH. +# SPDX-License-Identifier: Apache-2.0 + +"""mLSTM layer copied from xLSTM Large, with TiRex-specific tweaks.""" + +from dataclasses import dataclass +from typing import Literal + +import torch +import torch.nn as nn + +from iotdb.ainode.core.model.tirex2.component import kernel +from iotdb.ainode.core.model.tirex2.component.norm import MultiHeadLayerNorm + +# soft_cap, bias_linspace_init_ and the mLSTMLayerConfig fields are copied from xlstm. + + +def soft_cap( + values: torch.Tensor, cap_value: float | torch.Tensor | None +) -> torch.Tensor: + """Soft caps a tensor to cap_value with a scaled tanh; no cap when cap_value is None.""" + if cap_value is None: + return values + return cap_value * torch.tanh(values / cap_value) + + +def bias_linspace_init_( + param: torch.Tensor, start: float = 3.4, end: float = 6.0 +) -> torch.Tensor: + """Linearly spaced bias init across dimensions.""" + assert ( + param.dim() == 1 + ), f"param must be 1-dimensional (typically a bias), got {param.dim()}" + n_dims = param.shape[0] + init_vals = torch.linspace(start, end, n_dims) + with torch.no_grad(): + param.copy_(init_vals) + return param + + +@dataclass +class mLSTMLayerConfig: + """Configuration of the xLSTM Large mLSTM layer.""" + + embedding_dim: int + num_heads: int + use_bias: bool = False + norm_eps: float = 1e-6 + qk_dim_factor: float = 0.5 + v_dim_factor: float = 1.0 + gate_soft_cap: float = 15.0 + # "auto" runs the Triton kernel for CUDA inputs when Triton is available and the native one otherwise. + mlstm_kernel: Literal["auto", "native", "triton"] = "auto" + + +class mLSTMLayer(nn.Module): + """mLSTM implementation copied from xLSTM 7B.""" + + def __init__(self, config: mLSTMLayerConfig): + super().__init__() + self.config = config + + self.v_dim = int(config.embedding_dim * config.v_dim_factor) + self.qk_dim = int(config.embedding_dim * config.qk_dim_factor) + + # Fused input projection of q, k, v and the o, i, f gate pre-activations, in that order. + # The i and f gates always have a bias; without use_bias the rest of the bias starts at zero. + self.in_proj_split = ( + self.qk_dim, + self.qk_dim, + self.v_dim, + self.v_dim, + self.config.num_heads, + self.config.num_heads, + ) + self.in_proj = nn.Linear( + self.config.embedding_dim, sum(self.in_proj_split), bias=True + ) + if not self.config.use_bias: + nn.init.zeros_(self.in_proj.bias) + + self.ogate_act_fn = nn.Sigmoid() + if config.mlstm_kernel not in ("auto", "native", "triton"): + raise ValueError( + f"mLSTM kernel must be 'auto', 'native' or 'triton', got {config.mlstm_kernel!r}." + ) + + self.multihead_norm = MultiHeadLayerNorm( + num_heads=self.config.num_heads, + head_dim=self.v_dim // self.config.num_heads, + eps=self.config.norm_eps, + ) + self.out_proj = nn.Linear( + in_features=self.v_dim, + out_features=self.config.embedding_dim, + bias=self.config.use_bias, + ) + + # Match mLSTMBlock.reset_parameters gate initialisation; i and f are the last rows of in_proj. + i_rows = slice(-2 * config.num_heads, -config.num_heads) + f_rows = slice(-config.num_heads, None) + with torch.no_grad(): + torch.nn.init.zeros_(self.in_proj.weight[f_rows]) + bias_linspace_init_(self.in_proj.bias[f_rows], start=3.0, end=6.0) + torch.nn.init.zeros_(self.in_proj.weight[i_rows]) + torch.nn.init.normal_(self.in_proj.bias[i_rows], mean=0.0, std=0.1) + + def _load_from_state_dict( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ): + # Older checkpoints store a separate linear layer per projection + legacy = ("q", "k", "v", "ogate_preact", "igate_preact", "fgate_preact") + if all(f"{prefix}{name}.weight" in state_dict for name in legacy): + # The i and f gates always have a bias, the other projections only with use_bias. + # A missing bias that the checkpoint must have is reported instead of zero-filled. + required = legacy if self.config.use_bias else legacy[-2:] + missing_keys.extend( + f"{prefix}{name}.bias" + for name in required + if f"{prefix}{name}.bias" not in state_dict + ) + weights = [state_dict.pop(f"{prefix}{name}.weight") for name in legacy] + biases = [state_dict.pop(f"{prefix}{name}.bias", None) for name in legacy] + biases = [ + w.new_zeros(w.shape[0]) if b is None else b + for w, b in zip(weights, biases) + ] + state_dict[f"{prefix}in_proj.weight"] = torch.cat(weights) + state_dict[f"{prefix}in_proj.bias"] = torch.cat(biases) + super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) + + def _kernel(self, x: torch.Tensor): + """The chunkwise mLSTM kernel for inputs on the device of ``x``.""" + use_triton = self.config.mlstm_kernel == "triton" or ( + self.config.mlstm_kernel == "auto" and kernel.use_triton(x) + ) + # Imported lazily, as triton is not available on every platform. + if use_triton: + from iotdb.ainode.core.model.tirex2.component.kernel.mlstm_triton import ( + mlstm_chunkwise_triton, + ) + + return mlstm_chunkwise_triton + from iotdb.ainode.core.model.tirex2.component.kernel.mlstm_native import ( + mlstm_chunkwise_native, + ) + + return mlstm_chunkwise_native + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Process a full sequence through the mlstm.""" + assert x.ndim == 3, f"Input must have shape [B, S, D], got {x.shape}" + B, S, _ = x.shape + + q, k, v, o_preact, i_preact, f_preact = self.in_proj(x).split( + self.in_proj_split, dim=-1 + ) + i_preact = soft_cap(i_preact, cap_value=self.config.gate_soft_cap) + f_preact = soft_cap(f_preact, cap_value=self.config.gate_soft_cap) + + q = q.reshape(B, S, self.config.num_heads, -1).transpose(1, 2) + k = k.reshape(B, S, self.config.num_heads, -1).transpose(1, 2) + v = v.reshape(B, S, self.config.num_heads, -1).transpose(1, 2) + i_preact = i_preact.transpose(1, 2) + f_preact = f_preact.transpose(1, 2) + h = self._kernel(q)( + q=q, + k=k, + v=v, + i=i_preact, + f=f_preact, + ) + expected_h_shape = ( + B, + self.config.num_heads, + S, + self.v_dim // self.config.num_heads, + ) + assert ( + h.shape == expected_h_shape + ), f"Got {h.shape}, expected {expected_h_shape}" + + h = h.transpose(1, 2) + h_norm = self.multihead_norm(h) + h_norm = h_norm.reshape(B, S, -1) + + h_out = self.ogate_act_fn(o_preact) * h_norm + + y = self.out_proj(h_out) + return y diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/norm.py b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/norm.py new file mode 100644 index 0000000000000..cc47488e4ef06 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/norm.py @@ -0,0 +1,62 @@ +# 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 (c) NXAI GmbH. +# SPDX-License-Identifier: Apache-2.0 + +import torch +from torch import nn + + +class RMSNorm(nn.Module): + def __init__(self, num_features: int, eps: float = 1e-6): + super().__init__() + self.num_features = num_features + self.eps = eps + self.weight = nn.Parameter(torch.ones(num_features)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + in_dtype = x.dtype + x = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps) + return self.weight * x.to(in_dtype) + + +class MultiHeadLayerNorm(nn.Module): + def __init__(self, num_heads: int, head_dim: int, eps: float = 1e-6): + super().__init__() + self.num_heads = num_heads + self.head_dim = head_dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(num_heads * head_dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + B, S, NH, DH = x.shape + assert ( + NH == self.num_heads + ), f"Expected {self.num_heads} heads, got {NH}, input shape: {x.shape}" + assert ( + DH == self.head_dim + ), f"Expected {self.head_dim} head dimension, got {DH}, input shape: {x.shape}" + + in_dtype = x.dtype + x = x.float() + x_centered = x - x.mean(dim=-1, keepdim=True) + x = x_centered * torch.rsqrt( + x.var(dim=-1, keepdim=True, unbiased=False) + self.eps + ) + x = x.to(in_dtype).reshape(B, S, -1) + return self.weight * x diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/patch_tokenizer.py b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/patch_tokenizer.py new file mode 100644 index 0000000000000..4a65e76626781 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/patch_tokenizer.py @@ -0,0 +1,115 @@ +# 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 (c) NXAI GmbH. +# SPDX-License-Identifier: Apache-2.0 + +"""Patch-based tokenization utilities for TiRex models.""" + +import torch +from torch import nn + + +class Tokenizer: + """Tokenizer that applies sliding-window patching with optional left padding. + + Parameters + ---------- + output_patch_size : int + Number of timesteps produced per output patch. + input_patch_size : int + Window size used when unfolding the input sequence. + input_patch_stride : int + Stride between successive input patches. + """ + + def __init__( + self, + output_patch_size: int, + input_patch_size: int, + input_patch_stride: int, + ): + super().__init__() + self.output_patch_size = output_patch_size + self.input_patch_size = input_patch_size + self.input_patch_stride = input_patch_stride + self.patcher = Patch( + self.input_patch_size, self.input_patch_stride, left_pad=True + ) + + def input_transform(self, data: torch.Tensor): + """Patchify the 2D tensor and record patch count in state.""" + assert data.ndim == 2 + data = self.patcher(data) + return data, {"len": data.shape[-2]} + + def output_transform(self, data: torch.Tensor, tokenizer_state: dict): + """Reshape patched outputs back to the expected temporal layout.""" + data = torch.reshape( + data, (data.shape[0], -1, self.output_patch_size * tokenizer_state["len"]) + ) + return data + + +class Patch(nn.Module): + """Utility module that unfolds 1D tensors into patches with optional padding. + + Parameters + ---------- + patch_size : int + Size of each extracted patch. + patch_stride : int + Step between adjacent patches; must divide ``patch_size``. + left_pad : bool + If ``True``, pad on the left; otherwise pad on the right. + """ + + def __init__( + self, patch_size: int, patch_stride: int, left_pad: bool = False + ) -> None: + super().__init__() + self.patch_size = patch_size + self.patch_stride = patch_stride + self.left_pad = left_pad + assert self.patch_size % self.patch_stride == 0 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Pad if needed and unfold the last dimension into patches.""" + assert x.ndim == 2 + length = x.shape[-1] + + if length < self.patch_size or (length % self.patch_stride != 0): + if length < self.patch_size: + padding_size = ( + *x.shape[:-1], + self.patch_size - (length % self.patch_size), + ) + else: + padding_size = ( + *x.shape[:-1], + self.patch_stride - (length % self.patch_stride), + ) + padding = torch.full( + size=padding_size, fill_value=torch.nan, dtype=x.dtype, device=x.device + ) + if self.left_pad: + x = torch.concat((padding, x), dim=-1) + else: + x = torch.concat((x, padding), dim=-1) + + x = x.unfold(dimension=-1, size=self.patch_size, step=self.patch_stride) + return x diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/postprocessor.py b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/postprocessor.py new file mode 100644 index 0000000000000..5c2f16d265521 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/postprocessor.py @@ -0,0 +1,425 @@ +# 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 (c) NXAI GmbH. +# SPDX-License-Identifier: Apache-2.0 + +"""Single-pass, native-multivariate postprocessor with per-target differencing.""" + +import logging +from dataclasses import dataclass + +import numpy as np +import torch +from torch.nn.functional import pad + +logger = logging.getLogger(__name__) + + +@dataclass +class PostProcessorConfig: + """Postprocessor defaults for trend detection and band calibration. + + These constants are the differencing/calibration values that used to be + serialized in ``model/model-config.yaml``. Checkpoint configs now carry only + ``tta_diff`` to enable or disable the differencing path at inference time. + """ + + r2: float = 0.87 + trend_min_r2: float = 0.87 + trend_threshold: float = 1.5 + sigma: float = 1.5 + trend_window: float = 1.0 + diff_band_exponent: float = 0.51 + diff_band_scale: float = 1.68 + raw_band_exponent: float = 0.52 + raw_band_scale: float = 0.90 + + def __post_init__(self): + if not 0.0 < self.trend_window <= 1.0: + raise ValueError( + f"trend_window must be a fraction in (0, 1], got {self.trend_window!r}." + ) + if not self.diff_band_exponent >= 0.0: + raise ValueError( + f"diff_band_exponent must be >= 0, got {self.diff_band_exponent!r}." + ) + if not self.diff_band_scale > 0.0: + raise ValueError( + f"diff_band_scale must be strictly positive, got {self.diff_band_scale!r}." + ) + if not self.raw_band_exponent >= 0.0: + raise ValueError( + f"raw_band_exponent must be >= 0, got {self.raw_band_exponent!r}." + ) + if not self.raw_band_scale > 0.0: + raise ValueError( + f"raw_band_scale must be strictly positive, got {self.raw_band_scale!r}." + ) + if not (self.trend_min_r2 > 0.0 and self.trend_min_r2 <= 1.0): + raise ValueError( + f"trend_min_r2 ({self.trend_min_r2}) to be between 0 and 1." + ) + if not self.trend_threshold > 0.0: + raise ValueError( + f"trend_threshold has to be larger than 0. Got {self.trend_threshold}." + ) + + +class PostProcessor: + """Single-pass joint multivariate forecasting; trending targets are differenced.""" + + def __init__( + self, config: PostProcessorConfig | dict | None = None, *args, **kwargs + ) -> None: + if config is None: + self.cfg = PostProcessorConfig() + elif isinstance(config, dict): + self.cfg = PostProcessorConfig(**config) + else: + self.cfg = config + + def transform_input( + self, + target: list[torch.Tensor], + prediction_length: int, + past_covariates: list[torch.Tensor], + past_future_covariates: list[torch.Tensor], + *args, + tta_diff: bool = True, + **kwargs, + ): + """Pack each sample into one joint group and optionally difference trending targets.""" + # ``past_covariates`` / ``past_future_covariates`` are fed to the model + # but always on their raw scale - they are never differenced. + target = [t.unsqueeze(0) if t.ndim == 1 else t for t in target] + lengths = [t.shape[-1] for t in target for _ in range(t.shape[0])] + num_rows = [t.shape[0] for t in target] + max_len = max(lengths) + + # Pad if lengths of Targets is different + if all(length == max_len for length in lengths): + target_tensor = torch.cat(target, dim=0) + else: + target_tensor = torch.cat( + [pad(t, (max_len - t.shape[-1], 0), value=torch.nan) for t in target], + dim=0, + ) + + params = self._diff_params(target_tensor) + should_diff = ( + self._detect_trends(target_tensor, lengths) + if tta_diff + else torch.zeros(target_tensor.shape[0], dtype=torch.bool) + ) + + row_left_pad = torch.tensor( + [max_len - length for length in lengths], device=target_tensor.device + ) + diff_t = ( + self._diff_forward(target_tensor, row_left_pad) + if should_diff.any() + else target_tensor + ) + + target_rows = torch.where(should_diff[:, None], diff_t, target_tensor) + transform_params = [ + {"last_values": lv} for lv in params["last_values"].split(num_rows) + ] + diff_masks = [mask.tolist() for mask in should_diff.split(num_rows)] + + x = [] + group_vector = [] + target_mask = [] + group_sample_map = {} + group_target_indices = {} + group_row_is_diff = {} + + for sample_idx, (t_raw, sample_rows, diff_mask, pc, pfc) in enumerate( + zip( + target, + target_rows.split(num_rows), + diff_masks, + past_covariates, + past_future_covariates, + ) + ): + # One joint group per sample: its target rows (already left-padded to max_len, + # differenced where a trend was found) plus its raw covariates. + V_t = sample_rows.shape[0] + + pc_raw = ( + pc + if pc is not None + else torch.empty( + 0, t_raw.shape[-1], device=t_raw.device, dtype=t_raw.dtype + ) + ) + pfc_raw = ( + pfc + if pfc is not None + else torch.empty( + 0, + t_raw.shape[-1] + prediction_length, + device=t_raw.device, + dtype=t_raw.dtype, + ) + ) + if pc_raw.shape[-1] != t_raw.shape[-1]: + raise ValueError( + "Past covariates and targets have to have the same length " + f"(expected {t_raw.shape[-1]}, got {pc_raw.shape[-1]})." + ) + expected_future_covariate_length = t_raw.shape[-1] + prediction_length + if pfc_raw.shape[-1] < expected_future_covariate_length: + raise ValueError( + "Future known covariates must be at least as long as the target + the prediction length " + f"(expected at least {expected_future_covariate_length}, got {pfc_raw.shape[-1]})." + ) + if pfc_raw.shape[-1] > expected_future_covariate_length: + # Extra future-known covariate steps are harmless (e.g. when an + # overlong requested horizon is capped); keep only what the model + # consumes. + pfc_raw = pfc_raw[..., :expected_future_covariate_length] + V = V_t + pc_raw.shape[0] + pfc_raw.shape[0] + + next_group = (max(group_vector) if group_vector else 0) + 1 + group_sample_map[next_group] = sample_idx + group_target_indices[next_group] = list(range(V_t)) + group_row_is_diff[next_group] = list(diff_mask) + + left_pad = max_len - t_raw.shape[-1] + right_pad = prediction_length + item = torch.concat( + [ + pad(sample_rows, (0, right_pad), value=torch.nan), + pad(pc_raw, (left_pad, right_pad), value=torch.nan), + pad(pfc_raw, (left_pad, 0), value=torch.nan), + ], + dim=0, + ) + x.append(item) + + target_mask.extend([True] * V_t + [False] * (V - V_t)) + group_vector.extend([next_group] * V) + + device = target_tensor.device + x = torch.concat(x, dim=0) + group_vector = torch.tensor(group_vector, device=device, dtype=torch.float32) + target_mask = torch.tensor(target_mask, device=device, dtype=torch.bool) + + return ( + {"x": x, "group_vector": group_vector, "target_mask": target_mask}, + [], + { + "group_sample_map": group_sample_map, + "group_vector": group_vector, + "target_mask": target_mask, + "prediction_window_is_padded": True, + "single_pass": True, + "transform_params": transform_params, + "diff_masks": diff_masks, + "group_target_indices": group_target_indices, + "group_row_is_diff": group_row_is_diff, + }, + ) + + def transform_output( + self, + output: torch.Tensor, + *args, + group_sample_map: dict[int, int], + group_vector: torch.Tensor, + target_mask: torch.Tensor, + transform_params: list[dict], + diff_masks: list[list[bool]], + group_target_indices: dict[int, list[int]], + group_row_is_diff: dict[int, list[bool]], + preserve_grad: bool = False, + **kwargs, + ): + """Calibrate the band per row type and integrate differenced rows back.""" + if not preserve_grad: + output = output.detach().cpu() + group_vector = group_vector.to(output.device) + target_mask = target_mask.to(output.device) + + sample_rows = [[None] * len(mask) for mask in diff_masks] + group_id = int(group_vector[0].item()) + + while group_id <= group_vector.max().item(): + sample_idx = group_sample_map[group_id] + target_indices = group_target_indices[group_id] + mask = (group_vector == group_id) & target_mask + group_sample = output[mask] + + # ``row_is_diff`` is per target row: differenced and raw rows can mix. + row_is_diff = torch.tensor( + group_row_is_diff[group_id], device=output.device, dtype=torch.bool + ) + + # Raw-row band calibration: sqrt(raw_band_scale) * t^(raw_band_exponent-0.5) + # multiplies the model's per-step offset around its median. + not_diff = ~row_is_diff + if not_diff.any(): + q_med = group_sample.shape[1] // 2 + raw = group_sample[not_diff] + median_step = raw[:, q_med : q_med + 1, :] + offset = raw - median_step + t_idx = torch.arange( + 1, raw.shape[-1] + 1, device=raw.device, dtype=raw.dtype + ) + horizon_factor = t_idx ** (self.cfg.raw_band_exponent - 0.5) + rescaled = ( + median_step + + (self.cfg.raw_band_scale**0.5) * horizon_factor * offset + ) + group_sample = group_sample.clone() + group_sample[not_diff] = rescaled + + if row_is_diff.any(): + last_values = transform_params[sample_idx]["last_values"][ + target_indices + ].to(output.device) + if not preserve_grad: + last_values = last_values.detach() + group_sample = self._diff_inverse( + group_sample, + row_is_diff, + {"last_values": last_values}, + preserve_grad=preserve_grad, + ) + + for row, target_idx in zip(group_sample, target_indices): + sample_rows[sample_idx][target_idx] = row + + group_id += 1 + + return [torch.stack(rows, dim=0) for rows in sample_rows] + + def _detect_trends(self, x: torch.Tensor, lengths: list[int]) -> torch.Tensor: + """Per-row trend flags from an ``ols_sigma`` fit on the last ``trend_window`` fraction of each row. + + ``x`` is ``[rows, T]``, left-padded with NaN; ``lengths`` holds each row's length before padding. + Returns a bool tensor ``[rows]``. + """ + width = x.shape[-1] + windows = [max(1, round(self.cfg.trend_window * length)) for length in lengths] + windows = torch.tensor(windows, dtype=torch.float64, device=x.device) + + pos = ( + torch.arange(width, dtype=torch.float64, device=x.device)[None, :] + - (width - windows)[:, None] + ) + values = x.to(torch.float64) + valid = (pos >= 0) & ~torch.isnan(values) + score, r2, std, n = self._ols_sigma_score(values, valid, pos, windows) + return ( + (n >= 4) + & (std >= 1e-12) + & (score >= self.cfg.trend_threshold) + & (r2 >= self.cfg.trend_min_r2) + ) + + @staticmethod + def _ols_sigma_score( + values: torch.Tensor, + valid: torch.Tensor, + pos: torch.Tensor, + windows: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """``_ols_sigma_score`` for every row at once. + + values, pos: [rows, T] float64; valid: [rows, T] bool (in window and not NaN); windows: [rows]. + Returns score, r2, std and the number of valid values per row, each [rows]. + """ + n = valid.sum(-1, keepdim=True) # valid values per row + v = torch.where(valid, values, 0.0) # NaN / outside window -> 0 + p = torch.where(valid, pos, 0.0) + dv = torch.where( + valid, v - v.sum(-1, keepdim=True) / n, 0.0 + ) # centred values, 0 where invalid + dx = torch.where( + valid, p - p.sum(-1, keepdim=True) / n, 0.0 + ) # centred positions + s_xv = (dx * dv).sum(-1) + s_xx = (dx * dx).sum(-1) + s_vv = (dv * dv).sum(-1) + n = n.squeeze(-1) + std = torch.sqrt(s_vv / (n - 1)) # np.std(vals, ddof=1) + score = (s_xv / s_xx / std).abs() * windows # |slope_z| * len(row) + r2 = s_xv * s_xv / (s_xx * s_vv) # 1 - ss_res / ss_tot + return score, r2, std, n + + @staticmethod + def _diff_forward(x: torch.Tensor, left_pad: torch.Tensor) -> torch.Tensor: + """First-order row-wise differencing of left-padded rows ``[rows, T]``. + + ``left_pad`` holds each row's number of leading padding columns: the padding stays NaN and + each row's first real step is 0. + """ + diff = torch.full_like(x, torch.nan) + diff[:, 1:] = x[:, 1:] - x[:, :-1] + diff[torch.arange(x.shape[0], device=x.device), left_pad] = 0.0 + return diff + + @staticmethod + def _diff_params(x: torch.Tensor) -> dict: + """Retain last observed levels (NaN-safe, falling back to 0.0) for inversion.""" + not_nan = ~torch.isnan(x) + last_idx = not_nan.long().cumsum(dim=-1).argmax(dim=-1) + last_values = x.gather(-1, last_idx.unsqueeze(-1)).squeeze(-1) + return {"last_values": torch.nan_to_num(last_values, nan=0.0)} + + def _diff_inverse( + self, + sample: torch.Tensor, + mask: torch.Tensor, + params: dict, + preserve_grad: bool = False, + ) -> torch.Tensor: + """Invert first-order differencing on ``mask`` rows: median path + rss band term.""" + out = sample.clone() + if not mask.any(): + return out + constants = params["last_values"] + broadcast_shape = (-1,) + (1,) * (sample.ndim - 1) + diff = sample[mask] + + q_med = diff.shape[1] // 2 + median_step = diff[:, q_med : q_med + 1, :] + median_path = torch.cumsum(median_step, dim=-1) + offset = diff - median_step + + sign = offset.sum(dim=-1, keepdim=True).sign() + rss_squared = torch.cumsum(offset * offset, dim=-1) + if preserve_grad: + # The median offset is identically zero. Avoid sqrt'(0) in backward + # while keeping its forward value exactly zero. + rss = torch.sqrt(rss_squared.clamp_min(torch.finfo(rss_squared.dtype).tiny)) + rss = torch.where(rss_squared > 0, rss, torch.zeros_like(rss)) + else: + rss = torch.sqrt(rss_squared) + t_idx = torch.arange( + 1, diff.shape[-1] + 1, device=diff.device, dtype=diff.dtype + ) + horizon_factor = t_idx ** (self.cfg.diff_band_exponent - 0.5) + band = sign * rss * horizon_factor + + recon = median_path + (self.cfg.diff_band_scale**0.5) * band + out[mask] = recon + constants[mask].reshape(broadcast_shape) + return out diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/residual_block.py b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/residual_block.py new file mode 100644 index 0000000000000..d9ea6ac59c4fd --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/residual_block.py @@ -0,0 +1,67 @@ +# 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 (c) NXAI GmbH. +# SPDX-License-Identifier: Apache-2.0 + +"""Residual MLP building blocks used across TiRex modules.""" + +import torch +from torch import nn + +from iotdb.ainode.core.model.tirex2.component.norm import RMSNorm + + +class ResidualBlock(nn.Module): + """Two-layer MLP with residual projection and optional layer norm. + + Upstream compiles this block with ``torch.compile``. AINode runs it eagerly: Inductor needs a C++ + compiler and the Python headers at run time, and the packaged AINode executable has no headers. + """ + + def __init__( + self, + in_dim: int, + h_dim: int, + out_dim: int, + act_fn: nn.Module, + dropout_p: float = 0.0, + use_layer_norm: bool = False, + ) -> None: + super().__init__() + + self.dropout = nn.Dropout(dropout_p) + self.hidden_layer = nn.Linear(in_dim, h_dim) + self.act = act_fn + self.output_layer = nn.Linear(h_dim, out_dim) + self.residual_layer = nn.Linear(in_dim, out_dim) + + self.use_layer_norm = use_layer_norm + if use_layer_norm: + self.layer_norm = RMSNorm(out_dim) + + def forward(self, x: torch.Tensor): + """Apply hidden MLP, residual projection, and optional layer norm.""" + hid = self.act(self.hidden_layer(x)) + out = self.dropout(self.output_layer(hid)) + res = self.residual_layer(x) + + out = out + res + + if self.use_layer_norm: + return self.layer_norm(out) + return out diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/scaler.py b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/scaler.py new file mode 100644 index 0000000000000..4ffa477ee003a --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/scaler.py @@ -0,0 +1,97 @@ +# 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 (c) NXAI GmbH. +# SPDX-License-Identifier: Apache-2.0 + +"""Random past-window standardization for stochastic normalization.""" + +import torch + + +class Scaler: + """Normalizes using mean/variance computed over the full sequence. + + Parameters + ---------- + eps : float + Small constant added to avoid division by zero. + use_arcsinh : bool + If ``True``, apply an arcsinh squashing after standardizing. + binaryaware : bool + If ``True``, detect binary variates and bypass scaling for them. + """ + + def __init__( + self, + eps: float = 1e-8, + use_arcsinh: bool = False, + binaryaware: bool = False, + **kwargs + ) -> None: + super().__init__() + self.eps = eps + self.use_arcsinh = use_arcsinh + self.binaryaware = binaryaware + + def scale( + self, + x: torch.Tensor, + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + """Normalize ``x`` using mean/variance over the full sequence.""" + loc = torch.nan_to_num(torch.nanmean(x, dim=-1, keepdim=True), nan=0.0) + scale = torch.nan_to_num( + torch.nanmean((x - loc) ** 2, dim=-1, keepdim=True).sqrt(), nan=1.0 + ) + scale = torch.clamp_min(scale, min=self.eps) + + scaled_x = (x - loc) / scale + + # Detect binary variates (only NaN, 0.0, 1.0) and bypass scaling + if self.binaryaware: + valid = ~torch.isnan(x) + is_binary = ((x == 0.0) | (x == 1.0) | ~valid).all(dim=-1, keepdim=True) + loc = torch.where(is_binary, torch.zeros_like(loc), loc) + scale = torch.where(is_binary, torch.ones_like(scale), scale) + scaled_x = (x - loc) / scale + else: + is_binary = None + + if self.use_arcsinh: + if is_binary is not None and is_binary.any(): + scaled_x = torch.where(is_binary, scaled_x, torch.arcsinh(scaled_x)) + else: + scaled_x = torch.arcsinh(scaled_x) + + return scaled_x, (loc, scale, is_binary) + + def re_scale( + self, + x: torch.Tensor, + loc_scale: tuple[torch.Tensor, torch.Tensor, torch.Tensor | None], + ) -> torch.Tensor: + """Undo the normalization performed by :meth:`scale`.""" + loc, scale, is_binary = loc_scale + if self.use_arcsinh: + # clamp before sinh to prevent overflow: sinh(x) is stable for |x| < ~88 (float32) + sinh_x = torch.sinh(torch.clamp(x, -20.0, 20.0)) + if is_binary is not None and is_binary.any(): + # is_binary: (B, 1) -> broadcast over quantile and time dims + x = torch.where(is_binary[:, None], x, sinh_x) + else: + x = sinh_x + return x * scale[:, None] + loc[:, None] diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/slstm_block.py b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/slstm_block.py new file mode 100644 index 0000000000000..167a1c37d213e --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/slstm_block.py @@ -0,0 +1,310 @@ +# 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 (c) NXAI GmbH. +# SPDX-License-Identifier: Apache-2.0 + +"""sLSTM layers and configuration helpers. + +The recurrence runs on the vendored slstm-kernels: PyTorch on CPU/MPS (and on CUDA +without Triton) and Triton on CUDA. +""" + +import math +from dataclasses import dataclass, field +from math import sqrt +from typing import Literal + +import torch +from torch import nn + +from iotdb.ainode.core.model.tirex2.component import kernel +from iotdb.ainode.core.model.tirex2.component.norm import MultiHeadLayerNorm + + +class HeadwiseLinear(nn.Module): + """Per-head linear projections for all gates at once: [..., H, E] -> [..., G, H, O].""" + + def __init__(self, num_heads: int, head_dim: int, num_gates=4): + super().__init__() + self.weight = nn.Parameter( + torch.empty(num_gates, num_heads, head_dim, head_dim) + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.einsum("...HE,GHEO->...GHO", x, self.weight.transpose(-1, -2)) + + +# Copied from xlstm.components.init. +def small_init_init_(param: torch.Tensor, dim: int) -> torch.Tensor: + """Fills the input Tensor with values according to the method described in Transformers without Tears: Improving + the Normalization of Self-Attention - Nguyen, T. & Salazar, J. (2019), using a normal distribution. + Adopted from https://github.com/EleutherAI/gpt-neox/blob/main/megatron/model/init_functions.py. + """ + std = math.sqrt(2 / (5 * dim)) + torch.nn.init.normal_(param, mean=0.0, std=std) + return param + + +@dataclass +class sLSTMLayerConfig: + """Configuration for the sLSTM layers used inside TiRex.""" + + embedding_dim: int = -1 + num_heads: int = 4 # this must divide the embedding_dim + conv1d_kernel_size: int = 0 # 0 means no convolution included + dropout: float = 0.0 + + # Cell specific inits + recurrent_weight_init: str = "standard" + bias_init: str = "powerlaw_blockdependent" + + # The vendored slstm-kernels implementation: "triton" (CUDA) or "torch" (any device); "auto" picks + # "triton" for CUDA inputs when Triton is available and "torch" otherwise. + backend: Literal["auto", "triton", "torch"] = "auto" + + # The sLSTM has 4 gates and 4 states (y, c, n, m) + num_gates_i: int = field(default=4, init=False) + num_states: int = field(default=4, init=False) + + def __post_init__(self): + """Validate dimensions and derive head information.""" + if self.backend not in ("auto", "triton", "torch"): + raise ValueError( + f"sLSTM backend must be 'auto', 'triton' or 'torch', got {self.backend!r}." + ) + assert self.embedding_dim % self.num_heads == 0 + self.hidden_dim = self.embedding_dim + self.head_dim = self.embedding_dim // self.num_heads + + +class sLSTMLayer(nn.Module): + """sLSTM layer: gate projections, the recurrence (see ``slstm_kernels``) and a group norm.""" + + config_class = sLSTMLayerConfig + + def __init__(self, config: sLSTMLayerConfig, block_idx: int, num_blocks: int): + super().__init__() + self.config = config + + if self.config.conv1d_kernel_size > 0: + raise NotImplementedError( + "The sLSTM layer no longer supports a causal convolution (conv1d_kernel_size > 0)." + ) + + # Input projections of the f, i, z and o gates, in that order + self.gate_proj = HeadwiseLinear( + self.config.num_heads, self.config.head_dim, self.config.num_gates_i + ) + + self.group_norm = MultiHeadLayerNorm( + num_heads=self.config.num_heads, + head_dim=self.config.head_dim, + eps=1e-6, + ) + self.dropout = nn.Dropout(self.config.dropout) + + assert self.config.recurrent_weight_init in ["zeros", "standard"] + assert self.config.bias_init in [ + "powerlaw_blockdependent", + "zeros", + "standard", + "small_init", + ] + self._block_idx = block_idx + self._num_blocks = num_blocks + + self._recurrent_kernel_ = nn.Parameter( + torch.empty( + self.config.num_heads, + self.config.head_dim, + self.config.num_gates_i, + self.config.head_dim, + ) + ) + self._bias_ = nn.Parameter( + torch.empty( + self.config.num_heads, self.config.num_gates_i, self.config.head_dim + ) + ) + + self.reset_parameters() + + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + # Older checkpoints also store the recurrent kernel under this alias. + state_dict.pop(f"{prefix}recurrent_kernel", None) + # Older checkpoints store one [H, D_out, D_in] weight per gate + legacy_keys = [ + f"{prefix}{gate}.weight" for gate in ("fgate", "igate", "zgate", "ogate") + ] + if all(key in state_dict for key in legacy_keys): + state_dict[f"{prefix}gate_proj.weight"] = torch.stack( + [state_dict.pop(key) for key in legacy_keys] + ) + super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + + @torch.no_grad() + def reset_weights(self): + """Reset recurrent kernels according to the chosen scheme.""" + if self.config.recurrent_weight_init == "zeros": + nn.init.zeros_(self._recurrent_kernel_) + elif self.config.recurrent_weight_init == "standard": + bound = 1.0 / sqrt(self.config.hidden_dim) + nn.init.uniform_(self._recurrent_kernel_, -bound, bound) + + @torch.no_grad() + def reset_bias(self): + """Reset gate biases with power-law schedule for the forget gate.""" + if self.config.bias_init == "zeros": + nn.init.zeros_(self._bias_) + elif self.config.bias_init == "powerlaw_blockdependent": + ratio_0_to_1 = ( + self._block_idx / (self._num_blocks - 1) + if self._num_blocks > 1 + else 0.0 + ) + positions = torch.arange(self.config.head_dim) / (self.config.head_dim - 1) + nn.init.zeros_(self._bias_) + # gate order is i, f, z, o + self._bias_[:, 1, :] = -( + -5.0 + 12.0 * positions ** (0.3 + 1.3 * ratio_0_to_1) + ) + + def get_R(self): + """Return the recurrent kernel [H, D_in, G, D_out].""" + return self._recurrent_kernel_ + + def get_bias(self): + """Return the gate bias [H, G, D].""" + return self._bias_ + + def zero_state(self, batch_dim, input_): + """Allocate zero states for convolutional and recurrent parts.""" + return torch.zeros( + ( + self.config.num_states, + batch_dim, + 1, + self.config.num_heads, + self.config.head_dim, + ), + dtype=input_.dtype, + device=input_.device, + ) + + def recurrence(self, Wx: torch.Tensor, states: torch.Tensor | None = None): + return self.slstm_kernels( + Wx, self.get_R(), self.get_bias(), states, self.config.backend + ) + + @staticmethod + def slstm_kernels( + Wx: torch.Tensor, # [B, T, G, H, D] + R: torch.Tensor, # [H, D_in, G, D_out] ("HPGD") + b: torch.Tensor, # [H, G, D] + states: torch.Tensor | None = None, # [4, B, 1, H, D] + backend: Literal["auto", "triton", "torch"] = "auto", + ) -> tuple[torch.Tensor, torch.Tensor]: + """Run the recurrence on the vendored slstm-kernels, in FlashRNN's tensor layouts. + + Both match FlashRNN's sLSTM up to the layouts, keep the states in float32 + and run the recurrent matmul on bfloat16 operands. + """ + if backend == "auto": + backend = "triton" if kernel.use_triton(Wx) else "torch" + # Imported lazily, as triton is not available on every platform. + if backend == "triton": + from iotdb.ainode.core.model.tirex2.component.kernel.slstm_triton import ( + slstm_triton_fw as slstm_fw, + ) + else: + from iotdb.ainode.core.model.tirex2.component.kernel.slstm_torch import ( + slstm_torch as slstm_fw, + ) + + B, T, G, H, D = Wx.shape + x = Wx.permute(0, 1, 3, 2, 4).contiguous() # [B, T, H, G, D] + # Enforce layout even when R is already BF16 and .to() returns the same view. + weight = ( + R.permute(0, 2, 3, 1) + .to(torch.bfloat16, memory_format=torch.contiguous_format) + .contiguous() + ) # [H, G, D_out, D_in] + bias = b.to(torch.float32).contiguous() + if states is None: + states = torch.zeros((B, H, G, D), dtype=torch.float32, device=Wx.device) + else: + states = ( + states[:, :, 0].permute(1, 2, 0, 3).to(torch.float32).contiguous() + ) # [B, H, 4, D] + y, last_states = slstm_fw(x, states, weight, bias) + # Like FlashRNN: per-step states [S, B, T, H, D], of which callers read y[0]. + return y.unsqueeze(0), last_states.permute(2, 0, 1, 3).unsqueeze(2) + + def reset_parameters(self): + """Reset projections, recurrent weights, and biases.""" + small_init_init_(self.gate_proj.weight, dim=self.config.embedding_dim) + self.reset_weights() + self.reset_bias() + + def gate_preacts(self, x: torch.Tensor) -> torch.Tensor: + """Project [B, T, H * D] inputs to the gate pre-activations [B, T, G, H, D].""" + return self.gate_proj( + x.unflatten(-1, (self.config.num_heads, self.config.head_dim)) + ) + + def get_state(self, init_state, batch_dim, input_): + """Return ``init_state`` when provided, else allocate zeros.""" + if init_state is not None: + return init_state + return self.zero_state(batch_dim, input_) + + def step( + self, + x: torch.Tensor, + conv_state: torch.Tensor | None = None, + slstm_state: torch.Tensor | None = None, + ): + """Perform a single recurrent update.""" + batch_size, _, _ = x.shape + + Wx = self.gate_preacts(x) + y, slstm_state = self.recurrence( + Wx, states=self.get_state(slstm_state, batch_size, x) + ) + y = y[0] + + y = self.dropout(y) + + out = self.group_norm(y) + + return out, {"conv_state": conv_state, "slstm_state": slstm_state} + + def forward( + self, + x: torch.Tensor, + **kwargs, + ) -> torch.Tensor: + """Process a full sequence through the sLSTM.""" + Wx = self.gate_preacts(x) + y, _ = self.recurrence(Wx) + y = y[0] + + y = self.dropout(y) + + out = self.group_norm(y) + return out diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/variate_mixing_block.py b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/variate_mixing_block.py new file mode 100644 index 0000000000000..d4cdf49dd7a69 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/variate_mixing_block.py @@ -0,0 +1,300 @@ +# 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 (c) NXAI GmbH. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass +from typing import Any, Literal + +import torch +from torch import nn + +from iotdb.ainode.core.model.tirex2.component.attention_block import AttentionBlock +from iotdb.ainode.core.model.tirex2.component.bi_xlstm import BiXLSTM + + +def _unwrap_output( + output: torch.Tensor | tuple[torch.Tensor, Any], +) -> tuple[torch.Tensor, Any]: + """Unwrap module output that may be a tensor or (tensor, state) tuple. + + Parameters + ---------- + output : torch.Tensor or tuple[torch.Tensor, Any] + Module output, either a single tensor or (tensor, state) tuple + + Returns + ------- + tuple[torch.Tensor, Any] + Tuple of (tensor, state). If input is a tensor, state is None. + """ + if isinstance(output, tuple): + return output + return output, None + + +def variate_indices( + group_vector, target_mask, known_covariate_mask +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Variates of the reverse xLSTM cell and of non-singleton attention groups. + + Computed once per forward pass, as selecting by mask synchronizes with the + host, which would keep the block stack out of a CUDA graph. + """ + reverse_idx = torch.where( + ~target_mask if known_covariate_mask is None else known_covariate_mask + )[0] + attention_idx = None + if group_vector is not None: + _, inverse, counts = torch.unique( + group_vector.flatten(), return_inverse=True, return_counts=True + ) + attention_idx = torch.where(counts[inverse] > 1)[0] + return reverse_idx, attention_idx + + +@dataclass +class TimeMixerConfig: + """Configuration for time mixing model.""" + + model_type: Literal["bi-slstm", "bi-mlstm"] + + act_fn: nn.Module + + embedding_dim: int + num_heads: int = 4 + num_slstm_heads: int | None = None + + +@dataclass +class VariateMixerConfig: + """Configuration for variate mixing model.""" + + act_fn: nn.Module + + # dimensions + embedding_dim: int + num_heads: int = 4 + + # whether to RMS-normalize query/key vectors in the attention layer + use_qk_norm: bool = True + + # whether to use FlexAttention in the attention layer + use_flex_attention: bool = False + + +@dataclass +class MultivariateBlockConfig: + """Configuration for the MultivariateBlock. + + This block performs two-stage mixing: + 1. Time mixing: processes [B*V, L, P] - mixing across time dimension + 2. Variate mixing: processes [B*L, V, P] - mixing across variate dimension + + Normalization and residual connections are handled by the individual mixer components. + """ + + time_mixer: TimeMixerConfig + variate_mixer: VariateMixerConfig + block_order: list[str] | None = None + + # To initialize the xLSTM blocks + block_idx: int | None = None + num_blocks: int | None = None + + # Optional feed-forward network after variate mixing + dropout: float = 0.1 + eps: float = 1e-6 + + +class MultivariateBlock(nn.Module): + """Two-stage mixing block for multivariate time series. + + Stage 1 (Time Mixing): Processes each variate's time series independently [B*V, L, P] + Stage 2 (Variate Mixing): Reshapes to [L, B*V, P] and mixes across variates for each timestep. + Stage 3 (FFN): Final feed-forward transformation on tokens. + + Parameters + ---------- + config : MultivariateBlockConfig + Configuration specifying time mixer, variate mixer, and FFN settings + """ + + def __init__(self, config: MultivariateBlockConfig) -> None: + super().__init__() + + self.dropout = config.dropout if config.dropout is not None else 0.0 + + self.eps = config.eps + self.config = config + + # Build sub-blocks + self.time_mixer = self._build_time_mixer( + config.time_mixer, + block_idx=config.block_idx, + num_blocks=config.num_blocks, # type: ignore + ) + self.variate_mixer = self._build_variate_mixer( + config.variate_mixer, + block_idx=config.block_idx, + num_blocks=config.num_blocks, # type: ignore + ) + + def forward( + self, + x: torch.Tensor, + group_vector: torch.Tensor | None = None, + target_mask: torch.Tensor | None = None, + known_covariate_mask: torch.Tensor | None = None, + state: dict | None = None, + *, + indices: tuple | None = None, + return_svd: bool = False, + ) -> ( + tuple[torch.Tensor, dict[str, Any]] + | tuple[torch.Tensor, dict[str, Any], dict[str, torch.Tensor]] + ): + """Forward pass with three-stage mixing. + + Normalization and residual connections are handled by the individual mixer components. + + Parameters + ---------- + x : torch.Tensor + Input tensor of shape [B*V, L, P] where: + B = batch size + V = number of variates + L = sequence length + P = feature dimension + group_vector : torch.Tensor, optional + Optional group ids used when attention is the variate mixer. Should broadcast to the + variate dimension after reshaping (e.g., [B*V] or [B*V, 1]) + state : dict, optional + Optional state returned from a previous call. Keys: ``time_mixer`` and ``variate_mixer`` + indices : tuple, optional + :func:`variate_indices` of the masks, computed when omitted. + return_svd : bool, optional + Enables computation of the singular values on the first batch element in each sub-block. + `torch.svd_lowrank` is used to approximately compute the singular values (default: False) + + Returns + ------- + tuple[torch.Tensor, dict[str, Any]] or tuple[torch.Tensor, dict[str, Any], dict[str, torch.Tensor]] + If return_svd is False, returns (tokens, state) where tokens has shape [B*V, L, P] + and state is a dict with keys 'time_mixer' and 'variate_mixer'. + If return_svd is True, returns (tokens, state, svd_vals) where svd_vals contains + the singular values of the first batch element for each sub-block. + """ + BV, L, P = x.shape + reverse_idx, attention_idx = indices or variate_indices( + group_vector, target_mask, known_covariate_mask + ) + + # Stage 1: Time Mixing [B*V, L, P] -> [B*V, L, P] + # Each variate is processed independently across its time dimension + time_output, _ = _unwrap_output(self.time_mixer(x, reverse_idx)) + + # Stage 2: Variate Mixing + # Transpose for variate mixing: [B*V, L, P] -> [L, B*V, P] + # This treats L as the new batch dimension and mixes across B*V (all variates) + x_variate = time_output.permute(1, 0, 2) # [L, B*V, P] + + variate_output, _ = _unwrap_output( + self.variate_mixer( + x_variate, + group_vector=group_vector, + target_mask=target_mask, + attention_idx=attention_idx, + ) + ) + + # Transpose back: [L, B*V, P] -> [B*V, L, P] + x = variate_output.permute(1, 0, 2) + + return x, {} + + def _build_time_mixer( + self, config: TimeMixerConfig, block_idx: int, num_blocks: int + ) -> nn.Module: + """Instantiate time mixing model from configuration. + + Parameters + ---------- + config : TimeMixerConfig + Configuration for the time mixer + block_idx : int + Index of this block in the overall stack + num_blocks : int + Total number of blocks in the stack + + Returns + ------- + nn.Module + Time mixing module (sLSTM, mLSTM, or attention) + """ + match config.model_type: + case "bi-slstm": + cell_type, num_heads = ( + "slstm", + config.num_slstm_heads or config.num_heads, + ) + case "bi-mlstm": + cell_type, num_heads = "mlstm", config.num_heads + case _: + raise ValueError( + f"Unknown model_type for time mixer '{config.model_type}'. Allowed 'bi-slstm' and 'bi-mlstm'" + ) + + return BiXLSTM( + config.embedding_dim, + num_heads, + block_idx=block_idx, + num_blocks=num_blocks, + dropout=self.dropout, + cell_type=cell_type, + ) + + def _build_variate_mixer( + self, config: VariateMixerConfig, num_blocks: int, block_idx: int + ) -> nn.Module: + """Instantiate variate mixing model from configuration. + + Parameters + ---------- + config : VariateMixerConfig + Configuration for the variate mixer + num_blocks : int + Total number of blocks in the stack + block_idx : int + Index of this block in the overall stack + + Returns + ------- + nn.Module + Variate mixing module (attention, mLSTM, or identity) + """ + model = AttentionBlock( + input_dim=config.embedding_dim, + n_heads=config.num_heads, + act_fn=config.act_fn, + dropout=self.dropout, + use_qk_norm=config.use_qk_norm, + use_flex_attention=config.use_flex_attention, + ) + + return model diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/cuda_graph.py b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/cuda_graph.py new file mode 100644 index 0000000000000..004e39946b8a9 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/cuda_graph.py @@ -0,0 +1,90 @@ +# 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 (c) NXAI GmbH. +# SPDX-License-Identifier: Apache-2.0 + +"""Replay a forward pass from CUDA graphs: one launch instead of hundreds of small kernels.""" + +from collections import OrderedDict + +import torch + + +def _key(tensor: torch.Tensor | None): + return ( + None + if tensor is None + else (tuple(tensor.shape), tensor.dtype, tensor.cpu().numpy().tobytes()) + ) + + +class CUDAGraphRunner: + """Run ``build(*structure)(x)`` from a CUDA graph per shape of ``x`` and content of ``structure``. + + ``build`` runs eagerly and may synchronize with the host; the function it + returns must not. The ``max_graphs`` most recently used graphs are kept. + """ + + def __init__(self, build, max_graphs: int = 32): + self.build = build + self.max_graphs = max_graphs + self._graphs = OrderedDict() + # Graphs replay one at a time and their outputs are cloned at once, so they can share memory. + self._pool = torch.cuda.graph_pool_handle() + # One warm-up stream for all captures: cuBLAS keeps a workspace per stream. + self._stream = torch.cuda.Stream() + + def __len__(self) -> int: + return len(self._graphs) + + @torch.no_grad() + def __call__( + self, x: torch.Tensor, *structure: torch.Tensor | None + ) -> torch.Tensor: + key = (tuple(x.shape), x.dtype, x.device, *map(_key, structure)) + if key not in self._graphs: + self._graphs[key], warmup_output = self._capture(x, structure) + if len(self._graphs) > self.max_graphs: + self._graphs.popitem(last=False) + # The warm-up already computed this input; capture only records operations. + # Protect its side-stream storage while the caller stream clones it. + warmup_output.record_stream(torch.cuda.current_stream(x.device)) + return warmup_output.clone() + self._graphs.move_to_end(key) + + graph, static_input, static_output, _ = self._graphs[key] + static_input.copy_(x) + graph.replay() + return static_output.clone() + + def _capture(self, x: torch.Tensor, structure: tuple) -> tuple: + static_input = x.clone() + constants = tuple(None if t is None else t.clone() for t in structure) + fn = self.build(*constants) + + # Warm up outside the capture, which also runs lazy compilation and autotuning. + self._stream.wait_stream(torch.cuda.current_stream(x.device)) + with torch.cuda.stream(self._stream): + warmup_output = fn(static_input) + torch.cuda.current_stream(x.device).wait_stream(self._stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=self._pool): + static_output = fn(static_input) + # The graph reads tensors by address: keep the structure and what fn closes over alive. + return (graph, static_input, static_output, (constants, fn)), warmup_output diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/model.py b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/model.py new file mode 100644 index 0000000000000..e27d5a145b058 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/model.py @@ -0,0 +1,670 @@ +# 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 (c) NXAI GmbH. +# SPDX-License-Identifier: Apache-2.0 + +import logging +from dataclasses import dataclass, replace +from typing import Any, Literal + +import torch +from torch import nn + +from iotdb.ainode.core.model.tirex2.component.norm import RMSNorm +from iotdb.ainode.core.model.tirex2.component.patch_tokenizer import Tokenizer +from iotdb.ainode.core.model.tirex2.component.postprocessor import PostProcessor +from iotdb.ainode.core.model.tirex2.component.residual_block import ResidualBlock +from iotdb.ainode.core.model.tirex2.component.scaler import Scaler +from iotdb.ainode.core.model.tirex2.component.variate_mixing_block import ( + MultivariateBlock, + MultivariateBlockConfig, + TimeMixerConfig, + VariateMixerConfig, + variate_indices, +) +from iotdb.ainode.core.model.tirex2.cuda_graph import CUDAGraphRunner +from iotdb.ainode.core.model.tirex2.types import TimeseriesType + +MatmulPrecision = Literal["highest", "high", "medium"] + +logger = logging.getLogger(__file__) + + +def _resolve_act_func(spec: str) -> nn.Module: + """Instantiate an activation module from its ``torch.nn`` name, such as ``SiLU`` or ``torch.nn.SiLU``. + + Only ``torch.nn`` modules are accepted, as the name comes from the checkpoint's config.json. + """ + cls_name = spec.removeprefix("torch.nn.") + act_cls = getattr(torch.nn, cls_name, None) + if "." in cls_name or not ( + isinstance(act_cls, type) and issubclass(act_cls, nn.Module) + ): + raise ValueError(f"act_func must name a torch.nn module, got {spec!r}.") + return act_cls() + + +@dataclass +class MultivariateStackConfig: + """Named block templates with a per-position recipe for the multivariate block stack. + + Parameters + ---------- + templates : dict[str, MultivariateBlockConfig] + Named block configurations. Each key is a template name (e.g. ``"standard"``, + ``"recurrent"``). ``block_idx`` and ``num_blocks`` must NOT be set here; they + are injected by the backbone at stack construction time. + recipe : list[str] + One template name per block position (length must equal ``num_blocks``). + Example: ``["standard", "recurrent", "standard", ...]`` + """ + + templates: dict[str, MultivariateBlockConfig] + recipe: list[str] + + @classmethod + def from_dict( + cls, + config: dict, + act_func: nn.Module, + use_qk_norm: bool = True, + ) -> "MultivariateStackConfig": + """Build a stack config from a (possibly serialized) dict, injecting a shared activation. + + Each template may be given either as a fully-built + :class:`MultivariateBlockConfig` (used as-is) or as a plain dict whose + ``time_mixer``/``variate_mixer`` carry no ``act_fn``; in the latter case + the provided ``act_func`` instance is injected into every mixer, matching + how the backbone is configured with a single activation referenced + throughout. + + Parameters + ---------- + config : dict + Stack config with keys ``templates`` (name -> :class:`MultivariateBlockConfig` + or a template dict with ``time_mixer``, ``variate_mixer``, ``dropout``, + ``eps``) and ``recipe`` (list of template names). + act_func : nn.Module + Activation instance shared across all mixers built from dict templates. + use_qk_norm : bool + Whether the attention variate mixers should RMS-normalize their + query/key vectors. Injected into every mixer built from a dict + template (default: True). + + The recurrent kernels follow the device of their inputs, so any + serialized time mixer backend setting is ignored. + """ + templates = {} + for name, template in config["templates"].items(): + if isinstance(template, MultivariateBlockConfig): + templates[name] = template + else: + time_mixer_config = dict(template["time_mixer"]) + time_mixer_config.pop("rnn_backend", None) + templates[name] = MultivariateBlockConfig( + time_mixer=TimeMixerConfig(act_fn=act_func, **time_mixer_config), + variate_mixer=VariateMixerConfig( + act_fn=act_func, + use_qk_norm=use_qk_norm, + **template["variate_mixer"], + ), + dropout=template["dropout"], + eps=template["eps"], + ) + return cls(templates=templates, recipe=list(config["recipe"])) + + +class TiRex2(nn.Module): + def __init__( + self, + stack_config: dict, + num_blocks: int, + embedding_dim: int, + input_patch_size: int, + output_patch_size: int, + quantiles: list[float], + tokenizer_cfg: dict, + scaler_cfg: dict, + h_expand: float, + context_len: int, + future_len: int, + input_ff_dim: int, + act_func: str, + dropout: float, + *args, + use_qk_norm: bool = True, + stack_out_norm_config: dict | None = None, + tta_sign_flip: bool = False, + tta_diff: bool = True, + matmul_precision: MatmulPrecision | None = None, + **kwargs, + ) -> None: + if matmul_precision is not None: + torch.set_float32_matmul_precision(matmul_precision) + + kwargs.pop("postprocessor_cfg", None) + super().__init__(*args, **kwargs) + act_func = _resolve_act_func(act_func) + self.postprocessor = PostProcessor() + self.stack_config = MultivariateStackConfig.from_dict( + stack_config, act_func, use_qk_norm + ) + + if num_blocks == 0: + raise ValueError("Cannot create a model without any blocks in the stack") + + assert ( + future_len % input_patch_size == 0 and input_patch_size == output_patch_size + ), ( + "The future_len has to be integer divisible by the patch_size but " + f"future_len%input_patch_size={future_len % input_patch_size} and input_patch_size has to " + "be the same as output_patch_size (for inference) but " + f"input_patch_size == output_patch_size={input_patch_size == output_patch_size}" + ) + + self.input_patch_size = input_patch_size + self.output_patch_size = output_patch_size + self.input_ff_dim = input_ff_dim + self.act_func = act_func + self.h_expand = h_expand + self.context_len = context_len + self.future_len = future_len + self.embedding_dim = embedding_dim + self.dropout = dropout + + # Input layer + self.nan_mask_value = 0 + self.tokenizer = Tokenizer(**tokenizer_cfg) + self.scaler = Scaler(**scaler_cfg) + self.input_patch_embedding = ResidualBlock( + in_dim=self.input_patch_size * 2, + h_dim=self.input_ff_dim, + out_dim=embedding_dim, + act_fn=self.act_func, + dropout_p=self.dropout, + ) + + if stack_out_norm_config is not None: + eps = stack_out_norm_config.get("eps", 1e-7) + use_rmsnorm = stack_out_norm_config.get("use_rmsnorm", False) + if use_rmsnorm: + self.stack_out_norm = torch.nn.RMSNorm(embedding_dim, eps=eps) + else: + self.stack_out_norm = RMSNorm(embedding_dim, eps=eps) + else: + self.stack_out_norm = torch.nn.Identity() + + # Multivariate stack + self.stack: nn.ModuleList = self._create_stack(self.stack_config, num_blocks) + + # Output Layer + self.num_quantiles = len(quantiles) + quantiles_tensor = torch.tensor(quantiles) + self.register_buffer("quantiles", quantiles_tensor, persistent=False) + # TTA defaults for this checkpoint (carried in config.json). + # ``predict`` uses them whenever the matching argument is left as None. + self.tta_sign_flip = tta_sign_flip + self.tta_diff = tta_diff + # Lazily-built q -> 1-q index map for sign-flip TTA; see _complement_indices. + self._complement_idx_cache: torch.Tensor | None = None + self.output_patch_embedding = ResidualBlock( + in_dim=embedding_dim, + h_dim=int(embedding_dim * self.h_expand), + out_dim=self.num_quantiles * self.output_patch_size, + act_fn=self.act_func, + dropout_p=self.dropout, + ) + self.use_cuda_graphs = False + self._cuda_graphs: CUDAGraphRunner | None = None + + def forward( + self, batch: dict[str, Any] + ) -> torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Forward pass producing quantile predictions for all variates in batch.""" + ( + x, + group_vector, + target_mask, + known_covariate_mask, + scaler_state, + tokenizer_state, + ) = self._prepare_forward(batch) + + if self._uses_cuda_graphs(x): + x = self._run_cuda_graph(x, group_vector, target_mask, known_covariate_mask) + else: + x = self._build_stack_fn(group_vector, target_mask, known_covariate_mask)(x) + + x = torch.unflatten( + x, -1, (self.num_quantiles, self.output_patch_size) + ) # [B*V, L, D_out] -> [B*V, L, Q, P] + x = torch.transpose( + x, 1, 2 + ) # switch quantile and num_token_dimension [B*V, L, Q, P] -> [B*V, Q, L, P] + + # reverse tokenization and scaling + x = self.tokenizer.output_transform( + x, tokenizer_state + ) # [B*V, Q, L, P] -> [B*V, Q, T'], T' = T + padding + x = self.scaler.re_scale(x, scaler_state) + + return x + + def _build_stack_fn(self, group_vector, target_mask, known_covariate_mask): + """The stack as a function of the tokens; the mask indices are computed here, outside any CUDA graph.""" + indices = variate_indices(group_vector, target_mask, known_covariate_mask) + return lambda x: self._stack( + x, group_vector, target_mask, known_covariate_mask, indices + ) + + def _stack(self, x, group_vector, target_mask, known_covariate_mask, indices): + """Block stack and output layer.""" + # forward through actual mLSTM/sLSTM layers, we select only the output of the last layer; + # hidden representations are not relevant for forecasting + for x in self._forward_layerwise( + x, group_vector, target_mask, known_covariate_mask, indices + ): + pass + + # Normalise at the end of the stack + x = self.stack_out_norm(x) + x = nn.functional.dropout(x, self.dropout, training=self.training) + return self.output_patch_embedding(x) # [B*V, L, D] -> [B*V, L, D_out] + + def _run_cuda_graph( + self, x, group_vector, target_mask, known_covariate_mask + ) -> torch.Tensor: + """Run the stack from a CUDA graph, captured per input shape. + + To share graphs, several variates are padded to a multiple of 4 with zero + rows, each a target in a new group, so the real variates never see them. + """ + if self._cuda_graphs is None: + self._cuda_graphs = CUDAGraphRunner(self._build_stack_fn) + + num_rows = x.shape[0] + num_pad = -num_rows % 4 if num_rows > 1 and group_vector is not None else 0 + if num_pad: + x = nn.functional.pad(x, (0, 0, 0, 0, 0, num_pad)) + group_vector = torch.cat( + [ + group_vector, + group_vector.max() + 1 + torch.arange(num_pad, device=x.device), + ] + ) + target_mask = nn.functional.pad(target_mask, (0, num_pad), value=True) + if known_covariate_mask is not None: + known_covariate_mask = nn.functional.pad( + known_covariate_mask, (0, num_pad), value=False + ) + return self._cuda_graphs(x, group_vector, target_mask, known_covariate_mask)[ + :num_rows + ] + + def _uses_cuda_graphs(self, x: torch.Tensor) -> bool: + """CUDA graphs serve gradient-free CUDA inference; FlexAttention's block mask can't be replayed. + + Graphs only pay off while kernel launches dominate, so larger batches run + eagerly rather than holding a graph's buffers per shape. + """ + return ( + self.use_cuda_graphs + and x.is_cuda + and x.shape[0] <= 32 # variate rows + and not self.training + and not torch.is_grad_enabled() + and not any( + block.variate_mixer.attn.use_flex_attention for block in self.stack + ) + ) + + def _apply(self, fn, *args, **kwargs): + # The graphs read the parameters by address, which moving the model changes. + self._cuda_graphs = None + return super()._apply(fn, *args, **kwargs) + + def _forward_layerwise( + self, x, group_vector, target_mask, known_covariate_mask, indices=None + ): + # Run the input through the stack. + state = {i: None for i in range(len(self.stack))} + for i, block in enumerate(self.stack): # type: ignore[reportArgumentType] + block_state = state[i] + x, new_block_state = block( + x, + group_vector=group_vector, + target_mask=target_mask, + known_covariate_mask=known_covariate_mask, + state=block_state, + indices=indices, + ) + + if block_state is None: + state[i] = new_block_state + yield x + + def _prepare_forward(self, batch): + x: torch.Tensor = batch["x"] + group_vector: torch.Tensor | None = batch.get("group_vector", None) + target_mask: torch.Tensor | None = batch.get("target_mask", None) + + # Determine which covariates are "known" (have non-NaN values in the future window). + # Past-only covariates (all-NaN future) will be treated as forward-only by reverse_known_only. + if target_mask is not None and x.shape[-1] >= self.future_len: + has_future = ~torch.all(torch.isnan(x[:, -self.future_len :]), dim=-1) + known_covariate_mask = has_future & ~target_mask + else: + known_covariate_mask = None + + # apply scaling and tokenize + x, scaler_state = self.scaler.scale(x) + x, tokenizer_state = self.tokenizer.input_transform(x) + + # embed the tokens with the "is-present" indication (i.e., x_mask) + x_mask = (~torch.isnan(x)).to(x.dtype) + x = torch.nan_to_num(x, nan=self.nan_mask_value) + x = torch.cat((x, x_mask), dim=-1) + x = self.input_patch_embedding(x) + return ( + x, + group_vector, + target_mask, + known_covariate_mask, + scaler_state, + tokenizer_state, + ) + + @torch.no_grad + def predict( + self, + timeseries: list[TimeseriesType], + prediction_length: int, + *args, + tta_sign_flip: bool | None = None, + tta_diff: bool | None = None, + **kwargs, + ): + """Return quantile forecasts aligned with the input sequence length. + + Parameters + ---------- + timeseries : list[TimeseriesType] + A batch of multivariate timeseries. Each :class:`TimeseriesType` holds + the target and optional covariate tensors; the target has shape + ``(num_variates, sequence_length)``. + prediction_length : int + The forecast horizon. + tta_sign_flip : bool, optional + Opt-in sign-flip test-time augmentation. Leave as ``None`` (the + default) to use the checkpoint's configured setting + (``self.tta_sign_flip``, from ``config.json``); pass an + explicit ``True``/``False`` to override it for this call. When + enabled, the model is run a second time on the sign-flipped input + (target *and* every covariate negated); the flipped forecast is + mapped back to level space (values negated and the quantile axis + reversed via the ``q -> 1-q`` complement map) and the two passes + are averaged in level space. Requires a symmetric quantile set + and roughly doubles inference cost. When it resolves to disabled, + the output is byte-identical to a single pass. + tta_diff : bool, optional + Opt-in differencing path inside the postprocessor. Leave as ``None`` + (the default) to use the checkpoint's configured setting + (``self.tta_diff``, from ``config.json``); pass an explicit + ``True``/``False`` to override trend differencing for this call. + pad_context : bool, optional + If ``True`` (the default), contexts shorter than ``context_len`` are + left-padded with NaN up to the full model length. Pass ``False`` to + skip padding to the full model length, reducing the number of input + patches and inference time for short contexts at the cost of a small + loss in forecast accuracy. Padding to the longest context in each + batch and to a multiple of 4 patches still applies. + """ + if tta_sign_flip is None: + tta_sign_flip = self.tta_sign_flip + if tta_diff is None: + tta_diff = self.tta_diff + + forecasts = self._predict( + timeseries, prediction_length, *args, tta_diff=tta_diff, **kwargs + ) + # An empty batch has nothing to average (and no tensor to read a device + # off of below), so short-circuit before the augmentation pass. + if not tta_sign_flip or not forecasts: + return forecasts + + # Treat predict end-to-end as a black box run twice: a second pass on + # the sign-flipped series, mapped back to level space, then averaged. + flipped = [self._sign_flip(ts) for ts in timeseries] + forecasts_flip = self._predict( + flipped, prediction_length, *args, tta_diff=tta_diff, **kwargs + ) + + complement = self._complement_indices().to(forecasts_flip[0].device) + forecasts_flip = [-f.index_select(1, complement) for f in forecasts_flip] + return [(a + b) / 2 for a, b in zip(forecasts, forecasts_flip)] + + def _prepare_input( + self, + timeseries, + prediction_length, + tta_diff, + *args, + pad_context: bool = True, + **kwargs, + ): + """Run a single (un-augmented) forecast pass over the batch of series.""" + if self.postprocessor is None: + raise RuntimeError("Cannot predict without postprocessor configured.") + + if prediction_length < 1: + raise ValueError( + f"prediction_length must be >= 1, got {prediction_length}." + ) + + if prediction_length > self.future_len: + logger.warning( + f"prediction_length={prediction_length} exceeds the supported maximum of " + f"{self.future_len}. It will be truncated to {self.future_len}.", + ) + prediction_length = self.future_len + + device = next(self.parameters()).device + + context = [ts.target.cpu() for ts in timeseries] + past_covariates = [ + ts.past_covariates.cpu() if ts.past_covariates is not None else None + for ts in timeseries + ] + past_future_covariates = [] + for ts in timeseries: + if ts.future_covariates is None: + past_future_covariates.append(None) + continue + expected_future_covariate_length = ts.target.shape[-1] + prediction_length + future_covariates = ts.future_covariates + if future_covariates.shape[-1] > expected_future_covariate_length: + future_covariates = future_covariates[ + ..., :expected_future_covariate_length + ] + past_future_covariates.append(future_covariates.cpu()) + + batch, args, kwargs = self.postprocessor.transform_input( + context, + prediction_length, + *args, + past_covariates=past_covariates, + past_future_covariates=past_future_covariates, + tta_diff=tta_diff, + **kwargs, + ) + + batch = {k: v.to(device) for k, v in batch.items()} + # ensure that context size matches the model's expectations + batch["x"] = self._fix_context_size(prediction_length, batch["x"], pad_context) + return prediction_length, batch, args, kwargs + + def _fix_context_size(self, prediction_length, context, pad_context): + right_pad = self.future_len - prediction_length + context = nn.functional.pad(context, (0, right_pad), value=torch.nan) + + max_ts_len = self.context_len + self.future_len + if context.shape[-1] < max_ts_len: + # Without pad_context, pad to a multiple of 4 patches: fewer shapes to compile and capture. + bucket = 4 * self.input_patch_size + pad_to = ( + max_ts_len + if pad_context + else min(-(-context.shape[-1] // bucket) * bucket, max_ts_len) + ) + context = nn.functional.pad( + context, (pad_to - context.shape[-1], 0), value=torch.nan + ) + elif context.shape[-1] > max_ts_len: + context = context[..., -max_ts_len:] + return context + + @staticmethod + def _sign_flip(ts: TimeseriesType) -> TimeseriesType: + """Negate the target and every covariate of a series for sign-flip TTA.""" + return TimeseriesType( + target=-ts.target, + past_covariates=None if ts.past_covariates is None else -ts.past_covariates, + future_covariates=( + None if ts.future_covariates is None else -ts.future_covariates + ), + ) + + def _complement_indices(self, tol: float = 1e-6) -> torch.Tensor: + """Index map sending each quantile ``q[i]`` to the one closest to ``1 - q[i]``. + + Sign-flip TTA reverses the quantile axis of the flipped forecast: the + ``q`` quantile of ``-x`` is the ``1 - q`` quantile of ``x``. This builds + ``complement`` with ``complement[i] = argmin_j |q[j] - (1 - q[i])|``; + ``0.5`` maps to itself. Requires a symmetric quantile set - each + complement must exist within ``tol`` or a :class:`ValueError` is raised. + Cached on first use (it depends only on ``self.quantiles``). + """ + if self._complement_idx_cache is not None: + return self._complement_idx_cache + + q = self.quantiles.detach().to(dtype=torch.float32, device="cpu") + complement = 1.0 - q + indices = [] + for i in range(q.shape[0]): + dist = torch.abs(q - complement[i]) + idx = int(torch.argmin(dist).item()) + if dist[idx].item() > tol: + raise ValueError( + "tta_sign_flip requires a symmetric quantile set: quantile " + f"{q[i].item():.4f} has no complement (1 - q) within {tol} in " + f"{self.quantiles.tolist()}." + ) + indices.append(idx) + self._complement_idx_cache = torch.tensor(indices, dtype=torch.long) + return self._complement_idx_cache + + def _predict( + self, + timeseries: list[TimeseriesType], + prediction_length: int, + *args, + tta_diff: bool = True, + preserve_grad: bool = False, + pad_context: bool = True, + **kwargs, + ): + """Run the model on all series inside context in parallel.""" + prediction_length, batch, args, kwargs = self._prepare_input( + timeseries, + prediction_length, + tta_diff, + *args, + pad_context=pad_context, + **kwargs, + ) + + pred = self(batch) + # Drop the last token (predicts beyond sequence end); the remaining + # tail of length future_len covers exactly the 10 future patches + # that were directly supervised during training. + pred = pred[..., : -self.output_patch_size] + pred = pred[:, :, -self.future_len :] + pred = pred[:, :, :prediction_length] + + result = self._prepare_output( + timeseries, + pred, + prediction_length, + *args, + preserve_grad=preserve_grad, + **kwargs, + ) + + return result + + def _prepare_output( + self, + original_timeseries, + pred, + prediction_length, + *args, + preserve_grad, + **kwargs, + ): + + output = self.postprocessor.transform_output( + pred, + prediction_length, + *args, + preserve_grad=preserve_grad, + **kwargs, + ) # type: ignore + result = [] + for ts, out in zip(original_timeseries, output): + result.append(out.sort(dim=1).values.to(ts.target)) + return result + + def _create_stack( + self, + config: MultivariateBlockConfig | MultivariateStackConfig, + num_blocks: int, + ) -> nn.ModuleList: + """Create stack of multivariate blocks from configuration.""" + if isinstance(config, MultivariateStackConfig): + templates = config.templates + recipe = config.recipe + if len(recipe) != num_blocks: + raise ValueError( + f"block_recipe length ({len(recipe)}) must equal num_blocks ({num_blocks})" + ) + unknown = set(recipe) - set(templates) + if unknown: + raise ValueError( + f"Recipe references unknown templates: {unknown}. Available: {set(templates)}" + ) + else: + # Backward-compatible path: repeat single config for all positions + templates = {"default": config} + recipe = ["default"] * num_blocks + + stack = [] + for block_idx, name in enumerate(recipe): + template = templates[name] + block_config = replace(template, block_idx=block_idx, num_blocks=num_blocks) + stack.append(MultivariateBlock(block_config)) + return nn.ModuleList(stack) diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/modeling_tirex2.py b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/modeling_tirex2.py new file mode 100644 index 0000000000000..587dfd2d894f4 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/modeling_tirex2.py @@ -0,0 +1,44 @@ +# 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.tirex2.model import TiRex2 + + +class TiRex2ForPrediction(TiRex2, PyTorchModelHubMixin): + """ + TiRex-2 as an AINode model, loaded from a directory with the ``config.json`` and + ``model.safetensors`` of the Hugging Face checkpoint ``NX-AI/TiRex-2``. + + The keys of ``config.json`` are the constructor arguments of :class:`TiRex2`. The recurrent + kernels follow the device the model is moved to: Triton kernels on CUDA when Triton is + available, PyTorch kernels otherwise. + """ + + @classmethod + def _load_as_safetensor( + cls, model: "TiRex2ForPrediction", model_file: str, map_location: str, strict + ) -> "TiRex2ForPrediction": + # Load through load_state_dict, whose hooks convert the per-gate weights of the published + # checkpoint to the fused layout of the recurrent layers; every weight must match. + state_dict = load_file(model_file, device=str(torch.device(map_location))) + model.load_state_dict(state_dict, strict=True) + return model.eval() diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/pipeline_tirex2.py b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/pipeline_tirex2.py new file mode 100644 index 0000000000000..9158645f14dc9 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/pipeline_tirex2.py @@ -0,0 +1,130 @@ +# 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.tirex2.types import TimeseriesType + + +class TiRex2Pipeline(ForecastPipeline): + """ + Forecasting pipeline of TiRex-2. + + Each input is one series of one or more targets, forecast jointly, with optional past + covariates and future-known covariates. The model predicts nine quantiles (0.1 to 0.9) of up + to 320 future steps in one pass; the pipeline returns the median. + """ + + def __init__(self, model_info: ModelInfo, **model_kwargs): + super().__init__(model_info, **model_kwargs) + + @property + def max_output_length(self) -> int: + return self.model.future_len + + def _preprocess(self, inputs, **infer_kwargs) -> list[TimeseriesType]: + """ + Convert the validated inputs to TiRex-2 series. + + 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'. + + infer_kwargs: + - `output_length` (int): The number of steps to forecast, at most 320. + + Returns + ------- + list of TimeseriesType + One series per input: the targets as a (target_count, input_length) tensor, the + covariates known only in the past as a (count, input_length) tensor, and the + covariates also known in the future as a (count, input_length + output_length) + tensor of their past and future values. + """ + output_length = infer_kwargs.get("output_length", 96) + if output_length > self.max_output_length: + raise ValueError( + f"TiRex-2 forecasts at most {self.max_output_length} steps, " + f"but output_length is {output_length}." + ) + 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( + TimeseriesType( + target=targets.float(), + past_covariates=( + torch.stack([past_covariates[k] for k in past_only]).float() + if past_only + else None + ), + future_covariates=( + torch.stack( + [ + torch.cat([past_covariates[k], future_covariates[k]]) + for k in known + ] + ).float() + if known + else None + ), + ) + ) + return series + + def forecast( + self, inputs: list[TimeseriesType], **infer_kwargs + ) -> list[torch.Tensor]: + """ + Forecast all series in one batch. + + Returns + ------- + list of torch.Tensor + One tensor per series, of shape (target_count, quantile_count, output_length). + """ + output_length = infer_kwargs.get("output_length", 96) + with torch.no_grad(): + return self.model.predict(inputs, output_length) + + 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.model.quantiles.tolist().index(0.5) + return [output[:, median, :].float().cpu() for output in outputs] diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/types.py b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/types.py new file mode 100644 index 0000000000000..b8aa6b7a1ed87 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/model/tirex2/types.py @@ -0,0 +1,81 @@ +# 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 (c) NXAI GmbH. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass + +import torch + + +@dataclass +class TimeseriesType: + """A single (possibly multivariate) series with optional covariates, as passed to + :meth:`TiRex2.predict`. + + Parameters + ---------- + target : torch.Tensor + Target history, shape ``[V_t, T]`` (``V_t`` target variates, context length ``T``). + A univariate series is still 2D, with ``V_t == 1``. + past_covariates : torch.Tensor or None + Covariates known only up to the current time, shape ``[V_p, T]``, matching the + target's context length. ``None`` if there are no past covariates. + future_covariates : torch.Tensor or None + Covariates known ahead of time for the whole forecast horizon, shape + ``[V_f, >=T+H]`` (``H`` is the requested ``prediction_length``); extra trailing + steps beyond ``T+H`` are ignored. ``None`` if there are no future covariates. + + Examples + -------- + >>> import torch + >>> from iotdb.ainode.core.model.tirex2.types import TimeseriesType + >>> ts = TimeseriesType( + ... target=torch.randn(1, 128), + ... past_covariates=None, + ... future_covariates=None, + ... ) + >>> ts.past_length + 128 + """ + + target: torch.Tensor # [V_t, T] + past_covariates: torch.Tensor | None # [V_p, T] + future_covariates: ( + torch.Tensor | None + ) # [V_f, >=T+H]; extra future steps are ignored + + @property + def n_past_covariates(self) -> int: + return 0 if self.past_covariates is None else len(self.past_covariates) + + @property + def n_future_covariates(self) -> int: + return 0 if self.future_covariates is None else len(self.future_covariates) + + @property + def past_length(self) -> int: + return self.target.shape[-1] + + @property + def future_length(self) -> int: + return ( + 0 + if self.future_covariates is None + else self.future_covariates.shape[-1] - self.past_length + ) diff --git a/iotdb-core/ainode/src/assembly/resources/NOTICE-binary b/iotdb-core/ainode/src/assembly/resources/NOTICE-binary index e30b3ffd5b396..24f6740b469a1 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 TiRex-2 project: + + TiRex-2 + Copyright 2026 NXAI GmbH + Licensed under the Apache License, Version 2.0 + https://github.com/NX-AI/tirex-2 + +============================================================================ + Apache Commons Collections Copyright 2001-2019 The Apache Software Foundation diff --git a/iotdb-core/ainode/tests/test_builtin_model_limits.py b/iotdb-core/ainode/tests/test_builtin_model_limits.py new file mode 100644 index 0000000000000..3847acb8ddf75 --- /dev/null +++ b/iotdb-core/ainode/tests/test_builtin_model_limits.py @@ -0,0 +1,96 @@ +# 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 unittest +from unittest import mock + +import torch + +from iotdb.ainode.core.constant import TSStatusCode +from iotdb.ainode.core.manager.inference_manager import InferenceManager +from iotdb.ainode.core.model.model_info import BUILTIN_HF_TRANSFORMERS_MODEL_MAP +from iotdb.ainode.core.util.serde import convert_tensor_to_tsblock +from iotdb.thrift.ainode.ttypes import TForecastReq + + +class BuiltinModelLimitsTest(unittest.TestCase): + """output_length is checked against a model's own limit before pool dispatch.""" + + def _manager(self, model_info): + manager = object.__new__(InferenceManager) + manager._model_manager = mock.Mock() + manager._model_manager.is_model_registered.return_value = True + manager._model_manager.get_model_info.return_value = model_info + return manager + + def test_model_limit_lowers_the_ainode_limit(self): + model_info = BUILTIN_HF_TRANSFORMERS_MODEL_MAP["tirex2"] + self.assertEqual(model_info.max_output_length, 320) + self.assertEqual(self._manager(model_info)._max_output_length("tirex2"), 320) + + def test_models_without_a_limit_keep_the_ainode_limit(self): + model_info = BUILTIN_HF_TRANSFORMERS_MODEL_MAP["chronos2"] + self.assertIsNone(model_info.max_output_length) + self.assertGreater( + self._manager(model_info)._max_output_length("chronos2"), 320 + ) + + def test_loaded_model_rejects_invalid_horizon_and_serves_next_request(self): + # With tirex2 loaded, a request outside its horizon must fail without reaching the + # inference pool, and the next valid request must still be served by the pool. + manager = self._manager(BUILTIN_HF_TRANSFORMERS_MODEL_MAP["tirex2"]) + manager._pool_controller = mock.Mock() + manager._pool_controller.has_running_pools.return_value = True + dispatched = [] + + def process_request(request): + dispatched.append(request) + return [torch.zeros(1, request.output_length)] + + manager._process_request = process_request + input_data = convert_tensor_to_tsblock(torch.randn(1, 64)) + + for output_length in (400, 0): + with self.subTest(output_length=output_length): + dispatched.clear() + rejected = manager.forecast( + TForecastReq( + modelId="tirex2", + inputData=input_data, + outputLength=output_length, + ) + ) + self.assertEqual( + rejected.status.code, TSStatusCode.AINODE_INTERNAL_ERROR.value + ) + self.assertIn( + f"between 1 and 320, got {output_length}", rejected.status.message + ) + self.assertEqual(dispatched, []) + + served = manager.forecast( + TForecastReq( + modelId="tirex2", inputData=input_data, outputLength=24 + ) + ) + self.assertEqual(served.status.code, TSStatusCode.SUCCESS_STATUS.value) + self.assertEqual([r.output_length for r in dispatched], [24]) + + +if __name__ == "__main__": + unittest.main() diff --git a/iotdb-core/ainode/tests/test_tirex2.py b/iotdb-core/ainode/tests/test_tirex2.py new file mode 100644 index 0000000000000..70a0fc3f82aab --- /dev/null +++ b/iotdb-core/ainode/tests/test_tirex2.py @@ -0,0 +1,254 @@ +# 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 re +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.tirex2.component import kernel +from iotdb.ainode.core.model.tirex2.modeling_tirex2 import TiRex2ForPrediction +from iotdb.ainode.core.model.tirex2.pipeline_tirex2 import TiRex2Pipeline + +QUANTILES = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] + + +def _template(model_type): + return { + "time_mixer": {"model_type": model_type, "embedding_dim": 32, "num_heads": 2}, + "variate_mixer": {"embedding_dim": 32, "num_heads": 2}, + "dropout": 0.0, + "eps": 1e-7, + } + + +# A small TiRex-2 with the structure of NX-AI/TiRex-2: alternating mLSTM and sLSTM blocks +TINY_CONFIG = { + "num_blocks": 2, + "embedding_dim": 32, + "input_patch_size": 8, + "output_patch_size": 8, + "input_ff_dim": 64, + "h_expand": 2, + "context_len": 64, + "future_len": 32, + "dropout": 0.0, + "quantiles": QUANTILES, + "act_func": "torch.nn.SiLU", + "use_qk_norm": True, + "tokenizer_cfg": { + "input_patch_size": 8, + "output_patch_size": 8, + "input_patch_stride": 8, + }, + "scaler_cfg": {"use_arcsinh": True, "binaryaware": True}, + "stack_out_norm_config": {"eps": 1e-7}, + "tta_diff": True, + "stack_config": { + "templates": { + "mlstm_attn": _template("bi-mlstm"), + "slstm_attn": _template("bi-slstm"), + }, + "recipe": ["mlstm_attn", "slstm_attn"], + }, +} + + +def _published_layout(state_dict): + """The state dict in the layout of the published checkpoint: one weight per gate.""" + state_dict = dict(state_dict) + for key in [k for k in state_dict if k.endswith("fwd_cell.gate_proj.weight")]: + prefix = key[: -len("gate_proj.weight")] + weights = state_dict.pop(key) + for i, gate in enumerate(("fgate", "igate", "zgate", "ogate")): + state_dict[f"{prefix}{gate}.weight"] = weights[i].clone() + for key in [k for k in state_dict if k.endswith("fwd_cell.in_proj.weight")]: + prefix = key[: -len("in_proj.weight")] + weights = state_dict.pop(key) + biases = state_dict.pop(f"{prefix}in_proj.bias") + qk_dim, v_dim, heads = 16, 32, 2 + split = (qk_dim, qk_dim, v_dim, v_dim, heads, heads) + names = ("q", "k", "v", "ogate_preact", "igate_preact", "fgate_preact") + for name, w, b in zip(names, weights.split(split), biases.split(split)): + state_dict[f"{prefix}{name}.weight"] = w.clone() + if name in ("igate_preact", "fgate_preact"): + state_dict[f"{prefix}{name}.bias"] = b.clone() + return state_dict + + +class TiRex2Test(unittest.TestCase): + @classmethod + def setUpClass(cls): + torch.manual_seed(0) + cls.model = TiRex2ForPrediction(**copy.deepcopy(TINY_CONFIG)).eval() + cls.checkpoint = tempfile.TemporaryDirectory() + with open(os.path.join(cls.checkpoint.name, "config.json"), "w") as f: + json.dump(TINY_CONFIG, f) + save_file( + _published_layout(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["tirex2"]) + 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 TiRex2Pipeline(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_published_checkpoint_layout_loads(self): + loaded = TiRex2ForPrediction.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): + published = _published_layout(self.model.state_dict()) + gate_biases = [ + k + for k in published + if k.endswith(("igate_preact.bias", "fgate_preact.bias")) + ] + self.assertTrue(gate_biases) + for missing in ["output_patch_embedding.output_layer.weight", *gate_biases[:2]]: + with self.subTest( + missing=missing + ), tempfile.TemporaryDirectory() as directory: + with open(os.path.join(directory, "config.json"), "w") as f: + json.dump(TINY_CONFIG, f) + state_dict = dict(published) + state_dict.pop(missing) + save_file(state_dict, os.path.join(directory, "model.safetensors")) + with self.assertRaisesRegex(RuntimeError, re.escape(missing)): + TiRex2ForPrediction.from_pretrained(directory) + + def test_builtin_pipeline_forecasts_median(self): + pipeline = self.load_pipeline() + self.assertIsInstance(pipeline.model, TiRex2ForPrediction) + torch.manual_seed(1) + inputs = [ + {"targets": torch.randn(100).cumsum(0)}, + {"targets": torch.randn(3, 40).cumsum(-1)}, + ] + outputs = self.run_pipeline(pipeline, copy.deepcopy(inputs), 24) + self.assertEqual([tuple(o.shape) for o in outputs], [(1, 24), (3, 24)]) + for output in outputs: + self.assertEqual(output.dtype, torch.float32) + self.assertTrue(torch.isfinite(output).all()) + + # the median quantile of the model's own forecast + series = pipeline.preprocess(copy.deepcopy(inputs), output_length=24) + quantiles = self.model.predict(series, 24) + for output, forecast in zip(outputs, quantiles): + torch.testing.assert_close(output, forecast[:, QUANTILES.index(0.5), :]) + + 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(16)}, + } + ] + series = pipeline.preprocess(copy.deepcopy(inputs), output_length=16) + self.assertEqual(tuple(series[0].target.shape), (2, 60)) + self.assertEqual(tuple(series[0].past_covariates.shape), (1, 60)) + self.assertEqual(tuple(series[0].future_covariates.shape), (1, 76)) + torch.testing.assert_close( + series[0].future_covariates[0], + torch.cat( + [ + inputs[0]["past_covariates"]["known"], + inputs[0]["future_covariates"]["known"], + ] + ), + ) + outputs = self.run_pipeline(pipeline, copy.deepcopy(inputs), 16) + self.assertEqual(tuple(outputs[0].shape), (2, 16)) + + def test_output_length_above_model_horizon_is_rejected(self): + pipeline = self.load_pipeline() + inputs = [{"targets": torch.randn(50)}] + self.run_pipeline(pipeline, copy.deepcopy(inputs), TINY_CONFIG["future_len"]) + with self.assertRaises(ValueError): + self.run_pipeline( + pipeline, copy.deepcopy(inputs), TINY_CONFIG["future_len"] + 1 + ) + + def test_forecast_runs_without_torch_compile(self): + # Inductor needs a C++ compiler and the Python headers, which the packaged AINode lacks. + from torch._dynamo.utils import counters + + pipeline = self.load_pipeline() + frames = counters["frames"]["total"] + self.run_pipeline(pipeline, [{"targets": torch.randn(2, 64)}], 8) + self.assertEqual(counters["frames"]["total"], frames) + + def test_activation_must_be_a_torch_nn_module(self): + for act_func in ("os.system", "torch.nn.functional.silu", "Linear.x"): + config = copy.deepcopy(TINY_CONFIG) + config["act_func"] = act_func + with self.subTest(act_func=act_func), self.assertRaises(ValueError): + TiRex2ForPrediction(**config) + + def test_cpu_inputs_use_pytorch_kernels(self): + self.assertFalse(kernel.use_triton(torch.zeros(1))) + + def test_pre_ampere_gpus_use_pytorch_kernels(self): + cuda_input = mock.Mock(is_cuda=True, device=torch.device("cuda", 0)) + for capability, expected in (((7, 0), False), ((7, 5), False), ((8, 0), True)): + kernel._supports_triton_kernels.cache_clear() + with self.subTest(capability=capability), mock.patch.object( + kernel, "triton_available", return_value=True + ), mock.patch.object( + torch.cuda, "get_device_capability", return_value=capability + ): + self.assertEqual(kernel.use_triton(cuda_input), expected) + kernel._supports_triton_kernels.cache_clear() + + +if __name__ == "__main__": + unittest.main()