Merge pull request #316 from google-research/rajat-dev

Faster decoding related changes

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(

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)

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
@@ -107,6 +109,7 @@ class TimesFM_2p5_200M_Definition:
       use_bias=False,
       use_rotary_position_embeddings=True,
       ff_activation="swish",
+      fuse_qkv=True,
     ),
   )
   output_projection_point: ResidualBlockConfig = ResidualBlockConfig(
@@ -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 = []

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:

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
@@ -86,14 +85,13 @@ class RotaryPositionalEmbedding(nn.Module):
       / 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]
@@ -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."""
 
@@ -157,8 +177,9 @@ class MultiHeadAttention(nn.Module):
     use_per_dim_scale: bool = True,
     use_rotary_position_embeddings: bool = True,
     use_bias: bool = False,
-      attention_fn: Callable[..., torch.Tensor] = _dot_product_attention,
+    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,6 +188,7 @@ 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(
@@ -174,6 +196,9 @@ class MultiHeadAttention(nn.Module):
         f"'num_heads' heads ({self.num_heads})."
       )
 
+    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)
@@ -205,25 +230,25 @@ class MultiHeadAttention(nn.Module):
   ) -> 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
-    )
+    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
-    )
+      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()
 
@@ -244,13 +269,18 @@ 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]
+
+      start = decode_cache.next_index[0]
       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()
+
+      # 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(
@@ -260,9 +290,7 @@ class MultiHeadAttention(nn.Module):
         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,
@@ -270,6 +298,7 @@ class MultiHeadAttention(nn.Module):
       value,
       mask=attn_mask,
     )
+
     x = x.reshape(b, n_patches, self.in_features)
     out = self.out(x)
     return out, decode_cache
@@ -294,6 +323,7 @@ class Transformer(nn.Module):
       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":
@@ -334,9 +364,7 @@ class Transformer(nn.Module):
     )
     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))))
-        )
+      self.post_ff_ln(self.ff1(self.activation(self.ff0(self.pre_ff_ln(attn_output)))))
       + attn_output
     )
     return output_embeddings, decode_cache