stackformers 4.2.1__tar.gz → 4.3.1__tar.gz

{stackformers-4.2.1 → stackformers-4.3.1}/CHANGELOG.md

@@ -6,6 +6,58 @@ The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and
 adheres to [Semantic Versioning](https://semver.org/): MAJOR for breaking public API changes,
 MINOR for backwards-compatible features, PATCH for bug fixes and internal changes.
 
+## [4.3.1] — 2026-07-17
+
+### Fixed
+
+- **`MLMWrapper`'s masked pass leaked gradient to whatever produced `input.x`, despite
+  the target being detached.** `corrupted_x` was built from the live (undetached)
+  `input.x` at unmasked positions; self-attention mixes those into the masked
+  positions' predictions, so `mlm_loss` still reached back through every unmasked
+  position — confirmed empirically at 432/512 nonzero gradient entries under a
+  realistic 15% mask ratio. The existing test only covered 100% masking, which zeroes
+  that path as a side effect of `torch.where`'s backward and never exercised the
+  realistic partial-masking case.
+
+  Fixed by detaching `input.x` once, up front
+  (`input = input._replace(x=input.x.detach())`), so both `corrupted_x` and `target`
+  are built from the same severed copy. `mlm_loss` now trains only `mask_token`,
+  `encoder`, and `head` — never whatever produced `input.x` — regardless of mask
+  ratio. Added a regression test at the realistic mask ratio
+  (`test_wrapper_severs_gradient_to_upstream_input_under_partial_masking`) alongside
+  the existing full-masking one.
+
+## [4.3.0] — 2026-07-17
+
+### Added
+
+- **`MLMWrapper`** (`stackformers/mlm/`) — a masked-token-reconstruction auxiliary loss over any
+  encoder satisfying `EncoderLike`, domain-agnostic about what a token represents. Takes the
+  encoder at call time (`forward(input, encoder)`) rather than owning it, so nothing is
+  registered as a submodule and `mlm_wrapper.parameters()` never includes the encoder's weights.
+- `MLMWrapperConfig`, `MaskingStrategy`/`RandomMasking`, `ReconstructionHead`/`RegressionHead`,
+  and `MLMOutput` (`out`, `mlm_loss`).
+
+### Notes
+
+- **`out` is always the encoder's clean, unmasked output; only `mlm_loss` ever reflects
+  masking.** In training, `forward` runs the encoder twice — once clean (returned as `out`) and
+  once on a separately-masked copy (used only to compute `mlm_loss`) — so the main pipeline sees
+  byte-identical output whether or not this loss is being trained alongside it. In eval, only the
+  clean pass runs and `mlm_loss` is a constant zero, gated on `self.training` the same way
+  `nn.Dropout` and `nn.BatchNorm` already are — so callers can invoke it unconditionally in both
+  modes with no `if training` branch of their own.
+
+- **The reconstruction target is always `input.x.detach()`.** This severs the gradient path from
+  the loss back to whatever produced `input.x`, removing the collapse shortcut a trainable
+  tokenizer would otherwise have (drive every token toward one constant vector to make
+  reconstruction trivial).
+
+- **`RandomMasking` needs no packed-sequence boundary awareness** — each token's masking decision
+  is independent of every other's, so document identity (`cu_seqlens`) never enters it. That only
+  becomes necessary for a contiguous-span (blockwise) strategy, which could otherwise straddle
+  two packed documents.
+
 ## [4.2.1] — 2026-07-14
 
 ### Removed

{stackformers-4.2.1 → stackformers-4.3.1}/CLAUDE.md

@@ -25,6 +25,7 @@ Behavior comes from injected dependencies, not constructor flags.
 | dh     | dim per head |
 | w      | window size |
 | nt     | total tokens in packed sequence |
+| m      | number of masked tokens (gathered across batch/pack) |
 
 ## File organisation
 

{stackformers-4.2.1 → stackformers-4.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stackformers
-Version: 4.2.1
+Version: 4.3.1
 Summary: Typed, composable, SOLID transformer library for PyTorch
 Project-URL: Homepage, https://github.com/Red-Eyed/stackformers
 Project-URL: Repository, https://github.com/Red-Eyed/stackformers

{stackformers-4.2.1 → stackformers-4.3.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "stackformers"
-version = "4.2.1"
+version = "4.3.1"
 description = "Typed, composable, SOLID transformer library for PyTorch"
 readme = "README.md"
 license = { file = "LICENSE" }

{stackformers-4.2.1 → stackformers-4.3.1}/stackformers/__init__.py

@@ -17,6 +17,11 @@ from stackformers.feedforward.factory import build_ff
 from stackformers.feedforward.protocols import FeedForward
 from stackformers.feedforward.swiglu import SwiGLU
 from stackformers.layers import TransformerLayer
+from stackformers.mlm.config import MLMWrapperConfig
+from stackformers.mlm.head import RegressionHead
+from stackformers.mlm.masking import RandomMasking
+from stackformers.mlm.protocols import EncoderLike, MaskingStrategy, ReconstructionHead
+from stackformers.mlm.wrapper import MLMOutput, MLMWrapper
 from stackformers.norm.config import LayerNormConfig, RMSNormConfig
 from stackformers.norm.factory import NormConfig, build_norm
 from stackformers.norm.protocols import Norm
@@ -85,6 +90,9 @@ __all__ = [
     "CrossAttn",
     "FeedForward",
     "Norm",
+    "EncoderLike",
+    "MaskingStrategy",
+    "ReconstructionHead",
     # configs — attention
     "SelfAttentionConfig",
     "CrossAttentionConfig",
@@ -133,4 +141,10 @@ __all__ = [
     "CrossAttenderConfig",
     "CrossAttender",
     "plain_cross_attender_config",
+    # mlm
+    "MLMWrapperConfig",
+    "RandomMasking",
+    "RegressionHead",
+    "MLMWrapper",
+    "MLMOutput",
 ]

stackformers-4.3.1/stackformers/mlm/README.md

@@ -0,0 +1,54 @@
+# mlm
+
+A masked-token-reconstruction auxiliary loss, wrapped around any encoder that satisfies `EncoderLike`. Domain-agnostic — it has no idea what a token represents.
+
+## Usage
+
+```python
+import torch
+from stackformers import (
+    MLMWrapper,
+    MLMWrapperConfig,
+    TransformerEncoder,
+    make_padded_input,
+    plain_encoder_config,
+)
+
+encoder = TransformerEncoder(plain_encoder_config(dim=512, heads=8, num_layers=6))
+mlm_wrapper = MLMWrapper(MLMWrapperConfig(dim=512, mask_ratio=0.15))
+
+x = torch.randn(2, 128, 512)
+mask = torch.ones(2, 128, dtype=torch.bool)
+input = make_padded_input(x, mask)
+
+res = mlm_wrapper(input, encoder)
+
+task_loss = my_task_head(res.out)  # res.out is always the clean, unmasked encoder output
+loss = task_loss + 0.1 * res.mlm_loss  # zero automatically in eval — nothing to branch on
+```
+
+One call site covers both modes: in training, `mlm_wrapper` runs the encoder twice internally — once clean (returned as `res.out`) and once on a separately-masked copy (used only to compute `res.mlm_loss`) — so masking never reaches whatever consumes `res.out`. In eval, it runs the encoder once and `res.mlm_loss` is a constant zero. Pretraining with no main task yet — just drop `task_loss` and train on `res.mlm_loss` alone.
+
+`mlm_wrapper` and `encoder` are separate modules — `encoder` is passed in at call time, never stored, so `mlm_wrapper.parameters()` holds only `mask_token`, `masking_strategy`, and `head`. Build one optimizer over both: `itertools.chain(encoder.parameters(), mlm_wrapper.parameters())`. And since they're separate modules, `.train()`/`.eval()` on one does not propagate to the other — call both, or register both as submodules of a shared parent so one call does.
+
+## Design
+
+**`out` is always clean; only `mlm_loss` ever reflects masking.** `forward(input, encoder)` runs `encoder(input)` unconditionally for `out`. In training it additionally runs `encoder` a second time on a masked copy, purely to produce `mlm_loss` — that corrupted copy is never returned. This means the main pipeline sees byte-identical output whether or not the aux loss is being trained alongside it, and callers never need to special-case which pass produced `out`.
+
+**Training/eval is the one branch `forward()` is allowed, because it isn't tensor control flow.** `self.training` gates whether the second (masked) pass runs — the same mechanism `nn.Dropout` and `nn.BatchNorm` already use. It lets callers invoke `mlm_wrapper(input, encoder)` unconditionally in both modes, instead of an external `if training: ...` at every call site.
+
+**A wrapper over an encoder supplied at call time, not one it owns.** `MLMWrapper` never stores an encoder — `forward` takes one as an argument, so nothing is registered as a submodule and `mlm_wrapper.parameters()` never includes the encoder's weights.
+
+**The mask token is owned state, not an injected collaborator.** It's the wrapper's own learned parameter, analogous to a `nn.Linear`'s weight — not swappable behavior, so it isn't a `Protocol`.
+
+**`input.x` is detached once, before it is used for anything — not only as the target.** The masked pass replaces `input` with `input._replace(x=input.x.detach())` up front, so both the unmasked context fed to the encoder and the reconstruction target come from that same detached copy. Detaching only the target is not sufficient: self-attention still mixes the live, undetached unmasked positions into the masked positions' predictions, so a gradient path back to whatever produced `input.x` would otherwise survive at every unmasked position — worse, it grows with the fraction of tokens left unmasked. Severing the whole input removes the standard collapse shortcut for self-supervised regression targets (drive every token toward one constant vector, and reconstruction becomes trivial) — not just the degenerate all-masked case, but the realistic partial-masking case a target-only detach would silently miss. See `tests/mlm/test_wrapper.py::test_wrapper_severs_gradient_to_upstream_input_under_partial_masking` for the regression test.
+
+**No layout dispatch in the wrapper.** `MaskingStrategy` returns a boolean tensor shaped like `input.x`'s leading dims — `(b, n)` for `PaddedInput`, `(nt,)` for `PackedInput`. `torch.where` and boolean advanced indexing both broadcast/gather correctly against either shape, so `MLMWrapper.forward` never matches on the sequence variant; only `RandomMasking` does, internally.
+
+**Random masking needs no packed-sequence boundary awareness.** Each token's masking decision is independent of every other token's, so document identity (`cu_seqlens`) never enters the decision. Boundary-awareness only becomes necessary for a contiguous-span (blockwise) strategy, which could otherwise straddle two packed documents — add that logic when that strategy is added, not before.
+
+## Extending
+
+To add a masking strategy (e.g. blockwise), write an `nn.Module` satisfying `MaskingStrategy` and pass it as `MLMWrapper(..., masking_strategy=...)`.
+
+To add a reconstruction target type (e.g. a discretized/codebook target for BEiT-style cross-entropy, or a contrastive target), write an `nn.Module` satisfying `ReconstructionHead` and pass it as `MLMWrapper(..., head=...)`. Prediction and loss are scored together by one call — they're a coupled choice (a cross-entropy loss over a raw regression output isn't meaningful), not two independently pluggable axes.

stackformers-4.3.1/stackformers/mlm/config.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+
+
+class MLMWrapperConfig(BaseModel):
+    """Config for MLMWrapper: token dimension and corruption ratio."""
+
+    dim: int = Field(gt=0)
+    mask_ratio: float = Field(default=0.15, gt=0.0, lt=1.0)

stackformers-4.3.1/stackformers/mlm/head.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import torch.nn as nn
+import torch.nn.functional as F
+from jaxtyping import Float
+from torch import Tensor
+
+
+class RegressionHead(nn.Module):
+    """Reconstruct masked tokens via a linear projection, scored with MSE.
+
+    The projection gives the encoder a dedicated place to specialise for
+    reconstruction, so the shared representation isn't forced to double as the
+    literal token vector at every layer.
+    """
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(
+        self,
+        encoder_output_at_masked: Float[Tensor, "m d"],
+        target_at_masked: Float[Tensor, "m d"],
+    ) -> Float[Tensor, ""]:
+        prediction = self.proj(encoder_output_at_masked)
+        return F.mse_loss(prediction, target_at_masked)

stackformers-4.3.1/stackformers/mlm/masking.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from jaxtyping import Bool
+from torch import Tensor
+
+from stackformers.sequence import PackedInput, PaddedInput, SequenceInput
+
+
+class RandomMasking(nn.Module):
+    """Mask each valid token independently with probability mask_ratio.
+
+    Every token's masking decision is drawn independently, so packed sequences need
+    no document-boundary awareness here — that only becomes necessary for a
+    contiguous-span (blockwise) strategy, which could otherwise straddle two packed
+    documents.
+    """
+
+    def __init__(self, mask_ratio: float = 0.15) -> None:
+        super().__init__()
+        self.mask_ratio = mask_ratio
+
+    def forward(self, input: SequenceInput) -> Bool[Tensor, "*batch"]:
+        match input:
+            case PaddedInput(mask=mask):
+                candidate = torch.rand(mask.shape, device=mask.device) < self.mask_ratio
+                return candidate & mask
+            case PackedInput(x=x):
+                return torch.rand(x.shape[0], device=x.device) < self.mask_ratio

stackformers-4.3.1/stackformers/mlm/protocols.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+from jaxtyping import Bool, Float
+from torch import Tensor
+
+from stackformers.sequence import SequenceInput
+
+
+@runtime_checkable
+class EncoderLike(Protocol):
+    """Any encoder mapping a sequence to per-token embeddings of the same layout.
+
+    Implementations: Encoder, or any TransformerEncoderBase subclass.
+    """
+
+    def __call__(self, input: SequenceInput) -> Tensor: ...
+
+
+@runtime_checkable
+class MaskingStrategy(Protocol):
+    """Select which token positions to corrupt for reconstruction.
+
+    Returns True at positions to mask, shaped like the leading (non-feature) dims of
+    input.x — (b, n) for PaddedInput, (nt,) for PackedInput. Implementations dispatch
+    on the SequenceInput variant internally, so callers never need to know the layout.
+    Implementation: RandomMasking.
+    """
+
+    def __call__(self, input: SequenceInput) -> Bool[Tensor, "*batch"]: ...
+
+
+@runtime_checkable
+class ReconstructionHead(Protocol):
+    """Predict masked tokens from encoder output and score against the clean target.
+
+    Both arguments are already gathered down to just the masked positions (m = number
+    of masked tokens across the batch or pack). Returns a scalar loss.
+    Implementation: RegressionHead.
+    """
+
+    def __call__(
+        self,
+        encoder_output_at_masked: Float[Tensor, "m d"],
+        target_at_masked: Float[Tensor, "m d"],
+    ) -> Float[Tensor, ""]: ...

stackformers-4.3.1/stackformers/mlm/wrapper.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import torch
+import torch.nn as nn
+from jaxtyping import Float
+from torch import Tensor
+
+from stackformers.mlm.config import MLMWrapperConfig
+from stackformers.mlm.head import RegressionHead
+from stackformers.mlm.masking import RandomMasking
+from stackformers.mlm.protocols import EncoderLike, MaskingStrategy, ReconstructionHead
+from stackformers.sequence import SequenceInput
+
+
+class MLMOutput(NamedTuple):
+    out: Float[Tensor, "*batch d"]
+    mlm_loss: Float[Tensor, ""]
+
+
+class MLMWrapper(nn.Module):
+    """Masked-token-reconstruction auxiliary loss over an encoder supplied at call time.
+
+    Does not store or own an encoder — none is passed to __init__, and forward() takes
+    one as an argument instead of holding a reference across calls. This module's own
+    parameter tree (mask_token, masking_strategy, head) is exactly what it owns; the
+    encoder stays wherever the caller already keeps it, with no duplicate registration
+    under this wrapper.
+
+    `out` is always the encoder's clean, unmasked output — masking is invisible to
+    whatever consumes it, so the main pipeline behaves identically whether or not this
+    aux loss is being trained alongside it. Only `mlm_loss` ever reflects that masking
+    happened. In training mode that costs a second encoder forward pass on a
+    separately-masked copy of input; the two passes share weights but never interfere,
+    since the corrupted copy built for the loss is never returned as `out`. In eval mode
+    only the one clean pass runs and `mlm_loss` reports a constant zero, so callers can
+    invoke this unconditionally in both modes without an if-training branch of their
+    own — the same role self.training already plays in nn.Dropout or nn.BatchNorm.
+
+    The masked pass detaches input.x once, up front — both the unmasked context fed to
+    the encoder and the reconstruction target come from that same detached copy — so
+    mlm_loss trains the encoder, mask_token, and head, never whatever produced input.x.
+    Detaching only the target is not sufficient: self-attention still mixes the
+    unmasked (undetached) positions into the masked positions' predictions, so a live
+    path back to input.x would otherwise survive at every unmasked position. Severing
+    it entirely removes the representation-collapse shortcut a trainable tokenizer
+    would otherwise have available (drive every token toward a constant vector to make
+    reconstruction trivial).
+    """
+
+    def __init__(
+        self,
+        config: MLMWrapperConfig,
+        masking_strategy: MaskingStrategy | None = None,
+        head: ReconstructionHead | None = None,
+    ) -> None:
+        super().__init__()
+        self.mask_token = nn.Parameter(torch.empty(config.dim))
+        nn.init.trunc_normal_(self.mask_token, std=0.02, a=-0.04, b=0.04)
+        self.masking_strategy = (
+            masking_strategy if masking_strategy is not None else RandomMasking(config.mask_ratio)
+        )
+        self.head = head if head is not None else RegressionHead(config.dim)
+
+    def _masked_loss(self, input: SequenceInput, encoder: EncoderLike) -> Tensor:
+        # Detached once, up front, and reused for both corrupted_x and target below —
+        # unmasked positions must not carry a live path back to whatever produced
+        # input.x either, or self-attention mixes them into the masked positions'
+        # predictions and mlm_loss reaches upstream anyway despite target being detached.
+        input = input._replace(x=input.x.detach())
+        target = input.x
+        should_mask = self.masking_strategy(input)
+        # should_mask's shape always matches input.x's leading dims, so this select and
+        # the boolean indexing below both work for PaddedInput and PackedInput alike.
+        corrupted_x = torch.where(should_mask.unsqueeze(-1), self.mask_token, input.x)
+        masked_output = encoder(input._replace(x=corrupted_x))
+        return self.head(masked_output[should_mask], target[should_mask])
+
+    def forward(self, input: SequenceInput, encoder: EncoderLike) -> MLMOutput:
+        clean_output = encoder(input)
+        if self.training:
+            mlm_loss = self._masked_loss(input, encoder)
+            return MLMOutput(out=clean_output, mlm_loss=mlm_loss)
+        return MLMOutput(out=clean_output, mlm_loss=clean_output.new_zeros(()))

stackformers-4.3.1/tests/mlm/test_head.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from stackformers.mlm.head import RegressionHead
+
+M, D = 12, 16
+
+
+def test_regression_head_output_is_scalar(device_dtype: tuple[torch.device, torch.dtype]) -> None:
+    device, dtype = device_dtype
+    head = RegressionHead(dim=D).to(device=device, dtype=dtype)
+    prediction_in = torch.randn(M, D, device=device, dtype=dtype)
+    target = torch.randn(M, D, device=device, dtype=dtype)
+    loss = head(prediction_in, target)
+    assert loss.shape == ()
+
+
+def test_regression_head_zero_when_projection_matches_target(device: torch.device) -> None:
+    head = RegressionHead(dim=D).to(device)
+    nn.init.zeros_(head.proj.weight)
+    nn.init.zeros_(head.proj.bias)
+    target = torch.zeros(M, D, device=device)
+    prediction_in = torch.randn(M, D, device=device)  # proj(x) == 0 regardless, weight/bias are 0
+    loss = head(prediction_in, target)
+    assert torch.allclose(loss, torch.zeros_like(loss))
+
+
+def test_regression_head_gradients_flow(device: torch.device) -> None:
+    head = RegressionHead(dim=D).to(device)
+    prediction_in = torch.randn(M, D, device=device, requires_grad=True)
+    target = torch.randn(M, D, device=device)
+    loss = head(prediction_in, target)
+    loss.backward()
+    assert prediction_in.grad is not None
+    assert head.proj.weight.grad is not None

stackformers-4.3.1/tests/mlm/test_masking.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import torch
+
+from stackformers.mlm.masking import RandomMasking
+from stackformers.sequence import make_packed_input, make_padded_input
+
+B, N, D = 4, 64, 16
+NT = 20
+
+
+def test_random_masking_padded_shape(device_dtype: tuple[torch.device, torch.dtype]) -> None:
+    device, dtype = device_dtype
+    x = torch.randn(B, N, D, device=device, dtype=dtype)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    input = make_padded_input(x, mask)
+    should_mask = RandomMasking(mask_ratio=0.5)(input)
+    assert should_mask.shape == mask.shape
+    assert should_mask.dtype == torch.bool
+
+
+def test_random_masking_packed_shape(device_dtype: tuple[torch.device, torch.dtype]) -> None:
+    device, dtype = device_dtype
+    x = torch.randn(NT, D, device=device, dtype=dtype)
+    cu = torch.tensor([0, 12, 20], dtype=torch.int32, device=device)
+    input = make_packed_input(x, cu, max_seqlen=12)
+    should_mask = RandomMasking(mask_ratio=0.5)(input)
+    assert should_mask.shape == (NT,)
+    assert should_mask.dtype == torch.bool
+
+
+def test_random_masking_never_masks_padding(device_dtype: tuple[torch.device, torch.dtype]) -> None:
+    device, dtype = device_dtype
+    x = torch.randn(B, N, D, device=device, dtype=dtype)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    mask[:, N // 2 :] = False  # second half is padding
+    input = make_padded_input(x, mask)
+    should_mask = RandomMasking(mask_ratio=0.9)(input)
+    assert not (should_mask & ~mask).any()
+
+
+def test_random_masking_respects_ratio_roughly(
+    device_dtype: tuple[torch.device, torch.dtype],
+) -> None:
+    """Statistical check over a large sample — not exact, but far enough from any other ratio."""
+    device, dtype = device_dtype
+    torch.manual_seed(0)
+    x = torch.randn(1, 10_000, D, device=device, dtype=dtype)
+    mask = torch.ones(1, 10_000, dtype=torch.bool, device=device)
+    input = make_padded_input(x, mask)
+    should_mask = RandomMasking(mask_ratio=0.3)(input)
+    fraction = should_mask.float().mean().item()
+    assert 0.25 < fraction < 0.35

stackformers-4.3.1/tests/mlm/test_wrapper.py

@@ -0,0 +1,222 @@
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+
+from stackformers.attention.config import SelfAttentionConfig
+from stackformers.attention.self_attn import SelfAttention
+from stackformers.encoder import Encoder
+from stackformers.feedforward.config import SwiGLUConfig
+from stackformers.feedforward.swiglu import SwiGLU
+from stackformers.layers import TransformerLayer
+from stackformers.mlm.config import MLMWrapperConfig
+from stackformers.mlm.wrapper import MLMOutput, MLMWrapper
+from stackformers.norm.config import RMSNormConfig
+from stackformers.norm.factory import build_norm
+from stackformers.positional.none import NoPosEncoding
+from stackformers.sequence import (
+    PackedInput,
+    PaddedInput,
+    SequenceInput,
+    make_packed_input,
+    make_padded_input,
+)
+
+B, N, D, H, DH = 2, 16, 16, 4, 4
+NUM_LAYERS = 2
+NT = 10  # two packed seqs: 6 + 4
+
+
+class AllMasking(nn.Module):
+    """Test double: marks every valid position for masking."""
+
+    def forward(self, input: SequenceInput) -> torch.Tensor:
+        match input:
+            case PaddedInput(mask=mask):
+                return mask.clone()
+            case PackedInput(x=x):
+                return torch.ones(x.shape[0], dtype=torch.bool, device=x.device)
+
+
+def _build_encoder(device: torch.device, dtype: torch.dtype) -> Encoder:
+    attn_cfg = SelfAttentionConfig(dim=D, heads=H, dim_head=DH)
+    ff_cfg = SwiGLUConfig(dim=D)
+    norm_cfg = RMSNormConfig(dim=D)
+    layers = [
+        TransformerLayer(
+            self_attn=SelfAttention(attn_cfg, NoPosEncoding()),
+            ff=SwiGLU(ff_cfg),
+            norm_attn=build_norm(norm_cfg),
+            norm_ff=build_norm(norm_cfg),
+        )
+        for _ in range(NUM_LAYERS)
+    ]
+    return Encoder(layers=layers, final_norm=build_norm(norm_cfg)).to(device=device, dtype=dtype)
+
+
+@pytest.fixture
+def config() -> MLMWrapperConfig:
+    return MLMWrapperConfig(dim=D, mask_ratio=0.5)
+
+
+@pytest.fixture
+def x_pad(device_dtype: tuple[torch.device, torch.dtype]) -> PaddedInput:
+    device, dtype = device_dtype
+    x = torch.randn(B, N, D, device=device, dtype=dtype)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    return make_padded_input(x, mask)
+
+
+@pytest.fixture
+def x_packed(device_dtype: tuple[torch.device, torch.dtype]) -> PackedInput:
+    device, dtype = device_dtype
+    if not device.type == "cuda" or dtype not in (torch.float16, torch.bfloat16):
+        pytest.skip("packed attention requires CUDA with float16 or bfloat16")
+    x = torch.randn(NT, D, device=device, dtype=dtype)
+    cu = torch.tensor([0, 6, 10], dtype=torch.int32, device=device)
+    return make_packed_input(x, cu, max_seqlen=6)
+
+
+def test_wrapper_padded_output_shapes(
+    config: MLMWrapperConfig,
+    x_pad: PaddedInput,
+    device_dtype: tuple[torch.device, torch.dtype],
+) -> None:
+    device, dtype = device_dtype
+    encoder = _build_encoder(device, dtype)
+    wrapper = MLMWrapper(config).to(device=device, dtype=dtype)
+    res = wrapper(x_pad, encoder)
+    assert isinstance(res, MLMOutput)
+    assert res.mlm_loss.shape == ()
+    assert res.out.shape == (B, N, D)
+
+
+def test_wrapper_packed_output_shapes(
+    config: MLMWrapperConfig,
+    x_packed: PackedInput,
+    device_dtype: tuple[torch.device, torch.dtype],
+) -> None:
+    device, dtype = device_dtype
+    encoder = _build_encoder(device, dtype)
+    wrapper = MLMWrapper(config).to(device=device, dtype=dtype)
+    res = wrapper(x_packed, encoder)
+    assert res.mlm_loss.shape == ()
+    assert res.out.shape == (NT, D)
+
+
+def test_wrapper_mask_token_receives_gradient(device: torch.device) -> None:
+    config = MLMWrapperConfig(dim=D, mask_ratio=0.5)
+    encoder = _build_encoder(device, torch.float32)
+    wrapper = MLMWrapper(config).to(device)
+    x = torch.randn(B, N, D, device=device)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    res = wrapper(make_padded_input(x, mask), encoder)
+    res.mlm_loss.backward()
+    assert wrapper.mask_token.grad is not None
+
+
+def test_wrapper_does_not_own_encoder(device: torch.device) -> None:
+    """MLMWrapper takes the encoder at call time, not at construction — its own
+    parameter tree must never include the encoder's weights, since the encoder is
+    owned and registered wherever the caller already keeps it.
+    """
+    config = MLMWrapperConfig(dim=D, mask_ratio=0.5)
+    encoder = _build_encoder(device, torch.float32)
+    wrapper = MLMWrapper(config).to(device)
+    encoder_param_ids = {id(p) for p in encoder.parameters()}
+    wrapper_param_ids = {id(p) for p in wrapper.parameters()}
+    assert encoder_param_ids.isdisjoint(wrapper_param_ids)
+
+
+def test_wrapper_severs_gradient_to_upstream_input_under_full_masking(
+    device: torch.device,
+) -> None:
+    """Collapse guard (design doc §5): a trainable tokenizer upstream of x must get no
+    gradient from this loss, or the optimizer can collapse every token to one constant
+    vector to make reconstruction trivial.
+
+    With every position masked, x is detached before corrupted_x is built from it, so
+    corrupted_x is pure mask_token — independent of x entirely. x never appears in the
+    masked pass's graph at all, so x.grad must be exactly None after backward, not just
+    zero-valued.
+    """
+    config = MLMWrapperConfig(dim=D, mask_ratio=0.5)
+    encoder = _build_encoder(device, torch.float32)
+    wrapper = MLMWrapper(config, masking_strategy=AllMasking()).to(device)
+    x = torch.randn(B, N, D, device=device, requires_grad=True)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    res = wrapper(make_padded_input(x, mask), encoder)
+    res.mlm_loss.backward()
+    assert x.grad is None
+
+
+def test_wrapper_severs_gradient_to_upstream_input_under_partial_masking(
+    device: torch.device,
+) -> None:
+    """Regression test: detaching only the target is not enough to close the collapse
+    shortcut. With the default (partial) mask ratio, corrupted_x still holds the live x
+    at unmasked positions, and self-attention mixes those into the masked positions'
+    predictions — so mlm_loss could still reach x through that path unless x is
+    detached before it is used for anything, not only before it is used as the target.
+    This is the realistic training case; test_..._full_masking above is the degenerate
+    edge case that a target-only detach could still pass.
+    """
+    config = MLMWrapperConfig(dim=D, mask_ratio=0.15)
+    encoder = _build_encoder(device, torch.float32)
+    wrapper = MLMWrapper(config).to(device)
+    x = torch.randn(B, N, D, device=device, requires_grad=True)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    res = wrapper(make_padded_input(x, mask), encoder)
+    res.mlm_loss.backward()
+    assert x.grad is None
+
+
+def test_wrapper_accepts_custom_masking_strategy(device: torch.device) -> None:
+    config = MLMWrapperConfig(dim=D, mask_ratio=0.5)
+    encoder = _build_encoder(device, torch.float32)
+    wrapper = MLMWrapper(config, masking_strategy=AllMasking()).to(device)
+    x = torch.randn(B, N, D, device=device)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    res = wrapper(make_padded_input(x, mask), encoder)
+    assert res.mlm_loss.shape == ()
+
+
+def test_wrapper_training_out_is_always_clean(device: torch.device) -> None:
+    """out must never carry masking, even in training — only mlm_loss does. This keeps
+    the main pipeline byte-identical whether or not the MLM aux loss trains alongside
+    it, so callers never need to special-case which pass produced `out`. AllMasking
+    maximizes the contrast: if masking leaked into `out`, every token would be replaced
+    by mask_token, making it maximally different from a clean encoder(input) call.
+    """
+    config = MLMWrapperConfig(dim=D, mask_ratio=0.5)
+    encoder = _build_encoder(device, torch.float32)
+    wrapper = MLMWrapper(config, masking_strategy=AllMasking()).to(device)
+    x = torch.randn(B, N, D, device=device)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    input = make_padded_input(x, mask)
+
+    res = wrapper(input, encoder)
+    expected = encoder(input)
+
+    assert torch.equal(res.out, expected)
+
+
+def test_wrapper_eval_mode_passes_through_encoder_unmasked(device: torch.device) -> None:
+    """In eval mode, forward() must skip masking entirely: out matches a clean encoder
+    call exactly, and mlm_loss reports zero — so callers can invoke the wrapper
+    unconditionally in both modes without an if-training branch of their own.
+    """
+    config = MLMWrapperConfig(dim=D, mask_ratio=0.5)
+    encoder = _build_encoder(device, torch.float32)
+    wrapper = MLMWrapper(config).to(device)
+    wrapper.eval()
+    x = torch.randn(B, N, D, device=device)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    input = make_padded_input(x, mask)
+
+    res = wrapper(input, encoder)
+    expected = encoder(input)
+
+    assert torch.equal(res.out, expected)
+    assert torch.equal(res.mlm_loss, torch.zeros_like(res.mlm_loss))

{stackformers-4.2.1 → stackformers-4.3.1}/uv.lock

@@ -811,7 +811,7 @@ wheels = [
 
 [[package]]
 name = "stackformers"
-version = "4.2.1"
+version = "4.3.1"
 source = { editable = "." }
 dependencies = [
     { name = "beartype" },