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 ### Code Example
```python ```python
import torch
import numpy as np import numpy as np
import timesfm import timesfm
torch.set_float32_matmul_precision("high")
model = TimesFM_2p5_200M_torch.from_pretrained("google/timesfm-2.5-200m-pytorch") model = TimesFM_2p5_200M_torch.from_pretrained("google/timesfm-2.5-200m-pytorch")
model.compile( model.compile(
src/timesfm/configs.py
+1
View File
@@ -94,6 +94,7 @@ class TransformerConfig:
use_bias: bool use_bias: bool
use_rotary_position_embeddings: bool use_rotary_position_embeddings: bool
ff_activation: Literal["relu", "swish", "none"] ff_activation: Literal["relu", "swish", "none"]
fuse_qkv: bool
@dataclasses.dataclass(frozen=True) @dataclasses.dataclass(frozen=True)
src/timesfm/timesfm_2p5/timesfm_2p5_base.py
+81 -30
View File
@@ -16,7 +16,9 @@
import dataclasses import dataclasses
from typing import Any, Callable from typing import Any, Callable
import numpy as np import numpy as np
from .. import configs from .. import configs
ResidualBlockConfig = configs.ResidualBlockConfig ResidualBlockConfig = configs.ResidualBlockConfig
@@ -85,43 +87,94 @@ class TimesFM_2p5_200M_Definition:
output_patch_len: int = 128 output_patch_len: int = 128
output_quantile_len: int = 1024 output_quantile_len: int = 1024
quantiles: list[float] = dataclasses.field( 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 decode_index: int = 5
tokenizer: ResidualBlockConfig = ResidualBlockConfig( tokenizer: ResidualBlockConfig = ResidualBlockConfig(
input_dims=64, input_dims=64,
hidden_dims=1280, hidden_dims=1280,
output_dims=1280, output_dims=1280,
use_bias=True, use_bias=True,
activation="swish", activation="swish",
) )
stacked_transformers: StackedTransformersConfig = StackedTransformersConfig( stacked_transformers: StackedTransformersConfig = StackedTransformersConfig(
num_layers=20, num_layers=20,
transformer=TransformerConfig( transformer=TransformerConfig(
model_dims=1280, model_dims=1280,
hidden_dims=1280, hidden_dims=1280,
num_heads=16, num_heads=16,
attention_norm="rms", attention_norm="rms",
feedforward_norm="rms", feedforward_norm="rms",
qk_norm="rms", qk_norm="rms",
use_bias=False, use_bias=False,
use_rotary_position_embeddings=True, use_rotary_position_embeddings=True,
ff_activation="swish", ff_activation="swish",
), fuse_qkv=False,
),
) )
output_projection_point: ResidualBlockConfig = ResidualBlockConfig( output_projection_point: ResidualBlockConfig = ResidualBlockConfig(
input_dims=1280, input_dims=1280,
hidden_dims=1280, hidden_dims=1280,
output_dims=1280, output_dims=1280,
use_bias=False, use_bias=False,
activation="swish", activation="swish",
) )
output_projection_quantiles: ResidualBlockConfig = ResidualBlockConfig( 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, hidden_dims=1280,
output_dims=10240, num_heads=16,
attention_norm="rms",
feedforward_norm="rms",
qk_norm="rms",
use_bias=False, 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() raise NotImplementedError()
def forecast( def forecast(
self, horizon: int, inputs: list[np.ndarray] self, horizon: int, inputs: list[np.ndarray]
) -> tuple[np.ndarray, np.ndarray]: ) -> tuple[np.ndarray, np.ndarray]:
"""Forecasts the time series.""" """Forecasts the time series."""
if self.compiled_decode is None: if self.compiled_decode is None:
@@ -179,9 +232,7 @@ class TimesFM_2p5:
idx += 1 idx += 1
if idx == self.global_batch_size: if idx == self.global_batch_size:
idx = 0 idx = 0
point_forecast, quantile_forecast = self.compiled_decode( point_forecast, quantile_forecast = self.compiled_decode(horizon, values, masks)
horizon, values, masks
)
output_points.append(point_forecast) output_points.append(point_forecast)
output_quantiles.append(quantile_forecast) output_quantiles.append(quantile_forecast)
values = [] 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 numpy as np
import torch import torch
from huggingface_hub import ModelHubMixin, hf_hub_download 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 torch import nn
from .. import configs from .. import configs
@@ -78,7 +78,7 @@ class TimesFM_2p5_200M_torch_module(nn.Module):
def load_checkpoint(self, path: str): def load_checkpoint(self, path: str):
"""Loads a PyTorch TimesFM model from a checkpoint.""" """Loads a PyTorch TimesFM model from a checkpoint."""
tensors = load_file(path) tensors = load_file(path)
self.load_state_dict(tensors) self.load_state_dict(tensors, strict=True)
self.to(self.device) self.to(self.device)
def forward( def forward(
@@ -113,9 +113,6 @@ class TimesFM_2p5_200M_torch_module(nn.Module):
def decode(self, horizon: int, inputs, masks): def decode(self, horizon: int, inputs, masks):
"""Decodes the time series.""" """Decodes the time series."""
inputs = inputs.to(self.device)
masks = masks.to(self.device)
with torch.no_grad(): with torch.no_grad():
batch_size, context = inputs.shape[0], inputs.shape[1] batch_size, context = inputs.shape[0], inputs.shape[1]
num_decode_steps = (horizon - 1) // self.o num_decode_steps = (horizon - 1) // self.o
@@ -257,10 +254,16 @@ class TimesFM_2p5_200M_torch_module(nn.Module):
return outputs 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): class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5, ModelHubMixin):
"""PyTorch implementation of TimesFM 2.5 with 200M parameters.""" """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 @classmethod
def _from_pretrained( 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) instance.model.load_checkpoint(model_file_path)
return instance 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: def compile(self, forecast_config: configs.ForecastConfig, **kwargs) -> None:
"""Attempts to compile the model for fast decoding. """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(). **kwargs: Additional keyword arguments to pass to model.compile().
""" """
if kwargs.get("backend", None) is not None: if forecast_config.torch_compile:
self.model.compile(**kwargs) self.model = torch.compile(self.model)
self.model.eval()
self.global_batch_size = ( self.global_batch_size = (
forecast_config.per_core_batch_size * self.model.device_count 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}." f"Horizon must be less than the max horizon. {horizon} > {fc.max_horizon}."
) )
inputs = torch.Tensor(np.array(inputs)).to(self.model.device) inputs = (
masks = torch.Tensor(np.array(masks)).to(self.model.device).to(torch.bool) 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] batch_size = inputs.shape[0]
if fc.infer_is_positive: if fc.infer_is_positive:
src/timesfm/torch/transformer.py
+115 -109
View File
@@ -18,12 +18,11 @@ import math
from typing import Callable from typing import Callable
import torch import torch
from torch import nn
import torch.nn.functional as F import torch.nn.functional as F
from torch import nn
from .. import configs from .. import configs
from . import normalization from . import normalization, util
from . import util
LayerNorm = nn.LayerNorm LayerNorm = nn.LayerNorm
RMSNorm = normalization.RMSNorm RMSNorm = normalization.RMSNorm
@@ -31,26 +30,26 @@ DecodeCache = util.DecodeCache
def make_attn_mask( def make_attn_mask(
query_length: int, query_length: int,
num_all_masked_kv: torch.Tensor, num_all_masked_kv: torch.Tensor,
query_index_offset: torch.Tensor | None = None, query_index_offset: torch.Tensor | None = None,
kv_length: int = 0, kv_length: int = 0,
) -> torch.Tensor: ) -> torch.Tensor:
"""Makes attention mask.""" """Makes attention mask."""
if kv_length == 0: if kv_length == 0:
kv_length = query_length kv_length = query_length
q_index = torch.arange(query_length, device=num_all_masked_kv.device)[ 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: if query_index_offset is not None:
q_index = q_index + query_index_offset[:, None, None, None] q_index = q_index + query_index_offset[:, None, None, None]
kv_index = torch.arange(kv_length, device=num_all_masked_kv.device)[ kv_index = torch.arange(kv_length, device=num_all_masked_kv.device)[
None, None, None, : None, None, None, :
] ]
return torch.logical_and( return torch.logical_and(
q_index >= kv_index, q_index >= kv_index,
kv_index >= num_all_masked_kv[:, None, None, None], kv_index >= num_all_masked_kv[:, None, None, None],
) )
@@ -58,10 +57,10 @@ class RotaryPositionalEmbedding(nn.Module):
"""Rotary positional embedding.""" """Rotary positional embedding."""
def __init__( def __init__(
self, self,
embedding_dims: int, embedding_dims: int,
min_timescale: float = 1.0, min_timescale: float = 1.0,
max_timescale: float = 10000.0, max_timescale: float = 10000.0,
): ):
super().__init__() super().__init__()
self.embedding_dims = embedding_dims self.embedding_dims = embedding_dims
@@ -69,31 +68,30 @@ class RotaryPositionalEmbedding(nn.Module):
self.max_timescale = max_timescale self.max_timescale = max_timescale
def forward( def forward(
self, self,
inputs: torch.Tensor, inputs: torch.Tensor,
position: torch.Tensor | None = None, position: torch.Tensor | None = None,
): ):
"""Generates a JTensor of sinusoids with different frequencies.""" """Generates a JTensor of sinusoids with different frequencies."""
if self.embedding_dims != inputs.shape[-1]: if self.embedding_dims != inputs.shape[-1]:
raise ValueError( raise ValueError(
"The embedding dims of the rotary position embedding" "The embedding dims of the rotary position embedding"
"must match the hidden dimension of the inputs." "must match the hidden dimension of the inputs."
) )
half_embedding_dim = self.embedding_dims // 2 half_embedding_dim = self.embedding_dims // 2
fraction = ( fraction = (
2 2
* torch.arange(0, half_embedding_dim, device=inputs.device) * torch.arange(0, half_embedding_dim, device=inputs.device)
/ self.embedding_dims / self.embedding_dims
) )
timescale = ( timescale = (
self.min_timescale self.min_timescale * (self.max_timescale / self.min_timescale) ** fraction
* (self.max_timescale / self.min_timescale) ** fraction
).to(inputs.device) ).to(inputs.device)
if position is None: if position is None:
seq_length = inputs.shape[1] seq_length = inputs.shape[1]
position = torch.arange( position = torch.arange(seq_length, dtype=torch.float32, device=inputs.device)[
seq_length, dtype=torch.float32, device=inputs.device None, :
)[None, :] ]
if len(inputs.shape) == 4: if len(inputs.shape) == 4:
position = position[..., None, None] position = position[..., None, None]
@@ -114,16 +112,16 @@ class RotaryPositionalEmbedding(nn.Module):
def _dot_product_attention( def _dot_product_attention(
query, query,
key, key,
value, value,
mask=None, mask=None,
): ):
"""Computes dot-product attention given query, key, and value.""" """Computes dot-product attention given query, key, and value."""
attn_weights = torch.einsum("...qhd,...khd->...hqk", query, key) attn_weights = torch.einsum("...qhd,...khd->...hqk", query, key)
if mask is not None: if mask is not None:
attn_weights = torch.where( 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) attn_weights = F.softmax(attn_weights, dim=-1)
@@ -141,7 +139,7 @@ class PerDimScale(nn.Module):
def forward(self, x: torch.Tensor) -> torch.Tensor: def forward(self, x: torch.Tensor) -> torch.Tensor:
scale_factor = ( 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 return x * scale_factor
@@ -150,15 +148,16 @@ class MultiHeadAttention(nn.Module):
"""Multi-head attention.""" """Multi-head attention."""
def __init__( def __init__(
self, self,
num_heads: int, num_heads: int,
in_features: int, in_features: int,
*, *,
use_per_dim_scale: bool = True, use_per_dim_scale: bool = True,
use_rotary_position_embeddings: bool = True, use_rotary_position_embeddings: bool = True,
use_bias: bool = False, use_bias: bool = False,
attention_fn: Callable[..., torch.Tensor] = _dot_product_attention, attention_fn: Callable[..., torch.Tensor] = _dot_product_attention,
qk_norm: str = "rms", qk_norm: str = "rms",
fuse_qkv: bool = False,
): ):
super().__init__() super().__init__()
self.num_heads = num_heads self.num_heads = num_heads
@@ -167,16 +166,20 @@ class MultiHeadAttention(nn.Module):
self.use_bias = use_bias self.use_bias = use_bias
self.attention_fn = attention_fn self.attention_fn = attention_fn
self.qk_norm = qk_norm self.qk_norm = qk_norm
self.fuse_qkv = fuse_qkv
if self.in_features % self.num_heads != 0: if self.in_features % self.num_heads != 0:
raise ValueError( raise ValueError(
f"Memory dimension ({self.in_features}) must be divisible by " f"Memory dimension ({self.in_features}) must be divisible by "
f"'num_heads' heads ({self.num_heads})." f"'num_heads' heads ({self.num_heads})."
) )
self.query = nn.Linear(self.in_features, self.in_features, bias=use_bias) if self.fuse_qkv:
self.key = nn.Linear(self.in_features, self.in_features, bias=use_bias) self.qkv_proj = nn.Linear(self.in_features, 3 * self.in_features, bias=use_bias)
self.value = nn.Linear(self.in_features, 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) self.out = nn.Linear(self.in_features, self.in_features, bias=use_bias)
if self.qk_norm == "rms": if self.qk_norm == "rms":
@@ -189,7 +192,7 @@ class MultiHeadAttention(nn.Module):
self.use_rotary_position_embeddings = use_rotary_position_embeddings self.use_rotary_position_embeddings = use_rotary_position_embeddings
if self.use_rotary_position_embeddings: if self.use_rotary_position_embeddings:
self.rotary_position_embedding = RotaryPositionalEmbedding( self.rotary_position_embedding = RotaryPositionalEmbedding(
embedding_dims=self.head_dim, embedding_dims=self.head_dim,
) )
self.use_per_dim_scale = use_per_dim_scale 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) self.per_dim_scale = PerDimScale(num_dims=self.head_dim)
def forward( def forward(
self, self,
inputs_q: torch.Tensor, inputs_q: torch.Tensor,
*, *,
decode_cache: DecodeCache | None = None, decode_cache: DecodeCache | None = None,
patch_mask: torch.Tensor | None = None, patch_mask: torch.Tensor | None = None,
) -> tuple[torch.Tensor, DecodeCache | None]: ) -> tuple[torch.Tensor, DecodeCache | None]:
b, n_patches, _ = inputs_q.shape b, n_patches, _ = inputs_q.shape
if patch_mask is None: if patch_mask is None:
patch_mask = torch.zeros( patch_mask = torch.zeros(b, n_patches, dtype=torch.bool, device=inputs_q.device)
b, n_patches, dtype=torch.bool, device=inputs_q.device
)
query = self.query(inputs_q).view( if self.fuse_qkv:
b, n_patches, self.num_heads, self.head_dim qkv = self.qkv_proj(inputs_q)
) query, key, value = torch.chunk(qkv, 3, dim=-1)
key = self.key(inputs_q).view(b, n_patches, self.num_heads, self.head_dim) query = query.view(b, n_patches, self.num_heads, self.head_dim)
value = self.value(inputs_q).view( key = key.view(b, n_patches, self.num_heads, self.head_dim)
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: if decode_cache is None:
num_masked = torch.sum(patch_mask.to(torch.int32), dim=-1) num_masked = torch.sum(patch_mask.to(torch.int32), dim=-1)
next_index = torch.zeros_like(num_masked, dtype=torch.int32) next_index = torch.zeros_like(num_masked, dtype=torch.int32)
else: else:
num_masked = ( num_masked = (
torch.sum(patch_mask.to(torch.int32), dim=-1) torch.sum(patch_mask.to(torch.int32), dim=-1) + decode_cache.num_masked
+ decode_cache.num_masked
) )
next_index = decode_cache.next_index.clone() next_index = decode_cache.next_index.clone()
if self.use_rotary_position_embeddings: if self.use_rotary_position_embeddings:
position = ( position = (
torch.arange(n_patches, device=inputs_q.device)[None, :] torch.arange(n_patches, device=inputs_q.device)[None, :]
+ next_index[:, None] + next_index[:, None]
- num_masked[:, None] - num_masked[:, None]
) )
query = self.rotary_position_embedding(query, position) query = self.rotary_position_embedding(query, position)
key = self.rotary_position_embedding(key, position) key = self.rotary_position_embedding(key, position)
@@ -244,32 +247,36 @@ class MultiHeadAttention(nn.Module):
if decode_cache is not None: if decode_cache is not None:
_, decode_cache_size, _, _ = decode_cache.value.shape _, 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 end = start + n_patches
decode_cache.key[i, start:end] = key[i].clone()
decode_cache.value[i, start:end] = value[i].clone() # Perform a single, vectorized slice assignment for the entire batch.
key = decode_cache.key.clone() # This is vastly more efficient than a Python for-loop.
value = decode_cache.value.clone()
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.next_index += n_patches
decode_cache.num_masked = num_masked decode_cache.num_masked = num_masked
attn_mask = make_attn_mask( attn_mask = make_attn_mask(
query_length=n_patches, query_length=n_patches,
num_all_masked_kv=num_masked, num_all_masked_kv=num_masked,
query_index_offset=next_index, query_index_offset=next_index,
kv_length=decode_cache_size, kv_length=decode_cache_size,
) )
else: else:
attn_mask = make_attn_mask( attn_mask = make_attn_mask(query_length=n_patches, num_all_masked_kv=num_masked)
query_length=n_patches, num_all_masked_kv=num_masked
)
x = self.attention_fn( x = self.attention_fn(
query, query,
key, key,
value, value,
mask=attn_mask, mask=attn_mask,
) )
x = x.reshape(b, n_patches, self.in_features) x = x.reshape(b, n_patches, self.in_features)
out = self.out(x) out = self.out(x)
return out, decode_cache return out, decode_cache
@@ -289,11 +296,12 @@ class Transformer(nn.Module):
raise ValueError(f"Layer norm: {config.attention_norm} not supported.") raise ValueError(f"Layer norm: {config.attention_norm} not supported.")
self.attn = MultiHeadAttention( self.attn = MultiHeadAttention(
num_heads=config.num_heads, num_heads=config.num_heads,
in_features=config.model_dims, in_features=config.model_dims,
use_per_dim_scale=True, use_per_dim_scale=True,
use_rotary_position_embeddings=config.use_rotary_position_embeddings, use_rotary_position_embeddings=config.use_rotary_position_embeddings,
qk_norm=config.qk_norm, qk_norm=config.qk_norm,
fuse_qkv=config.fuse_qkv,
) )
if config.feedforward_norm == "rms": if config.feedforward_norm == "rms":
@@ -303,14 +311,14 @@ class Transformer(nn.Module):
raise ValueError(f"Layer norm: {config.feedforward_norm} not supported.") raise ValueError(f"Layer norm: {config.feedforward_norm} not supported.")
self.ff0 = nn.Linear( self.ff0 = nn.Linear(
in_features=config.model_dims, in_features=config.model_dims,
out_features=config.hidden_dims, out_features=config.hidden_dims,
bias=config.use_bias, bias=config.use_bias,
) )
self.ff1 = nn.Linear( self.ff1 = nn.Linear(
in_features=config.hidden_dims, in_features=config.hidden_dims,
out_features=config.model_dims, out_features=config.model_dims,
bias=config.use_bias, bias=config.use_bias,
) )
if config.ff_activation == "relu": if config.ff_activation == "relu":
self.activation = nn.ReLU() self.activation = nn.ReLU()
@@ -322,21 +330,19 @@ class Transformer(nn.Module):
raise ValueError(f"Activation: {config.ff_activation} not supported.") raise ValueError(f"Activation: {config.ff_activation} not supported.")
def forward( def forward(
self, self,
input_embeddings: torch.Tensor, input_embeddings: torch.Tensor,
patch_mask: torch.Tensor, patch_mask: torch.Tensor,
decode_cache: DecodeCache | None = None, decode_cache: DecodeCache | None = None,
) -> tuple[torch.Tensor, DecodeCache | None]: ) -> tuple[torch.Tensor, DecodeCache | None]:
attn_output, decode_cache = self.attn( attn_output, decode_cache = self.attn(
inputs_q=self.pre_attn_ln(input_embeddings), inputs_q=self.pre_attn_ln(input_embeddings),
decode_cache=decode_cache, decode_cache=decode_cache,
patch_mask=patch_mask, patch_mask=patch_mask,
) )
attn_output = self.post_attn_ln(attn_output) + input_embeddings attn_output = self.post_attn_ln(attn_output) + input_embeddings
output_embeddings = ( output_embeddings = (
self.post_ff_ln( self.post_ff_ln(self.ff1(self.activation(self.ff0(self.pre_ff_ln(attn_output)))))
self.ff1(self.activation(self.ff0(self.pre_ff_ln(attn_output)))) + attn_output
)
+ attn_output
) )
return output_embeddings, decode_cache return output_embeddings, decode_cache