From e7233bc10e4a4ade4386b3dac1eb93b352846d33 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Wed, 21 Oct 2026 09:43:34 -0700 Subject: [PATCH 1/2] Add `KumoAnomaly` reconstruction and scoring Co-authored-by: Aditi Gautam --- sdm/models/kumo/timeseries/anomaly/model.py | 189 +++++++++++++++++- .../kumo/timeseries/anomaly/test_model.py | 88 ++++++++ 2 files changed, 274 insertions(+), 3 deletions(-) create mode 100644 test/models/kumo/timeseries/anomaly/test_model.py diff --git a/sdm/models/kumo/timeseries/anomaly/model.py b/sdm/models/kumo/timeseries/anomaly/model.py index 46a3534a7b..51178bc975 100644 --- a/sdm/models/kumo/timeseries/anomaly/model.py +++ b/sdm/models/kumo/timeseries/anomaly/model.py @@ -5,16 +5,25 @@ from __future__ import annotations +import copy from collections.abc import Sequence -from typing import ClassVar +from functools import partial +from typing import Any, ClassVar, cast import torch +import torch.nn.functional as F from torch import Tensor from sdm import Recipe, Stype, TableTensor +from sdm._inference import inference_mode +from sdm._warnings import warn_once from sdm.models import Callback, Model from sdm.models._huggingface import download_checkpoint +from sdm.models.kumo.timeseries.anomaly.denoiser import DiffusionDenoiser +from sdm.models.kumo.timeseries.anomaly.embedding import SideEmbedding from sdm.models.kumo.timeseries.anomaly.recipe import default_recipe +from sdm.models.kumo.timeseries.anomaly.sampling import DPMSolver +from sdm.models.kumo.timeseries.anomaly.schedule import DiffusionSchedule class KumoAnomaly(Model): @@ -96,7 +105,7 @@ def forward( Args: x: The feature tensor with shape ``[..., R, D]`` with ``R`` rows and ``D`` columns. - recipe: The custom recipe for pre- and post-processing. + recipe: The custom recipe for pre-processing. num_samples: The number of reconstruction samples per complementary masking strategy. Reconstructions are combined using the sample median. @@ -111,13 +120,187 @@ def forward( absolute reconstruction error across preprocessed features at each time step. """ - raise NotImplementedError + callbacks = () if callbacks is None else callbacks + requires_grad = self.training + requires_grad |= any(callback.requires_grad for callback in callbacks) + + param = next(self.parameters()) + if not isinstance(x, TableTensor): + x = TableTensor.from_tensor(x) + x = cast(TableTensor, x.to(param.device)) + + recipe = ( + self.default_recipe() if recipe is None else copy.deepcopy(recipe) + ) + with ( + torch.amp.autocast(param.device.type, enabled=False), + inference_mode("no_grad" if requires_grad else "inference"), + ): + x = recipe.features.fit_transform(x, generator=generator) + + for callback in callbacks: + x, _ = callback.on_query_preprocessing_end(self, x, None) + + invalid = x.active_stypes - self.supported_feature_stypes - {Stype.id} + if len(invalid) > 0: + stypes = ", ".join(f"{str(stype)!r}" for stype in invalid) + warn_once( + key="model-unsupported-feature-stypes", + message=( + f"{self.__class__.__name__!r} received unsupported " + f"feature stypes {stypes}. Columns with unsupported " + f"feature stypes will not be consumed by the model." + ), + ) + if x.numerical.size(-1) > self.model.num_features: + raise ValueError( + f"{self.__class__.__name__!r} supports at most " + f"{self.model.num_features} features after pre-processing " + f"(got {x.numerical.size(-1)})" + ) + + with inference_mode("grad" if requires_grad else "inference"): + out = self.model( + x.numerical.to(param.dtype), + num_samples=num_samples, + num_steps=num_steps, + generator=generator, + ) + out = TableTensor( + columns={Stype.numerical: ["score"]}, + numerical=out.unsqueeze(-1), + ) + for callback in callbacks: + out = callback.on_model_forward_end(self, out) + return cast(TableTensor, out.to(x.dtype)) class _KumoAnomaly(torch.nn.Module): def __init__( self, + num_features: int = 40, + time_channels: int = 256, + feature_channels: int = 64, + channels: int = 128, + num_layers: int = 6, + num_heads: int = 8, + embedding_channels: int = 256, + num_steps: int = 500, + window_size: int = 100, + num_blocks: int = 4, + batch_size: int = 32, device: torch.device | str | None = None, dtype: torch.dtype | None = None, ) -> None: super().__init__() + factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype} + + self.num_features = num_features + self.num_steps = num_steps + self.window_size = window_size + self.num_blocks = num_blocks + self.batch_size = batch_size + + self.side_embedding = SideEmbedding( + num_features=num_features, + time_channels=time_channels, + feature_channels=feature_channels, + **factory_kwargs, + ) + self.denoiser = DiffusionDenoiser( + channels=channels, + side_channels=time_channels + feature_channels, + num_steps=num_steps, + num_layers=num_layers, + num_heads=num_heads, + embedding_channels=embedding_channels, + **factory_kwargs, + ) + + def forward( + self, + x: Tensor, # [..., R, D] + num_samples: int, + num_steps: int, + generator: torch.Generator | None, + ) -> Tensor: # [..., R] + + *size, length, features = x.size() + x = x.reshape(-1, length, features) + x = F.pad(x, (0, self.num_features - features)) + + # Pad the last window by repeating the last row: + padding = -length % self.window_size + x = torch.cat([x, x[:, -1:].expand(-1, padding, -1)], dim=1) + windows = x.reshape(-1, self.window_size, self.num_features) + + schedule = DiffusionSchedule( + num_steps=self.num_steps, + beta_start=0.0001, + beta_end=0.01, + schedule="linear", + device="cpu", + ) + sampler = DPMSolver(schedule.beta.float(), num_steps).to(x.device) + + out = torch.cat( + [ + self.reconstruct(batch, sampler, num_samples, generator) + for batch in windows.split(self.batch_size) + ] + ) + out = out.reshape(x.size())[:, :length, :features] + out = (out - x[:, :length, :features]).abs().mean(dim=-1) + return out.reshape(*size, length) + + def reconstruct( + self, + x: Tensor, # [B, L, K] + sampler: DPMSolver, + num_samples: int, + generator: torch.Generator | None, + ) -> Tensor: # [B, L, K] + + x = x.transpose(1, 2) # [B, K, L] + batch, _, length = x.size() + mean = x.mean(dim=(1, 2), keepdim=True) + std = x.std(dim=(1, 2), keepdim=True) + 1e-5 + x = ((x - mean) / std).clamp(-10, 10) + + time = torch.arange(length, device=x.device).expand(batch, -1) + blocks = time // (length // self.num_blocks) + + out = x.new_zeros((num_samples, *x.size())) + for strategy in range(2): + mask = ((blocks % 2) == strategy).to(x.dtype) + mask = mask.unsqueeze(1).expand_as(x) # [B, K, L] + predict_noise = partial( + self.predict_noise, + cond=mask * x, + target_mask=1 - mask, + side_info=self.side_embedding(time, mask), + strategy=torch.tensor([strategy], device=x.device), + ) + for i in range(num_samples): + noise = torch.empty_like(x).normal_(generator=generator) + sample = sampler(noise, predict_noise) * std + mean + out[i] += (1 - mask) * sample + + return out.median(dim=0).values.transpose(1, 2) + + def predict_noise( + self, + x: Tensor, # [B, K, L] + step: Tensor, # [1] + *, + cond: Tensor, # [B, K, L] + target_mask: Tensor, # [B, K, L] + side_info: Tensor, # [B, S, K, L] + strategy: Tensor, # [1] + ) -> Tensor: # [B, K, L] + return self.denoiser( + x=torch.stack([cond, target_mask * x], dim=1), + side_info=side_info, + step=step, + strategy=strategy, + ) diff --git a/test/models/kumo/timeseries/anomaly/test_model.py b/test/models/kumo/timeseries/anomaly/test_model.py new file mode 100644 index 0000000000..04a46f0add --- /dev/null +++ b/test/models/kumo/timeseries/anomaly/test_model.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from sdm import CategoricalTensor, Recipe, Stype, TableTensor +from sdm.explain.gradient import CaptureInputs, EnableInputGradients +from sdm.models import Callback, KumoAnomaly +from sdm.testing import withCUDA + + +@withCUDA +def test_forward(device: torch.device) -> None: + model = KumoAnomaly(pretrained=False, device=device) + x = torch.randn(2, 130, 3, device=device) + + def generator() -> torch.Generator: + return torch.Generator(device).manual_seed(0) + + out = model(x, num_samples=1, num_steps=1, generator=generator()) + assert out.size() == (2, 130, 1) + assert out.columns["numerical"] == ("score",) + assert out.device == device + assert out.numerical.isfinite().all() + assert (out.numerical >= 0).all() + + again = model(x, num_samples=1, num_steps=1, generator=generator()) + assert torch.equal(again.numerical, out.numerical) + + +def test_wide_float64_forward() -> None: + model = KumoAnomaly(pretrained=False) + x = torch.randn(60, 50, dtype=torch.float64) + + out = model(x, num_samples=1, num_steps=1) + assert out.size() == (60, 1) + assert out.numerical.dtype == torch.float64 + + +class DoubleScores(Callback): + def on_model_forward_end( + self, + model: torch.nn.Module, + out: TableTensor, + ) -> TableTensor: + return out.replace_blocks(numerical=2 * out.numerical.double()) + + +def test_callbacks() -> None: + model = KumoAnomaly(pretrained=False) + x = torch.randn(30, 2) + capture = CaptureInputs() + callbacks = [EnableInputGradients(), capture, DoubleScores()] + + def generator() -> torch.Generator: + return torch.Generator().manual_seed(0) + + out = model(x, num_samples=1, num_steps=1, generator=generator()) + doubled = model( + x, + num_samples=1, + num_steps=1, + callbacks=callbacks, + generator=generator(), + ) + assert doubled.numerical.dtype == torch.float32 + torch.testing.assert_close(doubled.numerical, 2 * out.numerical) + + doubled.numerical.sum().backward() + ((x, _),) = capture.inputs + assert x.numerical.grad is not None + + +def test_invalid_inputs() -> None: + model = KumoAnomaly(pretrained=False) + x = TableTensor( + columns={Stype.numerical: ["a"], Stype.categorical: ["b"]}, + numerical=torch.randn(10, 1), + categorical=CategoricalTensor.from_tensor( + torch.randint(0, 5, (10, 1)) + ), + ) + + with pytest.warns(RuntimeWarning, match="unsupported feature stypes"): + model(x, num_samples=1, num_steps=1) + with pytest.raises(ValueError, match="at most 40 features"): + model(torch.randn(10, 41), recipe=Recipe()) From 7d3adf1169b00ea60dac3af946366fa6b9212286 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Mon, 19 Oct 2026 11:08:29 -0700 Subject: [PATCH 2/2] Support post-norm and dropout in `TransformerBlock` Add optional `attn_residual_norm` and `mlp_residual_norm` modules to `TransformerBlock`, and apply its dropout to the attention output before the residual addition. Build the Kumo-Anomaly `ResidualBlock` on a new `KumoAnomalyTransformerBlock` instead of `torch.nn.TransformerEncoderLayer`. --- sdm/models/kumo/timeseries/anomaly/block.py | 71 ++++++++++++++----- sdm/nn/attention.py | 30 +++++++- .../kumo/timeseries/anomaly/test_block.py | 51 ++++++++----- test/nn/test_attention.py | 58 +++++++++++++++ test/nn/test_compile.py | 26 +++++++ 5 files changed, 199 insertions(+), 37 deletions(-) diff --git a/sdm/models/kumo/timeseries/anomaly/block.py b/sdm/models/kumo/timeseries/anomaly/block.py index 975e50468e..4cf5f8d30c 100644 --- a/sdm/models/kumo/timeseries/anomaly/block.py +++ b/sdm/models/kumo/timeseries/anomaly/block.py @@ -6,7 +6,53 @@ import torch from torch import Tensor -from torch.nn import Conv1d, LayerNorm, Linear, TransformerEncoderLayer +from torch.nn import GELU, Conv1d, Dropout, LayerNorm, Linear, Sequential + +from sdm.nn import TransformerBlock + + +class KumoAnomalyTransformerBlock(TransformerBlock): + """Post-norm transformer block of Kumo-Anomaly. + + Args: + channels: Input and output width. + num_heads: Attention heads. + hidden_channels: Feedforward width. + dropout: Attention and feedforward dropout probability. + device: The device. + dtype: The dtype. + """ + + def __init__( + self, + channels: int, + num_heads: int, + hidden_channels: int, + dropout: float, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, + ) -> None: + factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype} + super().__init__( + channels=channels, + num_query_heads=num_heads, + mlp=Sequential( + Linear(channels, hidden_channels, **factory_kwargs), + GELU(), + Dropout(dropout), + Linear(hidden_channels, channels, **factory_kwargs), + Dropout(dropout), + ), + attn_residual_norm=LayerNorm(channels, **factory_kwargs), + mlp_residual_norm=LayerNorm(channels, **factory_kwargs), + dropout=dropout, + **factory_kwargs, + ) + # Match the initialization of `torch.nn.MultiheadAttention`. + torch.nn.init.xavier_uniform_(self.attn.qkv_lin.weight) + torch.nn.init.zeros_(self.attn.qkv_lin.bias) + self.attn.out_lin.reset_parameters() + torch.nn.init.zeros_(self.attn.out_lin.bias) class ResidualBlock(torch.nn.Module): @@ -84,25 +130,18 @@ def __init__( with torch.no_grad(): projection.weight.mul_(0.5) - # Preserve post-residual normalization and attention dropout, which - # SDM's TransformerBlock does not currently provide together. - # TODO: Switch to SDM's TransformerBlock once it supports both. - self.time_layer = TransformerEncoderLayer( - d_model=channels, - nhead=num_heads, - dim_feedforward=hidden_channels, + self.time_layer = KumoAnomalyTransformerBlock( + channels=channels, + num_heads=num_heads, + hidden_channels=hidden_channels, dropout=dropout, - activation="gelu", - batch_first=True, **factory_kwargs, ) - self.feature_layer = TransformerEncoderLayer( - d_model=channels, - nhead=num_heads, - dim_feedforward=hidden_channels, + self.feature_layer = KumoAnomalyTransformerBlock( + channels=channels, + num_heads=num_heads, + hidden_channels=hidden_channels, dropout=dropout, - activation="gelu", - batch_first=True, **factory_kwargs, ) self.norm_after_time = LayerNorm(channels, **factory_kwargs) diff --git a/sdm/nn/attention.py b/sdm/nn/attention.py index fa4721ff7e..c6446e2e15 100644 --- a/sdm/nn/attention.py +++ b/sdm/nn/attention.py @@ -458,6 +458,10 @@ class TransformerBlock(torch.nn.Module): attention. post_attn_norm: Normalization applied to the attention output before its residual addition. + attn_residual_norm: Normalization applied after the attention residual + addition. + mlp_residual_norm: Normalization applied after the feedforward + residual addition. Requires ``mlp``. query_transform: Transformation applied to projected query heads before scaled dot-product attention. key_transform: Transformation applied to projected key heads before @@ -468,7 +472,8 @@ class TransformerBlock(torch.nn.Module): scale: Scaling factor passed to :func:`torch.nn.functional.scaled_dot_product_attention`. ``None`` uses ``1 / sqrt(channels_per_head)``. - dropout: Attention dropout probability. + dropout: Dropout probability applied to attention weights and to the + attention output before its residual addition. bias: If set to ``False``, the module will not learn an additive bias. device: The device. dtype: The dtype. @@ -483,6 +488,8 @@ def __init__( query_norm: torch.nn.Module | None = None, key_value_norm: torch.nn.Module | None = None, post_attn_norm: torch.nn.Module | None = None, + attn_residual_norm: torch.nn.Module | None = None, + mlp_residual_norm: torch.nn.Module | None = None, query_transform: torch.nn.Module | None = None, key_transform: torch.nn.Module | None = None, query_scaling: QueryScaling | None = None, @@ -494,12 +501,18 @@ def __init__( dtype: torch.dtype | None = None, ) -> None: super().__init__() + if mlp is None and mlp_residual_norm is not None: + raise ValueError("'mlp_residual_norm' requires 'mlp'") + factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype} self.mlp = mlp self.query_norm = query_norm self.key_value_norm = key_value_norm self.post_attn_norm = post_attn_norm + self.attn_residual_norm = attn_residual_norm + self.mlp_residual_norm = mlp_residual_norm + self.dropout = dropout self.attn = Attention( channels=channels, @@ -811,6 +824,8 @@ def _forward( if self.post_attn_norm is not None: attn_out = self.post_attn_norm(attn_out) + if self.training and self.dropout > 0.0: + attn_out = F.dropout(attn_out, p=self.dropout) if ( out is not None @@ -818,13 +833,24 @@ def _forward( and not out.is_contiguous() ): tmp = (attn_out + query).to(out.dtype) + if self.attn_residual_norm is not None: + tmp = self.attn_residual_norm(tmp) if self.mlp is not None: - tmp += self.mlp(tmp) + tmp = (tmp + self.mlp(tmp)).to(out.dtype) + if self.mlp_residual_norm is not None: + tmp = self.mlp_residual_norm(tmp) out.copy_(tmp) else: tmp = torch.add(attn_out, query, out=out) + if self.attn_residual_norm is not None: + tmp = self.attn_residual_norm(tmp) if self.mlp is not None: tmp = torch.add(tmp, self.mlp(tmp), out=out) + if self.mlp_residual_norm is not None: + tmp = self.mlp_residual_norm(tmp) + # Normalization allocates, so copy its result into `out`. + if out is not None and tmp is not out: + tmp = out.copy_(tmp) out = tmp return (out, kv) if return_key_value else out diff --git a/test/models/kumo/timeseries/anomaly/test_block.py b/test/models/kumo/timeseries/anomaly/test_block.py index 92347cd882..c2269441be 100644 --- a/test/models/kumo/timeseries/anomaly/test_block.py +++ b/test/models/kumo/timeseries/anomaly/test_block.py @@ -2,35 +2,48 @@ # SPDX-License-Identifier: Apache-2.0 import math -from collections.abc import Iterator import pytest import torch -import torch.nn.functional as F from torch import Tensor -from sdm.models.kumo.timeseries.anomaly.block import ResidualBlock +from sdm.models.kumo.timeseries.anomaly.block import ( + KumoAnomalyTransformerBlock, + ResidualBlock, +) from sdm.testing import withCUDA -@pytest.fixture(autouse=True) -def exact_gelu() -> Iterator[None]: - # At inference time on GPU, PyTorch's encoder layer switches to a faster - # kernel with an approximate GELU. Disable it to compare exact outputs. - enabled = torch.backends.mha.get_fastpath_enabled() - torch.backends.mha.set_fastpath_enabled(False) - yield - torch.backends.mha.set_fastpath_enabled(enabled) - - def reference_attention( - layer: torch.nn.TransformerEncoderLayer, x: Tensor + layer: KumoAnomalyTransformerBlock, x: Tensor ) -> Tensor: - """Evaluate post-norm attention without the encoder's fused fast path.""" - attn, _ = layer.self_attn(x, x, x, need_weights=False) - x = layer.norm1(x + layer.dropout1(attn)) - hidden = layer.linear2(layer.dropout(F.gelu(layer.linear1(x)))) - return layer.norm2(x + layer.dropout2(hidden)) + """Evaluate a sequence-first PyTorch encoder layer with equal weights.""" + state = layer.state_dict() + reference = torch.nn.TransformerEncoderLayer( + d_model=x.size(-1), + nhead=layer.attn.num_query_heads, + dim_feedforward=state["mlp.0.weight"].size(0), + dropout=0.0, + activation="gelu", + device=x.device, + dtype=x.dtype, + ) + names = { + "self_attn.in_proj_": "attn.qkv_lin.", + "self_attn.out_proj.": "attn.out_lin.", + "linear1.": "mlp.0.", + "linear2.": "mlp.3.", + "norm1.": "attn_residual_norm.", + "norm2.": "mlp_residual_norm.", + } + reference.load_state_dict( + { + target + key: state[source + key] + for target, source in names.items() + for key in ("weight", "bias") + } + ) + return reference.eval()(x.transpose(0, 1)).transpose(0, 1) def reference_block( diff --git a/test/nn/test_attention.py b/test/nn/test_attention.py index 4e45506efd..3e7e14b4d7 100644 --- a/test/nn/test_attention.py +++ b/test/nn/test_attention.py @@ -751,3 +751,61 @@ def test_transformer_block_kv_cache() -> None: torch.testing.assert_close(cache_out, direct_out) torch.testing.assert_close(cached_out, direct_out) + + +@withCUDA +def test_transformer_block_residual_norms(device: torch.device) -> None: + channels = 8 + mlp = torch.nn.Linear(channels, channels, device=device) + attn_residual_norm = torch.nn.LayerNorm(channels, device=device) + mlp_residual_norm = torch.nn.LayerNorm(channels, device=device) + module = TransformerBlock( + channels=channels, + num_query_heads=2, + mlp=mlp, + attn_residual_norm=attn_residual_norm, + mlp_residual_norm=mlp_residual_norm, + device=device, + ) + torch.nn.init.normal_(module.attn.out_lin.weight) + query = torch.randn(2, 4, 3, channels, device=device).transpose(1, 2) + + expected = attn_residual_norm(query + module.attn(query)) + expected = mlp_residual_norm(expected + mlp(expected)) + torch.testing.assert_close(module(query), expected) + + with torch.no_grad(): + out = torch.empty(2, 3, 4, channels, device=device) + assert module(query, batch_size_limit=1, out=out) is out + torch.testing.assert_close(out, expected) + + alias = query.clone() + assert module(alias, out=alias) is alias + torch.testing.assert_close(alias, expected) + + with pytest.raises(ValueError, match="requires 'mlp'"): + TransformerBlock( + channels=channels, + num_query_heads=2, + mlp=None, + mlp_residual_norm=mlp_residual_norm, + ) + + +@withCUDA +def test_transformer_block_dropout(device: torch.device) -> None: + module = TransformerBlock( + channels=4, + num_query_heads=1, + mlp=None, + dropout=0.5, + device=device, + ) + torch.nn.init.eye_(module.attn.out_lin.weight) + # A single key receives the full attention weight. Dropping weights and + # outputs each scales the unit values by 0 or 2. + query = torch.zeros(64, 1, 4, device=device) + with torch.no_grad(): + module.attn.qkv_lin.bias[8:].fill_(1.0) + assert set(module(query).unique().tolist()) == {0.0, 4.0} + torch.testing.assert_close(module.eval()(query), query + 1) diff --git a/test/nn/test_compile.py b/test/nn/test_compile.py index d61f15a436..8a0f292aeb 100644 --- a/test/nn/test_compile.py +++ b/test/nn/test_compile.py @@ -280,6 +280,32 @@ def test_transformer_block_compile( torch.testing.assert_close(out, expected) +@withCUDA +def test_transformer_block_residual_norms_compile( + device: torch.device, +) -> None: + channels = 8 + module = TransformerBlock( + channels=channels, + num_query_heads=2, + mlp=torch.nn.Linear(channels, channels, device=device), + attn_residual_norm=torch.nn.LayerNorm(channels, device=device), + mlp_residual_norm=torch.nn.LayerNorm(channels, device=device), + device=device, + ) + query = torch.randn(2, 3, channels, device=device) + expected = torch.empty(2, 3, channels, device=device, dtype=torch.bfloat16) + out = expected.new_empty(3, 2, channels).transpose(0, 1) + + # CUDA autocast runs `LayerNorm` in float32, so residual sums must be + # rounded to the dtype of `out` as in eager mode. + with torch.no_grad(), torch.autocast(device.type, dtype=torch.bfloat16): + module(query=query, out=expected) + actual = fullgraph(module)(query=query, out=out) + assert actual is out + assert torch.equal(out, expected) + + @withCUDA @pytest.mark.parametrize("self_attn", [False, True]) def test_induced_transformer_block_compile(