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

faster decoding related changes

Browse Source
This commit is contained in:
Rajat Sen
2025-10-01 16:50:27 +00:00
parent 2fded92687
commit a1c61b6f5d
5 changed files with 228 additions and 149 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_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
+81 -30
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,94 @@ 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=False,
),
)
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,
input_dims=1280,
hidden_dims=1280,
output_dims=10240,
use_bias=False,
activation="swish",
)
@dataclasses.dataclass(frozen=True)
class TimesFM_2p5_new_200M_Definition:
"""Framework-agnostic config of TimesFM 2.5."""
context_limit = 16384
input_patch_len: int = 32
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]
)
decode_index: int = 5
tokenizer: ResidualBlockConfig = ResidualBlockConfig(
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,
output_dims=10240,
num_heads=16,
attention_norm="rms",
feedforward_norm="rms",
qk_norm="rms",
use_bias=False,
activation="swish",
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",
)
output_projection_quantiles: ResidualBlockConfig = ResidualBlockConfig(
input_dims=1280,
hidden_dims=1280,
output_dims=10240,
use_bias=False,
activation="swish",
)
@@ -147,7 +200,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 +232,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 -10
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
@@ -78,7 +78,7 @@ class TimesFM_2p5_200M_torch_module(nn.Module):
def load_checkpoint(self, path: str):
"""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)
def forward(
@@ -113,9 +113,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
@@ -257,10 +254,16 @@ class TimesFM_2p5_200M_torch_module(nn.Module):
return outputs
class TimesFM_2p5_200M_torch_new_module(TimesFM_2p5_200M_torch_module):
"""TimesFM 2.5 with 200M parameters."""
config = timesfm_2p5_base.TimesFM_2p5_new_200M_Definition()
class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5, ModelHubMixin):
"""PyTorch implementation of TimesFM 2.5 with 200M parameters."""
model: nn.Module = TimesFM_2p5_200M_torch_module()
model: nn.Module = TimesFM_2p5_200M_torch_new_module()
@classmethod
def _from_pretrained(
@@ -310,6 +313,17 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5, ModelHubMixin):
instance.model.load_checkpoint(model_file_path)
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.
@@ -320,8 +334,9 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5, ModelHubMixin):
**kwargs: Additional keyword arguments to pass to model.compile().
"""
if kwargs.get("backend", None) is not None:
self.model.compile(**kwargs)
if forecast_config.torch_compile:
self.model = torch.compile(self.model)
self.model.eval()
self.global_batch_size = (
forecast_config.per_core_batch_size * self.model.device_count
)
@@ -363,8 +378,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
+115 -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)
@@ -141,7 +139,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 +148,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] = _dot_product_attention,
qk_norm: str = "rms",
fuse_qkv: bool = False,
):
super().__init__()
self.num_heads = num_heads
@@ -167,16 +166,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 +192,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 +200,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 +247,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 +296,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 +311,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 +330,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