Add autocast to timeseries benchmarks

benchmark/timeseries/fev_bench/README.md

@@ -19,15 +19,17 @@ pip install . fev==0.10.0
   ```
 
   `TimesFM3` follows official evaluation with symmetric averaging over two estimators, contexts truncated to 15,360 time steps and up to 32 variates per forward pass.
+  Unlike the official `float32` evaluation, we use `bfloat16` mixed precision for improved efficiency with negligible differences in aggregate benchmark accuracy.
 
 > [!NOTE]
 > Weights of `TimesFM3` are distributed under the [TimesFM Non-Commercial License v1.0](https://huggingface.co/google/timesfm-3.0-pytorch/blob/main/LICENSE).
 > Review the license before running the `TimesFM3` benchmark, which will download its weights noninteractively.
 
+String-valued covariates are encoded into 100 numeric buckets using deterministic `CRC32` hashing.
 Results will be written to `benchmark/timeseries/fev_bench_out/<model>/`, and already completed configurations are skipped.
 
 Pass `--task` to run only selected fev-bench tasks, *e.g.*, `--task hospital --task proenfo_gfc12`.
-Pass `--mini` to run the 20-task fev-bench-mini subset:
+Pass `--mini` to run the 20-task fev-bench-mini subset.
 
 ## Evaluate
 

benchmark/timeseries/fev_bench/predictor.py

@@ -12,6 +12,7 @@
 
 import sdm
 from benchmark.timeseries.model import ModelConfig
+from sdm import TableTensor
 
 
 class SDMPredictor(fev.ForecastingModel):
@@ -135,12 +136,18 @@ def _fit_predict(self, task: fev.Task) -> list[datasets.DatasetDict]:
                                 [y_i, y[..., : size - y_i.size(-1)]], dim=-1
                             )
 
-                        out = self.model(
-                            x_context=sdm.TableTensor.from_tensor(x_context),
-                            y_context=sdm.TableTensor.from_tensor(y_i),
-                            x_query=sdm.TableTensor.from_tensor(x_query),
-                            num_estimators=self.config.num_estimators,
-                        )
+                        with torch.amp.autocast(
+                            self.device.type,
+                            dtype=self.config.autocast_dtype,
+                            enabled=x_context.is_cuda
+                            and self.config.autocast_dtype != torch.float32,
+                        ):
+                            out = self.model(
+                                x_context=TableTensor.from_tensor(x_context),
+                                y_context=TableTensor.from_tensor(y_i),
+                                x_query=TableTensor.from_tensor(x_query),
+                                num_estimators=self.config.num_estimators,
+                            )
                         pred = out.numerical
                         pred = pred.view(y.size(0), task.horizon, size, -1)
                         preds.append(pred)

benchmark/timeseries/gift_eval/README.md

@@ -20,6 +20,7 @@ pip install --no-deps "git+https://github.com/SalesforceAIResearch/gift-eval.git
   ```
 
   `TimesFM3` follows official evaluation with symmetric averaging over two estimators, contexts truncated to 15,360 time steps and up to 32 variates per forward pass.
+  Unlike the official `float32` evaluation, we use `bfloat16` mixed precision for improved efficiency with negligible differences in aggregate benchmark accuracy.
 
 > [!NOTE]
 > Weights of `TimesFM3` are distributed under the [TimesFM Non-Commercial License v1.0](https://huggingface.co/google/timesfm-3.0-pytorch/blob/main/LICENSE).

benchmark/timeseries/gift_eval/predictor.py

@@ -95,15 +95,21 @@ def predict(
                 if y_i.size(-1) < size:
                     y_i = torch.cat([y_i, y[..., : size - y_i.size(-1)]], -1)
 
-                out = self.model(
-                    x_context=sdm.TableTensor.from_tensor(x),
-                    y_context=sdm.TableTensor.from_tensor(y_i),
-                    x_query=sdm.TableTensor(
-                        size=(x.size(0), self.prediction_length),
-                        device=self.device,
-                    ),
-                    num_estimators=self.config.num_estimators,
-                )
+                with torch.amp.autocast(
+                    self.device.type,
+                    dtype=self.config.autocast_dtype,
+                    enabled=x.is_cuda
+                    and self.config.autocast_dtype != torch.float32,
+                ):
+                    out = self.model(
+                        x_context=sdm.TableTensor.from_tensor(x),
+                        y_context=sdm.TableTensor.from_tensor(y_i),
+                        x_query=sdm.TableTensor(
+                            size=(x.size(0), self.prediction_length),
+                            device=self.device,
+                        ),
+                        num_estimators=self.config.num_estimators,
+                    )
                 pred = out.numerical
                 pred = pred.view(y.size(0), self.prediction_length, size, -1)
                 preds.append(pred)

benchmark/timeseries/model.py

@@ -13,6 +13,7 @@
 class ModelConfig:
     name: str
     factory: Callable[[torch.device], sdm.models.ICLModel]
+    autocast_dtype: torch.dtype
     num_estimators: int = 1
     max_context_size: int | None = None
     max_variates: int | None = None
@@ -25,6 +26,7 @@ class ModelConfig:
             accept_license=True,
             device=device,
         ),
+        autocast_dtype=torch.bfloat16,
         num_estimators=2,
         max_context_size=15_360,
         max_variates=32,