diff --git a/README.md b/README.md index edfe46c..546c5bd 100644 --- a/README.md +++ b/README.md @@ -57,8 +57,12 @@ pip install -e . ### Code Example ```python +import torch import numpy as np import timesfm + +torch.set_float32_matmul_precision("high") + model = timesfm.TimesFM_2p5_200M_torch.from_pretrained("google/timesfm-2.5-200m-pytorch") model.compile( diff --git a/src/timesfm/configs.py b/src/timesfm/configs.py index f5ba9c5..104ea42 100644 --- a/src/timesfm/configs.py +++ b/src/timesfm/configs.py @@ -94,6 +94,7 @@ class TransformerConfig: use_bias: bool use_rotary_position_embeddings: bool ff_activation: Literal["relu", "swish", "none"] + fuse_qkv: bool @dataclasses.dataclass(frozen=True) diff --git a/src/timesfm/timesfm_2p5/timesfm_2p5_base.py b/src/timesfm/timesfm_2p5/timesfm_2p5_base.py index ba5902e..f438e40 100644 --- a/src/timesfm/timesfm_2p5/timesfm_2p5_base.py +++ b/src/timesfm/timesfm_2p5/timesfm_2p5_base.py @@ -16,7 +16,9 @@ import dataclasses from typing import Any, Callable + import numpy as np + from .. import configs ResidualBlockConfig = configs.ResidualBlockConfig @@ -85,43 +87,44 @@ class TimesFM_2p5_200M_Definition: output_patch_len: int = 128 output_quantile_len: int = 1024 quantiles: list[float] = dataclasses.field( - default_factory=lambda: [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] + default_factory=lambda: [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] ) decode_index: int = 5 tokenizer: ResidualBlockConfig = ResidualBlockConfig( - input_dims=64, - hidden_dims=1280, - output_dims=1280, - use_bias=True, - activation="swish", + input_dims=64, + hidden_dims=1280, + output_dims=1280, + use_bias=True, + activation="swish", ) stacked_transformers: StackedTransformersConfig = StackedTransformersConfig( - num_layers=20, - transformer=TransformerConfig( - model_dims=1280, - hidden_dims=1280, - num_heads=16, - attention_norm="rms", - feedforward_norm="rms", - qk_norm="rms", - use_bias=False, - use_rotary_position_embeddings=True, - ff_activation="swish", - ), + num_layers=20, + transformer=TransformerConfig( + model_dims=1280, + hidden_dims=1280, + num_heads=16, + attention_norm="rms", + feedforward_norm="rms", + qk_norm="rms", + use_bias=False, + use_rotary_position_embeddings=True, + ff_activation="swish", + fuse_qkv=True, + ), ) output_projection_point: ResidualBlockConfig = ResidualBlockConfig( - input_dims=1280, - hidden_dims=1280, - output_dims=1280, - use_bias=False, - activation="swish", + input_dims=1280, + hidden_dims=1280, + output_dims=1280, + use_bias=False, + activation="swish", ) output_projection_quantiles: ResidualBlockConfig = ResidualBlockConfig( - input_dims=1280, - hidden_dims=1280, - output_dims=10240, - use_bias=False, - activation="swish", + input_dims=1280, + hidden_dims=1280, + output_dims=10240, + use_bias=False, + activation="swish", ) @@ -147,7 +150,7 @@ class TimesFM_2p5: raise NotImplementedError() def forecast( - self, horizon: int, inputs: list[np.ndarray] + self, horizon: int, inputs: list[np.ndarray] ) -> tuple[np.ndarray, np.ndarray]: """Forecasts the time series.""" if self.compiled_decode is None: @@ -179,9 +182,7 @@ class TimesFM_2p5: idx += 1 if idx == self.global_batch_size: idx = 0 - point_forecast, quantile_forecast = self.compiled_decode( - horizon, values, masks - ) + point_forecast, quantile_forecast = self.compiled_decode(horizon, values, masks) output_points.append(point_forecast) output_quantiles.append(quantile_forecast) values = [] diff --git a/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py b/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py index 9d201d5..897dbe7 100644 --- a/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py +++ b/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py @@ -22,7 +22,7 @@ from typing import Dict, Optional, Sequence, Union import numpy as np import torch from huggingface_hub import ModelHubMixin, hf_hub_download -from safetensors.torch import load_file +from safetensors.torch import load_file, save_file from torch import nn from .. import configs @@ -75,11 +75,19 @@ class TimesFM_2p5_200M_torch_module(nn.Module): self.device = torch.device("cpu") self.device_count = 1 - def load_checkpoint(self, path: str): + def load_checkpoint(self, path: str, **kwargs): """Loads a PyTorch TimesFM model from a checkpoint.""" tensors = load_file(path) - self.load_state_dict(tensors) + self.load_state_dict(tensors, strict=True) self.to(self.device) + torch_compile = True + if "torch_compile" in kwargs: + torch_compile = kwargs["torch_compile"] + if torch_compile: + print("Compiling model...") + self = torch.compile(self) + + self.eval() def forward( self, @@ -113,9 +121,6 @@ class TimesFM_2p5_200M_torch_module(nn.Module): def decode(self, horizon: int, inputs, masks): """Decodes the time series.""" - inputs = inputs.to(self.device) - masks = masks.to(self.device) - with torch.no_grad(): batch_size, context = inputs.shape[0], inputs.shape[1] num_decode_steps = (horizon - 1) // self.o @@ -307,9 +312,20 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5, ModelHubMixin): logging.info("Loading checkpoint from: %s", model_file_path) # Load the weights into the model. - instance.model.load_checkpoint(model_file_path) + instance.model.load_checkpoint(model_file_path, **model_kwargs) return instance + def _save_pretrained(self, save_directory: Union[str, Path]): + """ + Saves the model's state dictionary to a safetensors file. This method + is called by the `save_pretrained` method from `ModelHubMixin`. + """ + if not os.path.exists(save_directory): + os.makedirs(save_directory) + + weights_path = os.path.join(save_directory, "model.safetensors") + save_file(self.model.state_dict(), weights_path) + def compile(self, forecast_config: configs.ForecastConfig, **kwargs) -> None: """Attempts to compile the model for fast decoding. @@ -319,9 +335,6 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5, ModelHubMixin): forecast_config: Configuration for forecasting flags. **kwargs: Additional keyword arguments to pass to model.compile(). """ - - if kwargs.get("backend", None) is not None: - self.model.compile(**kwargs) self.global_batch_size = ( forecast_config.per_core_batch_size * self.model.device_count ) @@ -363,8 +376,10 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5, ModelHubMixin): f"Horizon must be less than the max horizon. {horizon} > {fc.max_horizon}." ) - inputs = torch.Tensor(np.array(inputs)).to(self.model.device) - masks = torch.Tensor(np.array(masks)).to(self.model.device).to(torch.bool) + inputs = ( + torch.from_numpy(np.array(inputs)).to(self.model.device).to(torch.float32) + ) + masks = torch.from_numpy(np.array(masks)).to(self.model.device).to(torch.bool) batch_size = inputs.shape[0] if fc.infer_is_positive: diff --git a/src/timesfm/torch/transformer.py b/src/timesfm/torch/transformer.py index df7c6d0..a2c12da 100644 --- a/src/timesfm/torch/transformer.py +++ b/src/timesfm/torch/transformer.py @@ -18,12 +18,11 @@ import math from typing import Callable import torch -from torch import nn import torch.nn.functional as F +from torch import nn from .. import configs -from . import normalization -from . import util +from . import normalization, util LayerNorm = nn.LayerNorm RMSNorm = normalization.RMSNorm @@ -31,26 +30,26 @@ DecodeCache = util.DecodeCache def make_attn_mask( - query_length: int, - num_all_masked_kv: torch.Tensor, - query_index_offset: torch.Tensor | None = None, - kv_length: int = 0, + query_length: int, + num_all_masked_kv: torch.Tensor, + query_index_offset: torch.Tensor | None = None, + kv_length: int = 0, ) -> torch.Tensor: """Makes attention mask.""" if kv_length == 0: kv_length = query_length q_index = torch.arange(query_length, device=num_all_masked_kv.device)[ - None, None, :, None + None, None, :, None ] if query_index_offset is not None: q_index = q_index + query_index_offset[:, None, None, None] kv_index = torch.arange(kv_length, device=num_all_masked_kv.device)[ - None, None, None, : + None, None, None, : ] return torch.logical_and( - q_index >= kv_index, - kv_index >= num_all_masked_kv[:, None, None, None], + q_index >= kv_index, + kv_index >= num_all_masked_kv[:, None, None, None], ) @@ -58,10 +57,10 @@ class RotaryPositionalEmbedding(nn.Module): """Rotary positional embedding.""" def __init__( - self, - embedding_dims: int, - min_timescale: float = 1.0, - max_timescale: float = 10000.0, + self, + embedding_dims: int, + min_timescale: float = 1.0, + max_timescale: float = 10000.0, ): super().__init__() self.embedding_dims = embedding_dims @@ -69,31 +68,30 @@ class RotaryPositionalEmbedding(nn.Module): self.max_timescale = max_timescale def forward( - self, - inputs: torch.Tensor, - position: torch.Tensor | None = None, + self, + inputs: torch.Tensor, + position: torch.Tensor | None = None, ): """Generates a JTensor of sinusoids with different frequencies.""" if self.embedding_dims != inputs.shape[-1]: raise ValueError( - "The embedding dims of the rotary position embedding" - "must match the hidden dimension of the inputs." + "The embedding dims of the rotary position embedding" + "must match the hidden dimension of the inputs." ) half_embedding_dim = self.embedding_dims // 2 fraction = ( - 2 - * torch.arange(0, half_embedding_dim, device=inputs.device) - / self.embedding_dims + 2 + * torch.arange(0, half_embedding_dim, device=inputs.device) + / self.embedding_dims ) timescale = ( - self.min_timescale - * (self.max_timescale / self.min_timescale) ** fraction + self.min_timescale * (self.max_timescale / self.min_timescale) ** fraction ).to(inputs.device) if position is None: seq_length = inputs.shape[1] - position = torch.arange( - seq_length, dtype=torch.float32, device=inputs.device - )[None, :] + position = torch.arange(seq_length, dtype=torch.float32, device=inputs.device)[ + None, : + ] if len(inputs.shape) == 4: position = position[..., None, None] @@ -114,16 +112,16 @@ class RotaryPositionalEmbedding(nn.Module): def _dot_product_attention( - query, - key, - value, - mask=None, + query, + key, + value, + mask=None, ): """Computes dot-product attention given query, key, and value.""" attn_weights = torch.einsum("...qhd,...khd->...hqk", query, key) if mask is not None: attn_weights = torch.where( - mask, attn_weights, -torch.finfo(attn_weights.dtype).max / 2 + mask, attn_weights, -torch.finfo(attn_weights.dtype).max / 2 ) attn_weights = F.softmax(attn_weights, dim=-1) @@ -131,6 +129,28 @@ def _dot_product_attention( return torch.einsum("...hqk,...khd->...qhd", attn_weights, value) +def _torch_dot_product_attention(query, key, value, mask=None): + """ + Performs the exact same (unscaled) attention as the above function, + but using the fast and fused F.scaled_dot_product_attention kernel. + """ + + # 1. Permute inputs from (B, L, H, D) to the expected (B, H, L, D) + query = query.permute(0, 2, 1, 3) + key = key.permute(0, 2, 1, 3) + value = value.permute(0, 2, 1, 3) + + # 2. Call the fused attention kernel + # - Pass the mask to `attn_mask`. + # - Set `scale=1.0` to disable the default 1/sqrt(d_k) scaling. + output = F.scaled_dot_product_attention(query, key, value, attn_mask=mask, scale=1.0) + + # 3. Permute the output back to the original (B, L, H, D) layout + output = output.permute(0, 2, 1, 3) + + return output + + class PerDimScale(nn.Module): """Per-dimension scaling.""" @@ -141,7 +161,7 @@ class PerDimScale(nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: scale_factor = ( - 1.442695041 / math.sqrt(self.num_dims) * F.softplus(self.per_dim_scale) + 1.442695041 / math.sqrt(self.num_dims) * F.softplus(self.per_dim_scale) ) return x * scale_factor @@ -150,15 +170,16 @@ class MultiHeadAttention(nn.Module): """Multi-head attention.""" def __init__( - self, - num_heads: int, - in_features: int, - *, - use_per_dim_scale: bool = True, - use_rotary_position_embeddings: bool = True, - use_bias: bool = False, - attention_fn: Callable[..., torch.Tensor] = _dot_product_attention, - qk_norm: str = "rms", + self, + num_heads: int, + in_features: int, + *, + use_per_dim_scale: bool = True, + use_rotary_position_embeddings: bool = True, + use_bias: bool = False, + attention_fn: Callable[..., torch.Tensor] = _torch_dot_product_attention, + qk_norm: str = "rms", + fuse_qkv: bool = False, ): super().__init__() self.num_heads = num_heads @@ -167,16 +188,20 @@ class MultiHeadAttention(nn.Module): self.use_bias = use_bias self.attention_fn = attention_fn self.qk_norm = qk_norm + self.fuse_qkv = fuse_qkv if self.in_features % self.num_heads != 0: raise ValueError( - f"Memory dimension ({self.in_features}) must be divisible by " - f"'num_heads' heads ({self.num_heads})." + f"Memory dimension ({self.in_features}) must be divisible by " + f"'num_heads' heads ({self.num_heads})." ) - self.query = nn.Linear(self.in_features, self.in_features, bias=use_bias) - self.key = nn.Linear(self.in_features, self.in_features, bias=use_bias) - self.value = nn.Linear(self.in_features, self.in_features, bias=use_bias) + if self.fuse_qkv: + self.qkv_proj = nn.Linear(self.in_features, 3 * self.in_features, bias=use_bias) + else: + self.query = nn.Linear(self.in_features, self.in_features, bias=use_bias) + self.key = nn.Linear(self.in_features, self.in_features, bias=use_bias) + self.value = nn.Linear(self.in_features, self.in_features, bias=use_bias) self.out = nn.Linear(self.in_features, self.in_features, bias=use_bias) if self.qk_norm == "rms": @@ -189,7 +214,7 @@ class MultiHeadAttention(nn.Module): self.use_rotary_position_embeddings = use_rotary_position_embeddings if self.use_rotary_position_embeddings: self.rotary_position_embedding = RotaryPositionalEmbedding( - embedding_dims=self.head_dim, + embedding_dims=self.head_dim, ) self.use_per_dim_scale = use_per_dim_scale @@ -197,41 +222,41 @@ class MultiHeadAttention(nn.Module): self.per_dim_scale = PerDimScale(num_dims=self.head_dim) def forward( - self, - inputs_q: torch.Tensor, - *, - decode_cache: DecodeCache | None = None, - patch_mask: torch.Tensor | None = None, + self, + inputs_q: torch.Tensor, + *, + decode_cache: DecodeCache | None = None, + patch_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, DecodeCache | None]: b, n_patches, _ = inputs_q.shape if patch_mask is None: - patch_mask = torch.zeros( - b, n_patches, dtype=torch.bool, device=inputs_q.device - ) + patch_mask = torch.zeros(b, n_patches, dtype=torch.bool, device=inputs_q.device) - query = self.query(inputs_q).view( - b, n_patches, self.num_heads, self.head_dim - ) - key = self.key(inputs_q).view(b, n_patches, self.num_heads, self.head_dim) - value = self.value(inputs_q).view( - b, n_patches, self.num_heads, self.head_dim - ) + if self.fuse_qkv: + qkv = self.qkv_proj(inputs_q) + query, key, value = torch.chunk(qkv, 3, dim=-1) + query = query.view(b, n_patches, self.num_heads, self.head_dim) + key = key.view(b, n_patches, self.num_heads, self.head_dim) + value = value.view(b, n_patches, self.num_heads, self.head_dim) + else: + query = self.query(inputs_q).view(b, n_patches, self.num_heads, self.head_dim) + key = self.key(inputs_q).view(b, n_patches, self.num_heads, self.head_dim) + value = self.value(inputs_q).view(b, n_patches, self.num_heads, self.head_dim) if decode_cache is None: num_masked = torch.sum(patch_mask.to(torch.int32), dim=-1) next_index = torch.zeros_like(num_masked, dtype=torch.int32) else: num_masked = ( - torch.sum(patch_mask.to(torch.int32), dim=-1) - + decode_cache.num_masked + torch.sum(patch_mask.to(torch.int32), dim=-1) + decode_cache.num_masked ) next_index = decode_cache.next_index.clone() if self.use_rotary_position_embeddings: position = ( - torch.arange(n_patches, device=inputs_q.device)[None, :] - + next_index[:, None] - - num_masked[:, None] + torch.arange(n_patches, device=inputs_q.device)[None, :] + + next_index[:, None] + - num_masked[:, None] ) query = self.rotary_position_embedding(query, position) key = self.rotary_position_embedding(key, position) @@ -244,32 +269,36 @@ class MultiHeadAttention(nn.Module): if decode_cache is not None: _, decode_cache_size, _, _ = decode_cache.value.shape - for i in range(b): - start = decode_cache.next_index[i] - end = start + n_patches - decode_cache.key[i, start:end] = key[i].clone() - decode_cache.value[i, start:end] = value[i].clone() - key = decode_cache.key.clone() - value = decode_cache.value.clone() + + start = decode_cache.next_index[0] + end = start + n_patches + + # Perform a single, vectorized slice assignment for the entire batch. + # This is vastly more efficient than a Python for-loop. + + decode_cache.key[:, start:end] = key + decode_cache.value[:, start:end] = value + + key = decode_cache.key + value = decode_cache.value decode_cache.next_index += n_patches decode_cache.num_masked = num_masked attn_mask = make_attn_mask( - query_length=n_patches, - num_all_masked_kv=num_masked, - query_index_offset=next_index, - kv_length=decode_cache_size, + query_length=n_patches, + num_all_masked_kv=num_masked, + query_index_offset=next_index, + kv_length=decode_cache_size, ) else: - attn_mask = make_attn_mask( - query_length=n_patches, num_all_masked_kv=num_masked - ) + attn_mask = make_attn_mask(query_length=n_patches, num_all_masked_kv=num_masked) x = self.attention_fn( - query, - key, - value, - mask=attn_mask, + query, + key, + value, + mask=attn_mask, ) + x = x.reshape(b, n_patches, self.in_features) out = self.out(x) return out, decode_cache @@ -289,11 +318,12 @@ class Transformer(nn.Module): raise ValueError(f"Layer norm: {config.attention_norm} not supported.") self.attn = MultiHeadAttention( - num_heads=config.num_heads, - in_features=config.model_dims, - use_per_dim_scale=True, - use_rotary_position_embeddings=config.use_rotary_position_embeddings, - qk_norm=config.qk_norm, + num_heads=config.num_heads, + in_features=config.model_dims, + use_per_dim_scale=True, + use_rotary_position_embeddings=config.use_rotary_position_embeddings, + qk_norm=config.qk_norm, + fuse_qkv=config.fuse_qkv, ) if config.feedforward_norm == "rms": @@ -303,14 +333,14 @@ class Transformer(nn.Module): raise ValueError(f"Layer norm: {config.feedforward_norm} not supported.") self.ff0 = nn.Linear( - in_features=config.model_dims, - out_features=config.hidden_dims, - bias=config.use_bias, + in_features=config.model_dims, + out_features=config.hidden_dims, + bias=config.use_bias, ) self.ff1 = nn.Linear( - in_features=config.hidden_dims, - out_features=config.model_dims, - bias=config.use_bias, + in_features=config.hidden_dims, + out_features=config.model_dims, + bias=config.use_bias, ) if config.ff_activation == "relu": self.activation = nn.ReLU() @@ -322,21 +352,19 @@ class Transformer(nn.Module): raise ValueError(f"Activation: {config.ff_activation} not supported.") def forward( - self, - input_embeddings: torch.Tensor, - patch_mask: torch.Tensor, - decode_cache: DecodeCache | None = None, + self, + input_embeddings: torch.Tensor, + patch_mask: torch.Tensor, + decode_cache: DecodeCache | None = None, ) -> tuple[torch.Tensor, DecodeCache | None]: attn_output, decode_cache = self.attn( - inputs_q=self.pre_attn_ln(input_embeddings), - decode_cache=decode_cache, - patch_mask=patch_mask, + inputs_q=self.pre_attn_ln(input_embeddings), + decode_cache=decode_cache, + patch_mask=patch_mask, ) attn_output = self.post_attn_ln(attn_output) + input_embeddings output_embeddings = ( - self.post_ff_ln( - self.ff1(self.activation(self.ff0(self.pre_ff_ln(attn_output)))) - ) - + attn_output + self.post_ff_ln(self.ff1(self.activation(self.ff0(self.pre_ff_ln(attn_output))))) + + attn_output ) return output_embeddings, decode_cache