freedak/timesfm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
Faster decoding related changes
This commit is contained in:
Rajat Sen
2025-10-01 15:23:52 -07:00
committed by GitHub
parent d5a3d5b872 a58bde615d
commit dc35e580d9
5 changed files with 202 additions and 153 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+4
View File
@@ -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
+1
View File
@@ -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
+33 -32
View File
@@ -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 = []
src/timesfm/timesfm_2p5/timesfm_2p5_torch.py
+27 -12
View File
@@ -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
+137 -109
View File
@@ -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