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`.

Author: aw471

sdm/models/kumo/timeseries/anomaly/block.py

@@ -6,7 +6,64 @@
 
 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)
+
+    def peak_bytes_per_example(
+        self,
+        element_size: int,
+        query_length: int,
+        key_value_length: int | None = None,
+    ) -> int:
+        r""":meta private:"""  # noqa: D415
+        length = max(query_length, key_value_length or 0)
+        factor = 10 if element_size <= 2 else 8
+        return factor * length * element_size * self.attn.q_dim
 
 
 class ResidualBlock(torch.nn.Module):
@@ -84,25 +141,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)
@@ -159,3 +209,17 @@ def forward(
             .chunk(2, dim=1)
         )
         return (x + residual) / math.sqrt(2.0), skip
+
+
+if __name__ == "__main__":
+    from sdm.testing.memory import benchmark_transformer_block_memory_peak
+
+    benchmark_transformer_block_memory_peak(
+        block=lambda channels, num_heads: KumoAnomalyTransformerBlock(
+            channels=channels,
+            num_heads=num_heads,
+            hidden_channels=64,
+            dropout=0.0,
+        ),
+        channels_and_heads=[(128, 8)],
+    )

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,20 +824,33 @@ 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
             and torch.compiler.is_compiling()
             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

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(

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)

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(