Timestral vibe

THIRD_PARTY_LICENSES.md

@@ -34,16 +34,6 @@ Its weights are not bundled with this project and may be downloaded only after t
 The [`sdm/models/timesfm3/`](sdm/models/timesfm3/) implementation contains code derived from [`TimesFM`](https://github.com/google-research/timesfm).
 Its weights are not bundled with this project and may be downloaded only after the user accepts their separate license.
 
-## Timestral
-
-- Source: https://github.com/mistralai/timestral
-- Code license: Apache License 2.0
-- License terms: [`sdm/models/timestral/LICENSE`](sdm/models/timestral/LICENSE)
-- Optional pretrained weights: https://huggingface.co/collections/mistralai/timestral
-- Weights license: [Modified MIT License](https://huggingface.co/mistralai/Timestral-209M/blob/main/LICENSE)
-
-The [`sdm/models/timestral/`](sdm/models/timestral/) implementation contain code derived from [`Timestral`](https://github.com/mistralai/timestral).
-
 ## DPM-Solver
 
 - Source: https://github.com/LuChengTHU/dpm-solver
@@ -67,3 +57,12 @@ Portions of [`CONTRIBUTING.md`](CONTRIBUTING.md) are adapted from the [PyTorch c
 ## Contributor Covenant
 
 [`CODE_OF_CONDUCT.md`](CODE_OF_CONDUCT.md) is adapted from [Contributor Covenant version 1.4](https://www.contributor-covenant.org/version/1/4/code-of-conduct/), distributed under the MIT License. The complete copyright notice and license terms are distributed in [`third_party/contributor-covenant/LICENSE`](third_party/contributor-covenant/LICENSE).
+
+## Timestral
+
+- Source: [mistralai/timestral](https://github.com/mistralai/timestral), revision `a534bdbd244d64f3d9d743f46c7bf6fc330ed4ff`.
+- Source license: [Apache-2.0](sdm/models/timestral/LICENSE).
+- Weights license: [Modified MIT](sdm/models/timestral/WEIGHTS_LICENSE), including a company revenue restriction.
+- Checkpoints: `mistralai/Timestral-5M` at `bdd20cfc15a3fea0e4cb5ccef37f8d3ca7aa335f`, `Timestral-28M` at `180698bbf5a50dda3516aa6419798437fb03a430`, `Timestral-123M` at `5a147287addda85f6acd8da3fb8858fe9af7e8f7`, and `Timestral-209M` at `9874d252ab0aae16f7f7c999f98f98d6ad7d8365`.
+
+The `sdm/models/timestral/` numerical implementation and the `sdm.nn.MonotoneQuantiles` and `sdm.nn.CausalPatchNorm` modules contain code derived from Timestral. The SDM wrapper preserves the upstream default forecast, including categorical encoding, sign averaging, advancing normalization frames, and nonnegative clipping. Prediction parity is checked with the pinned 5M weights; larger checkpoints share the implementation but have not been parity-tested.

benchmark/forecasting/timestral_parity.md

@@ -0,0 +1,71 @@
+# Timestral 5M parity
+
+Run from the repository root with the upstream reference installed:
+
+```bash
+pip install git+https://github.com/mistralai/timestral.git@a534bdbd244d64f3d9d743f46c7bf6fc330ed4ff
+python -m benchmark.forecasting.timestral_parity --device cpu --repeats 20
+python -m benchmark.forecasting.timestral_parity --device cuda --repeats 20
+```
+
+The script verifies the installed source revision and uses `mistralai/Timestral-5M` revision `bdd20cfc15a3fea0e4cb5ccef37f8d3ca7aa335f`. Downloading weights is subject to their [Modified MIT license](../../sdm/models/timestral/WEIGHTS_LICENSE).
+
+It compares upstream `TimestralForecaster.forecast` with SDM's public `Timestral.forward`, using default sign averaging, nonnegative clipping, and normalization refinement. It also compares SDM's fitted `predict` with its one-shot output. Deterministic inputs cover a short patch, univariate and multivariate targets, independent batches, missing and entirely missing targets, constant and positive targets, numerical and binary covariates, categorical past and known-future covariates, unseen and missing categories, missing target history with categorical covariates, a 1,030-step forecast, and truncation of a 32,775-step context.
+
+Recorded on 2026-10-18 with PyTorch `2.14.0+cu130`, a Cortex-X925 CPU using two PyTorch threads, and an NVIDIA GB10 GPU:
+
+| Comparison                    | Result                                                          |
+| ----------------------------- | --------------------------------------------------------------- |
+| CPU, horizons up to 65 steps  | Maximum absolute error `3.25e-5`                                |
+| CPU, 1,030-step horizon       | Maximum absolute error `0.0801`; maximum scaled error `1.82e-5` |
+| CPU, all 15 cases             | Pass `atol=rtol=2e-5`                                           |
+| CUDA, horizons up to 65 steps | Maximum absolute error `1.96e-5`                                |
+| CUDA, 1,030-step horizon      | Maximum absolute error `0.0547`; maximum scaled error `1.66e-5` |
+| CUDA, all 15 cases            | Pass `atol=rtol=1e-4`                                           |
+
+Scaled error is `abs(actual - reference) / (1 + abs(reference))`. The long-horizon example produces large tail quantiles; its absolute error should be interpreted alongside its scaled error. Timestral uses SDM's shared `TransformerBlock`, including `Attention`, `RotaryEmbedding` with explicit positions, and input-dependent sigmoid output gating. Time and variate attention are separate residual blocks, with one normalized `SwiGLU` MLP after the final attention step of each upstream layer. A single full mask replaces upstream's split context/forecast attention calls. Checkpoint loading permutes query/key channels to match SDM's partial-rotation layout while leaving values unchanged. This preserves the attention computation mathematically; reordered floating-point reductions introduce small differences on CPU and CUDA.
+
+Only 5M was checked against pretrained weights. The 28M, 123M, and 209M variants use pinned configurations and checkpoint loading but have not been parity-tested. Parity is scoped to float32; lower precision is unvalidated.
+
+The integration needs no upstream package at runtime. Optional pretrained pytest comparisons skip when upstream or cached 5M weights are unavailable. The ordinary public-contract and neural-block tests use random weights and run offline. SDM represents categorical values with category indices; for multivariate targets those indices become the raw numerical covariates, matching upstream when supplied the same codes.
+
+The default recipe uses `AlignCategories(min_frequency=0, unknown="append")` to align independently tensorized vocabularies. Context codes and unused dictionary entries are retained; new query categories are appended to the returned vocabulary without changing fitted state. Numeric vocabularies use tensor lookups; string vocabularies use SDM's batched relational join, with cuDF on CUDA when installed and the existing CPU fallback otherwise. Forecast tensor execution and target encoding stay on the input device.
+
+## Runtime versus upstream
+
+The [runtime benchmark](timestral_runtime.py) checks prediction parity before timing each workload. Reproduce it with:
+
+```bash
+python -m benchmark.forecasting.timestral_runtime --device cpu --output benchmark/forecasting/timestral_runtime_cpu.json
+python -m benchmark.forecasting.timestral_runtime --device cuda --output benchmark/forecasting/timestral_runtime_cuda.json
+```
+
+Measured before moving category alignment into the recipe, on the same CPU/GPU and PyTorch version described above, with two CPU threads, float32, eager execution, one series per call, and the default forecast options. Values are median milliseconds across 30 calls after five warmups. Each pair alternates execution order; CUDA is synchronized before and after each individual timed call. These are public API latencies, including preprocessing and output construction, with inputs already on the target device. Loading, initial table construction, and warmup are excluded. CPU and CUDA benchmarks ran sequentially.
+
+| Workload               | Context / horizon | CPU SDM | CPU upstream | CUDA SDM | CUDA upstream |
+| ---------------------- | ----------------- | ------- | ------------ | -------- | ------------- |
+| Univariate             | 128 / 96          | 79.03   | 68.79        | 9.25     | 9.09          |
+| Univariate             | 1,024 / 96        | 82.16   | 71.09        | 10.03    | 9.84          |
+| Three targets          | 1,024 / 96        | 85.71   | 73.62        | 10.56    | 10.32         |
+| Numerical covariates   | 1,024 / 96        | 92.55   | 80.50        | 12.06    | 11.47         |
+| Categorical covariates | 1,024 / 96        | 96.02   | 82.22        | 14.92    | 11.90         |
+| Long horizon           | 79 / 1,030        | 163.30  | 143.34       | 23.54    | 23.65         |
+
+SDM takes roughly 14–17% longer on CPU in these workloads. On CUDA, most workloads are within 5% of upstream; the categorical case takes roughly 25% longer. The categorical comparison includes SDM's dictionary alignment, while upstream accepts numeric category codes. Neither implementation is compiled or uses lower precision. No attribution of the remaining runtime difference to individual operations has been profiled.
+
+The [CPU report](timestral_runtime_cpu.json) and [CUDA report](timestral_runtime_cuda.json) include all individual samples, quartiles, revisions, and benchmark settings.
+
+## Shared transformer refactor
+
+Remeasured on 2026-10-18 after replacing the private attention and residual blocks with `sdm.nn.TransformerBlock`. The same benchmark settings apply, with 20 timed calls after five warmups (`--repeats 20`). The CUDA baseline was measured immediately before this refactor; CPU and CUDA runs were sequential. All values below are median milliseconds.
+
+| Workload                            | CPU SDM | CPU upstream | CUDA SDM before | CUDA SDM after | CUDA upstream after |
+| ----------------------------------- | ------- | ------------ | --------------- | -------------- | ------------------- |
+| Univariate (128 / 96)               | 80.70   | 68.87        | 9.91            | 9.43           | 9.12                |
+| Univariate (1,024 / 96)             | 84.05   | 71.73        | 10.72           | 10.33          | 9.84                |
+| Three targets (1,024 / 96)          | 85.74   | 74.24        | 11.37           | 10.84          | 10.31               |
+| Numerical covariates (1,024 / 96)   | 92.79   | 80.73        | 12.88           | 12.11          | 11.41               |
+| Categorical covariates (1,024 / 96) | 94.89   | 82.06        | 14.39           | 13.47          | 11.86               |
+| Long horizon (79 / 1,030)           | 163.31  | 139.79       | 24.77           | 23.81          | 23.58               |
+
+The full-mask implementation showed no latency regression in these workloads. Before/after timings are separate runs, and upstream CUDA timings also decreased, so these measurements do not isolate a speedup caused by the refactor. The public 5M parity checks still pass on both devices at the tolerances above. Larger batches, longer covariate timelines, and lower precision remain unbenchmarked.

benchmark/forecasting/timestral_parity.py

@@ -0,0 +1,217 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Compare public Timestral forecasts against the pinned upstream 5M model.
+
+Install git+https://github.com/mistralai/timestral.git@a534bdbd244d64f3d9d743f46c7bf6fc330ed4ff
+then run python -m benchmark.forecasting.timestral_parity --device cpu
+or --device cuda. This downloads the 5M weights under their Modified MIT terms.
+"""
+
+import argparse
+import json
+from importlib.metadata import distribution
+from time import perf_counter
+
+import torch
+
+from sdm import CategoricalTensor, Stype, TableTensor
+from sdm.models import Timestral
+
+SOURCE_REVISION = "a534bdbd244d64f3d9d743f46c7bf6fc330ed4ff"
+WEIGHT_REVISION = "bdd20cfc15a3fea0e4cb5ccef37f8d3ca7aa335f"
+
+
+def make_case(
+    name: str,
+    device: str,
+    *,
+    context_length: int = 79,
+    prediction_length: int = 65,
+) -> tuple[TableTensor, TableTensor, TableTensor, dict]:
+    t, h, v, batch = context_length, prediction_length, 1, 1
+    if name == "multivariate":
+        v = 3
+    if name == "batched":
+        batch = 2
+    if name == "long_horizon":
+        h = 1030
+    if name == "context_cap":
+        t, h = 32775, 1
+    if name == "single_patch":
+        t, h = 7, 1
+    time = torch.arange(t + h, device=device).float()
+    target = (time[:t] / 9).sin().expand(batch, v, t).clone()
+    target += torch.arange(v, device=device)[None, :, None] * 0.2
+    target += torch.arange(batch, device=device)[:, None, None] * 0.3
+    if name == "positive":
+        target = target.abs() + 1
+    if name == "constant":
+        target.fill_(2)
+    if name in {"missing", "categorical_missing"}:
+        target[..., :8] = torch.nan
+        target[..., 12:15] = torch.nan
+    if name == "all_missing":
+        target.fill_(torch.nan)
+    reference = {}
+    context = TableTensor(size=(batch, t), device=device)
+    query = TableTensor(size=(batch, h), device=device)
+    if name in {
+        "numeric",
+        "indicator",
+        "categorical",
+        "categorical_multi",
+        "categorical_missing",
+    }:
+        if name == "categorical_multi":
+            target = target.expand(batch, 2, t).clone()
+            v = 2
+        past = (time[:t] / 13).cos().expand(batch, 1, t).clone()
+        known = (time / 7).sin().expand(batch, 1, t + h).clone()
+        if name == "indicator":
+            known = (known > 0).float()
+        if name.startswith("categorical"):
+            past = (time[:t].long() % 3).expand(batch, 1, t).clone()
+            known = (time.long() % 4).expand(batch, 1, t + h).clone()
+            past[..., 3:6] = -1
+            known[..., 3:6] = -1
+            known[..., t + 4 :] = 4  # Category first seen in the horizon.
+
+            def table(codes: torch.Tensor, columns: list[str]) -> TableTensor:
+                return TableTensor(
+                    columns={Stype.categorical: columns},
+                    categorical=CategoricalTensor(
+                        codes.transpose(-2, -1),
+                        [torch.arange(5, device=device) for _ in columns],
+                    ),
+                )
+
+            context = table(
+                torch.cat([known[..., :t], past], dim=1), ["known", "past"]
+            )
+            query = table(known[..., t:], ["known"])
+            reference = {
+                "categorical_past_covariates": past.float().masked_fill(
+                    past < 0, torch.nan
+                ),
+                "categorical_known_covariates": known.float().masked_fill(
+                    known < 0, torch.nan
+                ),
+            }
+        else:
+            past[..., 4:6] = torch.nan
+            context = TableTensor.from_tensor(
+                torch.cat([known[..., :t], past], dim=1).transpose(-2, -1),
+                columns=["known", "past"],
+            )
+            query = TableTensor.from_tensor(
+                known[..., t:].transpose(-2, -1), columns=["known"]
+            )
+            reference = {"past_covariates": past, "known_covariates": known}
+    targets = TableTensor.from_tensor(
+        target.transpose(-2, -1), columns=[f"target{i}" for i in range(v)]
+    )
+    return context, targets, query, reference
+
+
+def main() -> None:
+    from timestral import TimestralForecaster  # noqa: PLC0415
+
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--device", default="cpu")
+    parser.add_argument("--repeats", type=int, default=5)
+    args = parser.parse_args()
+    installed = json.loads(
+        distribution("timestral").read_text("direct_url.json") or "{}"
+    )
+    assert installed.get("vcs_info", {}).get("commit_id") == SOURCE_REVISION
+    torch.set_num_threads(2)
+    model = Timestral(device=args.device, accept_license=True)
+    upstream = TimestralForecaster.from_pretrained(
+        "mistralai/Timestral-5M",
+        revision=WEIGHT_REVISION,
+        device=args.device,
+    )
+    cases = [
+        "single_patch",
+        "univariate",
+        "multivariate",
+        "batched",
+        "missing",
+        "all_missing",
+        "constant",
+        "positive",
+        "numeric",
+        "indicator",
+        "categorical",
+        "categorical_multi",
+        "categorical_missing",
+        "long_horizon",
+        "context_cap",
+    ]
+    report = {}
+    for name in cases:
+        x, y, query, covariates = make_case(name, args.device)
+        target = y.numerical.transpose(-2, -1)
+        expected = (
+            upstream.forecast(target, query.size(-2), **covariates)
+            .quantiles.transpose(-3, -2)
+            .flatten(-2)
+        )
+        actual = model(x, y, query, num_estimators=1).numerical
+        tol = 1e-4 if args.device.startswith("cuda") else 2e-5
+        torch.testing.assert_close(actual, expected, atol=tol, rtol=tol)
+        model.fit(
+            x,
+            y,
+            context_only_columns=["past"]
+            if "past" in x.column_names
+            else None,
+            num_estimators=1,
+        )
+        torch.testing.assert_close(model.predict(query).numerical, actual)
+        report[name] = {
+            "max_absolute_error": (actual - expected).abs().max().item(),
+            "max_scaled_error": (
+                (actual - expected).abs() / (1 + expected.abs())
+            )
+            .max()
+            .item(),
+        }
+        print(name, report[name], flush=True)
+    x, y, query, _ = make_case("numeric", args.device)
+    for _ in range(3):
+        model(x, y, query, num_estimators=1)
+    if args.device.startswith("cuda"):
+        start, end = (
+            torch.cuda.Event(enable_timing=True),
+            torch.cuda.Event(enable_timing=True),
+        )
+        start.record()
+        for _ in range(args.repeats):
+            model(x, y, query, num_estimators=1)
+        end.record()
+        end.synchronize()
+        elapsed = start.elapsed_time(end) / args.repeats
+    else:
+        start_time = perf_counter()
+        for _ in range(args.repeats):
+            model(x, y, query, num_estimators=1)
+        elapsed = 1000 * (perf_counter() - start_time) / args.repeats
+    print(
+        json.dumps(
+            {
+                "device": args.device,
+                "torch": torch.__version__,
+                "source_revision": SOURCE_REVISION,
+                "weight_revision": WEIGHT_REVISION,
+                "cases": report,
+                "numeric_forward_ms": elapsed,
+            },
+            indent=2,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()