From 4e1636b0e5b95c82da71483f461b04f9d9dd3118 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:04:37 +0200 Subject: [PATCH 01/72] docs(write/sequence_mixer): add module and class docstrings --- nvsubquadratic/modules/sequence_mixer.py | 195 +++++++++++++++++++---- 1 file changed, 163 insertions(+), 32 deletions(-) diff --git a/nvsubquadratic/modules/sequence_mixer.py b/nvsubquadratic/modules/sequence_mixer.py index 2a7f836c..b8585e56 100644 --- a/nvsubquadratic/modules/sequence_mixer.py +++ b/nvsubquadratic/modules/sequence_mixer.py @@ -1,7 +1,51 @@ # TODO: Add license header here -"""QKV-based sequence mixer implementation for ND signals.""" +"""Sequence mixer abstraction layer for ND signals. + +This module provides :class:`QKVSequenceMixer`, the **operator-agnostic dispatch +layer** that sits between the residual block and any concrete sequence-mixing +kernel (Hyena, attention, CKConv, Mamba, etc.). + +Architecture role +----------------- +Every residual block in the network needs a *sequence mixer* — a module that +lets tokens (or spatial positions) exchange information over long ranges. +Rather than hard-coding a specific operator into each block, the network passes +a ``mixer_cfg`` :class:`~nvsubquadratic.lazy_config.LazyConfig` down to +:class:`QKVSequenceMixer`, which instantiates the concrete mixer and wraps it +with shared QKV input/output projections. This means the rest of the network +(residual blocks, classifiers, diffusion heads) is **entirely agnostic** to +which sequence mixer is in use; swapping Hyena for attention, or attention for +CKConv, requires only a config change. + +Dispatch pattern +---------------- +The dispatch is performed by :func:`~nvsubquadratic.lazy_config.instantiate` +acting on ``mixer_cfg``. Any class whose constructor accepts ``(q, k, v, +cp_group, **kwargs)`` in its ``forward`` method can be used as the inner mixer. +The currently supported inner mixers are: + +* :class:`~nvsubquadratic.modules.hyena_nd.Hyena` — gated global-conv mixer + (subquadratic in sequence length). +* :class:`~nvsubquadratic.modules.attention.Attention` — multi-head + self-attention (quadratic in sequence length, but faster for short sequences + and easy to compose with RoPE). +* :class:`~nvsubquadratic.modules.ckconv_nd.CKConvND` — continuous-kernel conv + (any spatial rank, learned kernel parametrisation via an MLP). +* :class:`~nvsubquadratic.modules.mamba_nd.MambaNd` — Mamba SSM variant for + ND inputs. + +Input / output layout +--------------------- +All tensors flowing through this module use **channels-last** layout:: + + x: [B, *spatial, C] + +where ``B`` is batch size, ``spatial`` is one or more spatial axes (e.g. +``(T,)`` for sequences, ``(H, W)`` for images, ``(D, H, W)`` for volumes), +and ``C = hidden_dim``. +""" from typing import Callable @@ -11,10 +55,50 @@ class QKVSequenceMixer(torch.nn.Module): - """QKV sequence mixer with configurable projections and initialization. - - Wraps an inner mixer (e.g. Hyena) with linear QKV input and output - projections, mirroring the structure of ``ViT5Attention``. + """Operator-agnostic sequence mixer with shared QKV and output projections. + + :class:`QKVSequenceMixer` mirrors the structure of + :class:`~nvsubquadratic.modules.attention.Attention` (and ViT5Attention) + so that any inner mixer — Hyena, attention, CKConv, Mamba — can be dropped + in without changing the surrounding residual block. + + The forward pass is: + + .. code-block:: text + + x ─[Linear(C, 3C)]──► split ──► Q, K, V + │ + inner_mixer(Q, K, V, cp_group, **kwargs) + │ + [Linear(C, C)]──► y + + The QKV projection packs all three projections into a single + ``Linear(C, 3·C)`` call for efficiency; the output projection maps back to + ``C``. + + Attributes: + mixer (torch.nn.Module): The instantiated inner sequence-mixing + operator (e.g. :class:`~nvsubquadratic.modules.hyena_nd.Hyena`). + qkv_proj (torch.nn.Linear): Combined Q+K+V input projection, + shape ``(C, 3·C)``. + out_proj (torch.nn.Linear): Output projection, shape ``(C, C)``. + + Example:: + + from nvsubquadratic.lazy_config import LazyConfig + from nvsubquadratic.modules.hyena_nd import Hyena + + mixer_cfg = LazyConfig(Hyena)( + global_conv_cfg=..., + short_conv_cfg=..., + gate_nonlinear_cfg=..., + pixelhyena_norm_cfg=..., + qk_norm_cfg=None, + ) + block = QKVSequenceMixer(hidden_dim=256, mixer_cfg=mixer_cfg) + + x = torch.randn(2, 32, 32, 256) # [B, H, W, C] + y = block(x) # [B, H, W, C] """ def __init__( @@ -26,17 +110,39 @@ def __init__( init_method_in: Callable[[int], Callable[[torch.Tensor], torch.Tensor]] | None = None, init_method_out: Callable[[int], Callable[[torch.Tensor], torch.Tensor]] | None = None, ): - """Initialize the QKV sequence mixer. + """Initialise the QKV sequence mixer. Args: - hidden_dim: Hidden dimension. - mixer_cfg: LazyConfig for the inner sequence mixer layer. - qkv_bias: Whether the combined QKV projection has a bias term. - out_proj_bias: Whether the output projection has a bias term. - init_method_in: Optional curried initializer ``fn(dim) -> fn(tensor)`` - for the QKV projection weights (and zero-init for bias if present). - init_method_out: Optional curried initializer ``fn(dim) -> fn(tensor)`` - for the output projection weights (and zero-init for bias if present). + hidden_dim: Channel dimension ``C`` of the input / output tensor. + Both ``qkv_proj`` and ``out_proj`` are sized using this value. + mixer_cfg: :class:`~nvsubquadratic.lazy_config.LazyConfig` for the + inner sequence-mixing operator. The target class must accept + ``(q, k, v, cp_group, **kwargs)`` in its ``forward`` method. + Supported targets include + :class:`~nvsubquadratic.modules.hyena_nd.Hyena`, + :class:`~nvsubquadratic.modules.attention.Attention`, + :class:`~nvsubquadratic.modules.ckconv_nd.CKConvND`, and + :class:`~nvsubquadratic.modules.mamba_nd.MambaNd`. + qkv_bias: If ``True``, adds a learnable bias to the combined QKV + projection. The bias is zero-initialised when + ``init_method_in`` is provided. Defaults to ``False``. + out_proj_bias: If ``True``, adds a learnable bias to the output + projection. Zero-initialised when ``init_method_out`` is + provided. Defaults to ``False``. + init_method_in: Optional *curried* weight initialiser for + ``qkv_proj``. Must have the signature + ``fn(dim: int) -> fn(tensor: Tensor) -> None``. When provided, + ``fn(hidden_dim)`` is called and the returned callable is applied + to ``qkv_proj.weight.data``. If ``qkv_bias`` is also ``True``, + the bias is zero-initialised. Pass ``None`` to use PyTorch's + default (Kaiming uniform). + init_method_out: Same as ``init_method_in`` but applied to + ``out_proj.weight.data``. Typically a scaled initialiser (e.g. + ``1 / sqrt(num_layers)``) to control residual branch variance. + + Raises: + RuntimeError: Propagated from ``instantiate(mixer_cfg)`` if the + target class cannot be constructed (e.g. missing required args). """ super().__init__() @@ -57,25 +163,33 @@ def __init__( def flop_count(self, spatial_dims: tuple[int, ...], inference: bool = False) -> int: """Count FLOPs for QKV projections + inner mixer + output projection. - Let D = ``self.qkv_proj.in_features`` (hidden_dim), - T = prod(spatial_dims) (number of spatial positions). + Uses the standard multiply-accumulate convention where one FLOP = one + multiply + one add (i.e. the matrix-vector product ``y = Wx`` over + ``T`` tokens costs ``2 · T · in_dim · out_dim`` FLOPs). + + FLOPs breakdown (``D`` = ``hidden_dim``, ``T`` = ``prod(spatial_dims)``): - FLOPs breakdown: - 1. QKV projection (Linear(D, 3D)): 2 * T * D * 3D = 6 * T * D² - Three projections (query, key, value) packed into one linear. - 2. Inner mixer (e.g. Hyena): - Delegated to ``self.mixer.flop_count(spatial_dims, inference)``. - 3. Output projection (Linear(D, D)): 2 * T * D * D = 2 * T * D² + 1. **QKV projection** ``Linear(D, 3D)``: + ``2 · T · D · 3D = 6 · T · D²`` + 2. **Inner mixer** (e.g. Hyena, attention): + Delegated to ``self.mixer.flop_count(spatial_dims, inference)``. + For Hyena this is dominated by the FFT convolution + ``O(T log T · D)``; for attention it is ``O(T² · D)``. + 3. **Output projection** ``Linear(D, D)``: + ``2 · T · D² `` - Total: 8 * T * D² + inner_mixer_flops. + Total (excluding inner mixer): ``8 · T · D²``. Args: - spatial_dims: Spatial dimensions of the signal, e.g. (H, W). - Linear projections operate on T = prod(spatial_dims) tokens. - inference: Passed through to the inner mixer. + spatial_dims: Spatial extents of the input signal, e.g. ``(H, W)`` + for images or ``(T,)`` for 1D sequences. Linear projections + treat the flattened token count ``T = prod(spatial_dims)`` as + the sequence length. + inference: Forwarded to ``self.mixer.flop_count``. Some mixers + (e.g. autoregressive Mamba) have different inference-time costs. Returns: - Total FLOPs as an integer. + Total FLOPs as a non-negative integer. """ D = self.qkv_proj.in_features T = 1 @@ -94,15 +208,32 @@ def flop_count(self, spatial_dims: tuple[int, ...], inference: bool = False) -> def forward( self, x: torch.Tensor, cp_group: torch.distributed.ProcessGroup = None, **mixer_kwargs ) -> torch.Tensor: - """Forward pass of the QKV sequence mixer. + """Run the QKV-project → mix → output-project forward pass. Args: - x: torch.Tensor - The input tensor of shape [batch_size, *spatial_dims, hidden_dim]. - cp_group: torch.distributed.ProcessGroup - Context parallel process group. - **mixer_kwargs: Forwarded to the inner mixer (e.g. ``conditioning`` for FiLM). + x: Input tensor of shape ``(B, *spatial, C)`` where ``B`` is batch + size, ``spatial`` is one or more spatial axes (e.g. ``(T,)`` + for 1D, ``(H, W)`` for 2D, ``(D, H, W)`` for 3D), and + ``C = hidden_dim``. + cp_group: Optional context-parallel process group + (:class:`torch.distributed.ProcessGroup`). When provided, the + input is assumed to be already split across ranks along the + spatial axis, and the inner mixer is responsible for the + cross-rank communication (e.g. AllToAll for Hyena, + ring-attention for :class:`~nvsubquadratic.modules.attention.Attention`). + Pass ``None`` (default) for single-GPU / non-distributed runs. + **mixer_kwargs: Additional keyword arguments forwarded verbatim to + ``self.mixer.forward``. Common keys: + + * ``conditioning`` (:class:`torch.Tensor`, shape + ``(B, cond_dim)``): FiLM conditioning signal used by + :class:`~nvsubquadratic.modules.hyena_nd.Hyena` when a + ``condition_mixer`` is attached. + * Any other mixer-specific arguments. Returns: - torch.Tensor - The output tensor of shape [batch_size, *spatial_dims, hidden_dim]. + Output tensor of shape ``(B, *spatial, C)`` — same layout as the + input. """ # Q, K, V projections via single linear qkv = self.qkv_proj(x) From 89733c3c1b2473cd881ba3902317584fa5cf364d Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:05:21 +0200 Subject: [PATCH 02/72] docs(write/hyena_nd): add module and class docstrings with math context --- nvsubquadratic/modules/hyena_nd.py | 339 ++++++++++++++++++++++------- 1 file changed, 255 insertions(+), 84 deletions(-) diff --git a/nvsubquadratic/modules/hyena_nd.py b/nvsubquadratic/modules/hyena_nd.py index e9bbde90..8044f171 100644 --- a/nvsubquadratic/modules/hyena_nd.py +++ b/nvsubquadratic/modules/hyena_nd.py @@ -1,33 +1,79 @@ # TODO: Add license header here -"""Hyena-ND: gated global convolutional mixer for 1D/2D/3D signals. - -Computation graph (per-block): - - Q, K, V ← linear projections of input (done outside this module) +r"""Hyena-ND: gated global convolutional mixer for 1D/2D/3D signals. + +Background +---------- +The Hyena operator (Poli et al., "Hyena Hierarchy: Towards Larger Convolutional +Language Models", ICML 2023) replaces the quadratic attention map with a +**subquadratic gated convolution**: two multiplicative gates sandwich a +long-range (global-kernel) depthwise convolution whose kernel is generated +implicitly by a small neural network (see ``kernels_nd.py``). The operator +achieves O(N log N) time complexity in sequence length N rather than the O(N²) +of dense attention. + +This module generalises the original 1D Hyena to **arbitrary spatial rank** +(1D sequences, 2D images, 3D volumes) by routing the global convolution +through ``CKConvND`` / ``fftconv{1,2,3}d`` primitives from +``nvsubquadratic.ops.fftconv``. + +Computation graph (per forward call) +------------------------------------- +Given projections Q, K, V ∈ R^{B × C × *spatial} (channels-first internally): + +.. code-block:: none + + short_conv([Q; K; V]) — optional depthwise short conv on concat QKV │ - short_conv([Q; K; V]) depthwise short conv on concatenated QKV + QK-Norm(Q [, K]) — optional per-channel normalisation + │ (K is normalised only when gate_nonlinear = Identity) + z = Q ⊙ σ(K) — first multiplicative gate │ - QK-Norm(Q [, K]) optional per-channel normalization - │ (K is only normalized when gate_nonlinear is Identity) - z = Q ⊙ σ(K) first multiplicative gate + PixelHyena-Norm(z) — optional normalisation (GroupNorm / RMSNorm / …) │ - PixelHyena-Norm(z) optional normalization (GroupNorm / RMSNorm / ...) + h = GlobalConv(z) — long-range FFT convolution via CKConvND │ - h = GlobalConv(z) long-range convolution (FFTConv, etc.) + y = h ⊙ σ₂(V) — second multiplicative gate │ - y = h ⊙ σ₂(V) second multiplicative gate + Output-Norm(y) — optional normalisation before projection │ - Output-Norm(y) optional normalization before projection - │ - return y [B, *spatial, C] + return y — [B, *spatial, C] (channels-last on exit) -σ denotes `gate_nonlinear` (first gate) and σ₂ denotes `gate_nonlinear_2` -(second gate). By default σ₂ = σ. When both are Identity the gates +σ denotes ``gate_nonlinear`` (first gate) and σ₂ denotes ``gate_nonlinear_2`` +(second gate). By default σ₂ = σ. When both are ``Identity`` the gates reduce to plain element-wise products, recovering a linear variant closer -to Mamba's selective-scan formulation. Setting σ=SiLU, σ₂=Sigmoid follows -the gated attention formulation. +to Mamba's selective-scan formulation. Setting σ = SiLU, σ₂ = Sigmoid follows +the gated attention formulation from Hyena. + +ND generalisation +----------------- +The Hyena paper targets 1D autoregressive sequences with causal convolutions. +This implementation extends the design to spatial data: + +* The short conv is a standard ``torch.nn.Conv{1,2,3}d`` (or a distributed + equivalent) with a small kernel (e.g. 3×3 for images). +* The global conv is ``CKConvND``, which generates its kernel with a Random + Fourier Feature MLP (``kernels_nd.py``) and convolves it via + ``fftconv{1,2,3}d`` from ``nvsubquadratic.ops.fftconv``. For 2D/3D signals + the convolution is non-causal by default; causal 1D mode is preserved. + +Context parallelism +------------------- +When ``cp_group`` is supplied in ``forward``, the module uses +``AllToAllSingleFunction`` to shard the spatial dimension across devices +while gathering channels, applies the short conv globally, and then shards +back. The global conv receives only the local spatial slice (it must be +context-parallel-aware itself). + +Related modules +--------------- +* ``nvsubquadratic.modules.kernels_nd`` — implicit kernel parametrisation + (``CKConvKernelND``, ``RandomFourierPositionalEmbeddingND``) +* ``nvsubquadratic.ops.fftconv`` — FFT convolution primitives consumed by the + global conv +* ``nvsubquadratic.modules.ckconv_nd`` — ``CKConvND``, the usual choice for + ``global_conv_cfg`` """ from typing import Optional @@ -46,24 +92,67 @@ class Hyena(torch.nn.Module): - """Gated global convolutional mixer for ND signals. - - Two multiplicative gates sandwich a long-range (global) convolution: - - z = Q ⊙ σ(K) — first gate - h = GlobalConv(z) - y = h ⊙ σ₂(V) — second gate - - where σ is ``gate_nonlinear`` and σ₂ is ``gate_nonlinear_2`` (defaults - to σ when not provided). Setting both to Identity gives plain - element-wise products, recovering a linear gating variant. - - Optional components (each disabled by passing Identity or None): - - Short depthwise convolution on concatenated [Q, K, V] - - QK normalization (Q always; K only when σ = Identity) - - PixelHyena normalization between first gate and global conv - - Output normalization after second gate - - Context parallelism via AllToAll communication + r"""Gated global convolutional mixer for ND signals. + + The Hyena operator computes the following gated convolution (all tensors + channels-first internally, channels-last on the public interface): + + .. math:: + + z &= Q \odot \sigma(K) \\ + h &= \mathrm{GlobalConv}(z) \\ + y &= h \odot \sigma_2(V) + + where :math:`\sigma` is ``gate_nonlinear``, :math:`\sigma_2` is + ``gate_nonlinear_2`` (defaults to :math:`\sigma`), and + :math:`\mathrm{GlobalConv}` is a depthwise FFT convolution whose kernel is + generated on-the-fly by an implicit MLP (``CKConvND``). + + Setting both gates to ``Identity`` gives a **linear** gating variant + (element-wise products only). Setting :math:`\sigma = \mathrm{SiLU}` and + :math:`\sigma_2 = \mathrm{Sigmoid}` matches the gated attention formulation + used in the original Hyena paper. + + Paper reference + --------------- + Poli et al., "Hyena Hierarchy: Towards Larger Convolutional Language + Models", ICML 2023. The two-gate structure corresponds to the H3-style + decomposition described in Section 3. The ND extension replaces the + causal 1D FFT conv with a non-causal ND FFT conv (``CKConvND``). + + Optional components (each disabled by passing ``Identity`` or ``None``): + - Short depthwise convolution on concatenated ``[Q, K, V]`` + - QK normalisation (Q always; K only when :math:`\sigma = \mathrm{Identity}`) + - PixelHyena normalisation between first gate and global conv + - Output normalisation after second gate + - Context parallelism via AllToAll communication (``cp_group`` argument) + + Attributes: + global_conv (torch.nn.Module): Long-range global convolution, typically + ``CKConvND``. Must expose ``hidden_dim`` and + ``flop_count(spatial_dims, inference)`` for FLOP counting. + short_conv (torch.nn.Module): Short depthwise convolution applied to + the concatenated ``[Q, K, V]`` tensor (3·C input channels). + Must be one of ``torch.nn.Conv{1,2,3}d``, + ``DistributedDepthwiseConv{1,2,3}d``, or ``torch.nn.Identity``. + gate_nonlinear (torch.nn.Module): Activation :math:`\sigma` for the + first gate. Applied to K before multiplying with Q. + gate_nonlinear_2 (torch.nn.Module): Activation :math:`\sigma_2` for + the second gate. Applied to V before multiplying with h. + Shares the same object as ``gate_nonlinear`` when + ``gate_nonlinear_2_cfg`` is ``None``. + pixelhyena_norm (torch.nn.Module): Normalisation layer applied to + ``z = Q ⊙ σ(K)`` before the global conv. Parameters are + excluded from weight-decay via ``_no_weight_decay = True``. + output_norm (torch.nn.Module): Normalisation layer applied to + ``y = h ⊙ σ₂(V)`` before returning. Parameters are excluded + from weight-decay. + q_norm (torch.nn.Module | None): Per-channel normalisation for Q. + ``None`` when ``qk_norm_cfg`` is ``None``. + k_norm (torch.nn.Module): Per-channel normalisation for K. + ``Identity`` when the gate is nonlinear (magnitude already bounded + by :math:`\sigma`); a fresh instance of ``qk_norm_cfg`` otherwise. + ``None`` when ``qk_norm_cfg`` is ``None``. """ def __init__( @@ -76,22 +165,44 @@ def __init__( output_norm_cfg: LazyConfig = LazyConfig(torch.nn.Identity)(), gate_nonlinear_2_cfg: Optional[LazyConfig] = None, ): - """Constructor. + r"""Construct a Hyena gated global convolutional mixer. + + All ``*_cfg`` arguments are ``LazyConfig`` objects that are + instantiated inside ``__init__`` via ``nvsubquadratic.lazy_config.instantiate``. + This pattern allows full Python configurability without importing + module classes at config-definition time. Args: - global_conv_cfg: Global (long-range) convolutional layer. - short_conv_cfg: Short depthwise conv applied to concatenated [Q, K, V]. - Must produce a ConvNd, DistributedDepthwiseConvNd, or Identity. - gate_nonlinear_cfg: Activation for the first multiplicative gate (e.g. SiLU). - Use Identity for linear gating. - pixelhyena_norm_cfg: Normalization between first gate and global conv. - Use Identity to disable. - qk_norm_cfg: Per-channel normalization for Q (and K when gate is Identity). - None to disable. Separate instances are created for Q and K to - support stateful norms (e.g. RMSNorm with learnable scale). - output_norm_cfg: Normalization after the second gate. Defaults to Identity. - gate_nonlinear_2_cfg: Activation for the second multiplicative gate. - If None (default), reuses gate_nonlinear_cfg for both gates. + global_conv_cfg: ``LazyConfig`` for the long-range global + convolution (e.g. ``CKConvND``). The instantiated module must + expose ``hidden_dim: int`` and + ``flop_count(spatial_dims, inference) -> int``. + short_conv_cfg: ``LazyConfig`` for the short depthwise conv applied + to the concatenated ``[Q; K; V]`` tensor (3·C input channels). + Must instantiate to one of ``torch.nn.Conv{1,2,3}d``, + ``DistributedDepthwiseConv{1,2,3}d``, or ``torch.nn.Identity``. + Use ``Identity`` to skip the short conv entirely. + gate_nonlinear_cfg: ``LazyConfig`` for the first-gate activation + :math:`\sigma(K)` (e.g. ``SiLU``). Use ``Identity`` for + linear gating. + pixelhyena_norm_cfg: ``LazyConfig`` for the normalisation applied + between the first gate and the global conv. Use ``Identity`` + to disable. Parameters receive ``_no_weight_decay = True``. + qk_norm_cfg: ``LazyConfig`` for per-channel normalisation of Q (and + K when the gate is ``Identity``). Pass ``None`` to disable + QK-norm entirely. Two separate instances are created (one for + Q, one for K) so that stateful norms (e.g. ``RMSNorm`` with a + learnable scale) keep independent parameters. + output_norm_cfg: ``LazyConfig`` for the normalisation applied after + the second gate. Defaults to ``Identity`` (no normalisation). + Parameters receive ``_no_weight_decay = True``. + gate_nonlinear_2_cfg: ``LazyConfig`` for the second-gate activation + :math:`\sigma_2(V)`. If ``None`` (default), both gates share + the same activation object (``self.gate_nonlinear``). + + Raises: + AssertionError: If the instantiated ``short_conv`` is not one of + the supported Conv / DistributedDepthwiseConv / Identity types. """ super().__init__() @@ -145,7 +256,16 @@ def __init__( self.k_norm = None def extra_repr(self) -> str: - """Return extra representation string for the module.""" + """Return a compact summary of key configuration choices. + + Included fields: + - ``q_norm`` / ``k_norm`` class names (or ``"None"``). + - ``gates=<σ>/<σ₂>`` when the two gate activations differ. + - ``is_causal`` when the global conv exposes that attribute. + + Returns: + Comma-separated string suitable for ``repr(module)`` output. + """ is_causal = getattr(self.global_conv, "is_causal", None) q_norm_str = self.q_norm.__class__.__name__ if self.q_norm is not None else "None" k_norm_str = self.k_norm.__class__.__name__ if self.k_norm is not None else "None" @@ -159,34 +279,53 @@ def extra_repr(self) -> str: return ", ".join(parts) def flop_count(self, spatial_dims: tuple[int, ...], inference: bool = False) -> int: - """Count FLOPs for the Hyena gated global convolutional mixer. - - Let C = hidden_dim (per projection), S = prod(spatial_dims). - - FLOPs breakdown: - 1. Short depthwise conv on concatenated [Q, K, V] (3C channels): - 2 * 3C * S * k_prod, where k_prod = product of kernel sizes. - Each output element: k_prod MACs for 1 depthwise filter. - Skipped when short_conv is Identity. - 2. QK-Norm (when ``self.q_norm is not None``): - Q: 3 * C * S (RMSNorm-like). - K: 3 * C * S only when ``self.gate_nonlinear`` is Identity - (linear gating); a nonlinear σ(K) already bounds magnitude. - 3. First gate Q ⊙ σ(K): C * S (multiply). - + C * S for activation on K if gate_nonlinear is not Identity. - 4. PixelHyena norm (if not Identity): 3 * C * S. - 5. Global convolution (CKConvND): - Delegated to ``self.global_conv.flop_count(spatial_dims, inference)``. - 6. Second gate h ⊙ σ₂(V): C * S (multiply). - + C * S for activation on V if gate_nonlinear_2 is not Identity. - 7. Output norm (if not Identity): 3 * C * S. + r"""Count FLOPs for one forward pass of the Hyena mixer. + + Let ``C = self.global_conv.hidden_dim`` (the per-head channel count) + and ``S = prod(spatial_dims)`` (total number of spatial positions). + All counts use the **multiply-add = 1 FLOP** convention (i.e. a MAC + counts as 1). + + FLOP breakdown: + + 1. **Short depthwise conv** on concatenated ``[Q; K; V]`` + (``3·C`` input channels): + + .. math:: + + 2 \cdot \frac{in\_ch}{groups} \cdot out\_ch \cdot S \cdot k\_prod + + where :math:`k\_prod = \prod_d kernel\_size_d`. Skipped when + ``short_conv`` is ``Identity``. + + 2. **QK-Norm** (when ``self.q_norm is not None``): + ``3·C·S`` for Q; additional ``3·C·S`` for K only when + ``gate_nonlinear`` is ``Identity`` (linear gating). + + 3. **First gate** :math:`z = Q \odot \sigma(K)`: + ``C·S`` for the elementwise multiply, plus ``C·S`` for the + activation :math:`\sigma` when it is not ``Identity``. + + 4. **PixelHyena norm** (when not ``Identity``): ``3·C·S``. + + 5. **Global convolution**: delegated to + ``self.global_conv.flop_count(spatial_dims, inference)``. + + 6. **Second gate** :math:`y = h \odot \sigma_2(V)`: + ``C·S`` for the multiply, plus ``C·S`` for :math:`\sigma_2` when + not ``Identity``. + + 7. **Output norm** (when not ``Identity``): ``3·C·S``. Args: - spatial_dims: Spatial dimensions of the input, e.g. (H, W) for 2D. - inference: Passed through to CKConvND for kernel generation caching. + spatial_dims: Spatial extent of the input per axis, e.g. ``(H, W)`` + for a 2-D feature map of shape ``[B, C, H, W]``. + inference: Forwarded to ``self.global_conv.flop_count``; some + implementations skip re-generating the kernel at inference time + when it is cached. Returns: - Total FLOPs as an integer. + Total FLOP count as an integer (multiply-add = 1 FLOP convention). """ C = self.global_conv.hidden_dim S = 1 @@ -242,19 +381,51 @@ def forward( cp_group: torch.distributed.ProcessGroup = None, **mixer_kwargs, ) -> torch.Tensor: - """Compute y = OutputNorm( GlobalConv( Norm( Q ⊙ σ(K) ) ) ⊙ σ(V) ). + r"""Compute the Hyena gated global convolution. + + Implements: + + .. math:: + + y = \mathrm{OutputNorm}\!\bigl( + \mathrm{GlobalConv}\!\bigl( + \mathrm{Norm}(Q \odot \sigma(K)) + \bigr) \odot \sigma_2(V) + \bigr) + + Tensors enter and leave in **channels-last** layout + ``[B, *spatial, C]``. Internally the module works channels-first + ``[B, C, *spatial]`` for the short conv and global conv. + + Context parallelism (``cp_group``) + ----------------------------------- + When ``cp_group`` is provided and has size > 1, the method applies two + AllToAll communications around the short conv so that each device sees + the full spatial extent during the convolution: + + 1. Before short conv: ``split_to_full`` — gather spatial shards, + split along the channel dim. + 2. After short conv: ``full_to_split`` — scatter spatial, gather + channels back. - All tensors are channel-last on entry and exit. + The global conv receives only the local spatial slice and is expected + to handle its own CP communication internally. Args: - query: ``[B, *spatial, C]`` query tensor (from linear projection of input). - key: ``[B, *spatial, C]`` key tensor. - value: ``[B, *spatial, C]`` value tensor. - cp_group: Context-parallel process group. None disables CP. - **mixer_kwargs: Forwarded to the global conv (e.g. ``conditioning`` for FiLM). + query: ``[B, *spatial, C]`` — query tensor, typically the output + of a linear projection ``W_Q · x``. + key: ``[B, *spatial, C]`` — key tensor, typically ``W_K · x``. + value: ``[B, *spatial, C]`` — value tensor, typically ``W_V · x``. + cp_group: ``torch.distributed.ProcessGroup`` for context + parallelism. ``None`` disables CP (the default for single-GPU + runs). + **mixer_kwargs: Extra keyword arguments forwarded verbatim to + ``self.global_conv`` (e.g. ``conditioning`` for FiLM-conditioned + ``CKConvND``). Returns: - ``[B, *spatial, C]`` output tensor. + ``[B, *spatial, C]`` — output tensor in channels-last layout, + same shape as the inputs. """ # Reshape query, key, and value to [B, C, * spatial_dims] (Required for short convolutional projections). query = rearrange(query, "b ... c -> b c ...") From ed7f1d983bc2408b0b45d1dee15b9d30d4f8b353 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:05:31 +0200 Subject: [PATCH 03/72] docs(review/sequence_mixer): reviewer feedback --- docs/reviews/sequence_mixer_review.md | 146 ++++++++++++++++++++++++++ 1 file changed, 146 insertions(+) create mode 100644 docs/reviews/sequence_mixer_review.md diff --git a/docs/reviews/sequence_mixer_review.md b/docs/reviews/sequence_mixer_review.md new file mode 100644 index 00000000..711d74b2 --- /dev/null +++ b/docs/reviews/sequence_mixer_review.md @@ -0,0 +1,146 @@ +# Review: `nvsubquadratic/modules/sequence_mixer.py` + +Reviewed as a new external collaborator reading alongside the research paper. + +______________________________________________________________________ + +## Issues + +### 1. Module docstring: "dispatch" contract is stated but not defined precisely + +**Quoted text:** "Any class whose constructor accepts `(q, k, v, cp_group, **kwargs)` in its +`forward` method can be used as the inner mixer." + +**Fix:** This describes the `forward` signature but says "constructor" — confusing wording. +Change "constructor" to "the `forward` method". Also, the positional nature of `cp_group` +(it is passed positionally, not as a keyword) should be stated explicitly, because an inner +mixer that defines `cp_group` as `**kwargs` would silently fail to receive it. A one-liner +protocol note would help: + +``` +The inner mixer forward must accept: forward(q, k, v, cp_group, **kwargs) +where cp_group is passed positionally as the fourth argument. +``` + +### 2. Module docstring: `MambaNd` class name is wrong + +**Quoted text:** `:class:`~nvsubquadratic.modules.mamba_nd.MambaNd\`\` + +**Fix:** Check the actual class name in `mamba_nd.py`. If it is `MambaND` (uppercase D), fix +the cross-reference. A broken `:class:` reference silently produces a literal string in Sphinx. + +### 3. Class docstring: Attributes block uses wrong "shape" notation for Linear weights + +**Quoted text:** "qkv_proj (torch.nn.Linear): Combined Q+K+V input projection, shape `(C, 3·C)`." + +**Fix:** `torch.nn.Linear` stores weights with shape `(out_features, in_features)` = `(3C, C)`, +not `(C, 3C)`. The docstring currently states the transpose. Either correct to `(3C, C)` or +just drop the weight-shape hint and write "maps `C` → `3C`" to avoid the ambiguity. + +### 4. `__init__` docstring: `Raises` block is under-specified + +**Quoted text:** "RuntimeError: Propagated from `instantiate(mixer_cfg)` if the target class +cannot be constructed (e.g. missing required args)." + +**Fix:** In practice `instantiate` raises `omegaconf.errors.InstantiationException` or +`hydra._internal.utils.HydraException` (not plain `RuntimeError`) depending on the backend. +If `LazyConfig` uses a different backend, document the actual exception type. At minimum, +note that the exception originates from `LazyConfig.instantiate` and suggest the user check +the `mixer_cfg` target and arguments. + +### 5. `__init__` docstring: initialiser signature example is missing a concrete use case + +**Quoted text:** "Typically a scaled initialiser (e.g. `1 / sqrt(num_layers)`) to control +residual branch variance." + +**Fix:** Add a one-line concrete example showing what `init_method_out` looks like in practice, +so a reader knows what curried form to write: + +```python +import math + +init_method_out = lambda dim: lambda w: torch.nn.init.normal_( + w, std=1 / math.sqrt(num_layers) +) +``` + +This follows the GPT / Megatron pattern and is the primary use case; without an example the +curried signature is hard to guess. + +### 6. `flop_count` docstring: trailing whitespace in formula line + +**Quoted text:** "`2 · T · D² `" (note the trailing space before the closing backticks). + +**Fix:** Remove the trailing space: ``` "``2 · T · D²``" ```. Minor, but looks sloppy in rendered +Sphinx HTML. + +### 7. `flop_count` docstring: does not mention bias FLOPs or that biases are ignored + +**Quoted text:** "Uses the standard multiply-accumulate convention…" + +**Fix:** State explicitly that bias additions are excluded from the count (standard in ML FLOP +counting). If biases are eventually included, callers would get wrong numbers. One line suffices: +"Note: bias additions are excluded, following the standard ML FLOP-counting convention." + +### 8. `flop_count` docstring: inner mixer delegation may raise `AttributeError` + +**Quoted text:** "Delegated to `self.mixer.flop_count(spatial_dims, inference)`." + +**Fix:** Not all inner mixers implement `flop_count`. If `self.mixer` is an `Attention` module +that lacks this method, the call raises `AttributeError` at runtime. Add a `Raises` entry: + +``` +Raises: + AttributeError: If the inner mixer does not implement ``flop_count``. +``` + +### 9. `forward` docstring: the `cp_group=None` type annotation in the signature is wrong + +**Quoted text (function signature):** + +```python +def forward(self, x: torch.Tensor, cp_group: torch.distributed.ProcessGroup = None, ... +``` + +**Fix:** The type annotation should be `torch.distributed.ProcessGroup | None = None`, not just +`torch.distributed.ProcessGroup = None`. The current annotation misleads static analysers and +readers into thinking `None` is not a valid value. This is a code change, but fixing it is +better than papering over it in the docstring. + +### 10. `forward` docstring: `conditioning` kwarg shape is under-specified + +**Quoted text:** "`conditioning` (torch.Tensor, shape `(B, cond_dim)`): FiLM conditioning signal" + +**Fix:** The conditioning tensor shape depends on how the `condition_mixer` inside Hyena +processes it — it may be `(B, cond_dim)` before FiLM but the actual contract at the +`QKVSequenceMixer.forward` boundary should be clarified. Also document that passing +`conditioning` to a mixer that does not use it is silently ignored (it falls into +`**mixer_kwargs`), so callers don't need to guard against this. + +### 11. Module docstring: no note on how to add a new mixer type + +**Fix:** A new external collaborator trying to plug in a new operator (say, RWKV) needs to +know exactly what interface to implement. Add a short paragraph or `Note:` block: + +``` +Note: + To register a new mixer type, implement a ``torch.nn.Module`` whose + ``forward(q, k, v, cp_group, **kwargs)`` method follows the channels-last + convention ``[B, *spatial, C]`` and, optionally, implement + ``flop_count(spatial_dims, inference) -> int``. Then pass its + ``LazyConfig`` as ``mixer_cfg`` to :class:`QKVSequenceMixer` — no other + changes are needed. +``` + +### 12. Class docstring: the ASCII diagram does not show where biases are applied + +**Quoted text:** + +``` +x ─[Linear(C, 3C)]──► split ──► Q, K, V +``` + +**Fix:** The diagram is clear, but does not indicate that `qkv_bias` and `out_proj_bias` +optionally add bias terms. Since these default to `False` this is minor, but a parenthetical +`[+ bias?]` on each linear step would make the diagram fully accurate for the non-default +case and prevent confusion when a reader sees biases in a checkpoint. From 5f030204176404bce202e70fb0c3d9c9d00a18f4 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:06:33 +0200 Subject: [PATCH 04/72] docs(review/hyena_nd): reviewer feedback --- docs/reviews/hyena_nd_review.md | 240 ++++++++++++++++++++++++++++++++ 1 file changed, 240 insertions(+) create mode 100644 docs/reviews/hyena_nd_review.md diff --git a/docs/reviews/hyena_nd_review.md b/docs/reviews/hyena_nd_review.md new file mode 100644 index 00000000..427f3cd7 --- /dev/null +++ b/docs/reviews/hyena_nd_review.md @@ -0,0 +1,240 @@ +# Reviewer feedback: `nvsubquadratic/modules/hyena_nd.py` + +Reviewed against the Phase-1 docstrings added in commit `docs(write/hyena_nd)`. + +______________________________________________________________________ + +## 1. Module docstring — paper citation is incomplete and imprecise + +**Text:** `"Poli et al., 'Hyena Hierarchy: Towards Larger Convolutional Language Models', ICML 2023."` + +**Issue:** No arXiv ID or DOI. Collaborators without institutional access cannot look the paper up quickly. Also, the paper was published at ICML 2023 but first appeared on arXiv — give both. + +**Fix:** Change to: + +``` +Poli et al., "Hyena Hierarchy: Towards Larger Convolutional Language Models", +ICML 2023. arXiv:2302.10866. +``` + +______________________________________________________________________ + +## 2. Module docstring — "Section 3" reference is too vague + +**Text (class docstring):** `"The two-gate structure corresponds to the H3-style decomposition described in Section 3."` + +**Issue:** The Hyena paper is 30+ pages; "Section 3" does not pin the exact location. H3 is described in a *related-work* paper (Fu et al., 2023), not in the Hyena paper itself. An external reader will spend time searching. + +**Fix:** Replace with: + +``` +The two-gate structure follows the H3 block (Fu et al., "Hungry Hungry Hippos", +ICLR 2023, arXiv:2212.14052, Section 3.2) and is generalised in Hyena +(Poli et al., arXiv:2302.10866, Section 3, "The Hyena Recurrence"). +``` + +______________________________________________________________________ + +## 3. Module docstring — ND generalisation section does not mention boundary conditions + +**Text:** `"For 2D/3D signals the convolution is non-causal by default; causal 1D mode is preserved."` + +**Issue:** A new collaborator will not know whether 2D/3D uses zero-padding (linear conv) or circular conv. This matters because the choice affects the FLOP count and which `fftconv` function is invoked. The `mixed_fftconv` path (#120) also exists now. + +**Fix:** Add after the sentence: + +``` +By default the 2D/3D path uses zero-padded (linear) FFT convolution +(``fftconv2d`` / ``fftconv3d``), matching ``torch.nn.ConvNd(padding='same')`` +semantics. Set the ``circular`` flag on ``CKConvND`` to switch to periodic +boundary conditions, or use ``mixed_fftconv`` for per-axis mixed BCs +(see ``nvsubquadratic.ops.mixed_fftconv``). +``` + +______________________________________________________________________ + +## 4. Class docstring — `Attributes` block: `k_norm` type annotation is contradictory + +**Text:** + +``` +k_norm (torch.nn.Module): Per-channel normalisation for K. + ``Identity`` when the gate is nonlinear (magnitude already bounded + by :math:`\sigma`); a fresh instance of ``qk_norm_cfg`` otherwise. + ``None`` when ``qk_norm_cfg`` is ``None``. +``` + +**Issue:** The type is listed as `torch.nn.Module` but the final sentence says it can be `None`. The actual code sets `self.k_norm = None` when `qk_norm_cfg is None`. The type should be `torch.nn.Module | None`. + +**Fix:** Change to `k_norm (torch.nn.Module | None):` and restate the None case first: + +``` +k_norm (torch.nn.Module | None): Per-channel normalisation for K. + ``None`` when ``qk_norm_cfg`` is ``None`` (QK-norm entirely disabled). + ``torch.nn.Identity`` when the gate is nonlinear (σ already bounds + K's magnitude); a fresh instance of ``qk_norm_cfg`` when the gate is + ``Identity`` (linear gating). +``` + +______________________________________________________________________ + +## 5. `__init__` docstring — `output_norm_cfg` default value is misleading + +**Text:** ``` "Defaults to ``Identity`` (no normalisation)." ``` + +**Issue:** The actual default is `LazyConfig(torch.nn.Identity)()`, not a plain `torch.nn.Identity`. This distinction matters because `LazyConfig(torch.nn.Identity)()` is a frozen lazy config object, not a module; the module is only created by `instantiate` inside `__init__`. An external caller who tries to pass `torch.nn.Identity()` (an already-instantiated module) will get a confusing error from `instantiate`. + +**Fix:** Clarify the type contract: + +``` +output_norm_cfg: ``LazyConfig`` for the normalisation applied after the second + gate. Defaults to a ``LazyConfig`` wrapping ``torch.nn.Identity`` (no + normalisation). Do **not** pass an already-instantiated module — pass a + ``LazyConfig`` object that wraps the class. +``` + +______________________________________________________________________ + +## 6. `flop_count` docstring — FLOP count for QK-Norm assumes RMSNorm always + +**Text:** `"3·C·S for Q; additional 3·C·S for K"` + +**Issue:** The factor of 3 (mean, variance, scale) is correct for RMSNorm / LayerNorm but is silently assumed here. A GroupNorm or a simple scalar multiply would have different counts. The docstring presents this as exact without flagging the assumption. + +**Fix:** Add a note: + +``` +The factor of 3 assumes an RMSNorm-like norm (sum-of-squares + rsqrt + +elementwise scale). Other norm types will differ; this is an approximation. +``` + +______________________________________________________________________ + +## 7. `flop_count` docstring — missing explanation of why `out_ch` appears in the depthwise conv formula + +**Text:** + +``` +2 · (in_ch / groups) · out_ch · S · k_prod +``` + +**Issue:** For a true depthwise conv `in_ch == out_ch == groups`, so `(in_ch / groups) = 1` and the expression collapses to `2 · out_ch · S · k_prod`. The formula as written is actually the general grouped-conv formula, and it happens to be correct for depthwise but the derivation is unclear. An external reader will wonder why `out_ch` appears when they expect `in_ch // groups = 1`. + +**Fix:** Add an inline note: + +``` +For a pure depthwise conv (groups == in_ch == out_ch) this simplifies to +``2 · out_ch · S · k_prod``; the grouped formula is written here to handle +partially-grouped convolutions (e.g. ``DistributedDepthwiseConvNd``). +``` + +______________________________________________________________________ + +## 8. `forward` docstring — CP AllToAll semantics need more precision + +**Text:** + +``` +1. Before short conv: ``split_to_full`` — gather spatial shards, split along the channel dim. +2. After short conv: ``full_to_split`` — scatter spatial, gather channels back. +``` + +**Issue:** The terms `split_to_full` and `full_to_split` are internal string constants from `AllToAllSingleFunction`. A new collaborator does not know which spatial axis is sharded or what "split along channel" means quantitatively. Is it the *first* spatial axis? What is the shard size? + +**Fix:** Add: + +``` +The AllToAll shards along ``dim=2`` (the first spatial axis) and gathers +along ``dim=1`` (the channel axis). After the AllToAll, each device holds +the full spatial extent but only ``C / cp_size`` channels. After the reverse +AllToAll, the original ``C`` channels are restored with ``spatial / cp_size`` +positions per device. +``` + +______________________________________________________________________ + +## 9. `forward` docstring — `query` variable re-use is confusing and undocumented + +**Text:** (no docstring for this; it is inline code) + +In the body, `query` is overwritten twice: + +```python +query = query * self.gate_nonlinear(key) # now z, not Q +... +# then passed to global_conv as z +``` + +The variable is named `query` but after the first gate it represents the *gated intermediate* `z`. This is not mentioned anywhere. + +**Fix:** Add a note in the `forward` docstring under a "Implementation note" or "Variable naming" paragraph: + +``` +Implementation note: ``query`` is overwritten in-place (semantically) after +the first gate to hold the gated intermediate ``z = Q ⊙ σ(K)``. The original +Q tensor is no longer accessible after that point. This is intentional to +avoid an extra allocation. +``` + +______________________________________________________________________ + +## 10. `extra_repr` — `k_norm` can be `None` but is accessed unconditionally + +**Text:** + +```python +k_norm_str = self.k_norm.__class__.__name__ if self.k_norm is not None else "None" +``` + +**Issue:** This is correct code, but the docstring does not mention that `k_norm` may be `None` here (consistent with Issue 4). The docstring for `extra_repr` currently says nothing about when fields are omitted. + +**Fix:** Add to the `extra_repr` docstring: + +``` +When ``self.q_norm`` and ``self.k_norm`` are ``None`` (QK-norm disabled), +the strings ``"q_norm=None"`` and ``"k_norm=None"`` are still included so +the disabled state is explicit in ``repr(module)``. +``` + +______________________________________________________________________ + +## 11. Missing usage example anywhere in the file + +**Issue:** There is no `Example:` block anywhere in the module or class docstring. An external collaborator cannot quickly see how to wire up a `Hyena` block. The `kernels_nd.py` module even has a `# For test, please run:` note, showing that examples are expected. + +**Fix:** Add an `Example:` section to the `Hyena` class docstring showing the minimal construction with `CKConvND`, e.g.: + +```python +Example: + >>> import torch + >>> from nvsubquadratic.lazy_config import LazyConfig + >>> from nvsubquadratic.modules.hyena_nd import Hyena + >>> from nvsubquadratic.modules.ckconv_nd import CKConvND + >>> # Minimal 2D Hyena block (non-causal, no normalisation) + >>> hyena = Hyena( + ... global_conv_cfg=LazyConfig(CKConvND, hidden_dim=64, data_dim=2, ...), + ... short_conv_cfg=LazyConfig(torch.nn.Conv2d, 192, 192, 3, padding=1, groups=192), + ... gate_nonlinear_cfg=LazyConfig(torch.nn.SiLU), + ... pixelhyena_norm_cfg=LazyConfig(torch.nn.Identity), + ... qk_norm_cfg=None, + ... ) + >>> B, H, W, C = 2, 16, 16, 64 + >>> q = k = v = torch.randn(B, H, W, C) + >>> y = hyena(q, k, v) # [2, 16, 16, 64] +``` + +______________________________________________________________________ + +## 12. Module docstring — context parallelism section does not clarify which spatial axis is sharded + +**Text:** `"shard the spatial dimension across devices"` + +**Issue:** For 2D/3D inputs there are multiple spatial axes. Which one? The code shards `dim=2` (the first spatial axis), but this is never stated in the docstring. + +**Fix:** Change to: + +``` +The module shards along ``dim=2`` (the first spatial axis of the +channels-first ``[B, C, *spatial]`` tensor) while gathering the channel dim. +For 2D inputs ``[B, C, H, W]`` this means row-wise sharding (across H). +``` From b9506703f1ddf49e5b4630692290f402accfdf77 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:06:42 +0200 Subject: [PATCH 05/72] docs(integrate/sequence_mixer): apply reviewer feedback --- docs/reviews/sequence_mixer_review.md | 146 ----------------------- nvsubquadratic/modules/sequence_mixer.py | 93 ++++++++++----- 2 files changed, 66 insertions(+), 173 deletions(-) delete mode 100644 docs/reviews/sequence_mixer_review.md diff --git a/docs/reviews/sequence_mixer_review.md b/docs/reviews/sequence_mixer_review.md deleted file mode 100644 index 711d74b2..00000000 --- a/docs/reviews/sequence_mixer_review.md +++ /dev/null @@ -1,146 +0,0 @@ -# Review: `nvsubquadratic/modules/sequence_mixer.py` - -Reviewed as a new external collaborator reading alongside the research paper. - -______________________________________________________________________ - -## Issues - -### 1. Module docstring: "dispatch" contract is stated but not defined precisely - -**Quoted text:** "Any class whose constructor accepts `(q, k, v, cp_group, **kwargs)` in its -`forward` method can be used as the inner mixer." - -**Fix:** This describes the `forward` signature but says "constructor" — confusing wording. -Change "constructor" to "the `forward` method". Also, the positional nature of `cp_group` -(it is passed positionally, not as a keyword) should be stated explicitly, because an inner -mixer that defines `cp_group` as `**kwargs` would silently fail to receive it. A one-liner -protocol note would help: - -``` -The inner mixer forward must accept: forward(q, k, v, cp_group, **kwargs) -where cp_group is passed positionally as the fourth argument. -``` - -### 2. Module docstring: `MambaNd` class name is wrong - -**Quoted text:** `:class:`~nvsubquadratic.modules.mamba_nd.MambaNd\`\` - -**Fix:** Check the actual class name in `mamba_nd.py`. If it is `MambaND` (uppercase D), fix -the cross-reference. A broken `:class:` reference silently produces a literal string in Sphinx. - -### 3. Class docstring: Attributes block uses wrong "shape" notation for Linear weights - -**Quoted text:** "qkv_proj (torch.nn.Linear): Combined Q+K+V input projection, shape `(C, 3·C)`." - -**Fix:** `torch.nn.Linear` stores weights with shape `(out_features, in_features)` = `(3C, C)`, -not `(C, 3C)`. The docstring currently states the transpose. Either correct to `(3C, C)` or -just drop the weight-shape hint and write "maps `C` → `3C`" to avoid the ambiguity. - -### 4. `__init__` docstring: `Raises` block is under-specified - -**Quoted text:** "RuntimeError: Propagated from `instantiate(mixer_cfg)` if the target class -cannot be constructed (e.g. missing required args)." - -**Fix:** In practice `instantiate` raises `omegaconf.errors.InstantiationException` or -`hydra._internal.utils.HydraException` (not plain `RuntimeError`) depending on the backend. -If `LazyConfig` uses a different backend, document the actual exception type. At minimum, -note that the exception originates from `LazyConfig.instantiate` and suggest the user check -the `mixer_cfg` target and arguments. - -### 5. `__init__` docstring: initialiser signature example is missing a concrete use case - -**Quoted text:** "Typically a scaled initialiser (e.g. `1 / sqrt(num_layers)`) to control -residual branch variance." - -**Fix:** Add a one-line concrete example showing what `init_method_out` looks like in practice, -so a reader knows what curried form to write: - -```python -import math - -init_method_out = lambda dim: lambda w: torch.nn.init.normal_( - w, std=1 / math.sqrt(num_layers) -) -``` - -This follows the GPT / Megatron pattern and is the primary use case; without an example the -curried signature is hard to guess. - -### 6. `flop_count` docstring: trailing whitespace in formula line - -**Quoted text:** "`2 · T · D² `" (note the trailing space before the closing backticks). - -**Fix:** Remove the trailing space: ``` "``2 · T · D²``" ```. Minor, but looks sloppy in rendered -Sphinx HTML. - -### 7. `flop_count` docstring: does not mention bias FLOPs or that biases are ignored - -**Quoted text:** "Uses the standard multiply-accumulate convention…" - -**Fix:** State explicitly that bias additions are excluded from the count (standard in ML FLOP -counting). If biases are eventually included, callers would get wrong numbers. One line suffices: -"Note: bias additions are excluded, following the standard ML FLOP-counting convention." - -### 8. `flop_count` docstring: inner mixer delegation may raise `AttributeError` - -**Quoted text:** "Delegated to `self.mixer.flop_count(spatial_dims, inference)`." - -**Fix:** Not all inner mixers implement `flop_count`. If `self.mixer` is an `Attention` module -that lacks this method, the call raises `AttributeError` at runtime. Add a `Raises` entry: - -``` -Raises: - AttributeError: If the inner mixer does not implement ``flop_count``. -``` - -### 9. `forward` docstring: the `cp_group=None` type annotation in the signature is wrong - -**Quoted text (function signature):** - -```python -def forward(self, x: torch.Tensor, cp_group: torch.distributed.ProcessGroup = None, ... -``` - -**Fix:** The type annotation should be `torch.distributed.ProcessGroup | None = None`, not just -`torch.distributed.ProcessGroup = None`. The current annotation misleads static analysers and -readers into thinking `None` is not a valid value. This is a code change, but fixing it is -better than papering over it in the docstring. - -### 10. `forward` docstring: `conditioning` kwarg shape is under-specified - -**Quoted text:** "`conditioning` (torch.Tensor, shape `(B, cond_dim)`): FiLM conditioning signal" - -**Fix:** The conditioning tensor shape depends on how the `condition_mixer` inside Hyena -processes it — it may be `(B, cond_dim)` before FiLM but the actual contract at the -`QKVSequenceMixer.forward` boundary should be clarified. Also document that passing -`conditioning` to a mixer that does not use it is silently ignored (it falls into -`**mixer_kwargs`), so callers don't need to guard against this. - -### 11. Module docstring: no note on how to add a new mixer type - -**Fix:** A new external collaborator trying to plug in a new operator (say, RWKV) needs to -know exactly what interface to implement. Add a short paragraph or `Note:` block: - -``` -Note: - To register a new mixer type, implement a ``torch.nn.Module`` whose - ``forward(q, k, v, cp_group, **kwargs)`` method follows the channels-last - convention ``[B, *spatial, C]`` and, optionally, implement - ``flop_count(spatial_dims, inference) -> int``. Then pass its - ``LazyConfig`` as ``mixer_cfg`` to :class:`QKVSequenceMixer` — no other - changes are needed. -``` - -### 12. Class docstring: the ASCII diagram does not show where biases are applied - -**Quoted text:** - -``` -x ─[Linear(C, 3C)]──► split ──► Q, K, V -``` - -**Fix:** The diagram is clear, but does not indicate that `qkv_bias` and `out_proj_bias` -optionally add bias terms. Since these default to `False` this is minor, but a parenthetical -`[+ bias?]` on each linear step would make the diagram fully accurate for the non-default -case and prevent confusion when a reader sees biases in a checkpoint. diff --git a/nvsubquadratic/modules/sequence_mixer.py b/nvsubquadratic/modules/sequence_mixer.py index b8585e56..999b22df 100644 --- a/nvsubquadratic/modules/sequence_mixer.py +++ b/nvsubquadratic/modules/sequence_mixer.py @@ -22,8 +22,11 @@ Dispatch pattern ---------------- The dispatch is performed by :func:`~nvsubquadratic.lazy_config.instantiate` -acting on ``mixer_cfg``. Any class whose constructor accepts ``(q, k, v, -cp_group, **kwargs)`` in its ``forward`` method can be used as the inner mixer. +acting on ``mixer_cfg``. Any class whose ``forward`` method accepts +``(q, k, v, cp_group, **kwargs)`` can be used as the inner mixer. Note that +``cp_group`` is passed **positionally** as the fourth argument; an inner mixer +that captures it only via ``**kwargs`` would silently not receive it. + The currently supported inner mixers are: * :class:`~nvsubquadratic.modules.hyena_nd.Hyena` — gated global-conv mixer @@ -33,9 +36,17 @@ and easy to compose with RoPE). * :class:`~nvsubquadratic.modules.ckconv_nd.CKConvND` — continuous-kernel conv (any spatial rank, learned kernel parametrisation via an MLP). -* :class:`~nvsubquadratic.modules.mamba_nd.MambaNd` — Mamba SSM variant for +* :class:`~nvsubquadratic.modules.mamba_nd.Mamba` — Mamba SSM variant for ND inputs. +Note: + To add a new mixer type, implement a :class:`torch.nn.Module` whose + ``forward(q, k, v, cp_group, **kwargs)`` method follows the channels-last + convention ``[B, *spatial, C]`` and, optionally, implement + ``flop_count(spatial_dims, inference) -> int``. Then pass its + :class:`~nvsubquadratic.lazy_config.LazyConfig` as ``mixer_cfg`` to + :class:`QKVSequenceMixer` — no other changes are needed. + Input / output layout --------------------- All tensors flowing through this module use **channels-last** layout:: @@ -66,22 +77,24 @@ class QKVSequenceMixer(torch.nn.Module): .. code-block:: text - x ─[Linear(C, 3C)]──► split ──► Q, K, V - │ - inner_mixer(Q, K, V, cp_group, **kwargs) - │ - [Linear(C, C)]──► y + x ─[Linear(C → 3C, + bias?)]──► split ──► Q, K, V + │ + inner_mixer(Q, K, V, cp_group, **kwargs) + │ + [Linear(C → C, + bias?)]──► y The QKV projection packs all three projections into a single ``Linear(C, 3·C)`` call for efficiency; the output projection maps back to - ``C``. + ``C``. Both projections optionally include a bias term (disabled by + default; see ``qkv_bias`` and ``out_proj_bias``). Attributes: mixer (torch.nn.Module): The instantiated inner sequence-mixing operator (e.g. :class:`~nvsubquadratic.modules.hyena_nd.Hyena`). - qkv_proj (torch.nn.Linear): Combined Q+K+V input projection, - shape ``(C, 3·C)``. - out_proj (torch.nn.Linear): Output projection, shape ``(C, C)``. + qkv_proj (torch.nn.Linear): Combined Q+K+V input projection; + maps ``C`` → ``3·C`` (weight shape ``(3C, C)``). + out_proj (torch.nn.Linear): Output projection; maps ``C`` → ``C`` + (weight shape ``(C, C)``). Example:: @@ -116,13 +129,14 @@ def __init__( hidden_dim: Channel dimension ``C`` of the input / output tensor. Both ``qkv_proj`` and ``out_proj`` are sized using this value. mixer_cfg: :class:`~nvsubquadratic.lazy_config.LazyConfig` for the - inner sequence-mixing operator. The target class must accept - ``(q, k, v, cp_group, **kwargs)`` in its ``forward`` method. - Supported targets include + inner sequence-mixing operator. The target class's ``forward`` + method must accept ``(q, k, v, cp_group, **kwargs)`` where + ``cp_group`` is the fourth positional argument. Supported + targets include :class:`~nvsubquadratic.modules.hyena_nd.Hyena`, :class:`~nvsubquadratic.modules.attention.Attention`, :class:`~nvsubquadratic.modules.ckconv_nd.CKConvND`, and - :class:`~nvsubquadratic.modules.mamba_nd.MambaNd`. + :class:`~nvsubquadratic.modules.mamba_nd.Mamba`. qkv_bias: If ``True``, adds a learnable bias to the combined QKV projection. The bias is zero-initialised when ``init_method_in`` is provided. Defaults to ``False``. @@ -137,12 +151,25 @@ def __init__( the bias is zero-initialised. Pass ``None`` to use PyTorch's default (Kaiming uniform). init_method_out: Same as ``init_method_in`` but applied to - ``out_proj.weight.data``. Typically a scaled initialiser (e.g. - ``1 / sqrt(num_layers)``) to control residual branch variance. + ``out_proj.weight.data``. Typically a scaled initialiser that + controls residual-branch variance (GPT/Megatron style), e.g.:: + + import math + init_method_out = ( + lambda dim: lambda w: torch.nn.init.normal_( + w, std=1 / math.sqrt(num_layers) + ) + ) Raises: - RuntimeError: Propagated from ``instantiate(mixer_cfg)`` if the - target class cannot be constructed (e.g. missing required args). + Exception: Propagated from + :func:`~nvsubquadratic.lazy_config.instantiate` if the target + class cannot be constructed (e.g. missing required arguments or + an invalid ``mixer_cfg``). The exact exception type depends on + the ``LazyConfig`` backend (typically an + ``omegaconf.errors.InstantiationException`` or similar). + Check ``mixer_cfg._target_`` and its keyword arguments if this + is raised. """ super().__init__() @@ -165,7 +192,9 @@ def flop_count(self, spatial_dims: tuple[int, ...], inference: bool = False) -> Uses the standard multiply-accumulate convention where one FLOP = one multiply + one add (i.e. the matrix-vector product ``y = Wx`` over - ``T`` tokens costs ``2 · T · in_dim · out_dim`` FLOPs). + ``T`` tokens costs ``2 · T · in_dim · out_dim`` FLOPs). Bias + additions are excluded, following the standard ML FLOP-counting + convention. FLOPs breakdown (``D`` = ``hidden_dim``, ``T`` = ``prod(spatial_dims)``): @@ -176,7 +205,7 @@ def flop_count(self, spatial_dims: tuple[int, ...], inference: bool = False) -> For Hyena this is dominated by the FFT convolution ``O(T log T · D)``; for attention it is ``O(T² · D)``. 3. **Output projection** ``Linear(D, D)``: - ``2 · T · D² `` + ``2 · T · D²`` Total (excluding inner mixer): ``8 · T · D²``. @@ -190,6 +219,10 @@ def flop_count(self, spatial_dims: tuple[int, ...], inference: bool = False) -> Returns: Total FLOPs as a non-negative integer. + + Raises: + AttributeError: If the inner mixer does not implement + ``flop_count``. """ D = self.qkv_proj.in_features T = 1 @@ -206,7 +239,10 @@ def flop_count(self, spatial_dims: tuple[int, ...], inference: bool = False) -> return flops def forward( - self, x: torch.Tensor, cp_group: torch.distributed.ProcessGroup = None, **mixer_kwargs + self, + x: torch.Tensor, + cp_group: torch.distributed.ProcessGroup | None = None, + **mixer_kwargs, ) -> torch.Tensor: """Run the QKV-project → mix → output-project forward pass. @@ -223,13 +259,16 @@ def forward( ring-attention for :class:`~nvsubquadratic.modules.attention.Attention`). Pass ``None`` (default) for single-GPU / non-distributed runs. **mixer_kwargs: Additional keyword arguments forwarded verbatim to - ``self.mixer.forward``. Common keys: + ``self.mixer.forward``. Mixers that do not recognise a key + must accept and ignore it via their own ``**kwargs``. Common + keys: * ``conditioning`` (:class:`torch.Tensor`, shape - ``(B, cond_dim)``): FiLM conditioning signal used by + ``(B, cond_dim)``): FiLM conditioning vector consumed by :class:`~nvsubquadratic.modules.hyena_nd.Hyena` when a - ``condition_mixer`` is attached. - * Any other mixer-specific arguments. + ``condition_mixer`` is attached. Ignored by mixers that do + not have a ``condition_mixer`` (it passes through + ``**mixer_kwargs`` and is discarded). Returns: Output tensor of shape ``(B, *spatial, C)`` — same layout as the From c6a2d1f7aa21e3f2848b64fe969478f7b26f69db Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:08:40 +0200 Subject: [PATCH 06/72] docs(write/mixed_fftconv): add module and function docstrings --- docs/reviews/hyena_nd_review.md | 240 -------------- nvsubquadratic/modules/hyena_nd.py | 97 ++++-- nvsubquadratic/ops/mixed_fftconv.py | 488 +++++++++++++++++++++++----- 3 files changed, 474 insertions(+), 351 deletions(-) delete mode 100644 docs/reviews/hyena_nd_review.md diff --git a/docs/reviews/hyena_nd_review.md b/docs/reviews/hyena_nd_review.md deleted file mode 100644 index 427f3cd7..00000000 --- a/docs/reviews/hyena_nd_review.md +++ /dev/null @@ -1,240 +0,0 @@ -# Reviewer feedback: `nvsubquadratic/modules/hyena_nd.py` - -Reviewed against the Phase-1 docstrings added in commit `docs(write/hyena_nd)`. - -______________________________________________________________________ - -## 1. Module docstring — paper citation is incomplete and imprecise - -**Text:** `"Poli et al., 'Hyena Hierarchy: Towards Larger Convolutional Language Models', ICML 2023."` - -**Issue:** No arXiv ID or DOI. Collaborators without institutional access cannot look the paper up quickly. Also, the paper was published at ICML 2023 but first appeared on arXiv — give both. - -**Fix:** Change to: - -``` -Poli et al., "Hyena Hierarchy: Towards Larger Convolutional Language Models", -ICML 2023. arXiv:2302.10866. -``` - -______________________________________________________________________ - -## 2. Module docstring — "Section 3" reference is too vague - -**Text (class docstring):** `"The two-gate structure corresponds to the H3-style decomposition described in Section 3."` - -**Issue:** The Hyena paper is 30+ pages; "Section 3" does not pin the exact location. H3 is described in a *related-work* paper (Fu et al., 2023), not in the Hyena paper itself. An external reader will spend time searching. - -**Fix:** Replace with: - -``` -The two-gate structure follows the H3 block (Fu et al., "Hungry Hungry Hippos", -ICLR 2023, arXiv:2212.14052, Section 3.2) and is generalised in Hyena -(Poli et al., arXiv:2302.10866, Section 3, "The Hyena Recurrence"). -``` - -______________________________________________________________________ - -## 3. Module docstring — ND generalisation section does not mention boundary conditions - -**Text:** `"For 2D/3D signals the convolution is non-causal by default; causal 1D mode is preserved."` - -**Issue:** A new collaborator will not know whether 2D/3D uses zero-padding (linear conv) or circular conv. This matters because the choice affects the FLOP count and which `fftconv` function is invoked. The `mixed_fftconv` path (#120) also exists now. - -**Fix:** Add after the sentence: - -``` -By default the 2D/3D path uses zero-padded (linear) FFT convolution -(``fftconv2d`` / ``fftconv3d``), matching ``torch.nn.ConvNd(padding='same')`` -semantics. Set the ``circular`` flag on ``CKConvND`` to switch to periodic -boundary conditions, or use ``mixed_fftconv`` for per-axis mixed BCs -(see ``nvsubquadratic.ops.mixed_fftconv``). -``` - -______________________________________________________________________ - -## 4. Class docstring — `Attributes` block: `k_norm` type annotation is contradictory - -**Text:** - -``` -k_norm (torch.nn.Module): Per-channel normalisation for K. - ``Identity`` when the gate is nonlinear (magnitude already bounded - by :math:`\sigma`); a fresh instance of ``qk_norm_cfg`` otherwise. - ``None`` when ``qk_norm_cfg`` is ``None``. -``` - -**Issue:** The type is listed as `torch.nn.Module` but the final sentence says it can be `None`. The actual code sets `self.k_norm = None` when `qk_norm_cfg is None`. The type should be `torch.nn.Module | None`. - -**Fix:** Change to `k_norm (torch.nn.Module | None):` and restate the None case first: - -``` -k_norm (torch.nn.Module | None): Per-channel normalisation for K. - ``None`` when ``qk_norm_cfg`` is ``None`` (QK-norm entirely disabled). - ``torch.nn.Identity`` when the gate is nonlinear (σ already bounds - K's magnitude); a fresh instance of ``qk_norm_cfg`` when the gate is - ``Identity`` (linear gating). -``` - -______________________________________________________________________ - -## 5. `__init__` docstring — `output_norm_cfg` default value is misleading - -**Text:** ``` "Defaults to ``Identity`` (no normalisation)." ``` - -**Issue:** The actual default is `LazyConfig(torch.nn.Identity)()`, not a plain `torch.nn.Identity`. This distinction matters because `LazyConfig(torch.nn.Identity)()` is a frozen lazy config object, not a module; the module is only created by `instantiate` inside `__init__`. An external caller who tries to pass `torch.nn.Identity()` (an already-instantiated module) will get a confusing error from `instantiate`. - -**Fix:** Clarify the type contract: - -``` -output_norm_cfg: ``LazyConfig`` for the normalisation applied after the second - gate. Defaults to a ``LazyConfig`` wrapping ``torch.nn.Identity`` (no - normalisation). Do **not** pass an already-instantiated module — pass a - ``LazyConfig`` object that wraps the class. -``` - -______________________________________________________________________ - -## 6. `flop_count` docstring — FLOP count for QK-Norm assumes RMSNorm always - -**Text:** `"3·C·S for Q; additional 3·C·S for K"` - -**Issue:** The factor of 3 (mean, variance, scale) is correct for RMSNorm / LayerNorm but is silently assumed here. A GroupNorm or a simple scalar multiply would have different counts. The docstring presents this as exact without flagging the assumption. - -**Fix:** Add a note: - -``` -The factor of 3 assumes an RMSNorm-like norm (sum-of-squares + rsqrt + -elementwise scale). Other norm types will differ; this is an approximation. -``` - -______________________________________________________________________ - -## 7. `flop_count` docstring — missing explanation of why `out_ch` appears in the depthwise conv formula - -**Text:** - -``` -2 · (in_ch / groups) · out_ch · S · k_prod -``` - -**Issue:** For a true depthwise conv `in_ch == out_ch == groups`, so `(in_ch / groups) = 1` and the expression collapses to `2 · out_ch · S · k_prod`. The formula as written is actually the general grouped-conv formula, and it happens to be correct for depthwise but the derivation is unclear. An external reader will wonder why `out_ch` appears when they expect `in_ch // groups = 1`. - -**Fix:** Add an inline note: - -``` -For a pure depthwise conv (groups == in_ch == out_ch) this simplifies to -``2 · out_ch · S · k_prod``; the grouped formula is written here to handle -partially-grouped convolutions (e.g. ``DistributedDepthwiseConvNd``). -``` - -______________________________________________________________________ - -## 8. `forward` docstring — CP AllToAll semantics need more precision - -**Text:** - -``` -1. Before short conv: ``split_to_full`` — gather spatial shards, split along the channel dim. -2. After short conv: ``full_to_split`` — scatter spatial, gather channels back. -``` - -**Issue:** The terms `split_to_full` and `full_to_split` are internal string constants from `AllToAllSingleFunction`. A new collaborator does not know which spatial axis is sharded or what "split along channel" means quantitatively. Is it the *first* spatial axis? What is the shard size? - -**Fix:** Add: - -``` -The AllToAll shards along ``dim=2`` (the first spatial axis) and gathers -along ``dim=1`` (the channel axis). After the AllToAll, each device holds -the full spatial extent but only ``C / cp_size`` channels. After the reverse -AllToAll, the original ``C`` channels are restored with ``spatial / cp_size`` -positions per device. -``` - -______________________________________________________________________ - -## 9. `forward` docstring — `query` variable re-use is confusing and undocumented - -**Text:** (no docstring for this; it is inline code) - -In the body, `query` is overwritten twice: - -```python -query = query * self.gate_nonlinear(key) # now z, not Q -... -# then passed to global_conv as z -``` - -The variable is named `query` but after the first gate it represents the *gated intermediate* `z`. This is not mentioned anywhere. - -**Fix:** Add a note in the `forward` docstring under a "Implementation note" or "Variable naming" paragraph: - -``` -Implementation note: ``query`` is overwritten in-place (semantically) after -the first gate to hold the gated intermediate ``z = Q ⊙ σ(K)``. The original -Q tensor is no longer accessible after that point. This is intentional to -avoid an extra allocation. -``` - -______________________________________________________________________ - -## 10. `extra_repr` — `k_norm` can be `None` but is accessed unconditionally - -**Text:** - -```python -k_norm_str = self.k_norm.__class__.__name__ if self.k_norm is not None else "None" -``` - -**Issue:** This is correct code, but the docstring does not mention that `k_norm` may be `None` here (consistent with Issue 4). The docstring for `extra_repr` currently says nothing about when fields are omitted. - -**Fix:** Add to the `extra_repr` docstring: - -``` -When ``self.q_norm`` and ``self.k_norm`` are ``None`` (QK-norm disabled), -the strings ``"q_norm=None"`` and ``"k_norm=None"`` are still included so -the disabled state is explicit in ``repr(module)``. -``` - -______________________________________________________________________ - -## 11. Missing usage example anywhere in the file - -**Issue:** There is no `Example:` block anywhere in the module or class docstring. An external collaborator cannot quickly see how to wire up a `Hyena` block. The `kernels_nd.py` module even has a `# For test, please run:` note, showing that examples are expected. - -**Fix:** Add an `Example:` section to the `Hyena` class docstring showing the minimal construction with `CKConvND`, e.g.: - -```python -Example: - >>> import torch - >>> from nvsubquadratic.lazy_config import LazyConfig - >>> from nvsubquadratic.modules.hyena_nd import Hyena - >>> from nvsubquadratic.modules.ckconv_nd import CKConvND - >>> # Minimal 2D Hyena block (non-causal, no normalisation) - >>> hyena = Hyena( - ... global_conv_cfg=LazyConfig(CKConvND, hidden_dim=64, data_dim=2, ...), - ... short_conv_cfg=LazyConfig(torch.nn.Conv2d, 192, 192, 3, padding=1, groups=192), - ... gate_nonlinear_cfg=LazyConfig(torch.nn.SiLU), - ... pixelhyena_norm_cfg=LazyConfig(torch.nn.Identity), - ... qk_norm_cfg=None, - ... ) - >>> B, H, W, C = 2, 16, 16, 64 - >>> q = k = v = torch.randn(B, H, W, C) - >>> y = hyena(q, k, v) # [2, 16, 16, 64] -``` - -______________________________________________________________________ - -## 12. Module docstring — context parallelism section does not clarify which spatial axis is sharded - -**Text:** `"shard the spatial dimension across devices"` - -**Issue:** For 2D/3D inputs there are multiple spatial axes. Which one? The code shards `dim=2` (the first spatial axis), but this is never stated in the docstring. - -**Fix:** Change to: - -``` -The module shards along ``dim=2`` (the first spatial axis of the -channels-first ``[B, C, *spatial]`` tensor) while gathering the channel dim. -For 2D inputs ``[B, C, H, W]`` this means row-wise sharding (across H). -``` diff --git a/nvsubquadratic/modules/hyena_nd.py b/nvsubquadratic/modules/hyena_nd.py index 8044f171..213dfd3e 100644 --- a/nvsubquadratic/modules/hyena_nd.py +++ b/nvsubquadratic/modules/hyena_nd.py @@ -6,7 +6,8 @@ Background ---------- The Hyena operator (Poli et al., "Hyena Hierarchy: Towards Larger Convolutional -Language Models", ICML 2023) replaces the quadratic attention map with a +Language Models", ICML 2023, arXiv:2302.10866) replaces the quadratic attention +map with a **subquadratic gated convolution**: two multiplicative gates sandwich a long-range (global-kernel) depthwise convolution whose kernel is generated implicitly by a small neural network (see ``kernels_nd.py``). The operator @@ -57,13 +58,20 @@ Fourier Feature MLP (``kernels_nd.py``) and convolves it via ``fftconv{1,2,3}d`` from ``nvsubquadratic.ops.fftconv``. For 2D/3D signals the convolution is non-causal by default; causal 1D mode is preserved. + By default the 2D/3D path uses zero-padded (linear) FFT convolution + (``fftconv2d`` / ``fftconv3d``), matching ``torch.nn.ConvNd(padding='same')`` + semantics. Set the ``circular`` flag on ``CKConvND`` to switch to periodic + boundary conditions, or use ``mixed_fftconv`` for per-axis mixed BCs + (see ``nvsubquadratic.ops.mixed_fftconv``). Context parallelism ------------------- When ``cp_group`` is supplied in ``forward``, the module uses -``AllToAllSingleFunction`` to shard the spatial dimension across devices -while gathering channels, applies the short conv globally, and then shards -back. The global conv receives only the local spatial slice (it must be +``AllToAllSingleFunction`` to shard along ``dim=2`` (the first spatial axis of +the channels-first ``[B, C, *spatial]`` tensor) while gathering the channel dim. +For 2D inputs ``[B, C, H, W]`` this means row-wise sharding (across H). The +short conv is applied globally after the gather, and the result is sharded back. +The global conv receives only the local spatial slice (it must be context-parallel-aware itself). Related modules @@ -113,12 +121,14 @@ class Hyena(torch.nn.Module): :math:`\sigma_2 = \mathrm{Sigmoid}` matches the gated attention formulation used in the original Hyena paper. - Paper reference - --------------- - Poli et al., "Hyena Hierarchy: Towards Larger Convolutional Language - Models", ICML 2023. The two-gate structure corresponds to the H3-style - decomposition described in Section 3. The ND extension replaces the - causal 1D FFT conv with a non-causal ND FFT conv (``CKConvND``). + Paper references + ---------------- + The two-gate structure follows the H3 block (Fu et al., "Hungry Hungry + Hippos", ICLR 2023, arXiv:2212.14052, Section 3.2) and is generalised in + Hyena (Poli et al., "Hyena Hierarchy: Towards Larger Convolutional Language + Models", ICML 2023, arXiv:2302.10866, Section 3 "The Hyena Recurrence"). + The ND extension replaces the causal 1D FFT conv with a non-causal ND FFT + conv (``CKConvND``). Optional components (each disabled by passing ``Identity`` or ``None``): - Short depthwise convolution on concatenated ``[Q, K, V]`` @@ -127,6 +137,27 @@ class Hyena(torch.nn.Module): - Output normalisation after second gate - Context parallelism via AllToAll communication (``cp_group`` argument) + Example:: + + # Minimal 2D Hyena block (non-causal, no normalisation). + # In practice global_conv_cfg wraps a fully-configured CKConvND. + import torch + from nvsubquadratic.lazy_config import LazyConfig + from nvsubquadratic.modules.hyena_nd import Hyena + + hyena = Hyena( + global_conv_cfg=..., # LazyConfig wrapping CKConvND + short_conv_cfg=LazyConfig(torch.nn.Conv2d)( + 192, 192, 3, padding=1, groups=192 + ), + gate_nonlinear_cfg=LazyConfig(torch.nn.SiLU)(), + pixelhyena_norm_cfg=LazyConfig(torch.nn.Identity)(), + qk_norm_cfg=None, + ) + B, H, W, C = 2, 16, 16, 64 + q = k = v = torch.randn(B, H, W, C) + y = hyena(q, k, v) # [2, 16, 16, 64] + Attributes: global_conv (torch.nn.Module): Long-range global convolution, typically ``CKConvND``. Must expose ``hidden_dim`` and @@ -149,10 +180,11 @@ class Hyena(torch.nn.Module): from weight-decay. q_norm (torch.nn.Module | None): Per-channel normalisation for Q. ``None`` when ``qk_norm_cfg`` is ``None``. - k_norm (torch.nn.Module): Per-channel normalisation for K. - ``Identity`` when the gate is nonlinear (magnitude already bounded - by :math:`\sigma`); a fresh instance of ``qk_norm_cfg`` otherwise. - ``None`` when ``qk_norm_cfg`` is ``None``. + k_norm (torch.nn.Module | None): Per-channel normalisation for K. + ``None`` when ``qk_norm_cfg`` is ``None`` (QK-norm entirely + disabled). ``torch.nn.Identity`` when the gate is nonlinear + (:math:`\sigma` already bounds K's magnitude); a fresh instance of + ``qk_norm_cfg`` when the gate is ``Identity`` (linear gating). """ def __init__( @@ -194,8 +226,10 @@ def __init__( Q, one for K) so that stateful norms (e.g. ``RMSNorm`` with a learnable scale) keep independent parameters. output_norm_cfg: ``LazyConfig`` for the normalisation applied after - the second gate. Defaults to ``Identity`` (no normalisation). - Parameters receive ``_no_weight_decay = True``. + the second gate. Defaults to a ``LazyConfig`` wrapping + ``torch.nn.Identity`` (no normalisation). Do **not** pass an + already-instantiated module — pass a ``LazyConfig`` object that + wraps the class. Parameters receive ``_no_weight_decay = True``. gate_nonlinear_2_cfg: ``LazyConfig`` for the second-gate activation :math:`\sigma_2(V)`. If ``None`` (default), both gates share the same activation object (``self.gate_nonlinear``). @@ -260,6 +294,9 @@ def extra_repr(self) -> str: Included fields: - ``q_norm`` / ``k_norm`` class names (or ``"None"``). + When QK-norm is disabled both are ``None``; the strings + ``"q_norm=None"`` and ``"k_norm=None"`` are still emitted so + the disabled state is explicit in ``repr(module)``. - ``gates=<σ>/<σ₂>`` when the two gate activations differ. - ``is_causal`` when the global conv exposes that attribute. @@ -296,11 +333,18 @@ def flop_count(self, spatial_dims: tuple[int, ...], inference: bool = False) -> 2 \cdot \frac{in\_ch}{groups} \cdot out\_ch \cdot S \cdot k\_prod where :math:`k\_prod = \prod_d kernel\_size_d`. Skipped when - ``short_conv`` is ``Identity``. + ``short_conv`` is ``Identity``. For a pure depthwise conv + (``groups == in_ch == out_ch``) this simplifies to + ``2 · out_ch · S · k_prod``; the grouped formula is written here to + handle partially-grouped convolutions (e.g. + ``DistributedDepthwiseConvNd``). 2. **QK-Norm** (when ``self.q_norm is not None``): ``3·C·S`` for Q; additional ``3·C·S`` for K only when ``gate_nonlinear`` is ``Identity`` (linear gating). + The factor of 3 assumes an RMSNorm-like norm (sum-of-squares + + rsqrt + elementwise scale). Other norm types will differ; this is + an approximation. 3. **First gate** :math:`z = Q \odot \sigma(K)`: ``C·S`` for the elementwise multiply, plus ``C·S`` for the @@ -403,13 +447,24 @@ def forward( AllToAll communications around the short conv so that each device sees the full spatial extent during the convolution: - 1. Before short conv: ``split_to_full`` — gather spatial shards, - split along the channel dim. + 1. Before short conv: ``split_to_full`` — gather spatial shards along + ``dim=2`` (the first spatial axis), split along ``dim=1`` (channels). 2. After short conv: ``full_to_split`` — scatter spatial, gather channels back. - The global conv receives only the local spatial slice and is expected - to handle its own CP communication internally. + After step 1, each device holds the full spatial extent but only + ``C / cp_size`` channels. After step 2, the original ``C`` channels + are restored and each device holds ``spatial_0 / cp_size`` positions + along the first spatial axis. The global conv receives only the local + spatial slice and is expected to handle its own CP communication + internally. + + Implementation note + ------------------- + The ``query`` tensor is overwritten after the first gate to hold the + gated intermediate ``z = Q ⊙ σ(K)``; the original Q tensor is no + longer accessible after that point. This is intentional to avoid an + extra allocation. Args: query: ``[B, *spatial, C]`` — query tensor, typically the output diff --git a/nvsubquadratic/ops/mixed_fftconv.py b/nvsubquadratic/ops/mixed_fftconv.py index 8334c1ae..c6e13a29 100644 --- a/nvsubquadratic/ops/mixed_fftconv.py +++ b/nvsubquadratic/ops/mixed_fftconv.py @@ -3,43 +3,70 @@ r"""Mixed boundary-condition FFT-based convolution operators (fp32). -This module implements N-D depthwise FFT convolutions that allow each -spatial axis to independently use either: - -- **Periodic** (circular) boundary conditions → frequency-domain phase - ramp, no padding, no crop on that axis. -- **Non-periodic** (zero-padded "same") boundary conditions → FFT length - is padded up so wrap-around cancels out, centered crop on that axis. - -The choice is made per axis via a ``periodic: tuple[bool, ...]`` argument -of length equal to the number of spatial dimensions. - -Why this op? ------------- -Many PDE datasets (e.g. Well's ``rayleigh_benard``, -``viscoelastic_instability``, ``turbulent_radiative_layer``, -``rayleigh_taylor_instability``) have boundaries that are **periodic on -some axes and non-periodic on others**. A single global ``"zero"`` or -``"circular"`` mode is incorrect for all of these: zero-padding leaks -the wall/open boundary into periodic axes, and circular wraps the -non-periodic ones. This module is the FFT-conv-side fix. - -Relation to existing ops ------------------------- -- All ``periodic = False`` → bit-equivalent to - :mod:`nvsubquadratic.ops.fftconv` (linear / zero-padded "same"). -- All ``periodic = True`` → bit-equivalent to - :mod:`nvsubquadratic.ops.circular_fftconv` (circular / periodic). -- Mixed → new per-axis recipe (see :func:`_mixed_recipe`). - -The op routes the all-False / all-True cases through the existing -non-mixed paths automatically, so adopters pay no overhead for the -legacy modes. +What mixed boundary conditions mean +------------------------------------ +Many PDE datasets have boundaries that are **periodic on some spatial axes +and non-periodic (wall or open) on others**. For example, a Rayleigh-Bénard +simulation is typically periodic in the horizontal (x) direction and bounded +by no-slip walls in the vertical (y) direction. The standard FFT convolution +operators in :mod:`nvsubquadratic.ops.fftconv` (linear / zero-padded) and +:mod:`nvsubquadratic.ops.circular_fftconv` (circular / periodic) each apply a +**global** boundary mode that is wrong for at least one axis in these mixed +cases. + +This module is the fix: it lets each spatial axis independently use either + +- **Periodic** (circular) boundary conditions — no padding on that axis, the + convolution wraps around, and the "same"-alignment is achieved by a + frequency-domain phase ramp ``exp(-i 2π f_d s_d)`` with integer shift + ``s_d = -((K_d - 1) // 2)``. +- **Non-periodic** (zero-padded "same") boundary conditions — the FFT length + is padded so that wrap-around cancels out, and the output is aligned by a + centered crop rather than a phase ramp. + +The choice is expressed per axis via ``periodic: Sequence[bool]`` of length +equal to the number of spatial dimensions: + +.. code-block:: python + + # 2-D: x-axis periodic, y-axis zero-padded + y = mixed_fftconv2d_fp32_bhl(x, kernel, periodic=(True, False)) + + # 3-D: x and y periodic, z zero-padded (e.g. turbulent_radiative_layer_3D) + y = mixed_fftconv3d_fp32_bhl(x, kernel, periodic=(True, True, False)) + +When to use this vs. ``fftconv.py`` or ``circular_fftconv.py`` +-------------------------------------------------------------- +Use this module when **different spatial axes require different boundary +treatments**. For the all-same cases, prefer the dedicated modules: + +- All axes non-periodic → :mod:`nvsubquadratic.ops.fftconv`. +- All axes periodic → :mod:`nvsubquadratic.ops.circular_fftconv`. + +Both degenerate cases are **automatically routed** through the legacy ops +at runtime (zero overhead), so it is safe to use this module as a single +entry point even when ``periodic`` happens to be uniform. + +See ``docs/ops/MIXED_BC_PLAN.md`` for the per-axis algorithm, dataset +motivation table, and deferred work (fp16, multi-head, per-face BCs). + +Algorithm overview +------------------ +The N-D mixed FFT convolution is computed in **one** ``rfftn`` / ``irfftn`` +call. The per-axis variation is encoded entirely in the arguments: + +1. ``s=fft_shape`` — per-axis FFT lengths: ``N_d`` (periodic, no padding) or + ``min(N_d + (K_d+1)//2, 2*N_d)`` (non-periodic, zero-pad headroom). +2. Post-IFFT crop — per-axis slice: ``[0, N_d)`` (periodic) or + ``[K_d//2, K_d//2 + N_d)`` (non-periodic centered crop). +3. Phase ramp — applied to ``fft_k`` before the frequency-domain product: + ``exp(-i 2π f_d s_d)`` on periodic axes; no ramp (``s_d = 0``) on + non-periodic axes (alignment is handled by the crop). Layouts and shapes ------------------ -- **BHL** (channels-first, the fast path): ``[B, H, * spatial_dims]``; - kernel ``[1|B, H, * K_dims]``; output ``[B, H, * spatial_dims]``. +- **BHL** (channels-first, the fast path): ``[B, H, *spatial_dims]``; + kernel ``[1|B, H, *K_dims]``; output ``[B, H, *spatial_dims]``. - **BLH** wrappers (``*_w_reshape``) transparently reshape BLH → BHL → BLH. The leading kernel dimension may be ``1`` (kernel shared across the batch) @@ -53,23 +80,27 @@ .. math:: y \leftarrow y + \text{shortcut} \odot x -Phase ramps ------------ -On each periodic axis we align the output to "same" convolution by an -integer pixel shift of :math:`s_d = -\lfloor (K_d - 1) / 2 \rfloor`, -implemented as a frequency-domain phase ramp -:math:`\exp(-i 2\pi f_d s_d)`. Non-periodic axes use ``s_d = 0`` (the -alignment is absorbed into the centered crop). When -``use_phase_shift=False`` we instead apply ``torch.roll`` along the -periodic axes after the inverse transform. +This is not a generic skip connection — it fuses a specific algebraic +shortcut from Hyena-style gating to avoid a separate kernel launch. + +Phase ramps and the ``use_phase_shift`` flag +-------------------------------------------- +On each periodic axis, "same" alignment requires shifting the output by +``s_d = -((K_d - 1) // 2)`` samples. This can be done two ways: + +- ``use_phase_shift=True`` (default): multiply ``fft_k`` by the complex + ramp ``exp(-i 2π f_d s_d)`` before the IFFT. One fused frequency-domain + op, no data movement after the IFFT. +- ``use_phase_shift=False``: apply :func:`torch.roll` along the periodic + axes *after* the IFFT. Mathematically identical; useful as a reference + or when torch.compile cannot handle complex ops. Caching ------- -Per-axis 1-D phase ramps are cached in a small module-level LRU -(``_MIXED_PHASE_RAMP_1D_CACHE``). The N-D ramp is constructed by -broadcasted multiplication of the relevant 1-D ramps on demand — no -N-D ramp is materialised in the cache. Axes with shift 0 contribute -nothing (they are skipped entirely). +Per-axis 1-D phase ramps are cached in a module-level LRU +(``_MIXED_PHASE_RAMP_1D_CACHE``). The N-D ramp is constructed on demand by +broadcasted multiplication of the relevant 1-D ramps — no N-D tensor is +stored in the cache. Axes with shift 0 contribute nothing (skipped). """ from __future__ import annotations @@ -106,17 +137,42 @@ class _MixedPhaseRamp1DCache: - """LRU cache of 1-D frequency-domain phase ramps used by mixed FFT conv. + r"""LRU cache of 1-D frequency-domain phase ramps used by mixed FFT conv. + + On each **periodic** axis we need to shift the convolution output by + ``s_d = -((K_d - 1) // 2)`` samples to obtain "same"-aligned output. + In the frequency domain this shift corresponds to multiplying ``fft_k`` + by the complex exponential + + .. math:: + R_d[f] = \exp\!\left(-i\, 2\pi \frac{f}{F_d}\, s_d\right) + + where ``f`` ranges over the DFT frequencies for an FFT of length ``F_d`` + and ``s_d`` is the integer pixel shift. Non-periodic axes use ``s_d = 0`` + and contribute no ramp. + + This cache stores the 1-D ramps keyed by ``(F, s, is_rfft_axis, device, + dtype)`` in an ordered-dict LRU so that repeated forward passes with the + same spatial shape reuse the cached tensor without recomputation. - Each cached entry is a complex tensor of shape ``[F]`` (regular FFT axis) - or ``[F // 2 + 1]`` (rfft last axis) that encodes - ``exp(-i 2π f · s)`` for the given size ``F`` and integer shift ``s``. + The N-D ramp is **not** cached here — it is assembled by broadcasting the + relevant 1-D ramps in :func:`_build_nd_phase_ramp` so that each 1-D entry + is shared across all spatial configurations that share an axis. - The N-D ramp is built lazily by broadcasting the relevant 1-D ramps; - no N-D tensor is stored in the cache. + Attributes: + maxsize: Maximum number of entries before the LRU evicts the oldest. + _cache: Ordered dict mapping cache keys to complex ramp tensors. """ def __init__(self, maxsize: int = 256): + """Initialise the cache. + + Args: + maxsize: Maximum number of (F, s, axis, device, dtype) entries to + keep. Once exceeded, the least-recently-used entry is evicted. + Default 256 covers many typical spatial / kernel size combos + without meaningful memory cost (each 1-D ramp is small). + """ self.maxsize = maxsize self._cache: "OrderedDict[tuple, torch.Tensor]" = OrderedDict() @@ -128,6 +184,23 @@ def _key( device: torch.device, real_dtype: torch.dtype, ) -> tuple: + """Build a hashable cache key for a 1-D phase ramp. + + The key encodes every parameter that affects the ramp tensor so that + ramps from different devices or dtypes are never confused. + + Args: + F: FFT length along this axis. + s: Integer pixel shift encoded by the ramp. + is_rfft_axis: True when this is the last axis (rfft frequencies, + length ``F // 2 + 1``); False for intermediate axes (full DFT + frequencies, length ``F``). + device: Target device. + real_dtype: Real-valued input dtype (``float32`` or ``float64``). + + Returns: + A hashable tuple suitable as a dict key. + """ complex_dtype = torch.complex64 if real_dtype == torch.float32 else torch.complex128 dev_type = device.type dev_idx = device.index if device.index is not None else -1 @@ -141,20 +214,36 @@ def get( device: torch.device, real_dtype: torch.dtype, ) -> torch.Tensor: - """Return the 1-D phase ramp for the given (FFT length, shift) on this axis. + r"""Return the 1-D phase ramp for the given (FFT length, shift) on this axis. + + Computes (or retrieves from the LRU cache) the complex tensor + + .. math:: + R[f] = \cos(-2\pi f s) + i\,\sin(-2\pi f s) + + where ``f`` iterates over ``torch.fft.rfftfreq(F)`` (if + ``is_rfft_axis``) or ``torch.fft.fftfreq(F)`` (otherwise). The ramp + is computed under ``torch.no_grad()`` and ``torch.inference_mode`` + so it does not pollute the autograd graph. Args: F: FFT length along this axis (padded length for non-periodic, input length for periodic — caller decides which). - s: Integer pixel shift to apply via the ramp. ``s == 0`` callers - are expected to skip the multiply entirely; this method - still supports ``s == 0`` (returns all-ones for completeness). + s: Integer pixel shift to apply via the ramp. Callers with + ``s == 0`` are expected to skip the multiply entirely; this + method still handles ``s == 0`` (returns all-ones). is_rfft_axis: If True, this is the last spatial axis where the - rfft is taken; we build a ramp of length ``F // 2 + 1`` using - :func:`torch.fft.rfftfreq`. Otherwise we build a length-``F`` - ramp using :func:`torch.fft.fftfreq`. - device: Target device. - real_dtype: Real-valued dtype of the inputs (float32 or float64). + rfft is taken; a ramp of length ``F // 2 + 1`` is built + using :func:`torch.fft.rfftfreq`. Otherwise a full ramp of + length ``F`` is built using :func:`torch.fft.fftfreq`. + device: Target device for the ramp tensor. + real_dtype: Real-valued dtype of the inputs (``float32`` or + ``float64``); the ramp is stored in the corresponding complex + dtype (``complex64`` or ``complex128``). + + Returns: + 1-D complex tensor of shape ``[F // 2 + 1]`` (rfft axis) or + ``[F]`` (non-rfft axis) on ``device``. """ key = self._key(F, s, is_rfft_axis, device, real_dtype) cached = self._cache.get(key) @@ -237,19 +326,44 @@ def _build_nd_phase_ramp( device: torch.device, real_dtype: torch.dtype, ) -> torch.Tensor | None: - """Build (or fetch from cache) the broadcast N-D phase-ramp tensor. + r"""Build (or fetch from cache) the broadcast N-D phase-ramp tensor. + + Assembles the N-D phase ramp - The returned tensor has rfft-style shape - ``(F_0, F_1, ..., F_{D-2}, F_{D-1} // 2 + 1)`` with complex dtype, - suitable to be multiplied with ``torch.fft.rfftn(..., s=fft_shape)``. + .. math:: + R[f_0, \ldots, f_{D-1}] = + \prod_{d:\, s_d \ne 0} + \exp\!\left(-i\, 2\pi \frac{f_d}{F_d}\, s_d\right) - Axes with ``shift == 0`` contribute a length-1 broadcast factor (no - multiply along that axis). If every shift is zero, this function - returns ``None`` — callers should skip the multiply entirely. + by broadcasting individual 1-D ramps fetched from + ``_MIXED_PHASE_RAMP_1D_CACHE``. The result has rfft-style shape + ``(F_0, F_1, ..., F_{D-2}, F_{D-1} // 2 + 1)`` and complex dtype, + so it can be multiplied directly with + ``torch.fft.rfftn(x, s=fft_shape)``. - The N-D ramp itself is not cached; only the 1-D per-axis ramps are. - The product of 1-D ramps over broadcasted dims is materialised here - so that the downstream multiply with ``fft_x`` is a single op. + Axes with ``shift == 0`` (non-periodic axes, and periodic axes where the + kernel has size 1) contribute no ramp — they are skipped entirely rather + than adding a length-1 all-ones factor. If **all** shifts are zero the + function returns ``None`` so callers can skip the multiply with a single + branch. + + The N-D ramp is **not** cached; only the 1-D per-axis ramps are. The + product is materialised here (via broadcasted multiplication) so that the + downstream multiply with ``fft_x`` is a single fused op. + + Args: + fft_shape: Per-axis FFT lengths ``(F_0, ..., F_{D-1})``, as returned + by :func:`_mixed_recipe`. + shifts: Per-axis integer pixel shifts ``(s_0, ..., s_{D-1})``. + Negative values shift left (the typical case for periodic axes). + device: Device on which to materialise the ramp. + real_dtype: Real dtype of the input (determines the complex dtype of + the ramp: ``float32`` → ``complex64``, ``float64`` → ``complex128``). + + Returns: + Broadcast-ready complex tensor of shape + ``(F_0, ..., F_{D-2}, F_{D-1} // 2 + 1)``, or ``None`` if every + shift is zero. """ if all(s == 0 for s in shifts): return None @@ -281,7 +395,23 @@ def _build_nd_phase_ramp( def _normalize_periodic(periodic: Sequence[bool] | tuple[bool, ...], data_dim: int) -> tuple[bool, ...]: - """Normalise the ``periodic`` argument to a length-``data_dim`` tuple of bools.""" + """Normalise the ``periodic`` argument to a length-``data_dim`` tuple of bools. + + Accepts any sequence (list, tuple, generator) of truthy/falsy values and + returns a typed tuple of Python ``bool`` values, validating that the + length matches ``data_dim``. + + Args: + periodic: Per-axis periodicity flags. Length must equal ``data_dim``. + Any truthy value is treated as ``True`` (periodic). + data_dim: Expected number of spatial dimensions (1, 2, or 3). + + Returns: + A ``tuple[bool, ...]`` of length ``data_dim``. + + Raises: + AssertionError: If ``len(periodic) != data_dim``. + """ periodic_t = tuple(bool(p) for p in periodic) assert len(periodic_t) == data_dim, ( f"periodic must have length {data_dim} (data_dim), got length {len(periodic_t)}" @@ -299,11 +429,29 @@ def _dispatch_legacy_if_uniform( ) -> torch.Tensor | None: """If ``periodic`` is uniformly all-True or all-False, route to the existing op. - Returns the output tensor if the call was dispatched, or ``None`` to - indicate the caller should run the mixed path itself. + This ensures the all-False and all-True degenerate cases produce + bit-identical results to :mod:`nvsubquadratic.ops.fftconv` and + :mod:`nvsubquadratic.ops.circular_fftconv` respectively, and incur no + overhead from the mixed-axis logic. + + - ``all(periodic) == True`` → dispatches to the corresponding + ``circular_fftconv{1,2,3}d_fp32_bhl`` function, forwarding + ``use_phase_shift``. + - ``not any(periodic)`` (all False) → dispatches to the corresponding + ``fftconv{1,2,3}d_fp32_bhl`` function. The ``use_phase_shift`` flag + is not forwarded (it does not apply to zero-padded linear conv). + + Args: + periodic: Normalised per-axis periodicity tuple (length ``data_dim``). + data_dim: Number of spatial dimensions (1, 2, or 3). + x: Input tensor (BHL layout). + kernel: Kernel tensor (BHL layout). + shortcut: Optional per-channel residual scale ``[H]``. + use_phase_shift: Forwarded to the circular op when dispatching. - The legacy linear ops do not take ``use_phase_shift`` (it does not - apply there); we always dispatch when periodic is all-False. + Returns: + Output tensor if the call was dispatched to a legacy op, or ``None`` + if ``periodic`` is mixed and the caller must run the mixed path. """ if all(periodic): from nvsubquadratic.ops import circular_fftconv as _circ @@ -341,8 +489,42 @@ def _mixed_fftconv_nd_fp32_bhl( """Shared N-D mixed-BC FFT conv body (BHL, fp32). The 1D/2D/3D public entry points delegate here after shape validation. - All padding/cropping/phase-ramp logic is per axis according to - :func:`_mixed_recipe`. + Implements the following steps: + + 1. Validate shapes (ndim, batch dim, hidden dim, per-axis kernel-size + limits). + 2. Dispatch to the legacy linear / circular op if ``periodic`` is uniform + (via :func:`_dispatch_legacy_if_uniform`). + 3. Cast ``x`` and ``kernel`` to fp32, compute the per-axis FFT parameters + via :func:`_mixed_recipe`, and take ``rfftn`` over all spatial dims in + one call. + 4. Optionally apply the N-D phase ramp to ``fft_k`` (if + ``use_phase_shift=True``). + 5. Multiply ``fft_x * fft_k`` and apply the inverse ``irfftn``. + 6. If ``use_phase_shift=False``, apply ``torch.roll`` on the periodic + axes to achieve "same" alignment. + 7. Crop the IFFT output to the input spatial shape via the per-axis + crop windows from :func:`_mixed_recipe`. + 8. Cast back to the original dtype of ``x`` and add the optional + shortcut term ``y += shortcut * x``. + + Args: + x: Input tensor of shape ``[B, H, *spatial_dims]`` (any dtype). + kernel: Kernel tensor of shape ``[1|B, H, *K_dims]`` (any dtype). + periodic: Normalised per-axis periodicity tuple, length ``data_dim``. + shortcut: Optional per-channel residual scale ``[H]``. Must have the + same dtype as ``x``. + use_phase_shift: If True, apply the "same" shift in the frequency + domain (faster). If False, use ``torch.roll`` after the IFFT + (useful as a reference implementation). + data_dim: Number of spatial dimensions (1, 2, or 3). + + Returns: + Output tensor of shape ``[B, H, *spatial_dims]`` in the original + dtype of ``x``. + + Raises: + AssertionError: On shape mismatches or out-of-range kernel sizes. """ x_shape = x.shape assert x.ndim == 2 + data_dim, f"Expected {2 + data_dim}D input, got {x.ndim}D" @@ -511,8 +693,20 @@ def mixed_fftconv1d_fp32_bhl_w_reshape( ) -> torch.Tensor: """1D mixed-BC FFT conv wrapper for BLH layout (batch, length, hidden). - Reshapes BLH ↔ BHL around :func:`mixed_fftconv1d_fp32_bhl`. See that - function for argument semantics. + Reshapes BLH ↔ BHL around :func:`mixed_fftconv1d_fp32_bhl`. Prefer this + wrapper over operating in BLH natively — internally the FFT runs on + contiguous spatial axes (BHL), so the reshape cost is negligible compared + to the FFT itself. + + Args: + x: Input tensor of shape ``[B, L, H]`` (BLH, channels-last). + kernel: Kernel tensor of shape ``[B, K, H]`` (BLH). + periodic: Length-1 sequence of bools. + shortcut: Optional per-channel scale ``[H]``. + use_phase_shift: See :func:`mixed_fftconv1d_fp32_bhl`. + + Returns: + Tensor of shape ``[B, L, H]`` in the original dtype of ``x``. """ x_bhl = rearrange(x, "b l h -> b h l") k_bhl = rearrange(kernel, "b k h -> b h k") @@ -529,7 +723,19 @@ def mixed_fftconv2d_fp32_bhl_w_reshape( ) -> torch.Tensor: """2D mixed-BC FFT conv wrapper for BLH layout (batch, X, Y, hidden). - Reshapes BLH ↔ BHL around :func:`mixed_fftconv2d_fp32_bhl`. + Reshapes BLH ↔ BHL around :func:`mixed_fftconv2d_fp32_bhl`. Prefer this + over operating in BLH natively — the FFT runs faster on contiguous spatial + axes (BHL). + + Args: + x: Input tensor of shape ``[B, X, Y, H]`` (BLH, channels-last). + kernel: Kernel tensor of shape ``[B, K_x, K_y, H]`` (BLH). + periodic: Length-2 sequence ``(periodic_x, periodic_y)``. + shortcut: Optional per-channel scale ``[H]``. + use_phase_shift: See :func:`mixed_fftconv1d_fp32_bhl`. + + Returns: + Tensor of shape ``[B, X, Y, H]`` in the original dtype of ``x``. """ x_bhl = rearrange(x, "b x y h -> b h x y") k_bhl = rearrange(kernel, "b kx ky h -> b h kx ky") @@ -546,7 +752,18 @@ def mixed_fftconv3d_fp32_bhl_w_reshape( ) -> torch.Tensor: """3D mixed-BC FFT conv wrapper for BLH layout (batch, X, Y, Z, hidden). - Reshapes BLH ↔ BHL around :func:`mixed_fftconv3d_fp32_bhl`. + Reshapes BLH ↔ BHL around :func:`mixed_fftconv3d_fp32_bhl`. Prefer this + over operating in BLH natively. + + Args: + x: Input tensor of shape ``[B, X, Y, Z, H]`` (BLH, channels-last). + kernel: Kernel tensor of shape ``[B, K_x, K_y, K_z, H]`` (BLH). + periodic: Length-3 sequence ``(periodic_x, periodic_y, periodic_z)``. + shortcut: Optional per-channel scale ``[H]``. + use_phase_shift: See :func:`mixed_fftconv1d_fp32_bhl`. + + Returns: + Tensor of shape ``[B, X, Y, Z, H]`` in the original dtype of ``x``. """ x_bhl = rearrange(x, "b x y z h -> b h x y z") k_bhl = rearrange(kernel, "b kx ky kz h -> b h kx ky kz") @@ -569,7 +786,27 @@ def _chunked_along_channels( chunk_size: int, core_fn, ) -> torch.Tensor: - """Apply ``core_fn`` to ``[x, kernel, shortcut]`` slices of size ``chunk_size`` along H.""" + """Apply ``core_fn`` to ``[x, kernel, shortcut]`` slices of size ``chunk_size`` along H. + + Reduces peak GPU memory by avoiding materialising the full FFT-intermediate + tensors (of size ``[B, H, *fft_shape]``) for all channels at once. When + ``H <= chunk_size`` the call is a pass-through with no slicing overhead. + + Args: + x: Input tensor of shape ``[B, H, *spatial_dims]`` (BHL layout). + kernel: Kernel tensor of shape ``[1|B, H, *K_dims]``. The leading dim + is expected to be either 1 (shared kernel) or ``B``; slicing along + dim 1 is applied uniformly regardless. + shortcut: Optional per-channel residual scale ``[H]``, sliced to + match each chunk. ``None`` is forwarded unchanged. + chunk_size: Maximum number of channels to process per chunk. + core_fn: Callable with signature ``(x_chunk, k_chunk, sc_chunk) -> + Tensor`` that performs the actual convolution on a channel slice. + + Returns: + Output tensor of shape ``[B, H, *spatial_dims]`` reassembled from + chunk outputs concatenated along the channel dimension (dim 1). + """ H = x.shape[1] if H <= chunk_size: return core_fn(x, kernel, shortcut) @@ -596,9 +833,21 @@ def mixed_fftconv1d_fp32_bhl_chunked( ) -> torch.Tensor: """Memory-efficient 1D mixed-BC FFT conv (BHL) via channel chunking. - See :func:`mixed_fftconv1d_fp32_bhl` for the core semantics. The work is - identical, but it is performed on at most ``chunk_size`` channels at a - time, lowering peak FFT-intermediate memory. + Produces the same output as :func:`mixed_fftconv1d_fp32_bhl` but processes + at most ``chunk_size`` channels at a time, lowering peak FFT-intermediate + memory at the cost of multiple kernel launches. + + Args: + x: Input tensor of shape ``[B, H, L]`` (any dtype, cast to fp32). + kernel: Kernel tensor of shape ``[1|B, H, K]``. + periodic: Length-1 sequence of bools. ``periodic[0] == True`` → circular. + shortcut: Optional per-channel scale ``[H]`` added as ``y += shortcut * x``. + use_phase_shift: See :func:`mixed_fftconv1d_fp32_bhl`. + chunk_size: Number of channels per chunk. Defaults to + ``_DEFAULT_MIXED_CHUNK_SIZE`` (128). Pass ``None`` to use the default. + + Returns: + Tensor of shape ``[B, H, L]`` in the original dtype of ``x``. """ chunk = chunk_size if chunk_size is not None else _DEFAULT_MIXED_CHUNK_SIZE periodic_t = _normalize_periodic(periodic, data_dim=1) @@ -621,7 +870,19 @@ def mixed_fftconv2d_fp32_bhl_chunked( ) -> torch.Tensor: """Memory-efficient 2D mixed-BC FFT conv (BHL) via channel chunking. - See :func:`mixed_fftconv2d_fp32_bhl` for semantics. + Produces the same output as :func:`mixed_fftconv2d_fp32_bhl` but processes + at most ``chunk_size`` channels at a time to limit peak memory. + + Args: + x: Input tensor of shape ``[B, H, X, Y]`` (any dtype, cast to fp32). + kernel: Kernel tensor of shape ``[1|B, H, K_x, K_y]``. + periodic: Length-2 sequence ``(periodic_x, periodic_y)``. + shortcut: Optional per-channel scale ``[H]``. + use_phase_shift: See :func:`mixed_fftconv1d_fp32_bhl`. + chunk_size: Channels per chunk. Defaults to 128. + + Returns: + Tensor of shape ``[B, H, X, Y]`` in the original dtype of ``x``. """ chunk = chunk_size if chunk_size is not None else _DEFAULT_MIXED_CHUNK_SIZE periodic_t = _normalize_periodic(periodic, data_dim=2) @@ -644,7 +905,19 @@ def mixed_fftconv3d_fp32_bhl_chunked( ) -> torch.Tensor: """Memory-efficient 3D mixed-BC FFT conv (BHL) via channel chunking. - See :func:`mixed_fftconv3d_fp32_bhl` for semantics. + Produces the same output as :func:`mixed_fftconv3d_fp32_bhl` but processes + at most ``chunk_size`` channels at a time to limit peak memory. + + Args: + x: Input tensor of shape ``[B, H, X, Y, Z]`` (any dtype, cast to fp32). + kernel: Kernel tensor of shape ``[1|B, H, K_x, K_y, K_z]``. + periodic: Length-3 sequence ``(periodic_x, periodic_y, periodic_z)``. + shortcut: Optional per-channel scale ``[H]``. + use_phase_shift: See :func:`mixed_fftconv1d_fp32_bhl`. + chunk_size: Channels per chunk. Defaults to 128. + + Returns: + Tensor of shape ``[B, H, X, Y, Z]`` in the original dtype of ``x``. """ chunk = chunk_size if chunk_size is not None else _DEFAULT_MIXED_CHUNK_SIZE periodic_t = _normalize_periodic(periodic, data_dim=3) @@ -673,6 +946,19 @@ def mixed_fftconv1d_fp32_bhl_w_reshape_chunked( """Chunked 1D mixed-BC FFT conv wrapper for BLH layout. Reshapes BLH ↔ BHL around :func:`mixed_fftconv1d_fp32_bhl_chunked`. + Combines channel chunking (for memory savings) with the BLH → BHL reshape + (for FFT efficiency). + + Args: + x: Input tensor of shape ``[B, L, H]`` (BLH, channels-last). + kernel: Kernel tensor of shape ``[B, K, H]`` (BLH). + periodic: Length-1 sequence of bools. + shortcut: Optional per-channel scale ``[H]``. + use_phase_shift: See :func:`mixed_fftconv1d_fp32_bhl`. + chunk_size: Channels per chunk. Defaults to 128. + + Returns: + Tensor of shape ``[B, L, H]`` in the original dtype of ``x``. """ x_bhl = rearrange(x, "b l h -> b h l") k_bhl = rearrange(kernel, "b k h -> b h k") @@ -698,6 +984,17 @@ def mixed_fftconv2d_fp32_bhl_w_reshape_chunked( """Chunked 2D mixed-BC FFT conv wrapper for BLH layout. Reshapes BLH ↔ BHL around :func:`mixed_fftconv2d_fp32_bhl_chunked`. + + Args: + x: Input tensor of shape ``[B, X, Y, H]`` (BLH, channels-last). + kernel: Kernel tensor of shape ``[B, K_x, K_y, H]`` (BLH). + periodic: Length-2 sequence ``(periodic_x, periodic_y)``. + shortcut: Optional per-channel scale ``[H]``. + use_phase_shift: See :func:`mixed_fftconv1d_fp32_bhl`. + chunk_size: Channels per chunk. Defaults to 128. + + Returns: + Tensor of shape ``[B, X, Y, H]`` in the original dtype of ``x``. """ x_bhl = rearrange(x, "b x y h -> b h x y") k_bhl = rearrange(kernel, "b kx ky h -> b h kx ky") @@ -723,6 +1020,17 @@ def mixed_fftconv3d_fp32_bhl_w_reshape_chunked( """Chunked 3D mixed-BC FFT conv wrapper for BLH layout. Reshapes BLH ↔ BHL around :func:`mixed_fftconv3d_fp32_bhl_chunked`. + + Args: + x: Input tensor of shape ``[B, X, Y, Z, H]`` (BLH, channels-last). + kernel: Kernel tensor of shape ``[B, K_x, K_y, K_z, H]`` (BLH). + periodic: Length-3 sequence ``(periodic_x, periodic_y, periodic_z)``. + shortcut: Optional per-channel scale ``[H]``. + use_phase_shift: See :func:`mixed_fftconv1d_fp32_bhl`. + chunk_size: Channels per chunk. Defaults to 128. + + Returns: + Tensor of shape ``[B, X, Y, Z, H]`` in the original dtype of ``x``. """ x_bhl = rearrange(x, "b x y z h -> b h x y z") k_bhl = rearrange(kernel, "b kx ky kz h -> b h kx ky kz") From 68650e72de0357c6bd2deeea412bbbc14427329b Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:09:33 +0200 Subject: [PATCH 07/72] docs(write/kernels_nd): add module and class docstrings with math context --- nvsubquadratic/modules/kernels_nd.py | 523 ++++++++++++++++++++++++--- 1 file changed, 482 insertions(+), 41 deletions(-) diff --git a/nvsubquadratic/modules/kernels_nd.py b/nvsubquadratic/modules/kernels_nd.py index f95a69fb..86315465 100644 --- a/nvsubquadratic/modules/kernels_nd.py +++ b/nvsubquadratic/modules/kernels_nd.py @@ -1,7 +1,78 @@ # TODO: Add license header here -"""Implicit Kernel Implementations for ND signals (based on Random Fourier Feature Networks). +"""Implicit / learned kernel parametrisations for N-dimensional convolutional filters. + +Overview +-------- +Standard convolutional sequence models fix the filter bank at construction time +(e.g. a Gabor filter or a learnable lookup table). The classes in this module +instead **parametrise the kernel implicitly**: a small MLP maps spatial +coordinates to kernel values, so the filter shape is determined by the MLP's +learned weights rather than by an explicit table of size proportional to the +signal length. This approach — sometimes called an *implicit neural +representation* (INR) of the kernel — has two key advantages: + +1. **Resolution-independence**: the same parameter count describes kernels of + any length. A model trained at resolution N can be evaluated at a + different resolution without any weight surgery. + +2. **Continuous inductive bias**: the MLP is smooth almost everywhere, which + acts as a spectral regulariser and avoids the aliasing artefacts that arise + when a fixed filter is up-sampled or sub-sampled. + +The main consumer of this module is ``nvsubquadratic.modules.hyena_nd.Hyena`` +(via ``nvsubquadratic.modules.ckconv_nd.CKConvND``), where the generated kernel +is passed directly to the FFT convolution primitives in ``nvsubquadratic.ops``. + +Two positional-encoding families are provided, each yielding a matching kernel +class: + +* **Random Fourier Features (RFF)**: the first layer is a random (fixed) + frequency matrix; activations are cosine+sine concatenated. The result is + that the MLP's effective prior is a stationary (shift-invariant) kernel + corresponding to the RBF kernel. Controlled by ``omega_0`` (bandwidth). + +* **SIREN** (sinusoidal representation network, Sitzmann et al. 2020): every + layer uses ``sin`` activation. The frequency of the first layer is set by + ``omega_0``; subsequent layers use ``hidden_omega_0``. Produces smoother + high-frequency content than RFF and is amenable to multi-frequency + initialisations (see ``MultiOmegaSIRENKernelND`` and the block-diagonal + variants below). + +ND here refers to the spatial dimensionality of the signal: 1D (sequences), +2D (images), 3D (video / volumetric data), or higher. The coordinate grid +covers ``[-1, 1]^D`` normalised across all axes; the ``L_cache`` parameter +controls the number of discrete grid positions cached per axis, and the cache +grows automatically at runtime whenever a larger input is encountered. + +Kernel classes in this module +------------------------------ +``RandomFourierKernelND`` + RFF-based kernel. Recommended when a well-understood stationary prior is + desired. The ``omega_0`` parameter directly controls the bandwidth of the + implied RBF kernel. + +``SIRENKernelND`` + SIREN-based kernel. Supports optional FiLM conditioning (input-dependent + kernels). The recommended default for Hyena-ND models. + +``MultiOmegaSIRENKernelND`` + SIREN kernel with a per-row ``omega_0`` in the first layer, allowing the + model to represent multiple frequency bands simultaneously. + +``BlockDiagonalMultiOmegaSIRENKernelND`` + Multi-omega SIREN with block-diagonal weight masking at init, so each + frequency block starts as an independent narrow-band SIREN. + +``LearnableOmegaSIRENKernelND`` + SIREN kernel whose per-row ``omega_0`` multiplier is a learnable parameter + (clamped to a configurable range), enabling frequency adaptation during + training. + +``BlockDiagonalLearnableOmegaSIRENKernelND`` + Combines block-diagonal MLP init with learnable per-row omega scaling — + the most expressive variant in the family. For test, please run: PYTHONPATH=. python nvsubquadratic/modules/kernels_nd.py @@ -55,11 +126,60 @@ def _normalize_l_cache(L_cache: int | Sequence[int], data_dim: int) -> tuple[int class RandomFourierPositionalEmbeddingND(torch.nn.Module): - """Implements a N-dimensional positional embedding using Random Fourier Features. + """N-dimensional positional embedding using Random Fourier Features (RFF). - This module generates positional embeddings by applying a linear transformation - with randomized Fourier frequencies followed by sine and cosine functions. - It is suitable for tasks where positional information needs to be encoded. + Mathematical form + ----------------- + Given a coordinate grid ``x`` of shape ``[1, *spatial_dims, data_dim]`` with + values normalised to ``[-1, 1]`` per axis, the embedding is: + + phi(x) = [ cos(W x + b), sin(W x + b) ] shape [..., embedding_dim] + + where: + + * ``W`` is the first-layer weight matrix of shape + ``[embedding_dim//2, data_dim]``, drawn once at construction from + ``N(0, (2*pi*omega_0)^2)`` and then **frozen** (not trained). + * ``b`` is a bias vector of shape ``[embedding_dim//2]``, initialised to + zero. It is also frozen. + * The concatenation of cosine and sine doubles the embedding dimension. + + The resulting features approximate the feature map of a stationary RBF + (Gaussian) kernel with bandwidth ``omega_0`` — the larger ``omega_0``, the + higher the dominant spatial frequency encoded in the embedding. + + Grid caching + ------------ + To avoid rebuilding the meshgrid on every forward pass, the module + maintains a ``grid_cache`` buffer (a pre-computed coordinate tensor of + shape ``[1, 2*L_0-1, ..., 2*L_{d-1}-1, data_dim]`` in float32). On each + forward call the central ``[2*seq_len_i - 1]`` points are sliced per axis. + When a larger ``seq_len`` is seen at runtime the cache grows automatically + via ``_maybe_extend_grid_cache``, preserving the original step size on + each axis. + + Note: the ``W`` and ``b`` parameters have ``_no_weight_decay = True`` set + so that any weight-decay optimizer does not shrink the random projection. + + Attributes: + data_dim (int): Number of spatial / temporal input dimensions. + embedding_dim (int): Output embedding size (must be even; split equally + between cos and sin features). + L_cache_per_axis (tuple[int, ...]): Current per-axis cache extents. + May grow at runtime; the original value at construction is stored in + ``self.L_cache``. + L_cache (int | Sequence[int]): Original ``L_cache`` argument (for + diagnostics and external read-back). + omega_0 (float): Bandwidth / frequency scaling factor used for weight + init and for diagnostics. + use_bias (bool): Whether a bias is present in the linear projection. + linear (torch.nn.Linear): The frozen random frequency projection + ``W`` (and optionally ``b``), shape ``[embedding_dim//2, data_dim]``. + grid_cache (torch.Tensor): Non-persistent float32 buffer of shape + ``[1, 2*L_0-1, ..., 2*L_{d-1}-1, data_dim]``. + step_sizes (tuple[float, ...]): Per-axis grid step + ``1/(L_i - 1)`` at construction; used by cache extensions to + keep spacing constant. """ def __init__( @@ -223,10 +343,58 @@ def forward(self, seq_lens: tuple[int, ...]) -> tuple[torch.Tensor, torch.Tensor class RandomFourierKernelND(torch.nn.Module): - """Implements a learnable ND-dimensional freeform convolutional kernel using implicit neural representations. + """Learned convolutional kernel parametrised via Random Fourier Features and an MLP. + + Mathematical form + ----------------- + The kernel at grid coordinate ``x`` is: + + k(x) = Linear_out( MLP( phi(x) ) ) + + where: - This module combines positional embeddings, a feedforward neural network, and optional modulation - and normalization to compute freeform filters over a grid of spatial dimensions. + * ``phi(x) = [cos(W x + b), sin(W x + b)]`` is the RFF positional + embedding (see ``RandomFourierPositionalEmbeddingND``). + * ``MLP`` is a stack of ``num_layers - 1`` fully-connected layers each + followed by the ``nonlinear_cfg`` activation. + * ``Linear_out`` is a final linear layer that maps to ``out_dim`` channels. + + The output is a kernel tensor of shape ``[1, *spatial_dims, out_dim]`` + suitable for passing directly to the FFT convolution primitives in + ``nvsubquadratic.ops`` (after rearranging to channels-first layout via the + consuming ``CKConvND`` module). + + Hyperparameters controlling bandwidth / smoothness + --------------------------------------------------- + * ``omega_0``: Controls the frequency content of the RFF features. Higher + values concentrate the kernel's spectral energy at higher frequencies, + producing a narrower, higher-bandwidth filter. Typical range: 1.0–100.0. + * ``mlp_hidden_dim``: Width of the hidden MLP layers. Larger values allow + more expressive kernel shapes. + * ``num_layers``: Depth of the MLP. Must be >= 2 (one hidden layer + + output layer minimum). + + Initialisation + -------------- + * Hidden MLP layers are initialised with the user-supplied ``init_method`` + if provided; otherwise use PyTorch defaults. + * The output layer applies **Wang initialisation**: weights are scaled by + ``sqrt(1 / kernel_volume)`` where ``kernel_volume = prod(L_cache_per_axis)`` + (collapses to ``L_cache**data_dim`` for an isotropic grid). This + normalises the kernel's initial energy to be independent of grid size. + + Attributes: + out_dim (int): Number of output channels (kernel depth). + data_dim (int): Number of spatial / temporal input dimensions. + mlp_hidden_dim (int): Hidden width of the MLP. + num_layers (int): Total number of MLP layers (>= 2). + embedding_dim (int): RFF embedding dimensionality (must be even). + omega_0 (float): Bandwidth scaling factor for the positional embedding. + L_cache_per_axis (tuple[int, ...]): Per-axis cache extents (canonical form). + L_cache (int | Sequence[int]): Original ``L_cache`` argument (diagnostics). + positional_embedding (RandomFourierPositionalEmbeddingND): RFF encoder. + kernel_network (torch.nn.Sequential): Hidden MLP layers. + out_linear (torch.nn.Linear): Final projection to ``out_dim`` channels. """ def __init__( @@ -329,6 +497,31 @@ def forward(self, seq_lens: tuple[int, ...], conditioning: torch.Tensor | None = def _init_siren_weights(layer: torch.nn.Linear, is_first_layer: bool, w0: float) -> None: + """Initialise a ``nn.Linear`` layer with the SIREN uniform distribution. + + From Sitzmann et al. 2020 ("Implicit Neural Representations with Periodic + Activation Functions"), the weight init that keeps the distribution of + pre-activations stationary across layers of ``sin`` activations is: + + first layer: W ~ U(-1/d, +1/d) (note: scaled by 2pi*w0 below) + hidden layers: W ~ U(-sqrt(6/d)/(2pi*w0), +sqrt(6/d)/(2pi*w0)) + + where ``d = layer.in_features``. The factor ``2*pi*w0`` is absorbed into + the init bound so that the frequency content of each layer matches ``w0`` + at initialisation without applying the factor at every forward pass. + + Bias is always zero-initialised, matching the SIREN paper's convention. + + Args: + layer: The ``nn.Linear`` layer to initialise (modified in-place). + is_first_layer: If True, use the first-layer bound ``1/d``; otherwise + use the hidden-layer bound ``sqrt(6/d)/(2pi*w0)``. + w0: Frequency scaling factor. Larger values produce higher-frequency + features at initialisation. + + Returns: + None. Modifies ``layer.weight`` (and ``layer.bias`` if present) in-place. + """ with torch.no_grad(): # Compute the bound for the weights based on the SIREN paper. in_features = layer.in_features @@ -351,19 +544,83 @@ def _init_siren_weights(layer: torch.nn.Linear, is_first_layer: bool, w0: float) class Sine(torch.nn.Module): - """Sine activation used in SIREN with configurable frequency scaling.""" + """Sine activation function used in SIREN networks. + + Computes ``sin(x)`` element-wise. No frequency scaling is applied here; + the ``omega_0`` / ``hidden_omega_0`` factors are absorbed into the weight + initialisation (see ``_init_siren_weights``) so that the effective + frequency at each layer is determined at init without altering the forward + pass arithmetic. + + This design choice follows the SIREN paper and keeps the forward pass + free of any scalar multiplications that might interact poorly with + mixed-precision training. + + Attributes: + (none beyond the base nn.Module bookkeeping) + """ def forward(self, x: torch.Tensor) -> torch.Tensor: - """Forward pass of the Sine activation.""" + """Apply sine element-wise. + + Args: + x: Input tensor of any shape and dtype. + + Returns: + Tensor of the same shape and dtype as ``x`` with values + ``sin(x_i)`` for each element ``x_i``. + """ return torch.sin(x) class SIRENPositionalEmbeddingND(torch.nn.Module): - """Implements a N-dimensional positional embedding using Sine features. + """N-dimensional positional embedding using a SIREN first layer. + + Mathematical form + ----------------- + Given a coordinate grid ``x`` of shape ``[1, *spatial_dims, data_dim]`` + with values normalised to ``[-1, 1]`` per axis, the embedding is: + + phi(x) = sin( W x + b ) shape [..., embedding_dim] - This module generates positional embeddings by applying a linear transformation - with randomized frequencies followed by sine activation. - It is suitable for tasks where positional information needs to be encoded. + where: + + * ``W`` is a learned weight matrix of shape ``[embedding_dim, data_dim]``, + initialised from ``U(-2*pi*omega_0/d, +2*pi*omega_0/d)`` (first-layer + SIREN bound, see ``_init_siren_weights``). Unlike the RFF counterpart, + this weight **is** trainable. + * ``b`` is an optional bias vector, zero-initialised. + + The ``omega_0`` parameter controls the frequency content at init: higher + values bias the embedding toward higher spatial frequencies, giving the + downstream MLP a head-start in representing rapid kernel variations. + During training the weight can drift away from the init distribution. + + Grid caching + ------------ + Identical to ``RandomFourierPositionalEmbeddingND``: a coordinate tensor of + shape ``[1, 2*L_0-1, ..., 2*L_{d-1}-1, data_dim]`` is pre-computed in + float32 and cached as a non-persistent buffer. The forward pass slices the + central ``[2*seq_len_i - 1]`` entries per axis and calls + ``_maybe_extend_grid_cache`` if any axis is larger than the current cache. + + Note: the linear projection is forced to float32 internally (even under + autocast) to avoid quantisation errors in the SIREN's high-frequency sine. + The output is cast back to the weight's dtype before return. + + Attributes: + data_dim (int): Number of spatial / temporal input dimensions. + embedding_dim (int): Output embedding size. + L_cache_per_axis (tuple[int, ...]): Current per-axis cache extents. + L_cache (int | Sequence[int]): Original ``L_cache`` argument (diagnostics). + omega_0 (float): Frequency scaling factor used for SIREN init. + use_bias (bool): Whether a bias is present in the linear projection. + linear (torch.nn.Linear): Trainable SIREN first-layer projection, + shape ``[embedding_dim, data_dim]``. + grid_cache (torch.Tensor): Non-persistent float32 buffer of shape + ``[1, 2*L_0-1, ..., 2*L_{d-1}-1, data_dim]``. + step_sizes (tuple[float, ...]): Per-axis grid step ``1/(L_i - 1)`` + at construction; kept frozen for consistent cache extension. """ def __init__( @@ -518,30 +775,105 @@ def forward(self, seq_lens: tuple[int, ...]) -> tuple[torch.Tensor, torch.Tensor class SIRENKernelND(torch.nn.Module): - """Kernel parameterized by a SIREN (sinusoidal representation network) MLP. + """Convolutional kernel parametrised by a SIREN (sinusoidal representation network) MLP. - The network maps coordinates in an N-D grid directly to kernel values. - Optionally supports FiLM (Feature-wise Linear Modulation) conditioning: - when ``film_cfg`` is provided, a ``KernelFiLMGenerator`` produces per-layer - (gamma, beta) pairs that modulate hidden activations, making the kernel - input-dependent. + Mathematical form + ----------------- + The kernel at coordinate ``x`` is: + + k(x) = Linear_out( SIREN_MLP( phi(x) ) ) + + where: + + * ``phi(x) = sin(W_0 x + b_0)`` is the SIREN positional embedding + (``SIRENPositionalEmbeddingND``) with first-layer frequency ``omega_0``. + * ``SIREN_MLP`` is a stack of ``num_layers - 1`` layers, each computing + ``sin(W_i h + b_i)`` with weights initialised at frequency + ``hidden_omega_0``. + * ``Linear_out`` is a linear readout to ``out_dim`` channels, scaled by + Wang init (``sqrt(1 / kernel_volume)``) to normalise initial kernel energy. + + The full pipeline (without FiLM conditioning) is therefore: + + h_0 = sin(W_0 x + b_0) -- pos embedding + h_i = sin(W_i h_{i-1} + b_i) for i=1..N-1 -- hidden layers + k = W_out h_{N-1} + b_out -- output layer + + Hyperparameters controlling bandwidth / smoothness + --------------------------------------------------- + * ``omega_0``: Frequency of the first SIREN layer. Higher values produce + higher-frequency positional features at init. Typical range: 1.0–30.0. + * ``hidden_omega_0``: Frequency of the hidden SIREN layers. Usually set to + 1.0 (default) following the recommendation in the SIREN paper. + * ``mlp_hidden_dim``: Width of all hidden layers; wider networks can + express more complex kernel shapes. + + FiLM conditioning + ----------------- + When ``film_cfg`` is provided, a ``KernelFiLMGenerator`` is instantiated + and called on the ``conditioning`` tensor (shape ``[B, C]``) to produce + per-layer ``(gamma, beta)`` pairs (each of shape ``[B, mlp_hidden_dim]``). + The hidden activations are then modulated as: + + h_i <- gamma_i * h_i + beta_i + + When ``film_after_pos_embed=True``, an *extra* FiLM layer is applied to + the output of the positional embedding (before the first hidden layer), + making the positional features themselves input-dependent. This requires + ``embedding_dim == mlp_hidden_dim`` and one additional film layer in the + generator (``num_film_layers = num_layers``). + + When conditioning is present, the output kernel has shape + ``[B, *spatial, out_dim]``; otherwise it is ``[1, *spatial, out_dim]``. + + Initialisation + -------------- + * All ``hidden_linears`` are SIREN-initialised with ``hidden_omega_0``. + * ``out_linear`` is SIREN-initialised with ``hidden_omega_0``, then + additionally Wang-scaled by ``sqrt(1 / prod(L_cache_per_axis))``. + * Hidden linear weights and output bias get ``_no_weight_decay = True`` + so that weight-decay optimizers do not destroy the SIREN spectrum. + + Attributes: + out_dim (int): Number of output channels (kernel depth). + data_dim (int): Number of spatial / temporal input dimensions. + mlp_hidden_dim (int): Hidden width of the SIREN MLP. + num_layers (int): Total number of SIREN layers (>= 2). + embedding_dim (int): SIREN positional-embedding dimensionality. + omega_0 (float): First-layer frequency scaling. + hidden_omega_0 (float): Hidden-layer frequency scaling. + L_cache_per_axis (tuple[int, ...]): Per-axis cache extents (canonical form). + L_cache (int | Sequence[int]): Original ``L_cache`` argument (diagnostics). + positional_embedding (SIRENPositionalEmbeddingND): First SIREN layer. + hidden_linears (torch.nn.ModuleList): Hidden linear layers (length + ``num_layers - 1``). Interleaved with ``self.sine`` in the forward + pass; stored separately so FiLM can be inserted between them. + sine (Sine): Shared sine activation applied after every hidden linear. + out_linear (torch.nn.Linear): Final readout to ``out_dim`` channels. + num_film_layers (int): Number of hidden layers eligible for FiLM + modulation (equal to ``len(hidden_linears)``). + film_generator: ``KernelFiLMGenerator`` instance or ``None``. + film_after_pos_embed (bool): Whether the first FiLM pair modulates the + positional embedding output. Args: out_dim: Number of output channels for the generated kernel. data_dim: Number of spatial/temporal input dimensions (size of coordinate vector). mlp_hidden_dim: Hidden width of the SIREN network. num_layers: Total number of layers including the first and hidden layers (>= 2). - L_cache: Cache extent controlling the maximum supported grid size before cache growth. - use_bias: Whether to include biases in linear layers. + embedding_dim: Dimensionality of the SIREN positional embedding. omega_0: Frequency scaling for the first SIREN layer. - hidden_omega_0: Frequency scaling for subsequent SIREN layers. + L_cache: Cache extent controlling the maximum supported grid size before + cache growth. Either a scalar int (isotropic, same extent on all axes) + or a sequence of length ``data_dim`` (anisotropic, per-axis extents). + use_bias: Whether to include biases in linear layers. + hidden_omega_0: Frequency scaling for subsequent SIREN layers (default 1.0). film_cfg: Optional LazyConfig for KernelFiLMGenerator. When provided, enables input-dependent FiLM conditioning of all hidden SIREN layers. film_after_pos_embed: If True, the first FiLM (gamma, beta) pair modulates - the positional embedding *after* the sine activation (i.e. scales/shifts - the ``sin(omega_0 * x)`` output). Requires - ``embedding_dim == mlp_hidden_dim`` and one extra FiLM layer in ``film_cfg`` - (i.e. ``num_film_layers = num_layers - 1 + 1 = num_layers``). + the positional embedding *after* the sine activation. Requires + ``embedding_dim == mlp_hidden_dim`` and one extra FiLM layer in + ``film_cfg`` (i.e. ``num_film_layers = num_layers - 1 + 1 = num_layers``). """ def __init__( @@ -769,11 +1101,30 @@ def forward( def _as_float_tensor(values: Sequence[float] | torch.Tensor, *, name: str) -> torch.Tensor: - """Normalize an ω₀ schedule into a 1D float tensor. + """Normalise an omega_0 schedule argument into a 1-D float64 tensor. + + Accepts any sequence of floats (``list``, ``tuple``) or a 1-D tensor. + Used by the multi-omega SIREN classes so that LazyConfig-built + instantiations (which naturally represent schedules as lists of floats) + work out of the box without the caller having to convert manually. + + All entries must be strictly positive; this is validated here so that + every consumer (``MultiOmegaSIRENPositionalEmbeddingND``, + ``BlockDiagonalMultiOmegaSIRENKernelND``, etc.) gets a consistent error + message and need not repeat the check. + + Args: + values: A 1-D sequence of strictly-positive floats, or a 1-D + ``torch.Tensor``. Multi-dimensional tensors are rejected. + name: Human-readable parameter name used in error messages + (e.g. ``"omega_0_per_row"``). + + Returns: + 1-D ``torch.float64`` tensor with the same values. - Accepts any sequence of floats (``list``, ``tuple``) or a 1D tensor. Used - by the multi-ω₀ SIREN classes so that LazyConfig-built instantiations (which - naturally represent schedules as lists of floats) work out of the box. + Raises: + ValueError: If ``values`` is not 1-D, is empty, or contains a + non-positive entry. """ if isinstance(values, torch.Tensor): out = values.detach().to(dtype=torch.float64).flatten() @@ -795,11 +1146,33 @@ def _build_omega_0_per_block( omega_0_max: float, schedule: str, ) -> torch.Tensor: - """Build a per-block ω₀ vector from (min, max, num_blocks, schedule). + """Build a per-block omega_0 frequency schedule of length ``num_blocks``. - ``schedule`` is either ``"linear"`` (evenly spaced between min and max) or - ``"log"`` (log-spaced — evenly spaced in log-10). Returned tensor has - ``dtype=torch.float64`` and length ``num_blocks``. + Two schedule types are supported: + + * ``"linear"``: equally spaced between ``omega_0_min`` and ``omega_0_max``. + Block ``k`` gets ``omega_0_min + k * (omega_0_max - omega_0_min) / (num_blocks - 1)``. + * ``"log"``: equally spaced in log-10 between the two endpoints. + Block ``k`` gets ``10^(log10(min) + k * (log10(max) - log10(min)) / (num_blocks - 1))``. + + A ``"log"`` schedule is recommended when the frequency range spans more + than a decade (e.g. ``omega_0_min=1, omega_0_max=30``) because it gives + equal coverage to each octave rather than concentrating most blocks near + the high end. + + Args: + num_blocks: Number of frequency blocks. Must be >= 1. + omega_0_min: Lowest frequency in the schedule. Must be positive. + omega_0_max: Highest frequency in the schedule. Must be >= ``omega_0_min``. + schedule: Either ``"linear"`` or ``"log"``. + + Returns: + 1-D ``torch.float64`` tensor of shape ``[num_blocks]`` with the + per-block omega_0 values. + + Raises: + ValueError: If ``num_blocks < 1``, if the endpoints are non-positive, + if ``omega_0_max < omega_0_min``, or if ``schedule`` is unknown. """ if num_blocks < 1: raise ValueError(f"num_blocks must be >= 1, got {num_blocks}") @@ -831,12 +1204,23 @@ class MultiOmegaSIRENPositionalEmbeddingND(SIRENPositionalEmbeddingND): ``BlockDiagonalMultiOmegaSIRENKernelND`` for a variant that also block-masks the MLP to keep rows disjoint at init. + Attributes: + omega_0 (float): Mean of the ``omega_0_per_row`` schedule; stored for + parity with the scalar-``omega_0`` parent's diagnostic attribute. + omega_0_per_row (torch.Tensor): Non-persistent float32 buffer of shape + ``[embedding_dim]`` holding the per-row omega_0 values. + linear (torch.nn.Linear): First-layer weight with per-row SIREN init; + shape ``[embedding_dim, data_dim]``. Each row ``k`` is initialised + from ``U(-2*pi*omega_0_per_row[k]/d, +2*pi*omega_0_per_row[k]/d)``. + grid_cache, step_sizes, L_cache_per_axis, L_cache: + Inherited from ``SIRENPositionalEmbeddingND``; see that class. + Args: data_dim: Number of spatial/temporal input dimensions. embedding_dim: Dimensionality of the positional embedding. L_cache: Cache extent (controls the initial grid cache size). omega_0_per_row: Sequence of ``embedding_dim`` strictly-positive floats - (or a 1-D tensor) giving the ω₀ used for row *k* of the first + (or a 1-D tensor) giving the omega_0 used for row *k* of the first linear. use_bias: Whether to include a bias term. """ @@ -849,7 +1233,7 @@ def __init__( omega_0_per_row: Sequence[float] | torch.Tensor, use_bias: bool = True, ): - """Initialize the per-row multi-ω₀ SIREN positional embedding; see the class docstring.""" + """Initialize the per-row multi-omega SIREN positional embedding; see the class docstring.""" omega = _as_float_tensor(omega_0_per_row, name="omega_0_per_row") if omega.numel() != embedding_dim: raise ValueError(f"omega_0_per_row length ({omega.numel()}) must equal embedding_dim ({embedding_dim})") @@ -891,6 +1275,16 @@ class MultiOmegaSIRENKernelND(SIRENKernelND): The ``omega_0`` attribute reported on the module equals the mean of the schedule, purely for diagnostic purposes. + Attributes: + omega_0 (float): Mean of ``omega_0_per_row``; for diagnostics. + omega_0_per_row (torch.Tensor): Non-persistent float32 buffer of shape + ``[embedding_dim]`` holding the per-row omega_0 schedule. + positional_embedding (MultiOmegaSIRENPositionalEmbeddingND): Per-row + omega_0 positional encoder (replaces the scalar-omega parent's + ``SIRENPositionalEmbeddingND``). + hidden_linears, sine, out_linear, film_generator: + Inherited from :class:`SIRENKernelND`; see that class. + Args: out_dim: Number of output channels for the generated kernel. data_dim: Number of spatial/temporal input dimensions. @@ -898,7 +1292,7 @@ class MultiOmegaSIRENKernelND(SIRENKernelND): num_layers: Total number of SIREN layers (>= 2). embedding_dim: Positional-embedding dimensionality. omega_0_per_row: Sequence of ``embedding_dim`` strictly-positive floats - giving the per-row ω₀ in the first layer. + giving the per-row omega_0 in the first layer. L_cache: Cache extent (controls the initial grid cache size). use_bias: Whether to include biases in linear layers. hidden_omega_0: Frequency scaling for hidden SIREN layers (unchanged @@ -920,7 +1314,7 @@ def __init__( film_cfg: LazyConfig | None = None, film_after_pos_embed: bool = False, ): - """Initialize the per-row multi-ω₀ SIREN kernel; see the class docstring for argument semantics.""" + """Initialize the per-row multi-omega SIREN kernel; see the class docstring for argument semantics.""" omega = _as_float_tensor(omega_0_per_row, name="omega_0_per_row") if omega.numel() != embedding_dim: raise ValueError(f"omega_0_per_row length ({omega.numel()}) must equal embedding_dim ({embedding_dim})") @@ -1021,6 +1415,16 @@ class BlockDiagonalMultiOmegaSIRENKernelND(MultiOmegaSIRENKernelND): ``num_blocks``. When supplied, overrides ``omega_0_min``/``omega_0_max``/``schedule``. hidden_omega_0, film_cfg, film_after_pos_embed: Same as the parent. + + Attributes: + num_blocks (int): Number of frequency blocks. + off_block_scale (float): Off-diagonal weight scale applied at init. + omega_0_per_block (torch.Tensor): Non-persistent float32 buffer of + shape ``[num_blocks]`` holding the per-block omega_0 schedule. + positional_embedding (MultiOmegaSIRENPositionalEmbeddingND): Per-row + omega_0 positional encoder (constant within each block). + hidden_linears, out_linear, omega_0_per_row: + Inherited from :class:`MultiOmegaSIRENKernelND`; see that class. """ def __init__( @@ -1207,6 +1611,24 @@ class LearnableOmegaSIRENPositionalEmbeddingND(SIRENPositionalEmbeddingND): the optimizer behaviour matches the existing SIREN exactly, and the new classes can be A/B-tested with and without LR compensation). + + Attributes: + omega_0 (float): Constant part of the runtime multiplier (same as the + ``omega_0`` constructor argument), stored for diagnostics. + omega_0_scale_min (float): Lower clamp bound on ``omega_0_scale``. + omega_0_scale_max (float): Upper clamp bound on ``omega_0_scale``. + omega_0_const (torch.Tensor): Non-persistent float32 scalar buffer + holding ``2*pi*omega_0``; applied to the linear output at every + forward pass. + omega_0_scale (torch.nn.Parameter): Learnable per-row scale of shape + ``[embedding_dim]``. Clamped to + ``[omega_0_scale_min, omega_0_scale_max]`` by a forward pre-hook + before each forward call. + linear (torch.nn.Linear): First-layer weight ``W`` with *unscaled* + SIREN-1 init ``U(-1/d, +1/d)`` (no ``2*pi*omega_0`` factor in + the bound). Shape ``[embedding_dim, data_dim]``. + grid_cache, step_sizes, L_cache_per_axis, L_cache: + Inherited from ``SIRENPositionalEmbeddingND``; see that class. """ def __init__( @@ -1406,6 +1828,13 @@ class LearnableOmegaSIRENKernelND(SIRENKernelND): :class:`LearnableOmegaSIRENPositionalEmbeddingND` for details. omega_0_scale_min, omega_0_scale_max: Clamp bounds on the scale. apply_lr_scale: Forwarded to the positional embedding. Default False. + + Attributes: + positional_embedding (LearnableOmegaSIRENPositionalEmbeddingND): First + layer with learnable per-row omega_0 scale; replaces the parent's + ``SIRENPositionalEmbeddingND``. + hidden_linears, sine, out_linear, film_generator: + Inherited from :class:`SIRENKernelND`; see that class. """ def __init__( @@ -1426,7 +1855,7 @@ def __init__( film_cfg: LazyConfig | None = None, film_after_pos_embed: bool = False, ): - """Initialize the learnable-ω₀ SIREN kernel; see the class docstring for argument semantics.""" + """Initialize the learnable-omega SIREN kernel; see the class docstring for argument semantics.""" # Build the parent with the same ``omega_0`` so all internal # bookkeeping (``self.omega_0``, ``self.hidden_omega_0``) lines up. super().__init__( @@ -1522,6 +1951,18 @@ class BlockDiagonalLearnableOmegaSIRENKernelND(LearnableOmegaSIRENKernelND): so no row's first-layer sine collapses to a constant. apply_lr_scale: When True, attach ``_lr_scale = 1/(2π·omega_0_max)`` to the first-layer weight. Default False. + + Attributes: + num_blocks (int): Number of frequency blocks. + off_block_scale (float): Off-diagonal weight scale applied at init. + omega_0_per_block (torch.Tensor): Non-persistent float32 buffer of + shape ``[num_blocks]`` holding the per-block omega_0 schedule. + positional_embedding (LearnableOmegaSIRENPositionalEmbeddingND): + First layer with learnable per-row omega_0 scale; ``omega_0_const`` + is set to ``max(omega_0_per_block)`` and ``omega_0_scale`` is + initialised to ``omega_0_per_block / omega_0_const`` per row. + hidden_linears, out_linear, film_generator: + Inherited from :class:`SIRENKernelND`; see that class. """ def __init__( @@ -1546,7 +1987,7 @@ def __init__( film_cfg: LazyConfig | None = None, film_after_pos_embed: bool = False, ): - """Initialize the block-diagonal learnable-ω₀ SIREN kernel; see the class docstring.""" + """Initialize the block-diagonal learnable-omega SIREN kernel; see the class docstring.""" for name, dim in [ ("embedding_dim", embedding_dim), ("mlp_hidden_dim", mlp_hidden_dim), From 54eba08bf694b6e682dac43ab571431ec3673fb3 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:10:07 +0200 Subject: [PATCH 08/72] docs(review/mixed_fftconv): reviewer feedback --- docs/reviews/mixed_fftconv_review.md | 285 +++++++++++++++++++++++++++ 1 file changed, 285 insertions(+) create mode 100644 docs/reviews/mixed_fftconv_review.md diff --git a/docs/reviews/mixed_fftconv_review.md b/docs/reviews/mixed_fftconv_review.md new file mode 100644 index 00000000..da50f587 --- /dev/null +++ b/docs/reviews/mixed_fftconv_review.md @@ -0,0 +1,285 @@ +# Review: `nvsubquadratic/ops/mixed_fftconv.py` + +Reviewed as a new external collaborator reading alongside the research paper. +All item numbers below are concrete fixes to apply. + +______________________________________________________________________ + +## Module docstring + +**1. The "same" alignment example in the module-level "Algorithm overview" section +uses the tuple form `(True, False)` in the code snippet but the call signature +accepts `Sequence[bool]`, not a plain tuple label.** + +Quote: `y = mixed_fftconv2d_fp32_bhl(x, kernel, periodic=(True, False))` + +No issue with the Python itself, but the module docstring does not mention that +`CKConvND` uses the *string* form (`fft_padding=["circular", "zero"]`) and that +the `periodic` bool-tuple is the *internal* normalised form. A reader glancing at +this module and then looking at a config YAML will be confused about why their +config uses strings but the function takes bools. + +**Fix:** Add a one-sentence note at the end of the module-level "When to use this" +section, e.g.: + +> Note: `CKConvND` accepts `fft_padding=["circular", "zero"]` (string form). +> It internally converts this to the boolean `periodic` tuple before calling +> these ops. + +______________________________________________________________________ + +## `_mixed_recipe` + +**2. The docstring says the crop for a non-periodic axis is +`[K_d // 2, K_d // 2 + N_d)`, but the code uses integer floor division. For an +even kernel `K_d = 4`, the crop starts at index 2 — which is the "left-aligned" +half — but the docstring never explains the asymmetry for even kernels.** + +Quote from docstring: + +> crop `[K_d // 2, K_d // 2 + N_d)` + +The crop is correct for odd kernels (symmetric). For even kernels it applies the +"left half of even kernel" convention (matching `torch.nn.Conv*d(padding='same')` +with `dilation=1`), which is a non-obvious choice. If a collaborator tries to +reproduce results with manual `F.pad` they will get the wrong result unless they +know this. + +**Fix:** Add a note after the crop description: + +> For even `K_d`, floor division gives a left-biased crop (one extra sample on +> the right), matching the convention of `torch.nn.ConvNd(padding='same')`. +> Tests verify this against the spatial reference. + +______________________________________________________________________ + +**3. The function signature has `kshape` as a parameter name but the docstring +calls it "Kernel spatial dims". The abbreviation `kshape` is fine, but the +parameter name in the Args block should match the actual argument name exactly.** + +Quote from docstring: + +> `kshape: Kernel spatial dims `(K_0, ..., K\_\{D-1})`.` + +This is correct — just flag that the double space before "Kernel" is a minor +formatting inconsistency that ruff/mdformat may flag in future. + +**Fix:** Remove the extra space: `kshape: Kernel spatial dims ...` + +______________________________________________________________________ + +## `_MixedPhaseRamp1DCache` + +**4. The class docstring says "Non-periodic axes use `s_d = 0` and contribute +no ramp." but it does not explain *why* `s_d = 0` for non-periodic axes. A reader +unfamiliar with the algorithm will not know that non-periodic alignment is handled +by the centered crop instead.** + +Quote: + +> Non-periodic axes use `s_d = 0` and contribute no ramp. + +**Fix:** Append the explanation: + +> Non-periodic axes use `s_d = 0` because their "same" alignment is handled +> entirely by the centered crop `[K_d // 2, K_d // 2 + N_d)` in the IFFT +> output — no frequency-domain shift is needed. + +______________________________________________________________________ + +## `_MixedPhaseRamp1DCache.get` + +**5. The Returns line says "1-D complex tensor of shape `[F // 2 + 1]` (rfft +axis) or `[F]` (non-rfft axis)". But when `s == 0` the method still builds +and returns a tensor of all-ones (as the docstring says "returns all-ones for +completeness"). This is never actually reached in practice because +`_build_nd_phase_ramp` skips `s == 0` axes — so the "all-ones" note implies a +code path that never runs. This could mislead a reader into thinking `s == 0` +is a valid call site.** + +Quote: + +> Callers with `s == 0` are expected to skip the multiply entirely; this +> method still handles `s == 0` (returns all-ones). + +**Fix:** Remove the parenthetical "returns all-ones" claim (or move it to a +comment in the code), and instead say: + +> In practice `_build_nd_phase_ramp` skips axes where `s == 0` before +> calling this method. + +______________________________________________________________________ + +## `_build_nd_phase_ramp` + +**6. The docstring says "The N-D ramp is not cached; only the 1-D per-axis ramps +are." but does not mention that the returned tensor is a *view* of the 1-D cached +tensors (broadcast, not materialised as a separate allocation). A reader +concerned about memory may not realise that the N-D tensor is cheap — it is a +view chain, not a full materialisation.** + +Quote: + +> The N-D ramp is **not** cached; only the 1-D per-axis ramps are. The +> product is materialised here (via broadcasted multiplication) so that the +> downstream multiply with `fft_x` is a single fused op. + +The word "materialised" is ambiguous — it sounds like the whole N-D tensor is +allocated, but the product loop in the code multiplies the broadcast views, +producing a new (small, single-axis-varying) tensor each time. + +**Fix:** Clarify: + +> The N-D ramp itself is not cached. The loop multiplies broadcast 1-D views +> together, producing a compact tensor (not a full `(F_0, ..., F_{D-1})` +> allocation) that covers only the periodic axes with non-zero shifts. + +______________________________________________________________________ + +## `_mixed_fftconv_nd_fp32_bhl` + +**7. Step 6 in the docstring says "apply `torch.roll` on the periodic axes" +but the code actually filters to only the axes where `s != 0` (i.e., it +excludes size-1 periodic kernels):** + +Quote from code: + +```python +roll_shifts = tuple(s for s in shifts if s != 0) +roll_dims = tuple(2 + d for d, s in enumerate(shifts) if s != 0) +``` + +The docstring says "apply torch.roll on the periodic axes" which is slightly +inaccurate — it applies to periodic axes *with non-zero shifts*. + +**Fix:** Change step 6 to: + +> 6. If `use_phase_shift=False`, apply `torch.roll` on periodic axes that +> have a non-zero shift (`s_d != 0`); size-1 periodic kernels (shift 0) +> are skipped. + +______________________________________________________________________ + +**8. The kernel-size limit comments in the code (explaining why periodic axes +use `K <= N` and non-periodic use `K <= 2*N`) are inside `_mixed_fftconv_nd_fp32_bhl` +as inline comments but are not surfaced in the docstring's Raises section. +An external collaborator who passes an oversized kernel will get an +`AssertionError` with a message but no hint about what "double-grid" means.** + +Quote from Raises: + +> AssertionError: On shape mismatches or out-of-range kernel sizes. + +**Fix:** Expand the Raises block: + +> AssertionError: On shape mismatches or out-of-range kernel sizes. +> Specifically: `K_d > N_d` on a periodic axis (circular FFT length is +> `N_d`, so the kernel must fit); `K_d > 2*N_d` on a non-periodic axis +> (the padded FFT length is at most `2*N_d`, matching the standard +> "double-grid" SIREN kernel size `2*N_d - 1`). + +______________________________________________________________________ + +## Public 1D/2D/3D BHL entry points + +**9. `mixed_fftconv1d_fp32_bhl` and its 2D/3D siblings each reference +`use_phase_shift` but the 2D and 3D variants say "See +:func:`mixed_fftconv1d_fp32_bhl`" for its meaning. That cross-reference works, +but the reader must jump to the 1D function to understand the parameter — and +the 1D function's `use_phase_shift` description does not mention the performance +trade-off explicitly.** + +Quote from `mixed_fftconv1d_fp32_bhl`: + +> use_phase_shift: If True, align periodic axes via frequency-domain phase +> ramps. If False, align via :func:`torch.roll` on periodic axes after +> the inverse transform. The output is mathematically equivalent. + +**Fix:** Add a performance hint to the 1D description: + +> `use_phase_shift=True` (default) is faster — the shift is fused into the +> frequency-domain multiply with no additional data movement. Use `False` +> only as a reference or when `torch.compile` cannot handle complex ops. + +______________________________________________________________________ + +## BLH `_w_reshape` wrappers + +**10. The kernel argument docstring for `mixed_fftconv1d_fp32_bhl_w_reshape` +says ``` kernel: Kernel tensor of shape ``[B, K, H]`` (BLH) ```. But the BHL +wrappers also support a shared-kernel leading dim of 1 (`[1, K, H]`), which +is the standard depthwise case. The docstring is inconsistent with the BHL +entry points (which say `[1|B, H, K]`) and with the module-level shape +conventions.** + +Quote: + +> kernel: Kernel tensor of shape `[B, K, H]` (BLH). + +**Fix:** Change to `[1|B, K, H]` in all `_w_reshape` wrapper docstrings (1D, +2D, 3D, and their chunked counterparts). + +______________________________________________________________________ + +## Chunked variants + +**11. `_DEFAULT_MIXED_CHUNK_SIZE = 128` is a module-level constant referenced in +the chunked function docstrings but has no docstring or comment of its own. A +reader tuning memory usage will not know why 128 was chosen.** + +The related `fftconv_chunked.py` module mentions "~26% memory savings for ~11% +overhead" for chunk size 128 — that context is missing here. + +**Fix:** Add a comment above the constant: + +```python +# Default channel chunk size for the memory-efficient variants. +# A chunk of 128 channels typically gives ~26% peak-memory savings +# with ~11% throughput overhead relative to the non-chunked path +# (measured on H100; see fftconv_chunked.py for profiling details). +_DEFAULT_MIXED_CHUNK_SIZE = 128 +``` + +______________________________________________________________________ + +## Missing usage example + +**12. None of the public entry points include a minimal runnable example. The +module docstring has a snippet but it shows only the function call, not the +tensor shapes, so a new user cannot copy-paste-and-run it. A short `Example:` +block on `mixed_fftconv2d_fp32_bhl` (the most common use case) would help +significantly.** + +**Fix:** Add to `mixed_fftconv2d_fp32_bhl`: + +```python +Example: + >>> import torch + >>> from nvsubquadratic.ops.mixed_fftconv import mixed_fftconv2d_fp32_bhl + >>> B, H, X, Y, Kx, Ky = 2, 64, 32, 64, 63, 127 + >>> x = torch.randn(B, H, X, Y) + >>> kernel = torch.randn(1, H, Kx, Ky) + >>> # x periodic, y zero-padded + >>> y = mixed_fftconv2d_fp32_bhl(x, kernel, periodic=(True, False)) + >>> y.shape + torch.Size([2, 64, 32, 64]) +``` + +______________________________________________________________________ + +## Summary of fixes to apply + +| # | Location | Fix type | +| --- | ---------------------------- | -------------------------------------------- | +| 1 | Module docstring | Add note about string form in CKConvND | +| 2 | `_mixed_recipe` | Clarify even-kernel asymmetry | +| 3 | `_mixed_recipe` | Remove double space in Args | +| 4 | `_MixedPhaseRamp1DCache` | Explain why `s_d = 0` for non-periodic | +| 5 | `_MixedPhaseRamp1DCache.get` | Remove misleading "returns all-ones" claim | +| 6 | `_build_nd_phase_ramp` | Clarify "materialised" language | +| 7 | `_mixed_fftconv_nd_fp32_bhl` | Fix step 6 (roll only on `s != 0` axes) | +| 8 | `_mixed_fftconv_nd_fp32_bhl` | Expand Raises block with kernel-size context | +| 9 | `mixed_fftconv1d_fp32_bhl` | Add performance hint to `use_phase_shift` | +| 10 | All `_w_reshape` wrappers | Fix kernel shape to `[1\|B, K, H]` | +| 11 | `_DEFAULT_MIXED_CHUNK_SIZE` | Add comment with memory-savings context | +| 12 | `mixed_fftconv2d_fp32_bhl` | Add runnable `Example:` block | From 5fe4251a5cb42f4d354e9546a999d1099e9b1cfc Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:10:55 +0200 Subject: [PATCH 09/72] docs(review/kernels_nd): reviewer feedback --- docs/reviews/kernels_nd_review.md | 282 ++++++++++++++++++++++++++++++ 1 file changed, 282 insertions(+) create mode 100644 docs/reviews/kernels_nd_review.md diff --git a/docs/reviews/kernels_nd_review.md b/docs/reviews/kernels_nd_review.md new file mode 100644 index 00000000..1debef0d --- /dev/null +++ b/docs/reviews/kernels_nd_review.md @@ -0,0 +1,282 @@ +# Reviewer feedback: `nvsubquadratic/modules/kernels_nd.py` + +Second-pass review after Phase 1 docstring additions. Issues are numbered; each +quotes the offending text and says exactly what to fix. + +______________________________________________________________________ + +## 1. Module docstring: `RandomFourierKernelND` is not listed in the "Kernel classes" table + +**Location**: module-level docstring, "Kernel classes in this module" section. + +The table lists six consumer-facing kernel classes but omits +`RandomFourierKernelND` — the table starts with `SIRENKernelND`. The class is +described in the prose above, but a new reader scanning the table will not find +it. + +**Fix**: add `RandomFourierKernelND` as the first entry in the list, before +`SIRENKernelND`. + +______________________________________________________________________ + +## 2. `_normalize_l_cache`: missing `Raises` block + +**Location**: `_normalize_l_cache` docstring. + +The function raises `TypeError` (for bool or non-int/sequence input) and +`ValueError` (for wrong sequence length or values \< 2), but the docstring has +no `Raises:` section. An external collaborator reading this to decide whether +to guard against exceptions will not know what to catch. + +**Fix**: add a `Raises:` section listing both `TypeError` and `ValueError` +with the conditions that trigger each. + +______________________________________________________________________ + +## 3. `RandomFourierPositionalEmbeddingND.forward`: wrong shape annotation for grid return value + +**Location**: `forward` docstring, `Returns` block: + +> `torch.Tensor: The input positions normalized between [-1, 1] (shape: [1, * spatial_dims, 1]).` + +The trailing dimension is `data_dim` (the number of spatial axes), not `1`. +A 2D model has a grid of shape `[1, H, W, 2]`, not `[1, H, W, 1]`. The +`_build_grid_cache` code confirms: `dim=-1` of the stacked meshgrid has size +`data_dim`. + +**Fix**: change the shape annotation to `[1, *spatial_dims, data_dim]`. + +The same bug appears in `SIRENPositionalEmbeddingND.forward` — fix it there +too. + +______________________________________________________________________ + +## 4. `RandomFourierPositionalEmbeddingND.forward` and `SIRENPositionalEmbeddingND.forward`: stale `"concatenated sine and cosine values"` description for SIREN + +**Location**: `SIRENPositionalEmbeddingND.forward`, `Returns` block: + +> `torch.Tensor: The positional embeddings, concatenated sine and cosine values (shape: [1, * spatial_dims, embedding_dim]).` + +The SIREN embedding is `sin(Wx+b)` — a single sine, no cosine concatenation. +The description was copied verbatim from the RFF embedding, where `[cos, sin]` +concatenation is correct. + +**Fix**: change the SIREN `forward` return description to: +"The positional embeddings, `sin(W x + b)`, shape `[1, *spatial_dims, embedding_dim]`." + +______________________________________________________________________ + +## 5. `RandomFourierKernelND.forward`: `Returns` annotation says "tuple" but the function actually returns two tensors described as a single `tuple[torch.Tensor, torch.Tensor]` + +**Location**: `RandomFourierKernelND.forward`, `Returns` block: + +> `tuple[torch.Tensor, torch.Tensor]: The computed random Fourier kernel and the corresponding grid values. The kernel is a tensor of shape (1, * spatial_dims, out_dim) The grid is a tensor of shape (1, * spatial_dims, data_dim)` + +The format mixes the inline `tuple[…]` type-hint style with a multi-line +description, making it hard to parse. The return type annotation in the +`Returns:` block should follow Google style: one sub-item per returned object. + +**Fix**: rewrite as: + +``` +Returns: + tuple: + - torch.Tensor: Kernel values of shape ``[1, *spatial_dims, out_dim]``. + - torch.Tensor: Coordinate grid of shape ``[1, *spatial_dims, data_dim]``. +``` + +______________________________________________________________________ + +## 6. `SIRENKernelND` class docstring: `flop_count` method has no standalone docstring header describing its purpose in one sentence + +**Location**: `SIRENKernelND.flop_count`. + +The docstring jumps directly into the "At `inference=True`…" explanation +without a one-sentence summary. Google style requires a short summary line +first. + +**Fix**: add as the first line: + +> "Return an integer FLOP estimate for one kernel generation forward pass." + +______________________________________________________________________ + +## 7. `SIRENKernelND.flop_count`: `grid_lens` description is ambiguous + +**Location**: `flop_count` docstring, `Args:` block: + +> ``` grid_lens: Spatial extents passed to the positional embedding. The kernel grid has size ``(2 * L - 1)`` per dimension. ``` + +This does not say what unit `grid_lens` is in (seq_lens? cache extents?). The +code uses `for L in grid_lens: G *= 2 * L - 1`, so `L` here is the *output* +sequence length (each axis), not the cache extent. The description "Spatial +extents passed to the positional embedding" is correct but the parenthetical +`(2 * L - 1)` makes it look like `grid_lens` are the half-lengths. + +**Fix**: replace with: + +> ``` grid_lens: Per-axis output sequence lengths, i.e. the same tuple you would pass to ``forward`` as ``seq_lens``. The FLOP count uses the grid volume ``G = prod(2*L - 1 for L in grid_lens)`` as the number of coordinate points the MLP processes. ``` + +______________________________________________________________________ + +## 8. `_build_grid_cache`: missing Args/Returns in both `RandomFourierPositionalEmbeddingND` and `SIRENPositionalEmbeddingND` copies + +**Location**: `RandomFourierPositionalEmbeddingND._build_grid_cache` and +`SIRENPositionalEmbeddingND._build_grid_cache`. + +Both static methods have a one-paragraph description but no `Args:` or +`Returns:` sections. The parameters `L_per_axis`, `max_limits`, and `device` +are not described, and the return shape is not stated. + +**Fix**: add: + +``` +Args: + L_per_axis: Per-axis cache extents. The cache size along axis ``i`` + is ``2 * L_per_axis[i] - 1`` grid points. + max_limits: Per-axis coordinate limits. Each axis spans + ``[-max_limits[i], +max_limits[i]]``. Defaults to ``1.0`` on all + axes. + device: Target device for the returned tensor. Defaults to CPU. + +Returns: + Float32 tensor of shape ``[1, 2*L_0-1, ..., 2*L_{d-1}-1, data_dim]`` + representing the coordinate meshgrid, with a leading batch dimension of 1. +``` + +(The two copies are byte-identical; apply the same fix to both.) + +______________________________________________________________________ + +## 9. `_maybe_extend_grid_cache`: missing Args/Returns in both copies + +**Location**: `RandomFourierPositionalEmbeddingND._maybe_extend_grid_cache` and +`SIRENPositionalEmbeddingND._maybe_extend_grid_cache`. + +Same issue as item 8: no `Args:` or `Returns:` sections. + +**Fix**: add: + +``` +Args: + seq_lens: Requested per-axis output sequence lengths. Any axis where + ``seq_lens[i] > self.L_cache_per_axis[i]`` triggers a cache + extension. + +Returns: + None. Modifies ``self.grid_cache`` and ``self.L_cache_per_axis`` + in-place when an extension is needed. +``` + +______________________________________________________________________ + +## 10. `LearnableOmegaSIRENPositionalEmbeddingND._clamp_omega_scale_pre_hook`: missing Args/Returns + +**Location**: `_clamp_omega_scale_pre_hook`. + +The hook is registered as a forward pre-hook and takes `(module, inputs)` but +the docstring has no `Args:` block, which is required for any public/protected +method. + +**Fix**: add: + +``` +Args: + module: The module instance (``self``); provided by PyTorch's hook + mechanism. + inputs: The positional inputs tuple; not used. + +Returns: + None. Modifies ``self.omega_0_scale.data`` in-place. +``` + +______________________________________________________________________ + +## 11. `BlockDiagonalMultiOmegaSIRENKernelND._block_mask`: `out_dim` / `in_dim` shape contract not stated precisely enough + +**Location**: `_block_mask` docstring: + +> "`out_dim` and `in_dim` must both be divisible by `num_blocks`." + +This constraint is documented but neither a `Raises:` nor a note says what +happens if violated (currently it silently produces wrong-sized blocks because +integer division truncates). + +**Fix**: add a `Raises:` section: + +``` +Raises: + ZeroDivisionError: If ``num_blocks == 0``. + Note: if ``out_dim`` or ``in_dim`` is not divisible by ``num_blocks``, + the block sizes are silently truncated via integer division; the + caller (``BlockDiagonalMultiOmegaSIRENKernelND.__init__``) enforces + divisibility before calling this method. +``` + +and add `Args:` / `Returns:` sections (currently absent): + +``` +Args: + out_dim: Number of output features (rows of the weight matrix). + in_dim: Number of input features (columns of the weight matrix). + num_blocks: Number of blocks. + off_block_scale: Scalar value for off-diagonal block entries. + device: Target device. + dtype: Target dtype. + +Returns: + Float tensor of shape ``[out_dim, in_dim]`` with ``1.0`` on the + block diagonal and ``off_block_scale`` elsewhere. +``` + +______________________________________________________________________ + +## 12. `SIRENKernelND` class docstring: "Wang initialisation" is unexplained + +**Location**: `SIRENKernelND` class docstring, "Initialisation" section: + +> "The output layer applies **Wang initialisation**: weights are scaled by `sqrt(1 / kernel_volume)`…" + +An external collaborator unfamiliar with this term will not know where it comes +from. The SIREN paper (Sitzmann et al. 2020) does not use the name "Wang init" +— this is a practice from continuous convolution literature. + +**Fix**: add a brief parenthetical reference, e.g.: + +> "(introduced for implicit kernel networks by \[Wang et al.\]; equivalent to dividing the output layer's weights by the square root of the total number of kernel grid points, so that the initial filter energy is independent of grid size)" + +or cite the CKConv paper where this practice originates. + +______________________________________________________________________ + +## 13. `LearnableOmegaSIRENPositionalEmbeddingND` class docstring: `apply_lr_scale` mentions `_build_param_groups` without context + +**Location**: class docstring, `Args:` for `apply_lr_scale`: + +> "attach `_lr_scale = 1/(2π·omega_0)` to `self.linear.weight` so that `_build_param_groups` lowers its learning rate…" + +`_build_param_groups` is an internal optimizer helper not visible from this +file. A reader who does not know the codebase's optimizer infrastructure will +not understand how `_lr_scale` is consumed. + +**Fix**: replace with: + +> "When True, attach `_lr_scale = 1/(2*pi*omega_0)` to `self.linear.weight`. The optimizer utility `_build_param_groups` (in `experiments/`) reads this attribute and multiplies the layer's learning rate by `_lr_scale`, compensating for the missing `2*pi*omega_0` factor in the SIREN-1 init bound." + +______________________________________________________________________ + +## 14. Module docstring: no mention of where the `conditioning` tensor comes from + +**Location**: module docstring, no mention of FiLM conditioning flow. + +The module docstring describes the two kernel families (RFF and SIREN) but +does not mention that `SIRENKernelND` and its descendants support a +`conditioning` argument that makes the kernel input-dependent. This is the +FiLM path, and a new collaborator reading the module overview has no idea this +feature exists or how to trigger it. + +**Fix**: add a short paragraph in the Overview section: + +> "Optionally, `SIRENKernelND` (and all SIREN subclasses) accept a `conditioning` tensor of shape `[B, C]` via their `forward` method. When a `KernelFiLMGenerator` is wired in via `film_cfg`, the kernel becomes input-dependent — each sample in the batch gets a different filter. This is used in diffusion models and other conditional generation tasks where the kernel must respond to a global context signal." + +______________________________________________________________________ From 3117a38925ce47dd2da3f5363cfae8d8b7c0dcfa Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:12:26 +0200 Subject: [PATCH 10/72] docs(integrate/mixed_fftconv): apply reviewer feedback --- docs/reviews/mixed_fftconv_review.md | 285 --------------------------- nvsubquadratic/ops/mixed_fftconv.py | 82 ++++++-- 2 files changed, 62 insertions(+), 305 deletions(-) delete mode 100644 docs/reviews/mixed_fftconv_review.md diff --git a/docs/reviews/mixed_fftconv_review.md b/docs/reviews/mixed_fftconv_review.md deleted file mode 100644 index da50f587..00000000 --- a/docs/reviews/mixed_fftconv_review.md +++ /dev/null @@ -1,285 +0,0 @@ -# Review: `nvsubquadratic/ops/mixed_fftconv.py` - -Reviewed as a new external collaborator reading alongside the research paper. -All item numbers below are concrete fixes to apply. - -______________________________________________________________________ - -## Module docstring - -**1. The "same" alignment example in the module-level "Algorithm overview" section -uses the tuple form `(True, False)` in the code snippet but the call signature -accepts `Sequence[bool]`, not a plain tuple label.** - -Quote: `y = mixed_fftconv2d_fp32_bhl(x, kernel, periodic=(True, False))` - -No issue with the Python itself, but the module docstring does not mention that -`CKConvND` uses the *string* form (`fft_padding=["circular", "zero"]`) and that -the `periodic` bool-tuple is the *internal* normalised form. A reader glancing at -this module and then looking at a config YAML will be confused about why their -config uses strings but the function takes bools. - -**Fix:** Add a one-sentence note at the end of the module-level "When to use this" -section, e.g.: - -> Note: `CKConvND` accepts `fft_padding=["circular", "zero"]` (string form). -> It internally converts this to the boolean `periodic` tuple before calling -> these ops. - -______________________________________________________________________ - -## `_mixed_recipe` - -**2. The docstring says the crop for a non-periodic axis is -`[K_d // 2, K_d // 2 + N_d)`, but the code uses integer floor division. For an -even kernel `K_d = 4`, the crop starts at index 2 — which is the "left-aligned" -half — but the docstring never explains the asymmetry for even kernels.** - -Quote from docstring: - -> crop `[K_d // 2, K_d // 2 + N_d)` - -The crop is correct for odd kernels (symmetric). For even kernels it applies the -"left half of even kernel" convention (matching `torch.nn.Conv*d(padding='same')` -with `dilation=1`), which is a non-obvious choice. If a collaborator tries to -reproduce results with manual `F.pad` they will get the wrong result unless they -know this. - -**Fix:** Add a note after the crop description: - -> For even `K_d`, floor division gives a left-biased crop (one extra sample on -> the right), matching the convention of `torch.nn.ConvNd(padding='same')`. -> Tests verify this against the spatial reference. - -______________________________________________________________________ - -**3. The function signature has `kshape` as a parameter name but the docstring -calls it "Kernel spatial dims". The abbreviation `kshape` is fine, but the -parameter name in the Args block should match the actual argument name exactly.** - -Quote from docstring: - -> `kshape: Kernel spatial dims `(K_0, ..., K\_\{D-1})`.` - -This is correct — just flag that the double space before "Kernel" is a minor -formatting inconsistency that ruff/mdformat may flag in future. - -**Fix:** Remove the extra space: `kshape: Kernel spatial dims ...` - -______________________________________________________________________ - -## `_MixedPhaseRamp1DCache` - -**4. The class docstring says "Non-periodic axes use `s_d = 0` and contribute -no ramp." but it does not explain *why* `s_d = 0` for non-periodic axes. A reader -unfamiliar with the algorithm will not know that non-periodic alignment is handled -by the centered crop instead.** - -Quote: - -> Non-periodic axes use `s_d = 0` and contribute no ramp. - -**Fix:** Append the explanation: - -> Non-periodic axes use `s_d = 0` because their "same" alignment is handled -> entirely by the centered crop `[K_d // 2, K_d // 2 + N_d)` in the IFFT -> output — no frequency-domain shift is needed. - -______________________________________________________________________ - -## `_MixedPhaseRamp1DCache.get` - -**5. The Returns line says "1-D complex tensor of shape `[F // 2 + 1]` (rfft -axis) or `[F]` (non-rfft axis)". But when `s == 0` the method still builds -and returns a tensor of all-ones (as the docstring says "returns all-ones for -completeness"). This is never actually reached in practice because -`_build_nd_phase_ramp` skips `s == 0` axes — so the "all-ones" note implies a -code path that never runs. This could mislead a reader into thinking `s == 0` -is a valid call site.** - -Quote: - -> Callers with `s == 0` are expected to skip the multiply entirely; this -> method still handles `s == 0` (returns all-ones). - -**Fix:** Remove the parenthetical "returns all-ones" claim (or move it to a -comment in the code), and instead say: - -> In practice `_build_nd_phase_ramp` skips axes where `s == 0` before -> calling this method. - -______________________________________________________________________ - -## `_build_nd_phase_ramp` - -**6. The docstring says "The N-D ramp is not cached; only the 1-D per-axis ramps -are." but does not mention that the returned tensor is a *view* of the 1-D cached -tensors (broadcast, not materialised as a separate allocation). A reader -concerned about memory may not realise that the N-D tensor is cheap — it is a -view chain, not a full materialisation.** - -Quote: - -> The N-D ramp is **not** cached; only the 1-D per-axis ramps are. The -> product is materialised here (via broadcasted multiplication) so that the -> downstream multiply with `fft_x` is a single fused op. - -The word "materialised" is ambiguous — it sounds like the whole N-D tensor is -allocated, but the product loop in the code multiplies the broadcast views, -producing a new (small, single-axis-varying) tensor each time. - -**Fix:** Clarify: - -> The N-D ramp itself is not cached. The loop multiplies broadcast 1-D views -> together, producing a compact tensor (not a full `(F_0, ..., F_{D-1})` -> allocation) that covers only the periodic axes with non-zero shifts. - -______________________________________________________________________ - -## `_mixed_fftconv_nd_fp32_bhl` - -**7. Step 6 in the docstring says "apply `torch.roll` on the periodic axes" -but the code actually filters to only the axes where `s != 0` (i.e., it -excludes size-1 periodic kernels):** - -Quote from code: - -```python -roll_shifts = tuple(s for s in shifts if s != 0) -roll_dims = tuple(2 + d for d, s in enumerate(shifts) if s != 0) -``` - -The docstring says "apply torch.roll on the periodic axes" which is slightly -inaccurate — it applies to periodic axes *with non-zero shifts*. - -**Fix:** Change step 6 to: - -> 6. If `use_phase_shift=False`, apply `torch.roll` on periodic axes that -> have a non-zero shift (`s_d != 0`); size-1 periodic kernels (shift 0) -> are skipped. - -______________________________________________________________________ - -**8. The kernel-size limit comments in the code (explaining why periodic axes -use `K <= N` and non-periodic use `K <= 2*N`) are inside `_mixed_fftconv_nd_fp32_bhl` -as inline comments but are not surfaced in the docstring's Raises section. -An external collaborator who passes an oversized kernel will get an -`AssertionError` with a message but no hint about what "double-grid" means.** - -Quote from Raises: - -> AssertionError: On shape mismatches or out-of-range kernel sizes. - -**Fix:** Expand the Raises block: - -> AssertionError: On shape mismatches or out-of-range kernel sizes. -> Specifically: `K_d > N_d` on a periodic axis (circular FFT length is -> `N_d`, so the kernel must fit); `K_d > 2*N_d` on a non-periodic axis -> (the padded FFT length is at most `2*N_d`, matching the standard -> "double-grid" SIREN kernel size `2*N_d - 1`). - -______________________________________________________________________ - -## Public 1D/2D/3D BHL entry points - -**9. `mixed_fftconv1d_fp32_bhl` and its 2D/3D siblings each reference -`use_phase_shift` but the 2D and 3D variants say "See -:func:`mixed_fftconv1d_fp32_bhl`" for its meaning. That cross-reference works, -but the reader must jump to the 1D function to understand the parameter — and -the 1D function's `use_phase_shift` description does not mention the performance -trade-off explicitly.** - -Quote from `mixed_fftconv1d_fp32_bhl`: - -> use_phase_shift: If True, align periodic axes via frequency-domain phase -> ramps. If False, align via :func:`torch.roll` on periodic axes after -> the inverse transform. The output is mathematically equivalent. - -**Fix:** Add a performance hint to the 1D description: - -> `use_phase_shift=True` (default) is faster — the shift is fused into the -> frequency-domain multiply with no additional data movement. Use `False` -> only as a reference or when `torch.compile` cannot handle complex ops. - -______________________________________________________________________ - -## BLH `_w_reshape` wrappers - -**10. The kernel argument docstring for `mixed_fftconv1d_fp32_bhl_w_reshape` -says ``` kernel: Kernel tensor of shape ``[B, K, H]`` (BLH) ```. But the BHL -wrappers also support a shared-kernel leading dim of 1 (`[1, K, H]`), which -is the standard depthwise case. The docstring is inconsistent with the BHL -entry points (which say `[1|B, H, K]`) and with the module-level shape -conventions.** - -Quote: - -> kernel: Kernel tensor of shape `[B, K, H]` (BLH). - -**Fix:** Change to `[1|B, K, H]` in all `_w_reshape` wrapper docstrings (1D, -2D, 3D, and their chunked counterparts). - -______________________________________________________________________ - -## Chunked variants - -**11. `_DEFAULT_MIXED_CHUNK_SIZE = 128` is a module-level constant referenced in -the chunked function docstrings but has no docstring or comment of its own. A -reader tuning memory usage will not know why 128 was chosen.** - -The related `fftconv_chunked.py` module mentions "~26% memory savings for ~11% -overhead" for chunk size 128 — that context is missing here. - -**Fix:** Add a comment above the constant: - -```python -# Default channel chunk size for the memory-efficient variants. -# A chunk of 128 channels typically gives ~26% peak-memory savings -# with ~11% throughput overhead relative to the non-chunked path -# (measured on H100; see fftconv_chunked.py for profiling details). -_DEFAULT_MIXED_CHUNK_SIZE = 128 -``` - -______________________________________________________________________ - -## Missing usage example - -**12. None of the public entry points include a minimal runnable example. The -module docstring has a snippet but it shows only the function call, not the -tensor shapes, so a new user cannot copy-paste-and-run it. A short `Example:` -block on `mixed_fftconv2d_fp32_bhl` (the most common use case) would help -significantly.** - -**Fix:** Add to `mixed_fftconv2d_fp32_bhl`: - -```python -Example: - >>> import torch - >>> from nvsubquadratic.ops.mixed_fftconv import mixed_fftconv2d_fp32_bhl - >>> B, H, X, Y, Kx, Ky = 2, 64, 32, 64, 63, 127 - >>> x = torch.randn(B, H, X, Y) - >>> kernel = torch.randn(1, H, Kx, Ky) - >>> # x periodic, y zero-padded - >>> y = mixed_fftconv2d_fp32_bhl(x, kernel, periodic=(True, False)) - >>> y.shape - torch.Size([2, 64, 32, 64]) -``` - -______________________________________________________________________ - -## Summary of fixes to apply - -| # | Location | Fix type | -| --- | ---------------------------- | -------------------------------------------- | -| 1 | Module docstring | Add note about string form in CKConvND | -| 2 | `_mixed_recipe` | Clarify even-kernel asymmetry | -| 3 | `_mixed_recipe` | Remove double space in Args | -| 4 | `_MixedPhaseRamp1DCache` | Explain why `s_d = 0` for non-periodic | -| 5 | `_MixedPhaseRamp1DCache.get` | Remove misleading "returns all-ones" claim | -| 6 | `_build_nd_phase_ramp` | Clarify "materialised" language | -| 7 | `_mixed_fftconv_nd_fp32_bhl` | Fix step 6 (roll only on `s != 0` axes) | -| 8 | `_mixed_fftconv_nd_fp32_bhl` | Expand Raises block with kernel-size context | -| 9 | `mixed_fftconv1d_fp32_bhl` | Add performance hint to `use_phase_shift` | -| 10 | All `_w_reshape` wrappers | Fix kernel shape to `[1\|B, K, H]` | -| 11 | `_DEFAULT_MIXED_CHUNK_SIZE` | Add comment with memory-savings context | -| 12 | `mixed_fftconv2d_fp32_bhl` | Add runnable `Example:` block | diff --git a/nvsubquadratic/ops/mixed_fftconv.py b/nvsubquadratic/ops/mixed_fftconv.py index c6e13a29..ef57d3ae 100644 --- a/nvsubquadratic/ops/mixed_fftconv.py +++ b/nvsubquadratic/ops/mixed_fftconv.py @@ -47,6 +47,12 @@ at runtime (zero overhead), so it is safe to use this module as a single entry point even when ``periodic`` happens to be uniform. +Note: ``CKConvND`` accepts ``fft_padding=["circular", "zero"]`` (string form) +and internally converts this to the boolean ``periodic`` tuple via +``_resolve_periodic`` before calling these ops. The bool-tuple is the +*internal* normalised form used here; YAML configs should always use the +string list form. + See ``docs/ops/MIXED_BC_PLAN.md`` for the per-axis algorithm, dataset motivation table, and deferred work (fp16, multi-head, per-face BCs). @@ -149,7 +155,8 @@ class _MixedPhaseRamp1DCache: where ``f`` ranges over the DFT frequencies for an FFT of length ``F_d`` and ``s_d`` is the integer pixel shift. Non-periodic axes use ``s_d = 0`` - and contribute no ramp. + and contribute no ramp — their "same" alignment is handled entirely by the + centered crop ``[K_d // 2, K_d // 2 + N_d)`` applied to the IFFT output. This cache stores the 1-D ramps keyed by ``(F, s, is_rfft_axis, device, dtype)`` in an ordered-dict LRU so that repeated forward passes with the @@ -229,9 +236,9 @@ def get( Args: F: FFT length along this axis (padded length for non-periodic, input length for periodic — caller decides which). - s: Integer pixel shift to apply via the ramp. Callers with - ``s == 0`` are expected to skip the multiply entirely; this - method still handles ``s == 0`` (returns all-ones). + s: Integer pixel shift to apply via the ramp. In practice + ``_build_nd_phase_ramp`` skips axes where ``s == 0`` before + calling this method, so ``s == 0`` is not an expected input. is_rfft_axis: If True, this is the last spatial axis where the rfft is taken; a ramp of length ``F // 2 + 1`` is built using :func:`torch.fft.rfftfreq`. Otherwise a full ramp of @@ -288,11 +295,14 @@ def _mixed_recipe( - If ``periodic[d]`` is False (zero-padded "same"): ``F_d = min(N_d + (K_d + 1) // 2, 2 * N_d)``, crop ``[K_d // 2, K_d // 2 + N_d)``, shift ``s_d = 0`` (alignment is - handled by the centered crop). + handled by the centered crop). For even ``K_d``, floor division gives + a left-biased crop (one extra sample on the right), matching the + convention of ``torch.nn.ConvNd(padding='same')``. Tests verify this + against the spatial reference. Args: spatial: Input spatial dims ``(N_0, ..., N_{D-1})``. - kshape: Kernel spatial dims ``(K_0, ..., K_{D-1})``. + kshape: Kernel spatial dims ``(K_0, ..., K_{D-1})``. periodic: Per-axis periodicity flags. Returns: @@ -347,9 +357,10 @@ def _build_nd_phase_ramp( function returns ``None`` so callers can skip the multiply with a single branch. - The N-D ramp is **not** cached; only the 1-D per-axis ramps are. The - product is materialised here (via broadcasted multiplication) so that the - downstream multiply with ``fft_x`` is a single fused op. + The N-D ramp itself is not cached. The loop multiplies broadcast 1-D views + together, producing a compact tensor (not a full ``(F_0, ..., F_{D-1})`` + allocation) that covers only the periodic axes with non-zero shifts, so + that the downstream multiply with ``fft_x`` is a single fused op. Args: fft_shape: Per-axis FFT lengths ``(F_0, ..., F_{D-1})``, as returned @@ -501,8 +512,9 @@ def _mixed_fftconv_nd_fp32_bhl( 4. Optionally apply the N-D phase ramp to ``fft_k`` (if ``use_phase_shift=True``). 5. Multiply ``fft_x * fft_k`` and apply the inverse ``irfftn``. - 6. If ``use_phase_shift=False``, apply ``torch.roll`` on the periodic - axes to achieve "same" alignment. + 6. If ``use_phase_shift=False``, apply ``torch.roll`` on periodic axes + that have a non-zero shift (``s_d != 0``); size-1 periodic kernels + (shift 0) are skipped. 7. Crop the IFFT output to the input spatial shape via the per-axis crop windows from :func:`_mixed_recipe`. 8. Cast back to the original dtype of ``x`` and add the optional @@ -525,6 +537,11 @@ def _mixed_fftconv_nd_fp32_bhl( Raises: AssertionError: On shape mismatches or out-of-range kernel sizes. + Specifically: ``K_d > N_d`` on a periodic axis (the circular FFT + length is ``N_d``, so the kernel must fit within it); or + ``K_d > 2 * N_d`` on a non-periodic axis (the padded FFT length + is at most ``2 * N_d``, which accommodates the standard + "double-grid" SIREN kernel size of ``2 * N_d - 1``). """ x_shape = x.shape assert x.ndim == 2 + data_dim, f"Expected {2 + data_dim}D input, got {x.ndim}D" @@ -616,9 +633,13 @@ def mixed_fftconv1d_fp32_bhl( kernel: Kernel tensor of shape ``[1|B, H, K]`` (any dtype, cast to fp32). periodic: Length-1 sequence of bools. ``periodic[0] == True`` ⇒ circular. shortcut: Optional per-channel scale ``[H]`` added as ``y += shortcut * x``. - use_phase_shift: If True, align periodic axes via frequency-domain phase - ramps. If False, align via :func:`torch.roll` on periodic axes after - the inverse transform. The output is mathematically equivalent. + use_phase_shift: If True (default), align periodic axes via + frequency-domain phase ramps — the shift is fused into the + frequency-domain multiply with no extra data movement, making + this the faster path. If False, align via :func:`torch.roll` + on periodic axes after the inverse transform. The output is + mathematically equivalent; use ``False`` only as a reference or + when ``torch.compile`` cannot handle complex ops. Returns: Tensor of shape ``[B, H, L]`` in the original dtype of ``x``. @@ -648,6 +669,17 @@ def mixed_fftconv2d_fp32_bhl( Returns: Tensor of shape ``[B, H, X, Y]`` in the original dtype of ``x``. + + Example: + >>> import torch + >>> from nvsubquadratic.ops.mixed_fftconv import mixed_fftconv2d_fp32_bhl + >>> B, H, X, Y, Kx, Ky = 2, 64, 32, 64, 63, 127 + >>> x = torch.randn(B, H, X, Y) + >>> kernel = torch.randn(1, H, Kx, Ky) + >>> # x-axis periodic, y-axis zero-padded + >>> y = mixed_fftconv2d_fp32_bhl(x, kernel, periodic=(True, False)) + >>> y.shape + torch.Size([2, 64, 32, 64]) """ periodic_t = _normalize_periodic(periodic, data_dim=2) return _mixed_fftconv_nd_fp32_bhl(x, kernel, periodic_t, shortcut, use_phase_shift, data_dim=2) @@ -700,7 +732,8 @@ def mixed_fftconv1d_fp32_bhl_w_reshape( Args: x: Input tensor of shape ``[B, L, H]`` (BLH, channels-last). - kernel: Kernel tensor of shape ``[B, K, H]`` (BLH). + kernel: Kernel tensor of shape ``[1|B, K, H]`` (BLH). Leading dim 1 + for a shared kernel, ``B`` for per-sample kernels. periodic: Length-1 sequence of bools. shortcut: Optional per-channel scale ``[H]``. use_phase_shift: See :func:`mixed_fftconv1d_fp32_bhl`. @@ -729,7 +762,8 @@ def mixed_fftconv2d_fp32_bhl_w_reshape( Args: x: Input tensor of shape ``[B, X, Y, H]`` (BLH, channels-last). - kernel: Kernel tensor of shape ``[B, K_x, K_y, H]`` (BLH). + kernel: Kernel tensor of shape ``[1|B, K_x, K_y, H]`` (BLH). Leading + dim 1 for a shared kernel, ``B`` for per-sample kernels. periodic: Length-2 sequence ``(periodic_x, periodic_y)``. shortcut: Optional per-channel scale ``[H]``. use_phase_shift: See :func:`mixed_fftconv1d_fp32_bhl`. @@ -757,7 +791,8 @@ def mixed_fftconv3d_fp32_bhl_w_reshape( Args: x: Input tensor of shape ``[B, X, Y, Z, H]`` (BLH, channels-last). - kernel: Kernel tensor of shape ``[B, K_x, K_y, K_z, H]`` (BLH). + kernel: Kernel tensor of shape ``[1|B, K_x, K_y, K_z, H]`` (BLH). + Leading dim 1 for a shared kernel, ``B`` for per-sample kernels. periodic: Length-3 sequence ``(periodic_x, periodic_y, periodic_z)``. shortcut: Optional per-channel scale ``[H]``. use_phase_shift: See :func:`mixed_fftconv1d_fp32_bhl`. @@ -776,6 +811,10 @@ def mixed_fftconv3d_fp32_bhl_w_reshape( # ============================================================================= +# Default channel chunk size for the memory-efficient variants. +# A chunk of 128 channels typically gives ~26% peak-memory savings with ~11% +# throughput overhead relative to the non-chunked path (measured on H100; see +# fftconv_chunked.py for profiling details). _DEFAULT_MIXED_CHUNK_SIZE = 128 @@ -951,7 +990,8 @@ def mixed_fftconv1d_fp32_bhl_w_reshape_chunked( Args: x: Input tensor of shape ``[B, L, H]`` (BLH, channels-last). - kernel: Kernel tensor of shape ``[B, K, H]`` (BLH). + kernel: Kernel tensor of shape ``[1|B, K, H]`` (BLH). Leading dim 1 + for a shared kernel, ``B`` for per-sample kernels. periodic: Length-1 sequence of bools. shortcut: Optional per-channel scale ``[H]``. use_phase_shift: See :func:`mixed_fftconv1d_fp32_bhl`. @@ -987,7 +1027,8 @@ def mixed_fftconv2d_fp32_bhl_w_reshape_chunked( Args: x: Input tensor of shape ``[B, X, Y, H]`` (BLH, channels-last). - kernel: Kernel tensor of shape ``[B, K_x, K_y, H]`` (BLH). + kernel: Kernel tensor of shape ``[1|B, K_x, K_y, H]`` (BLH). Leading + dim 1 for a shared kernel, ``B`` for per-sample kernels. periodic: Length-2 sequence ``(periodic_x, periodic_y)``. shortcut: Optional per-channel scale ``[H]``. use_phase_shift: See :func:`mixed_fftconv1d_fp32_bhl`. @@ -1023,7 +1064,8 @@ def mixed_fftconv3d_fp32_bhl_w_reshape_chunked( Args: x: Input tensor of shape ``[B, X, Y, Z, H]`` (BLH, channels-last). - kernel: Kernel tensor of shape ``[B, K_x, K_y, K_z, H]`` (BLH). + kernel: Kernel tensor of shape ``[1|B, K_x, K_y, K_z, H]`` (BLH). + Leading dim 1 for a shared kernel, ``B`` for per-sample kernels. periodic: Length-3 sequence ``(periodic_x, periodic_y, periodic_z)``. shortcut: Optional per-channel scale ``[H]``. use_phase_shift: See :func:`mixed_fftconv1d_fp32_bhl`. From 9edd00e8e592b8668b9270b265816fccbf0069e4 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:15:04 +0200 Subject: [PATCH 11/72] docs(integrate/kernels_nd): apply reviewer feedback --- docs/reviews/kernels_nd_review.md | 282 --------------------------- nvsubquadratic/modules/kernels_nd.py | 194 +++++++++++++++--- 2 files changed, 163 insertions(+), 313 deletions(-) delete mode 100644 docs/reviews/kernels_nd_review.md diff --git a/docs/reviews/kernels_nd_review.md b/docs/reviews/kernels_nd_review.md deleted file mode 100644 index 1debef0d..00000000 --- a/docs/reviews/kernels_nd_review.md +++ /dev/null @@ -1,282 +0,0 @@ -# Reviewer feedback: `nvsubquadratic/modules/kernels_nd.py` - -Second-pass review after Phase 1 docstring additions. Issues are numbered; each -quotes the offending text and says exactly what to fix. - -______________________________________________________________________ - -## 1. Module docstring: `RandomFourierKernelND` is not listed in the "Kernel classes" table - -**Location**: module-level docstring, "Kernel classes in this module" section. - -The table lists six consumer-facing kernel classes but omits -`RandomFourierKernelND` — the table starts with `SIRENKernelND`. The class is -described in the prose above, but a new reader scanning the table will not find -it. - -**Fix**: add `RandomFourierKernelND` as the first entry in the list, before -`SIRENKernelND`. - -______________________________________________________________________ - -## 2. `_normalize_l_cache`: missing `Raises` block - -**Location**: `_normalize_l_cache` docstring. - -The function raises `TypeError` (for bool or non-int/sequence input) and -`ValueError` (for wrong sequence length or values \< 2), but the docstring has -no `Raises:` section. An external collaborator reading this to decide whether -to guard against exceptions will not know what to catch. - -**Fix**: add a `Raises:` section listing both `TypeError` and `ValueError` -with the conditions that trigger each. - -______________________________________________________________________ - -## 3. `RandomFourierPositionalEmbeddingND.forward`: wrong shape annotation for grid return value - -**Location**: `forward` docstring, `Returns` block: - -> `torch.Tensor: The input positions normalized between [-1, 1] (shape: [1, * spatial_dims, 1]).` - -The trailing dimension is `data_dim` (the number of spatial axes), not `1`. -A 2D model has a grid of shape `[1, H, W, 2]`, not `[1, H, W, 1]`. The -`_build_grid_cache` code confirms: `dim=-1` of the stacked meshgrid has size -`data_dim`. - -**Fix**: change the shape annotation to `[1, *spatial_dims, data_dim]`. - -The same bug appears in `SIRENPositionalEmbeddingND.forward` — fix it there -too. - -______________________________________________________________________ - -## 4. `RandomFourierPositionalEmbeddingND.forward` and `SIRENPositionalEmbeddingND.forward`: stale `"concatenated sine and cosine values"` description for SIREN - -**Location**: `SIRENPositionalEmbeddingND.forward`, `Returns` block: - -> `torch.Tensor: The positional embeddings, concatenated sine and cosine values (shape: [1, * spatial_dims, embedding_dim]).` - -The SIREN embedding is `sin(Wx+b)` — a single sine, no cosine concatenation. -The description was copied verbatim from the RFF embedding, where `[cos, sin]` -concatenation is correct. - -**Fix**: change the SIREN `forward` return description to: -"The positional embeddings, `sin(W x + b)`, shape `[1, *spatial_dims, embedding_dim]`." - -______________________________________________________________________ - -## 5. `RandomFourierKernelND.forward`: `Returns` annotation says "tuple" but the function actually returns two tensors described as a single `tuple[torch.Tensor, torch.Tensor]` - -**Location**: `RandomFourierKernelND.forward`, `Returns` block: - -> `tuple[torch.Tensor, torch.Tensor]: The computed random Fourier kernel and the corresponding grid values. The kernel is a tensor of shape (1, * spatial_dims, out_dim) The grid is a tensor of shape (1, * spatial_dims, data_dim)` - -The format mixes the inline `tuple[…]` type-hint style with a multi-line -description, making it hard to parse. The return type annotation in the -`Returns:` block should follow Google style: one sub-item per returned object. - -**Fix**: rewrite as: - -``` -Returns: - tuple: - - torch.Tensor: Kernel values of shape ``[1, *spatial_dims, out_dim]``. - - torch.Tensor: Coordinate grid of shape ``[1, *spatial_dims, data_dim]``. -``` - -______________________________________________________________________ - -## 6. `SIRENKernelND` class docstring: `flop_count` method has no standalone docstring header describing its purpose in one sentence - -**Location**: `SIRENKernelND.flop_count`. - -The docstring jumps directly into the "At `inference=True`…" explanation -without a one-sentence summary. Google style requires a short summary line -first. - -**Fix**: add as the first line: - -> "Return an integer FLOP estimate for one kernel generation forward pass." - -______________________________________________________________________ - -## 7. `SIRENKernelND.flop_count`: `grid_lens` description is ambiguous - -**Location**: `flop_count` docstring, `Args:` block: - -> ``` grid_lens: Spatial extents passed to the positional embedding. The kernel grid has size ``(2 * L - 1)`` per dimension. ``` - -This does not say what unit `grid_lens` is in (seq_lens? cache extents?). The -code uses `for L in grid_lens: G *= 2 * L - 1`, so `L` here is the *output* -sequence length (each axis), not the cache extent. The description "Spatial -extents passed to the positional embedding" is correct but the parenthetical -`(2 * L - 1)` makes it look like `grid_lens` are the half-lengths. - -**Fix**: replace with: - -> ``` grid_lens: Per-axis output sequence lengths, i.e. the same tuple you would pass to ``forward`` as ``seq_lens``. The FLOP count uses the grid volume ``G = prod(2*L - 1 for L in grid_lens)`` as the number of coordinate points the MLP processes. ``` - -______________________________________________________________________ - -## 8. `_build_grid_cache`: missing Args/Returns in both `RandomFourierPositionalEmbeddingND` and `SIRENPositionalEmbeddingND` copies - -**Location**: `RandomFourierPositionalEmbeddingND._build_grid_cache` and -`SIRENPositionalEmbeddingND._build_grid_cache`. - -Both static methods have a one-paragraph description but no `Args:` or -`Returns:` sections. The parameters `L_per_axis`, `max_limits`, and `device` -are not described, and the return shape is not stated. - -**Fix**: add: - -``` -Args: - L_per_axis: Per-axis cache extents. The cache size along axis ``i`` - is ``2 * L_per_axis[i] - 1`` grid points. - max_limits: Per-axis coordinate limits. Each axis spans - ``[-max_limits[i], +max_limits[i]]``. Defaults to ``1.0`` on all - axes. - device: Target device for the returned tensor. Defaults to CPU. - -Returns: - Float32 tensor of shape ``[1, 2*L_0-1, ..., 2*L_{d-1}-1, data_dim]`` - representing the coordinate meshgrid, with a leading batch dimension of 1. -``` - -(The two copies are byte-identical; apply the same fix to both.) - -______________________________________________________________________ - -## 9. `_maybe_extend_grid_cache`: missing Args/Returns in both copies - -**Location**: `RandomFourierPositionalEmbeddingND._maybe_extend_grid_cache` and -`SIRENPositionalEmbeddingND._maybe_extend_grid_cache`. - -Same issue as item 8: no `Args:` or `Returns:` sections. - -**Fix**: add: - -``` -Args: - seq_lens: Requested per-axis output sequence lengths. Any axis where - ``seq_lens[i] > self.L_cache_per_axis[i]`` triggers a cache - extension. - -Returns: - None. Modifies ``self.grid_cache`` and ``self.L_cache_per_axis`` - in-place when an extension is needed. -``` - -______________________________________________________________________ - -## 10. `LearnableOmegaSIRENPositionalEmbeddingND._clamp_omega_scale_pre_hook`: missing Args/Returns - -**Location**: `_clamp_omega_scale_pre_hook`. - -The hook is registered as a forward pre-hook and takes `(module, inputs)` but -the docstring has no `Args:` block, which is required for any public/protected -method. - -**Fix**: add: - -``` -Args: - module: The module instance (``self``); provided by PyTorch's hook - mechanism. - inputs: The positional inputs tuple; not used. - -Returns: - None. Modifies ``self.omega_0_scale.data`` in-place. -``` - -______________________________________________________________________ - -## 11. `BlockDiagonalMultiOmegaSIRENKernelND._block_mask`: `out_dim` / `in_dim` shape contract not stated precisely enough - -**Location**: `_block_mask` docstring: - -> "`out_dim` and `in_dim` must both be divisible by `num_blocks`." - -This constraint is documented but neither a `Raises:` nor a note says what -happens if violated (currently it silently produces wrong-sized blocks because -integer division truncates). - -**Fix**: add a `Raises:` section: - -``` -Raises: - ZeroDivisionError: If ``num_blocks == 0``. - Note: if ``out_dim`` or ``in_dim`` is not divisible by ``num_blocks``, - the block sizes are silently truncated via integer division; the - caller (``BlockDiagonalMultiOmegaSIRENKernelND.__init__``) enforces - divisibility before calling this method. -``` - -and add `Args:` / `Returns:` sections (currently absent): - -``` -Args: - out_dim: Number of output features (rows of the weight matrix). - in_dim: Number of input features (columns of the weight matrix). - num_blocks: Number of blocks. - off_block_scale: Scalar value for off-diagonal block entries. - device: Target device. - dtype: Target dtype. - -Returns: - Float tensor of shape ``[out_dim, in_dim]`` with ``1.0`` on the - block diagonal and ``off_block_scale`` elsewhere. -``` - -______________________________________________________________________ - -## 12. `SIRENKernelND` class docstring: "Wang initialisation" is unexplained - -**Location**: `SIRENKernelND` class docstring, "Initialisation" section: - -> "The output layer applies **Wang initialisation**: weights are scaled by `sqrt(1 / kernel_volume)`…" - -An external collaborator unfamiliar with this term will not know where it comes -from. The SIREN paper (Sitzmann et al. 2020) does not use the name "Wang init" -— this is a practice from continuous convolution literature. - -**Fix**: add a brief parenthetical reference, e.g.: - -> "(introduced for implicit kernel networks by \[Wang et al.\]; equivalent to dividing the output layer's weights by the square root of the total number of kernel grid points, so that the initial filter energy is independent of grid size)" - -or cite the CKConv paper where this practice originates. - -______________________________________________________________________ - -## 13. `LearnableOmegaSIRENPositionalEmbeddingND` class docstring: `apply_lr_scale` mentions `_build_param_groups` without context - -**Location**: class docstring, `Args:` for `apply_lr_scale`: - -> "attach `_lr_scale = 1/(2π·omega_0)` to `self.linear.weight` so that `_build_param_groups` lowers its learning rate…" - -`_build_param_groups` is an internal optimizer helper not visible from this -file. A reader who does not know the codebase's optimizer infrastructure will -not understand how `_lr_scale` is consumed. - -**Fix**: replace with: - -> "When True, attach `_lr_scale = 1/(2*pi*omega_0)` to `self.linear.weight`. The optimizer utility `_build_param_groups` (in `experiments/`) reads this attribute and multiplies the layer's learning rate by `_lr_scale`, compensating for the missing `2*pi*omega_0` factor in the SIREN-1 init bound." - -______________________________________________________________________ - -## 14. Module docstring: no mention of where the `conditioning` tensor comes from - -**Location**: module docstring, no mention of FiLM conditioning flow. - -The module docstring describes the two kernel families (RFF and SIREN) but -does not mention that `SIRENKernelND` and its descendants support a -`conditioning` argument that makes the kernel input-dependent. This is the -FiLM path, and a new collaborator reading the module overview has no idea this -feature exists or how to trigger it. - -**Fix**: add a short paragraph in the Overview section: - -> "Optionally, `SIRENKernelND` (and all SIREN subclasses) accept a `conditioning` tensor of shape `[B, C]` via their `forward` method. When a `KernelFiLMGenerator` is wired in via `film_cfg`, the kernel becomes input-dependent — each sample in the batch gets a different filter. This is used in diffusion models and other conditional generation tasks where the kernel must respond to a global context signal." - -______________________________________________________________________ diff --git a/nvsubquadratic/modules/kernels_nd.py b/nvsubquadratic/modules/kernels_nd.py index 86315465..bee08455 100644 --- a/nvsubquadratic/modules/kernels_nd.py +++ b/nvsubquadratic/modules/kernels_nd.py @@ -46,6 +46,23 @@ controls the number of discrete grid positions cached per axis, and the cache grows automatically at runtime whenever a larger input is encountered. +Input-dependent / conditional kernels +-------------------------------------- +``SIRENKernelND`` and all its subclasses accept an optional +``conditioning`` argument of shape ``[B, C]`` in their ``forward`` method. +When a ``KernelFiLMGenerator`` is wired in via the ``film_cfg`` constructor +argument, the generator maps this conditioning vector to a list of per-layer +``(gamma, beta)`` pairs that modulate the SIREN's hidden activations via +Feature-wise Linear Modulation (FiLM): + + h_i <- gamma_i * h_i + beta_i + +This makes the produced kernel batch-dependent: the output has shape +``[B, *spatial, out_dim]`` instead of the usual ``[1, *spatial, out_dim]``. +This feature is used in diffusion models and other conditional generation tasks +where each sample needs a different long-range filter. The ``conditioning`` +argument is ignored (no-op) when no ``film_cfg`` is provided. + Kernel classes in this module ------------------------------ ``RandomFourierKernelND`` @@ -106,6 +123,16 @@ def _normalize_l_cache(L_cache: int | Sequence[int], data_dim: int) -> tuple[int Returns: Tuple of length ``data_dim`` with one positive int per axis. + + Raises: + TypeError: If ``L_cache`` is a ``bool`` (which would silently cast to + ``0`` or ``1`` and fail the minimum-value check), or if it is + neither an ``int`` nor a sequence of ints. + ValueError: If ``L_cache`` is a sequence whose length differs from + ``data_dim``, or if any value in the resulting tuple is less + than 2 (required because the grid uses ``linspace(-1, 1, 2*L-1)`` + which needs at least 3 points, and the step ``1/(L-1)`` is + undefined at ``L=1``). """ if isinstance(L_cache, bool): raise TypeError("L_cache must be an int or a sequence of ints, got bool") @@ -256,6 +283,19 @@ def _build_grid_cache( Each axis spans ``[-max_limit_i, +max_limit_i]`` sampled at ``2 * L_i - 1`` points (default ``max_limit_i = 1.0`` at construction, possibly larger after a runtime extension to keep step size constant). + + Args: + L_per_axis: Per-axis cache extents. The number of grid points + along axis ``i`` is ``2 * L_per_axis[i] - 1``. + max_limits: Per-axis coordinate limits; axis ``i`` spans + ``[-max_limits[i], +max_limits[i]]``. Defaults to ``1.0`` + on all axes (the standard ``[-1, 1]`` normalised range). + device: Target device for the returned tensor. Defaults to CPU. + + Returns: + Float32 tensor of shape + ``[1, 2*L_0-1, ..., 2*L_{d-1}-1, data_dim]`` representing the + coordinate meshgrid, with a leading batch dimension of 1. """ if max_limits is None: max_limits = (1.0,) * len(L_per_axis) @@ -274,6 +314,15 @@ def _maybe_extend_grid_cache(self, seq_lens: tuple[int, ...]) -> None: ``max_limit = 1.0 + step_size * (seq_len - L_cache_orig)``. Axes that are not extended are rebuilt at their existing extent so the cache remains a single rectangular tensor. + + Args: + seq_lens: Requested per-axis output sequence lengths. Any axis + where ``seq_lens[i] > self.L_cache_per_axis[i]`` triggers a + cache extension for that axis. + + Returns: + None. Modifies ``self.grid_cache`` and + ``self.L_cache_per_axis`` in-place when an extension is needed. """ if all(L >= sl for L, sl in zip(self.L_cache_per_axis, seq_lens)): return @@ -291,19 +340,24 @@ def _maybe_extend_grid_cache(self, seq_lens: tuple[int, ...]) -> None: self.L_cache_per_axis = new_L_per_axis def forward(self, seq_lens: tuple[int, ...]) -> tuple[torch.Tensor, torch.Tensor]: - """Computes the positional embeddings for a sequence of the given length. + """Compute the RFF positional embeddings for a given spatial grid. Args: - seq_lens (tuple[int, ...]): Lengths of the input grid for which to compute the positional embeddings. + seq_lens: Per-axis output sequence lengths. Length must equal + ``self.data_dim``. For example, for a 2D signal of height H + and width W, pass ``(H, W)``. Returns: tuple: - - torch.Tensor: The positional embeddings, concatenated sine and cosine values (shape: [1, * spatial_dims, embedding_dim]). - - torch.Tensor: The input positions normalized between [-1, 1] (shape: [1, * spatial_dims, 1]). + - torch.Tensor: The positional embeddings, + ``[cos(Wx+b), sin(Wx+b)]`` concatenated along the last axis. + Shape ``[1, *spatial_dims, embedding_dim]``. + - torch.Tensor: The coordinate grid of positions normalised to + ``[-1, 1]`` per axis. Shape ``[1, *spatial_dims, data_dim]``. Raises: - AssertionError: If `seq_lens` is not of length `self.data_dim`. - AssertionError: If `self.grid_cache` is not of type `torch.float32`. + AssertionError: If ``len(seq_lens) != self.data_dim``. + AssertionError: If ``self.grid_cache`` is not ``float32``. """ # Check that the sequence lengths are of the correct length. assert len(seq_lens) == self.data_dim, ( @@ -478,16 +532,21 @@ def __init__( self.out_linear.weight.data *= math.sqrt(1.0 / kernel_volume) def forward(self, seq_lens: tuple[int, ...], conditioning: torch.Tensor | None = None) -> torch.Tensor: - """Computes the random Fourier kernel for a given grid of spatial dimensions. + """Compute the RFF kernel for a given grid of spatial dimensions. Args: - seq_lens (tuple[int, ...]): Lengths of the input grid for which to compute the positional embeddings. - conditioning: Unused. Accepted for API compatibility with FiLM-enabled kernels. + seq_lens: Per-axis output sequence lengths. Length must equal + ``self.data_dim``. + conditioning: Unused. Accepted for API compatibility with + FiLM-enabled kernels (e.g. ``SIRENKernelND`` with + ``film_cfg``). Returns: - tuple[torch.Tensor, torch.Tensor]: The computed random Fourier kernel and the corresponding grid values. - The kernel is a tensor of shape (1, * spatial_dims, out_dim) - The grid is a tensor of shape (1, * spatial_dims, data_dim) + tuple: + - torch.Tensor: Kernel values of shape + ``[1, *spatial_dims, out_dim]``. + - torch.Tensor: Coordinate grid of shape + ``[1, *spatial_dims, data_dim]``. """ # Generate positional embeddings and corresponding grid values pos_emb, grid = self.positional_embedding(seq_lens) @@ -688,6 +747,19 @@ def _build_grid_cache( ``2 * L_i - 1`` points. At construction every ``max_limit_i`` is ``1.0``; runtime extensions can pass per-axis values larger than 1 to preserve the original step size on extended axes. + + Args: + L_per_axis: Per-axis cache extents. The number of grid points + along axis ``i`` is ``2 * L_per_axis[i] - 1``. + max_limits: Per-axis coordinate limits; axis ``i`` spans + ``[-max_limits[i], +max_limits[i]]``. Defaults to ``1.0`` + on all axes. + device: Target device for the returned tensor. Defaults to CPU. + + Returns: + Float32 tensor of shape + ``[1, 2*L_0-1, ..., 2*L_{d-1}-1, data_dim]`` representing the + coordinate meshgrid, with a leading batch dimension of 1. """ if max_limits is None: max_limits = (1.0,) * len(L_per_axis) @@ -706,6 +778,15 @@ def _maybe_extend_grid_cache(self, seq_lens: tuple[int, ...]) -> None: ``max_limit = 1.0 + step_size * (seq_len - L_cache_orig)``. Axes that already cover their requested ``seq_len`` are rebuilt at their existing extent so the cache stays a single rectangular tensor. + + Args: + seq_lens: Requested per-axis output sequence lengths. Any axis + where ``seq_lens[i] > self.L_cache_per_axis[i]`` triggers a + cache extension for that axis. + + Returns: + None. Modifies ``self.grid_cache`` and + ``self.L_cache_per_axis`` in-place when an extension is needed. """ if all(L >= sl for L, sl in zip(self.L_cache_per_axis, seq_lens)): return @@ -723,19 +804,25 @@ def _maybe_extend_grid_cache(self, seq_lens: tuple[int, ...]) -> None: self.L_cache_per_axis = new_L_per_axis def forward(self, seq_lens: tuple[int, ...]) -> tuple[torch.Tensor, torch.Tensor]: - """Computes the positional embeddings for a sequence of the given length. + """Compute the SIREN positional embeddings for a given spatial grid. Args: - seq_lens (tuple[int, ...]): Lengths of the input grid for which to compute the positional embeddings. + seq_lens: Per-axis output sequence lengths. Length must equal + ``self.data_dim``. For example, for a 2D signal of height H + and width W, pass ``(H, W)``. Returns: tuple: - - torch.Tensor: The positional embeddings, concatenated sine and cosine values (shape: [1, * spatial_dims, embedding_dim]). - - torch.Tensor: The input positions normalized between [-1, 1] (shape: [1, * spatial_dims, 1]). + - torch.Tensor: The positional embeddings ``sin(W x + b)``, + where the linear projection is computed in float32 and the + result is cast back to the weight dtype. + Shape ``[1, *spatial_dims, embedding_dim]``. + - torch.Tensor: The coordinate grid of positions normalised to + ``[-1, 1]`` per axis. Shape ``[1, *spatial_dims, data_dim]``. Raises: - AssertionError: If `seq_lens` is not of length `self.data_dim`. - AssertionError: If `self.grid_cache` is not of type `torch.float32`. + AssertionError: If ``len(seq_lens) != self.data_dim``. + AssertionError: If ``self.grid_cache`` is not ``float32``. """ # Check that the sequence lengths are of the correct length. assert len(seq_lens) == self.data_dim, ( @@ -831,6 +918,10 @@ class SIRENKernelND(torch.nn.Module): * All ``hidden_linears`` are SIREN-initialised with ``hidden_omega_0``. * ``out_linear`` is SIREN-initialised with ``hidden_omega_0``, then additionally Wang-scaled by ``sqrt(1 / prod(L_cache_per_axis))``. + This "Wang init" (from the CKConv paper, Romero et al. 2021) divides + the output layer's weights by the square root of the total grid volume + (``L_cache**data_dim`` for isotropic grids), so the initial filter's + L2 energy is independent of the grid resolution. * Hidden linear weights and output bias get ``_no_weight_decay = True`` so that weight-decay optimizers do not destroy the SIREN spectrum. @@ -969,7 +1060,7 @@ def __init__( self.film_generator = None def flop_count(self, grid_lens: tuple[int, ...], inference: bool = False) -> int: - """Count FLOPs for SIREN kernel generation on the positional grid. + """Return an integer FLOP estimate for one kernel generation forward pass. At ``inference=True`` with no FiLM generator, returns 0 because the kernel is input-independent and can be precomputed once and cached. @@ -1003,8 +1094,10 @@ def flop_count(self, grid_lens: tuple[int, ...], inference: bool = False) -> int (Note: film_after_pos_embed requires embedding_dim == mlp_hidden_dim.) Args: - grid_lens: Spatial extents passed to the positional embedding. - The kernel grid has size ``(2 * L - 1)`` per dimension. + grid_lens: Per-axis output sequence lengths — the same tuple you + would pass to ``forward`` as ``seq_lens``. The total number + of coordinate points the MLP processes is + ``G = prod(2*L - 1 for L in grid_lens)``. inference: If True and no FiLM generator, return 0 (cacheable kernel). Returns: @@ -1527,8 +1620,29 @@ def _block_mask( """Build a block-mask: 1.0 on the diagonal blocks, ``off_block_scale`` elsewhere. The returned tensor has shape ``[out_dim, in_dim]``. Block sizes are - ``out_dim // num_blocks`` rows by ``in_dim // num_blocks`` columns; - ``out_dim`` and ``in_dim`` must both be divisible by ``num_blocks``. + ``out_dim // num_blocks`` rows by ``in_dim // num_blocks`` columns. + + Note: if ``out_dim`` or ``in_dim`` is not divisible by ``num_blocks``, + block sizes are silently truncated via integer division. The calling + ``__init__`` validates divisibility and raises ``ValueError`` before + reaching this method. + + Args: + out_dim: Number of output features (rows of the weight matrix). + in_dim: Number of input features (columns of the weight matrix). + num_blocks: Number of equal-sized blocks along both axes. + off_block_scale: Scalar fill value for off-diagonal block entries. + Use ``0.0`` for a strict block-diagonal and ``1.0`` to leave + the weights unmodified. + device: Target device for the mask tensor. + dtype: Target dtype for the mask tensor (should match the weight). + + Returns: + Tensor of shape ``[out_dim, in_dim]`` with ``1.0`` on the block + diagonal and ``off_block_scale`` elsewhere. + + Raises: + ZeroDivisionError: If ``num_blocks == 0``. """ rows_per_block = out_dim // num_blocks cols_per_block = in_dim // num_blocks @@ -1604,13 +1718,14 @@ class LearnableOmegaSIRENPositionalEmbeddingND(SIRENPositionalEmbeddingND): omega_0_scale_max: Upper clamp on ``ω₀_scale``. Default 2.0, giving a total multiplier inside the sine of up to ``4π·ω₀``. use_bias: Whether to include a bias term. - apply_lr_scale: When True, attach ``_lr_scale = 1/(2π·omega_0)`` to - ``self.linear.weight`` so that ``_build_param_groups`` lowers - its learning rate to compensate for the absent ``2π·ω₀`` factor - in the init bound. Default False (opt-in; off by default so - the optimizer behaviour matches the existing SIREN exactly, - and the new classes can be A/B-tested with and without LR - compensation). + apply_lr_scale: When True, attach ``_lr_scale = 1/(2*pi*omega_0)`` + to ``self.linear.weight``. The optimizer utility + ``_build_param_groups`` (in ``experiments/``) reads this attribute + and multiplies the layer's effective learning rate by + ``_lr_scale``, compensating for the missing ``2*pi*omega_0`` + factor in the SIREN-1 init bound so that the per-step update size + matches a standard SIREN. Default False (opt-in, so existing + runs are unaffected and the new classes can be A/B-tested). Attributes: omega_0 (float): Constant part of the runtime multiplier (same as the @@ -1727,7 +1842,24 @@ def __init__( self.linear.bias._lr_scale = 1.0 / (2.0 * math.pi * float(omega_0)) def _clamp_omega_scale_pre_hook(self, module, inputs): - """Clamp ``omega_0_scale`` into ``[scale_min, scale_max]`` in-place before forward.""" + """Clamp ``omega_0_scale`` into ``[scale_min, scale_max]`` in-place before forward. + + Registered as a ``register_forward_pre_hook`` so the clamp runs + automatically at the start of every forward call without requiring the + caller to call it explicitly. This is a "direct" parametrisation — the + clamping does not block gradient flow for values already inside the + valid range. + + Args: + module: The module instance (``self``); provided by PyTorch's hook + mechanism and not used directly. + inputs: The positional inputs tuple passed to ``forward``; not + inspected or modified. + + Returns: + None. Modifies ``self.omega_0_scale.data`` in-place via + ``clamp_``. + """ with torch.no_grad(): self.omega_0_scale.data.clamp_(min=self.omega_0_scale_min, max=self.omega_0_scale_max) From 88d5b814e73c50cda09751c2d9f8b0a0530075ff Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:16:15 +0200 Subject: [PATCH 12/72] docs(tracker): mark mixed_fftconv, hyena_nd, kernels_nd, sequence_mixer as done --- docs-tracker.md | 97 +++++++++++++++++++++++++------------------------ 1 file changed, 49 insertions(+), 48 deletions(-) diff --git a/docs-tracker.md b/docs-tracker.md index 7170c300..9c0c44d2 100644 --- a/docs-tracker.md +++ b/docs-tracker.md @@ -25,60 +25,61 @@ Work bottom-up: primitive ops → modules → networks → experiments. ### `nvsubquadratic/ops/` — FFT convolution primitives -| File | Status | Notes | -| -------------------------- | ------ | --------------------------------------------------------------------------------- | -| `README.md` | \[x\] | Folder overview, decision tree, math primer (new file) | -| `fftconv.py` | \[x\] | Module + key per-fn docstrings rewritten with math | -| `circular_fftconv.py` | \[x\] | Already strong; left as-is | -| `circular_fftconv_fp16.py` | \[x\] | Already strong; relies on FP16_FFTCONV_DERIVATION.md | -| `fftconv_fp16.py` | \[x\] | Already adequate; left as-is | -| `fftconv_multihead.py` | \[x\] | Module docstring expanded with multi-head/low-rank math | -| `fftconv_chunked.py` | \[x\] | Already strong; left as-is | -| `fftconv_custom.py` | \[x\] | Module docstring expanded with motivation; 1D causal wrappers added in 1D PR | -| `causal_conv1d_custom.py` | \[x\] | New (1D PR): thin wrappers for direct `causal_conv1d` + fused `b2b_causal_conv1d` | +| File | Status | Notes | +| -------------------------- | ------ | --------------------------------------------------------------------------------------- | +| `README.md` | \[x\] | Folder overview, decision tree, math primer (new file) | +| `fftconv.py` | \[x\] | Module + key per-fn docstrings rewritten with math | +| `circular_fftconv.py` | \[x\] | Already strong; left as-is | +| `circular_fftconv_fp16.py` | \[x\] | Already strong; relies on FP16_FFTCONV_DERIVATION.md | +| `fftconv_fp16.py` | \[x\] | Already adequate; left as-is | +| `fftconv_multihead.py` | \[x\] | Module docstring expanded with multi-head/low-rank math | +| `fftconv_chunked.py` | \[x\] | Already strong; left as-is | +| `fftconv_custom.py` | \[x\] | Module docstring expanded with motivation; 1D causal wrappers added in 1D PR | +| `causal_conv1d_custom.py` | \[x\] | New (1D PR): thin wrappers for direct `causal_conv1d` + fused `b2b_causal_conv1d` | +| `mixed_fftconv.py` | \[x\] | New (#120): per-axis mixed boundary-condition FFT conv; see `docs/ops/MIXED_BC_PLAN.md` | ### `nvsubquadratic/modules/` — Building blocks -| File | Status | Notes | -| ---------------------------------- | ------ | ------------------------------------------------------------------------------------- | -| `kernels_nd.py` | \[ \] | Learned kernel parametrisation | -| `hyena_nd.py` | \[ \] | Hyena operator (ND) — key paper component | -| `ckconv_nd.py` | \[ \] | CKConv (ND) | -| `ckconv_multihead_nd.py` | \[ \] | Multi-head CKConv | -| `mamba_nd.py` | \[ \] | Mamba SSM (ND) | -| `attention.py` | \[ \] | Standard attention | -| `vit5_attention.py` | \[ \] | ViT5 attention variant | -| `vit5_hyena_adapter.py` | \[ \] | Hyena adapter for ViT5 | -| `sequence_mixer.py` | \[ \] | Mixer abstraction | -| `condition_mixer.py` | \[ \] | Conditioning mixer | -| `residual_block.py` | \[ \] | Residual block | -| `vit5_residual_block.py` | \[ \] | ViT5 residual block | -| `patchify.py` | \[ \] | Patch embedding | -| `position_encoding.py` | \[ \] | Position encodings | -| `masks_nd.py` | \[ \] | ND masking utils | -| `mlp.py` | \[ \] | MLP block | -| `film.py` | \[ \] | FiLM conditioning | -| `grn.py` | \[ \] | GRN normalisation | -| `layer_scale.py` | \[ \] | LayerScale | -| `rms_norm.py` | \[ \] | RMS normalisation | -| `rms_norm_channel_first.py` | \[ \] | Channel-first RMS norm | -| `drop_path.py` | \[ \] | Stochastic depth | -| `causal_conv1d.py` | \[ \] | Causal 1D conv | -| `subq_ops_causal_conv1d.py` | \[x\] | New (1D PR): `nn.Conv1d`-compatible depthwise wrapper around `subq_ops.causal_conv1d` | -| `schedulers.py` | \[ \] | LR schedulers | -| `distributed_depthwise_conv_nd.py` | \[ \] | Distributed depthwise conv | +| File | Status | Notes | +| ---------------------------------- | ------ | ------------------------------------------------------------------------------------------- | +| `kernels_nd.py` | \[x\] | Learned kernel parametrisation — RFF + SIREN, FiLM-conditioned variants | +| `hyena_nd.py` | \[x\] | Hyena operator (ND) — two-gate sandwich, AllToAll CP, BC-aware convolution | +| `ckconv_nd.py` | \[ \] | CKConv (ND) | +| `ckconv_multihead_nd.py` | \[ \] | Multi-head CKConv | +| `mamba_nd.py` | \[ \] | Mamba SSM (ND) | +| `attention.py` | \[ \] | Standard attention | +| `vit5_attention.py` | \[ \] | ViT5 attention variant | +| `vit5_hyena_adapter.py` | \[ \] | Hyena adapter for ViT5 | +| `sequence_mixer.py` | \[x\] | Operator-agnostic dispatch layer (Hyena / Attention / CKConv / Mamba) | +| `condition_mixer.py` | \[ \] | Conditioning mixer | +| `residual_block.py` | \[ \] | Residual block | +| `vit5_residual_block.py` | \[ \] | ViT5 residual block | +| `patchify.py` | \[ \] | Patch embedding | +| `position_encoding.py` | \[ \] | Position encodings | +| `masks_nd.py` | \[ \] | ND masking utils | +| `mlp.py` | \[ \] | MLP block | +| `film.py` | \[ \] | FiLM conditioning | +| `grn.py` | \[ \] | GRN normalisation | +| `layer_scale.py` | \[ \] | LayerScale | +| `rms_norm.py` | \[ \] | RMS normalisation | +| `rms_norm_channel_first.py` | \[ \] | Channel-first RMS norm | +| `drop_path.py` | \[ \] | Stochastic depth | +| `causal_conv1d.py` | \[ \] | Causal 1D conv | +| `subq_ops_causal_conv1d.py` | \[x\] | New (1D PR): `nn.Conv1d`-compatible depthwise wrapper around `subq_ops.causal_conv1d` | +| `schedulers.py` | \[ \] | LR schedulers | +| `distributed_depthwise_conv_nd.py` | \[ \] | Distributed depthwise conv | +| `patch_merging.py` | \[ \] | Pending (feat/patch-merging PR): Swin-style 2×2 patch merging with register-row passthrough | ### `nvsubquadratic/networks/` — Full architectures -| File | Status | Notes | -| --------------------------- | ------ | ------------------------ | -| `general_purpose_resnet.py` | \[ \] | | -| `classification_resnet.py` | \[ \] | | -| `arc_resnet.py` | \[ \] | ARC task network | -| `arc_embedding.py` | \[ \] | ARC embedding | -| `vit5_classification.py` | \[ \] | ViT5 classification head | -| `huggingface_diffusers.py` | \[ \] | HF diffusers integration | -| `jit.py` / `jit_utils.py` | \[ \] | TorchScript utilities | +| File | Status | Notes | +| ------------------------------------- | ------ | ----------------------------------------------------------------------------------------------------------------------------------- | +| `general_purpose_resnet.py` | \[ \] | | +| `classification_resnet.py` | \[ \] | | +| `vit5_classification.py` | \[ \] | ViT5 classification head | +| `vit5_hierarchical_classification.py` | \[ \] | Pending (feat/patch-merging PR): Swin-style 4-stage hierarchical ViT-5 classifier with GAP readout and optional register-row layout | +| `huggingface_diffusers.py` | \[ \] | HF diffusers integration | +| `jit.py` / `jit_utils.py` | \[ \] | TorchScript utilities | ### `nvsubquadratic/parallel/` — Distributed primitives From 0259fbb151a74488dd3cf51f5904f9f4cb28436e Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:23:45 +0200 Subject: [PATCH 13/72] docs(write/patchify): add module and class docstrings --- nvsubquadratic/modules/patchify.py | 237 ++++++++++++++++++++++++----- 1 file changed, 195 insertions(+), 42 deletions(-) diff --git a/nvsubquadratic/modules/patchify.py b/nvsubquadratic/modules/patchify.py index 725d1483..c069fca3 100644 --- a/nvsubquadratic/modules/patchify.py +++ b/nvsubquadratic/modules/patchify.py @@ -1,8 +1,43 @@ # David W. Romero, 2025-09-09 -"""Patchify and Unpatchify layers as ConvND and ConvTransposeND layers. +"""Patch embedding and reconstruction layers for ND spatial signals. + +Overview +-------- +Patchification is the bridge between raw pixel space and the sequence of tokens +that a downstream mixer (Hyena, Attention, CKConv, …) operates on. Given an +input image (or any ND spatial signal) of shape ``[B, *spatial, C_in]``, the +layer divides the spatial axes into a grid of non-overlapping *patches* of size +``patch_size`` and linearly projects the flattened pixels within each patch into +a ``C_embed``-dimensional token embedding. The result is a spatially-ordered +grid of tokens ready for sequence mixing. + +Implementation approach +----------------------- +Both ``Patchify`` and ``Unpatchify`` are implemented as strided convolutions +(``Conv{1,2,3}d`` / ``ConvTranspose{1,2,3}d``) so that the unfold + linear +projection are fused into a single CUDA kernel. Setting +``kernel_size == stride == patch_size`` and ``padding == 0`` gives exactly the +ViT non-overlapping patch semantics. + +Output layout convention +------------------------ +All layers in this module use **channels-last** tensors externally:: + + input : [B, *spatial_dims, C_in] e.g. [B, H, W, C_in] for 2D + output : [B, *patch_grid, C_embed] e.g. [B, H/P, W/P, C_embed] + +Channels are reordered to channels-first only internally before the convolution +and back to channels-last before returning, which matches the layout expected by +``PositionEmbeddingND`` and the subsequent mixer blocks. + +Supported dimensionalities +-------------------------- +Both classes support ``data_dim ∈ {1, 2, 3}`` (time-series / images / volumes) +via the ``_CONV_CLASSES`` / ``_CONV_TRANSPOSE_CLASSES`` dispatch tables. + +Usage test:: -Usage test: PYTHONPATH=. python nvsubquadratic/modules/patchify.py """ @@ -28,13 +63,46 @@ class Patchify(torch.nn.Module): - """Conv-based image patchification (channels-last input). + """Conv-based patch embedding for ND spatial signals (channels-last I/O). - This mirrors the ViT/timm approach where a Conv with kernel_size=stride=patch_size - produces one embedding per patch location (non-overlapping patches). + Splits the spatial axes of the input into a regular grid of non-overlapping + patches and linearly projects each patch into an embedding vector. The + operation is equivalent to: - Input shape: [B, *spatial_dims, in_features] (channels-last, e.g., BHWC) - Output shape: [B, *spatial_dims // patch_size, out_features] + 1. Unfold every ``patch_size^data_dim`` pixel neighbourhood into a vector + of length ``C_in * patch_size^data_dim``. + 2. Apply a learned linear map from that vector to ``C_out`` dimensions. + + Because the unfold and linear projection can be fused into a single strided + convolution, this class simply wraps ``torch.nn.Conv{data_dim}d`` with + ``kernel_size = patch_size``, ``stride = stride``, and ``padding = 0``. + + **Output shape formula** (each spatial axis ``s`` independently):: + + out_s = floor((s - patch_size) / stride) + 1 + + For the default non-overlapping case (``stride == patch_size``) this + reduces to ``s / patch_size`` (assuming ``s`` is divisible by + ``patch_size``). + + **Layout convention** — inputs and outputs use *channels-last* ordering:: + + input : [B, *spatial_dims, C_in] (e.g. [B, H, W, C_in] for 2D) + output : [B, *patch_grid, C_out] (e.g. [B, H/P, W/P, C_out]) + + Internally, the tensor is transposed to channels-first before the Conv and + back to channels-last before returning, to match the layout expected by + ``PositionEmbeddingND`` and the mixer blocks. + + **Overlapping patches** — setting ``stride < patch_size`` produces + overlapping patches with the same formula above. This is less common in + ViT-style models but is supported. + + Attributes: + data_dim (int): Spatial dimensionality (1, 2, or 3). + patch_size (int): Receptive field size of each patch along every axis. + stride (int): Step between successive patch origins along every axis. + conv (torch.nn.Conv{data_dim}d): The underlying strided convolution. """ def __init__( @@ -46,16 +114,28 @@ def __init__( stride: int | None = None, bias: bool = True, ): - """Initialize the Patchify layer. + """Initialise the Patchify layer. Args: - in_features: The number of input channels. - out_features: The number of output channels (embedding dimension). - data_dim: The spatial dimensionality (1, 2, or 3). - patch_size: The size of each patch (kernel_size for the conv). - stride: The stride for the conv. Defaults to patch_size (non-overlapping). - bias: Whether the underlying conv has a learnable bias. Default True - for backward compatibility; set False for bias-free architectures. + in_features: Number of input channels ``C_in`` (e.g. 3 for RGB). + out_features: Embedding dimension ``C_out`` of each output token. + data_dim: Spatial dimensionality of the input signal. Must be 1 + (sequences), 2 (images), or 3 (volumes). + patch_size: Side length ``P`` of each patch. The convolution uses + ``kernel_size = patch_size`` along every spatial axis, so each + patch covers ``P^data_dim`` input pixels. + stride: Step size between consecutive patch origins along every + spatial axis. Defaults to ``patch_size``, giving + non-overlapping ViT-style patches. Set to a smaller value for + overlapping patches (denser token grids at the cost of more + tokens). + bias: If ``True`` (default), the projection conv includes a + learnable bias. Set to ``False`` for bias-free architectures + (e.g. when a subsequent normalisation layer makes bias + redundant). + + Raises: + ValueError: If ``data_dim`` is not 1, 2, or 3. """ super().__init__() if data_dim not in _CONV_CLASSES: @@ -79,13 +159,24 @@ def __init__( ) def forward(self, x: torch.Tensor) -> torch.Tensor: - """Forward pass of the Patchify layer. + """Embed the input tensor into a grid of patch tokens. Args: - x: The input tensor of shape [B, *spatial_dims, in_features]. + x: Input tensor in channels-last layout. Shape: + ``[B, *spatial_dims, C_in]``, e.g. ``[B, H, W, C_in]`` for + 2D images. Returns: - The output tensor of shape [B, *spatial_dims // stride, out_features]. + Patch-embedded tensor in channels-last layout. Shape: + ``[B, *patch_grid, C_out]``, where each spatial axis ``s`` is + reduced to ``floor((s - patch_size) / stride) + 1`` (equal to + ``s // patch_size`` when ``stride == patch_size`` and ``s`` is + divisible by ``patch_size``). + + Note: + ``.contiguous()`` is called after the channels-last → channels-first + rearrangement to avoid a stride-mismatch error in + ``torch.compile``'s ``convolution_backward``. """ # Channels-last -> channels-first for ConvNd # .contiguous() avoids a stride mismatch in torch.compile's convolution_backward @@ -100,14 +191,43 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class Unpatchify(torch.nn.Module): - """Inverse of Patchify for channels-last inputs (supports 1D/2D/3D). + """Inverse patch-embedding layer: reconstruct spatial signal from token grid. + + ``Unpatchify`` is the trainable inverse of ``Patchify``. Given a grid of + token embeddings at patch resolution, it reconstructs a signal at the + original spatial resolution using a transposed convolution + (``ConvTranspose{data_dim}d``). + + For non-overlapping patches (``stride == patch_size``), the default + transposed convolution is an exact spatial inverse: each output pixel is + produced by exactly one input token. When ``stride < patch_size`` + (overlapping), contributions from overlapping patches are *summed* by the + transposed convolution — this is the adjoint of the overlapping-patch + forward pass. + + **Output shape formula** (each spatial axis ``s`` of the patch-grid input):: + + out_s = (s - 1) * stride - 2 * padding + kernel_size + = (s - 1) * stride + patch_size (since padding == 0) + + For the non-overlapping case this gives ``s * patch_size``. + + **Layout convention** — inputs and outputs use *channels-last* ordering:: - Uses ConvTranspose to upsample from patch resolution back to original resolution. + input : [B, *patch_grid, C_embed] (e.g. [B, H/P, W/P, C_embed]) + output : [B, *spatial_dims, C_out] (e.g. [B, H, W, C_out]) - Input shape: [B, *spatial_dims, in_features] (channels-last) - Output shape: [B, *spatial_dims * stride, out_features] + **Weight initialisation** — PyTorch's default kaiming_uniform for + ``ConvTranspose`` uses ``fan_out = out_features * patch_size^data_dim``. + This is incorrect for large embedding dimensions; ``weight_init="fan_in"`` + corrects this by using the true fan-in + ``in_features * patch_size^data_dim``. - If exact spatial size control is required, pass output_spatial_shape to forward. + Attributes: + data_dim (int): Spatial dimensionality (1, 2, or 3). + patch_size (int): Kernel size of the transposed convolution. + stride (int): Stride of the transposed convolution. + deconv (torch.nn.ConvTranspose{data_dim}d): The underlying deconvolution. """ def __init__( @@ -120,24 +240,35 @@ def __init__( bias: bool = True, weight_init: Literal["default", "zeros", "fan_in"] = "default", ): - """Initialize the Unpatchify layer. + """Initialise the Unpatchify layer. Args: - in_features: The number of input channels (embedding dimension). - out_features: The number of output channels. - data_dim: The spatial dimensionality (1, 2, or 3). - patch_size: The size of each patch (kernel_size for the deconv). - stride: The stride for the deconv. Defaults to patch_size (inverse of non-overlapping). - bias: Whether the underlying deconv has a learnable bias. Default True - for backward compatibility; set False for bias-free architectures. + in_features: Embedding dimension ``C_embed`` of each input token. + out_features: Number of output channels ``C_out`` of the + reconstructed signal (e.g. 3 for RGB images). + data_dim: Spatial dimensionality. Must be 1, 2, or 3. + patch_size: Side length ``P`` of each patch. The transposed + convolution uses ``kernel_size = patch_size`` along every + spatial axis. + stride: Step between consecutive output patch origins. Defaults to + ``patch_size`` (non-overlapping, exact inverse of + ``Patchify`` with default stride). Must match the ``stride`` + used in the paired ``Patchify`` layer to recover the original + spatial resolution. + bias: If ``True`` (default), the deconvolution includes a learnable + bias term. weight_init: Weight initialisation strategy for the deconv kernel. - - ``"default"``: PyTorch default kaiming_uniform (fan based on out_channels — - incorrect for large in_features, can cause output variance blow-up). - - ``"zeros"``: Zero-initialise weights and bias (DiT-style; output is exactly - zero at init, safe residual-stream entry). - - ``"fan_in"``: Kaiming-uniform using the true fan-in - (``in_features * patch_size ** data_dim``), which correctly scales the output - variance to O(1) regardless of embedding dimension. + ``"default"`` uses PyTorch's built-in ``kaiming_uniform`` + (fan computed from ``out_features``; can cause output-variance + blow-up for large ``in_features``). ``"zeros"`` zero-inits + weights and bias (DiT-style; output is exactly zero at + initialisation, safe for residual-stream entry). ``"fan_in"`` + applies Kaiming-uniform with the corrected fan-in + ``in_features * patch_size^data_dim``, giving output variance + O(1) regardless of embedding dimension. + + Raises: + ValueError: If ``data_dim`` is not 1, 2, or 3. """ super().__init__() if data_dim not in _CONV_TRANSPOSE_CLASSES: @@ -175,14 +306,36 @@ def __init__( # "default": leave PyTorch's kaiming_uniform as-is def forward(self, x: torch.Tensor, output_spatial_shape: Tuple[int, ...] | None = None) -> torch.Tensor: - """Forward pass of the Unpatchify layer. + """Reconstruct the spatial signal from a grid of patch-token embeddings. Args: - x: The input tensor of shape [B, *spatial_dims, in_features]. - output_spatial_shape: The desired output spatial shape (optional). + x: Token-grid tensor in channels-last layout. Shape: + ``[B, *patch_grid, C_embed]``, e.g. + ``[B, H/P, W/P, C_embed]`` for 2D. The number of spatial + dimensions must equal ``data_dim``. + output_spatial_shape: Optional target spatial shape for the output. + When provided this is passed as ``output_size`` to the + transposed convolution, which resolves the output-size ambiguity + that arises when ``stride > 1`` (multiple input sizes can map to + the same output size via the floor in the forward formula). + Must have length ``data_dim``. When ``None``, PyTorch infers + the output size (may differ from the original input spatial size + if ``spatial_dim % patch_size != 0``). Returns: - The output tensor of shape [B, *spatial_dims * stride, out_features]. + Reconstructed signal tensor in channels-last layout. Shape: + ``[B, *spatial_dims, C_out]``. Without ``output_spatial_shape``, + each axis ``s`` of the patch grid expands to + ``(s - 1) * stride + patch_size``. + + Raises: + AssertionError: If the rank of ``x`` does not equal + ``data_dim + 2`` (batch + spatial + channel dims). + + Note: + ``.contiguous()`` is called after the rearrangement to channels-first + to avoid a stride-mismatch error in ``torch.compile``'s + ``convolution_backward``. """ expected_dims = self.data_dim + 2 # batch + spatial_dims + channels assert x.dim() == expected_dims, ( From 4e9ef62b1e7f07a37710cf08748b2eafeb2585c5 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:23:50 +0200 Subject: [PATCH 14/72] docs(write/residual_block): add module and class docstrings --- nvsubquadratic/modules/residual_block.py | 277 +++++++++++++++++++++-- 1 file changed, 259 insertions(+), 18 deletions(-) diff --git a/nvsubquadratic/modules/residual_block.py b/nvsubquadratic/modules/residual_block.py index 5056f8ca..5459346e 100644 --- a/nvsubquadratic/modules/residual_block.py +++ b/nvsubquadratic/modules/residual_block.py @@ -13,7 +13,36 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Residual block implementation for ND signals, composed of a sequence mixer and an MLP.""" +"""Residual block implementations for ND signals. + +This module provides the fundamental repeating unit of the nvsubquadratic +architecture: a residual block that wraps a *sequence mixer* and an *MLP* +with pre-norm and optional conditioning branches. Stacking many such blocks +forms the full depth of the network (``general_purpose_resnet`` / +``classification_resnet``). + +Two block variants are provided: + +:class:`ResidualBlock` + The standard pre-norm residual block used in most networks. Each forward + pass applies up to three gated sub-branches: + + 1. **Sequence mixer branch** — pre-norm → :class:`~nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer` → residual add. + 2. **Condition mixer branch** *(optional)* — pre-norm → cross-attention / conditioning mixer → residual add. + 3. **MLP branch** — pre-norm → :class:`~nvsubquadratic.modules.mlp.MLP` → residual add. + + Any branch can be disabled at config time by setting its config target to + ``torch.nn.Identity``; the corresponding norm must also be Identity. + +:class:`AdaLNZeroResidualBlock` + A DiT-style block that replaces LayerNorm with *Adaptive LayerNorm-Zero* + (AdaLN-Zero) modulation. A single zero-initialised linear layer maps the + conditioning vector to six affine parameters (shift + scale + gate for each + of the two branches), giving the conditioning signal fine-grained control + over both branches at initialisation-time stability (outputs ≈ 0). + +All tensors follow **channels-last** layout: ``(B, *spatial_dims, C)``. +""" from typing import Union @@ -23,7 +52,48 @@ class ResidualBlock(torch.nn.Module): - """Residual block for ND signals, composed of a sequence mixer and an MLP.""" + """Standard pre-norm residual block for ND signals. + + Each forward pass executes up to three residual sub-branches: + + .. code-block:: text + + x ──[input_norm]──► sequence_mixer ──► dropout ──►(+)──► x' + x'──[cond_norm]───► condition_mixer ──► dropout ──►(+)──► x'' (optional) + x''─[mlp_norm]────► mlp ──────────────► dropout ──►(+)──► output + + Any branch is *bypassed entirely* (not just zeroed) when its ``_cfg`` + target is ``torch.nn.Identity``. This design lets the same class serve + pure-sequence networks (condition branch disabled), cross-attention + encoder-decoders (condition branch enabled), and MLP-only ablations. + + All normalisation parameters (``input_norm``, ``condition_mixer_norm``, + ``mlp_norm``) are tagged with ``_no_weight_decay = True`` so that the + optimiser can exclude them from weight-decay groups. + + Attributes: + sequence_mixer (torch.nn.Module): The instantiated sequence-mixing + operator (e.g. :class:`~nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer` + wrapping Hyena, Attention, CKConv, or Mamba). May be + ``torch.nn.Identity`` when the mixer branch is disabled. + input_norm (torch.nn.Module): Pre-norm applied before the sequence + mixer (e.g. LayerNorm or RMSNorm). May be ``torch.nn.Identity``. + condition_mixer (torch.nn.Module): Cross-attention or conditioning + operator applied after the sequence mixer. May be + ``torch.nn.Identity`` to disable the conditioning branch entirely. + condition_mixer_norm (torch.nn.Module): Pre-norm applied before + ``condition_mixer``. Must be ``torch.nn.Identity`` when + ``condition_mixer`` is ``torch.nn.Identity``. + mlp (torch.nn.Module): Position-wise MLP + (e.g. :class:`~nvsubquadratic.modules.mlp.MLP`). May be + ``torch.nn.Identity`` to disable the MLP branch. + mlp_norm (torch.nn.Module): Pre-norm applied before the MLP. Must be + ``torch.nn.Identity`` when ``mlp`` is ``torch.nn.Identity``. + dropout (torch.nn.Module): Dropout (or stochastic depth) applied after + each active branch. Typically + :class:`~nvsubquadratic.modules.drop_path.DropPath` or + ``torch.nn.Dropout``. + """ def __init__( self, @@ -35,16 +105,47 @@ def __init__( mlp_norm_cfg: LazyConfig, dropout_cfg: LazyConfig, ): - """Initialize the ResidualBlock. + """Initialise the ResidualBlock. + + All positional sub-modules are supplied as + :class:`~nvsubquadratic.lazy_config.LazyConfig` objects and + instantiated here. Passing ``torch.nn.Identity`` as the target for a + ``*_cfg`` / ``*_norm_cfg`` pair disables that branch at zero cost (no + forward computation, no parameters). Args: - sequence_mixer_cfg: LazyConfig for the sequence mixer layer. - sequence_mixer_norm_cfg: LazyConfig for the sequence mixer norm. - condition_mixer_cfg: LazyConfig for the condition mixer layer. - condition_mixer_norm_cfg: LazyConfig for the condition mixer norm. - mlp_cfg: LazyConfig for the MLP layer. - mlp_norm_cfg: LazyConfig for the MLP norm. - dropout_cfg: LazyConfig for the dropout layer. + sequence_mixer_cfg: LazyConfig for the sequence mixer. Typical + targets: :class:`~nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer` + (which internally wraps Hyena, Attention, CKConv, or Mamba), + or ``torch.nn.Identity`` to skip the mixer branch entirely. + sequence_mixer_norm_cfg: LazyConfig for the pre-norm applied + before the sequence mixer. **Must** be ``torch.nn.Identity`` + when ``sequence_mixer_cfg`` targets ``torch.nn.Identity``. + condition_mixer_cfg: LazyConfig for the conditioning / cross- + attention operator. Pass ``torch.nn.Identity`` (the default + in most configs) to disable the conditioning branch. When + active, the module's ``forward`` must accept + ``(x, condition)`` positional arguments. + condition_mixer_norm_cfg: LazyConfig for the pre-norm applied + before ``condition_mixer``. **Must** be ``torch.nn.Identity`` + when ``condition_mixer_cfg`` targets ``torch.nn.Identity``. + mlp_cfg: LazyConfig for the position-wise MLP. Typical target: + :class:`~nvsubquadratic.modules.mlp.MLP`. Pass + ``torch.nn.Identity`` to skip the MLP branch. + mlp_norm_cfg: LazyConfig for the pre-norm applied before the MLP. + **Must** be ``torch.nn.Identity`` when ``mlp_cfg`` targets + ``torch.nn.Identity``. + dropout_cfg: LazyConfig for the dropout / stochastic-depth module + applied after each active branch. Typical targets: + ``torch.nn.Dropout``, + :class:`~nvsubquadratic.modules.drop_path.DropPath`, or + ``torch.nn.Identity`` for no dropout. + + Raises: + AssertionError: If a norm config does not match its corresponding + module config — i.e. if a mixer/MLP config is ``Identity`` but + the corresponding norm config is not (or vice versa in the + forward-only direction). """ if sequence_mixer_cfg.__target__ == torch.nn.Identity: assert sequence_mixer_norm_cfg.__target__ == torch.nn.Identity, ( @@ -86,14 +187,51 @@ def __init__( self.dropout = instantiate(dropout_cfg) def forward(self, x: torch.Tensor, condition: torch.Tensor) -> torch.Tensor: - """Forward pass of the residual block. + """Apply the residual block to the input tensor. + + Executes up to three residual sub-branches in order: sequence mixer, + conditioning mixer (optional), and MLP (optional). A branch is + **skipped entirely** when its corresponding module is + ``torch.nn.Identity``; the ``condition`` argument is only consumed + when the conditioning branch is active. + + **Path without conditioning** (``condition_mixer`` is Identity): + + .. code-block:: text + + x → input_norm → sequence_mixer → dropout →(+)→ x' + x' → mlp_norm → mlp → dropout →(+)→ output + + **Path with conditioning** (``condition_mixer`` is not Identity): + + .. code-block:: text + + x → input_norm → sequence_mixer → dropout →(+)→ x' + x' → condition_mixer_norm → condition_mixer(x', condition) + → dropout →(+)→ x'' + x'' → mlp_norm → mlp → dropout →(+)→ output Args: - x (torch.Tensor): Input tensor of shape (batch_size, *spatial_dims, num_hidden_channels) - condition (torch.Tensor): Condition tensor of shape (batch_size, *spatial_dims_condition, num_hidden_channels) + x: Input feature tensor of shape + ``(B, *spatial_dims, C)`` where ``B`` is the batch size, + ``spatial_dims`` are one or more spatial axes (e.g. ``(H, W)`` + for 2-D images or ``(T,)`` for 1-D sequences), and + ``C`` is the hidden channel dimension. + condition: Conditioning tensor used by ``condition_mixer``. Its + shape depends on the conditioning operator — a common choice + is ``(B, *spatial_dims_condition, C)`` for cross-attention, or + ``(B, C)`` for a global conditioning vector. This argument is + **ignored** (and may safely be ``None``) when + ``condition_mixer`` is ``torch.nn.Identity``. Returns: - torch.Tensor: Output tensor of shape (batch_size, *spatial_dims, num_hidden_channels) + torch.Tensor: Output tensor of the same shape as ``x``: + ``(B, *spatial_dims, C)``. + + Raises: + AssertionError: If ``condition`` is ``None`` but + ``condition_mixer`` is not ``torch.nn.Identity`` (i.e. a + conditioning tensor is required but was not provided). """ # Mixer branch if not isinstance(self.sequence_mixer, torch.nn.Identity): @@ -123,7 +261,61 @@ def forward(self, x: torch.Tensor, condition: torch.Tensor) -> torch.Tensor: class AdaLNZeroResidualBlock(torch.nn.Module): - """Residual block with inline AdaLN-Zero modulation for mixer and MLP branches.""" + """Pre-norm residual block with AdaLN-Zero conditioning (DiT-style). + + Replaces fixed LayerNorm with *Adaptive LayerNorm-Zero* (AdaLN-Zero) + modulation, following the Scalable Diffusion Transformers (DiT) recipe + (Peebles & Xie, 2023). A single zero-initialised linear projection maps + the conditioning vector to six affine parameters — shift, scale, and gate + for each of the two branches — so that at initialisation the block outputs + exactly zero (the residual stream is unchanged). + + **Forward computation** (one block): + + .. code-block:: text + + cond → [optional spatial mean] → condition_norm + → SiLU → Linear(C, 6C) → split into 6 × (B, C) + (shift_seq, scale_seq, gate_seq, shift_mlp, scale_mlp, gate_mlp) + + # Sequence mixer branch + x_norm = sequence_norm(x) # pre-norm + x_mod = x_norm * (1 + scale_seq) + shift_seq # AdaLN modulation + seq_out = sequence_mixer(x_mod, conditioning=cond) + seq_out = dropout(seq_out) * gate_seq # zero-init gate + x = x + seq_out + + # MLP branch + x_norm = mlp_norm(x) + x_mod = x_norm * (1 + scale_mlp) + shift_mlp + mlp_out = mlp(x_mod) + mlp_out = dropout(mlp_out) * gate_mlp + x = x + mlp_out + + The gate vectors are multiplied **after** dropout to provide per-token + scaling; at init they are zero (because ``condition_proj`` weights are + zero-initialised), so the block is a skip connection. + + Attributes: + sequence_mixer (torch.nn.Module): Instantiated sequence-mixing + operator. Its ``forward`` must accept a ``conditioning`` keyword + argument (forwarded from the pooled conditioning vector). + sequence_norm (torch.nn.Module): Pre-norm applied to ``x`` before + AdaLN modulation in the sequence mixer branch. + mlp (torch.nn.Module): Position-wise MLP instantiated from + ``mlp_cfg``. + mlp_norm (torch.nn.Module): Pre-norm applied to ``x`` before + AdaLN modulation in the MLP branch. + condition_norm (torch.nn.Module): Optional normalisation applied to + the conditioning vector before it is projected by + ``condition_proj``. Tagged ``_no_weight_decay``. + dropout (torch.nn.Module): Dropout / stochastic-depth applied after + each branch, before the gate multiply. + condition_proj (torch.nn.Sequential): ``SiLU → Linear(C, 6C)`` + with **zero-initialised** weights and biases, producing the six + AdaLN-Zero parameters. Zero init ensures the block is a pure + residual connection at the start of training. + """ def __init__( self, @@ -135,7 +327,28 @@ def __init__( dropout_cfg: LazyConfig, hidden_dim: int, ): - """Initialize the AdaLNZeroResidualBlock.""" + """Initialise the AdaLNZeroResidualBlock. + + Args: + sequence_mixer_cfg: LazyConfig for the sequence-mixing operator. + The instantiated module's ``forward`` must accept a + ``conditioning`` keyword argument (it receives the spatially- + pooled conditioning vector ``cond`` of shape ``(B, C)``). + sequence_mixer_norm_cfg: LazyConfig for the pre-norm applied + before the AdaLN modulation of the sequence mixer branch + (e.g. LayerNorm over ``hidden_dim``). + mlp_cfg: LazyConfig for the position-wise MLP. + mlp_norm_cfg: LazyConfig for the pre-norm applied before the + AdaLN modulation of the MLP branch. + condition_norm_cfg: LazyConfig for the normalisation applied to + the conditioning vector before ``condition_proj``. Use + ``torch.nn.Identity`` to skip conditioning normalisation. + dropout_cfg: LazyConfig for the dropout / stochastic-depth module + applied after each branch and before the gate multiply. + hidden_dim: Channel dimension ``C`` shared by all sub-modules. + Used to size the ``condition_proj`` linear layer + (``Linear(C, 6*C)``). + """ super().__init__() # Mixer branch handles spatial/temporal interactions on the residual stream. @@ -163,8 +376,36 @@ def __init__( torch.nn.init.zeros_(self.condition_proj[1].weight) torch.nn.init.zeros_(self.condition_proj[1].bias) - def forward(self, x: torch.Tensor, condition: Union[torch.Tensor | None]) -> torch.Tensor: - """Apply AdaLN-Zero residual mixing conditioned on the provided tensor.""" + def forward(self, x: torch.Tensor, condition: Union[torch.Tensor, None]) -> torch.Tensor: + """Apply AdaLN-Zero residual mixing conditioned on the provided tensor. + + The conditioning tensor is reduced to a single latent vector per + sample (spatial mean if it has spatial axes) before being projected + to six affine parameters. These parameters modulate the pre-norm + outputs of both branches via element-wise affine transforms, and gate + each branch's output before the residual add. + + Args: + x: Input feature tensor of shape ``(B, *spatial_dims, C)`` where + ``B`` is the batch size, ``spatial_dims`` are one or more + spatial axes, and ``C = hidden_dim``. + condition: Required conditioning tensor. Shape may be: + + * ``(B, C)`` — a pre-pooled global conditioning vector + (e.g. a timestep / class embedding from a diffusion model). + * ``(B, *spatial_dims_cond, C)`` — any spatial layout; the + forward pass reduces it to ``(B, C)`` via a mean over all + non-batch, non-channel axes. + + Must not be ``None``. + + Returns: + torch.Tensor: Output tensor of shape ``(B, *spatial_dims, C)``, + the same shape as ``x``. + + Raises: + ValueError: If ``condition`` is ``None``. + """ if condition is None: raise ValueError("AdaLNZeroResidualBlock requires a conditioning tensor.") From ad08ece5e1e1635129f0aec07e9876c0a137d8e0 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:24:08 +0200 Subject: [PATCH 15/72] docs(write/film): add module and class docstrings with math context --- nvsubquadratic/modules/film.py | 175 +++++++++++++++++++++++++++++---- 1 file changed, 156 insertions(+), 19 deletions(-) diff --git a/nvsubquadratic/modules/film.py b/nvsubquadratic/modules/film.py index f52fcd47..de16889b 100644 --- a/nvsubquadratic/modules/film.py +++ b/nvsubquadratic/modules/film.py @@ -1,5 +1,50 @@ """FiLM (Feature-wise Linear Modulation) components for kernel conditioning. +**What is FiLM?** + +FiLM conditions a feature map ``x`` on an external signal ``c`` (e.g. a +timestep embedding, class label, or physics parameter vector) via an affine +transformation: + + y = γ(c) ⊙ x + β(c) + +where ``γ(c)`` (scale) and ``β(c)`` (shift) are learned functions of the +conditioning vector and ``⊙`` denotes elementwise multiplication. The +parameters ``γ`` and ``β`` are **not** fixed — they are the outputs of a +small neural network (the *FiLM generator*) evaluated at runtime. + +*Reference*: Perez et al., "FiLM: Visual Reasoning with a General +Conditioning Layer", arXiv:1709.07871 (2017). + +**Why is FiLM useful?** + +* **Diffusion models** — inject the noising timestep *t* so that each + residual block adapts its scale and shift to the current noise level + (see DiT, Peebles & Xie, 2022). +* **Class conditioning** — modulate intermediate feature maps by a class + embedding, enabling a single model to behave differently across classes + without separate heads. +* **Physics / PDE parameter conditioning** — allow a learned PDE solver to + generalise across equation parameters (viscosity, Reynolds number, etc.) + by FiLM-conditioning the implicit neural operator kernel. +* **SIREN kernel modulation** — within :mod:`nvsubquadratic.modules.kernels_nd`, + FiLM adjusts the hidden activations of each SIREN layer, effectively + steering the learned convolution kernel per sample or per timestep. + +**Conditioning signal sources in this codebase** + +The conditioning vector ``c ∈ ℝ^{cond_dim}`` typically originates from +register tokens appended to the sequence. Two encoder modules are provided: + +* :class:`RegisterPooling` — learnable softmax-weighted average over all + register tokens; produces a ``[B, C]`` summary vector. +* :class:`RegisterCompressConcat` — compresses each register token with a + shared linear and concatenates; preserves per-register identity at the cost + of a larger ``cond_dim``. + +The encoded conditioning vector is then fed into :class:`KernelFiLMGenerator` +which maps it to ``(γ, β)`` pairs — one pair per SIREN hidden layer. + Provides: - KernelFiLMGenerator: MLP that maps a conditioning vector to per-layer (gamma, beta) pairs for modulating SIREN hidden layers. @@ -19,16 +64,41 @@ class KernelFiLMGenerator(nn.Module): - """Generates per-layer FiLM (gamma, beta) pairs from a conditioning vector. + """MLP that generates per-layer FiLM (γ, β) pairs from a conditioning vector. + + Given a conditioning signal ``c ∈ ℝ^{cond_dim}`` (e.g. from register tokens + processed by :class:`RegisterPooling` or :class:`RegisterCompressConcat`), + this module produces one ``(γ_l, β_l)`` pair per SIREN hidden layer ``l``: + + h_l ← γ_l(c) ⊙ h_l + β_l(c) + + The generator itself is a two-layer MLP with a GELU non-linearity: - Maps [B, cond_dim] -> list of num_film_layers x (gamma [B, kernel_hidden_dim], beta [B, kernel_hidden_dim]). + c → Linear(cond_dim, film_hidden_dim) → GELU + → Linear(film_hidden_dim, num_film_layers × 2 × kernel_hidden_dim) - The output bias is initialized to ``(gamma=1, beta=0)`` per layer so that - ``gamma * h + beta = h`` at init (identity modulation). All biases in the - MLP are **always** excluded from weight decay. + The flat output is split into ``num_film_layers`` chunks; each chunk is + further split in half to give ``(γ_l, β_l) ∈ ℝ^{kernel_hidden_dim}``. + + **Initialization strategy** — The output layer is initialized so that at + the start of training ``γ_l = 1`` and ``β_l = 0`` for every layer, making + FiLM an *identity* modulation. This prevents early instability when + the conditioning signal is still uninformative. The ``"small_random"`` + variant perturbs the output weights slightly to break weight-symmetry + while keeping the bias-induced identity. + + **Weight-decay handling** — All biases are permanently excluded from + weight decay (``_no_weight_decay = True``). Weight matrices can be + excluded entirely (``no_weight_decay=True``) or assigned a custom decay + value (``no_weight_decay=``). + + Attributes: + num_film_layers (int): Number of ``(γ, β)`` pairs produced. + kernel_hidden_dim (int): Feature dimension of each SIREN hidden layer. + mlp (nn.Sequential): Two-layer MLP (Linear → GELU → Linear). Args: - cond_dim: Dimensionality of the conditioning input. + cond_dim: Dimensionality of the conditioning input ``c``. kernel_hidden_dim: Hidden dimension of the SIREN layers to modulate. num_film_layers: Number of (gamma, beta) pairs to produce (one per SIREN hidden layer). film_hidden_dim: Hidden dimension of the FiLM generator MLP (bottleneck). @@ -96,6 +166,14 @@ def _init_weights(self, init_type: Literal["identity", "small_random"], init_std ``"small_random"`` additionally gives the output weights a small random perturbation to break the zero-weight saddle point. + + Args: + init_type: Initialization strategy — ``"identity"`` zeros the output + weights; ``"small_random"`` draws them from N(0, ``init_std``). + init_std: Standard deviation used when ``init_type="small_random"``. + + Raises: + ValueError: If ``init_type`` is not ``"identity"`` or ``"small_random"``. """ final_linear = self.mlp[-1] @@ -144,14 +222,20 @@ def flop_count(self) -> int: return flops def forward(self, conditioning: torch.Tensor) -> list[tuple[torch.Tensor, torch.Tensor]]: - """Generate FiLM parameters from the conditioning vector. + """Generate per-layer FiLM parameters from the conditioning vector. + + Runs the two-layer MLP on ``conditioning`` and splits the flat output + into ``num_film_layers`` ``(γ, β)`` pairs. Each pair should be applied + by the SIREN caller as ``h_l ← γ_l ⊙ h_l + β_l``. Args: - conditioning: [B, cond_dim] + conditioning: Conditioning vector of shape ``[B, cond_dim]``. Typically + produced by :class:`RegisterPooling` or :class:`RegisterCompressConcat`. Returns: - List of (gamma, beta) tuples, each [B, kernel_hidden_dim]. - Callers apply ``gamma * h + beta``. + A list of ``num_film_layers`` tuples ``(gamma, beta)``, where each + tensor has shape ``[B, kernel_hidden_dim]``. Index ``l`` of the list + corresponds to SIREN hidden layer ``l``. """ out = self.mlp(conditioning) # [B, num_film_layers * 2 * kernel_hidden_dim] chunks = out.chunk(self.num_film_layers, dim=-1) # num_film_layers x [B, 2 * kernel_hidden_dim] @@ -161,9 +245,26 @@ def forward(self, conditioning: torch.Tensor) -> list[tuple[torch.Tensor, torch. class RegisterPooling(nn.Module): - """Learnable weighted average over register tokens. + """Learnable softmax-weighted average over register tokens. + + Produces a single conditioning vector ``c ∈ ℝ^C`` per sample from a set of + ``R`` register tokens. The pooling weights are learned scalars ``w ∈ ℝ^R`` + passed through a softmax, so the contribution of each register is + non-negative and the weights sum to one: + + c = Σ_r softmax(w)_r · x_r - Produces a single conditioning vector per sample from multiple register tokens. + This is a lightweight alternative to :class:`RegisterCompressConcat` when a + scalar summary per register is sufficient. All register information is + blended into a single ``[B, C]`` vector that can be passed directly to + :class:`KernelFiLMGenerator`. + + The learned logit vector ``w`` is always excluded from weight decay via + ``_no_weight_decay = True``. + + Attributes: + logits (nn.Parameter): Learnable unnormalized pooling weights of shape + ``[num_registers]``, initialised to zero (uniform softmax at init). Args: num_registers: Number of register tokens to pool over. @@ -194,13 +295,21 @@ def flop_count(self, dim: int) -> int: return 3 * num_registers + 2 * num_registers * dim def forward(self, registers: torch.Tensor) -> torch.Tensor: - """Pool register tokens into a single vector. + """Pool register tokens into a single conditioning vector. + + Applies softmax to the learned logits and computes a weighted sum over + the register dimension: + + c_b = Σ_r softmax(logits)_r · registers_{b,r,:} Args: - registers: [B, num_registers, C] + registers: Register token tensor of shape ``[B, num_registers, C]``, + where ``B`` is the batch size and ``C`` is the channel dimension. + The number of registers along axis 1 must equal ``num_registers`` + passed to ``__init__``. Returns: - [B, C] pooled conditioning vector. + Pooled conditioning vector of shape ``[B, C]``. """ weights = F.softmax(self.logits, dim=0) # [num_registers] return torch.einsum("r, b r c -> b c", weights, registers) @@ -218,9 +327,24 @@ class RegisterCompressConcat(nn.Module): vector of size ``num_registers * compressed_dim`` that preserves per-register identity (unlike :class:`RegisterPooling` which averages them). + Formally, for a batch of register tensors ``X ∈ ℝ^{B × R × D}``: + + compressed_r = W · x_r, W ∈ ℝ^{compressed_dim × hidden_dim} (shared across r) + c = [compressed_0 ‖ compressed_1 ‖ … ‖ compressed_{R-1}] ∈ ℝ^{R · compressed_dim} + + The output ``c`` is then consumed by :class:`KernelFiLMGenerator` whose + ``cond_dim`` must equal ``num_registers * compressed_dim`` (see + :attr:`out_dim`). + Inspired by Mamba-Reg (Wang et al., 2024) which distributes and individually reads out register tokens rather than pooling them. + Attributes: + num_registers (int): Number of register tokens expected on the sequence axis. + compressed_dim (int): Output channel dimension per register after compression. + compress (nn.Linear): Shared weight-only (no bias) projection + ``hidden_dim → compressed_dim`` applied independently to each register. + Args: num_registers: Number of register tokens expected. hidden_dim: Channel dimension of each register token. @@ -236,7 +360,12 @@ def __init__(self, num_registers: int, hidden_dim: int, compressed_dim: int): # @property def out_dim(self) -> int: - """Dimensionality of the output conditioning vector.""" + """Dimensionality of the output conditioning vector. + + Returns: + ``num_registers * compressed_dim``, which must match ``cond_dim`` + of any downstream :class:`KernelFiLMGenerator`. + """ return self.num_registers * self.compressed_dim def flop_count(self, dim: int) -> int: @@ -255,13 +384,21 @@ def flop_count(self, dim: int) -> int: return self.num_registers * 2 * dim * self.compressed_dim def forward(self, registers: torch.Tensor) -> torch.Tensor: - """Compress each register and concatenate. + """Compress each register token and concatenate into a flat conditioning vector. + + Applies the shared linear projection to every register independently + (via broadcasting over the register axis), then flattens the register + and compressed-channel axes into a single vector per sample. Args: - registers: [B, num_registers, hidden_dim] + registers: Register token tensor of shape ``[B, num_registers, hidden_dim]``, + where ``B`` is the batch size. The number of registers along axis 1 + must equal ``num_registers`` passed to ``__init__``. Returns: - [B, num_registers * compressed_dim] conditioning vector. + Flat conditioning vector of shape ``[B, num_registers * compressed_dim]``. + Pass this directly as the ``conditioning`` argument of + :class:`KernelFiLMGenerator`. """ compressed = self.compress(registers) # [B, R, compressed_dim] return compressed.flatten(start_dim=1) # [B, R * compressed_dim] From 8ab96db88702b897661f2728730a2e123effbf0b Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:24:45 +0200 Subject: [PATCH 16/72] docs(review/residual_block): reviewer feedback --- docs/reviews/patchify_review.md | 138 ++++++++++++++++++++++++++ docs/reviews/residual_block_review.md | 137 +++++++++++++++++++++++++ 2 files changed, 275 insertions(+) create mode 100644 docs/reviews/patchify_review.md create mode 100644 docs/reviews/residual_block_review.md diff --git a/docs/reviews/patchify_review.md b/docs/reviews/patchify_review.md new file mode 100644 index 00000000..d73e9248 --- /dev/null +++ b/docs/reviews/patchify_review.md @@ -0,0 +1,138 @@ +# Review: `nvsubquadratic/modules/patchify.py` + +Reviewer pass on commit `docs(write/patchify): add module and class docstrings`. + +______________________________________________________________________ + +## Issues + +### 1. Module docstring — overlapping-patch semantics need a concrete example + +**Location:** Module docstring, "Overlapping patches" paragraph in `Patchify`. + +**Problem:** The docstring says "setting `stride < patch_size` produces overlapping patches +with the same formula above" but does not show the formula for that case. An external +collaborator who only reads the module docstring (not the class docstring) sees no formula at all. + +**Fix:** Add the general output-size formula `floor((s - patch_size) / stride) + 1` to +the module docstring overview, directly after the description of the output layout convention. + +______________________________________________________________________ + +### 2. `Patchify.__init__` — `patch_size` docstring does not mention the per-patch pixel count + +**Location:** `Patchify.__init__` docstring, `patch_size` arg entry. + +**Quote:** `"Side length P of each patch. The convolution uses kernel_size = patch_size along every spatial axis, so each patch covers P^data_dim input pixels."` + +**Problem:** The phrase "each patch covers P^data_dim input pixels" uses caret notation that +renders as plain text in most docstring viewers. Readers unfamiliar with the notation may +think `^` is a bitwise XOR. + +**Fix:** Replace `P^data_dim` with `P**data_dim` (Python exponentiation) or spell it out +as "`P × P` pixels (2D)" / "`P × P × P` voxels (3D)" with a note that the general +formula is `P ** data_dim`. + +______________________________________________________________________ + +### 3. `Patchify` class docstring — no mention of what happens when spatial size is not divisible by `patch_size` + +**Location:** `Patchify` class docstring, immediately after the output shape formula. + +**Problem:** The formula `s / patch_size` (non-overlapping case) implicitly assumes exact +divisibility, but the layer silently truncates if `s % patch_size != 0` (standard Conv +floor semantics). A new collaborator who passes a 224×224 image with `patch_size=16` +gets 14×14 tokens as expected, but a 225×224 image would give 14×14 tokens silently +dropping the last pixel column. This is a frequent source of bugs. + +**Fix:** Add a warning note under the output shape formula, e.g.: + +> **Note:** If `spatial_dim % patch_size != 0`, the last pixels in that axis are silently +> discarded (standard floor-division Conv semantics). Callers are responsible for ensuring +> the spatial dimensions are divisible by `patch_size` (e.g. via padding before this layer). + +______________________________________________________________________ + +### 4. `Unpatchify.forward` — `output_spatial_shape` description buries the most important use-case + +**Location:** `Unpatchify.forward` docstring, `output_spatial_shape` arg description. + +**Quote:** `"Must have length data_dim. When None, PyTorch infers the output size (may differ from the original input spatial size if spatial_dim % patch_size != 0)."` + +**Problem:** The caveat about size ambiguity is in the middle of a long sentence. It is the +most important reason to pass `output_spatial_shape`, but a skimming reader will miss it. + +**Fix:** Restructure the description to lead with the ambiguity problem: + +> When `stride > 1` multiple patch-grid sizes map to the same output size (floor in the +> forward direction discards remainders). Pass `output_spatial_shape` to resolve this +> ambiguity and guarantee recovery of the exact original spatial dimensions. Must have +> length `data_dim`. When `None`, the output size is inferred by PyTorch and may not +> match the original spatial size if `spatial_dim % patch_size != 0`. + +______________________________________________________________________ + +### 5. `Unpatchify` class docstring — overlapping-patch adjoint semantics are ambiguous + +**Location:** `Unpatchify` class docstring, second paragraph. + +**Quote:** `"When stride < patch_size (overlapping), contributions from overlapping patches are *summed* by the transposed convolution — this is the adjoint of the overlapping-patch forward pass."` + +**Problem:** The summation is correct mathematically, but "adjoint" may mislead readers into +thinking `Unpatchify` perfectly reconstructs the original signal. For overlapping patches +the transposed conv is only the linear adjoint (backward map), not an inverse; pixel values +are *accumulated*, not averaged, so the output scale grows with the overlap ratio. This is +especially unexpected for users who pair `Patchify` and `Unpatchify` in an autoencoder. + +**Fix:** Add an explicit note: + +> For overlapping patches this accumulation means `Unpatchify(Patchify(x))` does **not** +> recover `x` exactly; the output is a blurred, scaled version of `x`. Only for +> non-overlapping patches (`stride == patch_size`) does the round-trip preserve spatial +> alignment (up to the learned weights). + +______________________________________________________________________ + +### 6. `Unpatchify.__init__` — `weight_init` tradeoffs are missing for the `"default"` option + +**Location:** `Unpatchify.__init__` docstring, `weight_init` arg, `"default"` option. + +**Quote:** `'"default" uses PyTorch's built-in kaiming_uniform (fan computed from out_features; can cause output-variance blow-up for large in_features).'` + +**Problem:** The existing text warns about variance blow-up for `"default"` but does not tell +the reader *when* to choose it anyway. An external collaborator who has not read the +DiT/ViT literature will default to `"default"` without knowing the risk. + +**Fix:** Add a recommendation, e.g.: + +> Prefer `"fan_in"` for new architectures. `"default"` is retained for loading +> pre-trained checkpoints whose weights were saved under PyTorch's default init. + +______________________________________________________________________ + +### 7. Missing cross-reference to `PositionEmbeddingND` in module docstring + +**Location:** Module docstring. + +**Problem:** The module docstring mentions that the output layout matches what +"PositionEmbeddingND and the mixer blocks" expect, but does not tell the reader +where to find `PositionEmbeddingND`. + +**Fix:** Add an explicit cross-reference: + +> See `nvsubquadratic.modules.position_encoding.PositionEmbeddingND` for the positional +> encoding layer that is typically applied immediately after `Patchify`. + +______________________________________________________________________ + +### 8. `Patchify` and `Unpatchify` class docstrings — no ND usage example + +**Location:** Both class docstrings. + +**Problem:** The class-level docstrings show only implied 2D usage (the shapes use `H, W` +notation). A reader working on 1D sequences or 3D volumes must infer the generalisation +from the `data_dim` parameter alone, with no concrete example. + +**Fix:** Add a brief "Examples" section to at least one of the classes showing a 1D case +(`[B, L, C]` → `[B, L/P, C_embed]`) and referencing the `__main__` block for the 2D +demonstration. diff --git a/docs/reviews/residual_block_review.md b/docs/reviews/residual_block_review.md new file mode 100644 index 00000000..0a71d265 --- /dev/null +++ b/docs/reviews/residual_block_review.md @@ -0,0 +1,137 @@ +# Review: `nvsubquadratic/modules/residual_block.py` + +Reviewed as a second reader looking for gaps that would confuse an external collaborator +reading alongside the research paper. + +______________________________________________________________________ + +## Issues + +### 1. Module docstring does not clarify that `ResidualBlock` and `AdaLNZeroResidualBlock` are *alternative* block types, not composed together + +**Location**: module docstring, opening paragraph. + +**Quote**: "Two block variants are provided: ResidualBlock … AdaLNZeroResidualBlock …" + +**Fix**: Add one sentence explicitly stating that a network uses one *or* the other per +stage/config — they are not nested. New collaborators may assume `AdaLNZeroResidualBlock` +somehow wraps `ResidualBlock`. + +______________________________________________________________________ + +### 2. `ResidualBlock.__init__` docstring: `condition_mixer_cfg` description does not name a concrete supported type + +**Location**: `ResidualBlock.__init__` Args, `condition_mixer_cfg`. + +**Quote**: "When active, the module's `forward` must accept `(x, condition)` positional +arguments." + +**Fix**: Name the concrete supported type — e.g. `ConditionMixer` from +`nvsubquadratic.modules.condition_mixer` — or note that as of this writing no concrete +condition mixer class exists in the public API (so `torch.nn.Identity` is effectively the +only safe value). Without this, a reader does not know what to pass. + +______________________________________________________________________ + +### 3. `ResidualBlock.forward`: the `condition` parameter should be typed `Optional[torch.Tensor]` + +**Location**: `forward` signature, line `def forward(self, x: torch.Tensor, condition: torch.Tensor)`. + +**Fix**: Change the type annotation to `condition: torch.Tensor | None` (or +`Optional[torch.Tensor]`) to match actual runtime semantics: `condition` is +silently ignored (and therefore safely `None`) when `condition_mixer` is +`torch.nn.Identity`. The current signature implies it is always required, which +contradicts the docstring prose. + +______________________________________________________________________ + +### 4. `ResidualBlock.forward` docstring: tensor shape for `condition` is too vague + +**Location**: `ResidualBlock.forward` Args, `condition`. + +**Quote**: "Its shape depends on the conditioning operator — a common choice is +`(B, *spatial_dims_condition, C)` for cross-attention, or `(B, C)` for a global +conditioning vector." + +**Fix**: This is reasonable, but add a note that `condition` must have the *same* `C` +(hidden channel dimension) as `x`, since the condition mixer is expected to project into +the residual stream width. Otherwise readers may pass a conditioning tensor with a +different channel count and get a confusing runtime error. + +______________________________________________________________________ + +### 5. `ResidualBlock.__init__`: the Raises section is missing a description for the case where norm/mixer Identity constraint is violated + +**Location**: `ResidualBlock.__init__` Raises, currently reads "AssertionError: If a norm +config does not match…" + +**Fix**: The assertion message says "Sequence mixer norm must be Identity if sequence mixer +is Identity" — but the *converse* (norm is Identity while mixer is not) is silently +allowed and would cause the norm to become a no-op. Clarify the *exact* direction of the +constraint enforced: mixer is Identity ⟹ norm must also be Identity. The reverse is not +checked and is the user's responsibility. + +______________________________________________________________________ + +### 6. `AdaLNZeroResidualBlock.__init__`: `hidden_dim` relationship to sub-module configs is not explained + +**Location**: `AdaLNZeroResidualBlock.__init__` Args, `hidden_dim`. + +**Quote**: "Channel dimension `C` shared by all sub-modules. Used to size the +`condition_proj` linear layer (`Linear(C, 6*C)`)." + +**Fix**: Note that `hidden_dim` must match the `dim`/`hidden_dim` argument baked into the +instantiated `sequence_mixer_cfg`, `mlp_cfg`, etc. There is no runtime check; a mismatch +will produce a shape error deep inside `condition_proj`, which is hard to trace back to a +`hidden_dim` mismatch at block init. Suggest either adding an assertion or at minimum +documenting that they must agree. + +______________________________________________________________________ + +### 7. `AdaLNZeroResidualBlock.forward`: the `expand` inner function is undocumented as a local helper + +**Location**: `AdaLNZeroResidualBlock.forward`, `def expand(param, ref)`. + +**Current docstring**: "Broadcast a \[B, hidden_dim\] vector across ref's spatial axes." +(one line, present). + +**Fix**: The one-liner is fine, but add a note that this helper is defined *inside* +`forward` to avoid materialising a persistent buffer. It is a purely local broadcast +utility, not a module method. This avoids readers wondering why it is not a `staticmethod` +or class method. + +______________________________________________________________________ + +### 8. `AdaLNZeroResidualBlock.forward` docstring: does not mention that `conditioning=cond` is passed *into* the sequence mixer + +**Location**: `AdaLNZeroResidualBlock.forward`, Returns/prose section. + +**Fix**: The forward pass calls `self.sequence_mixer(seq_mod, conditioning=cond)` — the +pooled conditioning vector is forwarded into the inner mixer (e.g. Hyena with FiLM). This +is a non-obvious secondary conditioning path beyond the AdaLN modulation itself. Document +it explicitly in the forward docstring so readers understand there are two conditioning +signals: AdaLN-Zero affine modulation (shift/scale/gate) *and* the raw `cond` vector +passed to the inner mixer. + +______________________________________________________________________ + +### 9. Both classes: no cross-reference to the network classes that instantiate them + +**Location**: class-level docstrings. + +**Fix**: Add a `See Also` or `Note` pointing to +`nvsubquadratic.networks.general_purpose_resnet` and/or +`nvsubquadratic.networks.classification_resnet` as the canonical consumers of these +blocks. External collaborators reading this file in isolation will not know where to look +to understand how blocks are stacked. + +______________________________________________________________________ + +### 10. `AdaLNZeroResidualBlock`: class docstring does not state that the condition branch (cross-attention) of `ResidualBlock` is *absent* + +**Location**: `AdaLNZeroResidualBlock` class docstring. + +**Fix**: Add a sentence noting that `AdaLNZeroResidualBlock` has no separate +`condition_mixer` cross-attention branch (unlike `ResidualBlock`). All conditioning is +handled through the single AdaLN-Zero projection. Without this, readers familiar with +`ResidualBlock` will wonder what happened to the condition-mixer slot. From b54f66c8664a80e63b17ee012b3fbaee32288ac0 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:25:02 +0200 Subject: [PATCH 17/72] docs(review/film): reviewer feedback From f996a0f9e0989b6a267cc7f3130d77088ad6c11e Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:25:19 +0200 Subject: [PATCH 18/72] docs(review/patchify): reviewer feedback From e21ab2c1228bf1bad43852ff66af670507011987 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:26:41 +0200 Subject: [PATCH 19/72] docs(integrate/residual_block): apply reviewer feedback --- docs/reviews/residual_block_review.md | 137 ----------------------- nvsubquadratic/modules/film.py | 70 +++++++----- nvsubquadratic/modules/residual_block.py | 108 ++++++++++++++---- 3 files changed, 125 insertions(+), 190 deletions(-) delete mode 100644 docs/reviews/residual_block_review.md diff --git a/docs/reviews/residual_block_review.md b/docs/reviews/residual_block_review.md deleted file mode 100644 index 0a71d265..00000000 --- a/docs/reviews/residual_block_review.md +++ /dev/null @@ -1,137 +0,0 @@ -# Review: `nvsubquadratic/modules/residual_block.py` - -Reviewed as a second reader looking for gaps that would confuse an external collaborator -reading alongside the research paper. - -______________________________________________________________________ - -## Issues - -### 1. Module docstring does not clarify that `ResidualBlock` and `AdaLNZeroResidualBlock` are *alternative* block types, not composed together - -**Location**: module docstring, opening paragraph. - -**Quote**: "Two block variants are provided: ResidualBlock … AdaLNZeroResidualBlock …" - -**Fix**: Add one sentence explicitly stating that a network uses one *or* the other per -stage/config — they are not nested. New collaborators may assume `AdaLNZeroResidualBlock` -somehow wraps `ResidualBlock`. - -______________________________________________________________________ - -### 2. `ResidualBlock.__init__` docstring: `condition_mixer_cfg` description does not name a concrete supported type - -**Location**: `ResidualBlock.__init__` Args, `condition_mixer_cfg`. - -**Quote**: "When active, the module's `forward` must accept `(x, condition)` positional -arguments." - -**Fix**: Name the concrete supported type — e.g. `ConditionMixer` from -`nvsubquadratic.modules.condition_mixer` — or note that as of this writing no concrete -condition mixer class exists in the public API (so `torch.nn.Identity` is effectively the -only safe value). Without this, a reader does not know what to pass. - -______________________________________________________________________ - -### 3. `ResidualBlock.forward`: the `condition` parameter should be typed `Optional[torch.Tensor]` - -**Location**: `forward` signature, line `def forward(self, x: torch.Tensor, condition: torch.Tensor)`. - -**Fix**: Change the type annotation to `condition: torch.Tensor | None` (or -`Optional[torch.Tensor]`) to match actual runtime semantics: `condition` is -silently ignored (and therefore safely `None`) when `condition_mixer` is -`torch.nn.Identity`. The current signature implies it is always required, which -contradicts the docstring prose. - -______________________________________________________________________ - -### 4. `ResidualBlock.forward` docstring: tensor shape for `condition` is too vague - -**Location**: `ResidualBlock.forward` Args, `condition`. - -**Quote**: "Its shape depends on the conditioning operator — a common choice is -`(B, *spatial_dims_condition, C)` for cross-attention, or `(B, C)` for a global -conditioning vector." - -**Fix**: This is reasonable, but add a note that `condition` must have the *same* `C` -(hidden channel dimension) as `x`, since the condition mixer is expected to project into -the residual stream width. Otherwise readers may pass a conditioning tensor with a -different channel count and get a confusing runtime error. - -______________________________________________________________________ - -### 5. `ResidualBlock.__init__`: the Raises section is missing a description for the case where norm/mixer Identity constraint is violated - -**Location**: `ResidualBlock.__init__` Raises, currently reads "AssertionError: If a norm -config does not match…" - -**Fix**: The assertion message says "Sequence mixer norm must be Identity if sequence mixer -is Identity" — but the *converse* (norm is Identity while mixer is not) is silently -allowed and would cause the norm to become a no-op. Clarify the *exact* direction of the -constraint enforced: mixer is Identity ⟹ norm must also be Identity. The reverse is not -checked and is the user's responsibility. - -______________________________________________________________________ - -### 6. `AdaLNZeroResidualBlock.__init__`: `hidden_dim` relationship to sub-module configs is not explained - -**Location**: `AdaLNZeroResidualBlock.__init__` Args, `hidden_dim`. - -**Quote**: "Channel dimension `C` shared by all sub-modules. Used to size the -`condition_proj` linear layer (`Linear(C, 6*C)`)." - -**Fix**: Note that `hidden_dim` must match the `dim`/`hidden_dim` argument baked into the -instantiated `sequence_mixer_cfg`, `mlp_cfg`, etc. There is no runtime check; a mismatch -will produce a shape error deep inside `condition_proj`, which is hard to trace back to a -`hidden_dim` mismatch at block init. Suggest either adding an assertion or at minimum -documenting that they must agree. - -______________________________________________________________________ - -### 7. `AdaLNZeroResidualBlock.forward`: the `expand` inner function is undocumented as a local helper - -**Location**: `AdaLNZeroResidualBlock.forward`, `def expand(param, ref)`. - -**Current docstring**: "Broadcast a \[B, hidden_dim\] vector across ref's spatial axes." -(one line, present). - -**Fix**: The one-liner is fine, but add a note that this helper is defined *inside* -`forward` to avoid materialising a persistent buffer. It is a purely local broadcast -utility, not a module method. This avoids readers wondering why it is not a `staticmethod` -or class method. - -______________________________________________________________________ - -### 8. `AdaLNZeroResidualBlock.forward` docstring: does not mention that `conditioning=cond` is passed *into* the sequence mixer - -**Location**: `AdaLNZeroResidualBlock.forward`, Returns/prose section. - -**Fix**: The forward pass calls `self.sequence_mixer(seq_mod, conditioning=cond)` — the -pooled conditioning vector is forwarded into the inner mixer (e.g. Hyena with FiLM). This -is a non-obvious secondary conditioning path beyond the AdaLN modulation itself. Document -it explicitly in the forward docstring so readers understand there are two conditioning -signals: AdaLN-Zero affine modulation (shift/scale/gate) *and* the raw `cond` vector -passed to the inner mixer. - -______________________________________________________________________ - -### 9. Both classes: no cross-reference to the network classes that instantiate them - -**Location**: class-level docstrings. - -**Fix**: Add a `See Also` or `Note` pointing to -`nvsubquadratic.networks.general_purpose_resnet` and/or -`nvsubquadratic.networks.classification_resnet` as the canonical consumers of these -blocks. External collaborators reading this file in isolation will not know where to look -to understand how blocks are stacked. - -______________________________________________________________________ - -### 10. `AdaLNZeroResidualBlock`: class docstring does not state that the condition branch (cross-attention) of `ResidualBlock` is *absent* - -**Location**: `AdaLNZeroResidualBlock` class docstring. - -**Fix**: Add a sentence noting that `AdaLNZeroResidualBlock` has no separate -`condition_mixer` cross-attention branch (unlike `ResidualBlock`). All conditioning is -handled through the single AdaLN-Zero projection. Without this, readers familiar with -`ResidualBlock` will wonder what happened to the condition-mixer slot. diff --git a/nvsubquadratic/modules/film.py b/nvsubquadratic/modules/film.py index de16889b..811c9216 100644 --- a/nvsubquadratic/modules/film.py +++ b/nvsubquadratic/modules/film.py @@ -13,6 +13,11 @@ parameters ``γ`` and ``β`` are **not** fixed — they are the outputs of a small neural network (the *FiLM generator*) evaluated at runtime. +In this codebase ``γ(c)`` and ``β(c)`` are vectors in ``ℝ^{kernel_hidden_dim}`` +applied pointwise to each SIREN hidden activation (no spatial axis); the standard +spatial-feature-map interpretation from Perez et al. applies when the kernel +network is evaluated independently at every spatial coordinate. + *Reference*: Perez et al., "FiLM: Visual Reasoning with a General Conditioning Layer", arXiv:1709.07871 (2017). @@ -28,30 +33,19 @@ generalise across equation parameters (viscosity, Reynolds number, etc.) by FiLM-conditioning the implicit neural operator kernel. * **SIREN kernel modulation** — within :mod:`nvsubquadratic.modules.kernels_nd`, - FiLM adjusts the hidden activations of each SIREN layer, effectively + FiLM adjusts the hidden activations of each SIREN layer (SIREN = Sinusoidal + Representation Network, Sitzmann et al. 2020, arXiv:2006.09661), effectively steering the learned convolution kernel per sample or per timestep. -**Conditioning signal sources in this codebase** - -The conditioning vector ``c ∈ ℝ^{cond_dim}`` typically originates from -register tokens appended to the sequence. Two encoder modules are provided: - -* :class:`RegisterPooling` — learnable softmax-weighted average over all - register tokens; produces a ``[B, C]`` summary vector. -* :class:`RegisterCompressConcat` — compresses each register token with a - shared linear and concatenates; preserves per-register identity at the cost - of a larger ``cond_dim``. +**Conditioning signal sources** -The encoded conditioning vector is then fed into :class:`KernelFiLMGenerator` -which maps it to ``(γ, β)`` pairs — one pair per SIREN hidden layer. - -Provides: -- KernelFiLMGenerator: MLP that maps a conditioning vector to per-layer (gamma, beta) pairs - for modulating SIREN hidden layers. -- RegisterPooling: Learnable weighted average over register tokens to produce a single - conditioning vector per sample. -- RegisterCompressConcat: Compress each register token via a shared linear layer and - concatenate, producing a conditioning vector that preserves per-register identity. +The conditioning vector ``c ∈ ℝ^{cond_dim}`` can come from any source — +a timestep MLP, class embedding lookup, or physics parameter encoder. +In diffusion or class-conditioned settings any ``[B, cond_dim]`` tensor is +accepted; within this codebase it typically originates from register tokens +appended to the sequence. See :class:`RegisterPooling` and +:class:`RegisterCompressConcat` for the two register encoders provided here, and +:class:`KernelFiLMGenerator` for the FiLM generator that consumes the result. """ from __future__ import annotations @@ -68,7 +62,9 @@ class KernelFiLMGenerator(nn.Module): Given a conditioning signal ``c ∈ ℝ^{cond_dim}`` (e.g. from register tokens processed by :class:`RegisterPooling` or :class:`RegisterCompressConcat`), - this module produces one ``(γ_l, β_l)`` pair per SIREN hidden layer ``l``: + this module produces one ``(γ_l, β_l)`` pair per SIREN hidden layer ``l`` + (SIREN = Sinusoidal Representation Network, Sitzmann et al. 2020, + arXiv:2006.09661; see :mod:`nvsubquadratic.modules.kernels_nd`): h_l ← γ_l(c) ⊙ h_l + β_l(c) @@ -95,7 +91,9 @@ class KernelFiLMGenerator(nn.Module): Attributes: num_film_layers (int): Number of ``(γ, β)`` pairs produced. kernel_hidden_dim (int): Feature dimension of each SIREN hidden layer. - mlp (nn.Sequential): Two-layer MLP (Linear → GELU → Linear). + mlp (nn.Sequential): Two-layer MLP mapping ``[*, cond_dim]`` → + ``[*, num_film_layers × 2 × kernel_hidden_dim]`` via a + ``film_hidden_dim``-dimensional bottleneck (Linear → GELU → Linear). Args: cond_dim: Dimensionality of the conditioning input ``c``. @@ -131,6 +129,8 @@ def __init__( # noqa: D107 init_type: Literal["identity", "small_random"] = "identity", init_std: float = 1e-4, ): + if num_film_layers < 1: + raise ValueError(f"num_film_layers must be >= 1, got {num_film_layers}") super().__init__() self.num_film_layers = num_film_layers self.kernel_hidden_dim = kernel_hidden_dim @@ -234,8 +234,9 @@ def forward(self, conditioning: torch.Tensor) -> list[tuple[torch.Tensor, torch. Returns: A list of ``num_film_layers`` tuples ``(gamma, beta)``, where each - tensor has shape ``[B, kernel_hidden_dim]``. Index ``l`` of the list - corresponds to SIREN hidden layer ``l``. + tensor has shape ``[B, kernel_hidden_dim]``. Index ``0`` corresponds + to the first (shallowest) SIREN hidden layer and index + ``num_film_layers - 1`` to the deepest. """ out = self.mlp(conditioning) # [B, num_film_layers * 2 * kernel_hidden_dim] chunks = out.chunk(self.num_film_layers, dim=-1) # num_film_layers x [B, 2 * kernel_hidden_dim] @@ -306,7 +307,8 @@ def forward(self, registers: torch.Tensor) -> torch.Tensor: registers: Register token tensor of shape ``[B, num_registers, C]``, where ``B`` is the batch size and ``C`` is the channel dimension. The number of registers along axis 1 must equal ``num_registers`` - passed to ``__init__``. + passed to ``__init__``; a mismatch will raise a ``RuntimeError`` + from ``torch.einsum``. Returns: Pooled conditioning vector of shape ``[B, C]``. @@ -341,6 +343,7 @@ class RegisterCompressConcat(nn.Module): Attributes: num_registers (int): Number of register tokens expected on the sequence axis. + hidden_dim (int): Input channel dimension of each register token. compressed_dim (int): Output channel dimension per register after compression. compress (nn.Linear): Shared weight-only (no bias) projection ``hidden_dim → compressed_dim`` applied independently to each register. @@ -355,6 +358,7 @@ class RegisterCompressConcat(nn.Module): def __init__(self, num_registers: int, hidden_dim: int, compressed_dim: int): # noqa: D107 super().__init__() self.num_registers = num_registers + self.hidden_dim = hidden_dim self.compressed_dim = compressed_dim self.compress = nn.Linear(hidden_dim, compressed_dim, bias=False) @@ -368,20 +372,21 @@ def out_dim(self) -> int: """ return self.num_registers * self.compressed_dim - def flop_count(self, dim: int) -> int: + def flop_count(self, hidden_dim: int) -> int: """Count FLOPs for compress-and-concatenate (one sample). - Operations (R = num_registers, D_in = ``dim``, D_out = compressed_dim): + Operations (R = num_registers, D_in = ``hidden_dim``, D_out = compressed_dim): 1. Shared linear applied R times: R * 2 * D_in * D_out 2. Concatenation is a view/copy, not a FLOP. Args: - dim: Channel dimension of the register tokens (should match ``hidden_dim``). + hidden_dim: Channel dimension of the register tokens; must equal the + ``hidden_dim`` argument passed to ``__init__``. Returns: Total FLOPs as an integer. """ - return self.num_registers * 2 * dim * self.compressed_dim + return self.num_registers * 2 * hidden_dim * self.compressed_dim def forward(self, registers: torch.Tensor) -> torch.Tensor: """Compress each register token and concatenate into a flat conditioning vector. @@ -404,4 +409,7 @@ def forward(self, registers: torch.Tensor) -> torch.Tensor: return compressed.flatten(start_dim=1) # [B, R * compressed_dim] def extra_repr(self) -> str: # noqa: D102 - return f"num_registers={self.num_registers}, compressed_dim={self.compressed_dim}, out_dim={self.out_dim}" + return ( + f"num_registers={self.num_registers}, hidden_dim={self.hidden_dim}, " + f"compressed_dim={self.compressed_dim}, out_dim={self.out_dim}" + ) diff --git a/nvsubquadratic/modules/residual_block.py b/nvsubquadratic/modules/residual_block.py index 5459346e..ccad415e 100644 --- a/nvsubquadratic/modules/residual_block.py +++ b/nvsubquadratic/modules/residual_block.py @@ -18,10 +18,13 @@ This module provides the fundamental repeating unit of the nvsubquadratic architecture: a residual block that wraps a *sequence mixer* and an *MLP* with pre-norm and optional conditioning branches. Stacking many such blocks -forms the full depth of the network (``general_purpose_resnet`` / -``classification_resnet``). +forms the full depth of the network (see +``nvsubquadratic.networks.general_purpose_resnet`` and +``nvsubquadratic.networks.classification_resnet``). -Two block variants are provided: +Two block variants are provided. They are **alternative** building blocks — +a given network stage uses one *or* the other, never both nested together. +Choosing between them is a top-level config decision: :class:`ResidualBlock` The standard pre-norm residual block used in most networks. Each forward @@ -35,16 +38,18 @@ ``torch.nn.Identity``; the corresponding norm must also be Identity. :class:`AdaLNZeroResidualBlock` - A DiT-style block that replaces LayerNorm with *Adaptive LayerNorm-Zero* - (AdaLN-Zero) modulation. A single zero-initialised linear layer maps the - conditioning vector to six affine parameters (shift + scale + gate for each - of the two branches), giving the conditioning signal fine-grained control - over both branches at initialisation-time stability (outputs ≈ 0). + A DiT-style block that replaces fixed LayerNorm with *Adaptive + LayerNorm-Zero* (AdaLN-Zero) modulation. Unlike :class:`ResidualBlock`, + this block has **no separate cross-attention / condition-mixer branch**; + all conditioning is routed through a single zero-initialised projection + that produces shift + scale + gate parameters for both the sequence-mixer + and MLP branches. This gives the conditioning signal fine-grained control + over both branches at initialisation-time stability (block outputs ≈ 0). All tensors follow **channels-last** layout: ``(B, *spatial_dims, C)``. """ -from typing import Union +from typing import Optional, Union import torch @@ -67,10 +72,22 @@ class ResidualBlock(torch.nn.Module): pure-sequence networks (condition branch disabled), cross-attention encoder-decoders (condition branch enabled), and MLP-only ablations. + The Identity constraint is **one-directional**: if a mixer/MLP config + targets ``Identity``, the corresponding norm config *must also* be + ``Identity`` (enforced by assertion at init). The reverse — a norm set + to ``Identity`` while the mixer is active — is not checked and is the + caller's responsibility. + All normalisation parameters (``input_norm``, ``condition_mixer_norm``, ``mlp_norm``) are tagged with ``_no_weight_decay = True`` so that the optimiser can exclude them from weight-decay groups. + See Also: + :class:`~nvsubquadratic.networks.general_purpose_resnet.GeneralPurposeResnet` + and + :class:`~nvsubquadratic.networks.classification_resnet.ClassificationResnet` + for the canonical consumers of this block. + Attributes: sequence_mixer (torch.nn.Module): The instantiated sequence-mixing operator (e.g. :class:`~nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer` @@ -81,6 +98,9 @@ class ResidualBlock(torch.nn.Module): condition_mixer (torch.nn.Module): Cross-attention or conditioning operator applied after the sequence mixer. May be ``torch.nn.Identity`` to disable the conditioning branch entirely. + When not Identity, its ``forward(x, condition)`` signature must + accept the residual stream and a conditioning tensor whose channel + dimension matches ``C`` (the hidden channel dim of ``x``). condition_mixer_norm (torch.nn.Module): Pre-norm applied before ``condition_mixer``. Must be ``torch.nn.Identity`` when ``condition_mixer`` is ``torch.nn.Identity``. @@ -124,8 +144,12 @@ def __init__( condition_mixer_cfg: LazyConfig for the conditioning / cross- attention operator. Pass ``torch.nn.Identity`` (the default in most configs) to disable the conditioning branch. When - active, the module's ``forward`` must accept - ``(x, condition)`` positional arguments. + active, the instantiated module's ``forward`` must accept + ``(x, condition)`` positional arguments where ``condition`` + has the same channel dimension ``C`` as ``x``. As of this + writing, ``torch.nn.Identity`` is effectively the only + production-ready value; a concrete cross-attention class is + planned but not yet part of the public API. condition_mixer_norm_cfg: LazyConfig for the pre-norm applied before ``condition_mixer``. **Must** be ``torch.nn.Identity`` when ``condition_mixer_cfg`` targets ``torch.nn.Identity``. @@ -142,10 +166,10 @@ def __init__( ``torch.nn.Identity`` for no dropout. Raises: - AssertionError: If a norm config does not match its corresponding - module config — i.e. if a mixer/MLP config is ``Identity`` but - the corresponding norm config is not (or vice versa in the - forward-only direction). + AssertionError: If a mixer/MLP config targets ``torch.nn.Identity`` + but the corresponding norm config does not. The constraint is + one-directional: mixer=Identity ⟹ norm=Identity. The reverse + (norm=Identity while mixer is active) is not checked. """ if sequence_mixer_cfg.__target__ == torch.nn.Identity: assert sequence_mixer_norm_cfg.__target__ == torch.nn.Identity, ( @@ -186,7 +210,7 @@ def __init__( # Instantiate dropout self.dropout = instantiate(dropout_cfg) - def forward(self, x: torch.Tensor, condition: torch.Tensor) -> torch.Tensor: + def forward(self, x: torch.Tensor, condition: Optional[torch.Tensor]) -> torch.Tensor: """Apply the residual block to the input tensor. Executes up to three residual sub-branches in order: sequence mixer, @@ -220,7 +244,8 @@ def forward(self, x: torch.Tensor, condition: torch.Tensor) -> torch.Tensor: condition: Conditioning tensor used by ``condition_mixer``. Its shape depends on the conditioning operator — a common choice is ``(B, *spatial_dims_condition, C)`` for cross-attention, or - ``(B, C)`` for a global conditioning vector. This argument is + ``(B, C)`` for a global conditioning vector. The channel + dimension ``C`` must match that of ``x``. This argument is **ignored** (and may safely be ``None``) when ``condition_mixer`` is ``torch.nn.Identity``. @@ -270,6 +295,19 @@ class AdaLNZeroResidualBlock(torch.nn.Module): for each of the two branches — so that at initialisation the block outputs exactly zero (the residual stream is unchanged). + Unlike :class:`ResidualBlock`, this block has **no separate + cross-attention / condition-mixer branch**. All conditioning is routed + through the AdaLN-Zero projection plus a ``conditioning`` kwarg forwarded + directly into the sequence mixer (e.g. for FiLM inside Hyena). There are + therefore **two conditioning pathways** in a single forward pass: + + 1. **AdaLN-Zero affine modulation** — shift/scale/gate applied to the + pre-norm outputs of the sequence mixer and MLP branches. + 2. **Inner-mixer conditioning** — the pooled conditioning vector ``cond`` + is also passed as ``conditioning=cond`` to ``sequence_mixer.forward``, + allowing the inner operator (e.g. :class:`~nvsubquadratic.modules.hyena_nd.Hyena` + with FiLM kernel conditioning) to use it independently. + **Forward computation** (one block): .. code-block:: text @@ -281,7 +319,7 @@ class AdaLNZeroResidualBlock(torch.nn.Module): # Sequence mixer branch x_norm = sequence_norm(x) # pre-norm x_mod = x_norm * (1 + scale_seq) + shift_seq # AdaLN modulation - seq_out = sequence_mixer(x_mod, conditioning=cond) + seq_out = sequence_mixer(x_mod, conditioning=cond) # cond also forwarded seq_out = dropout(seq_out) * gate_seq # zero-init gate x = x + seq_out @@ -296,10 +334,16 @@ class AdaLNZeroResidualBlock(torch.nn.Module): scaling; at init they are zero (because ``condition_proj`` weights are zero-initialised), so the block is a skip connection. + See Also: + :class:`~nvsubquadratic.networks.general_purpose_resnet.GeneralPurposeResnet` + and + :class:`~nvsubquadratic.networks.classification_resnet.ClassificationResnet` + for the canonical consumers of this block. + Attributes: sequence_mixer (torch.nn.Module): Instantiated sequence-mixing operator. Its ``forward`` must accept a ``conditioning`` keyword - argument (forwarded from the pooled conditioning vector). + argument (forwarded from the pooled conditioning vector ``cond``). sequence_norm (torch.nn.Module): Pre-norm applied to ``x`` before AdaLN modulation in the sequence mixer branch. mlp (torch.nn.Module): Position-wise MLP instantiated from @@ -336,7 +380,7 @@ def __init__( pooled conditioning vector ``cond`` of shape ``(B, C)``). sequence_mixer_norm_cfg: LazyConfig for the pre-norm applied before the AdaLN modulation of the sequence mixer branch - (e.g. LayerNorm over ``hidden_dim``). + (e.g. ``LayerNorm(hidden_dim)``). mlp_cfg: LazyConfig for the position-wise MLP. mlp_norm_cfg: LazyConfig for the pre-norm applied before the AdaLN modulation of the MLP branch. @@ -347,7 +391,10 @@ def __init__( applied after each branch and before the gate multiply. hidden_dim: Channel dimension ``C`` shared by all sub-modules. Used to size the ``condition_proj`` linear layer - (``Linear(C, 6*C)``). + (``Linear(C, 6*C)``). **This value must match the channel + dimension baked into** ``sequence_mixer_cfg`` **and** + ``mlp_cfg`` — there is no runtime check, and a mismatch will + produce a shape error deep inside ``condition_proj``. """ super().__init__() @@ -385,6 +432,17 @@ def forward(self, x: torch.Tensor, condition: Union[torch.Tensor, None]) -> torc outputs of both branches via element-wise affine transforms, and gate each branch's output before the residual add. + There are **two conditioning pathways**: + + 1. **AdaLN-Zero** — shift/scale/gate parameters derived from + ``condition_proj(cond)`` are applied to the pre-norm outputs of + both the sequence-mixer and MLP branches. + 2. **Inner-mixer conditioning** — the same pooled ``cond`` vector is + also forwarded as ``conditioning=cond`` into + ``self.sequence_mixer``, allowing operators such as + :class:`~nvsubquadratic.modules.hyena_nd.Hyena` to apply + additional FiLM modulation to their kernel networks. + Args: x: Input feature tensor of shape ``(B, *spatial_dims, C)`` where ``B`` is the batch size, ``spatial_dims`` are one or more @@ -422,7 +480,12 @@ def forward(self, x: torch.Tensor, condition: Union[torch.Tensor, None]) -> torc ) # each (B, hidden_dim) def expand(param: torch.Tensor, ref: torch.Tensor) -> torch.Tensor: - """Broadcast a [B, hidden_dim] vector across ref's spatial axes.""" + """Broadcast a [B, hidden_dim] vector across ref's spatial axes. + + Defined as a local helper (rather than a method) to keep the + broadcast logic inline with the forward pass without allocating + a persistent buffer. + """ while param.ndim < ref.ndim: param = param.unsqueeze(1) # add singleton dims for broadcasting return param.expand(*ref.shape[:-1], param.shape[-1]) # match ref spatial layout @@ -432,6 +495,7 @@ def expand(param: torch.Tensor, ref: torch.Tensor) -> torch.Tensor: seq_mod = seq_norm * (1.0 + expand(scale_seq, seq_norm)) + expand( shift_seq, seq_norm ) # (B, *spatial_dims, hidden_dim) + # cond is also forwarded to the inner mixer (e.g. for Hyena FiLM kernel conditioning). seq_out = self.sequence_mixer(seq_mod, conditioning=cond) # (B, *spatial_dims, hidden_dim) seq_out = self.dropout(seq_out) # (B, *spatial_dims, hidden_dim) seq_out = seq_out * expand(gate_seq, seq_out) # (B, *spatial_dims, hidden_dim) From c3072c3def4ede20e536b70ef453c1092a3d2b6e Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:26:42 +0200 Subject: [PATCH 20/72] docs(integrate/patchify): apply reviewer feedback --- docs/reviews/patchify_review.md | 138 ----------------------------- nvsubquadratic/modules/patchify.py | 90 ++++++++++++++----- 2 files changed, 68 insertions(+), 160 deletions(-) delete mode 100644 docs/reviews/patchify_review.md diff --git a/docs/reviews/patchify_review.md b/docs/reviews/patchify_review.md deleted file mode 100644 index d73e9248..00000000 --- a/docs/reviews/patchify_review.md +++ /dev/null @@ -1,138 +0,0 @@ -# Review: `nvsubquadratic/modules/patchify.py` - -Reviewer pass on commit `docs(write/patchify): add module and class docstrings`. - -______________________________________________________________________ - -## Issues - -### 1. Module docstring — overlapping-patch semantics need a concrete example - -**Location:** Module docstring, "Overlapping patches" paragraph in `Patchify`. - -**Problem:** The docstring says "setting `stride < patch_size` produces overlapping patches -with the same formula above" but does not show the formula for that case. An external -collaborator who only reads the module docstring (not the class docstring) sees no formula at all. - -**Fix:** Add the general output-size formula `floor((s - patch_size) / stride) + 1` to -the module docstring overview, directly after the description of the output layout convention. - -______________________________________________________________________ - -### 2. `Patchify.__init__` — `patch_size` docstring does not mention the per-patch pixel count - -**Location:** `Patchify.__init__` docstring, `patch_size` arg entry. - -**Quote:** `"Side length P of each patch. The convolution uses kernel_size = patch_size along every spatial axis, so each patch covers P^data_dim input pixels."` - -**Problem:** The phrase "each patch covers P^data_dim input pixels" uses caret notation that -renders as plain text in most docstring viewers. Readers unfamiliar with the notation may -think `^` is a bitwise XOR. - -**Fix:** Replace `P^data_dim` with `P**data_dim` (Python exponentiation) or spell it out -as "`P × P` pixels (2D)" / "`P × P × P` voxels (3D)" with a note that the general -formula is `P ** data_dim`. - -______________________________________________________________________ - -### 3. `Patchify` class docstring — no mention of what happens when spatial size is not divisible by `patch_size` - -**Location:** `Patchify` class docstring, immediately after the output shape formula. - -**Problem:** The formula `s / patch_size` (non-overlapping case) implicitly assumes exact -divisibility, but the layer silently truncates if `s % patch_size != 0` (standard Conv -floor semantics). A new collaborator who passes a 224×224 image with `patch_size=16` -gets 14×14 tokens as expected, but a 225×224 image would give 14×14 tokens silently -dropping the last pixel column. This is a frequent source of bugs. - -**Fix:** Add a warning note under the output shape formula, e.g.: - -> **Note:** If `spatial_dim % patch_size != 0`, the last pixels in that axis are silently -> discarded (standard floor-division Conv semantics). Callers are responsible for ensuring -> the spatial dimensions are divisible by `patch_size` (e.g. via padding before this layer). - -______________________________________________________________________ - -### 4. `Unpatchify.forward` — `output_spatial_shape` description buries the most important use-case - -**Location:** `Unpatchify.forward` docstring, `output_spatial_shape` arg description. - -**Quote:** `"Must have length data_dim. When None, PyTorch infers the output size (may differ from the original input spatial size if spatial_dim % patch_size != 0)."` - -**Problem:** The caveat about size ambiguity is in the middle of a long sentence. It is the -most important reason to pass `output_spatial_shape`, but a skimming reader will miss it. - -**Fix:** Restructure the description to lead with the ambiguity problem: - -> When `stride > 1` multiple patch-grid sizes map to the same output size (floor in the -> forward direction discards remainders). Pass `output_spatial_shape` to resolve this -> ambiguity and guarantee recovery of the exact original spatial dimensions. Must have -> length `data_dim`. When `None`, the output size is inferred by PyTorch and may not -> match the original spatial size if `spatial_dim % patch_size != 0`. - -______________________________________________________________________ - -### 5. `Unpatchify` class docstring — overlapping-patch adjoint semantics are ambiguous - -**Location:** `Unpatchify` class docstring, second paragraph. - -**Quote:** `"When stride < patch_size (overlapping), contributions from overlapping patches are *summed* by the transposed convolution — this is the adjoint of the overlapping-patch forward pass."` - -**Problem:** The summation is correct mathematically, but "adjoint" may mislead readers into -thinking `Unpatchify` perfectly reconstructs the original signal. For overlapping patches -the transposed conv is only the linear adjoint (backward map), not an inverse; pixel values -are *accumulated*, not averaged, so the output scale grows with the overlap ratio. This is -especially unexpected for users who pair `Patchify` and `Unpatchify` in an autoencoder. - -**Fix:** Add an explicit note: - -> For overlapping patches this accumulation means `Unpatchify(Patchify(x))` does **not** -> recover `x` exactly; the output is a blurred, scaled version of `x`. Only for -> non-overlapping patches (`stride == patch_size`) does the round-trip preserve spatial -> alignment (up to the learned weights). - -______________________________________________________________________ - -### 6. `Unpatchify.__init__` — `weight_init` tradeoffs are missing for the `"default"` option - -**Location:** `Unpatchify.__init__` docstring, `weight_init` arg, `"default"` option. - -**Quote:** `'"default" uses PyTorch's built-in kaiming_uniform (fan computed from out_features; can cause output-variance blow-up for large in_features).'` - -**Problem:** The existing text warns about variance blow-up for `"default"` but does not tell -the reader *when* to choose it anyway. An external collaborator who has not read the -DiT/ViT literature will default to `"default"` without knowing the risk. - -**Fix:** Add a recommendation, e.g.: - -> Prefer `"fan_in"` for new architectures. `"default"` is retained for loading -> pre-trained checkpoints whose weights were saved under PyTorch's default init. - -______________________________________________________________________ - -### 7. Missing cross-reference to `PositionEmbeddingND` in module docstring - -**Location:** Module docstring. - -**Problem:** The module docstring mentions that the output layout matches what -"PositionEmbeddingND and the mixer blocks" expect, but does not tell the reader -where to find `PositionEmbeddingND`. - -**Fix:** Add an explicit cross-reference: - -> See `nvsubquadratic.modules.position_encoding.PositionEmbeddingND` for the positional -> encoding layer that is typically applied immediately after `Patchify`. - -______________________________________________________________________ - -### 8. `Patchify` and `Unpatchify` class docstrings — no ND usage example - -**Location:** Both class docstrings. - -**Problem:** The class-level docstrings show only implied 2D usage (the shapes use `H, W` -notation). A reader working on 1D sequences or 3D volumes must infer the generalisation -from the `data_dim` parameter alone, with no concrete example. - -**Fix:** Add a brief "Examples" section to at least one of the classes showing a 1D case -(`[B, L, C]` → `[B, L/P, C_embed]`) and referencing the `__main__` block for the 2D -demonstration. diff --git a/nvsubquadratic/modules/patchify.py b/nvsubquadratic/modules/patchify.py index c069fca3..f6f92eb7 100644 --- a/nvsubquadratic/modules/patchify.py +++ b/nvsubquadratic/modules/patchify.py @@ -27,9 +27,21 @@ input : [B, *spatial_dims, C_in] e.g. [B, H, W, C_in] for 2D output : [B, *patch_grid, C_embed] e.g. [B, H/P, W/P, C_embed] +The general output-size formula for each spatial axis ``s`` is:: + + out_s = floor((s - patch_size) / stride) + 1 + +For the default non-overlapping case (``stride == patch_size``) this reduces to +``s // patch_size`` when ``s`` is evenly divisible by ``patch_size``. + Channels are reordered to channels-first only internally before the convolution and back to channels-last before returning, which matches the layout expected by -``PositionEmbeddingND`` and the subsequent mixer blocks. +``nvsubquadratic.modules.position_encoding.PositionEmbeddingND`` and the +subsequent mixer blocks. + +See ``nvsubquadratic.modules.position_encoding.PositionEmbeddingND`` for the +positional encoding layer that is typically applied immediately after +``Patchify``. Supported dimensionalities -------------------------- @@ -69,8 +81,8 @@ class Patchify(torch.nn.Module): patches and linearly projects each patch into an embedding vector. The operation is equivalent to: - 1. Unfold every ``patch_size^data_dim`` pixel neighbourhood into a vector - of length ``C_in * patch_size^data_dim``. + 1. Unfold every ``patch_size ** data_dim`` pixel neighbourhood into a vector + of length ``C_in * patch_size ** data_dim``. 2. Apply a learned linear map from that vector to ``C_out`` dimensions. Because the unfold and linear projection can be fused into a single strided @@ -82,9 +94,15 @@ class Patchify(torch.nn.Module): out_s = floor((s - patch_size) / stride) + 1 For the default non-overlapping case (``stride == patch_size``) this - reduces to ``s / patch_size`` (assuming ``s`` is divisible by + reduces to ``s // patch_size`` (assuming ``s`` is divisible by ``patch_size``). + .. warning:: + If ``spatial_dim % patch_size != 0``, the last pixels in that axis are + silently discarded (standard floor-division Conv semantics). Callers + are responsible for ensuring spatial dimensions are divisible by + ``patch_size`` before calling this layer (e.g. by padding the input). + **Layout convention** — inputs and outputs use *channels-last* ordering:: input : [B, *spatial_dims, C_in] (e.g. [B, H, W, C_in] for 2D) @@ -98,6 +116,25 @@ class Patchify(torch.nn.Module): overlapping patches with the same formula above. This is less common in ViT-style models but is supported. + Examples: + 1D sequence (``data_dim=1``):: + + layer = Patchify(in_features=64, out_features=128, data_dim=1, patch_size=4) + x = torch.randn(2, 256, 64) # [B, L, C_in] + y = layer(x) # [B, L/4, 128] == [2, 64, 128] + + 2D image (``data_dim=2``) — see the ``__main__`` block for a runnable demo:: + + layer = Patchify(in_features=3, out_features=768, data_dim=2, patch_size=16) + x = torch.randn(8, 224, 224, 3) # [B, H, W, C_in] + y = layer(x) # [B, 14, 14, 768] + + 3D volume (``data_dim=3``):: + + layer = Patchify(in_features=1, out_features=256, data_dim=3, patch_size=8) + x = torch.randn(2, 64, 64, 64, 1) # [B, D, H, W, C_in] + y = layer(x) # [B, 8, 8, 8, 256] + Attributes: data_dim (int): Spatial dimensionality (1, 2, or 3). patch_size (int): Receptive field size of each patch along every axis. @@ -123,7 +160,8 @@ def __init__( (sequences), 2 (images), or 3 (volumes). patch_size: Side length ``P`` of each patch. The convolution uses ``kernel_size = patch_size`` along every spatial axis, so each - patch covers ``P^data_dim`` input pixels. + patch covers ``P ** data_dim`` input pixels (``P × P`` for 2D + images, ``P × P × P`` voxels for 3D volumes). stride: Step size between consecutive patch origins along every spatial axis. Defaults to ``patch_size``, giving non-overlapping ViT-style patches. Set to a smaller value for @@ -202,8 +240,13 @@ class Unpatchify(torch.nn.Module): transposed convolution is an exact spatial inverse: each output pixel is produced by exactly one input token. When ``stride < patch_size`` (overlapping), contributions from overlapping patches are *summed* by the - transposed convolution — this is the adjoint of the overlapping-patch - forward pass. + transposed convolution — this is the linear adjoint (backward map) of the + overlapping-patch forward pass, **not** a true inverse. Pixel values are + accumulated rather than averaged, so ``Unpatchify(Patchify(x))`` does not + recover ``x`` exactly for overlapping patches; the output is a blurred, + scaled version of ``x``. Only for non-overlapping patches + (``stride == patch_size``) does the round-trip preserve spatial alignment + (up to the learned weights). **Output shape formula** (each spatial axis ``s`` of the patch-grid input):: @@ -218,10 +261,10 @@ class Unpatchify(torch.nn.Module): output : [B, *spatial_dims, C_out] (e.g. [B, H, W, C_out]) **Weight initialisation** — PyTorch's default kaiming_uniform for - ``ConvTranspose`` uses ``fan_out = out_features * patch_size^data_dim``. + ``ConvTranspose`` uses ``fan_out = out_features * patch_size ** data_dim``. This is incorrect for large embedding dimensions; ``weight_init="fan_in"`` corrects this by using the true fan-in - ``in_features * patch_size^data_dim``. + ``in_features * patch_size ** data_dim``. Attributes: data_dim (int): Spatial dimensionality (1, 2, or 3). @@ -260,11 +303,14 @@ def __init__( weight_init: Weight initialisation strategy for the deconv kernel. ``"default"`` uses PyTorch's built-in ``kaiming_uniform`` (fan computed from ``out_features``; can cause output-variance - blow-up for large ``in_features``). ``"zeros"`` zero-inits - weights and bias (DiT-style; output is exactly zero at - initialisation, safe for residual-stream entry). ``"fan_in"`` - applies Kaiming-uniform with the corrected fan-in - ``in_features * patch_size^data_dim``, giving output variance + blow-up for large ``in_features``; retained primarily for + loading pre-trained checkpoints whose weights were saved under + PyTorch's default init — prefer ``"fan_in"`` for new + architectures). ``"zeros"`` zero-inits weights and bias + (DiT-style; output is exactly zero at initialisation, safe for + residual-stream entry). ``"fan_in"`` applies Kaiming-uniform + with the corrected fan-in + ``in_features * patch_size ** data_dim``, giving output variance O(1) regardless of embedding dimension. Raises: @@ -313,14 +359,14 @@ def forward(self, x: torch.Tensor, output_spatial_shape: Tuple[int, ...] | None ``[B, *patch_grid, C_embed]``, e.g. ``[B, H/P, W/P, C_embed]`` for 2D. The number of spatial dimensions must equal ``data_dim``. - output_spatial_shape: Optional target spatial shape for the output. - When provided this is passed as ``output_size`` to the - transposed convolution, which resolves the output-size ambiguity - that arises when ``stride > 1`` (multiple input sizes can map to - the same output size via the floor in the forward formula). - Must have length ``data_dim``. When ``None``, PyTorch infers - the output size (may differ from the original input spatial size - if ``spatial_dim % patch_size != 0``). + output_spatial_shape: When ``stride > 1``, multiple patch-grid + sizes map to the same output size (the floor in the forward + direction discards remainders). Pass ``output_spatial_shape`` + to resolve this ambiguity and guarantee recovery of the exact + original spatial dimensions. Must have length ``data_dim``. + When ``None``, PyTorch infers the output size and it may not + match the original spatial size if + ``spatial_dim % patch_size != 0``. Returns: Reconstructed signal tensor in channels-last layout. Shape: From 2eccb074319981c12e8e3926fa2efbd4a2a95922 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:27:05 +0200 Subject: [PATCH 21/72] docs(write/ckconv_nd): add module and class docstrings with math context --- nvsubquadratic/modules/ckconv_nd.py | 505 ++++++++++++++++++++++++---- 1 file changed, 431 insertions(+), 74 deletions(-) diff --git a/nvsubquadratic/modules/ckconv_nd.py b/nvsubquadratic/modules/ckconv_nd.py index afb6f427..804c148f 100644 --- a/nvsubquadratic/modules/ckconv_nd.py +++ b/nvsubquadratic/modules/ckconv_nd.py @@ -1,7 +1,109 @@ # TODO: Add license header here -"""CKConv (long-convolution) implementation for ND signals.""" +r"""Continuous Kernel Convolution (CKConv) for N-dimensional signals. + +Background +---------- +CKConv (Romero et al., "CKConv: Continuous Kernel Convolution With Arbitrary +Resolution", ICLR 2022, arXiv:2102.02611) is a long-range convolutional +operator whose filter is **not** a learnable lookup table but instead an +*implicit neural representation* of the kernel: a small MLP maps continuous +spatial coordinates to kernel values on the fly. + +Formally, for a D-dimensional input signal x ∈ R^{C × *spatial}, the CKConv +output is: + +.. code-block:: none + + y = x * k_θ + +where ``*`` denotes (multi-dimensional) convolution, and the kernel + +.. code-block:: none + + k_θ(p) = MLP_θ(pos_enc(p)) + +is evaluated at every position p on a continuous spatial grid normalised to +``[-1, 1]^D``. Because the kernel is defined by a fixed-size MLP rather than +an explicit filter bank, **the parameter count is independent of the signal +length**, and the same model can be evaluated at any resolution without weight +surgery (resolution-independence). Long-range dependencies are captured +because the MLP can, in principle, place non-zero weight at any relative +offset spanning the full input — there is no fixed receptive field. + +The convolution itself is computed in the frequency domain via FFT: + +.. code-block:: none + + y = IFFT( FFT(x) ⊙ FFT(k_θ) ) + +which costs O(N log N) per channel rather than O(N·K) for a spatial +convolution with kernel size K = N. This is the same subquadratic trick that +powers Hyena (see ``nvsubquadratic.modules.hyena_nd``). + +ND generalisation +----------------- +This module extends the original 1D CKConv formulation to **arbitrary spatial +rank** (1D sequences, 2D images, 3D volumes) by routing the global convolution +through the FFT primitives in ``nvsubquadratic.ops.fftconv``. The kernel +generator from ``nvsubquadratic.modules.kernels_nd`` evaluates the MLP on a +D-dimensional coordinate grid; the resulting filter is passed directly to the +FFT convolution ops. + +Boundary conditions +------------------- +Two boundary conditions are supported per spatial axis: + +* **Zero-padding** (``fft_padding="zero"``): standard linear convolution with + "same" output size. The kernel size is ``2*N`` (double-grid), covering the + full receptive field without wrap-around. Matches ``torch.nn.ConvNd(padding='same')`` + semantics. + +* **Circular / periodic** (``fft_padding="circular"``): wrap-around convolution + for signals with periodic boundary conditions (e.g. longitude in climate + models, ARC grid problems). The kernel size equals the input size + (single-grid, ``(N+1)//2`` grid points evaluated per axis). + +* **Per-axis mixed** (``fft_padding=["circular", "zero"]`` etc.): one mode + string per spatial axis. Routes through ``nvsubquadratic.ops.mixed_fftconv``. + Used for datasets such as Well's ``rayleigh_benard`` and + ``viscoelastic_instability``, where some axes are periodic and others are + not. + +* **Causal** (``is_causal=True``, 1D only): output at position ``n`` only + depends on inputs at positions ``0, …, n``. The kernel is evaluated on the + double-grid and then cropped to its causal (positive-lag) half before the + FFT convolution. + +Shortcut (skip connection) +-------------------------- +Every forward pass adds a per-channel residual term: + +.. code-block:: none + + y ← y + shortcut ⊙ x + +where ``shortcut`` is a learnable ``[hidden_dim]`` parameter vector. This +algebraic shortcut is fused into the FFT convolution op (no extra kernel +launch) and matches the design used throughout Hyena-style operators. + +Related modules +--------------- +* ``nvsubquadratic.modules.kernels_nd`` — implicit kernel parametrisation + (``SIRENKernelND``, ``RandomFourierKernelND``, FiLM-conditioned variants) +* ``nvsubquadratic.ops.fftconv`` — FFT convolution primitives consumed here +* ``nvsubquadratic.ops.mixed_fftconv`` — per-axis mixed-BC FFT convolution +* ``nvsubquadratic.modules.hyena_nd`` — Hyena operator that wraps CKConvND + as its global conv + +References: +---------- +* Romero et al. (2022). *CKConv: Continuous Kernel Convolution With Arbitrary + Resolution*. ICLR 2022. https://arxiv.org/abs/2102.02611 +* Sitzmann et al. (2020). *Implicit Neural Representations with Periodic + Activation Functions*. NeurIPS 2020. (SIREN kernel) +""" import copy import inspect @@ -203,7 +305,19 @@ def _parse_padding_mode(mode: str) -> bool: - """Map a single padding-mode string to its per-axis periodic flag.""" + """Map a single padding-mode string to its per-axis periodic flag. + + Args: + mode: A single padding-mode string, one of ``"zero"`` or ``"circular"`` + (case-insensitive, whitespace-stripped). + + Returns: + ``True`` if the mode is ``"circular"`` (periodic axis), ``False`` if + it is ``"zero"`` (zero-padded axis). + + Raises: + ValueError: If ``mode`` is not one of the recognised padding modes. + """ normalised = mode.strip().lower() if normalised not in _PADDING_MODE_TO_PERIODIC: valid = sorted(_PADDING_MODE_TO_PERIODIC) @@ -242,6 +356,19 @@ def _resolve_periodic( we keep a single canonical per-axis form to avoid two ways of saying the same thing. + Args: + fft_padding: Boundary-condition specification, either a single mode + string (``"zero"`` / ``"circular"``) or a sequence of mode strings + of length ``data_dim`` (one per spatial axis). + data_dim: Number of spatial dimensions in the input signal. Used to + validate the length of a sequence-form ``fft_padding`` and to + broadcast a scalar mode string. + + Returns: + A tuple of ``data_dim`` booleans, one per spatial axis. ``True`` + means that axis uses circular (periodic) convolution; ``False`` means + zero-padded (linear) convolution. + Raises: ValueError: on invalid mode strings, wrong number of axes, or disallowed input types. @@ -298,6 +425,18 @@ def _wrap_mixed_op(op_fn, periodic: tuple[bool, ...]): ``kernel`` and ``shortcut``. ``CKConvND.apply_convolution`` calls the FFT function as ``fn(x, kernel, shortcut)``; this wrapper binds ``periodic`` so the rest of the module is unchanged from the legacy ops. + + Args: + op_fn: A ``mixed_fftconv*`` function from + ``nvsubquadratic.ops.mixed_fftconv`` with signature + ``(x, kernel, periodic, shortcut)``. + periodic: Per-axis periodicity flags of length ``data_dim``. ``True`` + on an axis means circular (wrap-around) convolution; ``False`` + means zero-padded (linear) convolution. + + Returns: + A callable with signature ``(x, kernel, shortcut)`` that internally + forwards ``periodic`` to ``op_fn``. """ def _wrapped(x, kernel, shortcut): @@ -324,6 +463,17 @@ def _grid_is_single_per_axis( auto-derived as ``periodic[d]`` (periodic axis ⇒ single grid; non- periodic ⇒ double grid). This matches the recipe in :func:`nvsubquadratic.ops.mixed_fftconv._mixed_recipe`. + + Args: + grid_type: Either ``"single"`` (kernel size == input size, for circular + conv), ``"double"`` (kernel size == 2 × input size, for zero-padded + conv), or ``None`` to auto-derive per axis from ``periodic``. + periodic: Per-axis periodicity flags of length ``data_dim``. + + Returns: + A tuple of booleans of length ``len(periodic)``. ``True`` on axis + ``d`` means that axis uses the single grid (kernel size == input size); + ``False`` means double grid (kernel size == 2 × input size). """ if grid_type is None: return tuple(periodic) @@ -331,7 +481,82 @@ def _grid_is_single_per_axis( class CKConvND(torch.nn.Module): - """CKConv (long-convolution) implementation for ND signals.""" + """N-dimensional Continuous Kernel Convolution (CKConv) operator. + + CKConvND implements the CKConv operator (Romero et al., arXiv:2102.02611) + generalised to arbitrary spatial rank D ∈ {1, 2, 3}. The convolutional + kernel is **not** stored as an explicit lookup table; instead it is + produced on the fly by a small MLP evaluated on a continuous positional + grid: + + .. code-block:: none + + k_θ(p) = MLP_θ(pos_enc(p)), p ∈ [-1, 1]^D + + The convolution with the input signal x is then computed in the frequency + domain: + + .. code-block:: none + + y = IFFT( FFT(x) ⊙ FFT(k_θ) ) + shortcut ⊙ x + + at O(N log N) cost per channel, where N = prod(spatial_dims). + + The MLP and its positional encoding are provided through ``kernel_cfg`` + (typically a ``SIRENKernelND`` or ``RandomFourierKernelND`` lazy config + from ``nvsubquadratic.modules.kernels_nd``). An optional attenuation mask + ``mask_cfg`` (e.g. ``GaussianModulationND``) is applied to the kernel + values after the MLP forward pass to restrict the effective receptive + field at initialisation. + + **Boundary conditions** are controlled jointly by ``fft_padding`` and + ``grid_type``: + + * ``fft_padding="zero", grid_type="double"``: standard linear convolution + with "same" output size. The kernel spans the full input (double-grid: + ``2*N`` points) and is zero-padded before the FFT. + * ``fft_padding="circular", grid_type="single"``: periodic convolution, + kernel size == input size (single-grid: ``(N+1)//2`` grid points → kernel + of length N after the MLP). + * ``fft_padding=["circular", "zero"], grid_type=None``: per-axis mixed + boundary conditions; ``grid_type`` is auto-derived per axis. + + **Context parallelism**: when ``cp_group`` is supplied in ``forward``, the + kernel is sliced along the channel dimension to match the local slice of + the input, and the ``shortcut`` parameter is sliced accordingly. Causal + mode is not verified to be correct under CP. + + Attributes: + data_dim (int): Spatial rank of the input (1 for sequences, 2 for + images, 3 for volumes). + hidden_dim (int): Number of channels C processed by this operator. + fft_padding (str or Sequence[str]): Boundary condition specification + as supplied by the caller. The normalised per-axis representation + is in ``_periodic_per_axis``. + is_causal (bool): Whether the operator enforces causal (past-only) + convolution. Only valid when ``data_dim=1``. + use_chunked_fftconv (bool): Whether to process channels in chunks to + reduce peak GPU memory. + use_fp16_fft (bool): Whether to use fp16 FFT convolution ops. + fft_backend (str): FFT backend identifier, ``"torch_fft"`` or + ``"subq_ops"``. + grid_type (str or None): Kernel grid size mode (``"single"``, + ``"double"``, or ``None`` for per-axis auto-derivation). + kernel (nn.Module): Implicit kernel generator (produces + ``(kernel_values, grid)`` on each forward call). + mask (nn.Module): Attenuation mask applied to kernel values after + generation. ``nn.Identity`` when no mask is configured. + shortcut (nn.Parameter): Learnable per-channel skip-connection scale + of shape ``(hidden_dim,)``. Fused into the FFT convolution op. + fftconv_fn (callable): Selected FFT convolution function for + channels-last (BLH) input with internal reshape. + fftconv_fn_bhl_input (callable): Selected FFT convolution function for + channels-first (BHL) input. + _periodic_per_axis (tuple[bool, ...]): Per-axis periodicity flags + of length ``data_dim``, derived from ``fft_padding``. + _is_tuple_mode (bool): ``True`` when ``fft_padding`` was supplied as + a sequence of mode strings (mixed-BC path). + """ def __init__( self, @@ -346,61 +571,106 @@ def __init__( use_fp16_fft: bool = False, fft_backend: Literal["torch_fft", "subq_ops"] = "torch_fft", ): - """Initialize the CKConvND. + """Construct a CKConvND operator. + + Validates the combination of ``fft_padding``, ``grid_type``, + ``is_causal``, ``use_fp16_fft``, and ``fft_backend``, normalises the + per-axis boundary-condition representation, adjusts ``kernel_cfg`` + and ``mask_cfg`` to match the resolved kernel grid geometry, and + selects the appropriate FFT convolution function pair. Args: - data_dim: Dimension of input data (1D for sequences, 2D for images, 3D for videos, etc.). - hidden_dim: Hidden dimension. - kernel_cfg: LazyConfig for the kernel. - mask_cfg: LazyConfig for the mask. - grid_type: How the SIREN kernel grid relates to the input size. - ``"single"`` ⇒ kernel size == input size (paired with periodic - FFT convolution). ``"double"`` ⇒ kernel size == 2*input size - (paired with zero-padded FFT convolution). - **Must be ``None`` (or omitted)** when ``fft_padding`` is in - per-axis form — in that case the grid type is auto-derived per - axis (``"single"`` on periodic axes, ``"double"`` on non-periodic - axes). Required when ``fft_padding`` is a single mode string. - fft_padding: Boundary behavior of the FFT convolution. Accepts: - - - ``"zero"`` — every axis zero-padded ("same" linear conv). - - ``"circular"`` — every axis periodic (wrap-around conv). - - **List/tuple of mode strings** — one per spatial axis, in - order, e.g. ``["circular", "zero"]`` for a 2D config that is - periodic on x and zero-padded on y. Mode names are - case-insensitive and whitespace-stripped. - - Required for datasets with mixed periodic/non-periodic - boundary conditions (e.g. Well's ``rayleigh_benard``, - ``viscoelastic_instability``). When supplied as a per-axis - list, the ``grid_type`` argument must be ``None``. + data_dim: Spatial rank of the input signal. ``1`` for 1D + sequences, ``2`` for 2D images (H, W), ``3`` for 3D + volumes (D, H, W). + hidden_dim: Number of channels C. Determines the size of the + learnable ``shortcut`` parameter and the channel dimension of + every intermediate tensor. + kernel_cfg: Lazy config (``LazyConfig``) that instantiates the + implicit kernel generator when resolved. Typically points to + ``SIRENKernelND`` or ``RandomFourierKernelND``. The + ``L_cache`` field, if present, is adjusted to match the + resolved kernel grid size (single or double grid) before + instantiation. + mask_cfg: Lazy config for the attenuation mask applied to the + generated kernel values. Use ``torch.nn.Identity`` (or an + empty identity config) for no masking. If the mask class + accepts a ``grid_size`` parameter, it is set automatically + based on the largest per-axis kernel size. + grid_type: Relationship between the SIREN coordinate grid and the + input spatial size on each axis. + + * ``"single"``: grid spans ``(N+1)//2`` points → kernel size + equals input size N (for periodic / circular conv). + * ``"double"``: grid spans ``N`` points → kernel size is + ``2*N - 1 ≈ 2*N`` (for zero-padded conv). + * ``None``: **required** when ``fft_padding`` is a per-axis + list. The grid type is auto-derived per axis: ``"single"`` + on periodic axes, ``"double"`` on non-periodic axes. + + Must not be ``None`` when ``fft_padding`` is a single mode + string (``"zero"`` or ``"circular"``). + fft_padding: Boundary-condition mode. Accepted forms: + + * ``"zero"``: all axes zero-padded (linear "same" conv). + * ``"circular"``: all axes periodic (wrap-around conv). + Requires ``grid_type="single"`` and + ``use_chunked_fftconv=False``. + * ``["circular", "zero"]`` (list/tuple of mode strings, one + per spatial axis, length must equal ``data_dim``): per-axis + mixed boundary conditions. Requires ``grid_type=None`` and + ``fft_backend="torch_fft"`` and ``use_fp16_fft=False``. + Mode names are case-insensitive and whitespace-stripped. + Must be ``"zero"`` (or an all-``"zero"`` list) when ``is_causal=True``. - is_causal: If True, use causal (left-only) convolution where output at position i - only depends on inputs at positions 0, 1, ..., i. Only supported for 1D data. - use_chunked_fftconv: If True, use memory-efficient chunked FFT convolutions. - Processes channels in chunks to reduce peak memory from complex FFT - intermediates. Typical savings: ~26% memory with ~11% compute overhead. - Useful for memory-constrained training with large spatial dimensions - in 2D/3D. Default is False. - use_fp16_fft: If True, use fp16 FFT convolutions. Uses ortho - normalization to prevent overflow. Saves ~36% peak memory per - convolution with ~0.8% mean relative error vs f32. For zero/causal - padding, sizes are auto-padded to power-of-2. For circular padding, - the input spatial dimensions must already be powers of 2 (a runtime - assertion will fire otherwise). Default is False. - Not supported with a per-axis ``fft_padding`` in v1 (the fp16 - mixed op is a planned follow-up — see ``docs/ops/MIXED_BC_PLAN.md``). - fft_backend: FFT convolution backend to use. ``'torch_fft'`` (default) - uses the torch.fft-based implementations. ``'subq_ops'`` uses the - optimized CUDA kernels from ``subquadratic_ops_torch``. The subq_ops - backend currently supports: - - 2D, zero-padded, non-causal convolutions - - 1D causal convolutions (``data_dim=1`` + ``is_causal=True``) - It does not support fp16 FFT. It supports chunked convolutions via - channel-wise chunking. Per-sample (FiLM) weights are supported on - the 2D path only; the 1D causal CUDA kernel does not accept batched - weights. + is_causal: If ``True``, enforce causal (past-only) convolution so + that the output at position ``n`` only depends on inputs at + positions ``0, …, n``. Only valid for ``data_dim=1``. + Incompatible with periodic ``fft_padding`` and with the + per-axis list form of ``fft_padding``. Default: ``False``. + use_chunked_fftconv: If ``True``, process channels in groups to + reduce peak GPU memory from complex FFT intermediates. + Typical savings: ~26% memory at ~11% compute overhead. + Not supported with ``fft_padding="circular"``. + Default: ``False``. + use_fp16_fft: If ``True``, use fp16 FFT convolution ops. + Uses ``norm="ortho"`` internally to prevent overflow. Saves + ~36% peak memory per convolution at ~0.8% mean relative error + vs fp32. For zero/causal padding, spatial dims are + auto-padded to the next power of two. For circular padding, + the input dims must already be powers of two (a runtime + assertion fires otherwise). Not supported with a per-axis + ``fft_padding`` list (see ``docs/ops/MIXED_BC_PLAN.md``). + Not supported with ``fft_backend="subq_ops"``. + Default: ``False``. + fft_backend: Which FFT convolution backend to use. + + * ``"torch_fft"`` (default): torch.fft-based implementations + in ``nvsubquadratic.ops.fftconv`` and related modules. + * ``"subq_ops"``: optimised CUDA kernels from + ``subquadratic_ops_torch``. Supported configurations: + + - ``data_dim=2``, ``is_causal=False``, ``fft_padding="zero"`` + (2D non-causal zero-padded conv). + - ``data_dim=1``, ``is_causal=True`` (1D causal conv). + + Does not support fp16 FFT, per-axis ``fft_padding``, or + ``data_dim=3``. + + Raises: + AssertionError: If ``fft_backend`` is not one of the recognised + values, or if a constraint between ``grid_type``, + ``fft_padding``, ``is_causal``, ``use_fp16_fft``, and + ``fft_backend`` is violated. + ValueError: If ``fft_padding`` is invalid (wrong type, wrong + length, comma-separated string, boolean), if ``is_causal`` + is combined with a per-axis padding list or periodic padding, + or if the resolved ``(fft_padding, data_dim)`` combination + has no registered FFT function. + NotImplementedError: If ``use_fp16_fft=True`` is requested + together with a per-axis ``fft_padding`` list (fp16 mixed + ops are not yet implemented). """ assert fft_backend in ["torch_fft", "subq_ops"], ( f"Invalid fft_backend: {fft_backend!r}. Must be 'torch_fft' or 'subq_ops'." @@ -683,7 +953,15 @@ def __init__( self.grid_type = grid_type def extra_repr(self) -> str: - """Return extra representation string for the module.""" + """Return a concise summary string for ``print(module)`` and ``repr(module)``. + + Returns: + A human-readable string listing the key hyperparameters: + ``data_dim``, ``hidden_dim``, ``fft_padding``, + ``periodic_per_axis`` (only when in per-axis list mode), + ``grid_type``, ``is_causal``, ``use_chunked_fftconv``, + ``use_fp16_fft``, and ``fft_backend``. + """ bc_repr = f"fft_padding={self.fft_padding!r}" if self._is_tuple_mode: bc_repr += f", periodic_per_axis={self._periodic_per_axis}" @@ -737,12 +1015,15 @@ def flop_count(self, spatial_dims: tuple[int, ...], inference: bool = False) -> Shortcut (skip connection): C * prod(spatial_dims) (elementwise). Args: - spatial_dims: Spatial dimensions of the input signal, e.g. (H, W). - inference: If True and kernel has no FiLM, skip kernel generation - and kernel FFT (both are precomputable and cached). + spatial_dims: Spatial dimensions of the input signal, e.g. + ``(H, W)`` for a 2D image or ``(L,)`` for a 1D sequence. + Must have length equal to ``self.data_dim``. + inference: If ``True`` and the kernel has no FiLM conditioning, + skip the kernel generation and kernel FFT FLOPs (both can be + precomputed and cached at inference time). Returns: - Total FLOPs as an integer. + Total estimated FLOPs as an integer. """ C = self.hidden_dim has_film = getattr(self.kernel, "film_generator", None) is not None @@ -803,16 +1084,47 @@ def flop_count(self, spatial_dims: tuple[int, ...], inference: bool = False) -> def apply_convolution( self, x: torch.Tensor, conv_kernel: torch.Tensor, shortcut: torch.Tensor, is_bhl_input: bool ) -> torch.Tensor: - """Apply the convolution operation using the FFT-based convolution function. + """Apply the FFT-based depthwise convolution. + + Dispatches to the pre-selected ``fftconv_fn`` or + ``fftconv_fn_bhl_input`` depending on the memory layout of ``x``. + When ``is_bhl_input=True`` the kernel is first transposed from + channels-last ``(B, *spatial, C)`` to channels-first ``(B, C, *spatial)`` + to match the BHL-native FFT op. + + The output y is computed as: + + .. code-block:: none + + y = IFFT( FFT(x) ⊙ FFT(conv_kernel) ) + shortcut ⊙ x + + The ``shortcut`` term is fused inside the FFT op (no extra kernel + launch). Args: - x (torch.Tensor): Input tensor. - conv_kernel (torch.Tensor): Convolution kernel tensor. - shortcut (torch.Tensor): Shortcut tensor. - is_bhl_input (bool): Whether the input is in BHL format. + x: Input signal. + + * BLH layout (``is_bhl_input=False``): shape + ``(B, *spatial, C)`` where ``C = self.hidden_dim``. + * BHL layout (``is_bhl_input=True``): shape + ``(B, C, *spatial)``. + + conv_kernel: Kernel values produced by ``self.kernel`` and + optionally masked by ``self.mask``. Always in channels-last + (BLH) format on entry: shape ``(1_or_B, *kernel_spatial, C)``. + Transposed internally when ``is_bhl_input=True``. + shortcut: Per-channel skip-connection scale, shape ``(C,)``. + Typically ``self.shortcut`` or a CP-sliced view thereof. + is_bhl_input: If ``True``, treat ``x`` as channels-first + ``(B, C, *spatial)`` and use ``self.fftconv_fn_bhl_input``. + If ``False``, treat ``x`` as channels-last ``(B, *spatial, C)`` + and use ``self.fftconv_fn`` (which handles the reshape + internally). Returns: - torch.Tensor: Output tensor after applying convolution. + Output tensor in the same memory layout as the input ``x``: + ``(B, *spatial, C)`` when ``is_bhl_input=False``, or + ``(B, C, *spatial)`` when ``is_bhl_input=True``. """ if is_bhl_input: conv_kernel = rearrange( @@ -831,19 +1143,64 @@ def forward( cp_group: torch.distributed.ProcessGroup = None, **mixer_kwargs, ) -> torch.Tensor: - """Forward pass of the CKConvND. + """Run the CKConv forward pass. + + Generates the implicit kernel from the positional grid, optionally + applies the attenuation mask, crops the kernel for causal mode, + handles context-parallel channel slicing, and applies the FFT + convolution with the shortcut term. + + Computation (non-causal, no CP): + + .. code-block:: none + + grid_lens = [(s+1)//2 if single-grid axis else s for s in spatial_dims] + k_θ, grid = self.kernel(grid_lens, conditioning=conditioning) # (1, *grid_lens, C) + k_θ = self.mask(grid=grid, x=k_θ) # attenuation + y = IFFT(FFT(x) ⊙ FFT(k_θ)) + shortcut ⊙ x # FFT conv + + For causal mode (1D only), ``k_θ`` is cropped to its causal (positive- + lag) half before the FFT convolution: + + .. code-block:: none + + k_θ = k_θ[..., kernel_len // 2 :, :] # keep second half Args: - x (torch.Tensor): Input tensor of shape (batch_size, * spatial_dims, hidden_dim) or (batch_size, hidden_dim, * spatial_dims) - is_bhl_input (bool): Whether the input is in BHL format, i.e., (batch_size, hidden_dim, * spatial_dims). - Default is False. - cp_group (torch.distributed.ProcessGroup): Context parallel process group. - Default is None. - **mixer_kwargs: Additional keyword arguments forwarded to the kernel generator - (e.g. ``conditioning`` for FiLM-enabled SIRENKernelND). + x: Input signal tensor. Two supported layouts: + + * **Channels-last** (``is_bhl_input=False``, default): + shape ``(B, *spatial, C)`` where ``C = self.hidden_dim`` + and ``spatial`` has length ``self.data_dim``. + * **Channels-first** (``is_bhl_input=True``): + shape ``(B, C, *spatial)``. + + is_bhl_input: If ``True``, ``x`` is in channels-first + ``(B, C, *spatial)`` layout. Default: ``False`` (channels-last). + cp_group: Context-parallel process group. When provided and + ``cp_group.size() > 1``, the kernel and shortcut are sliced + along the channel dimension to match the local channel slice + held by this rank. The spatial slice of ``x`` is expected to + have already been distributed by the caller. Causal mode is + not verified to be correct under CP. Default: ``None`` + (single-device / no CP). + **mixer_kwargs: Additional keyword arguments forwarded to the + kernel generator. The following key is recognised: + + * ``conditioning`` (``torch.Tensor``, shape ``(B, cond_dim)``): + conditioning vector for FiLM-enabled kernels such as + ``SIRENKernelND`` with a ``film_cfg``. Ignored (no-op) when + the kernel has no FiLM generator. Returns: - torch.Tensor: Output tensor of shape (batch_size, * spatial_dims, hidden_dim) or (batch_size, hidden_dim, * spatial_dims) + Output tensor in the same memory layout as ``x``: + ``(B, *spatial, C)`` when ``is_bhl_input=False``, or + ``(B, C, *spatial)`` when ``is_bhl_input=True``. + + Raises: + ValueError: If ``cp_group`` is provided together with + ``is_causal=True`` (this combination has not been verified + and may produce incorrect results). """ # Get the spatial dimensions from the input tensor if is_bhl_input: From 6cc9d4b9f9765537ce7e5f03f5587e4ffbfeec88 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:28:13 +0200 Subject: [PATCH 22/72] docs(review/ckconv_nd): reviewer feedback --- docs/reviews/ckconv_nd_review.md | 209 +++++++++++++++++++++++++++++++ 1 file changed, 209 insertions(+) create mode 100644 docs/reviews/ckconv_nd_review.md diff --git a/docs/reviews/ckconv_nd_review.md b/docs/reviews/ckconv_nd_review.md new file mode 100644 index 00000000..b3f785f9 --- /dev/null +++ b/docs/reviews/ckconv_nd_review.md @@ -0,0 +1,209 @@ +# Reviewer feedback: `nvsubquadratic/modules/ckconv_nd.py` + +Second-reader pass against the Phase 1 docstrings. Issues are listed in +reading order (module docstring → module-level functions → class). + +______________________________________________________________________ + +## Module docstring + +**Issue 1 — `References` heading uses a non-Google-style RST adornment.** + +The heading reads: + +``` +References: +---------- +``` + +The colon is inside the heading text, which is non-standard Google style +(Google style omits the colon on section names like `Args:`, `Returns:`, +etc., but those are *field* sections, not free-form RST sections). More +critically, the separator `----------` has one more dash than the heading +text (`References:`), which breaks RST rendering. Fix: use a plain +`References` heading with the correct separator length, or follow the +pattern already used by `kernels_nd.py` and write the references as a +bullet list inside the `Related modules` section without a separate heading. + +______________________________________________________________________ + +## `_parse_padding_mode` + +**Issue 2 — Missing explanation of what the return value (`True`/`False`) +means in context.** + +The Returns block says: + +> `True` if the mode is `"circular"` (periodic axis), `False` if it is `"zero"` (zero-padded axis). + +This is correct, but a new reader has no idea why a boolean is returned +here or how it maps to the `periodic` flag used downstream. Add one +sentence linking the return value to `_PADDING_MODE_TO_PERIODIC` and +clarifying that `True` means "this axis uses circular (wrap-around) FFT +convolution." + +______________________________________________________________________ + +## `_resolve_periodic` + +**Issue 3 — `data_dim` arg doc does not mention OmegaConf / `ListConfig` +subtlety.** + +The `fft_padding` arg doc mentions that the list form works with +OmegaConf's `ListConfig`, but the *reason* this matters (why the +`isinstance(fft_padding, Sequence)` check is used instead of +`isinstance(fft_padding, list)`) is documented only as a comment in +`__init__`, not in `_resolve_periodic`. A reader calling `_resolve_periodic` +directly cannot see that context. Add a note to the `fft_padding` arg +docstring in `_resolve_periodic` that says: "OmegaConf `ListConfig` +objects are accepted because `Sequence` matching is used rather than +`isinstance(…, list)`." + +**Issue 4 — Return type annotation is missing.** + +The function signature declares `-> tuple[bool, ...]` but the Returns +block only says "a tuple of `data_dim` booleans." This is fine, but +the docstring should explicitly state the meaning of each boolean +(`True` = circular, `False` = zero-padded) to be self-contained — a +reader consulting just the docstring without seeing `_parse_padding_mode` +would not know this. + +______________________________________________________________________ + +## `_wrap_mixed_op` + +**Issue 5 — The inner `_wrapped` closure has no docstring.** + +`_wrap_mixed_op` is a factory that returns a callable. The returned +`_wrapped` function has no docstring of its own, so `help(_wrapped)` +(and any IDE inspection) shows nothing. Add a one-line docstring to +`_wrapped`: "Call `op_fn(x, kernel, periodic, shortcut)` with bound `periodic`." + +______________________________________________________________________ + +## `_grid_is_single_per_axis` + +**Issue 6 — The correspondence between `grid_type="single"` / `"double"` +and the actual kernel spatial sizes is not stated precisely.** + +The current docstring says `grid_type='single'` means "the SIREN kernel +grid spans `(N+1)//2` per axis so the produced kernel size equals the +input size on that axis," but it does not explain *why* `(N+1)//2` grid +points produce a kernel of size `N`. The reason is that `SIRENKernelND` +evaluates on a `(2*L - 1)`-point grid, so `L = (N+1)//2` → kernel size += `2*(N+1)//2 - 1 ≈ N`. Without this the paragraph is confusing — +"grid spans `(N+1)//2` … so kernel size equals input size" sounds like +a contradiction. Add the derivation: `kernel_size = 2*L - 1 = 2*(N+1)//2 - 1`. + +______________________________________________________________________ + +## `CKConvND` — class docstring + +**Issue 7 — The `shortcut` initialisation scheme is not documented.** + +The class Attributes block lists `shortcut (nn.Parameter)` but gives no +information about how it is initialised. Looking at the code: + +```python +bounds = math.sqrt(1.0 / hidden_dim) +self.shortcut.data.uniform_(-bounds, bounds) +``` + +This is a Kaiming-uniform-style init (scaled by `1/sqrt(hidden_dim)`). +The Attributes entry should mention this: "Initialised with +`uniform(-1/√C, 1/√C)` (Kaiming-uniform scale)." + +**Issue 8 — The `Attributes` block for `fftconv_fn` and +`fftconv_fn_bhl_input` does not document the call signature.** + +Both are described as "(callable): Selected FFT convolution function for …" +without stating what arguments they accept. A reader setting up a +subclass or unit test does not know the signature. Add: "Signature: +`(x, kernel, shortcut) → output`." + +______________________________________________________________________ + +## `CKConvND.__init__` + +**Issue 9 — `kernel_cfg` arg doc does not mention the `out_dim` +constraint.** + +`kernel_cfg` is described as "typically a `SIRENKernelND` or +`RandomFourierKernelND` lazy config." But the kernel must also have +`out_dim = hidden_dim` (its output channel count must match the +operator's channel count). This constraint is implicit in the code +(the convolution would silently produce wrong-shape tensors otherwise). +Add: "The kernel's `out_dim` must equal `hidden_dim`." + +**Issue 10 — `fft_backend="subq_ops"` constraint about per-sample +(FiLM) weights is mentioned in the original `__init__` docstring but +removed from the Phase 1 rewrite.** + +The original `__init__` docstring contained the sentence: "Per-sample +(FiLM) weights are supported on the 2D path only; the 1D causal CUDA +kernel does not accept batched weights." This was dropped in Phase 1. +Reinstate it under the `fft_backend` arg description for `"subq_ops"`. + +______________________________________________________________________ + +## `CKConvND.apply_convolution` + +**Issue 11 — `conv_kernel` shape annotation says `(1_or_B, *kernel_spatial, C)` but does not explain what `kernel_spatial` is relative to `spatial`.** + +For the zero-padded (double-grid) case `kernel_spatial ≈ 2 × spatial` +(the SIREN evaluates on `2*L - 1` points per axis); for the circular +(single-grid) case `kernel_spatial ≈ spatial`. A new reader inspecting +a live tensor would be confused by a kernel that is twice as wide as the +input. Add a one-sentence note: "`kernel_spatial` equals `spatial` +on single-grid (circular) axes and `2*N - 1` on double-grid (zero-padded) +axes." + +______________________________________________________________________ + +## `CKConvND.forward` + +**Issue 12 — Tensor shape annotation uses `C` but the class elsewhere +uses `hidden_dim`.** + +The Args block writes `C = self.hidden_dim` in the description but the +shape annotations say `(B, *spatial, C)`. For consistency with the +class Attributes (which use `hidden_dim`) the shape annotation should +use `hidden_dim`: `(B, *spatial, hidden_dim)`. This avoids a reader +having to cross-reference to find that `C` is a synonym. + +**Issue 13 — The `Raises` block documents a `ValueError` that is +actually a hard `raise ValueError` in a conditional but describes it as +an error about `cp_group + is_causal`, whereas the code uses the wording +"has not been verified to work" — the error is intentional and the +docstring should match: "Causal + CP is explicitly rejected because the +combination has not been verified for correctness."** + +Current text: + +> Raises: +> ValueError: If `cp_group` is provided together with `is_causal=True` +> (this combination has not been verified and may produce incorrect results). + +This is close but slightly misleading — the error says "has not been +verified to work" which could imply it *might* work. The docstring +should say it is **explicitly rejected** and the user must not rely on +the current behaviour silently passing through. + +______________________________________________________________________ + +## Formatting / style + +**Issue 14 — `flop_count` docstring uses non-Google-style free-form +section headers ("Phase 1", "Phase 2") inside the main description +without an `Args` separator.** + +The `flop_count` docstring was inherited verbatim from the original +code. Its structure mixes RST-style bullet headers with Google-style +field sections. In the Phase 1 docstring the `Args` block appears +after a long free-form body that uses `**Phase 1 — …**` bold headers. +This is acceptable in Google style (the convention allows a free-form +description before `Args:`), but the `**Phase 1**` / `**Phase 2**` +headers lose their bold rendering in plain-text help() output because +they use RST syntax. Convert them to plain headings or indented notes +consistent with the rest of the file (e.g. use `.. note::` blocks or +just indent as `Note:` Google-style). From 43033cd07a6b97fadbc6e1e7e3e03e01a0f9c56a Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:30:15 +0200 Subject: [PATCH 23/72] docs(integrate/ckconv_nd): apply reviewer feedback --- docs/reviews/ckconv_nd_review.md | 209 ---------------------------- nvsubquadratic/modules/ckconv_nd.py | 83 +++++++---- 2 files changed, 54 insertions(+), 238 deletions(-) delete mode 100644 docs/reviews/ckconv_nd_review.md diff --git a/docs/reviews/ckconv_nd_review.md b/docs/reviews/ckconv_nd_review.md deleted file mode 100644 index b3f785f9..00000000 --- a/docs/reviews/ckconv_nd_review.md +++ /dev/null @@ -1,209 +0,0 @@ -# Reviewer feedback: `nvsubquadratic/modules/ckconv_nd.py` - -Second-reader pass against the Phase 1 docstrings. Issues are listed in -reading order (module docstring → module-level functions → class). - -______________________________________________________________________ - -## Module docstring - -**Issue 1 — `References` heading uses a non-Google-style RST adornment.** - -The heading reads: - -``` -References: ----------- -``` - -The colon is inside the heading text, which is non-standard Google style -(Google style omits the colon on section names like `Args:`, `Returns:`, -etc., but those are *field* sections, not free-form RST sections). More -critically, the separator `----------` has one more dash than the heading -text (`References:`), which breaks RST rendering. Fix: use a plain -`References` heading with the correct separator length, or follow the -pattern already used by `kernels_nd.py` and write the references as a -bullet list inside the `Related modules` section without a separate heading. - -______________________________________________________________________ - -## `_parse_padding_mode` - -**Issue 2 — Missing explanation of what the return value (`True`/`False`) -means in context.** - -The Returns block says: - -> `True` if the mode is `"circular"` (periodic axis), `False` if it is `"zero"` (zero-padded axis). - -This is correct, but a new reader has no idea why a boolean is returned -here or how it maps to the `periodic` flag used downstream. Add one -sentence linking the return value to `_PADDING_MODE_TO_PERIODIC` and -clarifying that `True` means "this axis uses circular (wrap-around) FFT -convolution." - -______________________________________________________________________ - -## `_resolve_periodic` - -**Issue 3 — `data_dim` arg doc does not mention OmegaConf / `ListConfig` -subtlety.** - -The `fft_padding` arg doc mentions that the list form works with -OmegaConf's `ListConfig`, but the *reason* this matters (why the -`isinstance(fft_padding, Sequence)` check is used instead of -`isinstance(fft_padding, list)`) is documented only as a comment in -`__init__`, not in `_resolve_periodic`. A reader calling `_resolve_periodic` -directly cannot see that context. Add a note to the `fft_padding` arg -docstring in `_resolve_periodic` that says: "OmegaConf `ListConfig` -objects are accepted because `Sequence` matching is used rather than -`isinstance(…, list)`." - -**Issue 4 — Return type annotation is missing.** - -The function signature declares `-> tuple[bool, ...]` but the Returns -block only says "a tuple of `data_dim` booleans." This is fine, but -the docstring should explicitly state the meaning of each boolean -(`True` = circular, `False` = zero-padded) to be self-contained — a -reader consulting just the docstring without seeing `_parse_padding_mode` -would not know this. - -______________________________________________________________________ - -## `_wrap_mixed_op` - -**Issue 5 — The inner `_wrapped` closure has no docstring.** - -`_wrap_mixed_op` is a factory that returns a callable. The returned -`_wrapped` function has no docstring of its own, so `help(_wrapped)` -(and any IDE inspection) shows nothing. Add a one-line docstring to -`_wrapped`: "Call `op_fn(x, kernel, periodic, shortcut)` with bound `periodic`." - -______________________________________________________________________ - -## `_grid_is_single_per_axis` - -**Issue 6 — The correspondence between `grid_type="single"` / `"double"` -and the actual kernel spatial sizes is not stated precisely.** - -The current docstring says `grid_type='single'` means "the SIREN kernel -grid spans `(N+1)//2` per axis so the produced kernel size equals the -input size on that axis," but it does not explain *why* `(N+1)//2` grid -points produce a kernel of size `N`. The reason is that `SIRENKernelND` -evaluates on a `(2*L - 1)`-point grid, so `L = (N+1)//2` → kernel size -= `2*(N+1)//2 - 1 ≈ N`. Without this the paragraph is confusing — -"grid spans `(N+1)//2` … so kernel size equals input size" sounds like -a contradiction. Add the derivation: `kernel_size = 2*L - 1 = 2*(N+1)//2 - 1`. - -______________________________________________________________________ - -## `CKConvND` — class docstring - -**Issue 7 — The `shortcut` initialisation scheme is not documented.** - -The class Attributes block lists `shortcut (nn.Parameter)` but gives no -information about how it is initialised. Looking at the code: - -```python -bounds = math.sqrt(1.0 / hidden_dim) -self.shortcut.data.uniform_(-bounds, bounds) -``` - -This is a Kaiming-uniform-style init (scaled by `1/sqrt(hidden_dim)`). -The Attributes entry should mention this: "Initialised with -`uniform(-1/√C, 1/√C)` (Kaiming-uniform scale)." - -**Issue 8 — The `Attributes` block for `fftconv_fn` and -`fftconv_fn_bhl_input` does not document the call signature.** - -Both are described as "(callable): Selected FFT convolution function for …" -without stating what arguments they accept. A reader setting up a -subclass or unit test does not know the signature. Add: "Signature: -`(x, kernel, shortcut) → output`." - -______________________________________________________________________ - -## `CKConvND.__init__` - -**Issue 9 — `kernel_cfg` arg doc does not mention the `out_dim` -constraint.** - -`kernel_cfg` is described as "typically a `SIRENKernelND` or -`RandomFourierKernelND` lazy config." But the kernel must also have -`out_dim = hidden_dim` (its output channel count must match the -operator's channel count). This constraint is implicit in the code -(the convolution would silently produce wrong-shape tensors otherwise). -Add: "The kernel's `out_dim` must equal `hidden_dim`." - -**Issue 10 — `fft_backend="subq_ops"` constraint about per-sample -(FiLM) weights is mentioned in the original `__init__` docstring but -removed from the Phase 1 rewrite.** - -The original `__init__` docstring contained the sentence: "Per-sample -(FiLM) weights are supported on the 2D path only; the 1D causal CUDA -kernel does not accept batched weights." This was dropped in Phase 1. -Reinstate it under the `fft_backend` arg description for `"subq_ops"`. - -______________________________________________________________________ - -## `CKConvND.apply_convolution` - -**Issue 11 — `conv_kernel` shape annotation says `(1_or_B, *kernel_spatial, C)` but does not explain what `kernel_spatial` is relative to `spatial`.** - -For the zero-padded (double-grid) case `kernel_spatial ≈ 2 × spatial` -(the SIREN evaluates on `2*L - 1` points per axis); for the circular -(single-grid) case `kernel_spatial ≈ spatial`. A new reader inspecting -a live tensor would be confused by a kernel that is twice as wide as the -input. Add a one-sentence note: "`kernel_spatial` equals `spatial` -on single-grid (circular) axes and `2*N - 1` on double-grid (zero-padded) -axes." - -______________________________________________________________________ - -## `CKConvND.forward` - -**Issue 12 — Tensor shape annotation uses `C` but the class elsewhere -uses `hidden_dim`.** - -The Args block writes `C = self.hidden_dim` in the description but the -shape annotations say `(B, *spatial, C)`. For consistency with the -class Attributes (which use `hidden_dim`) the shape annotation should -use `hidden_dim`: `(B, *spatial, hidden_dim)`. This avoids a reader -having to cross-reference to find that `C` is a synonym. - -**Issue 13 — The `Raises` block documents a `ValueError` that is -actually a hard `raise ValueError` in a conditional but describes it as -an error about `cp_group + is_causal`, whereas the code uses the wording -"has not been verified to work" — the error is intentional and the -docstring should match: "Causal + CP is explicitly rejected because the -combination has not been verified for correctness."** - -Current text: - -> Raises: -> ValueError: If `cp_group` is provided together with `is_causal=True` -> (this combination has not been verified and may produce incorrect results). - -This is close but slightly misleading — the error says "has not been -verified to work" which could imply it *might* work. The docstring -should say it is **explicitly rejected** and the user must not rely on -the current behaviour silently passing through. - -______________________________________________________________________ - -## Formatting / style - -**Issue 14 — `flop_count` docstring uses non-Google-style free-form -section headers ("Phase 1", "Phase 2") inside the main description -without an `Args` separator.** - -The `flop_count` docstring was inherited verbatim from the original -code. Its structure mixes RST-style bullet headers with Google-style -field sections. In the Phase 1 docstring the `Args` block appears -after a long free-form body that uses `**Phase 1 — …**` bold headers. -This is acceptable in Google style (the convention allows a free-form -description before `Args:`), but the `**Phase 1**` / `**Phase 2**` -headers lose their bold rendering in plain-text help() output because -they use RST syntax. Convert them to plain headings or indented notes -consistent with the rest of the file (e.g. use `.. note::` blocks or -just indent as `Note:` Google-style). diff --git a/nvsubquadratic/modules/ckconv_nd.py b/nvsubquadratic/modules/ckconv_nd.py index 804c148f..1e118c14 100644 --- a/nvsubquadratic/modules/ckconv_nd.py +++ b/nvsubquadratic/modules/ckconv_nd.py @@ -312,8 +312,11 @@ def _parse_padding_mode(mode: str) -> bool: (case-insensitive, whitespace-stripped). Returns: - ``True`` if the mode is ``"circular"`` (periodic axis), ``False`` if - it is ``"zero"`` (zero-padded axis). + ``True`` if the mode is ``"circular"`` (the axis uses circular / + wrap-around FFT convolution, matching ``_PADDING_MODE_TO_PERIODIC``), + ``False`` if it is ``"zero"`` (the axis uses zero-padded linear FFT + convolution). The boolean is the ``periodic`` flag used by the + downstream FFT dispatch tables. Raises: ValueError: If ``mode`` is not one of the recognised padding modes. @@ -359,15 +362,20 @@ def _resolve_periodic( Args: fft_padding: Boundary-condition specification, either a single mode string (``"zero"`` / ``"circular"``) or a sequence of mode strings - of length ``data_dim`` (one per spatial axis). + of length ``data_dim`` (one per spatial axis). OmegaConf + ``ListConfig`` objects are accepted as sequences because ``Sequence`` + matching is used internally rather than ``isinstance(…, list)`` — + this matters when configs flow through ``LazyConfig``. data_dim: Number of spatial dimensions in the input signal. Used to validate the length of a sequence-form ``fft_padding`` and to broadcast a scalar mode string. Returns: A tuple of ``data_dim`` booleans, one per spatial axis. ``True`` - means that axis uses circular (periodic) convolution; ``False`` means - zero-padded (linear) convolution. + means that axis uses circular (wrap-around / periodic) FFT convolution; + ``False`` means zero-padded (linear) FFT convolution. This tuple is + the canonical ``periodic`` flag consumed by the FFT dispatch tables and + the ``mixed_fftconv*`` ops. Raises: ValueError: on invalid mode strings, wrong number of axes, or @@ -440,6 +448,7 @@ def _wrap_mixed_op(op_fn, periodic: tuple[bool, ...]): """ def _wrapped(x, kernel, shortcut): + """Call ``op_fn(x, kernel, periodic, shortcut)`` with bound ``periodic``.""" return op_fn(x, kernel, periodic, shortcut) return _wrapped @@ -452,10 +461,12 @@ def _grid_is_single_per_axis( """Return per-axis 'use single grid' flags for the SIREN kernel. In CKConvND, ``grid_type='single'`` means the SIREN kernel grid spans - ``(N+1)//2`` per axis so the produced kernel size equals the input size - on that axis (paired with periodic/circular FFT conv). ``'double'`` - means the kernel grid spans ``N`` so the kernel covers twice the input - (paired with zero-padded FFT conv). + ``L = (N+1)//2`` grid points per axis. ``SIRENKernelND`` evaluates on a + ``(2*L - 1)``-point grid, so the produced kernel has size + ``2*(N+1)//2 - 1 ≈ N`` — i.e. the kernel size equals the input size on + that axis (paired with periodic/circular FFT conv). ``'double'`` means + the kernel grid spans ``L = N`` points, giving kernel size ``2*N - 1 ≈ 2*N`` + (paired with zero-padded FFT conv), so the kernel covers twice the input. - **String mode** (``grid_type`` is ``'single'`` / ``'double'``): the same choice applies to every axis. @@ -548,10 +559,14 @@ class CKConvND(torch.nn.Module): generation. ``nn.Identity`` when no mask is configured. shortcut (nn.Parameter): Learnable per-channel skip-connection scale of shape ``(hidden_dim,)``. Fused into the FFT convolution op. + Initialised with ``uniform(-1/√hidden_dim, 1/√hidden_dim)`` + (Kaiming-uniform scale). fftconv_fn (callable): Selected FFT convolution function for - channels-last (BLH) input with internal reshape. + channels-last (BLH) input with internal reshape. Signature: + ``(x, kernel, shortcut) → output``. fftconv_fn_bhl_input (callable): Selected FFT convolution function for - channels-first (BHL) input. + channels-first (BHL) input. Signature: + ``(x, kernel, shortcut) → output``. _periodic_per_axis (tuple[bool, ...]): Per-axis periodicity flags of length ``data_dim``, derived from ``fft_padding``. _is_tuple_mode (bool): ``True`` when ``fft_padding`` was supplied as @@ -588,10 +603,11 @@ def __init__( every intermediate tensor. kernel_cfg: Lazy config (``LazyConfig``) that instantiates the implicit kernel generator when resolved. Typically points to - ``SIRENKernelND`` or ``RandomFourierKernelND``. The - ``L_cache`` field, if present, is adjusted to match the - resolved kernel grid size (single or double grid) before - instantiation. + ``SIRENKernelND`` or ``RandomFourierKernelND``. The kernel's + ``out_dim`` must equal ``hidden_dim``; a mismatch will produce + incorrect tensor shapes at runtime. The ``L_cache`` field, if + present, is adjusted to match the resolved kernel grid size + (single or double grid) before instantiation. mask_cfg: Lazy config for the attenuation mask applied to the generated kernel values. Use ``torch.nn.Identity`` (or an empty identity config) for no masking. If the mask class @@ -652,8 +668,11 @@ def __init__( ``subquadratic_ops_torch``. Supported configurations: - ``data_dim=2``, ``is_causal=False``, ``fft_padding="zero"`` - (2D non-causal zero-padded conv). + (2D non-causal zero-padded conv). Per-sample (FiLM) + batched kernel weights are supported on this path. - ``data_dim=1``, ``is_causal=True`` (1D causal conv). + The 1D causal CUDA kernel does not accept batched per-sample + weights; FiLM conditioning is unsupported on this path. Does not support fp16 FFT, per-axis ``fft_padding``, or ``data_dim=3``. @@ -977,12 +996,12 @@ def flop_count(self, spatial_dims: tuple[int, ...], inference: bool = False) -> Two phases: - **Phase 1 — Kernel generation** (via SIREN MLP): - Delegated to ``self.kernel.flop_count(grid_lens, inference)``. - At ``inference=True`` without FiLM, the kernel is input-independent - and can be precomputed, so this returns 0. + Phase 1 - Kernel generation (via SIREN MLP): + Delegated to ``self.kernel.flop_count(grid_lens, inference)``. + At ``inference=True`` without FiLM, the kernel is input-independent + and can be precomputed, so this returns 0. - **Phase 2 — FFT-based depthwise convolution** (C = ``self.hidden_dim``): + Phase 2 - FFT-based depthwise convolution (C = ``self.hidden_dim``): The convolution is computed in the frequency domain. Padded signal sizes Np_i depend on the padding mode: - ``"zero"`` non-causal ("same"-mode): @@ -1112,7 +1131,10 @@ def apply_convolution( conv_kernel: Kernel values produced by ``self.kernel`` and optionally masked by ``self.mask``. Always in channels-last (BLH) format on entry: shape ``(1_or_B, *kernel_spatial, C)``. - Transposed internally when ``is_bhl_input=True``. + ``kernel_spatial`` equals ``spatial`` on single-grid (circular) + axes and ``2*N - 1`` on double-grid (zero-padded) axes, where + ``N`` is the corresponding input spatial size. Transposed + internally when ``is_bhl_input=True``. shortcut: Per-channel skip-connection scale, shape ``(C,)``. Typically ``self.shortcut`` or a CP-sliced view thereof. is_bhl_input: If ``True``, treat ``x`` as channels-first @@ -1170,10 +1192,10 @@ def forward( x: Input signal tensor. Two supported layouts: * **Channels-last** (``is_bhl_input=False``, default): - shape ``(B, *spatial, C)`` where ``C = self.hidden_dim`` - and ``spatial`` has length ``self.data_dim``. + shape ``(B, *spatial, hidden_dim)`` where ``spatial`` has + length ``self.data_dim``. * **Channels-first** (``is_bhl_input=True``): - shape ``(B, C, *spatial)``. + shape ``(B, hidden_dim, *spatial)``. is_bhl_input: If ``True``, ``x`` is in channels-first ``(B, C, *spatial)`` layout. Default: ``False`` (channels-last). @@ -1194,13 +1216,16 @@ def forward( Returns: Output tensor in the same memory layout as ``x``: - ``(B, *spatial, C)`` when ``is_bhl_input=False``, or - ``(B, C, *spatial)`` when ``is_bhl_input=True``. + ``(B, *spatial, hidden_dim)`` when ``is_bhl_input=False``, or + ``(B, hidden_dim, *spatial)`` when ``is_bhl_input=True``. Raises: ValueError: If ``cp_group`` is provided together with - ``is_causal=True`` (this combination has not been verified - and may produce incorrect results). + ``is_causal=True``. This combination is **explicitly rejected** + because it has not been verified for correctness — the causal + kernel crop and CP channel slicing interact in ways that may + silently leak future positions. Do not rely on the error being + absent in future versions without re-verification. """ # Get the spatial dimensions from the input tensor if is_bhl_input: From a60ac6bb224ecf67f4c9d857366a4a334fd00729 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:32:02 +0200 Subject: [PATCH 24/72] docs(integrate/film): apply reviewer feedback (from worktree) --- nvsubquadratic/modules/film.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/nvsubquadratic/modules/film.py b/nvsubquadratic/modules/film.py index 811c9216..b8d9dd0e 100644 --- a/nvsubquadratic/modules/film.py +++ b/nvsubquadratic/modules/film.py @@ -98,7 +98,8 @@ class KernelFiLMGenerator(nn.Module): Args: cond_dim: Dimensionality of the conditioning input ``c``. kernel_hidden_dim: Hidden dimension of the SIREN layers to modulate. - num_film_layers: Number of (gamma, beta) pairs to produce (one per SIREN hidden layer). + num_film_layers: Number of (gamma, beta) pairs to produce (one per SIREN + hidden layer). Must be ≥ 1. film_hidden_dim: Hidden dimension of the FiLM generator MLP (bottleneck). no_weight_decay: Controls weight decay for FiLM **weight** parameters. All biases are always excluded from weight decay regardless of From dbe0c726adbd5f702ffeecadcaa55bd70e4543ed Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 16:32:35 +0200 Subject: [PATCH 25/72] docs(tracker): mark ckconv_nd, residual_block, patchify, film as done --- docs-tracker.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs-tracker.md b/docs-tracker.md index 9c0c44d2..6a5adf1c 100644 --- a/docs-tracker.md +++ b/docs-tracker.md @@ -44,7 +44,7 @@ Work bottom-up: primitive ops → modules → networks → experiments. | ---------------------------------- | ------ | ------------------------------------------------------------------------------------------- | | `kernels_nd.py` | \[x\] | Learned kernel parametrisation — RFF + SIREN, FiLM-conditioned variants | | `hyena_nd.py` | \[x\] | Hyena operator (ND) — two-gate sandwich, AllToAll CP, BC-aware convolution | -| `ckconv_nd.py` | \[ \] | CKConv (ND) | +| `ckconv_nd.py` | \[x\] | CKConv (ND) — implicit kernel `k_θ(p) = MLP_θ(pos_enc(p))`, FFT domain, BC modes | | `ckconv_multihead_nd.py` | \[ \] | Multi-head CKConv | | `mamba_nd.py` | \[ \] | Mamba SSM (ND) | | `attention.py` | \[ \] | Standard attention | @@ -52,13 +52,13 @@ Work bottom-up: primitive ops → modules → networks → experiments. | `vit5_hyena_adapter.py` | \[ \] | Hyena adapter for ViT5 | | `sequence_mixer.py` | \[x\] | Operator-agnostic dispatch layer (Hyena / Attention / CKConv / Mamba) | | `condition_mixer.py` | \[ \] | Conditioning mixer | -| `residual_block.py` | \[ \] | Residual block | +| `residual_block.py` | \[x\] | Residual block — pre-norm + mixer + MLP, optional FiLM/AdaLN-Zero conditioning | | `vit5_residual_block.py` | \[ \] | ViT5 residual block | -| `patchify.py` | \[ \] | Patch embedding | +| `patchify.py` | \[x\] | Patch embedding — strided conv, 1D/2D/3D, channels-last layout | | `position_encoding.py` | \[ \] | Position encodings | | `masks_nd.py` | \[ \] | ND masking utils | | `mlp.py` | \[ \] | MLP block | -| `film.py` | \[ \] | FiLM conditioning | +| `film.py` | \[x\] | FiLM conditioning — γ(c)⊙x + β(c), SIREN-based kernel generator | | `grn.py` | \[ \] | GRN normalisation | | `layer_scale.py` | \[ \] | LayerScale | | `rms_norm.py` | \[ \] | RMS normalisation | From 0fb4dea87439953af3cb58b9bb8313cd584ffaf0 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 17:33:45 +0200 Subject: [PATCH 26/72] docs(write/position_encoding): add module and class docstrings with math context --- nvsubquadratic/modules/position_encoding.py | 172 +++++++++++++++++++- 1 file changed, 168 insertions(+), 4 deletions(-) diff --git a/nvsubquadratic/modules/position_encoding.py b/nvsubquadratic/modules/position_encoding.py index 60a45cea..1ef97105 100644 --- a/nvsubquadratic/modules/position_encoding.py +++ b/nvsubquadratic/modules/position_encoding.py @@ -1,7 +1,60 @@ # TODO: Add license header here -"""Learned positional encodings for ND inputs.""" +"""Position encoding layers for ND spatial token grids. + +Overview +-------- +Sequence mixers (Hyena, Attention, CKConv, …) operate on a flat sequence of +token embeddings and are in principle permutation-equivariant. *Position +encodings* break that symmetry by injecting spatial-coordinate information +directly into the embedding space so the mixer can reason about where each +token sits in the original spatial grid. + +In this codebase position encodings are applied **after** patch embedding (see +``nvsubquadratic.modules.patchify.Patchify``): the ``Patchify`` layer maps the +raw input of shape ``[B, *spatial, C_in]`` to a patch-token grid of shape +``[B, *patch_grid, C_embed]``, and then a position encoding layer adds a +coordinate-dependent bias of the same shape before the tokens are passed into +the first mixer block. + +Variants +-------- +``PositionEmbeddingND`` + **Fully-learned / lookup-table** encoding. A separate ``nn.Embedding`` + table is trained for each spatial axis; the per-axis embeddings are + concatenated channel-wise. Requires the grid size to be bounded at + construction time (``max_dim_lengths``) but makes no structural assumption + about how position information should be represented. The embedding + parameters are tagged ``_no_weight_decay = True`` so they are excluded from + weight-decay optimiser groups. + + Best suited for fixed-resolution training where the spatial grid size does + not change between train and inference. For variable-resolution or + resolution-generalisation use-cases, prefer a sinusoidal or RFF-based + encoding (not yet included in this module). + +ND generalisation strategy +-------------------------- +A naive learned positional encoding for an ND grid would store +``prod(max_dim_lengths)`` parameters — exponential in the number of dimensions. +``PositionEmbeddingND`` avoids this by factorising the encoding across axes: a +separate embedding table of length ``max_dim_lengths[d]`` is kept for each +spatial axis ``d``, and the per-axis embeddings of dimension +``embedding_dim // data_dim`` are concatenated to produce the final +``embedding_dim``-dimensional token offset. The total parameter count is +therefore ``data_dim * max(max_dim_lengths) * (embedding_dim // data_dim) = +embedding_dim * max(max_dim_lengths)`` — linear in the embedding dimension and +the largest axis length. + +Cross-references +---------------- +* ``nvsubquadratic.modules.patchify.Patchify`` — produces the patch-token grid + that is the expected input to the ``forward`` method of position-encoding + layers in this module. +* ``nvsubquadratic.modules.kernels_nd.SIRENKernelND`` — uses a similar spatial + coordinate grid to parametrise implicit convolutional kernels. +""" from __future__ import annotations @@ -12,10 +65,80 @@ class PositionEmbeddingND(nn.Module): - """Learned positional encoding for 1D/2D/3D ... inputs.""" + """Axis-factorised learned positional encoding for ND spatial token grids. + + Each spatial axis has its own ``nn.Embedding`` table of shape + ``(max_dim_lengths[d], embedding_dim // data_dim)``. For a grid position + ``(i_0, i_1, ..., i_{D-1})`` the encoding is formed by looking up the + per-axis embeddings and *concatenating* them channel-wise:: + + PE(i_0, ..., i_{D-1}) = [E_0(i_0) ‖ E_1(i_1) ‖ ... ‖ E_{D-1}(i_{D-1})] + + where ``‖`` denotes concatenation along the channel dimension and each + ``E_d ∈ R^{max_dim_lengths[d] × (embedding_dim // data_dim)}`` is a + learned embedding matrix. The result has length ``embedding_dim`` + (requires ``embedding_dim % data_dim == 0``). + + This factorised form is a *separable* positional encoding: the position + information for axis ``d`` is captured entirely in the slice of channels + ``[d * per_dim, (d+1) * per_dim)``. It is not a *joint* encoding of the + full ND coordinate (unlike 2D sinusoidal encodings that mix axes), so + cross-axis interactions must be learned by the mixer layers themselves. + + **Output shape**:: + + forward(x) -> Tensor of shape [B, *spatial_dims, embedding_dim] + + The returned tensor is a broadcast-expanded encoding grid with the *same + shape as the input* ``x``, and is typically **added** to ``x`` before the + first mixer block:: + + x = x + position_embedding(x) + + **Parameter count** — ``data_dim`` embedding tables, each of size + ``max_dim_lengths[d] × (embedding_dim // data_dim)``:: + + total = sum(max_dim_lengths[d] * (embedding_dim // data_dim) + for d in range(data_dim)) + + **No weight decay** — all embedding parameters are tagged + ``param._no_weight_decay = True`` so that optimiser builders (e.g. those + using ``param._no_weight_decay`` to separate param groups) can exclude them + from L2 regularisation, following the standard ViT practice. + + Attributes: + embedding_dim (int): Total embedding dimension of the output encoding. + per_dim_embedding_dim (int): Per-axis slice width, + ``embedding_dim // data_dim``. + data_dim (int): Number of spatial axes (1, 2, or 3). + max_dim_lengths (tuple[int, ...]): Maximum supported grid size for each + spatial axis. Length equals ``data_dim``. + data_embeddings (nn.ModuleDict): Dictionary mapping axis keys + ``{"x"}`` / ``{"x", "y"}`` / ``{"x", "y", "z"}`` to the + corresponding ``nn.Embedding`` modules. + """ def __init__(self, embedding_dim: int, data_dim: int, max_dim_lengths: Sequence[int]): - """Configure positional embeddings for a fixed set of spatial dimensions.""" + """Initialise per-axis embedding tables. + + Args: + embedding_dim: Total number of channels in the output position + encoding. Must be divisible by ``data_dim`` so that the + budget can be split evenly across axes. + data_dim: Number of spatial axes of the input token grid. Must + satisfy ``1 <= data_dim <= 3``. + max_dim_lengths: Maximum grid size for each spatial axis. Must + have exactly ``data_dim`` entries. An ``nn.Embedding`` table + of length ``max_dim_lengths[d]`` is allocated for axis ``d``; + inputs whose spatial size along axis ``d`` exceeds this value + will raise a ``ValueError`` in ``forward``. + + Raises: + AssertionError: If ``data_dim < 1``. + AssertionError: If ``len(max_dim_lengths) != data_dim``. + AssertionError: If ``data_dim > 3``. + AssertionError: If ``embedding_dim % data_dim != 0``. + """ super().__init__() assert data_dim >= 1, "data_dim must be >= 1" assert len(max_dim_lengths) == data_dim, "max_dim_lengths must have length data_dim" @@ -37,7 +160,48 @@ def __init__(self, embedding_dim: int, data_dim: int, max_dim_lengths: Sequence[ param._no_weight_decay = True def forward(self, x: torch.Tensor) -> torch.Tensor: - """Return the concatenated per-axis embeddings for the input grid.""" + """Compute the position encoding grid for the given input token grid. + + For each spatial axis ``d``, position indices ``0, 1, ..., L_d - 1`` + are looked up in the corresponding ``nn.Embedding`` table to produce a + 1-D embedding slice of shape ``[L_d, per_dim_embedding_dim]``. That + slice is broadcast-expanded to the full spatial shape + ``[B, *spatial_dims, per_dim_embedding_dim]`` and the per-axis tensors + are concatenated along the last dimension to give the final encoding of + shape ``[B, *spatial_dims, embedding_dim]``. + + The returned tensor has the same shape as ``x`` and should be **added** + to ``x``:: + + x = x + position_embedding(x) + + Args: + x: Input token-grid tensor in channels-last layout. Shape: + ``[B, *spatial_dims, embedding_dim]``, where the number of + spatial axes must equal ``data_dim`` and the last dimension + must equal ``embedding_dim``. For example: + + - 1-D (sequences): ``[B, L, C]`` + - 2-D (images): ``[B, H, W, C]`` + - 3-D (volumes): ``[B, D, H, W, C]`` + + Returns: + Position encoding tensor of shape + ``[B, *spatial_dims, embedding_dim]`` — the same shape as ``x``. + Each spatial location ``(i_0, ..., i_{D-1})`` holds the + concatenation of the per-axis embedding lookups:: + + out[b, i_0, ..., i_{D-1}, :] = + [E_0(i_0) ‖ E_1(i_1) ‖ ... ‖ E_{D-1}(i_{D-1})] + + Raises: + ValueError: If ``x.ndim != data_dim + 2`` (wrong number of + dimensions — expected batch + ``data_dim`` spatial + channel). + ValueError: If ``x.shape[-1] != embedding_dim`` (channel dimension + mismatch). + ValueError: If any spatial dimension of ``x`` exceeds the + corresponding entry of ``max_dim_lengths``. + """ if x.ndim != self.data_dim + 2: raise ValueError( f"PositionEmbeddingND expects input of shape [B, *{self.data_dim} dims, C]; got {tuple(x.shape)}" From d4d36f72856c37b60ce18fbed37a667e0f95be48 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 20:58:25 +0200 Subject: [PATCH 27/72] docs(review/position_encoding): reviewer feedback --- docs/reviews/position_encoding_review.md | 89 ++++++++++++++++++++++++ 1 file changed, 89 insertions(+) create mode 100644 docs/reviews/position_encoding_review.md diff --git a/docs/reviews/position_encoding_review.md b/docs/reviews/position_encoding_review.md new file mode 100644 index 00000000..9f0a97cd --- /dev/null +++ b/docs/reviews/position_encoding_review.md @@ -0,0 +1,89 @@ +# Review: `nvsubquadratic/modules/position_encoding.py` + +Reviewed as a critical second reader targeting external collaborators reading alongside the paper. + +______________________________________________________________________ + +## Issue 1 — Module docstring mentions "sinusoidal, learned, RFF" variants but only one is implemented + +**Location**: Module docstring, "Variants" section. + +**Quoted text**: + +> "describe the variants present (sinusoidal, learned, RFF, etc.)" + +The module docstring (in the prose leading up to the `PositionEmbeddingND` variant block) does not explicitly acknowledge that sinusoidal and RFF encodings are *absent* from the module. The `PositionEmbeddingND` class description says "For variable-resolution or resolution-generalisation use-cases, prefer a sinusoidal or RFF-based encoding (not yet included in this module)." This is good, but the module-level "Variants" heading implies more variants are present. **Fix**: rename the "Variants" section to "Implemented variants" or explicitly add a "Planned / not yet implemented" sub-list (sinusoidal, RFF) so readers are not confused by the heading. + +______________________________________________________________________ + +## Issue 2 — The factorisation parameter-count formula in the module docstring has a misleading simplification + +**Location**: Module docstring, "ND generalisation strategy" section. + +**Quoted text**: + +> "The total parameter count is therefore `data_dim * max(max_dim_lengths) * (embedding_dim // data_dim) = embedding_dim * max(max_dim_lengths)` — linear in the embedding dimension and the largest axis length." + +The formula `data_dim * max(max_dim_lengths) * (embedding_dim // data_dim)` is only correct when all axes have the *same* maximum length. The true total is `sum(max_dim_lengths[d] * (embedding_dim // data_dim) for d in range(data_dim))`, which is already stated correctly in the class docstring's "Parameter count" section. The module-level simplification introduces a subtle error for non-square grids (e.g. a 2D 56×14 grid). **Fix**: replace the over-simplified formula with the exact sum, or explicitly state "assuming all axes have the same maximum length `M`". + +______________________________________________________________________ + +## Issue 3 — No mention of the `param._no_weight_decay` convention in the module docstring + +**Location**: Module docstring (overview / cross-references). + +The module docstring says nothing about the `_no_weight_decay` tagging convention. An external collaborator writing a custom optimiser builder will need to know that the codebase uses `param._no_weight_decay = True` as a signal. The class docstring explains it, but readers skimming the module overview will miss it. **Fix**: add a one-sentence note in the module docstring (e.g. in the `PositionEmbeddingND` variant description or in a new "Optimiser integration" note) that all embedding parameters are tagged with `_no_weight_decay = True` and that optimiser builders should respect this flag. + +______________________________________________________________________ + +## Issue 4 — `forward` docstring does not explain the broadcast-expand mechanism + +**Location**: `PositionEmbeddingND.forward` docstring, "Returns" section. + +The docstring says the returned tensor has "the same shape as `x`" but does not explain *how* the 1-D per-axis embedding slices are broadcast-expanded to the full `[B, *spatial_dims, per_dim_embedding_dim]` shape before concatenation. A reader inspecting the code will see `emb_axis.view(1, *shape, ...)` followed by `.expand(batch_size, *spatial_dims, ...)` and may not immediately understand why this is correct. **Fix**: add one sentence in the "Returns" section describing the broadcast-expand step: "Each 1-D per-axis embedding of shape `[L_d, per_dim]` is reshaped to `[1, ..., L_d, ..., per_dim]` (with 1 in all axes except `d`) and then broadcast-expanded to `[B, *spatial_dims, per_dim]`; the resulting tensors are concatenated on the last dimension." + +______________________________________________________________________ + +## Issue 5 — The `forward` method validates `x.shape[-1] == embedding_dim` but the docstring's `Args` does not mention this constraint explicitly + +**Location**: `PositionEmbeddingND.forward` docstring, "Args" block for `x`. + +**Quoted text**: + +> "Shape: `[B, *spatial_dims, embedding_dim]`, where the number of spatial axes must equal `data_dim` and the last dimension must equal `embedding_dim`." + +The constraint is present but buried in the shape annotation prose. The `Raises` block lists `ValueError` for channel-dimension mismatch, but does not state what the mismatch condition is ("`x.shape[-1] != embedding_dim`"). **Fix**: in the `Raises` entry for the channel mismatch error, state the condition explicitly: "If `x.shape[-1] != self.embedding_dim` (channel count does not match the embedding dimension passed at construction)." + +______________________________________________________________________ + +## Issue 6 — Missing `device` / `dtype` propagation note in `forward` + +**Location**: `PositionEmbeddingND.forward` docstring. + +The method internally calls `torch.arange(..., device=x.device)` but does not use `x.dtype` for the position indices (correctly using `torch.long`, since they are integer indices). However, the embedding lookup result inherits the dtype of `nn.Embedding.weight`. There is no note clarifying that the output dtype matches the embedding weight dtype (typically `float32`) and that callers working in mixed precision (e.g. `bfloat16`) may need to cast the result. **Fix**: add a `Note` block: "The returned tensor has the same `dtype` as the embedding weight parameters (typically `torch.float32`). In mixed-precision training the result should be cast to match `x.dtype` before addition: `x = x + pos_enc(x).to(x.dtype)`." + +______________________________________________________________________ + +## Issue 7 — Class docstring uses `‖` (U+2016 DOUBLE VERTICAL LINE) for concatenation; this is non-standard and may not render in all docstring viewers + +**Location**: `PositionEmbeddingND` class docstring. + +**Quoted text**: + +> `PE(i_0, ..., i_{D-1}) = [E_0(i_0) ‖ E_1(i_1) ‖ ... ‖ E_{D-1}(i_{D-1})]` + +The `‖` character is ambiguous (it is also the norm operator). Standard ML papers use `concat` or `[· ; ·]`. **Fix**: replace `‖` with `concat(...)` or add a brief note "(where `‖` denotes channel-wise concatenation)" immediately after the formula. + +______________________________________________________________________ + +## Issue 8 — `__init__` docstring lists `AssertionError` in `Raises` but the checks could be `ValueError` for better API ergonomics + +**Location**: `PositionEmbeddingND.__init__` docstring, `Raises` block. + +**Quoted text**: + +> "Raises: AssertionError: If `data_dim < 1`. AssertionError: If ..." + +The implementation uses `assert` statements, so technically the docstring is correct. However, documenting `AssertionError` as a public API contract is poor practice: `assert` statements are disabled under `python -O`, making the checks silently disappear in optimised deployments. The `Raises` documentation signals to callers that these are expected error conditions, not internal invariants. **Fix**: convert the four `assert` statements in `__init__` to `ValueError` raises (as is already done in `forward`) and update the `Raises` block accordingly. + +______________________________________________________________________ From 7704ed0cec473318cb681ca3b4951ed5468dca1d Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 20:59:10 +0200 Subject: [PATCH 28/72] docs(write/attention): add module and class docstrings with math context --- nvsubquadratic/modules/attention.py | 385 +++++++++++++++++++++++----- 1 file changed, 323 insertions(+), 62 deletions(-) diff --git a/nvsubquadratic/modules/attention.py b/nvsubquadratic/modules/attention.py index 6eb56bb7..ce7b28ed 100644 --- a/nvsubquadratic/modules/attention.py +++ b/nvsubquadratic/modules/attention.py @@ -1,7 +1,102 @@ # TODO: Add license header here -"""Multi-head self-attention with optional QK normalization and Rotary Positional Embeddings (RoPE).""" +r"""Standard scaled dot-product attention for ND spatial inputs. + +Background +---------- +This module implements the classic scaled dot-product attention (Vaswani et al., +"Attention Is All You Need", NeurIPS 2017): + +.. math:: + + \text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right) V + +where :math:`d_k` is the per-head dimension. Multi-head attention computes +:math:`H` independent attention heads in parallel by splitting the channel +dimension :math:`C` into :math:`H` heads each of size :math:`d_k = C / H`, +then concatenating the results: + +.. math:: + + \text{MultiHead}(Q,K,V) = \text{Concat}(\text{head}_1,\ldots,\text{head}_H) + + \text{head}_i = \text{softmax}\!\left(\frac{Q_i K_i^\top}{\sqrt{d_k}}\right) V_i + +The total FLOP count for a sequence of length :math:`L` with :math:`H` heads and +per-head dimension :math:`d_k` is dominated by the attention matrix products: + +.. math:: + + \text{FLOPs} \approx 4 \cdot B \cdot H \cdot L^2 \cdot d_k + +(two matrix multiplications of shape :math:`[L, d_k] \times [d_k, L]` each for +the logit computation and the value aggregation, summed over all heads and the +batch dimension). This is :math:`O(L^2)` in sequence length, which limits +scalability to long sequences — hence the availability of +:class:`~nvsubquadratic.modules.hyena_nd.Hyena` and +:class:`~nvsubquadratic.modules.ckconv_nd.CKConvND` as :math:`O(L \log L)` +alternatives. + +Role in the dispatch pattern +---------------------------- +:class:`Attention` is one of the concrete inner mixers recognised by +:class:`~nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer`. It receives +pre-projected Q, K, V tensors in **channels-last** layout ``[B, *spatial, C]`` +and returns an output of the same shape. Swapping :class:`Attention` for +:class:`~nvsubquadratic.modules.hyena_nd.Hyena` or +:class:`~nvsubquadratic.modules.ckconv_nd.CKConvND` requires only a config +change — the surrounding residual block is agnostic to the choice of mixer. + +The inner mixer contract is: ``forward(q, k, v, cp_group, **kwargs)`` where all +tensors are channels-last ``[B, *spatial, C]`` and ``cp_group`` is the fourth +positional argument. + +Positional encodings +-------------------- +Rotary Positional Embeddings (RoPE, Su et al., "RoFormer", 2021) can be +applied before the attention computation. For a pair of positions :math:`m` and +:math:`n`, RoPE encodes their relative distance through rotation matrices: + +.. math:: + + q_m^\top k_n = \text{Re}\!\left[\sum_{j} q_{m,j} k_{n,j}^* e^{i(m-n)\theta_j}\right] + +where :math:`\theta_j = \text{base}^{-2j/d_k}`. The cos/sin tables are +precomputed once at ``__init__`` time and stored as non-persistent registered +buffers so they survive ``.to(device)`` / ``.half()`` calls and are compatible +with ``torch.compile`` and CUDA-graph capture. 1D, 2D, and 3D spatial layouts +are supported with the following head-dim divisibility requirements: + +- 1D: ``head_dim`` divisible by 2 +- 2D: ``head_dim`` divisible by 4 (two half-dim tables, one per spatial axis) +- 3D: ``head_dim`` divisible by 6 (three one-third-dim tables, one per axis) + +Optional QK normalisation (cosine attention) is also supported; when enabled +the scale factor is set to 1.0 since the norms are already unit-normalised. + +Flash / memory-efficient attention +----------------------------------- +The forward pass delegates to ``torch.nn.functional.scaled_dot_product_attention``, +which automatically selects the most efficient backend available on the current +device (FlashAttention on A100/A10, cuDNN SDPA on H100, a memory-efficient +fallback otherwise). No manual dtype casting is performed; AMP / ``autocast`` +contexts are respected transparently. + +The ``dropout_p`` argument is passed only during training (``module.training`` +is ``True``); it is set to 0.0 at inference so dropout is never applied when +the model is in eval mode. + +Context parallelism +------------------- +A ``cp_group`` argument is accepted for context-parallel training. The current +implementation raises ``ValueError`` immediately — the zigzag all-gather/split +pathway is sketched for future compatibility but does not implement ring-attention +and would therefore materialise the full sequence on every rank. This limitation +means the memory cost scales as :math:`O(L^2 / \text{cp\_size})` per rank +rather than the desired :math:`O(L^2 / \text{cp\_size}^2)`. See +:class:`Attention` class docstring for details. +""" import torch import torch.nn.functional as F @@ -12,54 +107,109 @@ class Attention(torch.nn.Module): - """Multi-head attention (support for self and cross-attention) with optional QK normalization and Rotary Positional Embeddings (RoPE). - - **Input / Output:** - - - Accepts sequences ``[B, T, C]`` or images ``[B, H, W, C]``. - - Returns the same leading shape with channel dimension ``C``. - - **Dimensions:** - - - ``C`` must be divisible by ``num_heads``; per-head dim is ``D = C / num_heads``. - - **RoPE support:** - - - 1D (sequences): applied over length ``T``; requires ``D`` even. - - 2D (images): applied over ``(H, W)`` by splitting per-head dim into two halves - (Y-axis, X-axis); requires ``D`` divisible by 4. - - 3D (videos or other 3D data): applied over ``(H, W, D)`` by splitting per-head dim into three halves - (Z-axis, X-axis, Y-axis); requires ``D`` divisible by 8. - - **QK normalization:** - - - When enabled, queries and keys are L2-normalized per head along the last dimension. - - **Implementation notes:** - - - Uses PyTorch ``scaled_dot_product_attention`` and prefers FlashAttention kernels when available. - - **Context-parallel limitations:** - - - This implementation is for **illustration and compatibility only**. - - It uses zigzag all-gather/split for CP, which causes significant memory issues at long context - lengths because the attention layer does not implement internal ring attention. - - For production use with long contexts, use PyTorch's standard context-parallel attention - blocks (e.g., as in torchtitan: https://docs.pytorch.org/tutorials/unstable/context_parallel.html). - - Future work: migrate to PyTorch's standard CP attention API, which may also eliminate - the requirement for zigzag-style communication patterns. + r"""Multi-head scaled dot-product self-attention for 1D/2D/3D spatial inputs. + + Computes standard multi-head attention: + + .. math:: + + \text{head}_i = \text{softmax}\!\left(\frac{Q_i K_i^\top}{\sqrt{d_k}}\right) V_i + + \text{MultiHead}(Q,K,V) = \text{Concat}(\text{head}_1,\ldots,\text{head}_H) + + where :math:`d_k = C / H` is the per-head dimension, :math:`H` is the + number of heads, and :math:`C` is the hidden (channel) dimension. + + Spatial layout + -------------- + Inputs and outputs use **channels-last** layout: + + - 1D sequences: ``[B, T, C]`` + - 2D images: ``[B, H, W, C]`` + - 3D volumes: ``[B, D, H, W, C]`` + + Internally, spatial dimensions are flattened to a single sequence axis + ``L = prod(spatial_dims)`` for the SDPA kernel, then unflattened on output. + + Multi-head splitting + -------------------- + The channel axis ``C`` is split into ``H`` heads of size ``d_k = C / H``. + Internally the module works with the merged batch-head axis + ``(B * H, L, d_k)`` before the SDPA call and re-merges after. + + QK normalisation (cosine attention) + ------------------------------------ + When ``apply_qk_norm=True``, queries and keys are L2-normalised per head + along the last dimension before the attention logits are formed. This + replaces the ``1/sqrt(d_k)`` scaling with a fixed scale of 1.0 to avoid + flattening the already-normalised logits. + + Rotary Positional Embeddings (RoPE) + ------------------------------------ + RoPE is applied to Q and K before QK-normalisation and before the SDPA + call. The cos/sin buffers are precomputed once at ``__init__`` from + ``rope_spatial_dims`` and stored as non-persistent registered buffers + (``persistent=False``) so they are reconstructed from ``__init__`` args + and never serialised to checkpoints. Head-dim divisibility requirements: + + - 1D: ``head_dim`` divisible by 2 + - 2D: ``head_dim`` divisible by 4 (two half-dim RoPE tables, one per axis) + - 3D: ``head_dim`` divisible by 6 (three one-third-dim RoPE tables) + + Context parallelism (CP) + ------------------------- + When ``cp_group`` is passed at forward time and has size > 1, the module + gathers the full spatial sequence via zigzag all-gather before attention and + splits it back afterwards. **This approach is for illustration and + compatibility only** — it does not implement ring-attention and therefore + materialises the full sequence on every rank, causing memory pressure at + long context lengths. For production long-context training, use PyTorch's + built-in context-parallel SDPA (e.g. torchtitan style). + + Backend selection + ----------------- + Attention is computed with ``torch.nn.functional.scaled_dot_product_attention``, + which auto-selects FlashAttention (A100), cuDNN SDPA (H100), or a + memory-efficient fallback based on device capability. + + Attributes: + hidden_dim (int): Total channel dimension ``C``. + num_heads (int): Number of attention heads ``H``. In the current + implementation all heads are computed on every rank (there is no + head-parallel CP split). A ``# TODO(@farhad)`` in ``forward`` + flags that ``local_num_heads`` is always equal to ``num_heads``, + which may need revisiting for tensor-parallel training. + head_dim (int): Per-head dimension ``d_k = C / H``. + scale (float): Attention logit scale ``1 / sqrt(d_k)``; set to 1.0 + when ``apply_qk_norm=True``. + apply_qk_norm (bool): Whether L2 QK normalisation is active. + use_rope (bool): Whether RoPE positional encoding is active. + rope_base (float): Geometric base for RoPE frequency bands. + is_causal (bool): Whether to apply a causal (auto-regressive) mask. + attn_dropout (float): Dropout probability applied to attention weights + during training. Set to 0.0 automatically at inference regardless + of this value. Args: - hidden_dim (int): The dimension of the hidden states. - num_heads (int): The number of attention heads. - apply_qk_norm (bool): Whether to apply QK normalization. - use_rope (bool): Whether to apply RoPE. - is_causal (bool): Whether the attention is causal. Defaults to False. - attn_dropout (float): The dropout rate for the attention weights. - rope_base (float): The base of the RoPE. - rope_spatial_dims (tuple[int, ...]): Required when ``use_rope=True``. - Spatial dimensions used to precompute the RoPE cos/sin tables. - Examples: ``(4096,)`` for 1D, ``(64, 64)`` for 2D, ``(8, 64, 64)`` for 3D. + hidden_dim (int): Total hidden-state dimension ``C``. Must be + divisible by ``num_heads``. + num_heads (int): Number of parallel attention heads ``H``. + apply_qk_norm (bool): If ``True``, L2-normalise Q and K per head + along the last dimension (cosine attention). + use_rope (bool): If ``True``, apply Rotary Positional Embeddings to + Q and K before the attention logits. + is_causal (bool): If ``True``, apply a causal attention mask so each + position attends only to earlier positions. Defaults to ``False``. + attn_dropout (float): Dropout rate on attention weights (active only + during training). Defaults to ``0.0``. + rope_base (float): Base frequency for RoPE; controls how fast the + rotation frequency decays across head-dim pairs. + Defaults to ``10000.0``. + rope_spatial_dims (tuple[int, ...] | None): Spatial grid shape used + to precompute RoPE tables. Required when ``use_rope=True``. + Examples: ``(4096,)`` for 1D, ``(64, 64)`` for 2D, + ``(8, 64, 64)`` for 3D. Must match the spatial shape seen + during ``forward``. """ def __init__( @@ -73,7 +223,32 @@ def __init__( rope_base: float = 10000.0, rope_spatial_dims: tuple[int, ...] | None = None, ): - """Initialize the SelfAttention module.""" + """Initialise the Attention module and precompute RoPE buffers. + + Args: + hidden_dim (int): Total channel dimension ``C``. Must be + divisible by ``num_heads``. + num_heads (int): Number of attention heads ``H``. + apply_qk_norm (bool): Whether to L2-normalise Q and K per head. + use_rope (bool): Whether to apply Rotary Positional Embeddings. + is_causal (bool): Whether to use a causal attention mask. + Defaults to ``False``. + attn_dropout (float): Attention-weight dropout probability. + Defaults to ``0.0``. + rope_base (float): RoPE base frequency. Defaults to ``10000.0``. + rope_spatial_dims (tuple[int, ...] | None): Spatial grid shape + for RoPE table precomputation. Required when + ``use_rope=True``. + + Raises: + AssertionError: If ``hidden_dim % num_heads != 0``. + AssertionError: If ``use_rope=True`` and ``rope_spatial_dims`` + is ``None``. + AssertionError: If RoPE head-dim divisibility requirements are + not met (divisible by 2 for 1D, 4 for 2D, 6 for 3D). + ValueError: If ``rope_spatial_dims`` has length other than 1, 2, + or 3. + """ super().__init__() assert hidden_dim % num_heads == 0, "hidden_dim must be divisible by num_heads" self.hidden_dim = hidden_dim @@ -141,17 +316,37 @@ def __init__( raise ValueError(f"rope_spatial_dims must be 1D, 2D, or 3D, got {self._rope_ndim}D") def extra_repr(self) -> str: - """Extra repr for the Attention module.""" + """Return a concise string summary of this module's configuration. + + Returns: + str: Comma-separated key=value pairs for ``num_heads``, + ``apply_qk_norm``, ``is_causal``, ``attn_dropout``, + ``use_rope``, and ``rope_base``. + """ return f"num_heads={self.num_heads}, apply_qk_norm={self.apply_qk_norm}, is_causal={self.is_causal}, attn_dropout={self.attn_dropout}, use_rope={self.use_rope}, rope_base={self.rope_base}" def _flatten_spatial(self, x: torch.Tensor) -> tuple[torch.Tensor, tuple[int, ...]]: - """Flattens the spatial dimensions of the input and returns the original spatial shape. + """Flatten all spatial dimensions into a single sequence axis. + + Converts a channels-last tensor with 1–3 spatial dimensions into a + flat sequence tensor ``[B, L, C]`` where ``L = prod(spatial_dims)``. + The original spatial shape is returned so it can be restored by + :meth:`_unflatten_spatial`. Args: - x: torch.Tensor - The input tensor of shape [batch_size, *spatial_dims, hidden_dim]. + x (torch.Tensor): Input tensor of shape ``[B, *spatial_dims, C]`` + where ``len(spatial_dims)`` is 1, 2, or 3. Returns: - tuple[torch.Tensor, tuple[int, ...]]: The flattened tensor and the original spatial shape. + tuple[torch.Tensor, tuple[int, ...]]: + - Flattened tensor of shape ``[B, L, C]`` with + ``L = prod(spatial_dims)``. + - ``spatial_shape``: the original spatial dimensions as a + tuple, needed to invert the operation. + + Raises: + AssertionError: If ``x.ndim`` is not 3, 4, or 5 (i.e., if the + number of spatial dimensions is not 1, 2, or 3). """ assert x.ndim in [3, 4, 5], ( "Input must be tensors with shape (batch_size, *spatial_dims, hidden_dim), where len(spatial_dims) can be 1, 2, or 3." @@ -161,14 +356,38 @@ def _flatten_spatial(self, x: torch.Tensor) -> tuple[torch.Tensor, tuple[int, .. return x, spatial_shape def _unflatten_spatial(self, x: torch.Tensor, spatial_shape: tuple[int, ...]) -> torch.Tensor: - """Invert _flatten_spatial, based on the provided spatial shape. + """Restore spatial dimensions after attention, inverting :meth:`_flatten_spatial`. + + This is the inverse of :meth:`_flatten_spatial`: given a flat ``[B, L, C]`` + tensor and the saved ``spatial_shape``, it reshapes the sequence axis back + into the original spatial grid. + + .. note:: + + For 3D inputs the einops rearrange uses axis binding + ``"b (h w d) c -> b h w d c"`` with + ``h=spatial_shape[0], w=spatial_shape[1], d=spatial_shape[2]``. + The canonical convention (consistent with :meth:`_flatten_spatial`) + is ``spatial_shape = (D, H, W)`` (depth-first), so the three + variables map as: ``h`` ↔ D, ``w`` ↔ H, ``d`` ↔ W. + Callers must always pass the ``spatial_shape`` returned verbatim + by :meth:`_flatten_spatial` — never a hand-constructed tuple. Args: - x: torch.Tensor - The flattened tensor of shape [batch_size, *spatial_dims, hidden_dim]. - spatial_shape: tuple[int, ...] - The original spatial shape. + x (torch.Tensor): Flat sequence tensor of shape ``[B, L, C]`` + where ``L = prod(spatial_shape)``. + spatial_shape (tuple[int, ...]): Original spatial dimensions + ``(T,)``, ``(H, W)``, or ``(D, H, W)``. Must have length + 1, 2, or 3. Obtained from the second return value of + :meth:`_flatten_spatial`. Returns: - torch.Tensor - The unflattened tensor of shape [batch_size, *spatial_dims, hidden_dim]. + torch.Tensor: Tensor restored to ``[B, *spatial_shape, C]``. + For 1D inputs (``len(spatial_shape) == 1``) the tensor is + returned unchanged since the shape is already ``[B, T, C]``. + + Raises: + AssertionError: If ``len(spatial_shape)`` is not 1, 2, or 3. """ assert len(spatial_shape) in [1, 2, 3], "Spatial shape must be a tuple of length 1, 2, or 3." if len(spatial_shape) == 1: @@ -185,16 +404,58 @@ def forward( value: torch.Tensor, cp_group: torch.distributed.ProcessGroup = None, ) -> torch.Tensor: - """Apply multi-head self-attention with optional QK-norm and RoPE. + r"""Apply multi-head scaled dot-product attention. + + Computes: + + .. math:: + + \text{out} = \text{Concat}_{i=1}^{H} + \left[ + \text{softmax}\!\left( + \frac{Q_i K_i^\top}{\sqrt{d_k}} + \right) V_i + \right] + + where :math:`H` = ``num_heads`` and :math:`d_k` = ``head_dim``. + When ``apply_qk_norm=True``, Q and K are L2-normalised before the + logits are formed and the scale is 1.0 instead of ``1/sqrt(d_k)``. + + The forward pipeline is: + + 1. (Optional CP) Zigzag all-gather Q/K/V along the spatial axis. + 2. Split channel dim into heads: ``[B, *spatial, C] → [B*H, *spatial, d_k]``. + 3. (Optional) Apply RoPE to Q and K. + 4. (Optional) L2-normalise Q and K per head. + 5. Flatten spatial dims: ``[B*H, *spatial, d_k] → [B*H, L, d_k]``. + 6. Reshape to SDPA layout: ``[B*H, L, d_k] → [B, H, L, d_k]``. + 7. ``F.scaled_dot_product_attention`` (FlashAttention / cuDNN / fallback). + 8. Merge heads: ``[B, H, L, d_k] → [B, L, C]``. + 9. Unflatten spatial dims: ``[B, L, C] → [B, *spatial, C]``. + 10. (Optional CP) Zigzag-split output back to the local spatial slice. Args: - query: ``[B, *spatial_dims, hidden_dim]`` - key: ``[B, *spatial_dims, hidden_dim]`` - value: ``[B, *spatial_dims, hidden_dim]`` - cp_group: ``torch.distributed.ProcessGroup`` — context-parallel process group. + query (torch.Tensor): Query tensor of shape + ``[B, *spatial_dims, C]``. ``spatial_dims`` may be + ``(T,)``, ``(H, W)``, or ``(D, H, W)``. + key (torch.Tensor): Key tensor of shape + ``[B, *spatial_dims, C]``. Must match ``query`` shape. + value (torch.Tensor): Value tensor of shape + ``[B, *spatial_dims, C]``. Must match ``query`` shape. + cp_group (torch.distributed.ProcessGroup | None): Context-parallel + process group. When not ``None`` and ``cp_group.size() > 1``, + the full spatial sequence is gathered before attention and + split back afterwards. **Currently raises ``ValueError`` + as ring-attention is not yet implemented; provided for + future compatibility.** Defaults to ``None``. Returns: - out: ``[B, *spatial_dims, hidden_dim]`` (same as input) + torch.Tensor: Output of shape ``[B, *spatial_dims, C]``, the + same layout as the inputs. + + Raises: + ValueError: If ``cp_group`` is provided and has size > 1 (context + parallelism is not yet supported). """ # CP communication for self-attention (Megatron-style): # CP only splits sequence/spatial dimension, not hidden_dim or heads From 28130afa699c3aed6ce7686725cfa5d618555c89 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 20:59:50 +0200 Subject: [PATCH 29/72] docs(integrate/position_encoding): apply reviewer feedback --- docs/reviews/position_encoding_review.md | 89 --------------------- nvsubquadratic/modules/position_encoding.py | 80 +++++++++++++----- 2 files changed, 61 insertions(+), 108 deletions(-) delete mode 100644 docs/reviews/position_encoding_review.md diff --git a/docs/reviews/position_encoding_review.md b/docs/reviews/position_encoding_review.md deleted file mode 100644 index 9f0a97cd..00000000 --- a/docs/reviews/position_encoding_review.md +++ /dev/null @@ -1,89 +0,0 @@ -# Review: `nvsubquadratic/modules/position_encoding.py` - -Reviewed as a critical second reader targeting external collaborators reading alongside the paper. - -______________________________________________________________________ - -## Issue 1 — Module docstring mentions "sinusoidal, learned, RFF" variants but only one is implemented - -**Location**: Module docstring, "Variants" section. - -**Quoted text**: - -> "describe the variants present (sinusoidal, learned, RFF, etc.)" - -The module docstring (in the prose leading up to the `PositionEmbeddingND` variant block) does not explicitly acknowledge that sinusoidal and RFF encodings are *absent* from the module. The `PositionEmbeddingND` class description says "For variable-resolution or resolution-generalisation use-cases, prefer a sinusoidal or RFF-based encoding (not yet included in this module)." This is good, but the module-level "Variants" heading implies more variants are present. **Fix**: rename the "Variants" section to "Implemented variants" or explicitly add a "Planned / not yet implemented" sub-list (sinusoidal, RFF) so readers are not confused by the heading. - -______________________________________________________________________ - -## Issue 2 — The factorisation parameter-count formula in the module docstring has a misleading simplification - -**Location**: Module docstring, "ND generalisation strategy" section. - -**Quoted text**: - -> "The total parameter count is therefore `data_dim * max(max_dim_lengths) * (embedding_dim // data_dim) = embedding_dim * max(max_dim_lengths)` — linear in the embedding dimension and the largest axis length." - -The formula `data_dim * max(max_dim_lengths) * (embedding_dim // data_dim)` is only correct when all axes have the *same* maximum length. The true total is `sum(max_dim_lengths[d] * (embedding_dim // data_dim) for d in range(data_dim))`, which is already stated correctly in the class docstring's "Parameter count" section. The module-level simplification introduces a subtle error for non-square grids (e.g. a 2D 56×14 grid). **Fix**: replace the over-simplified formula with the exact sum, or explicitly state "assuming all axes have the same maximum length `M`". - -______________________________________________________________________ - -## Issue 3 — No mention of the `param._no_weight_decay` convention in the module docstring - -**Location**: Module docstring (overview / cross-references). - -The module docstring says nothing about the `_no_weight_decay` tagging convention. An external collaborator writing a custom optimiser builder will need to know that the codebase uses `param._no_weight_decay = True` as a signal. The class docstring explains it, but readers skimming the module overview will miss it. **Fix**: add a one-sentence note in the module docstring (e.g. in the `PositionEmbeddingND` variant description or in a new "Optimiser integration" note) that all embedding parameters are tagged with `_no_weight_decay = True` and that optimiser builders should respect this flag. - -______________________________________________________________________ - -## Issue 4 — `forward` docstring does not explain the broadcast-expand mechanism - -**Location**: `PositionEmbeddingND.forward` docstring, "Returns" section. - -The docstring says the returned tensor has "the same shape as `x`" but does not explain *how* the 1-D per-axis embedding slices are broadcast-expanded to the full `[B, *spatial_dims, per_dim_embedding_dim]` shape before concatenation. A reader inspecting the code will see `emb_axis.view(1, *shape, ...)` followed by `.expand(batch_size, *spatial_dims, ...)` and may not immediately understand why this is correct. **Fix**: add one sentence in the "Returns" section describing the broadcast-expand step: "Each 1-D per-axis embedding of shape `[L_d, per_dim]` is reshaped to `[1, ..., L_d, ..., per_dim]` (with 1 in all axes except `d`) and then broadcast-expanded to `[B, *spatial_dims, per_dim]`; the resulting tensors are concatenated on the last dimension." - -______________________________________________________________________ - -## Issue 5 — The `forward` method validates `x.shape[-1] == embedding_dim` but the docstring's `Args` does not mention this constraint explicitly - -**Location**: `PositionEmbeddingND.forward` docstring, "Args" block for `x`. - -**Quoted text**: - -> "Shape: `[B, *spatial_dims, embedding_dim]`, where the number of spatial axes must equal `data_dim` and the last dimension must equal `embedding_dim`." - -The constraint is present but buried in the shape annotation prose. The `Raises` block lists `ValueError` for channel-dimension mismatch, but does not state what the mismatch condition is ("`x.shape[-1] != embedding_dim`"). **Fix**: in the `Raises` entry for the channel mismatch error, state the condition explicitly: "If `x.shape[-1] != self.embedding_dim` (channel count does not match the embedding dimension passed at construction)." - -______________________________________________________________________ - -## Issue 6 — Missing `device` / `dtype` propagation note in `forward` - -**Location**: `PositionEmbeddingND.forward` docstring. - -The method internally calls `torch.arange(..., device=x.device)` but does not use `x.dtype` for the position indices (correctly using `torch.long`, since they are integer indices). However, the embedding lookup result inherits the dtype of `nn.Embedding.weight`. There is no note clarifying that the output dtype matches the embedding weight dtype (typically `float32`) and that callers working in mixed precision (e.g. `bfloat16`) may need to cast the result. **Fix**: add a `Note` block: "The returned tensor has the same `dtype` as the embedding weight parameters (typically `torch.float32`). In mixed-precision training the result should be cast to match `x.dtype` before addition: `x = x + pos_enc(x).to(x.dtype)`." - -______________________________________________________________________ - -## Issue 7 — Class docstring uses `‖` (U+2016 DOUBLE VERTICAL LINE) for concatenation; this is non-standard and may not render in all docstring viewers - -**Location**: `PositionEmbeddingND` class docstring. - -**Quoted text**: - -> `PE(i_0, ..., i_{D-1}) = [E_0(i_0) ‖ E_1(i_1) ‖ ... ‖ E_{D-1}(i_{D-1})]` - -The `‖` character is ambiguous (it is also the norm operator). Standard ML papers use `concat` or `[· ; ·]`. **Fix**: replace `‖` with `concat(...)` or add a brief note "(where `‖` denotes channel-wise concatenation)" immediately after the formula. - -______________________________________________________________________ - -## Issue 8 — `__init__` docstring lists `AssertionError` in `Raises` but the checks could be `ValueError` for better API ergonomics - -**Location**: `PositionEmbeddingND.__init__` docstring, `Raises` block. - -**Quoted text**: - -> "Raises: AssertionError: If `data_dim < 1`. AssertionError: If ..." - -The implementation uses `assert` statements, so technically the docstring is correct. However, documenting `AssertionError` as a public API contract is poor practice: `assert` statements are disabled under `python -O`, making the checks silently disappear in optimised deployments. The `Raises` documentation signals to callers that these are expected error conditions, not internal invariants. **Fix**: convert the four `assert` statements in `__init__` to `ValueError` raises (as is already done in `forward`) and update the `Raises` block accordingly. - -______________________________________________________________________ diff --git a/nvsubquadratic/modules/position_encoding.py b/nvsubquadratic/modules/position_encoding.py index 1ef97105..43bd74c3 100644 --- a/nvsubquadratic/modules/position_encoding.py +++ b/nvsubquadratic/modules/position_encoding.py @@ -18,8 +18,8 @@ coordinate-dependent bias of the same shape before the tokens are passed into the first mixer block. -Variants --------- +Implemented variants +-------------------- ``PositionEmbeddingND`` **Fully-learned / lookup-table** encoding. A separate ``nn.Embedding`` table is trained for each spatial axis; the per-axis embeddings are @@ -34,6 +34,15 @@ resolution-generalisation use-cases, prefer a sinusoidal or RFF-based encoding (not yet included in this module). +Planned variants (not yet implemented) +--------------------------------------- +* **Sinusoidal** — fixed, non-learned encodings of the form + ``PE(pos, 2i) = sin(pos / 10000^(2i/d))``; resolution-independent and + require no additional parameters. +* **Random Fourier Features (RFF)** — stochastic frequency sampling that + approximates an RBF kernel; a natural complement to the RFF-based implicit + kernel parametrisation in ``nvsubquadratic.modules.kernels_nd``. + ND generalisation strategy -------------------------- A naive learned positional encoding for an ND grid would store @@ -42,10 +51,23 @@ separate embedding table of length ``max_dim_lengths[d]`` is kept for each spatial axis ``d``, and the per-axis embeddings of dimension ``embedding_dim // data_dim`` are concatenated to produce the final -``embedding_dim``-dimensional token offset. The total parameter count is -therefore ``data_dim * max(max_dim_lengths) * (embedding_dim // data_dim) = -embedding_dim * max(max_dim_lengths)`` — linear in the embedding dimension and -the largest axis length. +``embedding_dim``-dimensional token offset. The total parameter count is:: + + sum(max_dim_lengths[d] * (embedding_dim // data_dim) + for d in range(data_dim)) + +which is linear in both the embedding dimension and the sum of the axis +lengths (not their product). For the common square-grid case where all axes +share the same maximum length ``M``, this simplifies to +``embedding_dim * M``. + +Optimiser integration +--------------------- +All embedding parameters in this module are tagged ``param._no_weight_decay = +True`` immediately after construction. Custom optimiser builders (e.g. those +that split parameters into weight-decay and no-weight-decay groups) should +inspect this attribute and exclude matching parameters from L2 regularisation, +following standard ViT practice. Cross-references ---------------- @@ -72,9 +94,9 @@ class PositionEmbeddingND(nn.Module): ``(i_0, i_1, ..., i_{D-1})`` the encoding is formed by looking up the per-axis embeddings and *concatenating* them channel-wise:: - PE(i_0, ..., i_{D-1}) = [E_0(i_0) ‖ E_1(i_1) ‖ ... ‖ E_{D-1}(i_{D-1})] + PE(i_0, ..., i_{D-1}) = concat(E_0(i_0), E_1(i_1), ..., E_{D-1}(i_{D-1})) - where ``‖`` denotes concatenation along the channel dimension and each + where ``concat`` denotes concatenation along the channel dimension and each ``E_d ∈ R^{max_dim_lengths[d] × (embedding_dim // data_dim)}`` is a learned embedding matrix. The result has length ``embedding_dim`` (requires ``embedding_dim % data_dim == 0``). @@ -134,16 +156,20 @@ def __init__(self, embedding_dim: int, data_dim: int, max_dim_lengths: Sequence[ will raise a ``ValueError`` in ``forward``. Raises: - AssertionError: If ``data_dim < 1``. - AssertionError: If ``len(max_dim_lengths) != data_dim``. - AssertionError: If ``data_dim > 3``. - AssertionError: If ``embedding_dim % data_dim != 0``. + ValueError: If ``data_dim < 1``. + ValueError: If ``data_dim > 3``. + ValueError: If ``len(max_dim_lengths) != data_dim``. + ValueError: If ``embedding_dim % data_dim != 0``. """ super().__init__() - assert data_dim >= 1, "data_dim must be >= 1" - assert len(max_dim_lengths) == data_dim, "max_dim_lengths must have length data_dim" - assert data_dim <= 3, "data_dim must be <= 3" - assert embedding_dim % data_dim == 0, "embedding_dim must be divisible by data_dim" + if data_dim < 1: + raise ValueError(f"data_dim must be >= 1, got {data_dim}") + if data_dim > 3: + raise ValueError(f"data_dim must be <= 3, got {data_dim}") + if len(max_dim_lengths) != data_dim: + raise ValueError(f"max_dim_lengths must have length data_dim={data_dim}, got {len(max_dim_lengths)}") + if embedding_dim % data_dim != 0: + raise ValueError(f"embedding_dim={embedding_dim} must be divisible by data_dim={data_dim}") self.embedding_dim = embedding_dim self.per_dim_embedding_dim = embedding_dim // data_dim @@ -192,13 +218,29 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: concatenation of the per-axis embedding lookups:: out[b, i_0, ..., i_{D-1}, :] = - [E_0(i_0) ‖ E_1(i_1) ‖ ... ‖ E_{D-1}(i_{D-1})] + concat(E_0(i_0), E_1(i_1), ..., E_{D-1}(i_{D-1})) + + Internally, for each axis ``d`` the 1-D embedding of shape + ``[L_d, per_dim_embedding_dim]`` is reshaped to + ``[1, ..., L_d, ..., 1, per_dim_embedding_dim]`` (singleton in all + axes except axis ``d``) and then broadcast-expanded to + ``[B, *spatial_dims, per_dim_embedding_dim]`` before the per-axis + tensors are concatenated along the last dimension. + + Note: + The returned tensor has the same ``dtype`` as the embedding weight + parameters (typically ``torch.float32`` regardless of the input + dtype). In mixed-precision training cast the result before adding + it to the token grid:: + + x = x + pos_enc(x).to(x.dtype) Raises: ValueError: If ``x.ndim != data_dim + 2`` (wrong number of dimensions — expected batch + ``data_dim`` spatial + channel). - ValueError: If ``x.shape[-1] != embedding_dim`` (channel dimension - mismatch). + ValueError: If ``x.shape[-1] != self.embedding_dim`` (channel count + of the input does not match the ``embedding_dim`` passed at + construction time). ValueError: If any spatial dimension of ``x`` exceeds the corresponding entry of ``max_dim_lengths``. """ From 005c2fbc888e9baabe2a4fd8014d181e4e61c1fb Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 21:00:06 +0200 Subject: [PATCH 30/72] docs(review/attention): reviewer feedback --- docs/reviews/attention_review.md | 168 +++++++++++++++++++++++++++++++ 1 file changed, 168 insertions(+) create mode 100644 docs/reviews/attention_review.md diff --git a/docs/reviews/attention_review.md b/docs/reviews/attention_review.md new file mode 100644 index 00000000..cb3642e8 --- /dev/null +++ b/docs/reviews/attention_review.md @@ -0,0 +1,168 @@ +# Reviewer Feedback: `nvsubquadratic/modules/attention.py` + +Reviewed after the Phase 1 write pass. Issues are ordered roughly by severity +(correctness hazard → missing information → style). + +______________________________________________________________________ + +## 1. `_unflatten_spatial` 3D rearrange has misleading variable names + +**Location:** `_unflatten_spatial`, line `return rearrange(x, "b (h w d) c -> b h w d c", ...)` + +**Problem:** The einops pattern uses variable names `h`, `w`, `d` but the +actual mapping — given `spatial_shape = (D, H, W)` — is `h` ↔ depth, +`w` ↔ height, `d` ↔ width. This is a semantic mismatch that could mislead +anyone reading the code. The docstring note added in Phase 1 documents the +mismatch but the correct fix is to rename the variables in the rearrange +pattern to `"b (d h w) c -> b d h w c"` with `d=spatial_shape[0], h=spatial_shape[1], w=spatial_shape[2]`, which makes the code self-consistent +with the established `(D, H, W)` convention used everywhere else. + +**Fix:** Change the rearrange string and keyword argument names: + +```python +return rearrange( + x, + "b (d h w) c -> b d h w c", + d=spatial_shape[0], + h=spatial_shape[1], + w=spatial_shape[2], +) +``` + +and remove the warning note from the docstring since the code will then be +unambiguous. + +______________________________________________________________________ + +## 2. Module-level docstring FLOP formula omits QKV projections + +**Location:** Module docstring, "Background" section, FLOP formula: + +> `FLOPs ≈ 4 · B · H · L² · d_k` + +**Problem:** The formula accounts only for the attention matrix products +(QK^T and attention·V). It omits the QKV input projections that live in +`QKVSequenceMixer`, which cost `3 · B · L · C²` additional FLOPs. The +docstring should either (a) state explicitly that these are only the +*attention kernel* FLOPs excluding projections, or (b) give the full formula. + +**Fix:** Add a clarifying parenthetical, e.g.: + +> "(attention kernel only; QKV input/output projections in +> `QKVSequenceMixer` add another `~6·B·L·C²` FLOPs)" + +______________________________________________________________________ + +## 3. Class docstring "Context parallelism (CP)" section contradicts forward docstring + +**Location:** Class docstring, "Context parallelism (CP)" section: + +> "When `cp_group` is passed at forward time and has size > 1, the module +> **gathers the full spatial sequence** via zigzag all-gather before attention +> and splits it back afterwards." + +**Problem:** The `forward` method immediately raises `ValueError("Context parallelism must be revisited.")` before any all-gather is attempted. The +class docstring says the gather-and-split *happens*, when in fact it never +does because the `raise` is unconditional. An external collaborator reading +only the class docstring will believe CP works. + +**Fix:** Rewrite the class docstring CP section to match the forward +docstring's wording — state that `cp_group.size() > 1` raises `ValueError` +and that the implementation is a sketch for future compatibility, not +functional code. + +______________________________________________________________________ + +## 4. `_rope_ndim` attribute is undocumented in the `Attributes` block + +**Location:** Class docstring, `Attributes` section. + +**Problem:** `self._rope_ndim` is set in `__init__` when `use_rope=True` and +used in every branch of the `forward` RoPE block, yet it does not appear in +the `Attributes` block. A collaborator reading the class docstring to +understand the module's state will miss it. + +**Fix:** Add to the Attributes block: + +``` +_rope_ndim (int | None): Spatial rank for which RoPE was initialised + (1, 2, or 3). Present only when ``use_rope=True``; not defined + otherwise. +``` + +______________________________________________________________________ + +## 5. `forward` step 1 description is misleading (CP is non-functional) + +**Location:** `forward` docstring, pipeline step list: + +> "1. (Optional CP) Zigzag all-gather Q/K/V along the spatial axis." + +**Problem:** Step 1 implies the all-gather actually executes when `cp_group` +is provided. In reality, the code raises immediately. This is inconsistent +and confusing. + +**Fix:** Replace step 1 with: + +> "1. (Optional CP, not yet implemented) Raises `ValueError` if +> `cp_group.size() > 1`; pass `cp_group=None` for all current use cases." + +______________________________________________________________________ + +## 6. `rope_spatial_dims` not listed as an attribute when `use_rope=True` + +**Location:** Class docstring, `Attributes` section; `__init__` docstring. + +**Problem:** `rope_spatial_dims` is accepted as an `__init__` argument but +not stored as `self.rope_spatial_dims`, which means it cannot be inspected +after construction. This is fine for internal use but the `__init__` +docstring says it "must match the spatial shape seen during `forward`" without +explaining how a caller should verify or recover this shape post-construction. +`extra_repr` also does not expose it. + +**Fix:** Either store `self.rope_spatial_dims = rope_spatial_dims` and add it +to `Attributes` + `extra_repr`, or add a note in the `__init__` docstring +explaining that `rope_spatial_dims` is intentionally not stored and that the +caller is responsible for tracking it. + +______________________________________________________________________ + +## 7. Missing shape annotation on `_flatten_spatial` return type for batch-head input + +**Location:** `_flatten_spatial` docstring, `Args` section: + +> "x (torch.Tensor): Input tensor of shape `[B, *spatial_dims, C]` where +> `len(spatial_dims)` is 1, 2, or 3." + +**Problem:** At the point `_flatten_spatial` is called in `forward`, `x` has +already been rearranged to `[B*H, *spatial_dims, d_k]` — the leading `B` +dimension is actually `B*H`. The docstring implies the batch axis is always +`B`, which is misleading to someone reading it in the context of `forward`. + +**Fix:** Add a note: + +> "Note: in `forward`, this is called after the channel → head split, so the +> leading dimension is `B * H` and `C` is `d_k = head_dim`." + +______________________________________________________________________ + +## 8. No `Example` block anywhere in the class + +**Problem:** External collaborators benefit from a minimal usage example +showing how to instantiate `Attention` for a typical 2D image use case. +The sibling module `hyena_nd.py` already has inline code examples. + +**Fix:** Add an `Example:` block to the class docstring, e.g.: + +```python +# 2D image attention with 8 heads, RoPE, and cosine-attention QK norm +attn = Attention( + hidden_dim=256, + num_heads=8, + apply_qk_norm=True, + use_rope=True, + rope_spatial_dims=(32, 32), +) +q = k = v = torch.randn(2, 32, 32, 256) # [B, H, W, C] +out = attn(q, k, v) # [B, H, W, C] +``` From fac9d3a1d2924fe032573fb6e70023af920be595 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 21:01:23 +0200 Subject: [PATCH 31/72] docs(integrate/attention): apply reviewer feedback --- docs/reviews/attention_review.md | 168 ---------------------------- nvsubquadratic/modules/attention.py | 67 +++++++---- 2 files changed, 45 insertions(+), 190 deletions(-) delete mode 100644 docs/reviews/attention_review.md diff --git a/docs/reviews/attention_review.md b/docs/reviews/attention_review.md deleted file mode 100644 index cb3642e8..00000000 --- a/docs/reviews/attention_review.md +++ /dev/null @@ -1,168 +0,0 @@ -# Reviewer Feedback: `nvsubquadratic/modules/attention.py` - -Reviewed after the Phase 1 write pass. Issues are ordered roughly by severity -(correctness hazard → missing information → style). - -______________________________________________________________________ - -## 1. `_unflatten_spatial` 3D rearrange has misleading variable names - -**Location:** `_unflatten_spatial`, line `return rearrange(x, "b (h w d) c -> b h w d c", ...)` - -**Problem:** The einops pattern uses variable names `h`, `w`, `d` but the -actual mapping — given `spatial_shape = (D, H, W)` — is `h` ↔ depth, -`w` ↔ height, `d` ↔ width. This is a semantic mismatch that could mislead -anyone reading the code. The docstring note added in Phase 1 documents the -mismatch but the correct fix is to rename the variables in the rearrange -pattern to `"b (d h w) c -> b d h w c"` with `d=spatial_shape[0], h=spatial_shape[1], w=spatial_shape[2]`, which makes the code self-consistent -with the established `(D, H, W)` convention used everywhere else. - -**Fix:** Change the rearrange string and keyword argument names: - -```python -return rearrange( - x, - "b (d h w) c -> b d h w c", - d=spatial_shape[0], - h=spatial_shape[1], - w=spatial_shape[2], -) -``` - -and remove the warning note from the docstring since the code will then be -unambiguous. - -______________________________________________________________________ - -## 2. Module-level docstring FLOP formula omits QKV projections - -**Location:** Module docstring, "Background" section, FLOP formula: - -> `FLOPs ≈ 4 · B · H · L² · d_k` - -**Problem:** The formula accounts only for the attention matrix products -(QK^T and attention·V). It omits the QKV input projections that live in -`QKVSequenceMixer`, which cost `3 · B · L · C²` additional FLOPs. The -docstring should either (a) state explicitly that these are only the -*attention kernel* FLOPs excluding projections, or (b) give the full formula. - -**Fix:** Add a clarifying parenthetical, e.g.: - -> "(attention kernel only; QKV input/output projections in -> `QKVSequenceMixer` add another `~6·B·L·C²` FLOPs)" - -______________________________________________________________________ - -## 3. Class docstring "Context parallelism (CP)" section contradicts forward docstring - -**Location:** Class docstring, "Context parallelism (CP)" section: - -> "When `cp_group` is passed at forward time and has size > 1, the module -> **gathers the full spatial sequence** via zigzag all-gather before attention -> and splits it back afterwards." - -**Problem:** The `forward` method immediately raises `ValueError("Context parallelism must be revisited.")` before any all-gather is attempted. The -class docstring says the gather-and-split *happens*, when in fact it never -does because the `raise` is unconditional. An external collaborator reading -only the class docstring will believe CP works. - -**Fix:** Rewrite the class docstring CP section to match the forward -docstring's wording — state that `cp_group.size() > 1` raises `ValueError` -and that the implementation is a sketch for future compatibility, not -functional code. - -______________________________________________________________________ - -## 4. `_rope_ndim` attribute is undocumented in the `Attributes` block - -**Location:** Class docstring, `Attributes` section. - -**Problem:** `self._rope_ndim` is set in `__init__` when `use_rope=True` and -used in every branch of the `forward` RoPE block, yet it does not appear in -the `Attributes` block. A collaborator reading the class docstring to -understand the module's state will miss it. - -**Fix:** Add to the Attributes block: - -``` -_rope_ndim (int | None): Spatial rank for which RoPE was initialised - (1, 2, or 3). Present only when ``use_rope=True``; not defined - otherwise. -``` - -______________________________________________________________________ - -## 5. `forward` step 1 description is misleading (CP is non-functional) - -**Location:** `forward` docstring, pipeline step list: - -> "1. (Optional CP) Zigzag all-gather Q/K/V along the spatial axis." - -**Problem:** Step 1 implies the all-gather actually executes when `cp_group` -is provided. In reality, the code raises immediately. This is inconsistent -and confusing. - -**Fix:** Replace step 1 with: - -> "1. (Optional CP, not yet implemented) Raises `ValueError` if -> `cp_group.size() > 1`; pass `cp_group=None` for all current use cases." - -______________________________________________________________________ - -## 6. `rope_spatial_dims` not listed as an attribute when `use_rope=True` - -**Location:** Class docstring, `Attributes` section; `__init__` docstring. - -**Problem:** `rope_spatial_dims` is accepted as an `__init__` argument but -not stored as `self.rope_spatial_dims`, which means it cannot be inspected -after construction. This is fine for internal use but the `__init__` -docstring says it "must match the spatial shape seen during `forward`" without -explaining how a caller should verify or recover this shape post-construction. -`extra_repr` also does not expose it. - -**Fix:** Either store `self.rope_spatial_dims = rope_spatial_dims` and add it -to `Attributes` + `extra_repr`, or add a note in the `__init__` docstring -explaining that `rope_spatial_dims` is intentionally not stored and that the -caller is responsible for tracking it. - -______________________________________________________________________ - -## 7. Missing shape annotation on `_flatten_spatial` return type for batch-head input - -**Location:** `_flatten_spatial` docstring, `Args` section: - -> "x (torch.Tensor): Input tensor of shape `[B, *spatial_dims, C]` where -> `len(spatial_dims)` is 1, 2, or 3." - -**Problem:** At the point `_flatten_spatial` is called in `forward`, `x` has -already been rearranged to `[B*H, *spatial_dims, d_k]` — the leading `B` -dimension is actually `B*H`. The docstring implies the batch axis is always -`B`, which is misleading to someone reading it in the context of `forward`. - -**Fix:** Add a note: - -> "Note: in `forward`, this is called after the channel → head split, so the -> leading dimension is `B * H` and `C` is `d_k = head_dim`." - -______________________________________________________________________ - -## 8. No `Example` block anywhere in the class - -**Problem:** External collaborators benefit from a minimal usage example -showing how to instantiate `Attention` for a typical 2D image use case. -The sibling module `hyena_nd.py` already has inline code examples. - -**Fix:** Add an `Example:` block to the class docstring, e.g.: - -```python -# 2D image attention with 8 heads, RoPE, and cosine-attention QK norm -attn = Attention( - hidden_dim=256, - num_heads=8, - apply_qk_norm=True, - use_rope=True, - rope_spatial_dims=(32, 32), -) -q = k = v = torch.randn(2, 32, 32, 256) # [B, H, W, C] -out = attn(q, k, v) # [B, H, W, C] -``` diff --git a/nvsubquadratic/modules/attention.py b/nvsubquadratic/modules/attention.py index ce7b28ed..862d01f2 100644 --- a/nvsubquadratic/modules/attention.py +++ b/nvsubquadratic/modules/attention.py @@ -32,7 +32,10 @@ (two matrix multiplications of shape :math:`[L, d_k] \times [d_k, L]` each for the logit computation and the value aggregation, summed over all heads and the -batch dimension). This is :math:`O(L^2)` in sequence length, which limits +batch dimension; **attention kernel only** — QKV input/output projections in +:class:`~nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer` add +:math:`\sim 6 \cdot B \cdot L \cdot C^2` additional FLOPs). This is +:math:`O(L^2)` in sequence length, which limits scalability to long sequences — hence the availability of :class:`~nvsubquadratic.modules.hyena_nd.Hyena` and :class:`~nvsubquadratic.modules.ckconv_nd.CKConvND` as :math:`O(L \log L)` @@ -158,13 +161,11 @@ class Attention(torch.nn.Module): Context parallelism (CP) ------------------------- - When ``cp_group`` is passed at forward time and has size > 1, the module - gathers the full spatial sequence via zigzag all-gather before attention and - splits it back afterwards. **This approach is for illustration and - compatibility only** — it does not implement ring-attention and therefore - materialises the full sequence on every rank, causing memory pressure at - long context lengths. For production long-context training, use PyTorch's - built-in context-parallel SDPA (e.g. torchtitan style). + **Not yet functional.** Passing a ``cp_group`` with ``size() > 1`` to + ``forward`` immediately raises ``ValueError("Context parallelism must be + revisited.")``. The zigzag all-gather/split code below the ``raise`` is + dead code retained as a sketch for a future ring-attention implementation. + Pass ``cp_group=None`` (the default) for all current use cases. Backend selection ----------------- @@ -189,6 +190,10 @@ class Attention(torch.nn.Module): attn_dropout (float): Dropout probability applied to attention weights during training. Set to 0.0 automatically at inference regardless of this value. + _rope_ndim (int): Spatial rank for which RoPE was initialised (1, 2, + or 3). Present only when ``use_rope=True``; not defined + otherwise. Used in ``forward`` to dispatch to the correct RoPE + apply function. Args: hidden_dim (int): Total hidden-state dimension ``C``. Must be @@ -210,6 +215,23 @@ class Attention(torch.nn.Module): Examples: ``(4096,)`` for 1D, ``(64, 64)`` for 2D, ``(8, 64, 64)`` for 3D. Must match the spatial shape seen during ``forward``. + + Example:: + + import torch + from nvsubquadratic.modules.attention import Attention + + # 2D image attention with 8 heads, RoPE, and cosine-attention QK norm + attn = Attention( + hidden_dim=256, + num_heads=8, + apply_qk_norm=True, + use_rope=True, + rope_spatial_dims=(32, 32), + ) + q = k = v = torch.randn(2, 32, 32, 256) # [B, H, W, C] + out = attn(q, k, v) # [B, H, W, C] + assert out.shape == q.shape """ def __init__( @@ -238,7 +260,12 @@ def __init__( rope_base (float): RoPE base frequency. Defaults to ``10000.0``. rope_spatial_dims (tuple[int, ...] | None): Spatial grid shape for RoPE table precomputation. Required when - ``use_rope=True``. + ``use_rope=True``. **Not stored** as an instance attribute; + the caller is responsible for tracking the spatial dims if they + need to recover them after construction (e.g. for serialisation + or ``extra_repr``). The corresponding cos/sin buffers are + stored as non-persistent registered buffers (``rope_cos``, + ``rope_sin``, etc.). Raises: AssertionError: If ``hidden_dim % num_heads != 0``. @@ -337,6 +364,12 @@ def _flatten_spatial(self, x: torch.Tensor) -> tuple[torch.Tensor, tuple[int, .. x (torch.Tensor): Input tensor of shape ``[B, *spatial_dims, C]`` where ``len(spatial_dims)`` is 1, 2, or 3. + .. note:: + + In :meth:`forward`, this is called after the channel → head + split, so the leading dimension is ``B * H`` (not ``B``) and + ``C`` is ``d_k = head_dim`` (not the full ``hidden_dim``). + Returns: tuple[torch.Tensor, tuple[int, ...]]: - Flattened tensor of shape ``[B, L, C]`` with @@ -362,17 +395,6 @@ def _unflatten_spatial(self, x: torch.Tensor, spatial_shape: tuple[int, ...]) -> tensor and the saved ``spatial_shape``, it reshapes the sequence axis back into the original spatial grid. - .. note:: - - For 3D inputs the einops rearrange uses axis binding - ``"b (h w d) c -> b h w d c"`` with - ``h=spatial_shape[0], w=spatial_shape[1], d=spatial_shape[2]``. - The canonical convention (consistent with :meth:`_flatten_spatial`) - is ``spatial_shape = (D, H, W)`` (depth-first), so the three - variables map as: ``h`` ↔ D, ``w`` ↔ H, ``d`` ↔ W. - Callers must always pass the ``spatial_shape`` returned verbatim - by :meth:`_flatten_spatial` — never a hand-constructed tuple. - Args: x (torch.Tensor): Flat sequence tensor of shape ``[B, L, C]`` where ``L = prod(spatial_shape)``. @@ -395,7 +417,7 @@ def _unflatten_spatial(self, x: torch.Tensor, spatial_shape: tuple[int, ...]) -> elif len(spatial_shape) == 2: return rearrange(x, "b (h w) c -> b h w c", h=spatial_shape[0], w=spatial_shape[1]) else: - return rearrange(x, "b (h w d) c -> b h w d c", h=spatial_shape[0], w=spatial_shape[1], d=spatial_shape[2]) + return rearrange(x, "b (d h w) c -> b d h w c", d=spatial_shape[0], h=spatial_shape[1], w=spatial_shape[2]) def forward( self, @@ -423,7 +445,8 @@ def forward( The forward pipeline is: - 1. (Optional CP) Zigzag all-gather Q/K/V along the spatial axis. + 1. (CP guard) Raises ``ValueError`` if ``cp_group.size() > 1``; + pass ``cp_group=None`` for all current use cases. 2. Split channel dim into heads: ``[B, *spatial, C] → [B*H, *spatial, d_k]``. 3. (Optional) Apply RoPE to Q and K. 4. (Optional) L2-normalise Q and K per head. From f371b4a11684f6cf56242c340a6c38c25129fba3 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 21:09:49 +0200 Subject: [PATCH 32/72] docs(integrate/vit5_residual_block+ckconv_multihead_nd): apply reviewer feedback --- nvsubquadratic/modules/ckconv_multihead_nd.py | 585 +++++++++++++++--- nvsubquadratic/modules/vit5_residual_block.py | 266 ++++++-- 2 files changed, 714 insertions(+), 137 deletions(-) diff --git a/nvsubquadratic/modules/ckconv_multihead_nd.py b/nvsubquadratic/modules/ckconv_multihead_nd.py index 41eb4acc..81e8e2c8 100644 --- a/nvsubquadratic/modules/ckconv_multihead_nd.py +++ b/nvsubquadratic/modules/ckconv_multihead_nd.py @@ -1,23 +1,121 @@ # TODO: Add license header here -"""CKConv Multi-Head (long-convolution with dense within-head mixing) for ND signals. +r"""Multi-head Continuous Kernel Convolution (CKConv) for 2D signals. -This module implements a multi-head variant of CKConv where each head performs -dense [head_dim x head_dim] channel mixing, similar to multi-head attention but -for convolutions. +Background +---------- +This module implements a **multi-head** extension of the CKConv operator +(Romero et al., "CKConv: Continuous Kernel Convolution With Arbitrary +Resolution", ICLR 2022, arXiv:2102.02611). For the single-head (depthwise) +variant see :mod:`nvsubquadratic.modules.ckconv_nd`. -Key difference from standard CKConvND (depthwise): -- CKConvND: Each channel convolved independently (depthwise) -- CKConvMultiheadND: Dense channel mixing within each head +In the single-head variant, every channel is convolved independently with its +own implicit kernel — the kernel has shape ``[C, K_x, K_y]`` and there is no +cross-channel mixing in the convolution layer. The multi-head variant enables +*dense channel mixing within each head*, analogous to multi-head attention but +for convolutions: -Supports an optional low-rank kernel factorization (``kernel_rank``) that -decomposes the full [head_dim x head_dim] kernel into two [head_dim x rank] -factors U and V, reducing SIREN output, FFT conv cost, and memory by -approximately ``2*rank / head_dim``. +- The hidden dimension ``C`` is split into ``H`` heads of ``d = C / H`` + channels each. +- Within head ``h``, the convolution kernel ``K^h`` has shape + ``[d, d, K_x, K_y]``, so every output channel sees all ``d`` input channels + of that head through a spatially-varying weight. +- Across heads, channels remain isolated (no cross-head mixing). -This enables cross-channel feature learning while maintaining computational -efficiency through head isolation. +Formally, for head ``h`` and a 2D input :math:`x^h \in \mathbb{R}^{d \times H \times W}`: + +.. code-block:: none + + y^h = K^h * x^h + shortcut^h ⊙ x^h + +where ``*`` denotes a dense 2D convolution (not depthwise), and the +concatenated output across all heads is: + +.. code-block:: none + + y = [y^0 | y^1 | ... | y^{H-1}] (concatenated along channel axis) + +Each per-head kernel :math:`K^h_\theta` is produced by a shared SIREN network +evaluated on a continuous positional grid normalised to ``[-1, 1]^2``: + +.. code-block:: none + + K^h_\theta(p) = MLP_\theta(pos_enc(p)) + +Because the MLP is small and the same for all heads, the number of parameters +scales with ``H * d^2 = C^2 / H`` (dense) rather than ``C^2`` (full +unstructured mixing), giving a parameter count between the depthwise and fully +dense extremes. + +Low-rank / factored kernel +-------------------------- +For large ``d``, the dense kernel ``[d, d, K_x, K_y]`` per head can be +expensive. An optional low-rank factorisation (``kernel_rank=r``) decomposes +the kernel as: + +.. code-block:: none + + K^h ≈ U^h · V^h, U^h ∈ R^{d × r × K_x × K_y}, + V^h ∈ R^{r × d × K_x × K_y} + +The SIREN outputs ``num_heads * 2 * r * d`` values per position rather than +``num_heads * d^2``, and the frequency-domain contraction is split into two +cheaper steps: + +.. code-block:: none + + z = V^h x (rank-to-d_in contraction) + y = U^h z (d_out-to-rank contraction) + +This reduces the per-head compute from :math:`O(d^2)` to :math:`O(2 d r)` and +the kernel parameter count by the same factor, while preserving the ``d × d`` +mixing capacity at rank ``r``. + +Boundary conditions +------------------- +Two boundary conditions are supported: + +* **Zero-padding** (``fft_padding="zero"``): standard linear convolution with + "same" output size. Kernel size is ``2*N``, covering the full receptive + field without wrap-around. +* **Circular / periodic** (``fft_padding="circular"``): wrap-around + convolution. The kernel size equals the input size (requires + ``grid_type="single"``). + +Shortcut (skip connection) +-------------------------- +Every forward pass adds a per-channel learnable shortcut term: + +.. code-block:: none + + y ← y + shortcut ⊙ x + +where ``shortcut`` is a ``[hidden_dim]`` parameter vector initialised with +Kaiming-uniform scale, identical in design to the shortcut in +:class:`nvsubquadratic.modules.ckconv_nd.CKConvND`. + +Current limitations +------------------- +* Only 2D inputs are supported (``data_dim`` must equal 2). +* Context parallelism (``cp_group``) is not yet implemented. +* Only ``"torch_fft"`` backend (no ``"subq_ops"`` backend support). + +Related modules +--------------- +* :mod:`nvsubquadratic.modules.ckconv_nd` — single-head (depthwise) CKConv, + supporting 1D / 2D / 3D inputs, mixed boundary conditions, and causal mode. +* :mod:`nvsubquadratic.modules.kernels_nd` — implicit kernel parametrisation + (``SIRENKernelND``, FiLM-conditioned variants) consumed by both variants. +* :mod:`nvsubquadratic.ops.fftconv_multihead` — the FFT convolution primitives + that implement the dense per-head spatial mixing called from ``forward``. + +References: +---------- +* Romero et al. (2022). *CKConv: Continuous Kernel Convolution With Arbitrary + Resolution*. ICLR 2022. https://arxiv.org/abs/2102.02611 +* Sitzmann et al. (2020). *Implicit Neural Representations with Periodic + Activation Functions*. NeurIPS 2020. (SIREN kernel used for ``MLP_θ``) """ import math @@ -36,14 +134,87 @@ class CKConvMultiheadND(torch.nn.Module): - """CKConv with multi-head dense channel mixing for 2D signals. - - Each head performs dense [head_dim x head_dim] convolution, enabling - cross-channel feature learning within heads while keeping heads isolated. - - The kernel is generated by a SIREN network that outputs - [num_heads * head_dim * head_dim] values per spatial position (full-rank), - or [num_heads * 2 * kernel_rank * head_dim] values (low-rank). + r"""Multi-head CKConv for 2D signals with dense within-head channel mixing. + + Extends :class:`nvsubquadratic.modules.ckconv_nd.CKConvND` from depthwise + convolution to *dense within-head* convolution by splitting the channel + dimension into ``num_heads`` groups and applying a separate + ``[head_dim × head_dim]`` implicit kernel to each group. + + **Mathematical description** + + Let :math:`x \\in \\mathbb{R}^{B \times C \times H \times W}` with + :math:`C = H_\text{heads} \\cdot d` (``num_heads × head_dim``). The input + is partitioned into heads: + + .. code-block:: none + + x^h ∈ R^{B × d × H × W}, h = 0, ..., H_heads - 1 + + Per head ``h``, the implicit kernel network produces + :math:`K^h_\theta \\in \\mathbb{R}^{d \times d \times K_x \times K_y}` + (full-rank) or its low-rank factorisation + :math:`U^h \\in \\mathbb{R}^{d \times r \times K_x \times K_y}`, + :math:`V^h \\in \\mathbb{R}^{r \times d \times K_x \times K_y}` (low-rank, + when ``kernel_rank`` is set). The output for each head is: + + .. code-block:: none + + y^h = K^h_θ * x^h + shortcut^h ⊙ x^h + + where ``*`` is a dense 2D linear (or circular) convolution computed via + FFT. The final output concatenates all head outputs: + + .. code-block:: none + + y = concat([y^0, y^1, ..., y^{H-1}], dim=channel) + + **Key differences from CKConvND (single-head)** + + * *Dense within-head mixing*: whereas :class:`CKConvND` uses a depthwise + kernel of shape ``[C, K_x, K_y]`` (one scalar kernel per channel), + ``CKConvMultiheadND`` uses a dense kernel of shape + ``[H, d, d, K_x, K_y]`` (a :math:`d \times d` mixing matrix per head + per spatial frequency). + * *No context parallelism*: CP support is not yet implemented; passing + ``cp_group`` to ``forward`` raises ``NotImplementedError``. + * *2D only*: the current implementation only supports ``data_dim=2`` + (images with two spatial axes). The single-head variant supports 1D, + 2D, and 3D inputs. + * *No causal mode, no fp16 FFT, no chunked FFT*: these features from the + single-head variant are not carried over here. + * *Optional low-rank kernel*: the ``kernel_rank`` parameter (absent in + ``CKConvND``) enables a factored ``U · V`` kernel decomposition that + reduces SIREN output size and FFT cost by approximately ``2r / d``. + + Attributes: + data_dim (int): Spatial rank of the input. Always 2 for this class. + hidden_dim (int): Total number of channels ``C = num_heads * head_dim``. + num_heads (int): Number of independent heads ``H``. + head_dim (int): Channels per head ``d = hidden_dim // num_heads``. + fft_padding (str): Boundary condition — ``"zero"`` (linear conv) or + ``"circular"`` (periodic conv). + grid_type (str): Kernel grid size mode — ``"single"`` (kernel size + equals input size, for circular conv) or ``"double"`` (kernel + size is ``2N``, for zero-padded conv). + kernel_rank (int or None): Rank of the low-rank kernel factorisation. + ``None`` means full-rank ``[d, d, K_x, K_y]`` kernels are used. + kernel (nn.Module): Implicit kernel generator (SIREN or similar). + Called as ``kernel(grid_lens, conditioning=...)`` and returns + ``(kernel_values, grid)`` where ``kernel_values`` has shape + ``[1_or_B, K_x, K_y, num_heads * d * d]`` (full-rank) or + ``[1_or_B, K_x, K_y, num_heads * 2 * r * d]`` (low-rank). + mask (nn.Module): Attenuation mask applied to kernel values after + generation. ``nn.Identity`` when no mask is configured. + shortcut (nn.Parameter): Learnable per-channel skip-connection scale + of shape ``(hidden_dim,)``. Added as ``shortcut ⊙ x`` after each + convolution. Initialised with Kaiming-uniform scale + ``uniform(-1/√hidden_dim, 1/√hidden_dim)``. + fftconv_fn (callable): Selected FFT convolution function, one of the + four functions from :mod:`nvsubquadratic.ops.fftconv_multihead`. + Signature varies by full-rank vs. low-rank path: + full-rank: ``(x, kernel, shortcut) → output``; + low-rank: ``(x, kernel_u, kernel_v, shortcut) → output``. """ def __init__( @@ -57,22 +228,79 @@ def __init__( fft_padding: Literal["zero", "circular"], kernel_rank: int | None = None, ): - """Initialize CKConvMultiheadND. + """Construct a CKConvMultiheadND operator. + + Validates parameter combinations, derives ``head_dim``, adjusts the + kernel output-scale for variance control, initialises the ``shortcut`` + parameter, and selects the appropriate FFT convolution function from + :mod:`nvsubquadratic.ops.fftconv_multihead`. Args: - data_dim: Dimension of input data (currently only 2D supported). - hidden_dim: Hidden dimension (must be divisible by num_heads). - num_heads: Number of heads for channel grouping. - kernel_cfg: LazyConfig for the SIREN kernel network. - Full-rank: outputs [num_heads * head_dim * head_dim] values. - Low-rank: outputs [num_heads * 2 * kernel_rank * head_dim] values. - mask_cfg: LazyConfig for the mask (typically Identity). - grid_type: Grid type for kernel generation ("single" or "double"). - fft_padding: FFT padding mode ("zero" or "circular"). - kernel_rank: Rank for low-rank kernel factorization. When None (default), - uses full-rank [head_dim x head_dim] kernels. When set, factorizes - into U @ V^T where U,V have shape [..., head_dim, rank]. Reduces - SIREN output by ``2*rank / head_dim`` and FFT conv cost similarly. + data_dim: Spatial rank of the input signal. Must be ``2``; a + value other than 2 raises ``AssertionError``. + hidden_dim: Total number of channels ``C``. Must be divisible by + ``num_heads``; a violation raises ``AssertionError``. + num_heads: Number of independent convolution heads ``H``. The + channels are split evenly: ``head_dim = hidden_dim // num_heads``. + kernel_cfg: ``LazyConfig`` that instantiates the implicit kernel + generator (e.g. ``SIRENKernelND``). The generator's output + dimension must be set externally to match the expected flat + kernel size: + + * Full-rank: ``out_dim = num_heads * head_dim * head_dim`` + * Low-rank: ``out_dim = num_heads * 2 * kernel_rank * head_dim`` + + The output-scale weight ``kernel.out_linear.weight`` is + multiplied in-place at construction time for variance control + (see Notes). + mask_cfg: ``LazyConfig`` for an optional attenuation mask applied + to the generated kernel values. Use ``torch.nn.Identity`` + for no masking. + grid_type: Relationship between the SIREN coordinate grid and the + input spatial size: + + * ``"single"``: grid spans ``(N+1)//2`` points per axis, + producing a kernel of size ``≈ N`` (for circular conv). + * ``"double"``: grid spans ``N`` points per axis, producing a + kernel of size ``2N - 1 ≈ 2N`` (for zero-padded conv). + + fft_padding: Boundary condition for the convolution: + + * ``"zero"``: zero-padded linear convolution, "same" output + size. + * ``"circular"``: periodic (wrap-around) convolution. Requires + ``grid_type="single"`` and ``K_x == H``, ``K_y == W`` at + runtime (enforced by an ``AssertionError``). + + kernel_rank: Rank ``r`` for the low-rank kernel factorisation. + When ``None`` (default), full-rank ``[d, d, K_x, K_y]`` + kernels are used and the SIREN outputs + ``num_heads * d^2`` values per spatial position. When set to + an integer ``r < d``, the kernel is factored as + :math:`K = U V` with ``U`` of shape ``[d, r, K_x, K_y]`` and + ``V`` of shape ``[r, d, K_x, K_y]``, and the SIREN outputs + ``num_heads * 2 * r * d`` values instead. + + Raises: + AssertionError: If ``data_dim != 2``, ``hidden_dim % num_heads != 0``, + ``grid_type`` is not ``"double"`` or ``"single"``, + ``fft_padding`` is not ``"zero"`` or ``"circular"``, or + ``fft_padding="circular"`` is combined with + ``grid_type != "single"``. + + Notes: + **Output-scale initialisation for variance control.** The SIREN's + final linear layer weight is rescaled in-place (via + ``torch.no_grad()``) so that the convolution output has unit + variance at initialisation: + + * Full-rank: each output channel sums over ``head_dim`` input + channels weighted by the kernel, so the weight is multiplied by + ``1 / √head_dim``. + * Low-rank: the two-step ``U @ V`` contraction has variance that + depends on both ``head_dim`` and ``kernel_rank``, so the weight + is multiplied by ``(1 / (head_dim * kernel_rank))^{1/4}`` — the + geometric mean of the per-factor scales. """ assert data_dim == 2, f"CKConvMultiheadND currently only supports 2D. Got {data_dim}D." assert hidden_dim % num_heads == 0, f"hidden_dim ({hidden_dim}) must be divisible by num_heads ({num_heads})" @@ -125,7 +353,13 @@ def __init__( self.fftconv_fn = fftconv2d_multihead_bhl def extra_repr(self) -> str: - """Return extra representation string for the module.""" + """Return a concise summary string for ``print(module)`` and ``repr(module)``. + + Returns: + A human-readable string listing ``data_dim``, ``hidden_dim``, + ``num_heads``, ``head_dim``, ``fft_padding``, ``grid_type``, and + (if set) ``kernel_rank``. + """ parts = ( f"data_dim={self.data_dim}, hidden_dim={self.hidden_dim}, " f"num_heads={self.num_heads}, head_dim={self.head_dim}, " @@ -136,20 +370,63 @@ def extra_repr(self) -> str: return parts def _reshape_lowrank_kernel(self, conv_kernel_flat: torch.Tensor, K_x: int, K_y: int, B: int | None = None): - """Reshape SIREN output into U and V low-rank factors. + """Reshape the flat SIREN output into the low-rank ``U`` and ``V`` factors. + + The SIREN network outputs a flat tensor whose last dimension encodes + all heads and both low-rank factors interleaved. This method splits + and permutes the tensor into the shapes expected by the low-rank FFT + convolution ops in :mod:`nvsubquadratic.ops.fftconv_multihead`. + + The flat layout (last dim of ``conv_kernel_flat``) is: + + .. code-block:: none + + [head_0_factor_U | head_0_factor_V | head_1_factor_U | ... ] + + More precisely, for each head the SIREN outputs + ``2 * rank * head_dim`` values arranged as + ``[rank, head_dim]`` for U followed by ``[rank, head_dim]`` for V, + with a factor index (``0`` = U, ``1`` = V) in the second-to-last + logical dimension after reshaping to + ``[..., num_heads, 2, rank, head_dim]``. Args: - conv_kernel_flat: [..., K_x, K_y, num_heads * 2 * rank * head_dim] - K_x: Spatial kernel height dimension. - K_y: Spatial kernel width dimension. - B: Batch size (for FiLM-batched kernels), or None for unbatched. + conv_kernel_flat: Flat SIREN output tensor. + + * Unbatched (``B=None``): shape + ``[K_x, K_y, num_heads * 2 * rank * head_dim]``. + * FiLM-batched (``B`` is not ``None``): shape + ``[B, K_x, K_y, num_heads * 2 * rank * head_dim]``. + + K_x: Kernel spatial height (first spatial axis of the SIREN grid). + K_y: Kernel spatial width (second spatial axis of the SIREN grid). + B: Batch size when each sample has its own kernel (FiLM + conditioning). ``None`` for the standard unbatched path where + one kernel is shared across all samples in the batch. Returns: - (kernel_u, kernel_v) tuple: - Unbatched: U [num_heads, head_dim, rank, K_x, K_y], - V [num_heads, rank, head_dim, K_x, K_y] - Batched: U [B, num_heads, head_dim, rank, K_x, K_y], - V [B, num_heads, rank, head_dim, K_x, K_y] + A ``(kernel_u, kernel_v)`` tuple of contiguous tensors: + + * Unbatched (``B=None``): + + * ``kernel_u``: shape ``[num_heads, head_dim, rank, K_x, K_y]`` + — the "output projection" factor. + * ``kernel_v``: shape ``[num_heads, rank, head_dim, K_x, K_y]`` + — the "input projection" factor. + + * FiLM-batched (``B`` is not ``None``): + + * ``kernel_u``: shape ``[B, num_heads, head_dim, rank, K_x, K_y]``. + * ``kernel_v``: shape ``[B, num_heads, rank, head_dim, K_x, K_y]``. + + Notes: + The contraction order in the FFT convolution is + ``z = V x`` (input projection) followed by ``y = U z`` + (output projection), i.e. the einsum chain is + ``(n, r, d_in) × (n, d_in) → (n, r)`` then + ``(n, d_out, r) × (n, r) → (n, d_out)``. The shape convention + for ``kernel_u`` and ``kernel_v`` matches the einsum indices used + in :func:`~CKConvMultiheadND.apply_convolution_batched_lowrank`. """ rank = self.kernel_rank head_dim = self.head_dim @@ -177,15 +454,36 @@ def _reshape_lowrank_kernel(self, conv_kernel_flat: torch.Tensor, K_x: int, K_y: return kernel_u, kernel_v def apply_convolution(self, x: torch.Tensor, conv_kernel: torch.Tensor, shortcut: torch.Tensor) -> torch.Tensor: - """Apply the multi-head convolution (full-rank). + """Apply the full-rank multi-head FFT convolution (shared kernel across batch). + + Calls the pre-selected ``fftconv_fn`` from + :mod:`nvsubquadratic.ops.fftconv_multihead`. The convolution is + performed in float32 regardless of the input dtype to avoid numerical + instability; the output is cast back to the input dtype before + returning. + + The operation per head ``h`` is: + + .. code-block:: none + + y^h = K^h * x^h + shortcut^h ⊙ x^h + + implemented via rFFT in the frequency domain as a dense + ``[head_dim × head_dim]`` matrix multiply at each spatial frequency. Args: - x: Input tensor [B, num_heads, head_dim, H, W], float32 - conv_kernel: Kernel [num_heads, head_dim, head_dim, K_x, K_y], float32 - shortcut: Shortcut tensor [hidden_dim], float32 + x: Input tensor of shape ``[B, num_heads, head_dim, H, W]``, + any floating-point dtype. + conv_kernel: Full-rank kernel tensor of shape + ``[num_heads, head_dim, head_dim, K_x, K_y]``, float32. + ``head_dim`` appears twice — first as ``d_out``, second as + ``d_in``. + shortcut: Learnable skip-connection scale of shape + ``[hidden_dim]``, float32. Fused into the FFT op. Returns: - Output tensor [B, num_heads, head_dim, H, W] + Output tensor of shape ``[B, num_heads, head_dim, H, W]`` in the + same dtype as the input ``x``. """ x_dtype = x.dtype out = self.fftconv_fn( @@ -198,16 +496,35 @@ def apply_convolution(self, x: torch.Tensor, conv_kernel: torch.Tensor, shortcut def apply_convolution_lowrank( self, x: torch.Tensor, kernel_u: torch.Tensor, kernel_v: torch.Tensor, shortcut: torch.Tensor ) -> torch.Tensor: - """Apply the multi-head convolution (low-rank). + """Apply the low-rank multi-head FFT convolution (shared kernel across batch). + + Calls the pre-selected low-rank ``fftconv_fn`` from + :mod:`nvsubquadratic.ops.fftconv_multihead`. The two-step contraction + avoids materialising the full ``[head_dim × head_dim]`` kernel spectrum + per spatial frequency: + + .. code-block:: none + + z = V x (shape: [B, num_heads, rank, H, W]) + y = U z (shape: [B, num_heads, head_dim, H, W]) + y += shortcut ⊙ x + + The convolution is performed in float32; the output is cast back to + the input dtype before returning. Args: - x: Input tensor [B, num_heads, head_dim, H, W] - kernel_u: U factor [num_heads, head_dim, rank, K_x, K_y] - kernel_v: V factor [num_heads, rank, head_dim, K_x, K_y] - shortcut: Shortcut tensor [hidden_dim] + x: Input tensor of shape ``[B, num_heads, head_dim, H, W]``, + any floating-point dtype. + kernel_u: Output-projection factor of shape + ``[num_heads, head_dim, rank, K_x, K_y]``, float32. + kernel_v: Input-projection factor of shape + ``[num_heads, rank, head_dim, K_x, K_y]``, float32. + shortcut: Learnable skip-connection scale of shape + ``[hidden_dim]``, float32. Fused into the FFT op. Returns: - Output tensor [B, num_heads, head_dim, H, W] + Output tensor of shape ``[B, num_heads, head_dim, H, W]`` in the + same dtype as the input ``x``. """ x_dtype = x.dtype out = self.fftconv_fn( @@ -221,18 +538,49 @@ def apply_convolution_lowrank( def apply_convolution_batched( self, x: torch.Tensor, conv_kernel: torch.Tensor, shortcut: torch.Tensor ) -> torch.Tensor: - """Apply multi-head convolution with per-sample (batched) kernels (full-rank). + """Apply the full-rank multi-head FFT convolution with per-sample kernels. + + Used when FiLM conditioning is active and each sample in the batch has + its own kernel (``conv_kernel.shape[0] == B``). The FFT convolution is + implemented directly via ``torch.fft.rfft2`` / ``irfft2`` and + ``torch.einsum``, bypassing the ``fftconv_fn`` (which expects a + shared-kernel layout). + + The zero-padded FFT size is computed per-axis as: + + .. code-block:: none + + fft_h = min(H + (K_x + 1) // 2, 2 * H) + fft_w = min(W + (K_y + 1) // 2, 2 * W) + + which matches the padding convention of the standard (non-batched) + FFT convolution ops in :mod:`nvsubquadratic.ops.fftconv_multihead`. - When FiLM conditioning is used, each sample in the batch has its own kernel. - We perform the FFT convolution with a batched einsum. + The frequency-domain operation per sample is: + + .. code-block:: none + + ŷ_{b,n,o,fx,fy} = Σ_i K̂_{b,n,o,i,fx,fy} · x̂_{b,n,i,fx,fy} + + implemented as a single einsum ``"bnihw,bnoihw->bnohw"``. + + After the inverse FFT, the output is cropped to ``[H, W]`` using a + centered crop starting at ``(K_x // 2, K_y // 2)``, and the shortcut + term is added. Args: - x: Input tensor [B, num_heads, head_dim, H, W], float32 - conv_kernel: Kernel [B, num_heads, head_dim, head_dim, K_x, K_y], float32 - shortcut: Shortcut tensor [hidden_dim], float32 + x: Input tensor of shape ``[B, num_heads, head_dim, H, W]``, + any floating-point dtype. Cast to float32 internally. + conv_kernel: Per-sample full-rank kernel of shape + ``[B, num_heads, head_dim, head_dim, K_x, K_y]``, float32. + ``head_dim`` appears twice (``d_out``, ``d_in``). + shortcut: Learnable skip-connection scale of shape + ``[hidden_dim]``, float32. Reshaped to + ``[1, num_heads, head_dim, 1, 1]`` and added as + ``shortcut ⊙ x`` after the inverse FFT. Returns: - Output tensor [B, num_heads, head_dim, H, W] + Output tensor of shape ``[B, num_heads, head_dim, H, W]``, float32. """ x = x.to(torch.float32) conv_kernel = conv_kernel.to(torch.float32) @@ -267,16 +615,36 @@ def apply_convolution_batched_lowrank( kernel_v: torch.Tensor, shortcut: torch.Tensor, ) -> torch.Tensor: - """Apply multi-head low-rank convolution with per-sample (batched) kernels. + """Apply the low-rank multi-head FFT convolution with per-sample kernels. + + Used when FiLM conditioning is active and each sample in the batch has + its own low-rank kernel pair ``(U, V)``. The two-step frequency-domain + contraction avoids materialising the full ``[head_dim × head_dim]`` + kernel spectrum per frequency bin: + + .. code-block:: none + + ẑ_{b,n,r,fx,fy} = Σ_i V̂_{b,n,r,i,fx,fy} · x̂_{b,n,i,fx,fy} ("bnihw,bnrihw->bnrhw") + ŷ_{b,n,o,fx,fy} = Σ_r Û_{b,n,o,r,fx,fy} · ẑ_{b,n,r,fx,fy} ("bnrhw,bnorhw->bnohw") + + The same zero-padded FFT size and centered-crop convention as + :meth:`apply_convolution_batched` are used. After the inverse FFT the + shortcut term is added element-wise. Args: - x: Input tensor [B, num_heads, head_dim, H, W] - kernel_u: U factor [B, num_heads, head_dim, rank, K_x, K_y] - kernel_v: V factor [B, num_heads, rank, head_dim, K_x, K_y] - shortcut: Shortcut tensor [hidden_dim] + x: Input tensor of shape ``[B, num_heads, head_dim, H, W]``, + any floating-point dtype. Cast to float32 internally. + kernel_u: Per-sample output-projection factor of shape + ``[B, num_heads, head_dim, rank, K_x, K_y]``, float32. + kernel_v: Per-sample input-projection factor of shape + ``[B, num_heads, rank, head_dim, K_x, K_y]``, float32. + shortcut: Learnable skip-connection scale of shape + ``[hidden_dim]``, float32. Reshaped to + ``[1, num_heads, head_dim, 1, 1]`` and added after the inverse + FFT. Returns: - Output tensor [B, num_heads, head_dim, H, W] + Output tensor of shape ``[B, num_heads, head_dim, H, W]``, float32. """ x = x.to(torch.float32) kernel_u = kernel_u.to(torch.float32) @@ -316,19 +684,78 @@ def forward( cp_group: torch.distributed.ProcessGroup = None, **mixer_kwargs, ) -> torch.Tensor: - """Forward pass of CKConvMultiheadND. + """Run the CKConvMultiheadND forward pass. + + Generates the per-head implicit kernel from the SIREN network, + optionally applies the attenuation mask, reshapes the flat SIREN output + into the per-head kernel layout, and applies the FFT convolution with + the shortcut term. + + The computation (non-FiLM path) is: + + .. code-block:: none + + # 1. Determine grid size + grid_lens = [(s + 1) // 2 for s in (H, W)] # if grid_type == "single" + = [H, W] # if grid_type == "double" + + # 2. Generate kernel from SIREN + k_flat, grid = self.kernel(grid_lens) # [1, K_x, K_y, C_flat] + + # 3. Apply mask (if not Identity) + k_flat = self.mask(grid=grid, x=k_flat) + + # 4. Reshape flat output into per-head kernel + # Full-rank: conv_kernel [num_heads, head_dim, head_dim, K_x, K_y] + # Low-rank: kernel_u [num_heads, head_dim, rank, K_x, K_y], + # kernel_v [num_heads, rank, head_dim, K_x, K_y] + + # 5. Apply FFT convolution + shortcut + out_heads = apply_convolution*(x_heads, kernel*, shortcut) + + When the kernel is FiLM-conditioned (``batched_kernel=True``), each + sample receives its own kernel and the batched convolution methods + (:meth:`apply_convolution_batched` or + :meth:`apply_convolution_batched_lowrank`) are used instead. Args: - x: Input tensor. Shape depends on is_bhl_input: - - is_bhl_input=False: [B, H, W, hidden_dim] (BLH format) - - is_bhl_input=True: [B, hidden_dim, H, W] (BHL format) - is_bhl_input: Whether input is in BHL format. - cp_group: Context parallel process group (not supported for multi-head). - **mixer_kwargs: Additional keyword arguments forwarded to the kernel generator - (e.g. ``conditioning`` for FiLM-enabled SIRENKernelND). + x: Input signal tensor. Two supported layouts: + + * **Channels-last** (``is_bhl_input=False``, default): shape + ``[B, H, W, hidden_dim]``. Rearranged internally to + ``[B, num_heads, head_dim, H, W]`` via einops. + * **Channels-first** (``is_bhl_input=True``): shape + ``[B, hidden_dim, H, W]``. Reshaped internally to + ``[B, num_heads, head_dim, H, W]`` via ``view``. + + is_bhl_input: If ``True``, treat ``x`` as channels-first + (BHL) layout. Default: ``False`` (channels-last / BLH). + cp_group: Context-parallel process group. **Not supported** — if + provided and ``cp_group.size() > 1`` a ``NotImplementedError`` + is raised immediately. Accepted as a keyword argument for + interface compatibility with :class:`CKConvND`. + **mixer_kwargs: Additional keyword arguments forwarded to the + kernel generator. Recognised key: + + * ``conditioning`` (``torch.Tensor``, shape ``[B, cond_dim]``): + conditioning vector for FiLM-enabled kernels such as + ``SIRENKernelND`` with a ``film_cfg``. When supplied, the + SIREN returns a per-sample kernel (batch dimension == B); + otherwise the kernel is shared across the batch (batch + dimension == 1). Returns: - Output tensor in same format as input. + Output tensor in the same memory layout as the input ``x``: + + * Channels-last: shape ``[B, H, W, hidden_dim]`` + (when ``is_bhl_input=False``). + * Channels-first: shape ``[B, hidden_dim, H, W]`` + (when ``is_bhl_input=True``). + + Raises: + NotImplementedError: If ``cp_group`` is provided with + ``cp_group.size() > 1``. Context parallelism is not yet + implemented for ``CKConvMultiheadND``. """ if cp_group is not None and cp_group.size() > 1: raise NotImplementedError("Context parallelism not yet supported for CKConvMultiheadND") diff --git a/nvsubquadratic/modules/vit5_residual_block.py b/nvsubquadratic/modules/vit5_residual_block.py index 0f34714f..070bcccf 100644 --- a/nvsubquadratic/modules/vit5_residual_block.py +++ b/nvsubquadratic/modules/vit5_residual_block.py @@ -1,8 +1,42 @@ -"""ViT-5 Residual Block: Pre-norm + Attention/MLP with LayerScale and DropPath. - -Architecture per the ViT-5 paper (Wang et al., 2026): - x = x + DropPath(LayerScale(Attention(Norm(x)))) - x = x + DropPath(LayerScale(MLP(Norm(x)))) +"""ViT-5 Residual Block: Pre-norm + Sequence Mixer/MLP with LayerScale and DropPath. + +This module provides the specialised residual block used throughout the ViT-5 +family of hierarchical vision transformers (Wang et al., 2026). It differs +from the generic :mod:`nvsubquadratic.modules.residual_block` in the following +ViT-5-specific design choices: + +**Structural differences vs. :class:`~nvsubquadratic.modules.residual_block.ResidualBlock`** + +1. **No condition-mixer branch** — :class:`ResidualBlock` supports an optional + cross-attention / conditioning branch between the sequence mixer and the MLP. + :class:`ViT5ResidualBlock` removes this branch entirely; ViT-5 routes + conditioning through register-token pooling instead (see point 3). + +2. **LayerScale on both branches** — ViT-5 wraps both the sequence-mixer and + MLP residual updates in independent :class:`~nvsubquadratic.modules.layer_scale.LayerScale` + modules (initialised to a small constant, typically 1e-4). The generic + :class:`ResidualBlock` uses a single shared dropout and no per-branch + learned scale. + +3. **Register-token FiLM conditioning** — When ``register_pooling_cfg`` is + provided, register tokens are extracted from the *pre-normalised* input, + pooled into a single conditioning vector per sample, and forwarded to the + sequence mixer as a ``conditioning`` keyword argument. This is the primary + mechanism by which ViT-5 communicates global context (e.g. class identity) + to local sequence operators such as Hyena with a SIREN kernel. + +4. **Optional Global Response Normalization (GRN)** — Following ConvNeXt V2, + an optional GRN layer can be inserted after the sequence mixer output to + promote inter-channel competition before the LayerScale + residual add. + +5. **Sequence layout** — Input is always ``[B, T, C]`` (batch, tokens, + channels); the token axis *T* concatenates patch tokens, an optional CLS + token, and register tokens in the order ``[patches, (CLS,) registers]``. + The generic block operates on arbitrary ``(B, *spatial_dims, C)`` tensors. + +For the generic pre-norm residual block (Hyena / Attention / CKConv / Mamba +with optional cross-attention conditioning), see +:mod:`nvsubquadratic.modules.residual_block`. """ import torch @@ -16,30 +50,70 @@ class ViT5ResidualBlock(nn.Module): """ViT-5 style residual block with LayerScale and stochastic depth. - Optionally owns a ``RegisterPooling`` module: when configured, registers - are extracted from the *normalized* input (after ``input_norm``) and pooled - into a conditioning vector that is threaded through ``**mixer_kwargs`` to - the sequence mixer (and ultimately to the SIREN kernel for FiLM). - - Args: - sequence_mixer_cfg: LazyConfig for the sequence mixer (QKVSequenceMixer wrapping Attention). - sequence_mixer_norm_cfg: LazyConfig for the pre-norm before the sequence mixer. - mlp_cfg: LazyConfig for the MLP. - mlp_norm_cfg: LazyConfig for the pre-norm before the MLP. - hidden_dim: Channel dimension (needed for LayerScale). - layer_scale_init: Initial value for LayerScale gammas. Set to 0 to disable LayerScale. - drop_path_rate: Stochastic depth drop probability. - register_pooling_cfg: Optional LazyConfig for RegisterPooling. When provided, - register tokens are extracted from the normalized input and pooled. - num_registers: Number of register tokens (needed for extraction). Only used - when register_pooling_cfg is provided. - register_start_idx: Start index of register tokens in the sequence. - With the standard layout [patches, CLS, registers, ...], this is - ``num_patches + 1`` (CLS readout) or ``num_patches`` (GAP readout). - Typically auto-computed by the network and injected at construction. - grn_cfg: Optional LazyConfig for GlobalResponseNorm (ConvNeXt V2). - When provided, GRN is applied after the sequence mixer output - to promote inter-channel feature competition. + Implements the two-branch pre-norm transformer block used in the ViT-5 + family. Each forward pass executes: + + .. code-block:: text + + # Branch 1 — sequence mixer + x_normed = input_norm(x) + [cond = register_pooling(x_normed[:, s:s+R, :]) # optional] + mixer_out = sequence_mixer(x_normed[, conditioning=cond]) + [mixer_out = grn(mixer_out) # optional] + x = x + drop_path(ls_attn(mixer_out)) + + # Branch 2 — MLP + x = x + drop_path(ls_mlp(mlp(mlp_norm(x)))) + + **Differences vs. the generic** :class:`~nvsubquadratic.modules.residual_block.ResidualBlock`: + + * No condition-mixer branch. Conditioning is handled by register pooling + inside branch 1 (see ``register_pooling_cfg``). + * Each branch has its own :class:`~nvsubquadratic.modules.layer_scale.LayerScale` + (``ls_attn`` / ``ls_mlp``) rather than a single shared dropout. + * A single :class:`~nvsubquadratic.modules.drop_path.DropPath` instance is + shared across both branches (``drop_path``). + * Input is always ``[B, T, C]`` — the spatial dimensions are fully flattened + into the token axis, and register tokens occupy known index positions. + + **Register-token conditioning**: + + When ``register_pooling_cfg`` is not ``None`` and ``num_registers > 0``, + register tokens are extracted from the *normalised* input at positions + ``[register_start_idx : register_start_idx + num_registers]``, pooled by + ``register_pooling`` into a ``(B, C)`` conditioning vector, and passed to + the sequence mixer as ``conditioning=``. This lets the mixer + (typically :class:`~nvsubquadratic.modules.vit5_attention.ViT5Attention` or + :class:`~nvsubquadratic.modules.vit5_hyena_adapter.ViT5HyenaAdapter`) apply + FiLM modulation to its internal kernel. + + Attributes: + input_norm (torch.nn.Module): Pre-norm applied before the sequence + mixer. Parameters are tagged ``_no_weight_decay = True``. + sequence_mixer (torch.nn.Module): Instantiated sequence-mixing + operator — typically a + :class:`~nvsubquadratic.modules.vit5_attention.ViT5Attention` or + :class:`~nvsubquadratic.modules.vit5_hyena_adapter.ViT5HyenaAdapter` + (both include their own QKV / output projections). + mlp_norm (torch.nn.Module): Pre-norm applied before the MLP. + Parameters are tagged ``_no_weight_decay = True``. + mlp (torch.nn.Module): Position-wise MLP applied after ``mlp_norm``. + ls_attn (LayerScale | torch.nn.Identity): Per-element learnable scale + for the sequence mixer branch. ``nn.Identity`` when + ``layer_scale_init == 0``. + ls_mlp (LayerScale | torch.nn.Identity): Per-element learnable scale + for the MLP branch. ``nn.Identity`` when ``layer_scale_init == 0``. + drop_path (DropPath | torch.nn.Identity): Stochastic depth applied + after both LayerScale modules. ``nn.Identity`` when + ``drop_path_rate == 0``. + register_pooling (torch.nn.Module | None): Optional module that maps + ``[B, num_registers, C]`` to a ``(B, C)`` conditioning vector. + ``None`` when register-based conditioning is disabled. + grn (torch.nn.Module | None): Optional Global Response Normalization + (ConvNeXt V2) applied to the sequence mixer output before + ``ls_attn``. ``None`` when disabled. + num_registers (int): Number of register tokens per sample. + register_start_idx (int): Token index at which register tokens begin. """ def __init__( @@ -56,7 +130,55 @@ def __init__( register_start_idx: int = 1, grn_cfg: LazyConfig | None = None, ): - """Instantiate norms, sequence mixer, MLP, and optional register pooling.""" + """Instantiate norms, sequence mixer, MLP, and optional register pooling. + + Args: + sequence_mixer_cfg: LazyConfig for the sequence mixer. Typical + targets are + :class:`~nvsubquadratic.modules.vit5_attention.ViT5Attention` + or + :class:`~nvsubquadratic.modules.vit5_hyena_adapter.ViT5HyenaAdapter`. + Both include QKV and output projections internally (unlike the + generic block where projections live in ``QKVSequenceMixer``). + sequence_mixer_norm_cfg: LazyConfig for the pre-norm applied before + the sequence mixer, e.g. ``RMSNorm(hidden_dim)``. + mlp_cfg: LazyConfig for the position-wise MLP, e.g. + :class:`~nvsubquadratic.modules.mlp.MLP`. + mlp_norm_cfg: LazyConfig for the pre-norm applied before the MLP, + e.g. ``RMSNorm(hidden_dim)``. + hidden_dim: Channel dimension ``C`` shared by all sub-modules. + Used to size :class:`~nvsubquadratic.modules.layer_scale.LayerScale` + (one learnable scalar per channel). + layer_scale_init: Initial value for both ``ls_attn`` and ``ls_mlp`` + LayerScale gammas. Set to ``0`` to replace LayerScale with + ``nn.Identity`` (disables per-channel learned scaling entirely). + Typical values: ``1e-4`` (early training) or ``1.0`` (fine-tune + from a strong checkpoint). + drop_path_rate: Stochastic depth probability. Set to ``0.0`` to + replace ``DropPath`` with ``nn.Identity`` (no drop during + training). A single ``DropPath`` instance is shared between + both branches. + register_pooling_cfg: Optional LazyConfig for a register pooling + module whose ``forward(regs)`` accepts ``[B, num_registers, C]`` + and returns ``(B, C)``. When ``None`` or when + ``num_registers == 0``, register conditioning is disabled and + the sequence mixer is called without a ``conditioning`` kwarg. + num_registers: Number of register tokens ``R`` in the sequence. + Must be consistent with the token layout baked into + ``sequence_mixer_cfg`` (e.g. ``ViT5Attention.num_registers``). + Only used when ``register_pooling_cfg`` is not ``None``. + register_start_idx: Zero-based token index at which the register + block begins. With the standard ViT-5 token layout + ``[patches, CLS, registers]``, this equals + ``num_patches_h * num_patches_w + 1`` for CLS-readout models + and ``num_patches_h * num_patches_w`` for GAP-readout models. + Typically injected by the network constructor. + grn_cfg: Optional LazyConfig for a + :class:`~nvsubquadratic.modules.grn.GlobalResponseNorm` module. + When provided, GRN is applied to the sequence mixer output + (``[B, T, C]``) before ``ls_attn``, promoting inter-channel + feature competition (ConvNeXt V2 recipe). + """ super().__init__() self.input_norm = instantiate(sequence_mixer_norm_cfg) @@ -89,34 +211,37 @@ def __init__( def flop_count(self, num_tokens: int, inference: bool = False) -> int: """Count FLOPs for one ViT-5 residual block. - Architecture per branch: - x = x + DropPath(LayerScale(Mixer(Norm(x)))) - x = x + DropPath(LayerScale(MLP(Norm(x)))) - - Let T = num_tokens. - - FLOPs breakdown: - 1. input_norm (RMSNorm): ``self.input_norm.flop_count(T)`` - 2. register_pooling: ``self.register_pooling.flop_count(D)`` - Only when FiLM conditioning is enabled (register_pooling is not None). - D is inferred from ``self.input_norm.weight.shape[0]``. - 3. sequence_mixer: ``self.sequence_mixer.flop_count(T, inference)`` - Dispatches to ViT5Attention or ViT5HyenaAdapter depending on config. - 4. grn: ``self.grn.flop_count(T)`` - Skipped when GRN is disabled. - 5. ls_attn (LayerScale): ``self.ls_attn.flop_count(T)`` - Skipped when LayerScale is replaced by Identity (init_value=0). - 6. drop_path: 0 (stochastic identity) - 7. mlp_norm (RMSNorm): ``self.mlp_norm.flop_count(T)`` - 8. mlp: ``self.mlp.flop_count(T)`` - 9. ls_mlp (LayerScale): ``self.ls_mlp.flop_count(T)`` + Counts MACs multiplied by 2 (multiply + add) for every sub-module that + exposes a ``flop_count`` method, and falls back to 0 for modules that + do not (e.g. GRN, when ``flop_count`` is absent). + + **Computation graph (one forward pass)**: + + .. code-block:: text + + input_norm → flop_count(T) + register_pooling → flop_count(D) [when enabled] + sequence_mixer → flop_count(T, inf) + grn → flop_count(T) [when enabled; 0 if absent] + ls_attn → flop_count(T) [when LayerScale, not Identity] + mlp_norm → flop_count(T) + mlp → flop_count(T) + ls_mlp → flop_count(T) [when LayerScale, not Identity] + + ``drop_path`` contributes 0 FLOPs (stochastic identity — no arithmetic + on active samples beyond the residual add, which is counted separately + by the caller if desired). Args: - num_tokens: Sequence length T. - inference: Passed through to the sequence mixer. + num_tokens: Sequence length ``T`` passed to each sub-module. This + should equal ``num_patches + (1 if has_cls else 0) + num_registers``. + inference: Passed through to ``self.sequence_mixer.flop_count``. + Some sequence mixers (e.g. Hyena with precomputed kernels) have + lower inference FLOPs than training FLOPs. Returns: - Total FLOPs as an integer. + Total FLOPs as an integer. Does not include the two residual adds + (``2 * T * C`` FLOPs each), which are typically negligible. """ flops = 0 @@ -153,14 +278,39 @@ def flop_count(self, num_tokens: int, inference: bool = False) -> int: return flops def forward(self, x: torch.Tensor, condition: torch.Tensor = None) -> torch.Tensor: - """Forward pass. + """Apply the ViT-5 residual block. + + Executes two residual branches in sequence: + + 1. **Sequence mixer branch** — normalise, optionally extract a register + conditioning vector, run the sequence mixer (and optional GRN), scale + with LayerScale, apply stochastic depth, add residual. + 2. **MLP branch** — normalise, run MLP, scale with LayerScale, apply + stochastic depth, add residual. + + Register conditioning detail: when ``self.register_pooling`` is not + ``None``, the slice ``x_normed[:, register_start_idx : register_start_idx + + num_registers, :]`` is extracted from the *normalised* input and + pooled to shape ``(B, C)``. This vector is forwarded to the sequence + mixer as ``conditioning=``, which the mixer uses for FiLM + modulation (e.g. scaling SIREN kernel features). Args: - x: [B, T, C] where T = num_tokens (patches + cls + registers [+ padding]). - condition: Unused, kept for API compatibility with ResidualBlock. + x: Input token sequence of shape ``[B, T, C]``, where: + - ``B`` — batch size, + - ``T = num_patches + (1 if has_cls else 0) + num_registers`` + — total token count following the ViT-5 layout + ``[patches, (CLS,) registers]``, + - ``C`` — channel (hidden) dimension. + condition: Accepted for API compatibility with + :class:`~nvsubquadratic.modules.residual_block.ResidualBlock` + but **always ignored** in this class. ViT-5 conditioning is + routed through register pooling, not through this argument. + Pass ``None`` (the default) when calling directly. Returns: - [B, T, C] + torch.Tensor: Output tensor of shape ``[B, T, C]``, the same shape + as ``x``. """ x_normed = self.input_norm(x) From 6d495a7a3c3c89bc83d46ea27c28dea32280f6f6 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 21:12:03 +0200 Subject: [PATCH 33/72] docs(tracker): mark attention, position_encoding, ckconv_multihead_nd, vit5_residual_block as done --- docs-tracker.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs-tracker.md b/docs-tracker.md index 6a5adf1c..bd89a0c5 100644 --- a/docs-tracker.md +++ b/docs-tracker.md @@ -45,17 +45,17 @@ Work bottom-up: primitive ops → modules → networks → experiments. | `kernels_nd.py` | \[x\] | Learned kernel parametrisation — RFF + SIREN, FiLM-conditioned variants | | `hyena_nd.py` | \[x\] | Hyena operator (ND) — two-gate sandwich, AllToAll CP, BC-aware convolution | | `ckconv_nd.py` | \[x\] | CKConv (ND) — implicit kernel `k_θ(p) = MLP_θ(pos_enc(p))`, FFT domain, BC modes | -| `ckconv_multihead_nd.py` | \[ \] | Multi-head CKConv | +| `ckconv_multihead_nd.py` | \[x\] | Multi-head CKConv — H heads, dense d×d kernel per head, low-rank U·V factorisation | | `mamba_nd.py` | \[ \] | Mamba SSM (ND) | -| `attention.py` | \[ \] | Standard attention | +| `attention.py` | \[x\] | Scaled dot-product attention — multi-head, RoPE, ND spatial, O(L²) FLOP formula | | `vit5_attention.py` | \[ \] | ViT5 attention variant | | `vit5_hyena_adapter.py` | \[ \] | Hyena adapter for ViT5 | | `sequence_mixer.py` | \[x\] | Operator-agnostic dispatch layer (Hyena / Attention / CKConv / Mamba) | | `condition_mixer.py` | \[ \] | Conditioning mixer | | `residual_block.py` | \[x\] | Residual block — pre-norm + mixer + MLP, optional FiLM/AdaLN-Zero conditioning | -| `vit5_residual_block.py` | \[ \] | ViT5 residual block | +| `vit5_residual_block.py` | \[x\] | ViT5 residual block — LayerScale, register-token conditioning, no condition-mixer branch | | `patchify.py` | \[x\] | Patch embedding — strided conv, 1D/2D/3D, channels-last layout | -| `position_encoding.py` | \[ \] | Position encodings | +| `position_encoding.py` | \[x\] | Axis-factorised learned PE — ND broadcast-expand, float32 output caveat | | `masks_nd.py` | \[ \] | ND masking utils | | `mlp.py` | \[ \] | MLP block | | `film.py` | \[x\] | FiLM conditioning — γ(c)⊙x + β(c), SIREN-based kernel generator | From e6efa18dedbbd91a8adc68fdde1823128e926b60 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 22:23:21 +0200 Subject: [PATCH 34/72] docs(write/vit5_hyena_adapter): add module and class docstrings --- nvsubquadratic/modules/vit5_hyena_adapter.py | 172 +++++++++++++++---- 1 file changed, 138 insertions(+), 34 deletions(-) diff --git a/nvsubquadratic/modules/vit5_hyena_adapter.py b/nvsubquadratic/modules/vit5_hyena_adapter.py index 61f1cbc4..f830ca68 100644 --- a/nvsubquadratic/modules/vit5_hyena_adapter.py +++ b/nvsubquadratic/modules/vit5_hyena_adapter.py @@ -1,13 +1,48 @@ -"""Adapter to plug 2D sequence mixers (e.g. Hyena) into the ViT5 token-sequence architecture. - -The ViT5 architecture processes [B, T, C] sequences. 2D mixers like Hyena expect -[B, H, W, C] spatial grids. This adapter reshapes the flat token sequence to a 2D -grid, applies the inner mixer, and reshapes back. - -Token ordering is handled upstream by the network (ViT5ClassificationNet). -The standard layout is [patches, CLS, registers, padding] where padding -ensures T is divisible by grid_w. This adapter is layout-agnostic: it treats -the entire sequence as a flat spatial grid reshaped to (T // grid_w, grid_w). +"""Adapter that plugs 2-D sequence mixers (e.g. Hyena) into the ViT-5 token-sequence architecture. + +Drop-in replacement interface for :class:`~nvsubquadratic.modules.vit5_attention.ViT5Attention`. + +Why an adapter is needed +------------------------ +:class:`~nvsubquadratic.modules.vit5_attention.ViT5Attention` expects a flat +``[B, T, C]`` token sequence where ``T = num_patches + (1 if has_cls) + num_registers``. +It owns its own QKV and output projections and produces ``[B, T, C]`` output — the +residual block calls ``mixer(x)`` and adds the result back to ``x``. + +2-D operators such as Hyena (wrapped in +:class:`~nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer`) expect a spatial +grid ``[B, H, W, C]`` — they cannot consume the flat sequence directly. Moreover, +``QKVSequenceMixer`` provides its own QKV and output projections that are already +part of the inner mixer; duplicating projections in the adapter would waste memory +and parameters. + +This module solves both issues with a *thin, stateless* reshape adapter: + +1. Receives ``[B, T, C]`` from the residual block. +2. Reshapes to ``[B, T // grid_w, grid_w, C]`` (a 2-D spatial grid). +3. Delegates entirely to the inner mixer (any ``[B, H, W, C]``-in / ``[B, H, W, C]``-out + module, e.g. ``QKVSequenceMixer(Hyena)``). +4. Reshapes back to ``[B, T, C]`` and returns. + +The adapter itself adds **no parameters** — all learnable weights (input projection, +output projection, Hyena kernel generator) live inside the inner mixer. Register +tokens and the CLS token are **not** handled specially here: they are treated as +ordinary spatial positions in the grid. The calling network +(:class:`~nvsubquadratic.networks.vit5_classification.ViT5ClassificationNet`) is +responsible for padding ``T`` so that it is exactly divisible by ``grid_w`` and for +arranging tokens into a layout that makes spatial sense to the mixer. + +Interface contract (same as ``ViT5Attention``) +---------------------------------------------- +``forward(x, **mixer_kwargs) -> Tensor`` + +* Input: ``x`` of shape ``[B, T, C]``. +* Output: tensor of shape ``[B, T, C]``. +* Optional kwargs (e.g. ``conditioning``) are forwarded verbatim to the inner mixer. + +The module also exposes a ``flop_count(num_tokens, inference)`` method that +delegates to the inner mixer's ``flop_count``, matching the API used by the network +for FLOPs accounting. """ import torch @@ -17,11 +52,48 @@ class ViT5HyenaAdapter(nn.Module): - """Bridges ViT5's [B, T, C] token sequences and Hyena's [B, H, W, C] spatial interface. - - Args: - inner_mixer_cfg: LazyConfig for the 2D sequence mixer (e.g. QKVSequenceMixer wrapping Hyena). - grid_w: Width of the 2D spatial grid. The height is inferred as T // grid_w. + """Bridges ViT-5's ``[B, T, C]`` token sequences and Hyena's ``[B, H, W, C]`` spatial interface. + + The adapter is a **parameter-free reshape wrapper**: it does not own any QKV + projection, output projection, or positional encoding. All learnable components + live inside ``inner_mixer`` (typically a + :class:`~nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer` wrapping a + :class:`~nvsubquadratic.modules.hyena_nd.Hyena` instance). + + Data flow:: + + x: [B, T, C] + │ + ▼ reshape (T → H × grid_w, where H = T // grid_w) + x: [B, H, grid_w, C] + │ + ▼ inner_mixer (QKVSequenceMixer → Hyena) + x: [B, H, grid_w, C] + │ + ▼ reshape back + x: [B, T, C] + + What the adapter handles vs. what the inner mixer handles: + + * **Adapter**: shape contract (flat ↔ 2-D), ``flop_count`` delegation. + * **Inner mixer**: input projection (C → 3C), Hyena global convolution, + gating, output projection (C → C), any normalisation, and optional + FiLM / AdaLN conditioning. + + Register-token handling: + Register tokens (and the CLS token, if present) are treated as ordinary + spatial positions within the reshaped grid — no masking or special-casing + is applied. The upstream network is responsible for: + + 1. Padding the sequence so that ``T % grid_w == 0``. + 2. Choosing a ``grid_w`` that places register/CLS tokens in a predictable + row (e.g. a dedicated "register row" at the bottom of the grid), so + that the spatial convolution inside Hyena sees a consistent layout. + + Attributes: + inner_mixer (nn.Module): The instantiated 2-D sequence mixer. + grid_w (int): Width of the 2-D spatial grid. The height is inferred + at runtime as ``T // grid_w``. """ def __init__( @@ -29,43 +101,75 @@ def __init__( inner_mixer_cfg: LazyConfig, grid_w: int, ): - """Store config and instantiate the inner 2D mixer.""" + """Instantiate the adapter and its inner 2-D mixer. + + Args: + inner_mixer_cfg: :class:`~nvsubquadratic.lazy_config.LazyConfig` + describing the 2-D sequence mixer to instantiate (e.g. + ``QKVSequenceMixer`` wrapping ``Hyena``). The instantiated module + must accept ``(x: Tensor[B, H, W, C], **kwargs)`` and return a + tensor of the same shape. + grid_w: Width of the 2-D spatial grid. Every call to ``forward`` + must supply a sequence length ``T`` that satisfies + ``T % grid_w == 0``; the grid height is computed as + ``H = T // grid_w``. + """ super().__init__() self.inner_mixer = instantiate(inner_mixer_cfg) self.grid_w = grid_w def flop_count(self, num_tokens: int, inference: bool = False) -> int: - """Count FLOPs for the Hyena adapter (reshape + inner mixer). + """Delegate FLOPs accounting to the inner mixer. - Reshapes the flat token sequence [B, T, C] into a 2D spatial grid - [B, H, W, C] (where W = ``self.grid_w``, H = T // W) and delegates - to the inner mixer (QKVSequenceMixer wrapping Hyena). - - The reshape itself is a free metadata operation — no FLOPs. + The adapter's reshape operations are pure metadata re-strides — zero + arithmetic FLOPs — so the total cost is entirely determined by + ``inner_mixer.flop_count``. Args: - num_tokens: Total sequence length T. Must be divisible by grid_w. - ``spatial_dims`` is computed as ``(T // grid_w, grid_w)``. - inference: Passed through to the inner mixer. + num_tokens: Total flat sequence length ``T``. Must satisfy + ``T % grid_w == 0``. The 2-D spatial dimensions passed to the + inner mixer are ``(T // grid_w, grid_w)``. + inference: Forwarded to the inner mixer. Some mixers (e.g. those + with cached Hyena kernels) report fewer FLOPs at inference time. Returns: - Total FLOPs from the inner mixer. + Total FLOPs reported by the inner mixer for a ``(T // grid_w, grid_w)`` + spatial grid. + + Raises: + AttributeError: If ``inner_mixer`` does not implement ``flop_count``. """ spatial_dims = (num_tokens // self.grid_w, self.grid_w) return self.inner_mixer.flop_count(spatial_dims, inference=inference) def forward(self, x: torch.Tensor, **mixer_kwargs) -> torch.Tensor: - """Forward pass. + """Reshape to 2-D grid, apply the inner mixer, reshape back. Args: - x: [B, T, C] token sequence. T must be divisible by grid_w. - Standard layout: [patches (H*W), CLS (1), registers (R), padding (P)]. - The adapter is layout-agnostic and reshapes the full sequence - to a 2D grid of shape (T // grid_w, grid_w). - **mixer_kwargs: Forwarded to the inner mixer (e.g. ``conditioning`` for FiLM). + x: Input token sequence of shape ``[B, T, C]`` where + + * ``B`` — batch size. + * ``T`` — total sequence length (must satisfy ``T % grid_w == 0``). + Typical layout (set by the network, not enforced here): + ``[patch_tokens (H_patch * W_patch), CLS (0 or 1), + register_tokens (R), padding (P)]``. + * ``C`` — channel / hidden dimension. + + **mixer_kwargs: Keyword arguments forwarded verbatim to + ``inner_mixer.forward``. Common keys include: + + * ``conditioning`` — FiLM/AdaLN conditioning tensor used by + some Hyena configurations. + * ``cp_group`` — process group for context-parallel (AllToAll) + sharding inside the Hyena operator. Returns: - [B, T, C] with tokens mixed via the 2D inner mixer. + Tensor of shape ``[B, T, C]`` — the token sequence after 2-D + Hyena mixing. The reshape is a view (no data copy) when the + tensor is contiguous. + + Raises: + RuntimeError: If ``T % grid_w != 0`` (implicit, from ``reshape``). """ B, T, C = x.shape x = x.reshape(B, T // self.grid_w, self.grid_w, C) @@ -74,5 +178,5 @@ def forward(self, x: torch.Tensor, **mixer_kwargs) -> torch.Tensor: return x def extra_repr(self) -> str: - """Return grid width for repr().""" + """Return a concise summary for ``repr()`` and ``print(model)``.""" return f"grid_w={self.grid_w}" From b02deb2d1e54ab70ef4aad954747b1d8ec786d32 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 22:23:32 +0200 Subject: [PATCH 35/72] docs(write/condition_mixer): add module and class docstrings --- nvsubquadratic/modules/condition_mixer.py | 184 ++++++++++++++++++++-- 1 file changed, 173 insertions(+), 11 deletions(-) diff --git a/nvsubquadratic/modules/condition_mixer.py b/nvsubquadratic/modules/condition_mixer.py index f3cdde7c..81ed4124 100644 --- a/nvsubquadratic/modules/condition_mixer.py +++ b/nvsubquadratic/modules/condition_mixer.py @@ -15,7 +15,59 @@ # Adapted from https://github.com/implicit-long-convs/ccnn_v2 -"""QKV condition mixer for conditioning.""" +"""QKV cross-attention condition mixer for injecting a conditioning signal into a feature map. + +**Role in the residual block** + +The conditioning mixer is the *middle* sub-branch of +:class:`~nvsubquadratic.modules.residual_block.ResidualBlock`. After the +sequence mixer has let spatial positions exchange information, the condition +mixer injects an external conditioning signal ``c`` — such as a diffusion +timestep embedding, a class label embedding, or a physics parameter vector — +into the residual stream ``x``: + +.. code-block:: text + + x ──[seq_mixer]──► x' + x'──[cond_mixer(x', c)]──► x'' ← this module + x''─[mlp]──────────────► output + +**Comparison with other conditioning strategies** + +* **FiLM / AdaLN-Zero** — applies a *per-channel* affine transform + ``y = γ(c) ⊙ x + β(c)`` to the feature map (see + :mod:`nvsubquadratic.modules.film`). This is fast and parameter-efficient + but does not allow the conditioning signal to attend selectively to specific + spatial positions. + +* **Cross-attention** — routes the conditioning signal through standard + Q/K/V attention so that every position in ``x`` can attend to all + conditioning tokens. Highly expressive but O(L_x · L_c) in cost. + +* **QKVConditionMixer** *(this module)* — implements a lightweight cross- + attention variant where queries come from the feature map ``x`` and keys/ + values come from the conditioning signal ``c``. The concrete attention + computation is delegated to a configurable inner ``mixer`` module (supplied + via ``mixer_cfg``), giving the same operator-agnostic dispatch pattern as + :class:`~nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer`. The + result is more expressive than FiLM (it can attend to individual conditioning + tokens) while keeping the projections and mixing strategy swappable. + +**Conditioning signal shapes** + +The condition tensor ``c`` may be: + +* ``(B, C)`` — a global (non-spatial) conditioning vector, e.g. a pre-pooled + timestep or class embedding. Internally unsqueezed to ``(B, 1, C)`` before + the K/V projection so that the inner mixer sees a single conditioning token. +* ``(B, *spatial_dims_cond, C)`` — spatially distributed conditioning tokens + (e.g. encoder output in an encoder-decoder architecture). Must have the + same number of dimensions as the feature map ``x``. + +In both cases the channel dimension ``C`` must equal ``hidden_dim``. + +All tensors use **channels-last** layout: ``(B, *spatial, C)``. +""" from typing import Callable @@ -25,7 +77,62 @@ class QKVConditionMixer(torch.nn.Module): - """QKV condition mixer.""" + """Cross-attention condition mixer that routes a conditioning signal into the feature map. + + This module implements the *condition mixer branch* of + :class:`~nvsubquadratic.modules.residual_block.ResidualBlock`. It injects + an external conditioning signal ``c`` (e.g. a timestep embedding, class + label, or physics parameter vector) into the residual stream ``x`` via + learned Q, K, V projections and a pluggable inner mixing operator: + + .. code-block:: text + + x ─[q_proj]──────────────────────────► Q ─┐ + c ─[kv_proj]──► split ──► K, V ────────────► inner_mixer(Q, K, V) ─[out_proj]──► y + + Queries are derived from the current feature map ``x`` so that each spatial + position can attend selectively to the conditioning tokens. Keys and values + are derived from the conditioning signal ``c``. This is therefore a form of + **cross-attention** conditioning — more expressive than FiLM + (which applies a uniform per-channel affine transform regardless of spatial + content) and more efficient than full self-attention between concatenated + feature and conditioning tokens. + + The inner mixing computation is delegated to ``mixer_cfg``, following the + same operator-agnostic dispatch pattern as + :class:`~nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer`. Any + module whose ``forward(q, k, v)`` conforms to channels-last tensors of + shape ``(B, *spatial, C)`` can be used as the inner mixer. + + **Weight initialisation** + + Optional curried initialisers ``init_method_in`` and ``init_method_out`` + allow the caller to supply per-projection weight schedules (e.g. + depth-scaled Gaussian init from GPT / Megatron). Both follow the + signature ``fn(dim: int) -> fn(tensor: Tensor) -> None``. When omitted, + PyTorch's default Kaiming-uniform init is used. + + **Weight decay** + + No weight-decay tags are set on the projections; the caller is responsible + for any per-parameter weight-decay grouping (see the analogous logic in + :class:`~nvsubquadratic.modules.film.KernelFiLMGenerator`). + + Attributes: + mixer (torch.nn.Module): The instantiated inner mixing operator. Its + ``forward(q, k, v)`` method receives channels-last tensors of shape + ``(B, *spatial, C)`` and must return a tensor of the same shape as + ``q``. + kv_proj (torch.nn.Linear): Combined K+V projection (no bias) that maps + the conditioning signal from ``C`` to ``2·C``. Weight shape: + ``(2·hidden_dim, hidden_dim)``. + q_proj (torch.nn.Linear): Query projection (no bias) that maps the + feature map from ``C`` to ``C``. Weight shape: + ``(hidden_dim, hidden_dim)``. + out_proj (torch.nn.Linear): Output projection (no bias) that maps the + mixer output back to ``C``. Weight shape: + ``(hidden_dim, hidden_dim)``. + """ def __init__( self, @@ -34,13 +141,28 @@ def __init__( init_method_in: Callable[[torch.Tensor], torch.Tensor] | None = None, init_method_out: Callable[[torch.Tensor], torch.Tensor] | None = None, ): - """Initialize the QKV condition mixer. + """Initialise the QKVConditionMixer. Args: - hidden_dim: Hidden dimension. - mixer_cfg: LazyConfig for the condition mixer layer. - init_method_in: Optional initialization method for the KV and Q projections. - init_method_out: Optional initialization method for the output projection. + hidden_dim: Channel dimension ``C`` shared by the feature map ``x`` + and the conditioning signal ``c``. All four linear projections + (``q_proj``, ``kv_proj``, ``out_proj``) are sized using this + value. + mixer_cfg: :class:`~nvsubquadratic.lazy_config.LazyConfig` for the + inner mixing operator. The instantiated module's ``forward`` + must accept ``(q, k, v)`` as positional arguments and return a + tensor of the same shape as ``q``. Any attention-compatible + module (e.g. a dot-product attention layer) can be used here. + init_method_in: Optional *curried* weight initialiser applied to + both ``q_proj.weight`` and ``kv_proj.weight``. Must have the + signature ``fn(dim: int) -> fn(tensor: Tensor) -> None``. + When provided, ``fn(hidden_dim)`` is called and the returned + callable is applied in-place to each weight matrix. Pass + ``None`` to keep PyTorch's default Kaiming-uniform init. + init_method_out: Same as ``init_method_in`` but applied to + ``out_proj.weight``. A common choice is a depth-scaled + Gaussian (GPT / Megatron style) to control the residual branch + variance at initialisation. Pass ``None`` to keep the default. """ super().__init__() @@ -61,14 +183,54 @@ def __init__( init_method_out(hidden_dim)(self.out_proj.weight.data) def forward(self, x: torch.Tensor, condition: torch.Tensor) -> torch.Tensor: - """Forward pass of the QKV condition mixer. + """Inject the conditioning signal into the feature map via cross-attention. + + Computes queries from the current feature map ``x`` and keys/values + from the conditioning signal ``condition``, then mixes them with the + inner ``mixer`` and projects the result back to ``hidden_dim``. + + The signal flow is: + + .. code-block:: text + + Q = q_proj(x) # (B, *spatial_dims, C) + K, V = split(kv_proj(condition)) # each (B, *spatial_dims_cond, C) + y = out_proj(mixer(Q, K, V)) # (B, *spatial_dims, C) + + A global (non-spatial) conditioning vector of shape ``(B, C)`` is + automatically unsqueezed to ``(B, 1, C)`` before the K/V projection + so that the inner mixer sees a single conditioning token per sample. Args: - x: Input tensor of shape [B, * spatial_dims, hidden_dim]. - condition: Condition tensor of shape [B, * spatial_dims_condition, hidden_dim] or [B, hidden_dim]. + x: Feature map tensor of shape ``(B, *spatial_dims, C)``, where + ``B`` is the batch size, ``spatial_dims`` is one or more + spatial axes (e.g. ``(H, W)`` for 2-D images or ``(T,)`` for + 1-D sequences), and ``C = hidden_dim``. Must have at least + three dimensions (batch + one spatial axis + channel). + condition: Conditioning signal tensor. Two shapes are accepted: + + * ``(B, C)`` — global conditioning vector (e.g. a timestep or + class embedding). Unsqueezed internally to ``(B, 1, C)`` + before projection. + * ``(B, *spatial_dims_cond, C)`` — spatially distributed + conditioning tokens (e.g. encoder output). Must have the + same number of dimensions as ``x``. The spatial extent + ``spatial_dims_cond`` need not match ``spatial_dims`` of + ``x``. + + The channel dimension ``C`` must equal ``hidden_dim`` in both + cases. Returns: - Output tensor of shape [B, * spatial_dims, hidden_dim]. + Output tensor of shape ``(B, *spatial_dims, C)`` — same spatial + layout as ``x``, with the conditioning signal blended in via + cross-attention. + + Raises: + ValueError: If ``x`` has fewer than three dimensions (i.e. is + missing at least one spatial axis). + ValueError: If ``condition.ndim`` is neither ``2`` (global vector) + nor equal to ``x.ndim`` (matching spatial rank). """ if x.ndim < 3: raise ValueError(f"x must have at least one spatial dimension; got shape {x.shape}.") From 6014acd8507ab430eae2e0ec424c2f2ecd0f7566 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 22:23:52 +0200 Subject: [PATCH 36/72] docs(write/mamba_nd): add module and class docstrings with math context --- nvsubquadratic/modules/mamba_nd.py | 227 +++++++++++++++++++++++++++-- 1 file changed, 214 insertions(+), 13 deletions(-) diff --git a/nvsubquadratic/modules/mamba_nd.py b/nvsubquadratic/modules/mamba_nd.py index 379abb81..ccc8622e 100644 --- a/nvsubquadratic/modules/mamba_nd.py +++ b/nvsubquadratic/modules/mamba_nd.py @@ -1,7 +1,118 @@ # TODO: Add license header here -"""Mamba mixer layer for ND signals.""" +r"""Mamba-ND: selective state-space mixer for 1D/2D/3D signals. + +Background +---------- +Mamba (Gu & Dao, "Mamba: Linear-Time Sequence Modeling with Selective State +Spaces", arXiv:2312.00752) is a **selective state-space model (SSM)** that +achieves linear time complexity in sequence length — O(N) — while retaining +the long-range modelling capacity of attention. + +The core SSM recurrence is: + +.. math:: + + h_t &= \bar{A}_t \, h_{t-1} + \bar{B}_t \, x_t \\ + y_t &= C_t \, h_t + +where :math:`h_t \in \mathbb{R}^{d \times N}` is a latent state, and +:math:`\bar{A}_t`, :math:`\bar{B}_t` are **input-dependent** discretised +transition matrices. The key departure from classical linear SSMs (e.g. +S4, S5) is that :math:`B`, :math:`C`, and the step size :math:`\Delta` are +*functions of the input* :math:`x_t` rather than fixed parameters: + +.. math:: + + \Delta_t,\ B_t,\ C_t = \text{Linear}(x_t) + +The continuous-time state matrix :math:`A` is discretised via the +**zero-order hold (ZOH)** rule: + +.. math:: + + \bar{A}_t = e^{\Delta_t A}, \qquad + \bar{B}_t = (e^{\Delta_t A} - I) A^{-1} B_t \approx \Delta_t B_t + +This selectivity allows Mamba to focus on relevant tokens and ignore +irrelevant context, giving it an advantage over fixed-kernel convolutions +(Hyena, CKConv) on tasks requiring content-based filtering, while remaining +subquadratic unlike attention. + +Comparison with other mixers +----------------------------- ++-----------+-------------------+--------------------+----------------------+ +| Mixer | Sequence-mixing | Kernel | Complexity (in N) | ++===========+===================+====================+======================+ +| Attention | pairwise dot-prod | input-dependent | O(N²) | ++-----------+-------------------+--------------------+----------------------+ +| Hyena | FFT convolution | fixed (learned MLP)| O(N log N) | ++-----------+-------------------+--------------------+----------------------+ +| Mamba | SSM recurrence | input-dependent | O(N) | ++-----------+-------------------+--------------------+----------------------+ + +ND generalisation strategy +-------------------------- +The Mamba recurrence is inherently sequential and 1D. This module extends it +to arbitrary spatial rank (1D sequences, 2D images, 3D volumes) by **flattening +all spatial axes into a single sequence dimension** before running the core Mamba +layer: + +.. code-block:: none + + [B, *spatial, C] + │ rearrange "b ... c -> b (...) c" + ▼ + [B, S, C] where S = prod(spatial_dims) + │ Mamba1D core (or bidirectional pair) + ▼ + [B, S, C] + │ reshape back to original spatial layout + ▼ + [B, *spatial, C] + +The scan order for multi-dimensional inputs follows the default PyTorch / +``einops`` row-major (C-contiguous) flattening: for a 2D ``[H, W]`` input the +tokens are visited in raster-scan order (row 0, col 0 → row 0, col W-1 → +row 1, col 0 → …). For 3D ``[D, H, W]`` inputs the outermost axis varies +slowest. This ordering is fixed and is not learned; future work could explore +Hilbert-curve or zigzag orderings for improved spatial locality. + +Bidirectional mode +------------------ +Setting ``bidirectional=True`` instantiates a second Mamba layer +(``core_layer_rev``) that processes the token sequence in **reverse order**. +The reversed output is flipped back and added to the forward output: + +.. math:: + + \text{out} = \text{Mamba}(x) + \text{flip}(\text{Mamba}_\text{rev}(\text{flip}(x))) + +This makes the effective receptive field of each position span the entire +sequence in both directions, at the cost of 2× parameters and compute. +For non-causal spatial tasks (images, volumes) bidirectional mode is strongly +recommended. + +Integration with the rest of the library +----------------------------------------- +:class:`Mamba` is designed to be used as the ``inner_mixer`` inside +:class:`~nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer` (which adds +shared QKV and output linear projections). It is also listed as a supported +mixer in :class:`~nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer`'s +dispatch table. + +Related modules +--------------- +* ``nvsubquadratic.modules.hyena_nd`` — Hyena (fixed-kernel gated conv) +* ``nvsubquadratic.modules.attention`` — multi-head self-attention +* ``nvsubquadratic.modules.sequence_mixer`` — operator-agnostic dispatch layer + +References: +---------- +Gu, A. & Dao, T. (2023). *Mamba: Linear-Time Sequence Modeling with Selective +State Spaces*. arXiv:2312.00752. +""" import torch from einops import rearrange @@ -10,14 +121,64 @@ class Mamba(torch.nn.Module): - """Mamba mixer layer for ND signals. + r"""Selective state-space mixer for ND signals. + + Wraps a 1D Mamba core layer (e.g. ``mamba_ssm.Mamba``) and extends it to + arbitrary spatial rank by flattening all spatial axes into a single + sequence dimension before the SSM recurrence and reshaping back afterward. + + The SSM recurrence computed by the core layer is: + + .. math:: + + h_t &= \bar{A}_t \, h_{t-1} + \bar{B}_t \, x_t \\ + y_t &= C_t \, h_t + + where the transition matrices :math:`\bar{A}_t`, :math:`\bar{B}_t` and + the readout matrix :math:`C_t` are all *functions of* :math:`x_t`, + derived via linear projections inside the core layer. The step size + :math:`\Delta_t` (also input-dependent) controls the discretisation via + the zero-order hold (ZOH) rule: + + .. math:: + + \bar{A}_t = e^{\Delta_t A}, \qquad + \bar{B}_t \approx \Delta_t B_t - ND signals are handled by reshaping the input to [B, flatten (* spatial_dims), hidden_dim] and then passing it to the core layer. + **Scan order for ND inputs**: spatial axes are flattened in row-major + (C-contiguous) order, i.e. for 2D ``[H, W]`` the sequence visits tokens + as (0,0), (0,1), …, (0,W-1), (1,0), … (raster-scan). For 3D ``[D,H,W]`` + the depth axis varies slowest. This ordering is fixed (not learned). - If bidirectionality is enabled, an additional Mamba layer is instantiated and used to process the reversed input. - The output of this layer is reversed and added to the output of the core (forward) layer. + **Bidirectional mode**: when ``bidirectional=True`` a second core layer + processes the flattened sequence in reverse, and its (re-reversed) output + is summed with the forward output. This gives every position a full-sequence + receptive field in both causal directions, which is beneficial for + non-causal spatial tasks such as image or volume modelling. - The output is then reshaped back to [B, * spatial_dims, hidden_dim]. + Attributes: + bidirectional (bool): Whether to apply a second reversed Mamba pass. + core_layer (torch.nn.Module): The forward (or only) Mamba core. + Must accept input of shape ``[B, S, C]`` and return ``[B, S, C]``. + core_layer_rev (torch.nn.Module): The reverse Mamba core. Only + present when ``bidirectional=True``; accessing this attribute + when ``bidirectional=False`` raises :class:`AttributeError`. + + Example:: + + import torch + from nvsubquadratic.lazy_config import LazyConfig + from nvsubquadratic.modules.mamba_nd import Mamba + from mamba_ssm import Mamba as MambaCore + + mamba = Mamba( + mamba_layer_cfg=LazyConfig(MambaCore)(d_model=128, d_state=16, d_conv=4, expand=2), + bidirectional=True, + ) + + # 2D input: batch=2, spatial=(16, 16), channels=128 + x = torch.randn(2, 16, 16, 128) + y = mamba(x) # [2, 16, 16, 128] """ def __init__( @@ -25,11 +186,23 @@ def __init__( mamba_layer_cfg: LazyConfig, bidirectional: bool = False, ): - """Initialize the Mamba mixer layer. + """Initialise the Mamba-ND wrapper. Args: - mamba_layer_cfg: LazyConfig - LazyConfig for the Mamba layer. - bidirectional: bool - Whether to use a bidirectional Mamba layer. + mamba_layer_cfg: :class:`~nvsubquadratic.lazy_config.LazyConfig` + for the underlying 1D Mamba core. The target class must + accept a 3-D tensor of shape ``[B, S, C]`` (batch, sequence + length, channels) and return a tensor of the same shape. + Typical targets include ``mamba_ssm.Mamba`` and + ``mamba_ssm.Mamba2``. The config is instantiated once for + ``core_layer`` and, when ``bidirectional=True``, a second + independent instantiation is created for ``core_layer_rev`` + so that the two directions have separate parameters. + bidirectional: If ``True``, run a second Mamba core on the + reversed sequence and sum both outputs. This doubles + parameter count and compute but gives non-causal coverage + of the full sequence — strongly recommended for spatial + tasks (images, volumes). Defaults to ``False``. """ super().__init__() self.bidirectional = bidirectional @@ -39,14 +212,42 @@ def __init__( if self.bidirectional: self.core_layer_rev = instantiate(mamba_layer_cfg) - def forward(self, x): - """Forward pass of the Mamba mixer layer. + def forward(self, x: torch.Tensor) -> torch.Tensor: + r"""Apply the Mamba SSM to an ND input signal. + + The forward pass performs the following steps: + + 1. **Flatten** all spatial axes into one sequence dimension: + ``[B, *spatial, C]`` → ``[B, S, C]``, where ``S = prod(spatial)``. + The flattening follows row-major (C-contiguous) order. + 2. **Forward SSM**: ``out = core_layer(x)`` — applies the selective + SSM recurrence :math:`y_t = C_t(\bar{A}_t h_{t-1} + \bar{B}_t x_t)`. + 3. **Reverse SSM** *(only when* ``bidirectional=True``): + ``out_rev = core_layer_rev(flip(x))`` — runs the SSM on the + reversed sequence, then flips back and adds to ``out``: + + .. math:: + + \text{out} \mathrel{+}= \text{flip}(\text{Mamba}_\text{rev}(\text{flip}(x))) + + 4. **Reshape** back to the original spatial layout: + ``[B, S, C]`` → ``[B, *spatial, C]``. Args: - x (torch.Tensor): Input tensor of shape (batch_size, * spatial_dims, hidden_dim) + x: Input tensor of shape ``(B, *spatial, C)`` where ``B`` is + batch size, ``spatial`` is one or more spatial dimensions + (e.g. ``(T,)`` for 1D sequences, ``(H, W)`` for 2D images, + ``(D, H, W)`` for 3D volumes), and ``C`` is the channel + (hidden) dimension. The tensor must be in channels-last + (BHC / BHWc) layout, consistent with the rest of the library. Returns: - torch.Tensor: Output tensor of shape (batch_size, * spatial_dims, hidden_dim) + Output tensor of shape ``(B, *spatial, C)`` — same shape and + layout as the input. When ``bidirectional=True`` the output is + the element-wise sum of the forward and reverse SSM outputs, + which doubles the effective output magnitude compared to a + unidirectional pass; downstream normalisation layers (e.g. + ``RMSNorm`` inside the residual block) absorb this scale. """ x_shape = x.shape # Reshape input to [B, flatten (* spatial_dims), hidden_dim From 32a153b95717d4635e4869a1cbb25b2105e5d73b Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 22:24:07 +0200 Subject: [PATCH 37/72] docs(review/condition_mixer): reviewer feedback --- docs/reviews/condition_mixer_review.md | 83 +++++++++++++++++++++++ docs/reviews/vit5_hyena_adapter_review.md | 67 ++++++++++++++++++ 2 files changed, 150 insertions(+) create mode 100644 docs/reviews/condition_mixer_review.md create mode 100644 docs/reviews/vit5_hyena_adapter_review.md diff --git a/docs/reviews/condition_mixer_review.md b/docs/reviews/condition_mixer_review.md new file mode 100644 index 00000000..fd5bbe97 --- /dev/null +++ b/docs/reviews/condition_mixer_review.md @@ -0,0 +1,83 @@ +# Review: `nvsubquadratic/modules/condition_mixer.py` + +Reviewed against the style references in `film.py`, `residual_block.py`, and `sequence_mixer.py`. + +______________________________________________________________________ + +## Issues + +### 1. `init_method_in` type annotation is wrong + +In `__init__`, the parameter is declared as: + +```python +init_method_in: Callable[[torch.Tensor], torch.Tensor] | None = (None,) +``` + +but the actual call-site is: + +```python +init_method_in(hidden_dim)(self.kv_proj.weight.data) +``` + +This is a **curried** function `fn(dim: int) -> fn(tensor: Tensor) -> None`, exactly the same as in `QKVSequenceMixer`. The annotation should be: + +```python +init_method_in: Callable[[int], Callable[[torch.Tensor], None]] | None = (None,) +``` + +Ditto for `init_method_out`. A new collaborator will write the wrong function and get a runtime error with no diagnostic help. + +______________________________________________________________________ + +### 2. Module docstring does not explain the relationship between `condition` channel dim and `hidden_dim` + +The module docstring says "the channel dimension `C` must equal `hidden_dim`" but never explains *why* — the `kv_proj` and `q_proj` are `Linear(hidden_dim, ...)` and operate directly on `condition`, so passing a conditioning tensor whose last dimension is not `hidden_dim` will silently produce wrong-shaped K/V tensors (no shape check is performed). This constraint should be stated as a hard requirement with the consequence of violating it. + +______________________________________________________________________ + +### 3. `forward` docstring does not mention what the inner `mixer` must return + +The docstring for `forward` describes the shapes of `q`, `k`, `v` but does not state that the inner `mixer(q, k, v)` must return a tensor with the **same spatial layout and channel dimension as `q`**. Because `mixer` is a black-box `LazyConfig`-instantiated module, a collaborator implementing a custom inner mixer would not know this contract without reading the code. + +Add a sentence such as: "The inner `mixer` must return a tensor of shape `(B, *spatial_dims, C)` matching `q`; the output projection `out_proj` will fail with a shape error otherwise." + +______________________________________________________________________ + +### 4. Missing interaction note: `condition_mixer` vs `AdaLNZeroResidualBlock` + +`residual_block.py` has two block types: `ResidualBlock` (which *has* a `condition_mixer` branch) and `AdaLNZeroResidualBlock` (which *does not*). A new collaborator reading `condition_mixer.py` has no pointer to this split. The module docstring should note that `QKVConditionMixer` is only used with `ResidualBlock`, not with `AdaLNZeroResidualBlock` (which routes all conditioning through the DiT AdaLN-Zero projection). + +______________________________________________________________________ + +### 5. The `condition_mixer_norm` interaction is not described + +`ResidualBlock.forward` applies `condition_mixer_norm` to `x` **before** calling `self.condition_mixer(x, condition)`. The condition mixer itself therefore receives a *normalised* feature map, not the raw residual. Nothing in `condition_mixer.py` documents this; a reader of only this file would assume `x` is the unnormalised residual stream. At minimum, add a note in the `forward` docstring: "In practice `x` has already been passed through `condition_mixer_norm` by the enclosing `ResidualBlock` before this module is called." + +______________________________________________________________________ + +### 6. `ValueError` message for `condition.ndim` mismatch quotes wrong expected value + +The error string reads: + +```python +f"Got condition.ndim={condition.ndim}, expected {x.ndim}." +``` + +The word "expected" is ambiguous — it sounds like *only* `x.ndim` is valid, but `2` is also valid (and is handled above this branch). Fix to: + +```python +f"Got condition.ndim={condition.ndim}, expected 2 (global vector) or {x.ndim} (matching spatial rank)." +``` + +______________________________________________________________________ + +### 7. Class docstring `Attributes` block omits the `hidden_dim` value + +Unlike `film.py` (which stores `num_film_layers` and `kernel_hidden_dim` as instance attributes), `QKVConditionMixer` does not persist `hidden_dim`. That is fine, but the `Attributes` block should make clear that `hidden_dim` can be recovered from `self.q_proj.in_features` if needed. This is a minor discoverability issue. + +______________________________________________________________________ + +### 8. No `See Also` cross-reference to `QKVSequenceMixer` + +The module docstring compares the condition mixer to FiLM and cross-attention, but does not point to `QKVSequenceMixer` in `sequence_mixer.py`, even though the two modules are structurally parallel (both share the `QKV + inner mixer + out_proj` skeleton). Add a `See Also` block to the class docstring pointing to `QKVSequenceMixer` and `ResidualBlock`. diff --git a/docs/reviews/vit5_hyena_adapter_review.md b/docs/reviews/vit5_hyena_adapter_review.md new file mode 100644 index 00000000..b3e29aa4 --- /dev/null +++ b/docs/reviews/vit5_hyena_adapter_review.md @@ -0,0 +1,67 @@ +# Review: `nvsubquadratic/modules/vit5_hyena_adapter.py` + +Reviewed after Phase 1 docstring pass. Issues are numbered and actionable. + +______________________________________________________________________ + +## 1. Module docstring: "drop-in replacement" claim needs qualification + +The second line reads: + +> "Drop-in replacement interface for `ViT5Attention`." + +This is misleading. `ViT5Attention` owns its own QKV and output projections; the adapter delegates all projections to `inner_mixer`. A reader replacing `ViT5Attention` with `ViT5HyenaAdapter` in config needs to know that `inner_mixer` (e.g. `QKVSequenceMixer`) must be configured with matching `hidden_dim` and the correct `num_heads`/`inner_dim` — the adapter itself accepts no `hidden_dim` argument. The docstring should add a sentence clarifying that the projection parameters must be configured **inside `inner_mixer_cfg`**, not on the adapter. + +______________________________________________________________________ + +## 2. Module docstring: `inner_mixer` contract is underspecified for the 2-D output reshape + +Step 4 of the "thin, stateless reshape adapter" list says: + +> "4. Reshapes back to `[B, T, C]` and returns." + +This silently assumes the inner mixer returns **exactly** `[B, H, W, C]` — the same spatial shape it received. If a mixer returns a different shape (e.g. a downsampling mixer), `reshape(B, T, C)` will raise a cryptic error. The docstring should state explicitly: "The inner mixer is assumed to be shape-preserving — its output must have the same `[B, H, W, C]` shape as its input." + +______________________________________________________________________ + +## 3. Class docstring: data-flow diagram uses `inner_mixer (QKVSequenceMixer → Hyena)` as a fixed example + +The class-level data-flow diagram labels the inner_mixer step as: + +``` +▼ inner_mixer (QKVSequenceMixer → Hyena) +``` + +This couples the class docstring to a specific implementation. The class is general — it wraps any `[B, H, W, C]`-preserving module. The parenthetical should be replaced with a more general description, e.g. `(any [B, H, W, C]-preserving mixer)`. + +______________________________________________________________________ + +## 4. `__init__` docstring: `inner_mixer_cfg` contract does not mention channels-last convention + +The docstring says the instantiated module must "accept `(x: Tensor[B, H, W, C], **kwargs)`" but does not mention that this is a **channels-last** layout. Hyena internally converts to channels-first `[B, C, H, W]` for convolution. A reader wiring a custom mixer needs to know the expected layout convention. Add: "The tensor is in channels-last layout `[B, H, W, C]`; any inner mixer that uses channels-first convolution (like `QKVSequenceMixer`) handles the permutation internally." + +______________________________________________________________________ + +## 5. `forward` docstring: second `reshape` output contiguity caveat is incomplete + +The Returns section states: "The reshape is a view (no data copy) when the tensor is contiguous." However, the inner mixer may return a non-contiguous tensor (e.g. after a transpose). In that case `reshape` falls back to a copy. This is fine correctness-wise but users writing CUDA-graph-safe code or tracing with `torch.compile` should know. Add: "If `inner_mixer` returns a non-contiguous tensor, the final `reshape` may trigger a contiguous copy; this does not affect correctness but can affect memory traffic." + +______________________________________________________________________ + +## 6. `flop_count` docstring: silent `AttributeError` is wrong error type in practice + +The Raises section says `AttributeError` if `inner_mixer` does not implement `flop_count`. In practice, if `inner_mixer` has no `flop_count`, Python raises `AttributeError` from the attribute lookup, not from a guarded check. This is correct but the docstring should also note that `flop_count` is a **de-facto protocol** not enforced by an interface — callers that want to guard against this should use `hasattr(adapter.inner_mixer, "flop_count")`. + +______________________________________________________________________ + +## 7. Missing: note about `grid_w` and the token-layout contract for the hierarchical case + +The module docstring mentions `ViT5ClassificationNet` as the network that pads and arranges tokens, but the current codebase also has `vit5_hierarchical_classification.py` (feat/patch-merging PR). The register-token handling note in the class docstring only mentions a single "register row" convention but does not mention that after patch merging the grid dimensions change. Add a brief note that `grid_w` must be consistent with the spatial width **after any patch-merging stage**, i.e. the calling network must supply the correct `grid_w` at each hierarchical stage. + +______________________________________________________________________ + +## 8. `extra_repr` docstring: wording is slightly off + +Current: "Return a concise summary for `repr()` and `print(model)`." + +`extra_repr` is called by PyTorch's `__repr__` machinery. The docstring should say what information it returns rather than how it is called: "Return `grid_w=` appended to PyTorch's default module repr." This matches the style in `vit5_attention.py`'s `extra_repr`. From c109f4b1ccf2ffd13f3f36c24be5f811d0eacae2 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 22:25:20 +0200 Subject: [PATCH 38/72] docs(integrate/condition_mixer): apply reviewer feedback --- docs/reviews/condition_mixer_review.md | 83 --------- docs/reviews/mamba_nd_review.md | 209 ++++++++++++++++++++++ nvsubquadratic/modules/condition_mixer.py | 103 ++++++++--- 3 files changed, 285 insertions(+), 110 deletions(-) delete mode 100644 docs/reviews/condition_mixer_review.md create mode 100644 docs/reviews/mamba_nd_review.md diff --git a/docs/reviews/condition_mixer_review.md b/docs/reviews/condition_mixer_review.md deleted file mode 100644 index fd5bbe97..00000000 --- a/docs/reviews/condition_mixer_review.md +++ /dev/null @@ -1,83 +0,0 @@ -# Review: `nvsubquadratic/modules/condition_mixer.py` - -Reviewed against the style references in `film.py`, `residual_block.py`, and `sequence_mixer.py`. - -______________________________________________________________________ - -## Issues - -### 1. `init_method_in` type annotation is wrong - -In `__init__`, the parameter is declared as: - -```python -init_method_in: Callable[[torch.Tensor], torch.Tensor] | None = (None,) -``` - -but the actual call-site is: - -```python -init_method_in(hidden_dim)(self.kv_proj.weight.data) -``` - -This is a **curried** function `fn(dim: int) -> fn(tensor: Tensor) -> None`, exactly the same as in `QKVSequenceMixer`. The annotation should be: - -```python -init_method_in: Callable[[int], Callable[[torch.Tensor], None]] | None = (None,) -``` - -Ditto for `init_method_out`. A new collaborator will write the wrong function and get a runtime error with no diagnostic help. - -______________________________________________________________________ - -### 2. Module docstring does not explain the relationship between `condition` channel dim and `hidden_dim` - -The module docstring says "the channel dimension `C` must equal `hidden_dim`" but never explains *why* — the `kv_proj` and `q_proj` are `Linear(hidden_dim, ...)` and operate directly on `condition`, so passing a conditioning tensor whose last dimension is not `hidden_dim` will silently produce wrong-shaped K/V tensors (no shape check is performed). This constraint should be stated as a hard requirement with the consequence of violating it. - -______________________________________________________________________ - -### 3. `forward` docstring does not mention what the inner `mixer` must return - -The docstring for `forward` describes the shapes of `q`, `k`, `v` but does not state that the inner `mixer(q, k, v)` must return a tensor with the **same spatial layout and channel dimension as `q`**. Because `mixer` is a black-box `LazyConfig`-instantiated module, a collaborator implementing a custom inner mixer would not know this contract without reading the code. - -Add a sentence such as: "The inner `mixer` must return a tensor of shape `(B, *spatial_dims, C)` matching `q`; the output projection `out_proj` will fail with a shape error otherwise." - -______________________________________________________________________ - -### 4. Missing interaction note: `condition_mixer` vs `AdaLNZeroResidualBlock` - -`residual_block.py` has two block types: `ResidualBlock` (which *has* a `condition_mixer` branch) and `AdaLNZeroResidualBlock` (which *does not*). A new collaborator reading `condition_mixer.py` has no pointer to this split. The module docstring should note that `QKVConditionMixer` is only used with `ResidualBlock`, not with `AdaLNZeroResidualBlock` (which routes all conditioning through the DiT AdaLN-Zero projection). - -______________________________________________________________________ - -### 5. The `condition_mixer_norm` interaction is not described - -`ResidualBlock.forward` applies `condition_mixer_norm` to `x` **before** calling `self.condition_mixer(x, condition)`. The condition mixer itself therefore receives a *normalised* feature map, not the raw residual. Nothing in `condition_mixer.py` documents this; a reader of only this file would assume `x` is the unnormalised residual stream. At minimum, add a note in the `forward` docstring: "In practice `x` has already been passed through `condition_mixer_norm` by the enclosing `ResidualBlock` before this module is called." - -______________________________________________________________________ - -### 6. `ValueError` message for `condition.ndim` mismatch quotes wrong expected value - -The error string reads: - -```python -f"Got condition.ndim={condition.ndim}, expected {x.ndim}." -``` - -The word "expected" is ambiguous — it sounds like *only* `x.ndim` is valid, but `2` is also valid (and is handled above this branch). Fix to: - -```python -f"Got condition.ndim={condition.ndim}, expected 2 (global vector) or {x.ndim} (matching spatial rank)." -``` - -______________________________________________________________________ - -### 7. Class docstring `Attributes` block omits the `hidden_dim` value - -Unlike `film.py` (which stores `num_film_layers` and `kernel_hidden_dim` as instance attributes), `QKVConditionMixer` does not persist `hidden_dim`. That is fine, but the `Attributes` block should make clear that `hidden_dim` can be recovered from `self.q_proj.in_features` if needed. This is a minor discoverability issue. - -______________________________________________________________________ - -### 8. No `See Also` cross-reference to `QKVSequenceMixer` - -The module docstring compares the condition mixer to FiLM and cross-attention, but does not point to `QKVSequenceMixer` in `sequence_mixer.py`, even though the two modules are structurally parallel (both share the `QKV + inner mixer + out_proj` skeleton). Add a `See Also` block to the class docstring pointing to `QKVSequenceMixer` and `ResidualBlock`. diff --git a/docs/reviews/mamba_nd_review.md b/docs/reviews/mamba_nd_review.md new file mode 100644 index 00000000..d20c95b8 --- /dev/null +++ b/docs/reviews/mamba_nd_review.md @@ -0,0 +1,209 @@ +# Review: `nvsubquadratic/modules/mamba_nd.py` + +Reviewed after Phase 1 docstring pass. Issues are ordered from highest to +lowest priority. Each item quotes the relevant text and states the exact fix. + +______________________________________________________________________ + +## 1. ZOH approximation presented as equality + +**Location**: module-level docstring, "Background" section. + +**Quoted text**: + +``` +\bar{B}_t = (e^{\Delta_t A} - I) A^{-1} B_t \approx \Delta_t B_t +``` + +**Issue**: The full ZOH expression is exact, but the approximation +`≈ Δ_t B_t` is only valid for small `Δ_t`. More precisely, the Mamba paper +(Eq. 4 in arXiv:2312.00752) uses the Euler discretisation `B̄_t = Δ_t B_t` +and notes that the ZOH formula is an alternative; the code in `mamba_ssm` +actually uses the Euler rule for `B`, not the full ZOH. Presenting the full +ZOH formula and then approximating it may confuse readers into thinking the +implementation uses ZOH for both `A` and `B`. + +**Fix**: Clarify that `Ā_t = exp(Δ_t A)` (ZOH) but `B̄_t = Δ_t B_t` (Euler +/ first-order), matching what `mamba_ssm` actually computes. Remove the +intermediate exact-then-approximate chain or clearly label it "alternative +ZOH formula (not used in practice)". + +______________________________________________________________________ + +## 2. Comparison table: Mamba complexity claim is misleading for training + +**Location**: module-level docstring, "Comparison with other mixers" section. + +**Quoted text**: + +``` +| Mamba | SSM recurrence | input-dependent | O(N) | +``` + +**Issue**: O(N) is the *inference* (autoregressive) complexity via the +recurrent form. During training Mamba uses a parallel associative scan whose +GPU-efficient implementation is O(N log N) in time, or O(N) with the +hardware-aware parallel scan (which requires special CUDA kernels — the whole +point of the `mamba_ssm` package). Listing O(N) without qualification makes +it look strictly cheaper than Hyena at training time, which is only true with +the custom CUDA kernels. + +**Fix**: Add a note distinguishing training (parallel scan, O(N) with custom +kernels or O(N log N) naively) from inference (recurrent, O(N) per step, O(1) +state size). E.g. add a "Notes" column or a footnote: "O(N) with hardware- +aware parallel scan (requires `mamba_ssm` CUDA extension); O(1) per step at +inference". + +______________________________________________________________________ + +## 3. Scan order: no mention of the implication for 2D spatial locality + +**Location**: module-level docstring, "ND generalisation strategy" section, and +class-level docstring paragraph starting "**Scan order for ND inputs**". + +**Quoted text**: + +``` +The scan order for multi-dimensional inputs follows the default PyTorch / +``einops`` row-major (C-contiguous) flattening: for a 2D ``[H, W]`` input the +tokens are visited in raster-scan order (row 0, col 0 → row 0, col W-1 → +row 1, col 0 → …). +``` + +**Issue**: This is accurate but omits the key practical consequence: tokens +that are spatially adjacent *vertically* (same column, adjacent rows) are far +apart in the flattened sequence (W steps away), so the SSM's effective +receptive field is anisotropic — it sees horizontal neighbours cheaply but +vertical neighbours only through W state-update steps. An external +collaborator reading this to decide whether to use Mamba for a 2D task needs +this information. + +**Fix**: Add one sentence explicitly warning about vertical anisotropy, e.g.: +"Note that vertically adjacent pixels (same column, adjacent rows) are W +tokens apart in the flattened sequence; for tall images this means the forward +SSM sees them only through many state transitions, potentially losing spatial +correlation. Bidirectional mode partially mitigates this." + +______________________________________________________________________ + +## 4. `forward` docstring: `x` is mutated (overwritten by `rearrange`) + +**Location**: `Mamba.forward`, Args section and implementation. + +**Quoted text** (implementation): + +```python +x = rearrange(x, "b ... c -> b (...) c") +``` + +**Issue**: The `rearrange` result is assigned back to `x`, which shadows the +original argument. The original spatial shape is captured in `x_shape` first, +so this is safe, but the docstring does not warn that the local variable `x` +changes meaning mid-function (it is the flattened view from line 254 onward). +This is a mild readability issue but analogous to the note in `hyena_nd.py`'s +`forward` docstring about `query` being overwritten — be consistent. + +**Fix**: Add an "Implementation note" paragraph to `forward`: +"The local variable `x` is rebound to the flattened `[B, S, C]` view after +the `rearrange` call; the original spatial shape is preserved in `x_shape` +for the final `reshape`." + +______________________________________________________________________ + +## 5. `__init__` docstring: no mention that `core_layer_rev` has independent parameters + +**Location**: `Mamba.__init__`, Args section for `mamba_layer_cfg`. + +**Quoted text**: + +``` +The config is instantiated once for ``core_layer`` and, when +``bidirectional=True``, a second independent instantiation is created for +``core_layer_rev`` so that the two directions have separate parameters. +``` + +**Issue**: The wording says "independent instantiation … separate parameters" +but does not say *how* independence is achieved — a reader unfamiliar with +`LazyConfig` might wonder if the second call shares weights via some internal +cache. + +**Fix**: Add a clarifying phrase: "`instantiate(mamba_layer_cfg)` is called +twice with the same config; each call constructs a fresh `nn.Module` with +newly initialised weights, so the two directions do not share parameters." + +______________________________________________________________________ + +## 6. Missing `Raises` section in `__init__` + +**Location**: `Mamba.__init__` docstring. + +**Issue**: If `mamba_layer_cfg` cannot be instantiated (wrong target class, +missing required arguments), `instantiate` will raise — most likely a +`RuntimeError`, `TypeError`, or an `omegaconf` exception, depending on the +`LazyConfig` backend. The `QKVSequenceMixer.__init__` docstring in +`sequence_mixer.py` includes a `Raises` section for this; `Mamba.__init__` +should match. + +**Fix**: Add: + +``` +Raises: + Exception: Propagated from + :func:`~nvsubquadratic.lazy_config.instantiate` if + ``mamba_layer_cfg`` cannot be constructed. Check that the target + class accepts ``[B, S, C]`` tensors and that all required constructor + arguments are provided in the config. +``` + +______________________________________________________________________ + +## 7. Class-level Attributes block: `core_layer_rev` uses vague conditional phrasing + +**Location**: `Mamba` class, Attributes block. + +**Quoted text**: + +``` +core_layer_rev (torch.nn.Module): The reverse Mamba core. Only + present when ``bidirectional=True``; accessing this attribute + when ``bidirectional=False`` raises :class:`AttributeError`. +``` + +**Issue**: Saying "raises `AttributeError`" is accurate but alarming without +context — it looks like a bug. It would be clearer to note the attribute is +intentionally absent (not registered) to keep `state_dict` / `parameters()` +clean when unused. + +**Fix**: Reword to: "The reverse Mamba core, instantiated only when +`bidirectional=True`. When `bidirectional=False` this attribute is not +registered and accessing it raises :class:`AttributeError` by design, keeping +the module's parameter count and `state_dict` unaffected." + +______________________________________________________________________ + +## 8. Module docstring: `References` section uses non-standard heading style + +**Location**: module-level docstring, last section. + +**Quoted text** (after ruff fix): + +``` +References: +---------- +``` + +**Issue**: Sphinx / NumPy / Google doc conventions all write `References` with +the underline immediately under the heading, not with a blank line. The ruff +linter already fixed a trailing colon issue here but the section may still +render oddly in Sphinx because of the mixed `References:` + underline style. +The rest of the module docstring uses RST underlines without trailing colons +on the heading (e.g. `Background\n----------`). + +**Fix**: Change to match the rest of the file: + +``` +References +---------- +``` + +(remove the trailing colon from `References:`). diff --git a/nvsubquadratic/modules/condition_mixer.py b/nvsubquadratic/modules/condition_mixer.py index 81ed4124..6424cc41 100644 --- a/nvsubquadratic/modules/condition_mixer.py +++ b/nvsubquadratic/modules/condition_mixer.py @@ -24,13 +24,24 @@ sequence mixer has let spatial positions exchange information, the condition mixer injects an external conditioning signal ``c`` — such as a diffusion timestep embedding, a class label embedding, or a physics parameter vector — -into the residual stream ``x``: +into the residual stream ``x``. In practice ``x`` is first passed through +``condition_mixer_norm`` by the enclosing +:class:`~nvsubquadratic.modules.residual_block.ResidualBlock` before this +module receives it; the condition mixer itself therefore operates on a +*normalised* feature map: .. code-block:: text - x ──[seq_mixer]──► x' - x'──[cond_mixer(x', c)]──► x'' ← this module - x''─[mlp]──────────────► output + x ──[seq_mixer]──────────────────────► x' + x'──[condition_mixer_norm]──► x'_norm + x'_norm──[cond_mixer(x'_norm, c)]──► x'' ← this module + x''─[mlp]───────────────────────────► output + +This module is **only** used with +:class:`~nvsubquadratic.modules.residual_block.ResidualBlock`. The +alternative :class:`~nvsubquadratic.modules.residual_block.AdaLNZeroResidualBlock` +has no condition-mixer branch at all; it routes all conditioning through its +zero-initialised AdaLN-Zero projection. **Comparison with other conditioning strategies** @@ -64,7 +75,11 @@ (e.g. encoder output in an encoder-decoder architecture). Must have the same number of dimensions as the feature map ``x``. -In both cases the channel dimension ``C`` must equal ``hidden_dim``. +In both cases the channel dimension ``C`` **must** equal ``hidden_dim``. +``kv_proj`` and ``q_proj`` are ``Linear(hidden_dim, ...)`` and operate +directly on ``condition``; passing a tensor whose last dimension differs +from ``hidden_dim`` will silently produce a wrong-shaped K/V and corrupt the +forward pass — no shape check is performed. All tensors use **channels-last** layout: ``(B, *spatial, C)``. """ @@ -102,7 +117,18 @@ class QKVConditionMixer(torch.nn.Module): same operator-agnostic dispatch pattern as :class:`~nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer`. Any module whose ``forward(q, k, v)`` conforms to channels-last tensors of - shape ``(B, *spatial, C)`` can be used as the inner mixer. + shape ``(B, *spatial, C)`` and returns a tensor of the same shape as ``q`` + can be used as the inner mixer. If the inner mixer returns a different + shape, ``out_proj`` will raise a shape error. + + .. note:: + + In practice ``x`` arriving at :meth:`forward` has already been passed + through ``condition_mixer_norm`` by the enclosing + :class:`~nvsubquadratic.modules.residual_block.ResidualBlock`. This + module is **not** used with + :class:`~nvsubquadratic.modules.residual_block.AdaLNZeroResidualBlock`, + which has no condition-mixer branch. **Weight initialisation** @@ -118,14 +144,23 @@ class QKVConditionMixer(torch.nn.Module): for any per-parameter weight-decay grouping (see the analogous logic in :class:`~nvsubquadratic.modules.film.KernelFiLMGenerator`). + See Also: + :class:`~nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer`: + Structurally parallel module for self-attention / sequence mixing + (``forward(x)`` rather than ``forward(x, condition)``). + :class:`~nvsubquadratic.modules.residual_block.ResidualBlock`: + The enclosing block that calls this module's ``forward`` after + applying ``condition_mixer_norm`` to the residual stream. + Attributes: mixer (torch.nn.Module): The instantiated inner mixing operator. Its - ``forward(q, k, v)`` method receives channels-last tensors of shape - ``(B, *spatial, C)`` and must return a tensor of the same shape as + ``forward(q, k, v)`` method receives channels-last tensors and must + return a tensor of the same spatial shape and channel dimension as ``q``. kv_proj (torch.nn.Linear): Combined K+V projection (no bias) that maps the conditioning signal from ``C`` to ``2·C``. Weight shape: - ``(2·hidden_dim, hidden_dim)``. + ``(2·hidden_dim, hidden_dim)``. The input channel dimension + ``hidden_dim`` can be recovered via ``self.kv_proj.in_features``. q_proj (torch.nn.Linear): Query projection (no bias) that maps the feature map from ``C`` to ``C``. Weight shape: ``(hidden_dim, hidden_dim)``. @@ -138,16 +173,19 @@ def __init__( self, hidden_dim: int, mixer_cfg: LazyConfig, - init_method_in: Callable[[torch.Tensor], torch.Tensor] | None = None, - init_method_out: Callable[[torch.Tensor], torch.Tensor] | None = None, + init_method_in: Callable[[int], Callable[[torch.Tensor], None]] | None = None, + init_method_out: Callable[[int], Callable[[torch.Tensor], None]] | None = None, ): """Initialise the QKVConditionMixer. Args: hidden_dim: Channel dimension ``C`` shared by the feature map ``x`` - and the conditioning signal ``c``. All four linear projections + and the conditioning signal ``c``. All linear projections (``q_proj``, ``kv_proj``, ``out_proj``) are sized using this - value. + value. The conditioning tensor ``c`` must have its last + dimension equal to ``hidden_dim``; mismatches are not checked + at init time and will silently produce wrong-shaped K/V tensors + during the forward pass. mixer_cfg: :class:`~nvsubquadratic.lazy_config.LazyConfig` for the inner mixing operator. The instantiated module's ``forward`` must accept ``(q, k, v)`` as positional arguments and return a @@ -155,14 +193,15 @@ def __init__( module (e.g. a dot-product attention layer) can be used here. init_method_in: Optional *curried* weight initialiser applied to both ``q_proj.weight`` and ``kv_proj.weight``. Must have the - signature ``fn(dim: int) -> fn(tensor: Tensor) -> None``. - When provided, ``fn(hidden_dim)`` is called and the returned - callable is applied in-place to each weight matrix. Pass - ``None`` to keep PyTorch's default Kaiming-uniform init. - init_method_out: Same as ``init_method_in`` but applied to - ``out_proj.weight``. A common choice is a depth-scaled - Gaussian (GPT / Megatron style) to control the residual branch - variance at initialisation. Pass ``None`` to keep the default. + signature ``fn(dim: int) -> fn(tensor: Tensor) -> None`` — + ``fn(hidden_dim)`` is called and the returned callable is + applied in-place to each weight matrix. Pass ``None`` to keep + PyTorch's default Kaiming-uniform init. + init_method_out: Same curried signature as ``init_method_in``, + applied to ``out_proj.weight``. A common choice is a + depth-scaled Gaussian (GPT / Megatron style) to control the + residual branch variance at initialisation. Pass ``None`` to + keep the default. """ super().__init__() @@ -185,9 +224,15 @@ def __init__( def forward(self, x: torch.Tensor, condition: torch.Tensor) -> torch.Tensor: """Inject the conditioning signal into the feature map via cross-attention. - Computes queries from the current feature map ``x`` and keys/values - from the conditioning signal ``condition``, then mixes them with the - inner ``mixer`` and projects the result back to ``hidden_dim``. + Computes queries from the (pre-normalised) feature map ``x`` and + keys/values from the conditioning signal ``condition``, then mixes them + with the inner ``mixer`` and projects the result back to ``hidden_dim``. + + .. note:: + + In the standard :class:`~nvsubquadratic.modules.residual_block.ResidualBlock` + usage, ``x`` has already been passed through ``condition_mixer_norm`` + before this method is called. The signal flow is: @@ -197,6 +242,10 @@ def forward(self, x: torch.Tensor, condition: torch.Tensor) -> torch.Tensor: K, V = split(kv_proj(condition)) # each (B, *spatial_dims_cond, C) y = out_proj(mixer(Q, K, V)) # (B, *spatial_dims, C) + The inner ``mixer`` must return a tensor of the same shape as ``Q`` + (i.e. ``(B, *spatial_dims, C)``); if it does not, ``out_proj`` will + raise a shape error. + A global (non-spatial) conditioning vector of shape ``(B, C)`` is automatically unsqueezed to ``(B, 1, C)`` before the K/V projection so that the inner mixer sees a single conditioning token per sample. @@ -218,8 +267,8 @@ class embedding). Unsqueezed internally to ``(B, 1, C)`` ``spatial_dims_cond`` need not match ``spatial_dims`` of ``x``. - The channel dimension ``C`` must equal ``hidden_dim`` in both - cases. + The channel dimension ``C`` **must** equal ``hidden_dim`` + (``self.kv_proj.in_features``); this is not checked at runtime. Returns: Output tensor of shape ``(B, *spatial_dims, C)`` — same spatial @@ -242,7 +291,7 @@ class embedding). Unsqueezed internally to ``(B, 1, C)`` elif condition.ndim != x.ndim: raise ValueError( f"Condition must have either 2 dimensions (global) or match x's spatial rank. " - f"Got condition.ndim={condition.ndim}, expected {x.ndim}." + f"Got condition.ndim={condition.ndim}, expected 2 (global vector) or {x.ndim} (matching spatial rank)." ) # Q projection from the current stream From f2c3b1f27fb3b35ac0b7bf9401a752ab53d69c20 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 22:26:07 +0200 Subject: [PATCH 39/72] docs(write/vit5_attention): add module and class docstrings --- nvsubquadratic/modules/vit5_attention.py | 410 +++++++++++++++++++---- 1 file changed, 354 insertions(+), 56 deletions(-) diff --git a/nvsubquadratic/modules/vit5_attention.py b/nvsubquadratic/modules/vit5_attention.py index b6346eb5..20605835 100644 --- a/nvsubquadratic/modules/vit5_attention.py +++ b/nvsubquadratic/modules/vit5_attention.py @@ -1,16 +1,66 @@ """ViT-5 Attention: Multi-head self-attention with RMSNorm QK-Norm and register-aware 2D RoPE. -Key differences from the base Attention module: -- QK-Norm uses RMSNorm (learnable, per-head) instead of L2 normalization. -- RoPE is applied only to patch tokens (not cls token). -- Register tokens get their own high-frequency RoPE. -- Operates on flattened sequences [B, T, C] where T = N (patches) + 1 (cls) + R (registers). - -Optimizations vs naive implementation: -- RoPE cos/sin are precomputed as registered buffers (CUDA-graph safe, no graph breaks). -- RoPE applied via a single broadcast multiply on [B, T, H, D] — no reshape to (B*H, T, D). -- SDPA backend auto-selected by PyTorch (CuDNN preferred on H100). -- No redundant dtype casts around SDPA (autocast handles precision). +This module implements the specialised self-attention block used throughout the +ViT-5 family of hierarchical vision transformers. It is a *self-contained* +alternative to the generic :class:`~nvsubquadratic.modules.attention.Attention` +module and is **not** interchangeable with it through the +:class:`~nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer` dispatch layer. +The key ViT-5-specific design choices are described below. + +**Structural differences vs.** :class:`~nvsubquadratic.modules.attention.Attention` + +1. **Self-contained projections** — :class:`ViT5Attention` owns its own QKV + projection (``nn.Linear(C, 3C)``) and output projection (``nn.Linear(C, C)``), + plus an optional output dropout. The generic :class:`Attention` is a *pure* + attention kernel that receives pre-projected Q, K, V tensors from the + surrounding :class:`~nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer`. + :class:`ViT5Attention` is therefore consumed directly by + :class:`~nvsubquadratic.modules.vit5_residual_block.ViT5ResidualBlock` without + any outer projection wrapper. + +2. **RMSNorm QK-Norm (per-head, learnable)** — When ``qk_norm`` is provided, + :class:`ViT5Attention` instantiates two independent norm modules (one for Q, + one for K) via :func:`~nvsubquadratic.lazy_config.instantiate`. The generic + :class:`Attention` uses a *shared* L2 (cosine) normalisation function + (:func:`~nvsubquadratic.utils.qk_norm.apply_qk_norm`) without learnable + parameters. The ViT-5 norm is applied after the per-head reshape so that + each head's Q and K are normalised independently. + +3. **Register-aware, dual-base 2D RoPE** — Patch tokens receive 2D RoPE with + base frequency ``rope_base`` (default 10000). Register tokens receive + their own 2D RoPE with a *different* base ``reg_rope_base`` (default 100, + i.e., much higher frequency), reflecting their role as global context carriers + rather than spatially localised patch representations. The CLS token, when + present, receives an identity rotation (cos=1, sin=0). The generic + :class:`Attention` applies a single shared RoPE base to all tokens uniformly. + +4. **Fixed token layout** — The input sequence ``[B, T, C]`` must follow the + strict ViT-5 token layout ``[patches, (CLS,) registers]`` where + ``T = H*W + (1 if has_cls else 0) + R``. There is no support for arbitrary + spatial shapes or causal masking. The generic :class:`Attention` accepts + 1D/2D/3D channels-last tensors and supports both causal and non-causal modes. + +5. **CUDA-graph-safe precomputed RoPE buffers** — Both patch and register + cos/sin tables are concatenated into a single ``[T, head_dim]`` buffer pair + (``rope_cos``, ``rope_sin``) stored as non-persistent ``register_buffer`` + entries, making the forward pass free of dynamic tensor creation and safe + for ``torch.compile(mode="max-autotune")`` and CUDA-graph capture. + +6. **RoPE rotation convention** — :func:`_rotate_half_per_axis` uses the + *split-half* convention from ``nvsubquadratic.utils.rope.rotate_half_blh`` + rather than the interleaved rotation used in some external codebases. + Existing checkpoints trained with this convention will produce ~4 pp accuracy + drops if the rotation function is swapped. See :func:`_rotate_half_per_axis` + for the full explanation. + +**Cross-references** + +* :class:`~nvsubquadratic.modules.attention.Attention` — generic 1D/2D/3D + scaled dot-product attention, used via ``QKVSequenceMixer``. +* :class:`~nvsubquadratic.modules.vit5_residual_block.ViT5ResidualBlock` — + consumes :class:`ViT5Attention` directly as its ``sequence_mixer``. +* :func:`_build_2d_rope_flat` — constructs the flattened 2D RoPE tables. +* :func:`_rotate_half_per_axis` — split-half rotation operator used in ``forward``. """ from typing import Callable, Optional @@ -28,7 +78,7 @@ def _build_2d_rope_flat( head_dim: int, rope_base: float, ) -> tuple[torch.Tensor, torch.Tensor]: - """Precompute flattened 2D RoPE cos/sin for a (height x width) grid. + r"""Precompute flattened 2D RoPE cos/sin for a (height x width) grid. The original codebase (``nvsubquadratic.utils.rope``) computed RoPE on the fly during every forward pass. We precompute the cos/sin tables at init @@ -42,14 +92,45 @@ def _build_2d_rope_flat( tables follow the module to the correct device/dtype via ``.to()`` without manual bookkeeping. - Channel layout (matches ``_rotate_half_per_axis``): - [Y_half | X_half], each of size head_dim/2. - Within each half, frequencies are ``repeat_interleave(2)`` so that the - paired-swap rotation in ``_rotate_half_per_axis`` operates on matching - frequency pairs. + **Frequency schedule** + + For each spatial axis, frequencies follow the standard RoPE schedule: + + .. math:: + + \theta_j = \text{rope\_base}^{-2j / (d_k/2)}, + \quad j = 0, 1, \ldots, d_k/4 - 1 + + where :math:`d_k` = ``head_dim``. Angles for position :math:`p` are + :math:`\phi_{p,j} = p \cdot \theta_j`, then ``repeat_interleave(2)`` + doubles each frequency so the paired-swap rotation in + :func:`_rotate_half_per_axis` operates on matching frequency pairs. + + **Channel layout** (matches :func:`_rotate_half_per_axis`): + ``[Y_half | X_half]``, each of size ``head_dim / 2``. + Within each half, frequencies are ``repeat_interleave(2)`` so that the + paired-swap rotation in :func:`_rotate_half_per_axis` operates on matching + frequency pairs. + + Note: + ``head_dim`` must be divisible by 4 (two halves of ``head_dim/2`` each, + further split into pairs of size ``head_dim/4`` for the repeat-interleave + step). This constraint is enforced by the caller + (:class:`ViT5Attention.__init__`). + + Args: + height: Number of patch rows ``H`` in the 2D grid. + width: Number of patch columns ``W`` in the 2D grid. + head_dim: Per-head channel dimension ``d_k``. Must be divisible by 4. + rope_base: Base frequency for the geometric frequency schedule. + Typical values: ``10000.0`` for patch tokens, ``100.0`` for + register tokens (higher frequency = denser positional encoding). Returns: - (cos, sin) each of shape [height * width, head_dim]. + A tuple ``(cos, sin)`` where each tensor has shape + ``[height * width, head_dim]``. These are meant to be stored as + non-persistent ``register_buffer`` entries and concatenated with the + CLS and register entries in :class:`ViT5Attention.__init__`. """ dim_half = head_dim // 2 theta = 1.0 / (rope_base ** (torch.arange(0, dim_half, 2).float() / dim_half)) @@ -77,19 +158,45 @@ def _build_2d_rope_flat( def _rotate_half_per_axis(x: torch.Tensor) -> torch.Tensor: - """Split-half rotation applied independently to Y and X channel halves. + r"""Split-half rotation applied independently to Y and X channel halves. + + This implements the rotation component of the RoPE formula + ``x' = x * cos + rotate(x) * sin``. The rotation maps each channel vector + to its 90-degree rotated counterpart using a *split-half* convention: + + .. math:: + + \text{rotate}(x) = \text{Concat}\bigl( + -x_{y,2},\, x_{y,1},\, -x_{x,2},\, x_{x,1} + \bigr) + + where :math:`x_{y,1}` and :math:`x_{y,2}` are the first and second + quarters of the head-dim vector (the Y-axis half, each of size + ``head_dim / 4``), and similarly :math:`x_{x,1}`, :math:`x_{x,2}` for + the X-axis half. + + **Why not the standard interleaved** ``rotate_half``? - **Why not the standard interleaved ``rotate_half``?** The original training run used ``rotate_half_blh`` from ``nvsubquadratic.utils.rope``, which splits each axis-half at the midpoint - and swaps with negation ([-x2, x1]). The commonly-seen interleaved - rotation ([-x1, x0, -x3, x2, ...]) is numerically *incompatible* with + and swaps with negation (``[-x2, x1]``). The commonly-seen interleaved + rotation (``[-x1, x0, -x3, x2, ...]``) is numerically *incompatible* with checkpoints trained under the split-half convention — using it causes a ~4 pp accuracy drop at validation. This function preserves exact numerical parity with the original rotation so existing checkpoints remain valid. - Channel layout: [Y_half | X_half], each half of size D/2. - Within each half: split at D/4 and swap with negation. + **Channel layout**: ``[Y_half | X_half]``, each half of size ``D/2``. + Within each half: split at ``D/4`` and swap with negation. + + Args: + x: Query or key tensor of shape ``[B, T, H, D]`` (before the SDPA + transpose) or any shape whose last dimension is the head dimension + ``D``. ``D`` must be divisible by 4 to allow the quarter-split. + + Returns: + torch.Tensor: Rotated tensor of the same shape as ``x``, representing + the 90-degree rotation of each frequency pair within the Y and X + channel halves. """ d = x.shape[-1] d_half = d // 2 @@ -102,33 +209,172 @@ def _rotate_half_per_axis(x: torch.Tensor) -> torch.Tensor: class ViT5Attention(nn.Module): - """ViT-5 multi-head self-attention with RMSNorm QK-Norm and register-aware RoPE. - - Expects input as [B, T, C] where T = num_patches + (1 if has_cls) + num_registers. - Token layout: [patches, (CLS), registers] -- no padding (stripped by the network). - The module includes its own QKV and output projections (unlike the base Attention - which is wrapped in QKVSequenceMixer). + r"""ViT-5 multi-head self-attention with RMSNorm QK-Norm and register-aware RoPE. + + This module is the primary sequence-mixing operator for the ViT-5 family of + hierarchical vision transformers. It computes standard scaled dot-product + attention: + + .. math:: + + \text{head}_i = \text{softmax}\!\left( + \frac{Q_i K_i^\top}{\sqrt{d_k}} + \right) V_i, \quad + \text{out} = \text{Concat}(\text{head}_1, \ldots, \text{head}_H) W_O + + where :math:`d_k = C / H` is the per-head dimension and :math:`W_O` is the + output projection. Q, K, V are obtained from a single fused linear + projection :math:`[Q, K, V] = x W_{QKV}`. + + **Token layout** + + Input shape: ``[B, T, C]`` where + ``T = num_patches_h * num_patches_w + (1 if has_cls else 0) + num_registers``. + Token ordering within the sequence axis: + + .. code-block:: text + + [ patch_0, patch_1, ..., patch_{H*W-1}, (CLS,) reg_0, ..., reg_{R-1} ] + <----- H*W patch tokens ---------> <--1--> <---- R registers ----> + + This ordering must be consistent with the token layout produced by the + network's patchify + register-injection layers (see + :class:`~nvsubquadratic.networks.vit5_classification.ViT5Classifier`). + + **Positional encoding** + + Three distinct positional encodings are applied: + + * **Patch tokens** — 2D RoPE with base frequency ``rope_base`` (default + 10000). The H×W grid is linearised in row-major (Y-then-X) order. + * **CLS token** — identity rotation: cos=1, sin=0. No positional bias is + imposed on the class token. + * **Register tokens** — 2D RoPE with base ``reg_rope_base`` (default 100), + treating the ``R`` registers as a ``sqrt(R) × sqrt(R)`` grid. The higher + base frequency (lower theta) gives denser angular spacing, reflecting the + role of registers as global context carriers without fixed spatial meaning. + + All three tables are concatenated into a single buffer pair (``rope_cos``, + ``rope_sin``) of shape ``[T, head_dim]`` and applied with a single broadcast + multiply in :meth:`forward`. + + **QK normalisation** + + When ``qk_norm`` is provided, two independent norm modules (``q_norm``, + ``k_norm``) are instantiated and applied to Q and K *after* the per-head + reshape but *before* RoPE. The norm is expected to be a learnable RMSNorm + or equivalent. Unlike the generic :class:`~nvsubquadratic.modules.attention.Attention` + module which uses a fixed L2 (cosine) normalisation, the learnable per-head + norm here allows the model to control the scale of the dot products. + + Note: + Norm is applied before RoPE in this module (``q_norm → rope``), whereas + the generic :class:`Attention` applies RoPE before L2-norm. The order + matters for checkpoint compatibility. + + **Differences vs.** :class:`~nvsubquadratic.modules.attention.Attention` + + * Self-contained QKV + output projections (generic uses outer + ``QKVSequenceMixer``). + * RMSNorm QK-Norm instead of L2 normalisation. + * Dual-base register-aware RoPE instead of single-base uniform RoPE. + * Fixed ``[B, T, C]`` input — no multi-dimensional spatial support, no + causal masking, no context-parallelism guard. + + Attributes: + hidden_dim (int): Total channel dimension ``C``. + num_heads (int): Number of attention heads ``H``. + head_dim (int): Per-head dimension ``d_k = C / H``. + scale (float): Attention logit scale, default ``head_dim ** -0.5``. + num_patches_h (int): Height of the patch grid used for 2D RoPE. + num_patches_w (int): Width of the patch grid used for 2D RoPE. + num_registers (int): Number of register tokens ``R``. + has_cls (bool): Whether the token sequence includes a CLS token + between the patch tokens and the register tokens. + attn_dropout (float): Dropout probability applied to attention weights + during training; set to 0.0 at inference. + qkv (nn.Linear): Fused QKV projection: ``Linear(C, 3C, bias=qkv_bias)``. + proj (nn.Linear): Output projection: ``Linear(C, C, bias=out_proj_bias)``. + proj_drop (nn.Dropout | nn.Identity): Dropout on the projected output. + q_norm (nn.Module): Per-head query normaliser. Present only when + ``qk_norm`` is provided (i.e. ``self.qk_norm is True``). + k_norm (nn.Module): Per-head key normaliser. Present only when + ``qk_norm`` is provided. + qk_norm (bool): Flag indicating whether QK normalisation is active. + rope_base (float): Base frequency for patch-token RoPE. + reg_rope_base (float): Base frequency for register-token RoPE. + reg_rope_h (int): Height dimension of the register RoPE grid + (``int(num_registers ** 0.5)``). + reg_rope_w (int): Width dimension of the register RoPE grid + (``int(num_registers ** 0.5)``). + rope_cos (torch.Tensor): Non-persistent buffer of shape ``[T, head_dim]`` + containing the concatenated patch + CLS + register cosine tables. + rope_sin (torch.Tensor): Non-persistent buffer of shape ``[T, head_dim]`` + containing the concatenated patch + CLS + register sine tables. Args: - hidden_dim: Total hidden dimension. - num_heads: Number of attention heads. - num_patches_h: Height of the patch grid (for 2D RoPE). - num_patches_w: Width of the patch grid (for 2D RoPE). - num_registers: Number of register tokens. - qk_norm: LazyConfig for the QK normalization layer, or None to disable. - rope_base: Base frequency for patch RoPE. - reg_rope_base: Base frequency for register RoPE (high frequency = 100 in paper). - attn_dropout: Attention dropout rate. - proj_dropout: Output projection dropout rate. - qkv_bias: Whether to use bias in QKV projection. - out_proj_bias: Whether to use bias in the output projection. - scale: Attention scaling factor. When None, defaults to ``head_dim ** -0.5``. - init_fn_qkv_proj: Optional callable ``fn(tensor) -> None`` applied to the - QKV projection weights. When None, weights keep PyTorch's default init. - has_cls: Whether the sequence contains a CLS token (between patches and - registers in the token layout). Defaults to True. - init_fn_out_proj: Optional callable ``fn(tensor) -> None`` applied to the - output projection weights. When None, weights keep PyTorch's default init. + hidden_dim: Total hidden dimension ``C``. Must be divisible by + ``num_heads``. + num_heads: Number of attention heads ``H``. + num_patches_h: Height of the patch grid (number of patch rows). + Used to build the patch 2D RoPE table. + num_patches_w: Width of the patch grid (number of patch columns). + Used to build the patch 2D RoPE table. + num_registers: Number of register tokens ``R`` appended after the + (optional) CLS token. Must be a perfect square when > 0 so that + the register RoPE grid is square (``sqrt(R) × sqrt(R)``). + Defaults to ``4``. + has_cls: If ``True``, the token sequence contains one CLS token + immediately after the patch tokens. The CLS token receives + identity RoPE (cos=1, sin=0). Defaults to ``True``. + qk_norm: :class:`~nvsubquadratic.lazy_config.LazyConfig` for the + per-head QK normalisation module (e.g. ``RMSNorm(head_dim)``). + When ``None``, QK normalisation is disabled. Defaults to ``None``. + rope_base: Base frequency :math:`\\theta_0` for the patch RoPE + frequency schedule. Defaults to ``10000.0``. + reg_rope_base: Base frequency for the register-token RoPE schedule. + A lower base (higher frequency) gives denser angular spacing. + Defaults to ``100.0``. + attn_dropout: Dropout rate on attention weights, applied only during + training (``module.training is True``). Defaults to ``0.0``. + proj_dropout: Dropout rate on the output projection. When ``0.0``, + ``proj_drop`` is an ``nn.Identity``. Defaults to ``0.0``. + qkv_bias: Whether to include a bias term in the fused QKV projection. + Defaults to ``False``. + out_proj_bias: Whether to include a bias term in the output projection. + Defaults to ``False``. + scale: Explicit attention logit scale. When ``None``, the scale + defaults to ``head_dim ** -0.5``. Defaults to ``None``. + init_fn_qkv_proj: Optional callable ``fn(weight: Tensor) -> None`` + applied to ``self.qkv.weight`` after construction. The bias, if + present, is zero-initialised. When ``None``, PyTorch's default + Xavier uniform initialisation is used. Defaults to ``None``. + init_fn_out_proj: Optional callable ``fn(weight: Tensor) -> None`` + applied to ``self.proj.weight`` after construction. The bias, if + present, is zero-initialised. When ``None``, PyTorch's default + initialisation is used. Defaults to ``None``. + + Example:: + + import torch + from nvsubquadratic.modules.vit5_attention import ViT5Attention + from nvsubquadratic.lazy_config import LazyConfig + import nvsubquadratic.modules.rms_norm as rms_norm_mod + + # 2D patch grid of 14×14 with 4 register tokens and 1 CLS token + attn = ViT5Attention( + hidden_dim=384, + num_heads=6, + num_patches_h=14, + num_patches_w=14, + num_registers=4, + has_cls=True, + qk_norm=LazyConfig(target=rms_norm_mod.RMSNorm, dim=64), + ) + T = 14 * 14 + 1 + 4 # patches + CLS + registers = 201 + x = torch.randn(2, T, 384) # [B, T, C] + out = attn(x) # [B, T, C] + assert out.shape == x.shape """ def __init__( # noqa: D107 @@ -258,9 +504,18 @@ def flop_count(self, num_tokens: int, inference: bool = False) -> int: Total: 8 * T * D² + 4 * T² * D + 4 * T * D + qk_norm_flops. + Note: + ``num_tokens`` should equal + ``num_patches_h * num_patches_w + (1 if has_cls else 0) + num_registers`` + to match the actual sequence length seen during :meth:`forward`. + Passing a different value will give a proportionally scaled estimate. + Args: - num_tokens: Total sequence length T (cls + patches + registers). - inference: Accepted for API consistency; does not affect the count. + num_tokens: Total sequence length T + (cls + patches + registers). Should equal + ``num_patches_h * num_patches_w + (1 if has_cls else 0) + num_registers``. + inference: Accepted for API consistency with other sequence-mixer + modules (e.g. Hyena); does not affect the FLOP count. Returns: Total FLOPs as an integer. @@ -284,14 +539,49 @@ def flop_count(self, num_tokens: int, inference: bool = False) -> int: return flops def forward(self, x: torch.Tensor) -> torch.Tensor: - """Forward pass. + """Apply ViT-5 multi-head self-attention to a token sequence. + + Executes the following pipeline: + + 1. **QKV projection** — ``x W_{QKV}`` reshaped to + ``[B, T, 3, H, d_k]``, then split into Q, K, V each of shape + ``[B, T, H, d_k]``. + 2. **(Optional) QK normalisation** — ``q_norm(Q)`` and ``k_norm(K)`` + applied independently along the last (head-dim) axis. + 3. **RoPE** — ``Q' = Q * cos + rotate(Q) * sin`` and + ``K' = K * cos + rotate(K) * sin``, where ``cos`` / ``sin`` are the + precomputed ``[T, head_dim]`` buffers broadcast to + ``[1, T, 1, head_dim]`` over the batch and head axes. Uses + :func:`_rotate_half_per_axis` (split-half convention). + 4. **Transpose for SDPA** — rearrange to ``[B, H, T, d_k]``. + 5. **Scaled dot-product attention** — delegates to + ``F.scaled_dot_product_attention``; PyTorch auto-selects the best + backend (CuDNN on H100, FlashAttention on A100, etc.). The + ``dropout_p`` is set to ``self.attn_dropout`` during training and + 0.0 at inference. + 6. **Merge heads** — ``out.transpose(1, 2).reshape(B, T, C)``. + 7. **Output projection + dropout** — ``proj_drop(proj(out))``. Args: - x: [B, T, C] where T = num_patches + (1 if has_cls) + num_registers. - Token layout: [patches, (CLS), registers]. + x: Input token sequence of shape ``[B, T, C]`` where: + + * ``B`` — batch size, + * ``T = num_patches_h * num_patches_w + (1 if has_cls else 0) + + num_registers`` — total token count following the ViT-5 + layout ``[patches, (CLS,) registers]``, + * ``C = hidden_dim`` — channel dimension. + + The spatial dimensions of the patch grid are baked into the + precomputed ``rope_cos`` / ``rope_sin`` buffers; ``T`` must + match ``rope_cos.shape[0]`` exactly. Returns: - [B, T, C] + torch.Tensor: Output tensor of shape ``[B, T, C]``, the same shape + as the input. + + Raises: + RuntimeError: If ``T`` does not match ``rope_cos.shape[0]``, + which would cause a shape mismatch in the broadcast multiply. """ B, T, C = x.shape @@ -330,7 +620,15 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: out = self.proj_drop(out) return out - def extra_repr(self) -> str: # noqa: D102 + def extra_repr(self) -> str: + """Return a concise string summary of this module's configuration. + + Returns: + str: Comma-separated key=value pairs covering the most important + hyperparameters: ``hidden_dim``, ``num_heads``, ``qk_norm``, + ``num_registers``, patch grid size, ``rope_base``, and + ``reg_rope_base``. + """ return ( f"hidden_dim={self.hidden_dim}, num_heads={self.num_heads}, " f"qk_norm={self.qk_norm}, num_registers={self.num_registers}, " From 3be1131bfacfae21542dd580794cb9ee24ca68ff Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 22:26:25 +0200 Subject: [PATCH 40/72] docs(write/vit5_hyena_adapter): add module and class docstrings --- nvsubquadratic/modules/vit5_hyena_adapter.py | 195 +++++++++++++++---- 1 file changed, 162 insertions(+), 33 deletions(-) diff --git a/nvsubquadratic/modules/vit5_hyena_adapter.py b/nvsubquadratic/modules/vit5_hyena_adapter.py index 61f1cbc4..00cc4b89 100644 --- a/nvsubquadratic/modules/vit5_hyena_adapter.py +++ b/nvsubquadratic/modules/vit5_hyena_adapter.py @@ -1,13 +1,56 @@ -"""Adapter to plug 2D sequence mixers (e.g. Hyena) into the ViT5 token-sequence architecture. - -The ViT5 architecture processes [B, T, C] sequences. 2D mixers like Hyena expect -[B, H, W, C] spatial grids. This adapter reshapes the flat token sequence to a 2D -grid, applies the inner mixer, and reshapes back. - -Token ordering is handled upstream by the network (ViT5ClassificationNet). -The standard layout is [patches, CLS, registers, padding] where padding -ensures T is divisible by grid_w. This adapter is layout-agnostic: it treats -the entire sequence as a flat spatial grid reshaped to (T // grid_w, grid_w). +"""Adapter that plugs 2-D sequence mixers (e.g. Hyena) into the ViT-5 token-sequence architecture. + +Drop-in replacement interface for :class:`~nvsubquadratic.modules.vit5_attention.ViT5Attention`. + +Why an adapter is needed +------------------------ +:class:`~nvsubquadratic.modules.vit5_attention.ViT5Attention` expects a flat +``[B, T, C]`` token sequence where ``T = num_patches + (1 if has_cls) + num_registers``. +It owns its own QKV and output projections and produces ``[B, T, C]`` output — the +residual block calls ``mixer(x)`` and adds the result back to ``x``. + +2-D operators such as Hyena (wrapped in +:class:`~nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer`) expect a spatial +grid ``[B, H, W, C]`` — they cannot consume the flat sequence directly. Moreover, +``QKVSequenceMixer`` provides its own QKV and output projections that are already +part of the inner mixer; duplicating projections in the adapter would waste memory +and parameters. + +This module solves both issues with a *thin, stateless* reshape adapter: + +1. Receives ``[B, T, C]`` from the residual block. +2. Reshapes to ``[B, T // grid_w, grid_w, C]`` (a 2-D spatial grid). +3. Delegates entirely to the inner mixer (any ``[B, H, W, C]``-in / ``[B, H, W, C]``-out + module, e.g. ``QKVSequenceMixer(Hyena)``). The inner mixer must be + **shape-preserving** — its output must have the same ``[B, H, W, C]`` shape as its + input; ``reshape`` will raise a cryptic error if the shape changes. +4. Reshapes back to ``[B, T, C]`` and returns. + +The adapter itself adds **no parameters** — all learnable weights (input projection, +output projection, Hyena kernel generator) live inside the inner mixer. Projection +dimensions (``hidden_dim``, ``num_heads``, etc.) must be configured inside +``inner_mixer_cfg``; the adapter accepts no ``hidden_dim`` argument itself. + +Register tokens and the CLS token are **not** handled specially here: they are treated +as ordinary spatial positions in the grid. The calling network +(:class:`~nvsubquadratic.networks.vit5_classification.ViT5ClassificationNet`) is +responsible for padding ``T`` so that it is exactly divisible by ``grid_w`` and for +arranging tokens into a layout that makes spatial sense to the mixer. In the +hierarchical (:class:`~nvsubquadratic.networks.vit5_hierarchical_classification.ViT5HierarchicalClassificationNet`) +setting, ``grid_w`` must be consistent with the spatial width **after any patch-merging +stage** — the network supplies the correct ``grid_w`` at each stage. + +Interface contract (same as ``ViT5Attention``) +---------------------------------------------- +``forward(x, **mixer_kwargs) -> Tensor`` + +* Input: ``x`` of shape ``[B, T, C]``. +* Output: tensor of shape ``[B, T, C]``. +* Optional kwargs (e.g. ``conditioning``) are forwarded verbatim to the inner mixer. + +The module also exposes a ``flop_count(num_tokens, inference)`` method that +delegates to the inner mixer's ``flop_count``, matching the API used by the network +for FLOPs accounting. """ import torch @@ -17,11 +60,52 @@ class ViT5HyenaAdapter(nn.Module): - """Bridges ViT5's [B, T, C] token sequences and Hyena's [B, H, W, C] spatial interface. - - Args: - inner_mixer_cfg: LazyConfig for the 2D sequence mixer (e.g. QKVSequenceMixer wrapping Hyena). - grid_w: Width of the 2D spatial grid. The height is inferred as T // grid_w. + """Bridges ViT-5's ``[B, T, C]`` token sequences and Hyena's ``[B, H, W, C]`` spatial interface. + + The adapter is a **parameter-free reshape wrapper**: it does not own any QKV + projection, output projection, or positional encoding. All learnable components + live inside ``inner_mixer`` (typically a + :class:`~nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer` wrapping a + :class:`~nvsubquadratic.modules.hyena_nd.Hyena` instance). + + Data flow:: + + x: [B, T, C] + │ + ▼ reshape (T → H × grid_w, where H = T // grid_w) + x: [B, H, grid_w, C] + │ + ▼ inner_mixer (any [B, H, W, C]-preserving mixer) + x: [B, H, grid_w, C] + │ + ▼ reshape back + x: [B, T, C] + + What the adapter handles vs. what the inner mixer handles: + + * **Adapter**: shape contract (flat ↔ 2-D), ``flop_count`` delegation. + * **Inner mixer**: input projection (C → 3C), Hyena global convolution, + gating, output projection (C → C), any normalisation, and optional + FiLM / AdaLN conditioning. The mixer receives tensors in channels-last + layout ``[B, H, W, C]``; if it uses channels-first convolution internally + (as ``QKVSequenceMixer`` does) it handles the permutation itself. + + Register-token handling: + Register tokens (and the CLS token, if present) are treated as ordinary + spatial positions within the reshaped grid — no masking or special-casing + is applied. The upstream network is responsible for: + + 1. Padding the sequence so that ``T % grid_w == 0``. + 2. Choosing a ``grid_w`` that places register/CLS tokens in a predictable + row (e.g. a dedicated "register row" at the bottom of the grid), so + that the spatial convolution inside Hyena sees a consistent layout. + 3. In the hierarchical case, supplying the correct ``grid_w`` at each + stage after patch merging changes the spatial width. + + Attributes: + inner_mixer (nn.Module): The instantiated 2-D sequence mixer. + grid_w (int): Width of the 2-D spatial grid. The height is inferred + at runtime as ``T // grid_w``. """ def __init__( @@ -29,43 +113,88 @@ def __init__( inner_mixer_cfg: LazyConfig, grid_w: int, ): - """Store config and instantiate the inner 2D mixer.""" + """Instantiate the adapter and its inner 2-D mixer. + + Args: + inner_mixer_cfg: :class:`~nvsubquadratic.lazy_config.LazyConfig` + describing the 2-D sequence mixer to instantiate (e.g. + ``QKVSequenceMixer`` wrapping ``Hyena``). The instantiated module + must accept ``(x: Tensor[B, H, W, C], **kwargs)`` in channels-last + layout and return a tensor of the same shape. Any inner mixer + that uses channels-first convolution (like ``QKVSequenceMixer``) + handles the permutation internally. Projection dimensions + (``hidden_dim``, ``num_heads``, etc.) must be set inside this config; + the adapter itself accepts no ``hidden_dim`` argument. + grid_w: Width of the 2-D spatial grid. Every call to ``forward`` + must supply a sequence length ``T`` that satisfies + ``T % grid_w == 0``; the grid height is computed as + ``H = T // grid_w``. In a hierarchical network, pass the + correct ``grid_w`` for each stage (after patch merging). + """ super().__init__() self.inner_mixer = instantiate(inner_mixer_cfg) self.grid_w = grid_w def flop_count(self, num_tokens: int, inference: bool = False) -> int: - """Count FLOPs for the Hyena adapter (reshape + inner mixer). + """Delegate FLOPs accounting to the inner mixer. - Reshapes the flat token sequence [B, T, C] into a 2D spatial grid - [B, H, W, C] (where W = ``self.grid_w``, H = T // W) and delegates - to the inner mixer (QKVSequenceMixer wrapping Hyena). + The adapter's reshape operations are pure metadata re-strides — zero + arithmetic FLOPs — so the total cost is entirely determined by + ``inner_mixer.flop_count``. - The reshape itself is a free metadata operation — no FLOPs. + Note: + ``flop_count`` is a de-facto protocol, not enforced by a formal + interface. To guard against missing implementations use + ``hasattr(adapter.inner_mixer, "flop_count")``. Args: - num_tokens: Total sequence length T. Must be divisible by grid_w. - ``spatial_dims`` is computed as ``(T // grid_w, grid_w)``. - inference: Passed through to the inner mixer. + num_tokens: Total flat sequence length ``T``. Must satisfy + ``T % grid_w == 0``. The 2-D spatial dimensions passed to the + inner mixer are ``(T // grid_w, grid_w)``. + inference: Forwarded to the inner mixer. Some mixers (e.g. those + with cached Hyena kernels) report fewer FLOPs at inference time. Returns: - Total FLOPs from the inner mixer. + Total FLOPs reported by the inner mixer for a ``(T // grid_w, grid_w)`` + spatial grid. + + Raises: + AttributeError: If ``inner_mixer`` does not implement ``flop_count``. """ spatial_dims = (num_tokens // self.grid_w, self.grid_w) return self.inner_mixer.flop_count(spatial_dims, inference=inference) def forward(self, x: torch.Tensor, **mixer_kwargs) -> torch.Tensor: - """Forward pass. + """Reshape to 2-D grid, apply the inner mixer, reshape back. Args: - x: [B, T, C] token sequence. T must be divisible by grid_w. - Standard layout: [patches (H*W), CLS (1), registers (R), padding (P)]. - The adapter is layout-agnostic and reshapes the full sequence - to a 2D grid of shape (T // grid_w, grid_w). - **mixer_kwargs: Forwarded to the inner mixer (e.g. ``conditioning`` for FiLM). + x: Input token sequence of shape ``[B, T, C]`` where + + * ``B`` — batch size. + * ``T`` — total sequence length (must satisfy ``T % grid_w == 0``). + Typical layout (set by the network, not enforced here): + ``[patch_tokens (H_patch * W_patch), CLS (0 or 1), + register_tokens (R), padding (P)]``. + * ``C`` — channel / hidden dimension. + + **mixer_kwargs: Keyword arguments forwarded verbatim to + ``inner_mixer.forward``. Common keys include: + + * ``conditioning`` — FiLM/AdaLN conditioning tensor used by + some Hyena configurations. + * ``cp_group`` — process group for context-parallel (AllToAll) + sharding inside the Hyena operator. Returns: - [B, T, C] with tokens mixed via the 2D inner mixer. + Tensor of shape ``[B, T, C]`` — the token sequence after 2-D + Hyena mixing. The first ``reshape`` (to ``[B, H, W, C]``) is a + zero-copy view when ``x`` is contiguous. If ``inner_mixer`` returns + a non-contiguous tensor, the final ``reshape`` (back to ``[B, T, C]``) + triggers a contiguous copy; this does not affect correctness but can + affect memory traffic in CUDA-graph or ``torch.compile`` contexts. + + Raises: + RuntimeError: If ``T % grid_w != 0`` (implicit, from ``reshape``). """ B, T, C = x.shape x = x.reshape(B, T // self.grid_w, self.grid_w, C) @@ -74,5 +203,5 @@ def forward(self, x: torch.Tensor, **mixer_kwargs) -> torch.Tensor: return x def extra_repr(self) -> str: - """Return grid width for repr().""" + """Return ``grid_w=`` appended to PyTorch's default module repr.""" return f"grid_w={self.grid_w}" From 27dabe1d0fb0bc8f14d3775692e0ecc8e6ec9763 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 22:26:31 +0200 Subject: [PATCH 41/72] docs(write/vit5_attention): add module and class docstrings --- nvsubquadratic/modules/vit5_attention.py | 410 +++++++++++++++++++---- 1 file changed, 354 insertions(+), 56 deletions(-) diff --git a/nvsubquadratic/modules/vit5_attention.py b/nvsubquadratic/modules/vit5_attention.py index b6346eb5..20605835 100644 --- a/nvsubquadratic/modules/vit5_attention.py +++ b/nvsubquadratic/modules/vit5_attention.py @@ -1,16 +1,66 @@ """ViT-5 Attention: Multi-head self-attention with RMSNorm QK-Norm and register-aware 2D RoPE. -Key differences from the base Attention module: -- QK-Norm uses RMSNorm (learnable, per-head) instead of L2 normalization. -- RoPE is applied only to patch tokens (not cls token). -- Register tokens get their own high-frequency RoPE. -- Operates on flattened sequences [B, T, C] where T = N (patches) + 1 (cls) + R (registers). - -Optimizations vs naive implementation: -- RoPE cos/sin are precomputed as registered buffers (CUDA-graph safe, no graph breaks). -- RoPE applied via a single broadcast multiply on [B, T, H, D] — no reshape to (B*H, T, D). -- SDPA backend auto-selected by PyTorch (CuDNN preferred on H100). -- No redundant dtype casts around SDPA (autocast handles precision). +This module implements the specialised self-attention block used throughout the +ViT-5 family of hierarchical vision transformers. It is a *self-contained* +alternative to the generic :class:`~nvsubquadratic.modules.attention.Attention` +module and is **not** interchangeable with it through the +:class:`~nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer` dispatch layer. +The key ViT-5-specific design choices are described below. + +**Structural differences vs.** :class:`~nvsubquadratic.modules.attention.Attention` + +1. **Self-contained projections** — :class:`ViT5Attention` owns its own QKV + projection (``nn.Linear(C, 3C)``) and output projection (``nn.Linear(C, C)``), + plus an optional output dropout. The generic :class:`Attention` is a *pure* + attention kernel that receives pre-projected Q, K, V tensors from the + surrounding :class:`~nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer`. + :class:`ViT5Attention` is therefore consumed directly by + :class:`~nvsubquadratic.modules.vit5_residual_block.ViT5ResidualBlock` without + any outer projection wrapper. + +2. **RMSNorm QK-Norm (per-head, learnable)** — When ``qk_norm`` is provided, + :class:`ViT5Attention` instantiates two independent norm modules (one for Q, + one for K) via :func:`~nvsubquadratic.lazy_config.instantiate`. The generic + :class:`Attention` uses a *shared* L2 (cosine) normalisation function + (:func:`~nvsubquadratic.utils.qk_norm.apply_qk_norm`) without learnable + parameters. The ViT-5 norm is applied after the per-head reshape so that + each head's Q and K are normalised independently. + +3. **Register-aware, dual-base 2D RoPE** — Patch tokens receive 2D RoPE with + base frequency ``rope_base`` (default 10000). Register tokens receive + their own 2D RoPE with a *different* base ``reg_rope_base`` (default 100, + i.e., much higher frequency), reflecting their role as global context carriers + rather than spatially localised patch representations. The CLS token, when + present, receives an identity rotation (cos=1, sin=0). The generic + :class:`Attention` applies a single shared RoPE base to all tokens uniformly. + +4. **Fixed token layout** — The input sequence ``[B, T, C]`` must follow the + strict ViT-5 token layout ``[patches, (CLS,) registers]`` where + ``T = H*W + (1 if has_cls else 0) + R``. There is no support for arbitrary + spatial shapes or causal masking. The generic :class:`Attention` accepts + 1D/2D/3D channels-last tensors and supports both causal and non-causal modes. + +5. **CUDA-graph-safe precomputed RoPE buffers** — Both patch and register + cos/sin tables are concatenated into a single ``[T, head_dim]`` buffer pair + (``rope_cos``, ``rope_sin``) stored as non-persistent ``register_buffer`` + entries, making the forward pass free of dynamic tensor creation and safe + for ``torch.compile(mode="max-autotune")`` and CUDA-graph capture. + +6. **RoPE rotation convention** — :func:`_rotate_half_per_axis` uses the + *split-half* convention from ``nvsubquadratic.utils.rope.rotate_half_blh`` + rather than the interleaved rotation used in some external codebases. + Existing checkpoints trained with this convention will produce ~4 pp accuracy + drops if the rotation function is swapped. See :func:`_rotate_half_per_axis` + for the full explanation. + +**Cross-references** + +* :class:`~nvsubquadratic.modules.attention.Attention` — generic 1D/2D/3D + scaled dot-product attention, used via ``QKVSequenceMixer``. +* :class:`~nvsubquadratic.modules.vit5_residual_block.ViT5ResidualBlock` — + consumes :class:`ViT5Attention` directly as its ``sequence_mixer``. +* :func:`_build_2d_rope_flat` — constructs the flattened 2D RoPE tables. +* :func:`_rotate_half_per_axis` — split-half rotation operator used in ``forward``. """ from typing import Callable, Optional @@ -28,7 +78,7 @@ def _build_2d_rope_flat( head_dim: int, rope_base: float, ) -> tuple[torch.Tensor, torch.Tensor]: - """Precompute flattened 2D RoPE cos/sin for a (height x width) grid. + r"""Precompute flattened 2D RoPE cos/sin for a (height x width) grid. The original codebase (``nvsubquadratic.utils.rope``) computed RoPE on the fly during every forward pass. We precompute the cos/sin tables at init @@ -42,14 +92,45 @@ def _build_2d_rope_flat( tables follow the module to the correct device/dtype via ``.to()`` without manual bookkeeping. - Channel layout (matches ``_rotate_half_per_axis``): - [Y_half | X_half], each of size head_dim/2. - Within each half, frequencies are ``repeat_interleave(2)`` so that the - paired-swap rotation in ``_rotate_half_per_axis`` operates on matching - frequency pairs. + **Frequency schedule** + + For each spatial axis, frequencies follow the standard RoPE schedule: + + .. math:: + + \theta_j = \text{rope\_base}^{-2j / (d_k/2)}, + \quad j = 0, 1, \ldots, d_k/4 - 1 + + where :math:`d_k` = ``head_dim``. Angles for position :math:`p` are + :math:`\phi_{p,j} = p \cdot \theta_j`, then ``repeat_interleave(2)`` + doubles each frequency so the paired-swap rotation in + :func:`_rotate_half_per_axis` operates on matching frequency pairs. + + **Channel layout** (matches :func:`_rotate_half_per_axis`): + ``[Y_half | X_half]``, each of size ``head_dim / 2``. + Within each half, frequencies are ``repeat_interleave(2)`` so that the + paired-swap rotation in :func:`_rotate_half_per_axis` operates on matching + frequency pairs. + + Note: + ``head_dim`` must be divisible by 4 (two halves of ``head_dim/2`` each, + further split into pairs of size ``head_dim/4`` for the repeat-interleave + step). This constraint is enforced by the caller + (:class:`ViT5Attention.__init__`). + + Args: + height: Number of patch rows ``H`` in the 2D grid. + width: Number of patch columns ``W`` in the 2D grid. + head_dim: Per-head channel dimension ``d_k``. Must be divisible by 4. + rope_base: Base frequency for the geometric frequency schedule. + Typical values: ``10000.0`` for patch tokens, ``100.0`` for + register tokens (higher frequency = denser positional encoding). Returns: - (cos, sin) each of shape [height * width, head_dim]. + A tuple ``(cos, sin)`` where each tensor has shape + ``[height * width, head_dim]``. These are meant to be stored as + non-persistent ``register_buffer`` entries and concatenated with the + CLS and register entries in :class:`ViT5Attention.__init__`. """ dim_half = head_dim // 2 theta = 1.0 / (rope_base ** (torch.arange(0, dim_half, 2).float() / dim_half)) @@ -77,19 +158,45 @@ def _build_2d_rope_flat( def _rotate_half_per_axis(x: torch.Tensor) -> torch.Tensor: - """Split-half rotation applied independently to Y and X channel halves. + r"""Split-half rotation applied independently to Y and X channel halves. + + This implements the rotation component of the RoPE formula + ``x' = x * cos + rotate(x) * sin``. The rotation maps each channel vector + to its 90-degree rotated counterpart using a *split-half* convention: + + .. math:: + + \text{rotate}(x) = \text{Concat}\bigl( + -x_{y,2},\, x_{y,1},\, -x_{x,2},\, x_{x,1} + \bigr) + + where :math:`x_{y,1}` and :math:`x_{y,2}` are the first and second + quarters of the head-dim vector (the Y-axis half, each of size + ``head_dim / 4``), and similarly :math:`x_{x,1}`, :math:`x_{x,2}` for + the X-axis half. + + **Why not the standard interleaved** ``rotate_half``? - **Why not the standard interleaved ``rotate_half``?** The original training run used ``rotate_half_blh`` from ``nvsubquadratic.utils.rope``, which splits each axis-half at the midpoint - and swaps with negation ([-x2, x1]). The commonly-seen interleaved - rotation ([-x1, x0, -x3, x2, ...]) is numerically *incompatible* with + and swaps with negation (``[-x2, x1]``). The commonly-seen interleaved + rotation (``[-x1, x0, -x3, x2, ...]``) is numerically *incompatible* with checkpoints trained under the split-half convention — using it causes a ~4 pp accuracy drop at validation. This function preserves exact numerical parity with the original rotation so existing checkpoints remain valid. - Channel layout: [Y_half | X_half], each half of size D/2. - Within each half: split at D/4 and swap with negation. + **Channel layout**: ``[Y_half | X_half]``, each half of size ``D/2``. + Within each half: split at ``D/4`` and swap with negation. + + Args: + x: Query or key tensor of shape ``[B, T, H, D]`` (before the SDPA + transpose) or any shape whose last dimension is the head dimension + ``D``. ``D`` must be divisible by 4 to allow the quarter-split. + + Returns: + torch.Tensor: Rotated tensor of the same shape as ``x``, representing + the 90-degree rotation of each frequency pair within the Y and X + channel halves. """ d = x.shape[-1] d_half = d // 2 @@ -102,33 +209,172 @@ def _rotate_half_per_axis(x: torch.Tensor) -> torch.Tensor: class ViT5Attention(nn.Module): - """ViT-5 multi-head self-attention with RMSNorm QK-Norm and register-aware RoPE. - - Expects input as [B, T, C] where T = num_patches + (1 if has_cls) + num_registers. - Token layout: [patches, (CLS), registers] -- no padding (stripped by the network). - The module includes its own QKV and output projections (unlike the base Attention - which is wrapped in QKVSequenceMixer). + r"""ViT-5 multi-head self-attention with RMSNorm QK-Norm and register-aware RoPE. + + This module is the primary sequence-mixing operator for the ViT-5 family of + hierarchical vision transformers. It computes standard scaled dot-product + attention: + + .. math:: + + \text{head}_i = \text{softmax}\!\left( + \frac{Q_i K_i^\top}{\sqrt{d_k}} + \right) V_i, \quad + \text{out} = \text{Concat}(\text{head}_1, \ldots, \text{head}_H) W_O + + where :math:`d_k = C / H` is the per-head dimension and :math:`W_O` is the + output projection. Q, K, V are obtained from a single fused linear + projection :math:`[Q, K, V] = x W_{QKV}`. + + **Token layout** + + Input shape: ``[B, T, C]`` where + ``T = num_patches_h * num_patches_w + (1 if has_cls else 0) + num_registers``. + Token ordering within the sequence axis: + + .. code-block:: text + + [ patch_0, patch_1, ..., patch_{H*W-1}, (CLS,) reg_0, ..., reg_{R-1} ] + <----- H*W patch tokens ---------> <--1--> <---- R registers ----> + + This ordering must be consistent with the token layout produced by the + network's patchify + register-injection layers (see + :class:`~nvsubquadratic.networks.vit5_classification.ViT5Classifier`). + + **Positional encoding** + + Three distinct positional encodings are applied: + + * **Patch tokens** — 2D RoPE with base frequency ``rope_base`` (default + 10000). The H×W grid is linearised in row-major (Y-then-X) order. + * **CLS token** — identity rotation: cos=1, sin=0. No positional bias is + imposed on the class token. + * **Register tokens** — 2D RoPE with base ``reg_rope_base`` (default 100), + treating the ``R`` registers as a ``sqrt(R) × sqrt(R)`` grid. The higher + base frequency (lower theta) gives denser angular spacing, reflecting the + role of registers as global context carriers without fixed spatial meaning. + + All three tables are concatenated into a single buffer pair (``rope_cos``, + ``rope_sin``) of shape ``[T, head_dim]`` and applied with a single broadcast + multiply in :meth:`forward`. + + **QK normalisation** + + When ``qk_norm`` is provided, two independent norm modules (``q_norm``, + ``k_norm``) are instantiated and applied to Q and K *after* the per-head + reshape but *before* RoPE. The norm is expected to be a learnable RMSNorm + or equivalent. Unlike the generic :class:`~nvsubquadratic.modules.attention.Attention` + module which uses a fixed L2 (cosine) normalisation, the learnable per-head + norm here allows the model to control the scale of the dot products. + + Note: + Norm is applied before RoPE in this module (``q_norm → rope``), whereas + the generic :class:`Attention` applies RoPE before L2-norm. The order + matters for checkpoint compatibility. + + **Differences vs.** :class:`~nvsubquadratic.modules.attention.Attention` + + * Self-contained QKV + output projections (generic uses outer + ``QKVSequenceMixer``). + * RMSNorm QK-Norm instead of L2 normalisation. + * Dual-base register-aware RoPE instead of single-base uniform RoPE. + * Fixed ``[B, T, C]`` input — no multi-dimensional spatial support, no + causal masking, no context-parallelism guard. + + Attributes: + hidden_dim (int): Total channel dimension ``C``. + num_heads (int): Number of attention heads ``H``. + head_dim (int): Per-head dimension ``d_k = C / H``. + scale (float): Attention logit scale, default ``head_dim ** -0.5``. + num_patches_h (int): Height of the patch grid used for 2D RoPE. + num_patches_w (int): Width of the patch grid used for 2D RoPE. + num_registers (int): Number of register tokens ``R``. + has_cls (bool): Whether the token sequence includes a CLS token + between the patch tokens and the register tokens. + attn_dropout (float): Dropout probability applied to attention weights + during training; set to 0.0 at inference. + qkv (nn.Linear): Fused QKV projection: ``Linear(C, 3C, bias=qkv_bias)``. + proj (nn.Linear): Output projection: ``Linear(C, C, bias=out_proj_bias)``. + proj_drop (nn.Dropout | nn.Identity): Dropout on the projected output. + q_norm (nn.Module): Per-head query normaliser. Present only when + ``qk_norm`` is provided (i.e. ``self.qk_norm is True``). + k_norm (nn.Module): Per-head key normaliser. Present only when + ``qk_norm`` is provided. + qk_norm (bool): Flag indicating whether QK normalisation is active. + rope_base (float): Base frequency for patch-token RoPE. + reg_rope_base (float): Base frequency for register-token RoPE. + reg_rope_h (int): Height dimension of the register RoPE grid + (``int(num_registers ** 0.5)``). + reg_rope_w (int): Width dimension of the register RoPE grid + (``int(num_registers ** 0.5)``). + rope_cos (torch.Tensor): Non-persistent buffer of shape ``[T, head_dim]`` + containing the concatenated patch + CLS + register cosine tables. + rope_sin (torch.Tensor): Non-persistent buffer of shape ``[T, head_dim]`` + containing the concatenated patch + CLS + register sine tables. Args: - hidden_dim: Total hidden dimension. - num_heads: Number of attention heads. - num_patches_h: Height of the patch grid (for 2D RoPE). - num_patches_w: Width of the patch grid (for 2D RoPE). - num_registers: Number of register tokens. - qk_norm: LazyConfig for the QK normalization layer, or None to disable. - rope_base: Base frequency for patch RoPE. - reg_rope_base: Base frequency for register RoPE (high frequency = 100 in paper). - attn_dropout: Attention dropout rate. - proj_dropout: Output projection dropout rate. - qkv_bias: Whether to use bias in QKV projection. - out_proj_bias: Whether to use bias in the output projection. - scale: Attention scaling factor. When None, defaults to ``head_dim ** -0.5``. - init_fn_qkv_proj: Optional callable ``fn(tensor) -> None`` applied to the - QKV projection weights. When None, weights keep PyTorch's default init. - has_cls: Whether the sequence contains a CLS token (between patches and - registers in the token layout). Defaults to True. - init_fn_out_proj: Optional callable ``fn(tensor) -> None`` applied to the - output projection weights. When None, weights keep PyTorch's default init. + hidden_dim: Total hidden dimension ``C``. Must be divisible by + ``num_heads``. + num_heads: Number of attention heads ``H``. + num_patches_h: Height of the patch grid (number of patch rows). + Used to build the patch 2D RoPE table. + num_patches_w: Width of the patch grid (number of patch columns). + Used to build the patch 2D RoPE table. + num_registers: Number of register tokens ``R`` appended after the + (optional) CLS token. Must be a perfect square when > 0 so that + the register RoPE grid is square (``sqrt(R) × sqrt(R)``). + Defaults to ``4``. + has_cls: If ``True``, the token sequence contains one CLS token + immediately after the patch tokens. The CLS token receives + identity RoPE (cos=1, sin=0). Defaults to ``True``. + qk_norm: :class:`~nvsubquadratic.lazy_config.LazyConfig` for the + per-head QK normalisation module (e.g. ``RMSNorm(head_dim)``). + When ``None``, QK normalisation is disabled. Defaults to ``None``. + rope_base: Base frequency :math:`\\theta_0` for the patch RoPE + frequency schedule. Defaults to ``10000.0``. + reg_rope_base: Base frequency for the register-token RoPE schedule. + A lower base (higher frequency) gives denser angular spacing. + Defaults to ``100.0``. + attn_dropout: Dropout rate on attention weights, applied only during + training (``module.training is True``). Defaults to ``0.0``. + proj_dropout: Dropout rate on the output projection. When ``0.0``, + ``proj_drop`` is an ``nn.Identity``. Defaults to ``0.0``. + qkv_bias: Whether to include a bias term in the fused QKV projection. + Defaults to ``False``. + out_proj_bias: Whether to include a bias term in the output projection. + Defaults to ``False``. + scale: Explicit attention logit scale. When ``None``, the scale + defaults to ``head_dim ** -0.5``. Defaults to ``None``. + init_fn_qkv_proj: Optional callable ``fn(weight: Tensor) -> None`` + applied to ``self.qkv.weight`` after construction. The bias, if + present, is zero-initialised. When ``None``, PyTorch's default + Xavier uniform initialisation is used. Defaults to ``None``. + init_fn_out_proj: Optional callable ``fn(weight: Tensor) -> None`` + applied to ``self.proj.weight`` after construction. The bias, if + present, is zero-initialised. When ``None``, PyTorch's default + initialisation is used. Defaults to ``None``. + + Example:: + + import torch + from nvsubquadratic.modules.vit5_attention import ViT5Attention + from nvsubquadratic.lazy_config import LazyConfig + import nvsubquadratic.modules.rms_norm as rms_norm_mod + + # 2D patch grid of 14×14 with 4 register tokens and 1 CLS token + attn = ViT5Attention( + hidden_dim=384, + num_heads=6, + num_patches_h=14, + num_patches_w=14, + num_registers=4, + has_cls=True, + qk_norm=LazyConfig(target=rms_norm_mod.RMSNorm, dim=64), + ) + T = 14 * 14 + 1 + 4 # patches + CLS + registers = 201 + x = torch.randn(2, T, 384) # [B, T, C] + out = attn(x) # [B, T, C] + assert out.shape == x.shape """ def __init__( # noqa: D107 @@ -258,9 +504,18 @@ def flop_count(self, num_tokens: int, inference: bool = False) -> int: Total: 8 * T * D² + 4 * T² * D + 4 * T * D + qk_norm_flops. + Note: + ``num_tokens`` should equal + ``num_patches_h * num_patches_w + (1 if has_cls else 0) + num_registers`` + to match the actual sequence length seen during :meth:`forward`. + Passing a different value will give a proportionally scaled estimate. + Args: - num_tokens: Total sequence length T (cls + patches + registers). - inference: Accepted for API consistency; does not affect the count. + num_tokens: Total sequence length T + (cls + patches + registers). Should equal + ``num_patches_h * num_patches_w + (1 if has_cls else 0) + num_registers``. + inference: Accepted for API consistency with other sequence-mixer + modules (e.g. Hyena); does not affect the FLOP count. Returns: Total FLOPs as an integer. @@ -284,14 +539,49 @@ def flop_count(self, num_tokens: int, inference: bool = False) -> int: return flops def forward(self, x: torch.Tensor) -> torch.Tensor: - """Forward pass. + """Apply ViT-5 multi-head self-attention to a token sequence. + + Executes the following pipeline: + + 1. **QKV projection** — ``x W_{QKV}`` reshaped to + ``[B, T, 3, H, d_k]``, then split into Q, K, V each of shape + ``[B, T, H, d_k]``. + 2. **(Optional) QK normalisation** — ``q_norm(Q)`` and ``k_norm(K)`` + applied independently along the last (head-dim) axis. + 3. **RoPE** — ``Q' = Q * cos + rotate(Q) * sin`` and + ``K' = K * cos + rotate(K) * sin``, where ``cos`` / ``sin`` are the + precomputed ``[T, head_dim]`` buffers broadcast to + ``[1, T, 1, head_dim]`` over the batch and head axes. Uses + :func:`_rotate_half_per_axis` (split-half convention). + 4. **Transpose for SDPA** — rearrange to ``[B, H, T, d_k]``. + 5. **Scaled dot-product attention** — delegates to + ``F.scaled_dot_product_attention``; PyTorch auto-selects the best + backend (CuDNN on H100, FlashAttention on A100, etc.). The + ``dropout_p`` is set to ``self.attn_dropout`` during training and + 0.0 at inference. + 6. **Merge heads** — ``out.transpose(1, 2).reshape(B, T, C)``. + 7. **Output projection + dropout** — ``proj_drop(proj(out))``. Args: - x: [B, T, C] where T = num_patches + (1 if has_cls) + num_registers. - Token layout: [patches, (CLS), registers]. + x: Input token sequence of shape ``[B, T, C]`` where: + + * ``B`` — batch size, + * ``T = num_patches_h * num_patches_w + (1 if has_cls else 0) + + num_registers`` — total token count following the ViT-5 + layout ``[patches, (CLS,) registers]``, + * ``C = hidden_dim`` — channel dimension. + + The spatial dimensions of the patch grid are baked into the + precomputed ``rope_cos`` / ``rope_sin`` buffers; ``T`` must + match ``rope_cos.shape[0]`` exactly. Returns: - [B, T, C] + torch.Tensor: Output tensor of shape ``[B, T, C]``, the same shape + as the input. + + Raises: + RuntimeError: If ``T`` does not match ``rope_cos.shape[0]``, + which would cause a shape mismatch in the broadcast multiply. """ B, T, C = x.shape @@ -330,7 +620,15 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: out = self.proj_drop(out) return out - def extra_repr(self) -> str: # noqa: D102 + def extra_repr(self) -> str: + """Return a concise string summary of this module's configuration. + + Returns: + str: Comma-separated key=value pairs covering the most important + hyperparameters: ``hidden_dim``, ``num_heads``, ``qk_norm``, + ``num_registers``, patch grid size, ``rope_base``, and + ``reg_rope_base``. + """ return ( f"hidden_dim={self.hidden_dim}, num_heads={self.num_heads}, " f"qk_norm={self.qk_norm}, num_registers={self.num_registers}, " From 24dc74ff1f6897988c865faffda0dccc2d0e9de3 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 22:27:02 +0200 Subject: [PATCH 42/72] docs(review/vit5_hyena_adapter): reviewer feedback --- docs/reviews/vit5_hyena_adapter_review.md | 88 +++++++++++++++++++++++ 1 file changed, 88 insertions(+) create mode 100644 docs/reviews/vit5_hyena_adapter_review.md diff --git a/docs/reviews/vit5_hyena_adapter_review.md b/docs/reviews/vit5_hyena_adapter_review.md new file mode 100644 index 00000000..935874e4 --- /dev/null +++ b/docs/reviews/vit5_hyena_adapter_review.md @@ -0,0 +1,88 @@ +# Review: `nvsubquadratic/modules/vit5_hyena_adapter.py` + +Reviewed after Phase 1 docstring pass. Issues are numbered and actionable. + +______________________________________________________________________ + +## 1. Module docstring: "drop-in replacement" claim needs stronger qualification + +The second summary line reads: + +> "Drop-in replacement interface for `ViT5Attention`." + +This is accurate but incomplete. `ViT5Attention` owns its own QKV/output projections; the +adapter delegates all projections to `inner_mixer`. A reader doing a config swap needs to +know that `inner_mixer_cfg` (e.g. `QKVSequenceMixer`) **must** be configured with the +correct `hidden_dim`/`num_heads` — the adapter accepts no such arguments itself. The +module docstring already says "Projection dimensions must be configured inside +`inner_mixer_cfg`" but this appears deep in the body. Move or repeat this caveat in the +opening summary paragraph so it is immediately visible. + +______________________________________________________________________ + +## 2. Module docstring: shape-preservation contract lives only in item 3, not in the interface section + +Item 3 of the adapter steps says the inner mixer must be shape-preserving, but the +"Interface contract" section at the bottom says nothing about this. A reader who reads +only that section will miss the constraint. Add a bullet there: "The inner mixer must +return a tensor of the same shape `[B, H, W, C]` it received; downsampling or strided +mixers are not supported." + +______________________________________________________________________ + +## 3. Class docstring: `Attributes` block omits `inner_mixer` type information + +The Attributes block reads: + +> `inner_mixer (nn.Module): The instantiated 2-D sequence mixer.` + +This is correct but too terse. Add the channels-last expectation: "Accepts and returns +`[B, H, W, C]` tensors (channels-last). Typically a `QKVSequenceMixer` wrapping +`Hyena`." + +______________________________________________________________________ + +## 4. `__init__` docstring: `grid_w` arg does not mention what happens at hierarchical stages specifically enough + +The docstring says "In a hierarchical network, pass the correct `grid_w` for each stage +(after patch merging)." but does not tell the reader where to find that value. Add: +"After a 2× patch-merging step, `grid_w` halves; the network's stage configuration +(e.g. `ViT5HierarchicalClassificationNet`) should be the source of truth for each +stage's `grid_w`." + +______________________________________________________________________ + +## 5. `forward` docstring: `mixer_kwargs` description lists only two keys but there may be others + +The docstring lists `conditioning` and `cp_group` as "common keys". This is fine, but +add a closing note: "Any additional kwargs accepted by the concrete inner mixer are +also forwarded; consult the inner mixer's docstring for the full list." + +______________________________________________________________________ + +## 6. `forward` docstring: the `Raises` entry is imprecise about what `reshape` actually raises + +The current text says `RuntimeError: If T % grid_w != 0 (implicit, from reshape)`. +PyTorch's `reshape` actually raises `RuntimeError` with the message "shape ... is +invalid for input of size ...". To be precise: "RuntimeError: Raised by `torch.Tensor.reshape` +if `T % grid_w != 0`, with a message reporting the mismatched total element count." + +______________________________________________________________________ + +## 7. `extra_repr` docstring: does not mention the return value format + +Current: "Return `grid_w=` appended to PyTorch's default module repr." + +"Appended" is slightly wrong — `extra_repr` returns a string that PyTorch *inserts* +inside the parentheses of the repr, not appended after. Fix to: "Return the string +`'grid_w='` inserted into PyTorch's module repr (inside the parentheses)." + +______________________________________________________________________ + +## 8. Missing: note about non-contiguous output from inner mixer and `reshape` fallback + +The `forward` Returns section mentions the non-contiguous-copy caveat for the final +`reshape`. However it does not say **which** inner mixers are likely to cause this. +Add: "In practice, `QKVSequenceMixer` returns a contiguous tensor (its output +projection is a `Linear` applied on the last axis), so the final `reshape` is +typically a free view." From 31ce6c5027775cb04210f6d353b31d93669af510 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 22:27:33 +0200 Subject: [PATCH 43/72] docs(review/vit5_attention): reviewer feedback --- docs/reviews/vit5_attention_review.md | 69 +++++++++++++++++++++++++++ 1 file changed, 69 insertions(+) create mode 100644 docs/reviews/vit5_attention_review.md diff --git a/docs/reviews/vit5_attention_review.md b/docs/reviews/vit5_attention_review.md new file mode 100644 index 00000000..3ff7b63d --- /dev/null +++ b/docs/reviews/vit5_attention_review.md @@ -0,0 +1,69 @@ +# Reviewer Feedback: `nvsubquadratic/modules/vit5_attention.py` + +Reviewed after Phase 1 docstring pass. Issues are numbered for easy tracking by the integrator. + +______________________________________________________________________ + +## Critical — Missing or Incorrect Information + +**1. `_build_2d_rope_flat`: `head_dim` divisibility constraint is wrong.** + +The docstring states: *"`head_dim` must be divisible by 4."* However, the code computes `dim_half = head_dim // 2` and uses `torch.arange(0, dim_half, 2)` — which only requires `head_dim` divisible by 2, not 4. The divisibility-by-4 requirement belongs to `_rotate_half_per_axis` (which needs `d_quarter = d // 4`), not to this function. Fix: change the Note to say "`head_dim` must be divisible by 2" and move the divisibility-by-4 note to `_rotate_half_per_axis`. + +**2. `_rotate_half_per_axis`: missing assertion or note that `head_dim % 4 == 0` is *not* enforced at runtime.** + +The docstring says *"`D` must be divisible by 4"* but there is no `assert` or guard in the code. Passing an odd `head_dim` causes silent integer-truncation in the quarter-splits. Fix: add "Note: this constraint is not checked at runtime; the caller is responsible for ensuring `head_dim % 4 == 0`." + +**3. `ViT5Attention` class docstring: the QK-norm ordering claim is wrong.** + +Under **QK normalisation** the docstring says: *"applied to Q and K `after` the per-head reshape but `before` RoPE."* But in `forward()`, after `qkv.unbind()` the code does `if self.qk_norm: q = self.q_norm(q); k = self.k_norm(k)` and *then* applies RoPE. So norm is applied **before** RoPE — the docstring is correct there. However the subsequent `Note:` says *"Norm is applied before RoPE in this module (`q_norm → rope`)"* which agrees. The class-level prose and the Note are consistent; but the phrase *"after the per-head reshape"* is ambiguous — Q and K come from `qkv.unbind(dim=2)` giving shape `[B, T, H, d_k]`, so the reshape is implicit inside the `unbind`. Fix: add the explicit intermediate shape `[B, T, H, d_k]` immediately after `unbind` to make "per-head reshape" concrete. + +**4. `ViT5Attention.__init__`: no `assert hidden_dim % num_heads == 0` message.** + +The code has `assert hidden_dim % num_heads == 0` with no message string. The class docstring does not document what happens when the constraint is violated. The generic `Attention.__init__` includes `"hidden_dim must be divisible by num_heads"`. Fix: add an `AssertionError` entry to the class docstring's Args or add a `Raises` block at the class level, and/or add the message to the `assert`. + +**5. `ViT5Attention.__init__`: `num_registers` perfect-square constraint is undocumented in the body.** + +The class-level Args block says *"Must be a perfect square when > 0"* but there is no `assert int(num_registers**0.5)**2 == num_registers` in the code. If a non-square value (e.g. `num_registers=6`) is passed, `reg_rope_h = reg_rope_w = 2` silently and `_build_2d_rope_flat(2, 2, ...)` produces only 4 rows instead of 6, causing a shape mismatch in `torch.cat`. Fix: either add the assert with an informative message, or document the silent truncation explicitly in the Args block. + +______________________________________________________________________ + +## Moderate — Incomplete or Confusing Documentation + +**6. `_build_2d_rope_flat`: Return shape annotation is ambiguous for register tokens.** + +The Returns section says shape is `[height * width, head_dim]`. This is correct but callers (in `__init__`) also pass `reg_rope_h, reg_rope_w` for register tokens. The docstring does not warn that the *register* call uses a square-root approximation for the grid dimensions, so if `num_registers` is not a perfect square the returned table has fewer rows than `num_registers`. Fix: add a sentence noting this function is used for both patch and register grids and that the grid dimensions must together equal the intended token count. + +**7. `ViT5Attention.forward`: `Raises` section says `RuntimeError` but that's not what PyTorch actually raises.** + +The docstring documents: *"`RuntimeError`: If `T` does not match `rope_cos.shape[0]`."* In practice PyTorch raises a `RuntimeError` on broadcast mismatch, but the message is not user-friendly. The docstring should additionally note what the actual `T` must equal (`num_patches_h * num_patches_w + int(has_cls) + num_registers`) so users can debug easily. Fix: expand the `Raises` description to include the expected value formula. + +**8. `ViT5Attention.flop_count`: `q_norm.flop_count(T)` signature mismatch with norm modules.** + +The docstring says "Delegated to `self.q_norm` / `self.k_norm`" but does not specify whether those norm modules are expected to accept a single `num_tokens` integer or a full shape tuple. The generic `Attention` FLOP counting differs: it uses global stats. Fix: add a sentence clarifying the expected signature of the norm's `flop_count` method, e.g. `flop_count(num_tokens: int) -> int`. + +**9. `ViT5Attention` class docstring: the `Example` block uses a non-existent import.** + +```python +import nvsubquadratic.modules.rms_norm as rms_norm_mod + +qk_norm = (LazyConfig(target=rms_norm_mod.RMSNorm, dim=64),) +``` + +There is no `rms_norm` module at `nvsubquadratic.modules.rms_norm` confirmed in the codebase (the tracker lists it as `[ ]` undocumented/unverified). The example may fail for readers trying to run it. Fix: either use an abstract placeholder comment, or use the actual path if confirmed, or replace with a note that any norm accepting `[B, T, H, d_k]` tensors works. + +**10. `ViT5Attention` class docstring: register RoPE grid description mixes up base-frequency intuition.** + +The docstring says: *"The higher base frequency (lower theta) gives denser angular spacing."* Standard RoPE convention: higher `rope_base` → lower frequency (slower rotation, sparser angular spacing). But `reg_rope_base=100` is *lower* than `rope_base=10000`, so `reg_rope_base` gives *higher* frequency rotations (denser spacing). The parenthetical "(lower theta)" is confusing because `theta_j ∝ rope_base^{-2j/d}` decreases with higher base. Fix: rewrite as *"A lower base value (`reg_rope_base=100` vs `rope_base=10000`) yields higher rotation frequencies (theta decays more slowly), giving denser angular spacing for register positions."* + +______________________________________________________________________ + +## Minor — Style and Consistency + +**11. `ViT5Attention.extra_repr`: docstring does not match `attention.py` style.** + +`attention.py`'s `extra_repr` docstring begins with "Return a concise string summary…" and lists the specific keys. The Phase 1 docstring for `ViT5Attention.extra_repr` says the same but omits `has_cls` and `scale` from the listed parameters even though `extra_repr` does not emit them — which is fine, but `has_cls` is absent from `extra_repr` output yet is a consequential hyperparameter. Fix (minor): add a note that `has_cls` and `scale` are not printed by `extra_repr`. + +**12. Module docstring: cross-reference to `_build_2d_rope_flat` and `_rotate_half_per_axis` uses `:func:` but these are module-private.** + +The module docstring cross-references `:func:_build_2d_rope_flat` and `:func:_rotate_half_per_axis`. These functions are private (leading underscore) and are not exported from the module. Sphinx will produce broken references unless configured for private member documentation. Fix: either add `.. autofunction::` directives to the Sphinx config or replace `:func:` with inline monospace references (double-backtick). From 508694a82fbc5b72f32704ce6da6bd09c69a465f Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 22:28:14 +0200 Subject: [PATCH 44/72] docs(integrate/vit5_hyena_adapter): apply reviewer feedback --- docs/reviews/vit5_hyena_adapter_review.md | 88 -------------------- nvsubquadratic/modules/vit5_hyena_adapter.py | 28 ++++++- 2 files changed, 24 insertions(+), 92 deletions(-) delete mode 100644 docs/reviews/vit5_hyena_adapter_review.md diff --git a/docs/reviews/vit5_hyena_adapter_review.md b/docs/reviews/vit5_hyena_adapter_review.md deleted file mode 100644 index 935874e4..00000000 --- a/docs/reviews/vit5_hyena_adapter_review.md +++ /dev/null @@ -1,88 +0,0 @@ -# Review: `nvsubquadratic/modules/vit5_hyena_adapter.py` - -Reviewed after Phase 1 docstring pass. Issues are numbered and actionable. - -______________________________________________________________________ - -## 1. Module docstring: "drop-in replacement" claim needs stronger qualification - -The second summary line reads: - -> "Drop-in replacement interface for `ViT5Attention`." - -This is accurate but incomplete. `ViT5Attention` owns its own QKV/output projections; the -adapter delegates all projections to `inner_mixer`. A reader doing a config swap needs to -know that `inner_mixer_cfg` (e.g. `QKVSequenceMixer`) **must** be configured with the -correct `hidden_dim`/`num_heads` — the adapter accepts no such arguments itself. The -module docstring already says "Projection dimensions must be configured inside -`inner_mixer_cfg`" but this appears deep in the body. Move or repeat this caveat in the -opening summary paragraph so it is immediately visible. - -______________________________________________________________________ - -## 2. Module docstring: shape-preservation contract lives only in item 3, not in the interface section - -Item 3 of the adapter steps says the inner mixer must be shape-preserving, but the -"Interface contract" section at the bottom says nothing about this. A reader who reads -only that section will miss the constraint. Add a bullet there: "The inner mixer must -return a tensor of the same shape `[B, H, W, C]` it received; downsampling or strided -mixers are not supported." - -______________________________________________________________________ - -## 3. Class docstring: `Attributes` block omits `inner_mixer` type information - -The Attributes block reads: - -> `inner_mixer (nn.Module): The instantiated 2-D sequence mixer.` - -This is correct but too terse. Add the channels-last expectation: "Accepts and returns -`[B, H, W, C]` tensors (channels-last). Typically a `QKVSequenceMixer` wrapping -`Hyena`." - -______________________________________________________________________ - -## 4. `__init__` docstring: `grid_w` arg does not mention what happens at hierarchical stages specifically enough - -The docstring says "In a hierarchical network, pass the correct `grid_w` for each stage -(after patch merging)." but does not tell the reader where to find that value. Add: -"After a 2× patch-merging step, `grid_w` halves; the network's stage configuration -(e.g. `ViT5HierarchicalClassificationNet`) should be the source of truth for each -stage's `grid_w`." - -______________________________________________________________________ - -## 5. `forward` docstring: `mixer_kwargs` description lists only two keys but there may be others - -The docstring lists `conditioning` and `cp_group` as "common keys". This is fine, but -add a closing note: "Any additional kwargs accepted by the concrete inner mixer are -also forwarded; consult the inner mixer's docstring for the full list." - -______________________________________________________________________ - -## 6. `forward` docstring: the `Raises` entry is imprecise about what `reshape` actually raises - -The current text says `RuntimeError: If T % grid_w != 0 (implicit, from reshape)`. -PyTorch's `reshape` actually raises `RuntimeError` with the message "shape ... is -invalid for input of size ...". To be precise: "RuntimeError: Raised by `torch.Tensor.reshape` -if `T % grid_w != 0`, with a message reporting the mismatched total element count." - -______________________________________________________________________ - -## 7. `extra_repr` docstring: does not mention the return value format - -Current: "Return `grid_w=` appended to PyTorch's default module repr." - -"Appended" is slightly wrong — `extra_repr` returns a string that PyTorch *inserts* -inside the parentheses of the repr, not appended after. Fix to: "Return the string -`'grid_w='` inserted into PyTorch's module repr (inside the parentheses)." - -______________________________________________________________________ - -## 8. Missing: note about non-contiguous output from inner mixer and `reshape` fallback - -The `forward` Returns section mentions the non-contiguous-copy caveat for the final -`reshape`. However it does not say **which** inner mixers are likely to cause this. -Add: "In practice, `QKVSequenceMixer` returns a contiguous tensor (its output -projection is a `Linear` applied on the last axis), so the final `reshape` is -typically a free view." diff --git a/nvsubquadratic/modules/vit5_hyena_adapter.py b/nvsubquadratic/modules/vit5_hyena_adapter.py index 00cc4b89..46dd9f30 100644 --- a/nvsubquadratic/modules/vit5_hyena_adapter.py +++ b/nvsubquadratic/modules/vit5_hyena_adapter.py @@ -1,6 +1,9 @@ """Adapter that plugs 2-D sequence mixers (e.g. Hyena) into the ViT-5 token-sequence architecture. Drop-in replacement interface for :class:`~nvsubquadratic.modules.vit5_attention.ViT5Attention`. +**Important**: unlike ``ViT5Attention``, the adapter owns no QKV or output projections. +All projection dimensions (``hidden_dim``, ``num_heads``, etc.) must be configured +inside ``inner_mixer_cfg`` (e.g. as part of ``QKVSequenceMixer``). Why an adapter is needed ------------------------ @@ -46,6 +49,8 @@ * Input: ``x`` of shape ``[B, T, C]``. * Output: tensor of shape ``[B, T, C]``. +* The inner mixer must return a tensor of the same shape ``[B, H, W, C]`` it received; + downsampling or strided mixers are not supported. * Optional kwargs (e.g. ``conditioning``) are forwarded verbatim to the inner mixer. The module also exposes a ``flop_count(num_tokens, inference)`` method that @@ -103,7 +108,10 @@ class ViT5HyenaAdapter(nn.Module): stage after patch merging changes the spatial width. Attributes: - inner_mixer (nn.Module): The instantiated 2-D sequence mixer. + inner_mixer (nn.Module): The instantiated 2-D sequence mixer. Accepts + and returns ``[B, H, W, C]`` tensors in channels-last layout. + Typically a :class:`~nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer` + wrapping :class:`~nvsubquadratic.modules.hyena_nd.Hyena`. grid_w (int): Width of the 2-D spatial grid. The height is inferred at runtime as ``T // grid_w``. """ @@ -129,7 +137,10 @@ def __init__( must supply a sequence length ``T`` that satisfies ``T % grid_w == 0``; the grid height is computed as ``H = T // grid_w``. In a hierarchical network, pass the - correct ``grid_w`` for each stage (after patch merging). + correct ``grid_w`` for each stage (after patch merging). After + a 2× patch-merging step, ``grid_w`` halves; the network's stage + configuration (e.g. ``ViT5HierarchicalClassificationNet``) is the + source of truth for each stage's ``grid_w``. """ super().__init__() self.inner_mixer = instantiate(inner_mixer_cfg) @@ -185,6 +196,10 @@ def forward(self, x: torch.Tensor, **mixer_kwargs) -> torch.Tensor: * ``cp_group`` — process group for context-parallel (AllToAll) sharding inside the Hyena operator. + Any additional kwargs accepted by the concrete inner mixer are + also forwarded; consult the inner mixer's docstring for the full + list. + Returns: Tensor of shape ``[B, T, C]`` — the token sequence after 2-D Hyena mixing. The first ``reshape`` (to ``[B, H, W, C]``) is a @@ -192,9 +207,14 @@ def forward(self, x: torch.Tensor, **mixer_kwargs) -> torch.Tensor: a non-contiguous tensor, the final ``reshape`` (back to ``[B, T, C]``) triggers a contiguous copy; this does not affect correctness but can affect memory traffic in CUDA-graph or ``torch.compile`` contexts. + In practice, ``QKVSequenceMixer`` returns a contiguous tensor (its + output projection is a ``Linear`` on the last axis), so the final + ``reshape`` is typically a free view. Raises: - RuntimeError: If ``T % grid_w != 0`` (implicit, from ``reshape``). + RuntimeError: Raised by ``torch.Tensor.reshape`` if + ``T % grid_w != 0``, with a message reporting the mismatched + total element count. """ B, T, C = x.shape x = x.reshape(B, T // self.grid_w, self.grid_w, C) @@ -203,5 +223,5 @@ def forward(self, x: torch.Tensor, **mixer_kwargs) -> torch.Tensor: return x def extra_repr(self) -> str: - """Return ``grid_w=`` appended to PyTorch's default module repr.""" + """Return the string ``'grid_w='`` inserted into PyTorch's module repr.""" return f"grid_w={self.grid_w}" From bf803a983c07fce6b72685009931cdf28e553ff7 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 22:28:21 +0200 Subject: [PATCH 45/72] docs(write/mamba_nd): add module and class docstrings with math context --- nvsubquadratic/modules/mamba_nd.py | 257 +++++++++++++++++++++++++++-- 1 file changed, 244 insertions(+), 13 deletions(-) diff --git a/nvsubquadratic/modules/mamba_nd.py b/nvsubquadratic/modules/mamba_nd.py index 379abb81..03cafdca 100644 --- a/nvsubquadratic/modules/mamba_nd.py +++ b/nvsubquadratic/modules/mamba_nd.py @@ -1,7 +1,134 @@ # TODO: Add license header here -"""Mamba mixer layer for ND signals.""" +r"""Mamba-ND: selective state-space mixer for 1D/2D/3D signals. + +Background +---------- +Mamba (Gu & Dao, "Mamba: Linear-Time Sequence Modeling with Selective State +Spaces", arXiv:2312.00752) is a **selective state-space model (SSM)** that +achieves linear time complexity in sequence length — O(N) — while retaining +the long-range modelling capacity of attention. + +The core SSM recurrence is: + +.. math:: + + h_t &= \bar{A}_t \, h_{t-1} + \bar{B}_t \, x_t \\ + y_t &= C_t \, h_t + +where :math:`h_t \in \mathbb{R}^{d \times N}` is a latent state, and +:math:`\bar{A}_t`, :math:`\bar{B}_t` are **input-dependent** discretised +transition matrices. The key departure from classical linear SSMs (e.g. +S4, S5) is that :math:`B`, :math:`C`, and the step size :math:`\Delta` are +*functions of the input* :math:`x_t` rather than fixed parameters: + +.. math:: + + \Delta_t,\ B_t,\ C_t = \mathrm{Linear}(x_t) + +The continuous-time state matrix :math:`A` is discretised using two rules: + +* :math:`\bar{A}_t = e^{\Delta_t A}` via the **zero-order hold (ZOH)** rule. +* :math:`\bar{B}_t = \Delta_t B_t` via the **Euler (first-order)** rule. + +The Euler rule for :math:`\bar{B}_t` is the discretisation actually used by +``mamba_ssm`` (Eq. 4 in arXiv:2312.00752); the full ZOH formula +:math:`(e^{\Delta A} - I) A^{-1} B` is an alternative that the paper mentions +but does not use in the default implementation. + +This selectivity allows Mamba to focus on relevant tokens and ignore +irrelevant context, giving it an advantage over fixed-kernel convolutions +(Hyena, CKConv) on tasks requiring content-based filtering, while remaining +subquadratic unlike attention. + +Comparison with other mixers +----------------------------- ++-----------+-------------------+--------------------+--------------------------------------------+ +| Mixer | Sequence-mixing | Kernel | Complexity (in N) | ++===========+===================+====================+============================================+ +| Attention | pairwise dot-prod | input-dependent | O(N^2) | ++-----------+-------------------+--------------------+--------------------------------------------+ +| Hyena | FFT convolution | fixed (learned MLP)| O(N log N) | ++-----------+-------------------+--------------------+--------------------------------------------+ +| Mamba | SSM recurrence | input-dependent | O(N) training; O(1)/step inference | ++-----------+-------------------+--------------------+--------------------------------------------+ + +O(N) training requires the hardware-aware parallel scan in ``mamba_ssm`` +(custom CUDA extension). A naive sequential or parallel-scan implementation +is O(N log N). At inference time the recurrent form costs O(1) per step with +a fixed-size state, making Mamba particularly attractive for autoregressive +generation. + +ND generalisation strategy +-------------------------- +The Mamba recurrence is inherently sequential and 1D. This module extends it +to arbitrary spatial rank (1D sequences, 2D images, 3D volumes) by **flattening +all spatial axes into a single sequence dimension** before running the core Mamba +layer: + +.. code-block:: none + + [B, *spatial, C] + | rearrange "b ... c -> b (...) c" + v + [B, S, C] where S = prod(spatial_dims) + | Mamba1D core (or bidirectional pair) + v + [B, S, C] + | reshape back to original spatial layout + v + [B, *spatial, C] + +The scan order for multi-dimensional inputs follows the default PyTorch / +``einops`` row-major (C-contiguous) flattening: for a 2D ``[H, W]`` input the +tokens are visited in raster-scan order (row 0, col 0 to row 0, col W-1 to +row 1, col 0 and so on). For 3D ``[D, H, W]`` inputs the outermost axis varies +slowest. This ordering is fixed and is not learned; future work could explore +Hilbert-curve or zigzag orderings for improved spatial locality. + +**Vertical anisotropy warning**: in a 2D ``[H, W]`` input, vertically adjacent +pixels (same column, adjacent rows) are ``W`` tokens apart in the flattened +sequence. The SSM must propagate information across ``W`` state-update steps +to relate them, which may lose spatial correlation for wide images. +Bidirectional mode partially mitigates this by letting the reverse scan see +vertical neighbours in the forward direction. + +Bidirectional mode +------------------ +Setting ``bidirectional=True`` instantiates a second Mamba layer +(``core_layer_rev``) that processes the token sequence in **reverse order**. +The reversed output is flipped back and added to the forward output: + +.. math:: + + \text{out} = \mathrm{Mamba}(x) + + \mathrm{flip}(\mathrm{Mamba}_\mathrm{rev}(\mathrm{flip}(x))) + +This makes the effective receptive field of each position span the entire +sequence in both directions, at the cost of 2x parameters and compute. +For non-causal spatial tasks (images, volumes) bidirectional mode is strongly +recommended. + +Integration with the rest of the library +----------------------------------------- +:class:`Mamba` is designed to be used as the ``inner_mixer`` inside +:class:`~nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer` (which adds +shared QKV and output linear projections). It is also listed as a supported +mixer in :class:`~nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer`'s +dispatch table. + +Related modules +--------------- +* ``nvsubquadratic.modules.hyena_nd`` - Hyena (fixed-kernel gated conv) +* ``nvsubquadratic.modules.attention`` - multi-head self-attention +* ``nvsubquadratic.modules.sequence_mixer`` - operator-agnostic dispatch layer + +References: +---------- +Gu, A. & Dao, T. (2023). *Mamba: Linear-Time Sequence Modeling with Selective +State Spaces*. arXiv:2312.00752. +""" import torch from einops import rearrange @@ -10,14 +137,64 @@ class Mamba(torch.nn.Module): - """Mamba mixer layer for ND signals. + r"""Selective state-space mixer for ND signals. + + Wraps a 1D Mamba core layer (e.g. ``mamba_ssm.Mamba``) and extends it to + arbitrary spatial rank by flattening all spatial axes into a single + sequence dimension before the SSM recurrence and reshaping back afterward. + + The SSM recurrence computed by the core layer is: + + .. math:: + + h_t &= \bar{A}_t \, h_{t-1} + \bar{B}_t \, x_t \\ + y_t &= C_t \, h_t + + where the transition matrices :math:`\bar{A}_t`, :math:`\bar{B}_t` and + the readout matrix :math:`C_t` are all *functions of* :math:`x_t`, + derived via linear projections inside the core layer. The step size + :math:`\Delta_t` (also input-dependent) controls the discretisation: + :math:`\bar{A}_t = e^{\Delta_t A}` (ZOH) and + :math:`\bar{B}_t = \Delta_t B_t` (Euler). + + **Scan order for ND inputs**: spatial axes are flattened in row-major + (C-contiguous) order, i.e. for 2D ``[H, W]`` the sequence visits tokens + as (0,0), (0,1), ..., (0,W-1), (1,0), ... (raster-scan). For 3D + ``[D,H,W]`` the depth axis varies slowest. This ordering is fixed + (not learned). Vertically adjacent pixels are ``W`` steps apart in the + flattened sequence; see the module docstring for the anisotropy implication. - ND signals are handled by reshaping the input to [B, flatten (* spatial_dims), hidden_dim] and then passing it to the core layer. + **Bidirectional mode**: when ``bidirectional=True`` a second core layer + processes the flattened sequence in reverse, and its (re-reversed) output + is summed with the forward output. This gives every position a full-sequence + receptive field in both causal directions, which is beneficial for + non-causal spatial tasks such as image or volume modelling. - If bidirectionality is enabled, an additional Mamba layer is instantiated and used to process the reversed input. - The output of this layer is reversed and added to the output of the core (forward) layer. + Attributes: + bidirectional (bool): Whether to apply a second reversed Mamba pass. + core_layer (torch.nn.Module): The forward (or only) Mamba core. + Must accept input of shape ``[B, S, C]`` and return ``[B, S, C]``. + core_layer_rev (torch.nn.Module): The reverse Mamba core, instantiated + only when ``bidirectional=True``. When ``bidirectional=False`` this + attribute is not registered and accessing it raises + :class:`AttributeError` by design, keeping the module's parameter + count and ``state_dict`` unaffected. - The output is then reshaped back to [B, * spatial_dims, hidden_dim]. + Example:: + + import torch + from nvsubquadratic.lazy_config import LazyConfig + from nvsubquadratic.modules.mamba_nd import Mamba + from mamba_ssm import Mamba as MambaCore + + mamba = Mamba( + mamba_layer_cfg=LazyConfig(MambaCore)(d_model=128, d_state=16, d_conv=4, expand=2), + bidirectional=True, + ) + + # 2D input: batch=2, spatial=(16, 16), channels=128 + x = torch.randn(2, 16, 16, 128) + y = mamba(x) # [2, 16, 16, 128] """ def __init__( @@ -25,11 +202,30 @@ def __init__( mamba_layer_cfg: LazyConfig, bidirectional: bool = False, ): - """Initialize the Mamba mixer layer. + """Initialise the Mamba-ND wrapper. Args: - mamba_layer_cfg: LazyConfig - LazyConfig for the Mamba layer. - bidirectional: bool - Whether to use a bidirectional Mamba layer. + mamba_layer_cfg: :class:`~nvsubquadratic.lazy_config.LazyConfig` + for the underlying 1D Mamba core. The target class must + accept a 3-D tensor of shape ``[B, S, C]`` (batch, sequence + length, channels) and return a tensor of the same shape. + Typical targets include ``mamba_ssm.Mamba`` and + ``mamba_ssm.Mamba2``. ``instantiate(mamba_layer_cfg)`` is + called twice when ``bidirectional=True``; each call constructs + a fresh ``nn.Module`` with newly initialised weights, so the + two directions do not share parameters. + bidirectional: If ``True``, run a second Mamba core on the + reversed sequence and sum both outputs. This doubles + parameter count and compute but gives non-causal coverage + of the full sequence -- strongly recommended for spatial + tasks (images, volumes). Defaults to ``False``. + + Raises: + Exception: Propagated from + :func:`~nvsubquadratic.lazy_config.instantiate` if + ``mamba_layer_cfg`` cannot be constructed. Check that the + target class accepts ``[B, S, C]`` tensors and that all + required constructor arguments are provided in the config. """ super().__init__() self.bidirectional = bidirectional @@ -39,14 +235,49 @@ def __init__( if self.bidirectional: self.core_layer_rev = instantiate(mamba_layer_cfg) - def forward(self, x): - """Forward pass of the Mamba mixer layer. + def forward(self, x: torch.Tensor) -> torch.Tensor: + r"""Apply the Mamba SSM to an ND input signal. + + The forward pass performs the following steps: + + 1. **Flatten** all spatial axes into one sequence dimension: + ``[B, *spatial, C]`` to ``[B, S, C]``, where ``S = prod(spatial)``. + The flattening follows row-major (C-contiguous) order. + 2. **Forward SSM**: ``out = core_layer(x)`` -- applies the selective + SSM recurrence :math:`y_t = C_t(\bar{A}_t h_{t-1} + \bar{B}_t x_t)`. + 3. **Reverse SSM** (only when ``bidirectional=True``): + ``out_rev = core_layer_rev(flip(x))`` -- runs the SSM on the + reversed sequence, then flips back and adds to ``out``: + + .. math:: + + \text{out} \mathrel{+}= + \mathrm{flip}(\mathrm{Mamba}_\mathrm{rev}(\mathrm{flip}(x))) + + 4. **Reshape** back to the original spatial layout: + ``[B, S, C]`` to ``[B, *spatial, C]``. + + Implementation note + ------------------- + The local variable ``x`` is rebound to the flattened ``[B, S, C]`` + view after the ``rearrange`` call; the original spatial shape is + preserved in ``x_shape`` for the final ``reshape``. Args: - x (torch.Tensor): Input tensor of shape (batch_size, * spatial_dims, hidden_dim) + x: Input tensor of shape ``(B, *spatial, C)`` where ``B`` is + batch size, ``spatial`` is one or more spatial dimensions + (e.g. ``(T,)`` for 1D sequences, ``(H, W)`` for 2D images, + ``(D, H, W)`` for 3D volumes), and ``C`` is the channel + (hidden) dimension. The tensor must be in channels-last + (BHC / BHWc) layout, consistent with the rest of the library. Returns: - torch.Tensor: Output tensor of shape (batch_size, * spatial_dims, hidden_dim) + Output tensor of shape ``(B, *spatial, C)`` -- same shape and + layout as the input. When ``bidirectional=True`` the output is + the element-wise sum of the forward and reverse SSM outputs, + which doubles the effective output magnitude compared to a + unidirectional pass; downstream normalisation layers (e.g. + ``RMSNorm`` inside the residual block) absorb this scale. """ x_shape = x.shape # Reshape input to [B, flatten (* spatial_dims), hidden_dim From 027953788a0fd4819a53555849f6c01e358a4c79 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 22:29:22 +0200 Subject: [PATCH 46/72] docs(review/mamba_nd): reviewer feedback --- docs/reviews/mamba_nd_review.md | 201 ++++++++++++++++++++++++++++++++ 1 file changed, 201 insertions(+) create mode 100644 docs/reviews/mamba_nd_review.md diff --git a/docs/reviews/mamba_nd_review.md b/docs/reviews/mamba_nd_review.md new file mode 100644 index 00000000..50840991 --- /dev/null +++ b/docs/reviews/mamba_nd_review.md @@ -0,0 +1,201 @@ +# Review: `nvsubquadratic/modules/mamba_nd.py` + +Reviewed after Phase 1 docstring pass. Issues are ordered from highest to +lowest priority. Each item quotes the relevant text and states the exact fix. + +______________________________________________________________________ + +## 1. ZOH approximation presented as equality + +**Location**: module-level docstring, "Background" section. + +**Quoted text**: + +``` +\bar{B}_t = (e^{\Delta_t A} - I) A^{-1} B_t \approx \Delta_t B_t +``` + +**Issue**: The full ZOH expression is exact, but the approximation +`approx Delta_t B_t` is only valid for small `Delta_t`. More precisely, the Mamba +paper (Eq. 4 in arXiv:2312.00752) uses the Euler discretisation `B_bar_t = Delta_t B_t` +and notes that the ZOH formula is an alternative; the code in `mamba_ssm` +actually uses the Euler rule for `B`, not the full ZOH. Presenting the full +ZOH formula and then approximating it may confuse readers into thinking the +implementation uses ZOH for both `A` and `B`. + +**Fix**: Clarify that `A_bar_t = exp(Delta_t A)` (ZOH) but `B_bar_t = Delta_t B_t` (Euler / +first-order), matching what `mamba_ssm` actually computes. Remove the +intermediate exact-then-approximate chain or clearly label it "alternative ZOH +formula (not used in practice)". + +______________________________________________________________________ + +## 2. Comparison table: Mamba complexity claim is misleading for training + +**Location**: module-level docstring, "Comparison with other mixers" section. + +**Quoted text**: + +``` +| Mamba | SSM recurrence | input-dependent | O(N) | +``` + +**Issue**: O(N) is the _inference_ (autoregressive) complexity via the +recurrent form. During training Mamba uses a parallel associative scan whose +GPU-efficient implementation is O(N log N) in time, or O(N) with the +hardware-aware parallel scan (which requires special CUDA kernels -- the whole +point of the `mamba_ssm` package). Listing O(N) without qualification makes +it look strictly cheaper than Hyena at training time, which is only true with +the custom CUDA kernels. + +**Fix**: Add a note distinguishing training (parallel scan, O(N) with custom +kernels or O(N log N) naively) from inference (recurrent, O(N) per step, O(1) +state size). Add a "Notes" column or a footnote: "O(N) with hardware-aware +parallel scan (requires `mamba_ssm` CUDA extension); O(1) per step at +inference". + +______________________________________________________________________ + +## 3. Scan order: no mention of the implication for 2D spatial locality + +**Location**: module-level docstring, "ND generalisation strategy" section, and +class-level docstring paragraph starting "**Scan order for ND inputs**". + +**Quoted text**: + +``` +The scan order for multi-dimensional inputs follows the default PyTorch / +``einops`` row-major (C-contiguous) flattening: for a 2D ``[H, W]`` input the +tokens are visited in raster-scan order ... +``` + +**Issue**: This is accurate but omits the key practical consequence: tokens +that are spatially adjacent _vertically_ (same column, adjacent rows) are far +apart in the flattened sequence (W steps away), so the SSM's effective +receptive field is anisotropic -- it sees horizontal neighbours cheaply but +vertical neighbours only through W state-update steps. An external +collaborator reading this to decide whether to use Mamba for a 2D task needs +this information. + +**Fix**: Add one sentence explicitly warning about vertical anisotropy: +"Note that vertically adjacent pixels (same column, adjacent rows) are W +tokens apart in the flattened sequence; for tall images this means the forward +SSM sees them only through many state transitions, potentially losing spatial +correlation. Bidirectional mode partially mitigates this." + +______________________________________________________________________ + +## 4. `forward` docstring: `x` is rebound mid-function without a note + +**Location**: `Mamba.forward`, Args section and implementation. + +**Quoted text** (implementation): + +```python +x = rearrange(x, "b ... c -> b (...) c") +``` + +**Issue**: The `rearrange` result is assigned back to `x`, which shadows the +original argument. The original spatial shape is captured in `x_shape` first, +so this is safe, but the docstring does not warn that the local variable `x` +changes meaning mid-function (it is the flattened view from that line onward). +This is a mild readability issue but analogous to the note in `hyena_nd.py`'s +`forward` docstring about `query` being overwritten -- be consistent. + +**Fix**: Add an "Implementation note" paragraph to `forward`: +"The local variable `x` is rebound to the flattened `[B, S, C]` view after +the `rearrange` call; the original spatial shape is preserved in `x_shape` +for the final `reshape`." + +______________________________________________________________________ + +## 5. `__init__` docstring: weight-sharing ambiguity for bidirectional instantiation + +**Location**: `Mamba.__init__`, Args section for `mamba_layer_cfg`. + +**Quoted text**: + +``` +The config is instantiated once for ``core_layer`` and, when +``bidirectional=True``, a second independent instantiation is created for +``core_layer_rev`` so that the two directions have separate parameters. +``` + +**Issue**: The wording says "independent instantiation ... separate parameters" +but does not say _how_ independence is achieved -- a reader unfamiliar with +`LazyConfig` might wonder if the second call shares weights via some internal +cache. + +**Fix**: Add a clarifying phrase: "`instantiate(mamba_layer_cfg)` is called +twice with the same config; each call constructs a fresh `nn.Module` with +newly initialised weights, so the two directions do not share parameters." + +______________________________________________________________________ + +## 6. Missing `Raises` section in `__init__` + +**Location**: `Mamba.__init__` docstring. + +**Issue**: If `mamba_layer_cfg` cannot be instantiated (wrong target class, +missing required arguments), `instantiate` will raise -- most likely a +`RuntimeError`, `TypeError`, or an `omegaconf` exception, depending on the +`LazyConfig` backend. The `QKVSequenceMixer.__init__` docstring in +`sequence_mixer.py` includes a `Raises` section for this; `Mamba.__init__` +should match for consistency. + +**Fix**: Add: + +``` +Raises: + Exception: Propagated from + :func:`~nvsubquadratic.lazy_config.instantiate` if + ``mamba_layer_cfg`` cannot be constructed. Check that the target + class accepts ``[B, S, C]`` tensors and that all required constructor + arguments are provided in the config. +``` + +______________________________________________________________________ + +## 7. Class-level Attributes block: `core_layer_rev` uses vague conditional phrasing + +**Location**: `Mamba` class, Attributes block. + +**Quoted text**: + +``` +core_layer_rev (torch.nn.Module): The reverse Mamba core. Only + present when ``bidirectional=True``; accessing this attribute + when ``bidirectional=False`` raises :class:`AttributeError`. +``` + +**Issue**: Saying "raises `AttributeError`" is accurate but alarming without +context -- it looks like a bug. It would be clearer to note the attribute is +intentionally absent (not registered) to keep `state_dict` / `parameters()` +clean when unused. + +**Fix**: Reword to: "The reverse Mamba core, instantiated only when +`bidirectional=True`. When `bidirectional=False` this attribute is not +registered and accessing it raises `AttributeError` by design, keeping +the module's parameter count and `state_dict` unaffected." + +______________________________________________________________________ + +## 8. Module docstring: `References` section uses non-standard heading style + +**Location**: module-level docstring, last section. + +**Quoted text** (before ruff fix): + +``` +References: +---------- +``` + +**Issue**: Sphinx / NumPy / Google doc conventions write `References` with the +underline immediately under the heading, without a trailing colon. The rest of +the module docstring uses RST underlines without trailing colons (e.g. +`Background\n----------`). The ruff linter fixes this automatically but it is +worth noting for future edits. + +**Fix**: Use `References\n----------` (no trailing colon), matching the style +of every other section in this file. From add27d517d3cf8adf0b8e01016451e0047881253 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 22:29:30 +0200 Subject: [PATCH 47/72] docs(integrate/mamba_nd): apply reviewer feedback --- docs/reviews/mamba_nd_review.md | 201 -------------------------------- 1 file changed, 201 deletions(-) delete mode 100644 docs/reviews/mamba_nd_review.md diff --git a/docs/reviews/mamba_nd_review.md b/docs/reviews/mamba_nd_review.md deleted file mode 100644 index 50840991..00000000 --- a/docs/reviews/mamba_nd_review.md +++ /dev/null @@ -1,201 +0,0 @@ -# Review: `nvsubquadratic/modules/mamba_nd.py` - -Reviewed after Phase 1 docstring pass. Issues are ordered from highest to -lowest priority. Each item quotes the relevant text and states the exact fix. - -______________________________________________________________________ - -## 1. ZOH approximation presented as equality - -**Location**: module-level docstring, "Background" section. - -**Quoted text**: - -``` -\bar{B}_t = (e^{\Delta_t A} - I) A^{-1} B_t \approx \Delta_t B_t -``` - -**Issue**: The full ZOH expression is exact, but the approximation -`approx Delta_t B_t` is only valid for small `Delta_t`. More precisely, the Mamba -paper (Eq. 4 in arXiv:2312.00752) uses the Euler discretisation `B_bar_t = Delta_t B_t` -and notes that the ZOH formula is an alternative; the code in `mamba_ssm` -actually uses the Euler rule for `B`, not the full ZOH. Presenting the full -ZOH formula and then approximating it may confuse readers into thinking the -implementation uses ZOH for both `A` and `B`. - -**Fix**: Clarify that `A_bar_t = exp(Delta_t A)` (ZOH) but `B_bar_t = Delta_t B_t` (Euler / -first-order), matching what `mamba_ssm` actually computes. Remove the -intermediate exact-then-approximate chain or clearly label it "alternative ZOH -formula (not used in practice)". - -______________________________________________________________________ - -## 2. Comparison table: Mamba complexity claim is misleading for training - -**Location**: module-level docstring, "Comparison with other mixers" section. - -**Quoted text**: - -``` -| Mamba | SSM recurrence | input-dependent | O(N) | -``` - -**Issue**: O(N) is the _inference_ (autoregressive) complexity via the -recurrent form. During training Mamba uses a parallel associative scan whose -GPU-efficient implementation is O(N log N) in time, or O(N) with the -hardware-aware parallel scan (which requires special CUDA kernels -- the whole -point of the `mamba_ssm` package). Listing O(N) without qualification makes -it look strictly cheaper than Hyena at training time, which is only true with -the custom CUDA kernels. - -**Fix**: Add a note distinguishing training (parallel scan, O(N) with custom -kernels or O(N log N) naively) from inference (recurrent, O(N) per step, O(1) -state size). Add a "Notes" column or a footnote: "O(N) with hardware-aware -parallel scan (requires `mamba_ssm` CUDA extension); O(1) per step at -inference". - -______________________________________________________________________ - -## 3. Scan order: no mention of the implication for 2D spatial locality - -**Location**: module-level docstring, "ND generalisation strategy" section, and -class-level docstring paragraph starting "**Scan order for ND inputs**". - -**Quoted text**: - -``` -The scan order for multi-dimensional inputs follows the default PyTorch / -``einops`` row-major (C-contiguous) flattening: for a 2D ``[H, W]`` input the -tokens are visited in raster-scan order ... -``` - -**Issue**: This is accurate but omits the key practical consequence: tokens -that are spatially adjacent _vertically_ (same column, adjacent rows) are far -apart in the flattened sequence (W steps away), so the SSM's effective -receptive field is anisotropic -- it sees horizontal neighbours cheaply but -vertical neighbours only through W state-update steps. An external -collaborator reading this to decide whether to use Mamba for a 2D task needs -this information. - -**Fix**: Add one sentence explicitly warning about vertical anisotropy: -"Note that vertically adjacent pixels (same column, adjacent rows) are W -tokens apart in the flattened sequence; for tall images this means the forward -SSM sees them only through many state transitions, potentially losing spatial -correlation. Bidirectional mode partially mitigates this." - -______________________________________________________________________ - -## 4. `forward` docstring: `x` is rebound mid-function without a note - -**Location**: `Mamba.forward`, Args section and implementation. - -**Quoted text** (implementation): - -```python -x = rearrange(x, "b ... c -> b (...) c") -``` - -**Issue**: The `rearrange` result is assigned back to `x`, which shadows the -original argument. The original spatial shape is captured in `x_shape` first, -so this is safe, but the docstring does not warn that the local variable `x` -changes meaning mid-function (it is the flattened view from that line onward). -This is a mild readability issue but analogous to the note in `hyena_nd.py`'s -`forward` docstring about `query` being overwritten -- be consistent. - -**Fix**: Add an "Implementation note" paragraph to `forward`: -"The local variable `x` is rebound to the flattened `[B, S, C]` view after -the `rearrange` call; the original spatial shape is preserved in `x_shape` -for the final `reshape`." - -______________________________________________________________________ - -## 5. `__init__` docstring: weight-sharing ambiguity for bidirectional instantiation - -**Location**: `Mamba.__init__`, Args section for `mamba_layer_cfg`. - -**Quoted text**: - -``` -The config is instantiated once for ``core_layer`` and, when -``bidirectional=True``, a second independent instantiation is created for -``core_layer_rev`` so that the two directions have separate parameters. -``` - -**Issue**: The wording says "independent instantiation ... separate parameters" -but does not say _how_ independence is achieved -- a reader unfamiliar with -`LazyConfig` might wonder if the second call shares weights via some internal -cache. - -**Fix**: Add a clarifying phrase: "`instantiate(mamba_layer_cfg)` is called -twice with the same config; each call constructs a fresh `nn.Module` with -newly initialised weights, so the two directions do not share parameters." - -______________________________________________________________________ - -## 6. Missing `Raises` section in `__init__` - -**Location**: `Mamba.__init__` docstring. - -**Issue**: If `mamba_layer_cfg` cannot be instantiated (wrong target class, -missing required arguments), `instantiate` will raise -- most likely a -`RuntimeError`, `TypeError`, or an `omegaconf` exception, depending on the -`LazyConfig` backend. The `QKVSequenceMixer.__init__` docstring in -`sequence_mixer.py` includes a `Raises` section for this; `Mamba.__init__` -should match for consistency. - -**Fix**: Add: - -``` -Raises: - Exception: Propagated from - :func:`~nvsubquadratic.lazy_config.instantiate` if - ``mamba_layer_cfg`` cannot be constructed. Check that the target - class accepts ``[B, S, C]`` tensors and that all required constructor - arguments are provided in the config. -``` - -______________________________________________________________________ - -## 7. Class-level Attributes block: `core_layer_rev` uses vague conditional phrasing - -**Location**: `Mamba` class, Attributes block. - -**Quoted text**: - -``` -core_layer_rev (torch.nn.Module): The reverse Mamba core. Only - present when ``bidirectional=True``; accessing this attribute - when ``bidirectional=False`` raises :class:`AttributeError`. -``` - -**Issue**: Saying "raises `AttributeError`" is accurate but alarming without -context -- it looks like a bug. It would be clearer to note the attribute is -intentionally absent (not registered) to keep `state_dict` / `parameters()` -clean when unused. - -**Fix**: Reword to: "The reverse Mamba core, instantiated only when -`bidirectional=True`. When `bidirectional=False` this attribute is not -registered and accessing it raises `AttributeError` by design, keeping -the module's parameter count and `state_dict` unaffected." - -______________________________________________________________________ - -## 8. Module docstring: `References` section uses non-standard heading style - -**Location**: module-level docstring, last section. - -**Quoted text** (before ruff fix): - -``` -References: ----------- -``` - -**Issue**: Sphinx / NumPy / Google doc conventions write `References` with the -underline immediately under the heading, without a trailing colon. The rest of -the module docstring uses RST underlines without trailing colons (e.g. -`Background\n----------`). The ruff linter fixes this automatically but it is -worth noting for future edits. - -**Fix**: Use `References\n----------` (no trailing colon), matching the style -of every other section in this file. From 6610f02168cc27fd58833367b67387b36ecef53e Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 22:29:40 +0200 Subject: [PATCH 48/72] docs(integrate/vit5_attention): apply reviewer feedback --- docs/reviews/vit5_attention_review.md | 69 ----------------- nvsubquadratic/modules/vit5_attention.py | 98 ++++++++++++++++-------- 2 files changed, 66 insertions(+), 101 deletions(-) delete mode 100644 docs/reviews/vit5_attention_review.md diff --git a/docs/reviews/vit5_attention_review.md b/docs/reviews/vit5_attention_review.md deleted file mode 100644 index 3ff7b63d..00000000 --- a/docs/reviews/vit5_attention_review.md +++ /dev/null @@ -1,69 +0,0 @@ -# Reviewer Feedback: `nvsubquadratic/modules/vit5_attention.py` - -Reviewed after Phase 1 docstring pass. Issues are numbered for easy tracking by the integrator. - -______________________________________________________________________ - -## Critical — Missing or Incorrect Information - -**1. `_build_2d_rope_flat`: `head_dim` divisibility constraint is wrong.** - -The docstring states: *"`head_dim` must be divisible by 4."* However, the code computes `dim_half = head_dim // 2` and uses `torch.arange(0, dim_half, 2)` — which only requires `head_dim` divisible by 2, not 4. The divisibility-by-4 requirement belongs to `_rotate_half_per_axis` (which needs `d_quarter = d // 4`), not to this function. Fix: change the Note to say "`head_dim` must be divisible by 2" and move the divisibility-by-4 note to `_rotate_half_per_axis`. - -**2. `_rotate_half_per_axis`: missing assertion or note that `head_dim % 4 == 0` is *not* enforced at runtime.** - -The docstring says *"`D` must be divisible by 4"* but there is no `assert` or guard in the code. Passing an odd `head_dim` causes silent integer-truncation in the quarter-splits. Fix: add "Note: this constraint is not checked at runtime; the caller is responsible for ensuring `head_dim % 4 == 0`." - -**3. `ViT5Attention` class docstring: the QK-norm ordering claim is wrong.** - -Under **QK normalisation** the docstring says: *"applied to Q and K `after` the per-head reshape but `before` RoPE."* But in `forward()`, after `qkv.unbind()` the code does `if self.qk_norm: q = self.q_norm(q); k = self.k_norm(k)` and *then* applies RoPE. So norm is applied **before** RoPE — the docstring is correct there. However the subsequent `Note:` says *"Norm is applied before RoPE in this module (`q_norm → rope`)"* which agrees. The class-level prose and the Note are consistent; but the phrase *"after the per-head reshape"* is ambiguous — Q and K come from `qkv.unbind(dim=2)` giving shape `[B, T, H, d_k]`, so the reshape is implicit inside the `unbind`. Fix: add the explicit intermediate shape `[B, T, H, d_k]` immediately after `unbind` to make "per-head reshape" concrete. - -**4. `ViT5Attention.__init__`: no `assert hidden_dim % num_heads == 0` message.** - -The code has `assert hidden_dim % num_heads == 0` with no message string. The class docstring does not document what happens when the constraint is violated. The generic `Attention.__init__` includes `"hidden_dim must be divisible by num_heads"`. Fix: add an `AssertionError` entry to the class docstring's Args or add a `Raises` block at the class level, and/or add the message to the `assert`. - -**5. `ViT5Attention.__init__`: `num_registers` perfect-square constraint is undocumented in the body.** - -The class-level Args block says *"Must be a perfect square when > 0"* but there is no `assert int(num_registers**0.5)**2 == num_registers` in the code. If a non-square value (e.g. `num_registers=6`) is passed, `reg_rope_h = reg_rope_w = 2` silently and `_build_2d_rope_flat(2, 2, ...)` produces only 4 rows instead of 6, causing a shape mismatch in `torch.cat`. Fix: either add the assert with an informative message, or document the silent truncation explicitly in the Args block. - -______________________________________________________________________ - -## Moderate — Incomplete or Confusing Documentation - -**6. `_build_2d_rope_flat`: Return shape annotation is ambiguous for register tokens.** - -The Returns section says shape is `[height * width, head_dim]`. This is correct but callers (in `__init__`) also pass `reg_rope_h, reg_rope_w` for register tokens. The docstring does not warn that the *register* call uses a square-root approximation for the grid dimensions, so if `num_registers` is not a perfect square the returned table has fewer rows than `num_registers`. Fix: add a sentence noting this function is used for both patch and register grids and that the grid dimensions must together equal the intended token count. - -**7. `ViT5Attention.forward`: `Raises` section says `RuntimeError` but that's not what PyTorch actually raises.** - -The docstring documents: *"`RuntimeError`: If `T` does not match `rope_cos.shape[0]`."* In practice PyTorch raises a `RuntimeError` on broadcast mismatch, but the message is not user-friendly. The docstring should additionally note what the actual `T` must equal (`num_patches_h * num_patches_w + int(has_cls) + num_registers`) so users can debug easily. Fix: expand the `Raises` description to include the expected value formula. - -**8. `ViT5Attention.flop_count`: `q_norm.flop_count(T)` signature mismatch with norm modules.** - -The docstring says "Delegated to `self.q_norm` / `self.k_norm`" but does not specify whether those norm modules are expected to accept a single `num_tokens` integer or a full shape tuple. The generic `Attention` FLOP counting differs: it uses global stats. Fix: add a sentence clarifying the expected signature of the norm's `flop_count` method, e.g. `flop_count(num_tokens: int) -> int`. - -**9. `ViT5Attention` class docstring: the `Example` block uses a non-existent import.** - -```python -import nvsubquadratic.modules.rms_norm as rms_norm_mod - -qk_norm = (LazyConfig(target=rms_norm_mod.RMSNorm, dim=64),) -``` - -There is no `rms_norm` module at `nvsubquadratic.modules.rms_norm` confirmed in the codebase (the tracker lists it as `[ ]` undocumented/unverified). The example may fail for readers trying to run it. Fix: either use an abstract placeholder comment, or use the actual path if confirmed, or replace with a note that any norm accepting `[B, T, H, d_k]` tensors works. - -**10. `ViT5Attention` class docstring: register RoPE grid description mixes up base-frequency intuition.** - -The docstring says: *"The higher base frequency (lower theta) gives denser angular spacing."* Standard RoPE convention: higher `rope_base` → lower frequency (slower rotation, sparser angular spacing). But `reg_rope_base=100` is *lower* than `rope_base=10000`, so `reg_rope_base` gives *higher* frequency rotations (denser spacing). The parenthetical "(lower theta)" is confusing because `theta_j ∝ rope_base^{-2j/d}` decreases with higher base. Fix: rewrite as *"A lower base value (`reg_rope_base=100` vs `rope_base=10000`) yields higher rotation frequencies (theta decays more slowly), giving denser angular spacing for register positions."* - -______________________________________________________________________ - -## Minor — Style and Consistency - -**11. `ViT5Attention.extra_repr`: docstring does not match `attention.py` style.** - -`attention.py`'s `extra_repr` docstring begins with "Return a concise string summary…" and lists the specific keys. The Phase 1 docstring for `ViT5Attention.extra_repr` says the same but omits `has_cls` and `scale` from the listed parameters even though `extra_repr` does not emit them — which is fine, but `has_cls` is absent from `extra_repr` output yet is a consequential hyperparameter. Fix (minor): add a note that `has_cls` and `scale` are not printed by `extra_repr`. - -**12. Module docstring: cross-reference to `_build_2d_rope_flat` and `_rotate_half_per_axis` uses `:func:` but these are module-private.** - -The module docstring cross-references `:func:_build_2d_rope_flat` and `:func:_rotate_half_per_axis`. These functions are private (leading underscore) and are not exported from the module. Sphinx will produce broken references unless configured for private member documentation. Fix: either add `.. autofunction::` directives to the Sphinx config or replace `:func:` with inline monospace references (double-backtick). diff --git a/nvsubquadratic/modules/vit5_attention.py b/nvsubquadratic/modules/vit5_attention.py index 20605835..ebfe1685 100644 --- a/nvsubquadratic/modules/vit5_attention.py +++ b/nvsubquadratic/modules/vit5_attention.py @@ -59,8 +59,8 @@ scaled dot-product attention, used via ``QKVSequenceMixer``. * :class:`~nvsubquadratic.modules.vit5_residual_block.ViT5ResidualBlock` — consumes :class:`ViT5Attention` directly as its ``sequence_mixer``. -* :func:`_build_2d_rope_flat` — constructs the flattened 2D RoPE tables. -* :func:`_rotate_half_per_axis` — split-half rotation operator used in ``forward``. +* ``_build_2d_rope_flat`` — constructs the flattened 2D RoPE tables (module-private). +* ``_rotate_half_per_axis`` — split-half rotation operator used in ``forward`` (module-private). """ from typing import Callable, Optional @@ -113,24 +113,31 @@ def _build_2d_rope_flat( frequency pairs. Note: - ``head_dim`` must be divisible by 4 (two halves of ``head_dim/2`` each, - further split into pairs of size ``head_dim/4`` for the repeat-interleave - step). This constraint is enforced by the caller - (:class:`ViT5Attention.__init__`). + ``head_dim`` must be divisible by 2 so that ``dim_half = head_dim // 2`` + is an integer and ``torch.arange(0, dim_half, 2)`` produces a valid + frequency vector. The stricter divisibility-by-4 requirement belongs + to ``_rotate_half_per_axis``, which needs the quarter-split. Args: height: Number of patch rows ``H`` in the 2D grid. width: Number of patch columns ``W`` in the 2D grid. - head_dim: Per-head channel dimension ``d_k``. Must be divisible by 4. + head_dim: Per-head channel dimension ``d_k``. Must be divisible by 2. rope_base: Base frequency for the geometric frequency schedule. Typical values: ``10000.0`` for patch tokens, ``100.0`` for - register tokens (higher frequency = denser positional encoding). + register tokens (lower base → higher frequency → denser angular spacing). Returns: A tuple ``(cos, sin)`` where each tensor has shape ``[height * width, head_dim]``. These are meant to be stored as non-persistent ``register_buffer`` entries and concatenated with the CLS and register entries in :class:`ViT5Attention.__init__`. + + Note: this function is also called for the register-token RoPE grid with + ``height = reg_rope_h``, ``width = reg_rope_w``. If ``num_registers`` is + not a perfect square, ``reg_rope_h * reg_rope_w < num_registers`` and the + returned table will have fewer rows than expected, causing a shape mismatch + in the subsequent ``torch.cat``. The caller is responsible for ensuring + ``height * width`` equals the intended token count. """ dim_half = head_dim // 2 theta = 1.0 / (rope_base ** (torch.arange(0, dim_half, 2).float() / dim_half)) @@ -197,6 +204,12 @@ def _rotate_half_per_axis(x: torch.Tensor) -> torch.Tensor: torch.Tensor: Rotated tensor of the same shape as ``x``, representing the 90-degree rotation of each frequency pair within the Y and X channel halves. + + Note: + The ``D % 4 == 0`` constraint is **not** checked at runtime. Passing a + ``head_dim`` that is not divisible by 4 will cause silent integer + truncation in the quarter-splits, producing incorrect rotations. The + caller (:class:`ViT5Attention`) is responsible for ensuring this holds. """ d = x.shape[-1] d_half = d // 2 @@ -250,9 +263,11 @@ class ViT5Attention(nn.Module): * **CLS token** — identity rotation: cos=1, sin=0. No positional bias is imposed on the class token. * **Register tokens** — 2D RoPE with base ``reg_rope_base`` (default 100), - treating the ``R`` registers as a ``sqrt(R) × sqrt(R)`` grid. The higher - base frequency (lower theta) gives denser angular spacing, reflecting the - role of registers as global context carriers without fixed spatial meaning. + treating the ``R`` registers as a ``sqrt(R) × sqrt(R)`` grid. A lower base + value (``reg_rope_base=100`` vs ``rope_base=10000``) yields higher rotation + frequencies (theta decays more slowly across head-dim pairs), giving denser + angular spacing for register positions. This reflects their role as global + context carriers without fixed spatial meaning. All three tables are concatenated into a single buffer pair (``rope_cos``, ``rope_sin``) of shape ``[T, head_dim]`` and applied with a single broadcast @@ -261,16 +276,20 @@ class ViT5Attention(nn.Module): **QK normalisation** When ``qk_norm`` is provided, two independent norm modules (``q_norm``, - ``k_norm``) are instantiated and applied to Q and K *after* the per-head - reshape but *before* RoPE. The norm is expected to be a learnable RMSNorm - or equivalent. Unlike the generic :class:`~nvsubquadratic.modules.attention.Attention` - module which uses a fixed L2 (cosine) normalisation, the learnable per-head - norm here allows the model to control the scale of the dot products. + ``k_norm``) are instantiated and applied to Q and K after ``qkv.unbind()`` + produces tensors of shape ``[B, T, H, d_k]``, and *before* RoPE. The norm + is expected to be a learnable RMSNorm or equivalent (accepting input of shape + ``[B, T, H, d_k]`` and normalising along the last axis). Unlike the generic + :class:`~nvsubquadratic.modules.attention.Attention` module which uses a fixed + L2 (cosine) normalisation, the learnable per-head norm here allows the model + to control the scale of the dot products. Note: - Norm is applied before RoPE in this module (``q_norm → rope``), whereas - the generic :class:`Attention` applies RoPE before L2-norm. The order - matters for checkpoint compatibility. + Norm is applied **before** RoPE in this module (order: ``unbind → + q_norm/k_norm → rope → SDPA``), whereas the generic :class:`Attention` + applies RoPE before L2-norm. The order matters for checkpoint + compatibility — swapping the two will change the effective positional + encoding applied to normalised queries and keys. **Differences vs.** :class:`~nvsubquadratic.modules.attention.Attention` @@ -321,9 +340,12 @@ class ViT5Attention(nn.Module): num_patches_w: Width of the patch grid (number of patch columns). Used to build the patch 2D RoPE table. num_registers: Number of register tokens ``R`` appended after the - (optional) CLS token. Must be a perfect square when > 0 so that - the register RoPE grid is square (``sqrt(R) × sqrt(R)``). - Defaults to ``4``. + (optional) CLS token. Should be a perfect square when > 0 so + that the register RoPE grid is exactly ``sqrt(R) × sqrt(R)``. + If ``R`` is not a perfect square, ``reg_rope_h = reg_rope_w = + int(R**0.5)`` silently truncates, producing only + ``reg_rope_h * reg_rope_w < R`` RoPE rows and causing a + ``torch.cat`` shape mismatch at init time. Defaults to ``4``. has_cls: If ``True``, the token sequence contains one CLS token immediately after the patch tokens. The CLS token receives identity RoPE (cos=1, sin=0). Defaults to ``True``. @@ -354,14 +376,15 @@ class ViT5Attention(nn.Module): present, is zero-initialised. When ``None``, PyTorch's default initialisation is used. Defaults to ``None``. + Raises: + AssertionError: If ``hidden_dim % num_heads != 0``. + Example:: import torch from nvsubquadratic.modules.vit5_attention import ViT5Attention - from nvsubquadratic.lazy_config import LazyConfig - import nvsubquadratic.modules.rms_norm as rms_norm_mod - # 2D patch grid of 14×14 with 4 register tokens and 1 CLS token + # 2D patch grid of 14x14 with 4 register tokens and 1 CLS token, no QK norm attn = ViT5Attention( hidden_dim=384, num_heads=6, @@ -369,12 +392,13 @@ class ViT5Attention(nn.Module): num_patches_w=14, num_registers=4, has_cls=True, - qk_norm=LazyConfig(target=rms_norm_mod.RMSNorm, dim=64), ) T = 14 * 14 + 1 + 4 # patches + CLS + registers = 201 x = torch.randn(2, T, 384) # [B, T, C] out = attn(x) # [B, T, C] assert out.shape == x.shape + # To enable QK-norm, pass a LazyConfig targeting any norm module + # that accepts [B, T, H, d_k] tensors and normalises along the last axis. """ def __init__( # noqa: D107 @@ -397,7 +421,7 @@ def __init__( # noqa: D107 init_fn_out_proj: Optional[Callable[[torch.Tensor], None]] = None, ): super().__init__() - assert hidden_dim % num_heads == 0 + assert hidden_dim % num_heads == 0, "hidden_dim must be divisible by num_heads" self.hidden_dim = hidden_dim self.num_heads = num_heads self.head_dim = hidden_dim // num_heads @@ -488,6 +512,9 @@ def flop_count(self, num_tokens: int, inference: bool = False) -> int: Three projections packed into one: 2 * T * D * 3D. 2. QK-Norm (2x RMSNorm on Q and K): Delegated to self.q_norm / self.k_norm. Only counted when ``self.qk_norm`` is True; 0 otherwise. + Each norm module must expose ``flop_count(num_tokens: int) -> int`` + returning the cost for a sequence of ``num_tokens`` tokens across + all heads (i.e. for the full ``[B, T, H, d_k]`` shaped input). 3. RoPE on Q and K: 4 * T * D Each of Q, K: x * cos + rotate(x) * sin = 2 elementwise multiplies per element, over T * D elements, for both Q and K. @@ -581,7 +608,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: Raises: RuntimeError: If ``T`` does not match ``rope_cos.shape[0]``, - which would cause a shape mismatch in the broadcast multiply. + causing a shape mismatch in the broadcast multiply + ``q * cos``. The expected value is + ``num_patches_h * num_patches_w + int(has_cls) + num_registers`` + as set at construction time. """ B, T, C = x.shape @@ -623,11 +653,15 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: def extra_repr(self) -> str: """Return a concise string summary of this module's configuration. + Note: + ``has_cls`` and ``scale`` are consequential hyperparameters that + are **not** included in the output string. Use + ``module.has_cls`` and ``module.scale`` to inspect them directly. + Returns: - str: Comma-separated key=value pairs covering the most important - hyperparameters: ``hidden_dim``, ``num_heads``, ``qk_norm``, - ``num_registers``, patch grid size, ``rope_base``, and - ``reg_rope_base``. + str: Comma-separated key=value pairs covering ``hidden_dim``, + ``num_heads``, ``qk_norm``, ``num_registers``, patch grid size, + ``rope_base``, and ``reg_rope_base``. """ return ( f"hidden_dim={self.hidden_dim}, num_heads={self.num_heads}, " From 421fc8ee946017e80ee3b08b31c07f7e23693603 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 22:31:22 +0200 Subject: [PATCH 49/72] docs(tracker): mark vit5_attention, vit5_hyena_adapter, condition_mixer, mamba_nd as done --- docs-tracker.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs-tracker.md b/docs-tracker.md index bd89a0c5..24bb3b0e 100644 --- a/docs-tracker.md +++ b/docs-tracker.md @@ -46,12 +46,12 @@ Work bottom-up: primitive ops → modules → networks → experiments. | `hyena_nd.py` | \[x\] | Hyena operator (ND) — two-gate sandwich, AllToAll CP, BC-aware convolution | | `ckconv_nd.py` | \[x\] | CKConv (ND) — implicit kernel `k_θ(p) = MLP_θ(pos_enc(p))`, FFT domain, BC modes | | `ckconv_multihead_nd.py` | \[x\] | Multi-head CKConv — H heads, dense d×d kernel per head, low-rank U·V factorisation | -| `mamba_nd.py` | \[ \] | Mamba SSM (ND) | +| `mamba_nd.py` | \[x\] | Mamba SSM (ND) — selective SSM, ZOH discretisation, raster-scan ND, bidirectional mode | | `attention.py` | \[x\] | Scaled dot-product attention — multi-head, RoPE, ND spatial, O(L²) FLOP formula | -| `vit5_attention.py` | \[ \] | ViT5 attention variant | -| `vit5_hyena_adapter.py` | \[ \] | Hyena adapter for ViT5 | +| `vit5_attention.py` | \[x\] | ViT5 attention — register-aware 2D RoPE, QK-norm, CUDA-graph-safe buffers | +| `vit5_hyena_adapter.py` | \[x\] | Hyena adapter for ViT5 — drop-in for vit5_attention, register-token + hierarchy support | | `sequence_mixer.py` | \[x\] | Operator-agnostic dispatch layer (Hyena / Attention / CKConv / Mamba) | -| `condition_mixer.py` | \[ \] | Conditioning mixer | +| `condition_mixer.py` | \[x\] | Cross-attention conditioning mixer — both global (B,C) and spatial (B,\*,C) signals | | `residual_block.py` | \[x\] | Residual block — pre-norm + mixer + MLP, optional FiLM/AdaLN-Zero conditioning | | `vit5_residual_block.py` | \[x\] | ViT5 residual block — LayerScale, register-token conditioning, no condition-mixer branch | | `patchify.py` | \[x\] | Patch embedding — strided conv, 1D/2D/3D, channels-last layout | From 7c9ce22e2d0cd57e3b71e55b4c411b9e2463bed2 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Tue, 26 May 2026 09:06:55 +0200 Subject: [PATCH 50/72] docs(write+integrate/mlp,grn,layer_scale,masks_nd): add module and class docstrings with math context --- nvsubquadratic/modules/grn.py | 175 ++++++++-- nvsubquadratic/modules/layer_scale.py | 150 ++++++++- nvsubquadratic/modules/masks_nd.py | 468 +++++++++++++++++++------- nvsubquadratic/modules/mlp.py | 356 ++++++++++++++++++-- 4 files changed, 967 insertions(+), 182 deletions(-) diff --git a/nvsubquadratic/modules/grn.py b/nvsubquadratic/modules/grn.py index 57523e35..5257f0cc 100644 --- a/nvsubquadratic/modules/grn.py +++ b/nvsubquadratic/modules/grn.py @@ -1,7 +1,50 @@ -"""Global Response Normalization (GRN) layer. +"""Global Response Normalisation (GRN) layer. -Promotes inter-channel feature competition via divisive normalization, -as proposed in ConvNeXt V2 (Woo et al., 2023, arXiv:2301.00808). +GRN is a channel-wise normalisation technique introduced in ConvNeXt V2 +(Woo et al., "ConvNeXt V2: Co-designing and Scaling ConvNets with Masked +Autoencoders", arXiv:2301.00808, 2023). It addresses *feature collapse* +observed when training ConvNets with masked-autoencoder pre-training: many +channels converge to identical activation patterns, so the network wastes +capacity. + +**Core operation** + +For an input ``x`` of shape ``[B, *spatial, C]`` (channels-last): + +1. Compute the per-channel global L2 norm across all spatial positions:: + + gx[b, c] = ||x[b, :, c]||_2 # shape [B, C] + +2. Normalise each channel norm by the *mean* norm across channels + (divisive / response normalisation step):: + + nx[b, c] = gx[b, c] / (mean_C(gx[b, c]) + eps) # shape [B, C] + +3. Rescale the input with learned scalar parameters ``γ`` (gamma) and + ``β`` (beta), both zero-initialised so the layer starts as identity:: + + y = γ * (x * nx) + β + x + + The ``+ x`` term is a residual connection that preserves the identity + mapping at initialisation. ``γ`` and ``β`` are 1-D tensors of size C + that broadcast over the batch and all spatial dimensions. + + Note: ``x * nx`` is equivalent to weighting each channel of ``x`` by + how strongly it activates relative to the cross-channel average — *not* + a simple pointwise division by the spatial L2 norm. The quantity ``nx`` + encodes *relative channel strength*, not normalisation to unit norm. + +**Why GRN instead of LayerNorm inside gated MLPs** + +LayerNorm standardises each token's feature vector independently, which +can suppress the *relative* differences between channels. GRN instead +promotes *inter-channel competition*: channels with strong global +activations are amplified relative to weaker channels, encouraging each +channel to specialise. This is particularly effective inside gated-linear +units (GLU / SwiGLU) where per-channel magnitude carries semantic weight. + +Reference: + Woo et al., "ConvNeXt V2", arXiv:2301.00808 (CVPR 2023). """ import torch @@ -9,19 +52,74 @@ class GlobalResponseNorm(nn.Module): - """Global Response Normalization. + """Global Response Normalisation (GRN) layer (Woo et al., arXiv:2301.00808). + + Computes a per-channel global L2 norm across all spatial positions, + normalises each channel norm by the cross-channel mean norm, then + rescales the input with learned ``γ`` / ``β`` parameters plus a + residual connection: + + .. code-block:: text + + gx = ||x||_{spatial, L2} # [B, 1, ..., 1, C] + nx = gx / (mean_C(gx) + eps) # [B, 1, ..., 1, C] + out = γ * (x * nx) + β + x # [B, *spatial, C] + + ``γ`` and ``β`` are 1-D tensors of length ``C`` that broadcast over the + batch dimension and all spatial dimensions. They are zero-initialised + so the layer starts as an identity (``out = x``) and the network can + learn to activate the normalisation only where it is beneficial. - Aggregates spatial activations per channel (L2 norm), then applies - divisive normalization across channels to promote feature diversity. - Gamma and beta are zero-initialized so the layer starts as identity. + **Inter-channel competition** - Args: - dim: Number of channels (last dimension of input). - eps: Small constant for numerical stability. + The divisive step ``gx / mean_C(gx)`` produces values > 1 for channels + whose global L2 norm exceeds the cross-channel average, and < 1 for + weaker channels. Multiplying the input by ``nx`` therefore amplifies + dominant channels and suppresses weak ones, enforcing a form of + *lateral inhibition* across the channel dimension. This is the key + mechanism by which GRN combats feature collapse (see module docstring). + + Unlike LayerNorm — which normalises each token's feature vector and + discards inter-channel magnitude differences — GRN preserves and + *amplifies* these differences, making it particularly effective inside + gated MLPs (GLU / SwiGLU) where per-channel activation strength + carries semantic weight. + + **Broadcast semantics** + + ``keepdim=True`` in the spatial reduction produces ``gx`` of shape + ``[B, 1, ..., 1, C]`` (one singleton per spatial axis). The mean is + then taken along the *channel* axis (``dim=-1, keepdim=True``) to yield + ``nx`` of the same shape. The subsequent multiplication ``x * nx`` + broadcasts over all spatial positions without an explicit tile, so GRN + is memory-efficient and agnostic to the number of spatial dimensions + (1-D sequences, 2-D images, 3-D volumes, etc.). + + Attributes: + dim (int): Number of channels C; must equal ``x.shape[-1]`` at + every forward call. + gamma (nn.Parameter): Learnable per-channel scale; shape ``(C,)``, + zero-initialised. + beta (nn.Parameter): Learnable per-channel bias; shape ``(C,)``, + zero-initialised. + eps (float): Small positive constant added to ``mean_C(gx)`` in the + denominator to prevent division by zero. + + Reference: + Woo et al., "ConvNeXt V2", arXiv:2301.00808 (CVPR 2023), + Sec. 3 "Global Response Normalization". """ def __init__(self, dim: int, eps: float = 1e-6): - """Initialize GRN with zero-initialized gamma and beta.""" + """Initialise GRN with zero-initialised gamma and beta. + + Args: + dim: Number of channels C. Must match the size of the last + dimension of every input tensor passed to ``forward``. + Determines the shape of ``gamma`` and ``beta``. + eps: Small positive constant added to ``mean_C(gx)`` in the + denominator for numerical stability. Defaults to ``1e-6``. + """ super().__init__() self.dim = dim self.gamma = nn.Parameter(torch.zeros(dim)) @@ -29,26 +127,61 @@ def __init__(self, dim: int, eps: float = 1e-6): self.eps = eps def flop_count(self, num_tokens: int) -> int: - """Count FLOPs for GRN. + """Return the approximate FLOP count for one forward pass. - Using T = num_tokens and C = dim: - - L2 norm per channel over spatial dims: T*C (square) + T*C (sum) + C (sqrt) = 2*T*C + C - - mean over channels: C (sum) - - division: C (div) - - apply to input: T*C (mult) + T*C (mult by gamma) + T*C (add beta) + T*C (residual) + Let T = ``num_tokens`` (total number of spatial positions summed + over the batch, i.e. ``B * prod(spatial_shape)``) and C = ``self.dim``. + The cost is dominated by element-wise operations over the T × C + activation grid: - Total: ~6 * T * C FLOPs + * **Squared L2 norm per channel** — element-wise square + reduction + sum over T positions per channel → 2 · T · C FLOPs. + * **Square root per channel** — C FLOPs (negligible vs T · C, + included for completeness). + * **Cross-channel mean** — C additions → C FLOPs. + * **Division** ``gx / (mean_C(gx) + eps)`` — C FLOPs. + * **Broadcast multiply** ``x * nx`` — T · C FLOPs. + * **Scale** ``γ * (x * nx)`` — T · C FLOPs. + * **Add beta and residual** — 2 · T · C FLOPs. + + Total: approximately **6 · T · C** FLOPs. + + Args: + num_tokens: Total number of spatial positions in the batch, + i.e. ``B * prod(spatial_shape)``. Note this includes the + batch dimension: for a batch of 8 images of size 32×32 the + value is ``8 * 32 * 32 = 8192``. + + Returns: + Estimated integer FLOP count for one forward pass. """ return 6 * num_tokens * self.dim def forward(self, x: torch.Tensor) -> torch.Tensor: - """Apply GRN to input tensor. + """Apply Global Response Normalisation to the input tensor. + + The input must be in **channels-last** layout: the channel axis is + the *last* dimension and all intermediate dimensions are spatial. + The batch axis is always the first dimension. Args: - x: Input of shape ``[B, *spatial, C]`` (channels-last). + x: Input activation tensor of shape ``[B, *spatial, C]``, where + + * ``B`` — batch size, + * ``*spatial`` — any number (≥ 1) of spatial dimensions + (e.g. ``(L,)`` for 1-D sequences, ``(H, W)`` for 2-D + images, ``(D, H, W)`` for 3-D volumes), + * ``C`` — number of channels; must equal ``self.dim``. Returns: - Tensor of same shape as input. + torch.Tensor: Output tensor of shape ``[B, *spatial, C]``, the + same dtype and device as ``x``, with GRN applied: + ``γ * (x * nx) + β + x``. + + Raises: + RuntimeError: If ``x.shape[-1]`` does not equal ``self.dim`` + (raised implicitly when broadcasting ``self.gamma`` / + ``self.beta`` against a mismatched channel dimension). """ spatial_dims = tuple(range(1, x.ndim - 1)) gx = torch.norm(x, p=2, dim=spatial_dims, keepdim=True) # [B, 1..., C] diff --git a/nvsubquadratic/modules/layer_scale.py b/nvsubquadratic/modules/layer_scale.py index 4690ea3d..a7250637 100644 --- a/nvsubquadratic/modules/layer_scale.py +++ b/nvsubquadratic/modules/layer_scale.py @@ -1,7 +1,36 @@ """LayerScale: learnable per-channel scaling of residual branch outputs. -Reference: Touvron et al., "Going deeper with Image Transformers" (CaiT), ICCV 2021. -Used in ViT-5 as a default component for training stability. +LayerScale (Touvron et al., "Going deeper with Image Transformers", arXiv:2103.17239, +ICCV 2021) is a lightweight technique for stabilising the training of very deep +vision transformers. In a standard residual network the update rule is + + x ← x + F(x) + +which means that the residual branch F can grow arbitrarily large at +initialisation, making depth 36+ networks difficult to train. LayerScale +modifies the update to + + x ← x + λ ⊙ F(x) + +where ``λ ∈ ℝ^C`` is a *learnable*, *per-channel* scalar vector initialised to +a small positive constant (e.g. ``1e-4``). At the start of training the gated +residual updates are therefore nearly zero, so the effective depth of the network +is small. As training proceeds ``λ`` grows, progressively incorporating the +residual branches until the full network capacity is exploited. + +The module is used by +:class:`~nvsubquadratic.modules.vit5_residual_block.ViT5ResidualBlock` where +**two independent** :class:`LayerScale` instances — one for the sequence-mixer +branch and one for the MLP branch — wrap each residual update: + +.. code-block:: text + + x = x + drop_path(ls_attn(mixer(norm(x)))) + x = x + drop_path(ls_mlp(mlp(mlp_norm(x)))) + +Reference: + Touvron, H., et al. "Going deeper with Image Transformers." + ICCV 2021. arXiv:2103.17239. """ import torch @@ -9,39 +38,124 @@ class LayerScale(nn.Module): - """Learnable diagonal scaling applied element-wise to the channel dimension. + """Learnable per-channel scalar gate for residual branch outputs. + + **Operation** + + Given an input tensor ``x`` of arbitrary leading batch / spatial dimensions + followed by a channel dimension ``C``, LayerScale computes + + output = x * γ + + where ``γ ∈ ℝ^C`` is broadcast element-wise along all axes except the last + one. Concretely, if ``x`` has shape ``(B, T, C)`` (the ViT-5 layout) then + ``γ`` is of shape ``(C,)`` and is broadcast to ``(1, 1, C)`` automatically + by PyTorch. The same broadcast rule applies to any channels-last layout: + ``(B, H, W, C)``, ``(B, T, H, W, C)``, etc. + + **Training dynamics** + + The ``init_values`` argument controls the initial magnitude of every element + of ``γ``. A small value (e.g. ``1e-4``) means the residual update is + almost entirely suppressed at the start of training, which: + + * Prevents gradient explosion in very deep networks (depth ≥ 24). + * Lets the skip connections carry most of the signal early on, and allows + the residual branches to activate gradually once they have learned useful + features. + + Using a larger ``init_values`` (e.g. ``1.0``) is appropriate when + fine-tuning from a checkpoint where the residual branches are already + well-trained and suppressing them would slow convergence. - Given input x of shape (..., dim), returns x * diag(gamma) where gamma - is a learnable vector initialized to ``init_value``. + The parameter ``γ`` is tagged with ``_no_weight_decay = True`` so that + optimiser weight-decay regularisation (L2) is **not** applied to it. This + is standard practice and matches the original CaiT training recipe. - Args: - dim: Number of channels. - init_value: Initial value for the scaling vector (typically 1e-4). + **How it differs from a plain nn.Linear / scalar gate** + + Unlike a ``nn.Linear(C, C)`` projection (which mixes channels), LayerScale + applies an independent scalar per channel. This is equivalent to a + *diagonal* linear map ``diag(γ)`` and requires only ``C`` parameters rather + than ``C²``. + + Attributes: + gamma (nn.Parameter): Learnable scale vector of shape ``(dim,)``, + initialised to ``init_values``. Tagged ``_no_weight_decay = True`` + to exclude it from L2 weight-decay in the optimiser. + + Example:: + + ls = LayerScale(dim=768, init_values=1e-4) + x = torch.randn(2, 196, 768) # [B, T, C] + out = ls(x) # [B, T, C], same shape as x """ - def __init__(self, dim: int, init_value: float = 1e-4): - """Initialize learnable scale vector to init_value.""" + def __init__(self, dim: int, init_values: float = 1e-4): + """Initialise the learnable scale vector. + + Args: + dim: Channel dimension ``C``. Determines the length of the + ``gamma`` parameter vector. Must match the channel (last) + dimension of tensors passed to :meth:`forward`. + init_values: Initial value for every element of ``gamma``. + All elements are set to this scalar at construction time. + Typical choices: + + * ``1e-4`` — recommended for training from scratch on deep + networks; effectively suppresses residual branches initially. + * ``1e-5`` — used in the original CaiT paper for the deepest + (depth-48) variants. + * ``1.0`` — effectively disables the gating at initialisation; + useful when fine-tuning from a strong pre-trained checkpoint. + """ super().__init__() - self.gamma = nn.Parameter(init_value * torch.ones(dim)) + self.gamma = nn.Parameter(init_values * torch.ones(dim)) self.gamma._no_weight_decay = True def flop_count(self, num_tokens: int) -> int: - """Count FLOPs for per-channel scaling of ``num_tokens`` token vectors. + """Count floating-point multiply operations for one forward pass. - Each token of dimension D = ``self.gamma.shape[0]`` is multiplied - element-wise by the learned gamma vector: D FLOPs per token. + Each element of ``x`` is multiplied by the corresponding element of + ``gamma``, so the total number of scalar multiplications is - Total: num_tokens * D. + FLOPs = num_tokens × dim + + where ``dim = self.gamma.shape[0]``. The broadcast of ``gamma`` is + free (no arithmetic), and the element-wise multiply is counted as one + FLOP per output element (following the convention used throughout this + codebase of counting multiply-only, not multiply-add pairs). Args: - num_tokens: Number of token vectors being scaled. + num_tokens: Number of token (or spatial) positions ``T`` in the + input tensor. The full tensor has ``T × dim`` elements, each + requiring one multiply. Returns: - Total FLOPs as an integer. + Integer FLOP count equal to ``num_tokens * dim``. """ dim = self.gamma.shape[0] return num_tokens * dim def forward(self, x: torch.Tensor) -> torch.Tensor: - """Scale input by per-channel gamma.""" + """Apply per-channel scaling to the input tensor. + + Multiplies every channel slice of ``x`` by the corresponding scalar in + ``gamma``. PyTorch broadcasts ``gamma`` of shape ``(C,)`` across all + leading dimensions of ``x`` automatically. + + Args: + x: Input tensor of shape ``(*leading_dims, C)`` where the last + dimension must equal ``dim`` passed to the constructor. + Common shapes: + + * ``(B, T, C)`` — ViT-5 / transformer token sequences. + * ``(B, H, W, C)`` — channels-last 2-D feature maps. + * ``(B, T, H, W, C)`` — channels-last 3-D (video) tensors. + + Returns: + torch.Tensor: Scaled tensor with the same shape and dtype as ``x``, + where ``output[..., c] = x[..., c] * gamma[c]`` for each channel + index ``c``. + """ return x * self.gamma diff --git a/nvsubquadratic/modules/masks_nd.py b/nvsubquadratic/modules/masks_nd.py index ed43000f..1fdebc59 100644 --- a/nvsubquadratic/modules/masks_nd.py +++ b/nvsubquadratic/modules/masks_nd.py @@ -1,9 +1,60 @@ # TODO: Add license header here -"""Modulation masks for N-dimensional data. - -These masks are used to modulate the input features of a convolutional kernel. +r"""Learnable spatial modulation masks for N-dimensional convolution kernels. + +Background +---------- +Long-range convolutional operators such as :class:`~nvsubquadratic.modules.ckconv_nd.CKConvND` +and :class:`~nvsubquadratic.modules.hyena_nd.Hyena` use implicit kernel networks that +produce a dense kernel defined over the full coordinate grid. Left unconstrained, +these kernels couple every spatial position to every other one — including pairs +that are spatially very distant — which can hurt both optimisation and +generalisation. + +The modules in this file implement **soft receptive-field windows**: differentiable +spatial masks :math:`m \in [0, 1]^{*\text{spatial} \times C}` that are multiplied +element-wise into the implicit kernel values *before* the FFT convolution step. +After masking, the effective kernel decays towards zero beyond a characteristic +spatial radius, concentrating each channel's receptive field. + +There are two families: + +* **Exponential** (:class:`ExponentialModulationND`) — fixed, non-learnable + decay rates initialised on a log-uniform ramp from slow to fast; used to + enforce a hard inductive bias without any gradient signal changing the + bandwidth. + +* **Gaussian** (:class:`GaussianModulationND`, + :class:`BlockAlignedGaussianModulationND`) — learnable standard-deviation + parameters per spatial axis and channel; the mask is a factorised product of + Gaussians, one per axis. During training the bandwidth can grow or shrink, + subject to ``[min_std, max_std]`` clamp bounds. + +How masks interact with the FFT convolution operators +------------------------------------------------------ +:class:`~nvsubquadratic.modules.ckconv_nd.CKConvND` evaluates its implicit +kernel network on a coordinate grid to obtain a dense kernel tensor +``k ∈ R^{* spatial × C}``. The modulation mask ``m`` is applied in coordinate +space *before* the FFT: + +.. code-block:: none + + k_masked = modulation(grid, k) # element-wise, coordinate space + output = fftconvNd(x, k_masked) + +Because the FFT convolution operates on the windowed kernel ``k_masked``, +the effective receptive field of the operator is bounded by the support of +the mask. Narrow masks (small std / fast decay) correspond to local +operators; wide masks (large std / slow decay) approach global convolutions. + +The convention used throughout this module is: + +* The coordinate grid has values in **[−1, 1]** along each axis (center = 0). +* Mask values are in **[0, 1]** where **1 = fully included** (center of the + kernel, near zero displacement) and **0 = fully excluded** (large + displacement, far from the center of the convolution kernel). Masking thus + *suppresses* the kernel at large displacements, not at zero displacement. For testing: PYTHONPATH=. python nvsubquadratic/modules/masks_nd.py @@ -31,6 +82,21 @@ def _normalize_init_extent(init_extent: float | Sequence[float] | None, data_dim axis. The clamp ``[min_std, max_std]`` enforces feasibility, so very large values simply saturate the entire ramp at ``max_std`` (= "this axis is essentially unmasked at init"). + + Args: + init_extent: A positive float (broadcast to all axes), a sequence of + ``data_dim`` positive floats (one per axis), or ``None`` (treated + as ``1.0`` on every axis). + data_dim: Number of spatial/temporal dimensions. Determines the + expected length of a sequence ``init_extent``. + + Returns: + A tuple of ``data_dim`` positive finite floats. + + Raises: + TypeError: If ``init_extent`` is a bool, or not a float or sequence. + ValueError: If ``init_extent`` is a sequence with length != ``data_dim``, + or if any element is non-positive or non-finite. """ if init_extent is None: init_extent = 1.0 @@ -51,16 +117,53 @@ def _normalize_init_extent(init_extent: float | Sequence[float] | None, data_dim class ExponentialModulationND(torch.nn.Module): - """Applies exponential decay modulation to input features. + r"""Fixed exponential-decay spatial window applied to implicit convolutional kernels. + + Geometry + -------- + Given a coordinate grid normalised to ``[−1, 1]`` (center = 0) and a + set of per-axis, per-channel decay rates ``w_{d,c} > 0``, the mask at + spatial position :math:`p = (p_0, \ldots, p_{D-1})` for channel ``c`` is: + + .. math:: + + m_c(p) = \prod_{d=0}^{D-1} \exp\!\bigl(-\lvert p_d \rvert \cdot \lvert w_{d,c} \rvert\bigr) + + All values lie in ``(0, 1]``. The mask equals **1** at the origin + (displacement 0) and decays towards **0** as any coordinate moves away from + the center. Channels with large ``w`` decay quickly (narrow receptive + field); channels with small ``w`` decay slowly (broad receptive field). + + The ND generalisation follows automatically from the product structure: + for a 2D image grid the mask is a 2D tent surface; for a 3D volume it + is a 3D "tent" shaped object. - This module modulates input features by applying an exponential decay function on each dimension of an N-dimensional input. - The decay rates are parameterized by a set of learned decay rates. + Decay rates are **not learnable** (they are registered as a parameter so + they travel with the module and appear in ``state_dict``, but they are + marked ``_no_weight_decay = True`` and are not updated by the optimizer). + The rates are initialised on a linear ramp from ``slow_decay_pct`` to + ``fast_decay_pct``, divided by ``data_dim`` so that the product across + axes has a consistent magnitude regardless of the number of dimensions. + + Role in CKConvND + ---------------- + :class:`~nvsubquadratic.modules.ckconv_nd.CKConvND` optionally passes the + output of its implicit kernel network through this module before the FFT + convolution step. The resulting masked kernel is then convolved with the + input signal via ``fftconvNd``. Args: - data_dim (int): Dimension of input data (1D for sequences, 2D for images, 3D for videos, etc.). - num_channels (int): Number of input channels to be modulated. - fast_decay_pct (float, optional): Percentage for the fastest decay rate. Default is 13.81. - slow_decay_pct (float, optional): Percentage for the slowest decay rate. Default is 2.3. + data_dim: Number of spatial/temporal dimensions (1 for sequences, 2 for + images, 3 for videos). + num_channels: Number of feature channels ``C``. Each channel receives + a distinct decay rate. + fast_decay_pct: Upper end of the decay-rate ramp (fastest / narrowest + channel). Default ``13.81`` (≈ :math:`\ln(10^6)`, so the + narrowest channel decays to near zero within a small fraction of + the grid). + slow_decay_pct: Lower end of the decay-rate ramp (slowest / broadest + channel). Default ``2.3`` (≈ :math:`\ln(10)`, so the broadest + channel retains ≈ 10 % of its value at the grid boundary). """ def __init__( @@ -70,13 +173,13 @@ def __init__( fast_decay_pct: float = 13.81, slow_decay_pct: float = 2.3, ): - """Initialize the ExponentialModulationND class. + """Initialise the exponential modulation module. Args: - data_dim: Dimension of input data. - num_channels: Number of input channels to be modulated. - fast_decay_pct: Percentage for the fastest decay rate. - slow_decay_pct: Percentage for the slowest decay rate. + data_dim: Number of spatial/temporal dimensions. + num_channels: Number of feature channels to modulate. + fast_decay_pct: Upper end of the per-channel decay-rate ramp (fastest channel). + slow_decay_pct: Lower end of the per-channel decay-rate ramp (slowest channel). """ super().__init__() self.data_dim = data_dim @@ -99,18 +202,33 @@ def extra_repr(self): return f"data_dim={self.data_dim}, num_channels={self.num_channels}, fast_decay_pct={self.fast_decay_pct}, slow_decay_pct={self.slow_decay_pct}" def forward(self, grid: torch.Tensor, x: torch.Tensor) -> torch.Tensor: - """Applies exponential modulation to the input tensor `x` based on the coordinates in `grid`. + r"""Apply exponential decay modulation element-wise to kernel features. + + For each spatial position :math:`p` and channel :math:`c` computes: + + .. math:: + + \text{out}[\ldots, c] = x[\ldots, c] \cdot + \prod_{d} \exp\!\bigl(-\lvert p_d \rvert \cdot \lvert w_{d,c} \rvert\bigr) + + The product is over all ``data_dim`` spatial axes. The mask value is + **1** at the origin and decreases monotonically towards 0 as the + displacement from the origin grows. Args: - grid (torch.Tensor): A tensor representing grid values (shape: [1, * spatial_dims, data_dim]). - Must have dtype `torch.float32`. - x (torch.Tensor): Input features to be modulated (shape: [batch_size, * spatial_dims, num_channels]). + grid: Coordinate grid of shape ``[1, *spatial_dims, data_dim]`` with + values in ``[−1, 1]``. Each entry ``grid[..., d]`` contains the + normalised coordinate along axis ``d``. Must be ``torch.float32`` + (lower precision collapses nearby coordinates together). + x: Kernel feature tensor of shape ``[B, *spatial_dims, num_channels]`` + to be modulated. ``B`` is the batch size; ``*spatial_dims`` must + match the spatial shape of ``grid``. Returns: - torch.Tensor: The modulated input features (same shape as `x`). + torch.Tensor: Modulated features with the same shape and dtype as ``x``. Raises: - AssertionError: If `grid` is not of type `torch.float32`. + AssertionError: If ``grid.dtype`` is not ``torch.float32``. """ # Ensure the grid tensor has the correct data type assert grid.dtype == torch.float32, ( @@ -134,10 +252,26 @@ def _std_from_attenuation(attenuation: float, position: float, data_dim: int) -> mask = exp(-0.5 * data_dim * (position / σ)²) = attenuation ⟹ σ = position * sqrt( -data_dim / (2 * ln(attenuation)) ) + This helper is used during initialisation of :class:`GaussianModulationND` + to derive ``min_std`` and ``max_std`` from user-specified attenuation targets. + All calls from that class pass ``data_dim=1`` because the attenuation targets + are defined as **single-axis** (1D) measurements; the product structure of + the full ND mask then gives an ND corner value of ``attenuation ** data_dim``. + Args: - attenuation: Desired mask value (0 < attenuation < 1). - position: Absolute grid coordinate (> 0). - data_dim: Number of spatial dimensions. + attenuation: Desired mask value at ``position`` (must be in ``(0, 1)``). + position: Absolute normalised grid coordinate at which the attenuation + target is measured (must be ``> 0``). + data_dim: Number of spatial dimensions assumed in the formula. Pass + ``1`` for single-axis targets (the typical case in this module). + + Returns: + float: The standard deviation ``σ > 0`` such that the Gaussian mask + equals ``attenuation`` at the given ``position``. + + Raises: + AssertionError: If ``attenuation`` is not in ``(0, 1)`` or ``position`` + is not ``> 0``. """ assert 0.0 < attenuation < 1.0, f"attenuation must be in (0, 1), got {attenuation}" assert position > 0.0, f"position must be > 0, got {position}" @@ -145,76 +279,104 @@ def _std_from_attenuation(attenuation: float, position: float, data_dim: int) -> class GaussianModulationND(torch.nn.Module): - """Gaussian decay modulation across N spatial/temporal dimensions. - - For each data dimension d and channel c we learn a (positive) standard deviation sigma_{d,c}. - Given a coordinate grid (centered around 0) we apply: - - mask_{..., c} = Π_d exp( - 0.5 * (grid_d / sigma_{d,c})^2 ) - - which is then multiplied elementwise with the input features. - - Mean is fixed (no learnable shift) so modulation remains symmetric around zero. - - **Initialization** — pass ``min_attenuation_at_step`` and - ``max_attenuation_at_limit`` (plus ``grid_size``, auto-injected by - CKConvND). These define the **clamp bounds** — the narrowest any - channel can get (``min_std``) and the widest (``max_std``). Optionally - pass ``init_extent`` to control the initial bandwidth scale **per - axis**. - - All attenuation values are **single-axis** (1D) measurements. Since - the mask is a product of per-dimension Gaussians, the 2D corner value - is ``attenuation ** 2``, and the 3D corner is ``attenuation ** 3``, - etc. - - - ``min_attenuation_at_step`` — 1D mask value at the first grid step - from center for the narrowest *possible* channel. Sets ``min_std`` - and the **reference** ``init_std_low`` (narrowest channel starts at - the clamp bound when ``init_extent = 1``). - - ``max_attenuation_at_limit`` — 1D mask value at the grid boundary - (position 1) for the widest *possible* channel. Sets ``max_std``. - - ``init_extent`` — per-axis bandwidth scale that multiplicatively - scales **both** ends of the per-axis logspace ramp: - ``init_std_low[d] = clamp(min_std * extent[d], min_std, max_std)`` - ``init_std_high[d] = clamp(init_std_high_unit * extent[d], min_std, max_std)`` - where ``init_std_high_unit ≈ 0.4724`` is the std at which a 1D - Gaussian reaches ``0.1`` at position 1. Pass a float (broadcast to - all axes) or a sequence of length ``data_dim``. Values must be - strictly ``> 0``; defaults to ``1.0`` on every axis (recovers the - reference ramp ``[min_std, init_std_high_unit]``). - - Examples on an anisotropic ``L_cache=(8, 64, 64)`` cube cache - (mask grid_size = 127): - - * ``init_extent = 1.0`` — all axes use the reference ramp from - ``min_std ≈ 0.0075`` to ``init_std_high ≈ 0.4724``. On the - depth axis the bottom of the ramp is unusably narrow (mask ≈ 0 - across depth for early channels). - * ``init_extent = (1.0, 0.25, 0.25)`` — depth uses the reference - ramp; H/W are 4× narrower at both ends (extreme localization). - * ``init_extent = (max_std/min_std, 1.0, 1.0)`` (≈ ``(416, 1, 1)`` - at defaults) — depth ramp saturates at ``max_std`` end-to-end; - every depth channel is initialised at the widest possible - Gaussian, i.e. depth axis is essentially unmasked at init. - Useful when ``L_cache_d`` is short and depth-axis frequency - content is naturally bounded by the small kernel grid. - - Only **initialization** is per-axis — ``min_std`` and ``max_std`` - (clamp bounds) remain scalar and shared across axes. + r"""Learnable Gaussian-window spatial mask for ND convolutional kernels. + + Geometry + -------- + For a coordinate grid normalised to ``[−1, 1]`` (center = 0) and + per-axis, per-channel standard deviations :math:`\sigma_{d,c} > 0`, the + mask value at position :math:`p = (p_0, \ldots, p_{D-1})` for channel + ``c`` is the product of per-axis Gaussians: + + .. math:: + + m_c(p) = \prod_{d=0}^{D-1} + \exp\!\Bigl(-\tfrac{1}{2}\bigl(p_d / \sigma_{d,c}\bigr)^2\Bigr) + + All values lie in ``(0, 1]``. The mask equals **1** at the origin + and decays symmetrically in all directions; the level set at value ``v`` + is an axis-aligned ellipsoid with semi-axes :math:`\sigma_{d,c}\sqrt{-2\ln v}`. + + **1 = fully included, 0 = fully excluded.** A narrow Gaussian (small + ``σ``) concentrates the effective kernel around the origin, making the + operator local; a wide Gaussian (large ``σ``) lets the full grid + contribute, approaching a global convolution. + + ND generalisation + ----------------- + The mask factorises over axes. In 2D the mask surface looks like a 2D + Gaussian bell (not a sphere — each axis has an independent ``σ``). In + 3D it is a trivariate axis-aligned Gaussian. Because the mask is a + *product*, the corner value at position ``(position, position, …, + position)`` is the product of the individual per-axis Gaussian values, + which equals ``single_axis_mask_value ** data_dim``. The attenuation + parameters (``min_attenuation_at_step``, ``max_attenuation_at_limit``) + are therefore defined as **single-axis (1D) measurements** — the + effective ND attenuation at the grid corner is stricter by a factor of + ``data_dim`` in the exponent. + + Parametrisation and clamping + ---------------------------- + The learned parameter ``std_param`` of shape ``[data_dim, num_channels]`` + stores raw values that are mapped to strictly-positive std values via + ``parametrization``: + + * ``'direct'`` — ``std_param`` IS the std; a ``register_forward_pre_hook`` + clamps it into ``[min_std, max_std]`` in-place (``torch.no_grad``), so + gradients at the boundary are preserved through the activation. + * ``'log'`` — ``std = exp(std_param)``; hard clamp applied after (breaks + boundary gradients — see inline warning). + * ``'softplus'`` — ``std = softplus(std_param)``; hard clamp applied after. + + Initialisation + -------------- + ``min_attenuation_at_step`` and ``max_attenuation_at_limit`` define the + **clamp bounds** ``[min_std, max_std]``. The initial ``std_param`` is a + logspace ramp from ``min_std`` to ``init_std_high_unit`` on every axis, + scaled per-axis by ``init_extent``: + + * ``init_std_low[d] = clamp(min_std * extent[d], min_std, max_std)`` + * ``init_std_high[d] = clamp(init_std_high_unit * extent[d], min_std, max_std)`` + + where ``init_std_high_unit ≈ 0.4724`` is the std at which a 1D Gaussian + reaches ``0.1`` at position 1. See ``init_extent`` below. + + All attenuation values are **single-axis (1D)** measurements; see the ND + generalisation note above. Args: - data_dim: Number of spatial/temporal dimensions. - num_channels: Number of feature channels to modulate. - min_attenuation_at_step: 1D mask value at first grid step (sets clamp - lower bound and the reference init lower bound). - max_attenuation_at_limit: 1D mask value at grid boundary (sets clamp - upper bound). + data_dim: Number of spatial/temporal dimensions (1 for sequences, 2 + for images, 3 for volumes). + num_channels: Number of feature channels ``C`` to modulate. + grid_size: Number of grid points per spatial dimension. Used to + compute the size of the smallest grid step + (``min_step = 2 / (grid_size - 1)``), which sets ``min_std``. + Auto-injected by :class:`~nvsubquadratic.modules.ckconv_nd.CKConvND`. + min_attenuation_at_step: Target 1D mask value at the first grid step + from the origin for the **narrowest** channel. Smaller values + → narrower minimum std → more local minimum channel. Default + ``0.1``. + max_attenuation_at_limit: Target 1D mask value at the grid boundary + (``position = 1``) for the **widest** channel. Larger values + → wider maximum std → less attenuation at the boundary. Default + ``0.95``. init_extent: Scalar or per-axis sequence controlling the initial - bandwidth scale on each axis. Strictly ``> 0``; default - ``1.0`` (reference ramp on every axis). - grid_size: Kernel grid points per dimension. Auto-injected by CKConvND. - parametrization: ``'log'``, ``'softplus'``, or ``'direct'``. + bandwidth scale on each axis. Must be strictly ``> 0``; defaults + to ``1.0`` on every axis (reference ramp on every axis). + + Examples for an anisotropic ``L_cache = (8, 64, 64)`` cache + with ``grid_size = 127``: + + * ``init_extent = 1.0`` — all axes use the reference ramp. On a + short depth axis the bottom of the ramp can be unusably narrow. + * ``init_extent = (max_std/min_std, 1.0, 1.0)`` — depth ramp + saturates at ``max_std`` (axis effectively unmasked at init). + * ``init_extent = (1.0, 0.25, 0.25)`` — H/W are 4× narrower than + the reference (extreme localisation on spatial axes). + + parametrization: One of ``'direct'``, ``'log'``, ``'softplus'``. + Controls the mapping from ``std_param`` to std values. Default + ``'direct'``. """ def __init__( @@ -227,7 +389,21 @@ def __init__( init_extent: float | Sequence[float] = 1.0, parametrization: str = "direct", ): - """Initialize the GaussianModulationND class.""" + """Initialise the Gaussian modulation module. + + Args: + data_dim: Number of spatial/temporal dimensions. + num_channels: Number of feature channels to modulate. + grid_size: Number of grid points per spatial dimension. + min_attenuation_at_step: 1D mask value at the first grid step from + the origin for the narrowest channel (sets ``min_std``). + max_attenuation_at_limit: 1D mask value at the grid boundary for the + widest channel (sets ``max_std``). + init_extent: Per-axis bandwidth scale for initialisation (> 0). + Pass a float to broadcast, or a sequence of length ``data_dim``. + parametrization: Mapping from ``std_param`` to std values. + One of ``'direct'``, ``'log'``, ``'softplus'``. + """ super().__init__() assert parametrization in {"log", "softplus", "direct"}, ( "parametrization must be 'log' or 'softplus' or 'direct'" @@ -300,10 +476,18 @@ def _clamp_direct_std_param_pre_hook(self, module, inputs): self.std_param.data.clamp_(min=self.min_std, max=self.max_std) def _compute_std(self) -> torch.Tensor: - """Computes the standard deviation for each channel based on the learned weights. + """Map the raw ``std_param`` to strictly-positive standard deviations. + + Applies the parametrization-specific mapping and clamps to + ``[min_std, max_std]``. For the ``'direct'`` parametrization the + clamp is applied by a pre-hook, so gradients at the boundary are + preserved; for ``'log'`` and ``'softplus'`` the clamp is applied + here (which zeros gradients at the boundary). Returns: - torch.Tensor: The standard deviation for each channel (shape: [data_dim, num_channels]). + torch.Tensor: Standard deviations of shape ``[data_dim, num_channels]`` + in ``float32``. Values are guaranteed to lie in + ``[min_std, max_std]``. """ std = self.std_param.float() # [data_dim, num_channels] if self.parametrization == "direct": @@ -322,7 +506,16 @@ def _compute_std(self) -> torch.Tensor: @staticmethod def _mask_value(std: float, position: float) -> float: - """1D Gaussian mask value: exp(-0.5 * (position / std)^2).""" + """1D Gaussian mask value: exp(-0.5 * (position / std)^2). + + Args: + std: Standard deviation of the Gaussian (> 0). + position: Absolute normalised grid coordinate at which to evaluate + the mask. + + Returns: + float: Mask value in ``(0, 1]``. + """ return math.exp(-0.5 * (position / std) ** 2) def extra_repr(self): @@ -349,18 +542,34 @@ def extra_repr(self): ) def forward(self, grid: torch.Tensor, x: torch.Tensor) -> torch.Tensor: - """Applies Gaussian modulation to the input tensor `x` based on the coordinates in `grid`. + r"""Apply Gaussian spatial modulation element-wise to kernel features. + + For each spatial position :math:`p` and channel :math:`c` computes: + + .. math:: + + \text{out}[\ldots, c] = x[\ldots, c] \cdot + \exp\!\Bigl(-\tfrac{1}{2} \sum_{d} \bigl(p_d / \sigma_{d,c}\bigr)^2\Bigr) + + The exponent is computed as a single einsum (sum over axes) for + efficiency, avoiding the intermediate ``prod(exp)`` formulation. + + Mask convention: **1 = fully included** (origin), **0 = fully + excluded** (large displacement). Args: - grid (torch.Tensor): A tensor representing grid values (shape: [1, * spatial_dims, data_dim]). - Must have dtype `torch.float32`. - x (torch.Tensor): Input features to be modulated (shape: [batch_size, * spatial_dims, num_channels]). + grid: Coordinate grid of shape ``[1, *spatial_dims, data_dim]`` with + values in ``[−1, 1]``. Must be ``torch.float32``. + x: Kernel feature tensor of shape ``[B, *spatial_dims, num_channels]``. + ``*spatial_dims`` must match the spatial shape of ``grid``. Returns: - torch.Tensor: The modulated input features (same shape as `x`). + torch.Tensor: Modulated features with the same shape and dtype as + ``x``. The internal Gaussian computation is always done in + ``float32`` and then cast to ``x.dtype``. Raises: - AssertionError: If `grid` is not of type `torch.float32`. + AssertionError: If ``grid.dtype`` is not ``torch.float32``. """ # Ensure the grid tensor has the correct data type assert grid.dtype == torch.float32, f"grid must be float32. Current dtype: {grid.dtype}" @@ -379,28 +588,51 @@ def forward(self, grid: torch.Tensor, x: torch.Tensor) -> torch.Tensor: class BlockAlignedGaussianModulationND(GaussianModulationND): - """Gaussian modulation with channel-reversed std_param for block-structured SIRENs. + r"""Gaussian modulation with channel-reversed std ordering for block-structured SIRENs. - The parent :class:`GaussianModulationND` initializes ``std_param`` so that - channel 0 has the narrowest Gaussian (``min_std``) and the last channel has - the widest (``init_std_high``). That ordering assumes the *narrow* mask - (short spatial support → broad spectral support) should be applied to - channels that carry high-frequency content. + Motivation + ---------- + :class:`GaussianModulationND` initialises ``std_param`` so that channel 0 + has the **narrowest** Gaussian (smallest ``σ``, shortest spatial support, + broadest spectral support) and the last channel has the **widest** Gaussian. + This ordering is natural when channel index 0 carries high-frequency content, + which should be localised in space. Block-structured SIRENs such as :class:`~nvsubquadratic.modules.kernels_nd.BlockDiagonalMultiOmegaSIRENKernelND` - with a ``linear`` or ``log`` schedule put the *lowest* ω₀ (low-frequency - content) on the first block, so the natural alignment is the opposite: - widest Gaussian on channel 0, narrowest on the last channel. - - This subclass just reverses ``std_param`` along the channel axis after the - parent's initialization — no other behaviour changes (forward pass, - clamping, parametrization, grad flow are all inherited unchanged). + with a ``'linear'`` or ``'log'`` frequency schedule assign the **lowest** + :math:`\omega_0` (low-frequency content) to the **first** block. Low + frequencies have long spatial support, so the natural alignment is the + opposite: **widest** Gaussian on channel 0 (lowest :math:`\omega_0`), + **narrowest** on the last channel (highest :math:`\omega_0`). + + Implementation + -------------- + This subclass reverses ``std_param`` along the channel axis (``dim=-1``) + immediately after the parent's ``__init__``. All other behaviour — + forward pass, pre-hook clamping, parametrisation, gradient flow — is + inherited unchanged from :class:`GaussianModulationND`. + + The channel ordering after reversal: + + * Channel 0: widest Gaussian (largest ``σ``), longest spatial support, + lowest effective frequency → matched to the lowest-:math:`\omega_0` block. + * Channel ``C-1``: narrowest Gaussian (smallest ``σ``), shortest spatial + support, highest effective frequency → matched to the highest-:math:`\omega_0` + block. Args: - data_dim, num_channels, grid_size, min_attenuation_at_step, - max_attenuation_at_limit, init_extent, parametrization: - Passed straight through to :class:`GaussianModulationND`. + data_dim: Number of spatial/temporal dimensions. + num_channels: Number of feature channels ``C`` to modulate. + grid_size: Number of grid points per spatial dimension. + min_attenuation_at_step: 1D mask value at the first grid step (sets + ``min_std`` clamp bound). Default ``0.1``. + max_attenuation_at_limit: 1D mask value at the grid boundary (sets + ``max_std`` clamp bound). Default ``0.95``. + init_extent: Per-axis bandwidth scale for initialisation (> 0). + Default ``1.0``. + parametrization: One of ``'direct'``, ``'log'``, ``'softplus'``. + Default ``'direct'``. """ def __init__( @@ -413,7 +645,7 @@ def __init__( init_extent: float = 1.0, parametrization: str = "direct", ): - """Initialize the block-aligned Gaussian mask; see the class docstring for argument semantics.""" + """Initialise the block-aligned Gaussian mask; see the class docstring for argument semantics.""" super().__init__( data_dim=data_dim, num_channels=num_channels, diff --git a/nvsubquadratic/modules/mlp.py b/nvsubquadratic/modules/mlp.py index b7db2f03..6a4cc4e1 100644 --- a/nvsubquadratic/modules/mlp.py +++ b/nvsubquadratic/modules/mlp.py @@ -1,15 +1,107 @@ # TODO: Add license header here -"""MLP implementation for ND signals. - -Supports two backends: - -* ``"torch"`` — pure PyTorch (``nn.Linear`` → activation → ``nn.Linear``). -* ``"quack"`` — QuACK fused GEMM+activation kernels (Hopper/Blackwell only). - Requires ``quack-kernels >= 0.3.0``, ``bias=False``, dims divisible by 8, - and a supported activation (``glu``, ``swiglu``). Raises at init if any - constraint is violated. +"""Channel-mixing MLP block for ND residual networks. + +Within the residual block architecture (see +:mod:`nvsubquadratic.modules.residual_block`), every repeating unit has two +complementary branches: + +1. **Sequence mixer** — captures long-range *spatial/temporal* interactions + (Hyena, Attention, CKConv, Mamba, …). +2. **MLP** *(this module)* — performs *point-wise channel mixing*, refining + each position's feature vector independently of its neighbours. + +The :class:`MLP` class provides a two-layer feed-forward network that acts +as the channel-mixing branch. Three activation variants are offered: + +* **Plain GELU / ReLU / SiLU MLP** (``activation in {"gelu", "relu", "silu"}``): + + .. code-block:: text + + y = W₂(act(W₁(x))) + + ``W₁ : C → H``, ``W₂ : H → C`` where ``H = floor(expansion_factor × C)``. + +* **Gated Linear Unit (GLU)** (``activation="glu"``): + + .. code-block:: text + + y = W₂(sigmoid(W₁ₐ(x)) ⊙ W₁ᵦ(x)) + + ``W₁`` projects to ``2H``, then is split at the midpoint into a gate half + and a value half. ``W₂ : H → C``. + +* **SwiGLU** (``activation="swiglu"``, the default in many modern configs): + + .. code-block:: text + + y = W₂(SiLU(W₁ₐ(x)) ⊙ W₁ᵦ(x)) + + Same gated structure as GLU but with SiLU in place of sigmoid; follows + Noam Shazeer's SwiGLU paper (arXiv:2002.05202). + +Expansion-ratio semantics +-------------------------- +``expansion_factor`` controls the width of the hidden (intermediate) layer +relative to the model dimension ``C``: + + H = floor(expansion_factor * C) + +For **plain** activations (GELU, ReLU, SiLU) the total linear-layer parameter +count is ``C*H + H*C = 2*C*H``. For **gated** activations (GLU, SwiGLU) the +first projection doubles to ``C → 2H``, so the count becomes +``C*2H + H*C = 3*C*H``. To keep the parameter budget comparable between +plain and gated MLPs, gated configs are typically paired with a smaller +``expansion_factor``: + +- Plain GELU at ``expansion_factor=4`` ≈ ``8C²`` parameters. +- SwiGLU at ``expansion_factor=8/3`` ≈ ``8C²`` parameters. +- A commonly used approximation is ``expansion_factor ≈ 4/3`` relative to the + plain target. + +Activation function summary +---------------------------- ++----------+----------------------------------------------+----------------------------+ +| Name | Formula | Notes | ++==========+==============================================+============================+ +| ``relu`` | ``max(0, x)`` | Plain; sparse activations | ++----------+----------------------------------------------+----------------------------+ +| ``gelu`` | ``x · Φ(x)`` (tanh approximation) | Plain; smooth, common in | +| | | Transformers / BERT-family | ++----------+----------------------------------------------+----------------------------+ +| ``silu`` | ``x · σ(x)`` | Plain; also called Swish | ++----------+----------------------------------------------+----------------------------+ +| ``glu`` | ``sigmoid(a) ⊙ b`` | Gated; ``W₁`` → 2H | ++----------+----------------------------------------------+----------------------------+ +| ``swiglu``| ``SiLU(a) ⊙ b`` | Gated; recommended modern | +| | | default (arXiv:2002.05202) | ++----------+----------------------------------------------+----------------------------+ + +Hardware back-ends +------------------ +By default (``backend="torch"``) the MLP uses standard ``nn.Linear`` layers +with a PyTorch activation. On Hopper/Blackwell GPUs with QuACK kernels +installed, setting ``backend="quack"`` enables fused GEMM+activation kernels +that reduce memory traffic. QuACK is **experimental** (forward correctness +verified; backward and benchmark still pending) and is currently blocked by a +``NotImplementedError`` in :func:`_validate_quack_backend`. + +GRN integration +--------------- +The :class:`MLP` itself does not embed a GRN layer; if Global Response +Normalization (see :mod:`nvsubquadratic.modules.grn`) is desired it should be +inserted between the two linear layers by the calling code or by wrapping +:class:`MLP` in a subclass. GRN is particularly effective inside gated-linear +units (GLU / SwiGLU) where per-channel magnitude carries semantic weight. + +Tensor layout +------------- +All tensors use **channels-last** layout: ``(B, *spatial_dims, C)``, where +``B`` is the batch size, ``*spatial_dims`` are one or more spatial axes +(length-1 for 1-D sequences, ``(H, W)`` for 2-D images, etc.), and ``C`` is +the channel dimension. The MLP is spatially agnostic — it applies the same +two-layer projection independently to every position. """ from typing import Callable, Literal @@ -77,6 +169,15 @@ def _split_to_interleaved(weight: torch.Tensor) -> torch.Tensor: This rearranges the *rows* of a ``(2N, D)`` weight matrix so that the GEMM output lands in QuACK's expected layout. Uses ``torch.stack`` + ``reshape`` so autograd can differentiate through it. + + Args: + weight: Weight matrix of shape ``(2N, D)`` in PyTorch split layout, + where rows ``[:N]`` are the gate weights and rows ``[N:]`` are the + value weights. + + Returns: + torch.Tensor: Weight matrix of the same shape ``(2N, D)`` re-ordered + into QuACK's interleaved layout (even rows = value, odd rows = gate). """ N = weight.shape[0] // 2 # QuACK interleaved layout: even rows = value, odd rows = gate. @@ -89,7 +190,27 @@ def _quack_mlp_forward( weight2: torch.Tensor, activation: str, ) -> torch.Tensor: - """Dispatch to the appropriate QuACK fused kernel (gated vs non-gated).""" + """Dispatch to the appropriate QuACK fused kernel (gated vs non-gated). + + Selects between the gated (GLU/SwiGLU) and non-gated (GELU/ReLU) QuACK + kernel families, converts weights to QuACK's layout if necessary, and + executes the fused forward pass. Intermediate pre-activations are stored + only when ``torch.is_grad_enabled()`` so that backward can reuse them. + + Args: + x: Input tensor of shape ``(B, *spatial_dims, C)`` (channels-last). + Must have the channel dimension divisible by 8 (QuACK constraint). + weight1: First linear layer weight of shape ``(H * glu_factor, C)`` + where ``H`` is the hidden dimension and ``glu_factor`` is 2 for + gated activations (GLU / SwiGLU) and 1 otherwise. + weight2: Second linear layer weight of shape ``(C, H)``. + activation: One of ``"glu"``, ``"swiglu"``, ``"gelu"``, ``"relu"``. + Must be present in :data:`_QUACK_ACT_MAP`. + + Returns: + torch.Tensor: Output tensor of shape ``(B, *spatial_dims, C)``, + matching the input shape. + """ quack_act = _QUACK_ACT_MAP[activation] if activation in _QUACK_GATED_ACTS: w1 = _split_to_interleaved(weight1) @@ -119,7 +240,36 @@ def _validate_quack_backend( dim: int, hidden_dim: int, ) -> None: - """Raise ``ValueError`` at init time if QuACK constraints are not met.""" + """Raise ``ValueError`` at init time if QuACK constraints are not met. + + Called during :class:`MLP.__init__` when ``backend="quack"`` is requested. + Fails fast with an actionable error message rather than letting a misuse + surface as an obscure CUDA error at first forward call. + + .. note:: + This function currently raises :exc:`NotImplementedError` unconditionally + because the QuACK backend is experimental (backward pass and benchmark + are not yet validated). Use ``backend="torch"`` for all production use. + + Args: + activation: Requested activation name. Must be a key in + :data:`_QUACK_ACT_MAP` (``"glu"``, ``"swiglu"``, ``"gelu"``, + ``"relu"``). + bias: Whether bias terms are enabled. QuACK fused GEMM kernels do not + support bias, so ``bias=True`` raises :exc:`ValueError`. + dim: Input / output channel dimension ``C``. Must be divisible by 8. + hidden_dim: Expanded hidden dimension ``H = expansion_factor * C``. + Must be divisible by 8. + + Raises: + NotImplementedError: Always, because the QuACK backend is not yet + production-ready. + ValueError: If ``quack-kernels`` is not installed or is below + :data:`_QUACK_MLP_MIN_VERSION`, if ``activation`` is not supported, + if ``bias=True``, or if ``dim`` / ``hidden_dim`` are not divisible + by 8. (These checks are unreachable until the + :exc:`NotImplementedError` is removed.) + """ raise NotImplementedError( "MLP backend='quack' is experimental and needs more testing " "(forward correctness verified, backward + benchmark pending). " @@ -156,17 +306,79 @@ def _validate_quack_backend( class MLP(nn.Module): - """Two-layer MLP with optional QuACK fused GEMM+activation backend. + """Point-wise two-layer MLP — the channel-mixing branch of each residual block. - Args: - dim: Input and output dimension. - activation: Activation function ('relu', 'gelu', 'silu', 'glu', 'swiglu'). - dropout_cfg: LazyConfig for the dropout layer. - expansion_factor: Hidden-dim multiplier. - bias: Whether to use bias in linear layers. - backend: ``"torch"`` for pure PyTorch, ``"quack"`` for fused kernels. - init_method_in: Optional init for first layer weights. - init_method_out: Optional init for second layer weights. + Acts on each spatial position independently (no cross-position interaction), + expanding the channel dimension by ``expansion_factor``, applying a + non-linearity, and projecting back: + + **Plain MLP** (``activation in {"relu", "gelu", "silu"}``): + + .. code-block:: text + + y = W₂( act( W₁(x) ) ) + + where ``W₁ : C → H``, ``W₂ : H → C``, and ``H = floor(expansion_factor × C)``. + + **Gated variants** (``activation in {"glu", "swiglu"}``): + + .. code-block:: text + + [a, b] = W₁(x) # W₁ : C → 2H, split at midpoint + y = W₂( gate(a) ⊙ b ) # W₂ : H → C + + - *GLU* (``"glu"``): ``gate(a) = sigmoid(a)`` + - *SwiGLU* (``"swiglu"``): ``gate(a) = SiLU(a)`` + + Because the first projection ``W₁`` must produce ``2H`` outputs for the + gate, ``layer1`` has shape ``(2H, C)`` when a gated activation is used + (see ``glu_factor`` in ``__init__``), while ``layer2`` remains ``(C, H)``. + This means the **parameter count differs** from the plain variant: + + - Plain: ``C*H + H*C = 2CH`` parameters in the two linear layers. + - Gated: ``C*2H + H*C = 3CH`` parameters in the two linear layers. + + At ``expansion_factor=2``: + + - Plain (GELU): ``H = 2C``, params ≈ ``4C²``. + - SwiGLU: ``H = 2C``, params ≈ ``6C²``. + + To keep parameter counts comparable, gated configs often use a smaller + ``expansion_factor`` (e.g. ``4/3`` rather than ``2``). + + Dropout is inserted between the two layers (applied *after* the + activation and gate product). + + The ``backend`` argument selects the compute kernel family: + + - ``"torch"`` (default): Standard ``nn.Linear`` + PyTorch activation. + Works everywhere. + - ``"quack"`` (*experimental, currently blocked*): QuACK fused GEMM + + activation kernels targeting Hopper / Blackwell GPUs. Requires + ``quack-kernels >= 0.3.0``, ``bias=False``, channel dimensions divisible + by 8, and a supported activation. Currently raises + :exc:`NotImplementedError` at init time pending backward validation. + + Attributes: + hidden_dim (int): Expanded hidden channel dimension + ``H = floor(expansion_factor * dim)``. For gated variants + this is the *post-gate* width — the actual ``layer1`` output + width is ``2 * hidden_dim``. + activation (str): Name of the activation function in use. + One of ``"relu"``, ``"gelu"``, ``"silu"``, ``"glu"``, + ``"swiglu"``. + backend (str): Compute backend, either ``"torch"`` or ``"quack"``. + is_glu_variant (bool): ``True`` when ``activation`` is ``"glu"`` or + ``"swiglu"``, indicating that ``layer1`` produces ``2 * hidden_dim`` + outputs and the forward path applies a gating product. + layer1 (nn.Linear): First linear projection. + Shape: ``(hidden_dim * glu_factor, dim)`` where ``glu_factor`` + is 2 for gated activations, 1 otherwise. + dropout (nn.Module): Dropout layer instantiated from ``dropout_cfg``. + Applied between the activation (or gate product) and ``layer2``. + layer2 (nn.Linear): Second linear projection. + Shape: ``(dim, hidden_dim)``. Projects back to the input + channel dimension ``C``. """ def __init__( @@ -180,7 +392,59 @@ def __init__( init_method_in: Callable[[torch.Tensor], torch.Tensor] | None = None, init_method_out: Callable[[torch.Tensor], torch.Tensor] | None = None, ): - """Initialize the MLP.""" + """Initialise the MLP. + + Args: + dim: Input and output channel dimension ``C``. Both ``layer1`` + and ``layer2`` preserve this outer dimension (the network's + residual stream width). + activation: Non-linearity to apply between the two linear layers. + Controls whether the MLP is plain or gated: + + - ``"relu"``, ``"gelu"``, ``"silu"``: plain MLP; ``layer1`` + maps ``C → H``. + - ``"glu"``: sigmoid-gated MLP; ``layer1`` maps ``C → 2H`` + and is split into gate + value halves. + - ``"swiglu"``: SiLU-gated MLP (recommended for modern + configs); same gating structure as ``"glu"`` but with + SiLU in place of sigmoid. + + dropout_cfg: :class:`~nvsubquadratic.lazy_config.LazyConfig` + specifying the dropout module inserted between the activation + and ``layer2``. Use ``torch.nn.Identity`` (or a + ``LazyConfig`` that targets it) for no dropout. + expansion_factor: Multiplier that sets the hidden dimension + ``H = floor(expansion_factor * dim)``. Defaults to ``2.0``. + For gated activations, the actual ``layer1`` output width is + ``2 * H``; to keep total parameter count similar to a plain + MLP with ``expansion_factor=2``, use ``expansion_factor ≈ 4/3`` + for SwiGLU / GLU. + bias: Whether to include additive bias terms in ``layer1`` and + ``layer2``. Defaults to ``False``. Must be ``False`` when + ``backend="quack"``. + backend: Compute kernel family. ``"torch"`` (default) uses + ``nn.Linear`` + PyTorch activation and runs on any hardware. + ``"quack"`` enables fused GEMM+activation kernels but is + currently experimental (raises :exc:`NotImplementedError` + at init). + init_method_in: Optional weight initialiser for ``layer1``. + Expected signature: ``init_method_in(out_features)(weight)`` + — a *curried* callable that first takes the number of output + features and returns an in-place initialiser. Bias is always + zero-initialised when present, regardless of this argument. + Pass ``None`` to use PyTorch's default Kaiming uniform init. + init_method_out: Optional weight initialiser for ``layer2``. + Same curried signature as ``init_method_in``, applied to + ``layer2.weight``. Pass ``None`` for PyTorch default init. + + Raises: + NotImplementedError: If ``backend="quack"`` (experimental; use + ``backend="torch"`` for now). + ValueError: If ``backend="quack"`` and any QuACK constraint is + violated (unsupported activation, ``bias=True``, or dimension + not divisible by 8). Currently unreachable due to the + :exc:`NotImplementedError` raised first. + """ super().__init__() self.hidden_dim = int(dim * expansion_factor) self.activation = activation @@ -212,7 +476,29 @@ def __init__( nn.init.zeros_(self.layer2.bias) def _apply_activation(self, x: torch.Tensor) -> torch.Tensor: - """Apply the activation function to the input.""" + """Apply the configured activation function to a pre-activation tensor. + + For plain activations (``"relu"``, ``"gelu"``, ``"silu"``) this is a + simple element-wise call. For gated variants (``"glu"``, ``"swiglu"``) + the tensor is split along the last dimension and a gate product is + applied, halving the channel count from ``2H`` to ``H``. + + Args: + x: Pre-activation tensor of shape ``(B, *spatial_dims, H')`` + where ``H' = hidden_dim`` for plain activations and + ``H' = 2 * hidden_dim`` for gated variants. + + Returns: + torch.Tensor: Post-activation tensor of shape + ``(B, *spatial_dims, hidden_dim)``. The channel dimension is + halved for gated activations (``"glu"``, ``"swiglu"``), unchanged + for plain ones. + + Raises: + ValueError: If ``self.activation`` is not one of the supported + values (``"relu"``, ``"gelu"``, ``"silu"``, ``"glu"``, + ``"swiglu"``). + """ if self.activation in ["relu", "gelu", "silu"]: return getattr(F, self.activation)(x) elif self.activation == "glu": @@ -243,9 +529,11 @@ def flop_count(self, num_tokens: int) -> int: Args: num_tokens: Number of tokens (positions) the MLP is applied to. + Equal to the product of all spatial dimensions, i.e. + ``H * W`` for 2-D images or ``T`` for 1-D sequences. Returns: - Total FLOPs as an integer. + int: Total FLOPs as an integer. """ T = num_tokens flops = 0 @@ -262,7 +550,25 @@ def flop_count(self, num_tokens: int) -> int: return flops def forward(self, x: torch.Tensor) -> torch.Tensor: - """Forward pass of the MLP.""" + """Apply the two-layer MLP to an ND feature tensor. + + Operates identically on every spatial position — no cross-position + interaction occurs in this module. The QuACK path bypasses dropout + (QuACK kernels fuse both linear layers and the activation into a + single kernel; dropout must be inserted by the caller if needed). + + Args: + x: Input feature tensor of shape ``(B, *spatial_dims, C)`` in + channels-last layout, where ``B`` is the batch size, + ``spatial_dims`` are one or more spatial axes (e.g. ``(H, W)`` + for 2-D images or ``(T,)`` for 1-D sequences), and ``C`` + is the channel dimension (must equal the ``dim`` passed at + construction time). + + Returns: + torch.Tensor: Output tensor of shape ``(B, *spatial_dims, C)``, + the same shape as ``x``. + """ if self.backend == "quack": return _quack_mlp_forward(x, self.layer1.weight, self.layer2.weight, self.activation) From 1435c4b07a84679429a66b46fa8756769435dcdf Mon Sep 17 00:00:00 2001 From: David Wessels Date: Tue, 26 May 2026 09:09:01 +0200 Subject: [PATCH 51/72] docs(tracker): mark mlp, grn, layer_scale, masks_nd as done --- docs-tracker.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs-tracker.md b/docs-tracker.md index 24bb3b0e..35581e51 100644 --- a/docs-tracker.md +++ b/docs-tracker.md @@ -56,11 +56,11 @@ Work bottom-up: primitive ops → modules → networks → experiments. | `vit5_residual_block.py` | \[x\] | ViT5 residual block — LayerScale, register-token conditioning, no condition-mixer branch | | `patchify.py` | \[x\] | Patch embedding — strided conv, 1D/2D/3D, channels-last layout | | `position_encoding.py` | \[x\] | Axis-factorised learned PE — ND broadcast-expand, float32 output caveat | -| `masks_nd.py` | \[ \] | ND masking utils | -| `mlp.py` | \[ \] | MLP block | +| `masks_nd.py` | \[x\] | Exponential + Gaussian receptive-field windows; mask convention 1=included, 0=excluded | +| `mlp.py` | \[x\] | Two-layer MLP — GELU/SwiGLU/GLU variants, expansion-ratio math, QuACK backend noted | | `film.py` | \[x\] | FiLM conditioning — γ(c)⊙x + β(c), SIREN-based kernel generator | -| `grn.py` | \[ \] | GRN normalisation | -| `layer_scale.py` | \[ \] | LayerScale | +| `grn.py` | \[x\] | GRN — per-channel L2 norm, inter-channel competition, ConvNeXt V2 reference | +| `layer_scale.py` | \[x\] | LayerScale — per-channel λ⊙F(x), init_values guidance, \_no_weight_decay tag | | `rms_norm.py` | \[ \] | RMS normalisation | | `rms_norm_channel_first.py` | \[ \] | Channel-first RMS norm | | `drop_path.py` | \[ \] | Stochastic depth | From ed5ee9c810108b3e6afcc7e02578ae1c4a0e498a Mon Sep 17 00:00:00 2001 From: David Wessels Date: Tue, 26 May 2026 09:12:28 +0200 Subject: [PATCH 52/72] docs(write+integrate/rms_norm,rms_norm_channel_first,drop_path,causal_conv1d): add module and class docstrings with math context --- nvsubquadratic/modules/causal_conv1d.py | 94 ++++++++--- nvsubquadratic/modules/drop_path.py | 101 ++++++++++-- nvsubquadratic/modules/rms_norm.py | 150 ++++++++++++++---- .../modules/rms_norm_channel_first.py | 112 +++++++++---- 4 files changed, 368 insertions(+), 89 deletions(-) diff --git a/nvsubquadratic/modules/causal_conv1d.py b/nvsubquadratic/modules/causal_conv1d.py index eae241c0..53881a35 100644 --- a/nvsubquadratic/modules/causal_conv1d.py +++ b/nvsubquadratic/modules/causal_conv1d.py @@ -2,11 +2,38 @@ """Causal 1D convolution module. -This module provides a 1D convolution with causal (left-only) padding, -ensuring that outputs at position i only depend on inputs at positions 0, 1, ..., i. +Provides :class:`CausalConv1D`, a :class:`torch.nn.Conv1d` subclass that pads +sequences on the **left only**, enforcing the causal constraint that the output at +position ``i`` depends solely on inputs at positions ``0 … i``: -Classes: - CausalConv1D: Conv1d subclass with configurable causal or symmetric padding. +.. code-block:: text + + Causal padding (kernel_size=K, dilation=D): + left_pad = (K - 1) * D right_pad = 0 + + For stride=1 the output length equals the input length: + L_out = L_in (same as "same" conv but strictly causal) + +The module also supports **symmetric** (non-causal) same-padding when +``is_causal=False``, making it a drop-in for standard convolutions in contexts +that need a runtime-switchable causality flag. + +**Use inside Mamba / SSM blocks** + +:class:`CausalConv1D` is used in :mod:`nvsubquadratic.modules.mamba_nd` as a +depthwise filter applied along each spatial axis before the selective scan. +The causal constraint along the time axis is essential for autoregressive +generation; for spatial axes (height, width) the symmetric variant is used +so the scan sees the full context. + +**Difference from** ``subq_ops_causal_conv1d`` **/ ``causal_conv1d_custom``** + +:class:`CausalConv1D` uses PyTorch's standard ``F.conv1d`` and works everywhere +(CPU, any GPU). The modules in +:mod:`nvsubquadratic.ops.causal_conv1d_custom` and +:mod:`nvsubquadratic.modules.subq_ops_causal_conv1d` are thin wrappers around +the hand-fused ``causal_conv1d`` CUDA kernel from ``mamba_ssm``; they are +faster but require a compatible CUDA installation. """ from __future__ import annotations @@ -18,28 +45,52 @@ class CausalConv1D(torch.nn.Conv1d): """1D convolution with configurable causal (left-only) or symmetric padding. - When ``is_causal=True`` (default), pads only on the left so outputs at - position *i* depend only on inputs at positions 0 … i. When - ``is_causal=False``, uses standard symmetric padding (same as Conv1d). + Subclasses :class:`torch.nn.Conv1d` and overrides :meth:`forward` to apply + explicit left-only padding rather than the built-in ``padding`` argument. + This guarantees that ``output[b, :, i]`` depends only on + ``input[b, :, 0 … i]`` (the causal constraint). - Example: - >>> conv = CausalConv1D(in_channels=16, out_channels=16, kernel_size=7, groups=16) - >>> x = torch.randn(2, 16, 100) # [B, C, L] - >>> y = conv(x) # [B, C, 100] - same length, causal + **Padding formulae** + + For a kernel of size ``K`` with dilation ``D``: + + - *Causal* (``is_causal=True``): + ``left_pad = (K-1)*D``, ``right_pad = 0``. + Output length equals input length for stride=1. + - *Symmetric* (``is_causal=False``): + ``sym_pad = ((K-1)*D) // 2`` on both sides. + Equivalent to ``Conv1d(padding='same')`` for odd ``K``. + + Pass ``padding=0`` to the parent constructor (the default) — padding is + handled explicitly in :meth:`forward`. + + Attributes: + is_causal (bool): If ``True``, left-only padding is applied. + + Example:: + + conv = CausalConv1D(in_channels=16, out_channels=16, kernel_size=7, groups=16) + x = torch.randn(2, 16, 100) # [B, C, L] + y = conv(x) # [B, 16, 100] — same length, strictly causal Notes: - - Stride > 1 reduces the sequence length as in standard convs. - - For grouped/depthwise convolutions, set ``groups=in_channels``. + - Stride > 1 reduces the output length as in standard convolutions. + - For depthwise convolutions set ``groups=in_channels``. + - The ``padding`` argument of the parent constructor is ignored; always + pass ``padding=0`` (or omit it) when constructing this class. """ def __init__(self, *args, is_causal: bool = True, **kwargs) -> None: - """Initialize CausalConv1D. + """Initialise CausalConv1D. Args: - *args: Positional arguments forwarded to ``torch.nn.Conv1d``. - is_causal: If True, apply left-only (causal) padding. If False, - use standard symmetric padding. - **kwargs: Keyword arguments forwarded to ``torch.nn.Conv1d``. + *args: Positional arguments forwarded to :class:`torch.nn.Conv1d` + (``in_channels``, ``out_channels``, ``kernel_size``, …). + is_causal: If ``True`` (default), apply left-only (causal) padding. + If ``False``, apply symmetric same-padding. + **kwargs: Keyword arguments forwarded to :class:`torch.nn.Conv1d`. + ``padding`` should be ``0`` (or omitted) since :meth:`forward` + handles padding explicitly. """ super().__init__(*args, **kwargs) self.is_causal = is_causal @@ -48,10 +99,13 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: # type: ignore[override """Apply 1D convolution with causal or symmetric padding. Args: - input: Input tensor of shape [B, C, L]. + input: Input tensor of shape ``[B, C_in, L]``. Returns: - Output tensor of shape [B, C_out, L] (same length for stride=1). + torch.Tensor: Output of shape ``[B, C_out, L_out]``. + + - ``L_out = L`` when ``stride=1`` (same-length output). + - ``L_out = ceil(L / stride)`` for ``stride > 1``. """ kernel_size = self.kernel_size[0] if isinstance(self.kernel_size, tuple) else self.kernel_size dilation = self.dilation[0] if isinstance(self.dilation, tuple) else self.dilation diff --git a/nvsubquadratic/modules/drop_path.py b/nvsubquadratic/modules/drop_path.py index 3e6d1a30..29cd461a 100644 --- a/nvsubquadratic/modules/drop_path.py +++ b/nvsubquadratic/modules/drop_path.py @@ -1,6 +1,37 @@ """DropPath (Stochastic Depth): randomly drop entire residual branches during training. -Reference: Huang et al., "Deep Networks with Stochastic Depth", ECCV 2016. +Stochastic Depth (Huang et al., "Deep Networks with Stochastic Depth", ECCV 2016) +is a regularisation technique that randomly bypasses entire residual branches with +per-sample probability ``p`` during training: + +.. code-block:: text + + # training + keep_prob = 1 - drop_prob + mask = Bernoulli(keep_prob) ∈ {0, 1}^B # one scalar per sample + out = x * mask / keep_prob # inverted dropout scaling + + # inference + out = x # pure identity + +The inverted-dropout scaling (dividing by ``keep_prob``) keeps the expected +output magnitude equal to the input magnitude, so no rescaling is needed at +inference time. + +Unlike standard Dropout — which drops individual *elements* — DropPath drops the +*entire residual contribution* of a sample, which gives a stronger regularisation +signal in deep residual networks. The effective network depth seen by each sample +is therefore drawn from a uniform distribution over ``[1, L]`` during training, +where ``L`` is the total depth. + +**Integration pattern** (``vit5_residual_block.py``):: + + x = x + drop_path(ls_attn(mixer(norm(x))), drop_prob, self.training) + x = x + drop_path(ls_mlp(mlp(mlp_norm(x))), drop_prob, self.training) + +Reference: + Huang, G., et al., "Deep Networks with Stochastic Depth", + ECCV 2016. arXiv:1603.09382. """ import torch @@ -8,7 +39,26 @@ def drop_path(x: torch.Tensor, drop_prob: float, training: bool) -> torch.Tensor: - """Per-sample stochastic depth.""" + """Apply per-sample stochastic depth (functional form). + + During training each sample in the batch is independently kept or dropped + with probability ``1 - drop_prob`` / ``drop_prob`` respectively. The + kept samples are rescaled by ``1 / (1 - drop_prob)`` to preserve the + expected magnitude. At inference time the function is an identity. + + Args: + x: Input tensor of shape ``[B, *]`` — any layout; the drop mask has + shape ``(B, 1, …, 1)`` and broadcasts over all non-batch dimensions. + drop_prob: Probability of dropping a sample's contribution. + ``0.0`` disables dropping; ``1.0`` would zero every sample. + training: Whether the model is currently in training mode. + Set to ``False`` (or call ``model.eval()``) to disable dropping. + + Returns: + torch.Tensor: Same shape and dtype as ``x``. During training, + approximately ``drop_prob * B`` samples are zeroed and the rest are + rescaled. During inference, returns ``x`` unchanged. + """ if drop_prob == 0.0 or not training: return x keep_prob = 1.0 - drop_prob @@ -20,33 +70,62 @@ def drop_path(x: torch.Tensor, drop_prob: float, training: bool) -> torch.Tensor class DropPath(nn.Module): - """Drop paths (stochastic depth) per sample. + """Drop paths (stochastic depth) per sample — ``nn.Module`` wrapper. + + Thin stateful wrapper around the functional :func:`drop_path` that stores + the drop probability and reads ``self.training`` automatically, making it + a plug-in replacement wherever an ``nn.Module`` is required. + + **Effect on training vs. inference** + + - **Training** (``model.train()``): each sample's residual branch output + is dropped with probability ``drop_prob`` and kept samples are rescaled + by ``1 / (1 - drop_prob)``. + - **Inference** (``model.eval()``): the module is a pure identity; no + Bernoulli sampling or scaling is performed. + + Attributes: + drop_prob (float): Probability of dropping a sample's residual output. + Typically set between ``0.0`` (no drop) and ``0.3`` for deep ViTs. Args: - drop_prob: Probability of dropping the path. 0.0 means no drop. + drop_prob: Drop probability. Defaults to ``0.0`` (disabled). """ def __init__(self, drop_prob: float = 0.0): - """Store drop probability.""" + """Initialise DropPath. + + Args: + drop_prob: Probability of dropping each sample's residual update. + ``0.0`` disables the module (pure identity). Default ``0.0``. + """ super().__init__() self.drop_prob = drop_prob def flop_count(self) -> int: - """Count FLOPs for stochastic depth. + """Return FLOP count — always zero. DropPath is a stochastic identity (training) or pure identity - (inference). Neither involves floating-point arithmetic beyond a - Bernoulli sample and a scalar divide, which are negligible. + (inference). The Bernoulli sampling and scalar division are + negligible and not counted as floating-point arithmetic. Returns: - Always 0. + Always ``0``. """ return 0 def forward(self, x: torch.Tensor) -> torch.Tensor: - """Apply stochastic depth to the input.""" + """Apply stochastic depth to the input tensor. + + Args: + x: Input tensor of shape ``[B, *]``. + + Returns: + torch.Tensor: Same shape and dtype as ``x``, with per-sample + dropping applied during training. + """ return drop_path(x, self.drop_prob, self.training) def extra_repr(self) -> str: - """Return drop_prob for repr().""" + """Return drop probability for ``repr()``.""" return f"drop_prob={self.drop_prob:.3f}" diff --git a/nvsubquadratic/modules/rms_norm.py b/nvsubquadratic/modules/rms_norm.py index a5b72bfc..0a06fc6c 100644 --- a/nvsubquadratic/modules/rms_norm.py +++ b/nvsubquadratic/modules/rms_norm.py @@ -1,12 +1,39 @@ """RMSNorm — Root Mean Square Layer Normalization. -Uses QuACK's fused kernel on CUDA when available (Hopper/Blackwell: H100, B200, B300 only) -and ``use_quack=True`` (the default). On other GPUs (e.g. Ampere), when quack is not -installed, on CPU, or when ``use_quack=False``, uses a pure-PyTorch fallback. +RMSNorm (Zhang & Sennrich, "Root Mean Square Layer Normalization", NeurIPS 2019) is a +simplified variant of LayerNorm that omits the mean-centering step. For a vector +``x ∈ ℝ^D`` it computes: -Set ``use_quack=False`` when running under ``torch.compile`` to let the compiler -fuse the norm with adjacent operations instead of treating the QuACK kernel as an -opaque barrier. +.. code-block:: text + + RMS(x) = sqrt( (1/D) Σ_d x_d² + ε ) + x̂ = x / RMS(x) + out = γ ⊙ x̂ + +where ``γ ∈ ℝ^D`` is a learned per-element scale and ``ε`` is a small stability +constant. Dropping the mean-centering from LayerNorm cuts the compute roughly in +half for large ``D`` while typically achieving comparable training stability. + +**Backends** + +Two backend implementations are selected at runtime: + +- **QuACK** (default, ``use_quack=True``): fused CUDA kernel from the ``quack`` + package. Active only on Hopper / Blackwell GPUs (SM ≥ 9.0 — H100, B200, B300). + On older GPUs (e.g. Ampere A100) QuACK's backward kernel is incompatible; the + module detects this and falls back to PyTorch automatically. The QuACK kernel is + opaque to ``torch.compile`` and acts as a fusion barrier. + +- **PyTorch** (``use_quack=False`` or fallback): pure-Python ``torch`` ops + (``pow``, ``mean``, ``rsqrt``), upcasted to float32 for numerical safety. Fully + visible to the compiler, enabling fusion with adjacent ops. + +Set ``use_quack=False`` explicitly when running under ``torch.compile`` so the compiler +can fuse the norm with surrounding operations. + +Reference: + Zhang, B. & Sennrich, R., "Root Mean Square Layer Normalization", + NeurIPS 2019. arXiv:1910.07467. """ import warnings @@ -36,29 +63,48 @@ def _rmsnorm_pytorch(x: torch.Tensor, weight: torch.nn.Parameter, eps: float) -> class RMSNorm(nn.Module): - """Root Mean Square Layer Normalization. + """Root Mean Square Layer Normalization (Zhang & Sennrich, arXiv:1910.07467). + + Normalises a tensor over its last dimension using RMS statistics, then + scales the result by a learned per-channel weight: - Normalizes over the last dimension and scales by a learnable weight. - Accepts tensors of any shape ``[..., dim]``. + .. code-block:: text - Two backends are available: + RMS(x) = sqrt( mean(x²) + ε ) # scalar per token + out = (x / RMS(x)) * γ # γ broadcast over leading dims - - **QuACK** (default): Fused CUDA kernel from the ``quack`` package. - Only runs on Hopper/Blackwell GPUs (SM 9.0+). Opaque to - ``torch.compile`` — acts as a fusion barrier. - - **PyTorch**: Pure ``torch`` ops (``pow``, ``mean``, ``rsqrt``). Fully - visible to ``torch.compile``, enabling fusion with adjacent operations. + Accepts tensors of any shape ``[*leading, D]``; only the last dimension is + normalised. The learned weight ``γ`` has shape ``(D,)`` and is **excluded + from weight decay** via the ``_no_weight_decay`` tag. + + **Backend selection** (see module docstring for full details): + + - ``use_quack=True`` (default): QuACK fused kernel on SM ≥ 9.0 GPUs. + Falls back to PyTorch automatically on older GPUs. + - ``use_quack=False``: Pure PyTorch; preferred under ``torch.compile``. + + Attributes: + weight (nn.Parameter): Learnable scale ``γ`` of shape ``(dim,)``, + ones-initialised. Tagged ``_no_weight_decay = True``. + eps (float): Stability constant added inside the square root. + use_quack (bool): Whether to attempt the QuACK kernel path. Args: - dim: Size of the last dimension to normalize over. - eps: Small constant added to the variance for numerical stability. + dim: Size of the last dimension ``D`` to normalise over. + eps: Small positive constant for numerical stability. Default ``1e-6``. use_quack: If ``True`` (default), use the QuACK fused kernel when available. Set to ``False`` to force the PyTorch path, which lets ``torch.compile`` fuse the norm with surrounding ops. """ def __init__(self, dim: int, eps: float = 1e-6, use_quack: bool = True): - """Initialize RMSNorm.""" + """Initialise RMSNorm with ones-initialised weight. + + Args: + dim: Channel dimension ``D``; determines the shape of ``weight``. + eps: Stability constant added to the RMS denominator. Default ``1e-6``. + use_quack: Enable QuACK kernel path. Default ``True``. + """ super().__init__() self.weight = nn.Parameter(torch.ones(dim)) self.weight._no_weight_decay = True @@ -85,7 +131,17 @@ def flop_count(self, num_tokens: int) -> int: return 3 * num_tokens * dim def forward(self, x: torch.Tensor) -> torch.Tensor: - """Normalize over last dimension and scale by weight.""" + """Apply RMS normalisation over the last dimension. + + Args: + x: Input tensor of shape ``[*leading, D]``. Any number of + leading dimensions is supported; only the last axis is + normalised. + + Returns: + torch.Tensor: Normalised and scaled tensor, same shape and dtype + as ``x``. + """ # Only use quack on GPUs that support it (SM 9+). On Ampere (8.x), quack backward fails; # we must not call quack at all so autograd never invokes its backward. if self.use_quack and _quack_available and x.is_cuda and _cuda_supports_quack(x.device): @@ -105,23 +161,48 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class PerHeadRMSNorm(nn.Module): - """RMSNorm applied independently to each head. + """RMSNorm applied independently to each attention head (QK-norm). + + Accepts a flat hidden representation of shape ``[*leading, H·D]``, + reshapes to ``[*leading, H, D]``, applies an independent :class:`RMSNorm` + to each head's ``D``-dimensional slice, then flattens back to + ``[*leading, H·D]``. Each head has its own learnable scale ``γ ∈ ℝ^D``. + + This is the *QK-norm* technique used in ViT-5 / vit5_attention.py to + stabilise attention logit magnitudes at large model sizes: - Accepts ``[..., hidden_dim]``, reshapes to ``[..., num_heads, head_dim]``, - normalizes over ``head_dim``, and flattens back. Each head has its own - learnable scale vector of size ``head_dim``. + .. code-block:: text + + x : [*leading, H*D] + x → reshape → [*leading, H, D] + x → RMSNorm per head → [*leading, H, D] + x → flatten → [*leading, H*D] + + Attributes: + num_heads (int): Number of attention heads ``H``. + head_dim (int): Dimension per head ``D``. + norm (RMSNorm): Shared :class:`RMSNorm` instance applied to each head + slice (the weight ``γ`` has shape ``(D,)``). Args: - num_heads: Number of attention heads. - head_dim: Dimension of each head. - eps: Small constant added to the variance for numerical stability. + num_heads: Number of attention heads ``H``. + head_dim: Dimension of each head ``D``. + eps: Small constant for numerical stability. Default ``1e-6``. use_quack: If ``True`` (default), use the QuACK fused kernel when - available. Set to ``False`` to force the PyTorch path, which - lets ``torch.compile`` fuse the norm with surrounding ops. + available. Set to ``False`` to force the PyTorch path so that + ``torch.compile`` can fuse the norm with surrounding ops. """ def __init__(self, num_heads: int, head_dim: int, eps: float = 1e-6, use_quack: bool = True): - """Initialize PerHeadRMSNorm.""" + """Initialise PerHeadRMSNorm. + + Args: + num_heads: Number of attention heads ``H``. + head_dim: Dimension per head ``D``; the inner :class:`RMSNorm` + normalises over this dimension. + eps: Stability constant. Default ``1e-6``. + use_quack: Enable QuACK kernel path. Default ``True``. + """ super().__init__() self.num_heads = num_heads self.head_dim = head_dim @@ -145,7 +226,16 @@ def flop_count(self, num_tokens: int) -> int: return 3 * num_tokens * self.num_heads * self.head_dim def forward(self, x: torch.Tensor) -> torch.Tensor: - """Normalize each head independently over head_dim.""" + """Apply per-head RMS normalisation. + + Args: + x: Input tensor of shape ``[*leading, num_heads * head_dim]``. + + Returns: + torch.Tensor: Normalised tensor of the same shape as ``x``, where + each head's ``head_dim``-slice has been RMS-normalised and scaled + by the shared learnable ``γ``. + """ shape = x.shape x = x.view(*shape[:-1], self.num_heads, self.head_dim) x = self.norm(x) diff --git a/nvsubquadratic/modules/rms_norm_channel_first.py b/nvsubquadratic/modules/rms_norm_channel_first.py index 31d59da2..ae982c82 100644 --- a/nvsubquadratic/modules/rms_norm_channel_first.py +++ b/nvsubquadratic/modules/rms_norm_channel_first.py @@ -1,21 +1,41 @@ -"""RMSNorm — Channel-First variant (normalizes over dim=1). +"""RMSNorm — Channel-First variant (normalises over ``dim=1``). -Uses QuACK's fused channel-first kernel on CUDA when available (Hopper/Blackwell: -H100, B200, B300 only) and ``use_quack=True`` (the default). On other GPUs, -when quack is not installed, on CPU, or when ``use_quack=False``, uses a -pure-PyTorch fallback. +This is the channel-first counterpart of :mod:`nvsubquadratic.modules.rms_norm`. +It applies Root Mean Square Layer Normalization (Zhang & Sennrich, NeurIPS 2019) +along the **channel axis** (``dim=1``) of tensors stored in channel-first layout +such as ``[B, C, H, W]`` or ``[B, C, T]``: -Set ``use_quack=False`` when running under ``torch.compile`` to let the compiler -fuse the norm with adjacent operations instead of treating the QuACK kernel as an -opaque barrier. +.. code-block:: text -Expects input tensors in channel-first layout, e.g. ``[B, C, H, W]`` or -``[B, C, T]``. Use this only at call sites where the data is *already* -channel-first (e.g. inside the Hyena mixer). For trunk code where tensors -are ``[B, T, C]``, use the regular channel-last ``RMSNorm`` instead. + RMS(x)_b = sqrt( (1/C) Σ_c x[b, c, ...]² + ε ) # scalar per sample + out[b, c, ...] = (x[b, c, ...] / RMS(x)_b) * γ_c -Modules that handle channel-first norms specially (e.g. ``Hyena``) can -duck-type on the ``channels_first`` attribute (always ``True`` here). +where ``γ ∈ ℝ^C`` is a learned per-channel scale initialised to ones. + +**When to use this module vs. :class:`~nvsubquadratic.modules.rms_norm.RMSNorm`** + +- Use :class:`RMSNormChannelFirst` where data is *already* channel-first + (e.g. inside :class:`~nvsubquadratic.modules.hyena_nd.HyenaOperatorND` which + uses ``[B, C, *spatial]`` tensors throughout). +- Use the regular channel-last :class:`~nvsubquadratic.modules.rms_norm.RMSNorm` + in trunk code where tensors are ``[B, T, C]`` (ViT / transformer layout). + +**Duck-typing sentinel** + +The class-level attribute ``channels_first = True`` lets callers (e.g. +``HyenaOperatorND``) detect the layout at construction time without an +``isinstance`` check. + +**Backends** (same policy as :mod:`nvsubquadratic.modules.rms_norm`): + +- **QuACK** (default, ``use_quack=True``): fused CUDA kernel + ``quack.rmsnorm_channel_first``; SM ≥ 9.0 only. +- **PyTorch** (``use_quack=False`` or fallback): pure ``torch`` ops, upcasts + to float32 for numerical safety; fully visible to ``torch.compile``. + +Reference: + Zhang, B. & Sennrich, R., "Root Mean Square Layer Normalization", + NeurIPS 2019. arXiv:1910.07467. """ import torch @@ -46,22 +66,32 @@ def _rmsnorm_channel_first_pytorch(x: torch.Tensor, weight: torch.nn.Parameter, class RMSNormChannelFirst(nn.Module): - """Root Mean Square Layer Normalization — channel-first layout. + """Root Mean Square Layer Normalization — channel-first layout (arXiv:1910.07467). + + Normalises a tensor along **``dim=1``** (the channel axis) and scales the + result by a learned per-channel weight. Accepts tensors of shape + ``[B, C, *spatial]``, e.g. ``[B, C, H, W]`` for 2-D or ``[B, C, L]`` for 1-D. - Normalizes over ``dim=1`` (the channel dimension) and scales by a learnable - weight. Accepts tensors of shape ``[B, C, *spatial]`` (e.g. ``[B, C, H, W]``). + .. code-block:: text - Two backends are available: + RMS(x)_b = sqrt( mean_C(x[b, :, ...]²) + ε ) # scalar per sample + out[b,:,…] = x[b,:,…] / RMS(x)_b * γ # γ: [C, 1, …, 1] - - **QuACK** (default): Fused CUDA kernel from the ``quack`` package - (``rmsnorm_channel_first``). Only runs on Hopper/Blackwell GPUs - (SM 9.0+). Opaque to ``torch.compile`` — acts as a fusion barrier. - - **PyTorch**: Pure ``torch`` ops (``pow``, ``mean``, ``rsqrt``). Fully - visible to ``torch.compile``, enabling fusion with adjacent operations. + **Duck-typing sentinel**: the class attribute ``channels_first = True`` + allows callers (e.g. ``HyenaOperatorND``) to detect the layout without + an ``isinstance`` check. + + Attributes: + channels_first (bool): Always ``True``; used by callers to detect + channel-first norm modules at construction time. + weight (nn.Parameter): Learnable scale ``γ`` of shape ``(C,)``, + ones-initialised. Tagged ``_no_weight_decay = True``. + eps (float): Stability constant added inside the square root. + use_quack (bool): Whether to attempt the QuACK kernel path. Args: - dim: Number of channels (size of dim=1) to normalize over. - eps: Small constant added to the variance for numerical stability. + dim: Number of channels ``C`` (size of ``dim=1``). + eps: Small positive constant for numerical stability. Default ``1e-6``. use_quack: If ``True`` (default), use the QuACK fused kernel when available. Set to ``False`` to force the PyTorch path, which lets ``torch.compile`` fuse the norm with surrounding ops. @@ -70,7 +100,13 @@ class RMSNormChannelFirst(nn.Module): channels_first: bool = True def __init__(self, dim: int, eps: float = 1e-6, use_quack: bool = True): - """Initialize RMSNormChannelFirst.""" + """Initialise RMSNormChannelFirst with ones-initialised weight. + + Args: + dim: Number of channels ``C``; determines the shape of ``weight``. + eps: Stability constant. Default ``1e-6``. + use_quack: Enable QuACK kernel path. Default ``True``. + """ super().__init__() self.weight = nn.Parameter(torch.ones(dim)) self.weight._no_weight_decay = True @@ -78,12 +114,32 @@ def __init__(self, dim: int, eps: float = 1e-6, use_quack: bool = True): self.use_quack = use_quack def flop_count(self, num_tokens: int) -> int: - """Count FLOPs — identical cost to channel-last RMSNorm.""" + """Return approximate FLOP count, identical in cost to channel-last RMSNorm. + + The three operations — square, mean+rsqrt, scale — each touch every + element once, giving ``3 * num_tokens * C`` FLOPs regardless of the + memory layout. + + Args: + num_tokens: Total number of spatial positions in the batch + (``B * prod(spatial_shape)``). + + Returns: + Integer FLOP estimate: ``3 * num_tokens * C``. + """ dim = self.weight.shape[0] return 3 * num_tokens * dim def forward(self, x: torch.Tensor) -> torch.Tensor: - """Normalize over channel dimension (dim=1) and scale by weight.""" + """Normalise over the channel dimension (``dim=1``) and scale by weight. + + Args: + x: Input tensor of shape ``[B, C, *spatial]``. + + Returns: + torch.Tensor: Normalised and scaled tensor of the same shape and + dtype as ``x``. + """ if self.use_quack and _quack_cf_available and x.is_cuda and _cuda_supports_quack(x.device): return _quack_rmsnorm_cf(x, self.weight, eps=self.eps) return _rmsnorm_channel_first_pytorch(x, self.weight, self.eps) From 096fcd7560f24c33ea887ce1541c74141bbdfe55 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Tue, 26 May 2026 09:12:43 +0200 Subject: [PATCH 53/72] docs(tracker): mark rms_norm, rms_norm_channel_first, drop_path, causal_conv1d as done --- docs-tracker.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs-tracker.md b/docs-tracker.md index 35581e51..8a738606 100644 --- a/docs-tracker.md +++ b/docs-tracker.md @@ -61,10 +61,10 @@ Work bottom-up: primitive ops → modules → networks → experiments. | `film.py` | \[x\] | FiLM conditioning — γ(c)⊙x + β(c), SIREN-based kernel generator | | `grn.py` | \[x\] | GRN — per-channel L2 norm, inter-channel competition, ConvNeXt V2 reference | | `layer_scale.py` | \[x\] | LayerScale — per-channel λ⊙F(x), init_values guidance, \_no_weight_decay tag | -| `rms_norm.py` | \[ \] | RMS normalisation | -| `rms_norm_channel_first.py` | \[ \] | Channel-first RMS norm | -| `drop_path.py` | \[ \] | Stochastic depth | -| `causal_conv1d.py` | \[ \] | Causal 1D conv | +| `rms_norm.py` | \[x\] | RMSNorm + PerHeadRMSNorm — QuACK/PyTorch backends, math formula, QK-norm usage | +| `rms_norm_channel_first.py` | \[x\] | Channel-first RMSNorm — normalises dim=1, `channels_first` sentinel, Hyena usage | +| `drop_path.py` | \[x\] | Stochastic depth — functional + Module, inverted-dropout scaling, causal vs training | +| `causal_conv1d.py` | \[x\] | CausalConv1D — left-only pad formula, symmetric mode, Mamba usage context | | `subq_ops_causal_conv1d.py` | \[x\] | New (1D PR): `nn.Conv1d`-compatible depthwise wrapper around `subq_ops.causal_conv1d` | | `schedulers.py` | \[ \] | LR schedulers | | `distributed_depthwise_conv_nd.py` | \[ \] | Distributed depthwise conv | From 70101c96695f0553acc6403f5f57587b7b8677e7 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Tue, 26 May 2026 09:46:56 +0200 Subject: [PATCH 54/72] docs(write+integrate/schedulers,distributed_depthwise_conv_nd,general_purpose_resnet,classification_resnet): add module and class docstrings with math/arch context --- .../modules/distributed_depthwise_conv_nd.py | 274 ++++++++++++------ nvsubquadratic/modules/schedulers.py | 38 ++- .../networks/classification_resnet.py | 81 ++++-- .../networks/general_purpose_resnet.py | 157 ++++++++-- 4 files changed, 400 insertions(+), 150 deletions(-) diff --git a/nvsubquadratic/modules/distributed_depthwise_conv_nd.py b/nvsubquadratic/modules/distributed_depthwise_conv_nd.py index 050f9dfe..9f5643b2 100644 --- a/nvsubquadratic/modules/distributed_depthwise_conv_nd.py +++ b/nvsubquadratic/modules/distributed_depthwise_conv_nd.py @@ -1,37 +1,56 @@ # TODO: Add license header here -"""Distributed depthwise convolution wrappers for Context Parallelism. - -This module provides distributed depthwise-only convolution wrappers that handle: - -1. **Context Parallel (CP) Communication**: Handles CP-aware communication patterns - with direct channel slicing. - -2. **Depthwise-only Design**: Distributed weight management where: - - Weight shape: [num_groups, kernel_size...] (no extra dimension) - - Direct repeat_interleave for weight expansion - - in_channels == out_channels (depthwise property) - -Key Features: -- Depthwise convolutions only (no standard grouped convolutions) -- Group-based parameter sharing to reduce memory usage -- Distributed CP slicing logic - -Example Usage: - # 1D depthwise convolution with 128 groups for weight sharing - conv1d = DistributedDepthwiseConv1d( - hidden_dim=384, kernel_size=3, num_groups=128, causal=True - ) - - # 2D depthwise convolution - conv2d = DistributedDepthwiseConv2d( - hidden_dim=384, kernel_size=3, num_groups=128 - ) - - # 3D depthwise convolution - conv3d = DistributedDepthwiseConv3d( - hidden_dim=1024, kernel_size=3, num_groups=None # None = no weight sharing - ) +"""Distributed depthwise convolution wrappers for Context Parallelism (CP). + +This module provides 1-D, 2-D, and 3-D *depthwise* convolution layers that work +transparently under **Context Parallelism** (CP), where the channel dimension is +split across ``cp_world_size`` ranks. + +**Context Parallelism and channel slicing** + +In the Hyena / sequence-mixer CP setup each rank holds a contiguous slice of the +channel axis: rank ``r`` processes channels ``[r·C/P, (r+1)·C/P)`` where ``C`` is +the total hidden dimension and ``P`` is the CP world size. The depthwise +convolution weight needs to match — if ``cp_group`` is provided, :meth:`forward` +expands the full weight tensor then slices the appropriate channel range before +calling ``F.convNd``. No inter-rank communication is required. + +**Group-based weight sharing** + +To reduce the parameter count, weights are stored for ``num_groups`` prototypes +and expanded to ``hidden_dim`` channels via ``repeat_interleave`` along the +channel axis at each forward pass: + +.. code-block:: text + + weight: [G, *kernel] G = num_groups ≤ C + expanded: [C, *kernel] via repeat_interleave(C//G, dim=0) + +Setting ``num_groups=None`` (or ``num_groups=hidden_dim``) gives each channel its +own independent filter (standard depthwise convolution). Smaller ``num_groups`` +reduces memory at the cost of expressivity. + +**Initialisation** + +Weights and biases are drawn from ``Uniform(-b, b)`` where +``b = 1 / sqrt(prod(kernel_size))``, matching the default ``nn.Conv*`` scheme. + +Classes: + DistributedDepthwiseConv1d: 1-D variant; supports ``causal=True`` for left-only + padding (autoregressive / time-axis use). + DistributedDepthwiseConv2d: 2-D variant; always uses ``padding="same"``. + DistributedDepthwiseConv3d: 3-D variant; always uses ``padding="same"``. + +Example:: + + # 1-D causal depthwise conv with 128 weight groups on 384 channels + conv1d = DistributedDepthwiseConv1d(hidden_dim=384, kernel_size=3, num_groups=128, causal=True) + + # 2-D depthwise conv + conv2d = DistributedDepthwiseConv2d(hidden_dim=384, kernel_size=3, num_groups=128) + + # 3-D depthwise conv, no weight sharing + conv3d = DistributedDepthwiseConv3d(hidden_dim=1024, kernel_size=3, num_groups=None) """ import math @@ -44,18 +63,39 @@ class DistributedDepthwiseConv1d(nn.Module): - """1D depthwise convolution for CP parallelism. + """1-D depthwise convolution with CP-aware channel slicing and weight sharing. - This implements a depthwise convolution where: - - in_channels == out_channels (depthwise property) - - Weights are shared across groups via num_groups parameter - - Context Parallel slicing is done directly on the channel dimension + Stores a compact weight of shape ``[G, K]`` (``G`` groups, kernel size ``K``) + and expands it to ``[C, K]`` at each forward pass via ``repeat_interleave``. + When ``cp_group`` is provided, an additional slice ``[r·(C/P) : (r+1)·(C/P)]`` + is taken before calling ``F.conv1d``, so each CP rank only processes its + local channel shard. + + Supports two padding modes: + + - ``causal=True``: left-only pad of ``(K-1)`` before the conv; output length + equals input length with no future dependency. + - ``causal=False`` (default): ``padding="same"`` (symmetric); suitable for + spatial axes in multi-dimensional Hyena. Attributes: - hidden_dim (int): Number of input/output channels - kernel_size (int): Size of the convolution kernel - num_groups (int): Number of groups for weight sharing (reduces parameters) - group_dim (int): Channels per group (hidden_dim // num_groups) + hidden_dim (int): Total number of input/output channels ``C``. + kernel_size (int): Convolution kernel size ``K``. + causal (bool): Whether left-only (causal) padding is used. + num_groups (int): Number of weight prototype groups ``G``. + group_dim (int): Channels per group ``C // G``. + weight (nn.Parameter): Filter weights of shape ``[G, K]``. + bias (nn.Parameter | None): Optional bias of shape ``[G]``. + + Args: + hidden_dim: Total number of input/output channels ``C``. + kernel_size: Convolution kernel size ``K``. + causal: Apply left-only padding. Default ``False``. + num_groups: Weight prototype groups ``G``. ``None`` → ``G = C`` + (standard depthwise, no sharing). + bias: Include a learnable bias. Default ``False``. + dtype: Parameter dtype. Default ``torch.float32``. + device: Parameter device. Defaults to ``cuda:current`` if available. """ def __init__( @@ -68,16 +108,16 @@ def __init__( dtype: Optional[torch.dtype] = None, device: Optional[torch.device] = None, ): - """Initialize DistributedDepthwiseConv1d. + """Initialise DistributedDepthwiseConv1d. Args: - hidden_dim: Number of input and output channels - kernel_size: Size of the convolving kernel - causal: If True, applies causal padding - num_groups: Number of groups for weight sharing (if None, uses hidden_dim - no sharing) - bias: If True, adds a learnable bias - dtype: Data type for parameters - device: Device for parameters + hidden_dim: Total number of input/output channels ``C``. + kernel_size: Convolution kernel size ``K``. + causal: Apply left-only causal padding. Default ``False``. + num_groups: Weight prototype groups ``G ≤ C``. ``None`` → ``G = C``. + bias: Include a learnable bias. Default ``False``. + dtype: Parameter dtype. Default ``torch.float32``. + device: Parameter device. Defaults to current CUDA device. """ super().__init__() @@ -121,16 +161,29 @@ def init_weights(self): torch.nn.init.uniform_(self.bias, a=-bounds, b=bounds) def forward(self, x: torch.Tensor, cp_group: Optional[torch.distributed.ProcessGroup] = None) -> torch.Tensor: - """Forward pass with optional context parallelism. + """Apply 1-D depthwise convolution with optional CP channel slicing. + + The full weight ``[G, K]`` is expanded to ``[C, K]`` via + ``repeat_interleave``. If ``cp_group`` is given, only the slice + for the current rank is retained before calling ``F.conv1d``. Args: - x: Input tensor of shape [batch, channels, length] - - Without CP: channels = hidden_dim - - With CP: channels = hidden_dim // cp_world_size - cp_group: Context parallel process group (None for no CP) + x: Input tensor of shape ``[B, C_local, L]`` where + + * ``C_local = hidden_dim`` when not using CP, or + * ``C_local = hidden_dim // cp_world_size`` on CP rank ``r``. + + cp_group: Context-parallel process group. ``None`` → single-device + mode (no slicing). Returns: - Output tensor of same shape as input + torch.Tensor: Output of shape ``[B, C_local, L]``; same length as + input for stride=1. + + Raises: + AssertionError: If ``x.ndim != 3``. + RuntimeError: If ``x.shape[1]`` does not match the expected local + channel count after CP slicing. """ assert x.ndim == 3, f"Input must be 3D [batch, channels, length], got {x.ndim}D" @@ -188,12 +241,30 @@ def forward(self, x: torch.Tensor, cp_group: Optional[torch.distributed.ProcessG class DistributedDepthwiseConv2d(nn.Module): - """2D depthwise convolution for CP parallelism. + """2-D depthwise convolution with CP-aware channel slicing and weight sharing. + + Stores weights of shape ``[G, Kh, Kw]`` and expands to ``[C, Kh, Kw]`` at + runtime via ``repeat_interleave``. When ``cp_group`` is provided, the + appropriate channel slice for the current rank is extracted before calling + ``F.conv2d(padding="same")``. - This implements a depthwise convolution where: - - in_channels == out_channels (depthwise property) - - Weights are shared across groups via num_groups parameter - - Context Parallel slicing is done directly on the channel dimension + Expects input in **channel-first** layout: ``[B, C_local, H, W]``. + + Attributes: + hidden_dim (int): Total number of channels ``C``. + kernel_size (Tuple[int, int]): 2-D kernel dimensions ``(Kh, Kw)``. + num_groups (int): Weight prototype groups ``G``. + group_dim (int): Channels per group ``C // G``. + weight (nn.Parameter): Shape ``[G, Kh, Kw]``. + bias (nn.Parameter | None): Shape ``[G]`` or ``None``. + + Args: + hidden_dim: Total number of channels ``C``. + kernel_size: Kernel size; ``int`` → ``(K, K)``. + num_groups: Weight prototype groups ``G``. ``None`` → ``G = C``. + bias: Include a learnable bias. Default ``False``. + dtype: Parameter dtype. Default ``torch.float32``. + device: Parameter device. Defaults to current CUDA device. """ def __init__( @@ -205,15 +276,15 @@ def __init__( dtype: Optional[torch.dtype] = None, device: Optional[torch.device] = None, ): - """Initialize DistributedDepthwiseConv2d. + """Initialise DistributedDepthwiseConv2d. Args: - hidden_dim: Number of input and output channels - kernel_size: Size of the convolving kernel - num_groups: Number of groups for weight sharing (if None, uses hidden_dim) - bias: If True, adds a learnable bias - dtype: Data type for parameters - device: Device for parameters + hidden_dim: Total number of input/output channels ``C``. + kernel_size: Kernel size; ``int`` is broadcast to ``(K, K)``. + num_groups: Weight prototype groups ``G ≤ C``. ``None`` → ``G = C``. + bias: Include a learnable bias. Default ``False``. + dtype: Parameter dtype. Default ``torch.float32``. + device: Parameter device. Defaults to current CUDA device. """ super().__init__() @@ -254,14 +325,20 @@ def init_weights(self): torch.nn.init.uniform_(self.bias, a=-bounds, b=bounds) def forward(self, x: torch.Tensor, cp_group: Optional[torch.distributed.ProcessGroup] = None) -> torch.Tensor: - """Forward pass with optional context parallelism. + """Apply 2-D depthwise convolution with optional CP channel slicing. Args: - x: Input tensor of shape [batch, channels, height, width] - cp_group: Context parallel process group (None for no CP) + x: Input tensor of shape ``[B, C_local, H, W]``. + cp_group: Context-parallel process group. ``None`` → single-device. Returns: - Output tensor of same shape as input + torch.Tensor: Output of shape ``[B, C_local, H, W]`` (same spatial + size because ``padding="same"``). + + Raises: + AssertionError: If ``x.ndim != 4``. + RuntimeError: If ``x.shape[1]`` does not match the expected local + channel count after CP slicing. """ assert x.ndim == 4, f"Input must be 4D [batch, channels, height, width], got {x.ndim}D" @@ -302,12 +379,29 @@ def forward(self, x: torch.Tensor, cp_group: Optional[torch.distributed.ProcessG class DistributedDepthwiseConv3d(nn.Module): - """3D depthwise convolution for CP parallelism. + """3-D depthwise convolution with CP-aware channel slicing and weight sharing. + + Stores weights of shape ``[G, Kd, Kh, Kw]`` and expands to ``[C, Kd, Kh, Kw]`` + at runtime via ``repeat_interleave``. When ``cp_group`` is provided, only + the current rank's channel slice is passed to ``F.conv3d(padding="same")``. - This implements a depthwise convolution where: - - in_channels == out_channels (depthwise property) - - Weights are shared across groups via num_groups parameter - - Context Parallel slicing is done directly on the channel dimension + Expects input in **channel-first** layout: ``[B, C_local, D, H, W]``. + + Attributes: + hidden_dim (int): Total number of channels ``C``. + kernel_size (Tuple[int, int, int]): 3-D kernel dimensions ``(Kd, Kh, Kw)``. + num_groups (int): Weight prototype groups ``G``. + group_dim (int): Channels per group ``C // G``. + weight (nn.Parameter): Shape ``[G, Kd, Kh, Kw]``. + bias (nn.Parameter | None): Shape ``[G]`` or ``None``. + + Args: + hidden_dim: Total number of channels ``C``. + kernel_size: Kernel size; ``int`` → ``(K, K, K)``. + num_groups: Weight prototype groups ``G``. ``None`` → ``G = C``. + bias: Include a learnable bias. Default ``False``. + dtype: Parameter dtype. Default ``torch.float32``. + device: Parameter device. Defaults to current CUDA device. """ def __init__( @@ -319,15 +413,15 @@ def __init__( dtype: Optional[torch.dtype] = None, device: Optional[torch.device] = None, ): - """Initialize DistributedDepthwiseConv3d. + """Initialise DistributedDepthwiseConv3d. Args: - hidden_dim: Number of input and output channels - kernel_size: Size of the convolving kernel - num_groups: Number of groups for weight sharing (if None, uses hidden_dim) - bias: If True, adds a learnable bias - dtype: Data type for parameters - device: Device for parameters + hidden_dim: Total number of input/output channels ``C``. + kernel_size: Kernel size; ``int`` is broadcast to ``(K, K, K)``. + num_groups: Weight prototype groups ``G ≤ C``. ``None`` → ``G = C``. + bias: Include a learnable bias. Default ``False``. + dtype: Parameter dtype. Default ``torch.float32``. + device: Parameter device. Defaults to current CUDA device. """ super().__init__() @@ -368,14 +462,20 @@ def init_weights(self): torch.nn.init.uniform_(self.bias, a=-bounds, b=bounds) def forward(self, x: torch.Tensor, cp_group: Optional[torch.distributed.ProcessGroup] = None) -> torch.Tensor: - """Forward pass with optional context parallelism. + """Apply 3-D depthwise convolution with optional CP channel slicing. Args: - x: Input tensor of shape [batch, channels, depth, height, width] - cp_group: Context parallel process group (None for no CP) + x: Input tensor of shape ``[B, C_local, D, H, W]``. + cp_group: Context-parallel process group. ``None`` → single-device. Returns: - Output tensor of same shape as input + torch.Tensor: Output of shape ``[B, C_local, D, H, W]`` (same + spatial size because ``padding="same"``). + + Raises: + AssertionError: If ``x.ndim != 5``. + RuntimeError: If ``x.shape[1]`` does not match the expected local + channel count after CP slicing. """ assert x.ndim == 5, f"Input must be 5D [batch, channels, depth, height, width], got {x.ndim}D" diff --git a/nvsubquadratic/modules/schedulers.py b/nvsubquadratic/modules/schedulers.py index bc6ba266..04cfdd43 100644 --- a/nvsubquadratic/modules/schedulers.py +++ b/nvsubquadratic/modules/schedulers.py @@ -1,6 +1,25 @@ # TODO: Add license header here -"""Custom LR scheduler utilities.""" +"""Custom LR scheduler utilities. + +Currently provides :class:`ResumableSequentialLR`, a bug-fixed subclass of +:class:`torch.optim.lr_scheduler.SequentialLR` that correctly restores the +learning rate to the optimizer's ``param_groups`` after a checkpoint resume. + +**Background** + +A typical training schedule consists of multiple phases — e.g. a linear warmup +followed by a cosine decay. PyTorch's ``SequentialLR`` chains together a list of +sub-schedulers and advances through them when configured milestone epochs are +reached. However, as of PyTorch 2.10 its ``load_state_dict`` method restores +the scheduler's internal bookkeeping but silently omits propagating the restored +LR back to the optimizer, causing the schedule to restart from the initial warmup +LR after every checkpoint resume. + +:class:`ResumableSequentialLR` patches this by calling +``optimizer.param_groups[i]["lr"] = _last_lr[i]`` immediately after the parent +``load_state_dict`` completes. +""" import torch @@ -29,8 +48,21 @@ class ResumableSequentialLR(torch.optim.lr_scheduler.SequentialLR): still exists. """ - def load_state_dict(self, state_dict): - """Load state and apply restored LR to optimizer param groups.""" + def load_state_dict(self, state_dict: dict) -> None: + """Load scheduler state and propagate restored LRs to the optimizer. + + Calls the parent ``SequentialLR.load_state_dict``, then immediately + copies each value in ``self._last_lr`` into the corresponding + ``optimizer.param_groups[i]["lr"]``. This ensures that the first + ``optimizer.step()`` after a resume uses the learning rate that was + active when the checkpoint was saved, rather than the freshly + initialised (warmup-start) value. + + Args: + state_dict: Scheduler state dictionary as produced by + :meth:`state_dict`. Typically loaded from a checkpoint with + ``torch.load`` and passed directly to this method. + """ super().load_state_dict(state_dict) if hasattr(self, "_last_lr") and self._last_lr: for param_group, lr in zip(self.optimizer.param_groups, self._last_lr): diff --git a/nvsubquadratic/networks/classification_resnet.py b/nvsubquadratic/networks/classification_resnet.py index bdb5b0c8..352c25e4 100644 --- a/nvsubquadratic/networks/classification_resnet.py +++ b/nvsubquadratic/networks/classification_resnet.py @@ -2,7 +2,32 @@ # Adapted from https://github.com/implicit-long-convs/ccnn_v2 -"""Simple implementation of a ResNet for classification.""" +"""Classification residual network — global-average-pool readout. + +:class:`ClassificationResNet` subclasses :class:`ResidualNetwork` and overrides +:meth:`forward` to add a **global average pooling** step before the output +projection, collapsing all spatial positions into a single per-sample vector: + +.. code-block:: text + + input [B, *spatial, in_channels] + ↓ (inherited) dropout_in, in_proj, blocks → [B, *spatial, hidden_dim] + ↓ GAP: reshape → [B, T, hidden_dim] → mean over T + ↓ out_norm → [B, hidden_dim] + ↓ out_proj → [B, num_classes] + output {"logits": [B, num_classes]} + +The pooling is layout-agnostic: it flattens all spatial dimensions into a +single token axis before taking the mean, so the same network class works for +1-D (sequence), 2-D (image), and 3-D (video / volume) inputs. + +All constructor arguments are inherited from +:class:`~nvsubquadratic.networks.general_purpose_resnet.ResidualNetwork`; +``target_size`` is not used by this subclass (the GAP step replaces the +readout-crop mechanism). + +Adapted from https://github.com/implicit-long-convs/ccnn_v2. +""" import torch @@ -10,40 +35,40 @@ class ClassificationResNet(ResidualNetwork): - """Simple implementation of a ResNet for classification. - - It assumes: - - the input tensor is of shape (batch_size, *spatial_dims, in_channels). - - the output tensor is of shape (batch_size, num_classes). - - Args: - in_channels (int): Number of input channels - out_channels (int): Number of output channels - num_blocks (int): Number of blocks - hidden_dim (int): Number of hidden dimensions - in_proj_cfg (LazyConfig): Configuration for the input projection - out_proj_cfg (LazyConfig): Configuration for the output projection - norm_cfg (LazyConfig): Configuration for the normalization - block_cfg (LazyConfig): Configuration for the residual block - dropout_in_cfg (LazyConfig): Configuration for the dropout in layer (applied to the input) - condition_in_proj_cfg (LazyConfig | None): Configuration for the condition input projection or None if no condition is used. - If provided, the condition tensor is of shape [B, * spatial_dims_condition, hidden_dim]. - If not provided, the condition tensor is None. + """Residual network with global-average-pool readout for classification. + + Inherits the full constructor and backbone from + :class:`~nvsubquadratic.networks.general_purpose_resnet.ResidualNetwork`. + Overrides only :meth:`forward` to replace the spatial output with a + single class-logit vector via global average pooling. + + **Output shape**: ``[B, out_channels]`` regardless of input spatial size — + the model is therefore resolution-agnostic at inference time. + + **No ``target_size``**: the inherited ``target_size`` attribute is ignored; + GAP serves as the spatial aggregation step. + + All constructor arguments are documented in + :class:`~nvsubquadratic.networks.general_purpose_resnet.ResidualNetwork`. + The typical value for ``out_channels`` here is the number of classes. """ - def forward(self, input_and_condition: dict[str, torch.Tensor]) -> torch.Tensor: - """Forward pass of the ClassificationResNet. + def forward(self, input_and_condition: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Run classification forward pass: backbone → GAP → norm → projection. Args: - input_and_condition: A dictionary containing the input and condition. - Keys: "input" and "condition". + input_and_condition: Dictionary with two keys: - - input: Input tensor of shape [B, * spatial_dims, hidden_dim]. - - condition: Condition tensor of shape [B, * spatial_dims_condition, hidden_dim] or None. + * ``"input"`` — signal tensor of shape + ``[B, *spatial, in_channels]``. + * ``"condition"`` — optional conditioning tensor of shape + ``[B, *spatial_cond, hidden_dim]``, or ``None``. Returns: - Dict[str, torch.Tensor]: - - "logits": tensor of shape [B, out_channels]. + dict[str, torch.Tensor]: Single-key dict: + + * ``"logits"`` — shape ``[B, out_channels]`` (one logit vector per + sample, all spatial information collapsed by global average pooling). """ # Extract the input and condition from the dictionary x, condition = input_and_condition["input"], input_and_condition["condition"] diff --git a/nvsubquadratic/networks/general_purpose_resnet.py b/nvsubquadratic/networks/general_purpose_resnet.py index 4a8b6c1a..29594e8d 100644 --- a/nvsubquadratic/networks/general_purpose_resnet.py +++ b/nvsubquadratic/networks/general_purpose_resnet.py @@ -2,7 +2,51 @@ # Adapted from https://github.com/implicit-long-convs/ccnn_v2 -"""Simple implementation of a ResNet for general purpose tasks.""" +"""General-purpose residual network backbone. + +:class:`ResidualNetwork` is a flexible, task-agnostic sequence-of-blocks backbone +used for regression, generation, and spatial-recall tasks. It wires together a +configurable stack of :class:`~nvsubquadratic.modules.residual_block.ResidualBlock` +instances (or any compatible block) via the :mod:`nvsubquadratic.lazy_config` +system, enabling operator swaps (Hyena / Attention / CKConv / Mamba) purely +through config without code changes. + +**Architecture** + +.. code-block:: text + + input [B, *spatial, in_channels] + ↓ dropout_in + ↓ in_proj → [B, *spatial, hidden_dim] + ↓ block × N (each block also receives the optional condition) + ↓ out_norm + ↓ out_proj → [B, *spatial, out_channels] + ↓ readout crop (optional, for spatial-recall tasks) + output {"logits": [B, *target_spatial, out_channels]} + +**Conditioning** + +An optional ``condition_in_proj`` linearly projects an external conditioning +signal (e.g. a class embedding or a global context vector) into ``hidden_dim`` +before it is fed to each block's conditioning branch (e.g. FiLM / cross-attention +in :class:`~nvsubquadratic.modules.residual_block.ResidualBlock`). + +**Readout crop** + +When ``target_size`` is set the network extracts the bottom-right +``target_size`` spatial region of the output before returning. This is used +for spatial-recall tasks where the model must predict a target patch embedded +inside a larger context window. A ``target_size`` element of ``1`` collapses +(squeezes) that spatial dimension, e.g. ``(1, H, W)`` for a 2-D slice of a +3-D input. + +**Gradient checkpointing** + +Set ``gradient_checkpointing=True`` to recompute activations during the backward +pass instead of storing them, trading compute for memory at large scale. + +Adapted from https://github.com/implicit-long-convs/ccnn_v2. +""" from typing import Sequence @@ -14,31 +58,60 @@ class ResidualNetwork(nn.Module): - """Simple implementation of a Residual Network for general purpose tasks. - - It assumes: - - the input tensor is of shape (batch_size, *spatial_dims, in_channels). - - the output tensor is of shape (batch_size, *spatial_dims, out_channels). + """General-purpose residual network backbone (see module docstring for architecture). + + All sub-modules (projections, norm, blocks) are instantiated from + :class:`~nvsubquadratic.lazy_config.LazyConfig` objects so the architecture + can be fully configured from YAML/JSON without subclassing. + + **Tensor layout** + + All tensors are in **channels-last** format: ``[B, *spatial, C]``. The + ``data_dim`` argument records the number of spatial axes (1 for sequences, + 2 for images, 3 for volumes) and is used to convert a scalar ``target_size`` + into a per-axis tuple. + + **Output format** + + :meth:`forward` always returns a ``dict`` with key ``"logits"`` whose value + has shape ``[B, *spatial_out, out_channels]``. When ``target_size=None``, + ``spatial_out = spatial``; otherwise it is the cropped target region. + + Attributes: + in_channels (int): Input channel count. + out_channels (int): Output channel count / number of classes. + num_blocks (int): Number of stacked residual blocks. + hidden_dim (int): Internal feature dimension used throughout the trunk. + data_dim (int): Number of spatial axes (1/2/3). + gradient_checkpointing (bool): Recompute activations on backward. + target_size (tuple | None): Per-axis readout crop size, or ``None``. + dropout_in (nn.Module): Input dropout / augmentation applied first. + in_proj (nn.Module): ``in_channels → hidden_dim`` linear projection. + condition_in_proj (nn.Module | None): Optional ``hidden_dim → hidden_dim`` + projection for the conditioning signal. + blocks (nn.ModuleList): Stack of ``num_blocks`` residual blocks. + out_norm (nn.Module): Post-trunk normalisation (weight-decay excluded). + out_proj (nn.Module): ``hidden_dim → out_channels`` readout projection. Args: - in_channels (int): Number of input channels - out_channels (int): Number of output channels - num_blocks (int): Number of blocks - hidden_dim (int): Number of hidden dimensions - in_proj_cfg (LazyConfig): Configuration for the input projection - out_proj_cfg (LazyConfig): Configuration for the output projection - norm_cfg (LazyConfig): Configuration for the normalization - block_cfg (LazyConfig): Configuration for the residual block - dropout_in_cfg (LazyConfig): Configuration for the dropout in layer (applied to the input) - condition_in_proj_cfg (LazyConfig | None): Configuration for the condition input projection or None if no condition is used. - If provided, the condition tensor is of shape [B, * spatial_dims_condition, hidden_dim]. - If not provided, the condition tensor is None. - target_size (int | tuple | None): Size of the readout region. Can be: - - int: Same size for all spatial dimensions (e.g., 16 for 16x16 in 2D) - - tuple: Different size per dimension. For 3D spatial recall where the target - is a 2D image on the last depth slice, use (1, H, W) to extract only the - last depth slice with HxW spatial region. - - None: No readout extraction (return full output) + in_channels: Number of input signal channels. + out_channels: Number of output channels (e.g. vocabulary / class count). + num_blocks: Depth of the residual tower. + hidden_dim: Width of the residual tower. + data_dim: Spatial dimensionality (1, 2, or 3). + in_proj_cfg: LazyConfig for the input projection (typically ``nn.Linear``). + out_proj_cfg: LazyConfig for the output projection. + norm_cfg: LazyConfig for the output normalisation layer. + block_cfg: LazyConfig for each residual block; instantiated ``num_blocks`` + times. + dropout_in_cfg: LazyConfig for the input dropout layer. + condition_in_proj_cfg: Optional LazyConfig for the condition projection. + Pass ``None`` for unconditional networks. + target_size: Readout crop size. ``int`` → same size on every spatial + axis. ``tuple`` → per-axis sizes (use ``1`` to squeeze that axis). + ``None`` → return the full output. + gradient_checkpointing: Enable activation recomputation in :meth:`forward` + to reduce peak memory at the cost of extra compute. """ def __init__( @@ -57,7 +130,23 @@ def __init__( target_size: int | Sequence[int] | None = None, gradient_checkpointing: bool = False, ): - """Initialize the ResidualNetwork.""" + """Instantiate all sub-modules from LazyConfig objects. + + Args: + in_channels: Number of input signal channels. + out_channels: Number of output channels. + num_blocks: Number of residual blocks to stack. + hidden_dim: Internal feature width. + data_dim: Spatial dimensionality (1 / 2 / 3). + in_proj_cfg: Config for the input projection. + out_proj_cfg: Config for the output projection. + norm_cfg: Config for the output norm layer. + block_cfg: Config for each residual block (instantiated N times). + dropout_in_cfg: Config for input dropout. + condition_in_proj_cfg: Optional config for condition projection. + target_size: Readout crop specification (see class docstring). + gradient_checkpointing: Recompute activations during backward pass. + """ super().__init__() self.in_channels = in_channels self.out_channels = out_channels @@ -136,18 +225,22 @@ def _get_readout_region(self, x: torch.Tensor) -> torch.Tensor: return x[tuple(slices)] def forward(self, input_and_condition: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """Forward pass of the ResidualNetwork. + """Run the full forward pass: project → blocks → norm → project → crop. Args: - input_and_condition: A dictionary containing the input and condition. - Keys: "input" and "condition". + input_and_condition: Dictionary with two keys: - - input: Input tensor of shape [B, * spatial_dims, hidden_dim]. - - condition: Condition tensor of shape [B, * spatial_dims_condition, hidden_dim]. + * ``"input"`` — signal tensor of shape ``[B, *spatial, in_channels]``. + * ``"condition"`` — optional conditioning tensor of shape + ``[B, *spatial_cond, hidden_dim]``, or ``None`` when + ``condition_in_proj_cfg`` was not provided. Returns: - Dict[str, torch.Tensor]: - - "logits": tensor of shape [B, * spatial_dims, out_channels]. + dict[str, torch.Tensor]: Single-key dict: + + * ``"logits"`` — shape ``[B, *spatial_out, out_channels]`` where + ``spatial_out`` equals ``spatial`` unless ``target_size`` is set, + in which case it is the cropped readout region. """ # Extract the input and condition from the dictionary x, condition = input_and_condition["input"], input_and_condition["condition"] From ee472eb7cdbdb150fbf9f8975f9e0fe3d3b4c88b Mon Sep 17 00:00:00 2001 From: David Wessels Date: Tue, 26 May 2026 09:47:23 +0200 Subject: [PATCH 55/72] docs(tracker): mark schedulers, distributed_depthwise_conv_nd, general_purpose_resnet, classification_resnet as done --- docs-tracker.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs-tracker.md b/docs-tracker.md index 8a738606..c22f199b 100644 --- a/docs-tracker.md +++ b/docs-tracker.md @@ -66,16 +66,16 @@ Work bottom-up: primitive ops → modules → networks → experiments. | `drop_path.py` | \[x\] | Stochastic depth — functional + Module, inverted-dropout scaling, causal vs training | | `causal_conv1d.py` | \[x\] | CausalConv1D — left-only pad formula, symmetric mode, Mamba usage context | | `subq_ops_causal_conv1d.py` | \[x\] | New (1D PR): `nn.Conv1d`-compatible depthwise wrapper around `subq_ops.causal_conv1d` | -| `schedulers.py` | \[ \] | LR schedulers | -| `distributed_depthwise_conv_nd.py` | \[ \] | Distributed depthwise conv | +| `schedulers.py` | \[x\] | ResumableSequentialLR — PyTorch ≤2.10 bug fix, load_state_dict LR propagation | +| `distributed_depthwise_conv_nd.py` | \[x\] | CP-aware 1D/2D/3D depthwise conv — group weight sharing, channel slicing, causal padding | | `patch_merging.py` | \[ \] | Pending (feat/patch-merging PR): Swin-style 2×2 patch merging with register-row passthrough | ### `nvsubquadratic/networks/` — Full architectures | File | Status | Notes | | ------------------------------------- | ------ | ----------------------------------------------------------------------------------------------------------------------------------- | -| `general_purpose_resnet.py` | \[ \] | | -| `classification_resnet.py` | \[ \] | | +| `general_purpose_resnet.py` | \[x\] | ResidualNetwork — LazyConfig blocks, conditioning, readout crop, gradient checkpointing | +| `classification_resnet.py` | \[x\] | ClassificationResNet — GAP readout, resolution-agnostic, inherits ResidualNetwork | | `vit5_classification.py` | \[ \] | ViT5 classification head | | `vit5_hierarchical_classification.py` | \[ \] | Pending (feat/patch-merging PR): Swin-style 4-stage hierarchical ViT-5 classifier with GAP readout and optional register-row layout | | `huggingface_diffusers.py` | \[ \] | HF diffusers integration | From 0920a0b5786bc8fd614a87fa6437ce54e6e10c4b Mon Sep 17 00:00:00 2001 From: David Wessels Date: Tue, 26 May 2026 10:31:44 +0200 Subject: [PATCH 56/72] docs(write+integrate/vit5_classification,a2a_comms,lazy_config,cleanfid,qk_norm,quack_utils): add module/class docstrings --- nvsubquadratic/lazy_config.py | 151 +++++++++++++++--- nvsubquadratic/metrics/cleanfid.py | 51 +++++- .../networks/vit5_classification.py | 43 ++++- nvsubquadratic/parallel/a2a_comms.py | 124 ++++++++++++-- nvsubquadratic/utils/qk_norm.py | 84 ++++++++-- nvsubquadratic/utils/quack_utils.py | 16 +- 6 files changed, 426 insertions(+), 43 deletions(-) diff --git a/nvsubquadratic/lazy_config.py b/nvsubquadratic/lazy_config.py index 646dd2af..f738b692 100644 --- a/nvsubquadratic/lazy_config.py +++ b/nvsubquadratic/lazy_config.py @@ -1,7 +1,47 @@ # TODO: Add license header here # Adapted from https://github.com/implicit-long-convs/ccnn_v2 -"""Lazy configuration class for lazy object instantiation.""" +"""Lazy configuration system for deferred object instantiation. + +This module provides a lightweight alternative to Hydra/detectron2-style lazy +configs, designed to let the entire network architecture be specified as a +plain Python or OmegaConf config dict that can be serialised, diffed, and +instantiated on demand — without importing heavy framework dependencies. + +**Core concept** + +A :class:`LazyConfig` holds a reference to a target class or callable together +with its keyword arguments. Calling the object produces an +:class:`~omegaconf.DictConfig` with a ``__target__`` key: + +.. code-block:: python + + cfg = LazyConfig(torch.nn.LayerNorm)(normalized_shape=768, eps=1e-6) + # cfg == DictConfig({"__target__": "torch.nn.LayerNorm", + # "normalized_shape": 768, "eps": 1e-6}) + + norm = instantiate(cfg) # → LayerNorm(768, eps=1e-6) + +:func:`instantiate` resolves ``__target__`` via :func:`importlib.import_module`, +merges any ``**kwargs`` overrides (useful for injecting per-block arguments like +``drop_path_rate``), recursively instantiates nested configs, and evaluates +simple arithmetic strings (e.g. ``"160 * 3"`` → ``480``) found in values. + +**Placeholder values** + +:data:`PLACEHOLDER` is a singleton sentinel that marks fields whose values are +not yet known at config-construction time (e.g. a hidden dimension resolved +later by OmegaConf interpolation). :func:`_contains_placeholder` prevents +premature instantiation of configs that still hold placeholders. + +**Config I/O** + +:func:`save_config` / :func:`load_config` serialise configs to YAML via +OmegaConf. :func:`to_config` reverse-engineers a ``__target__`` dict from an +already-instantiated object (best-effort; works for simple cases). + +Adapted from https://github.com/implicit-long-convs/ccnn_v2. +""" import ast import copy @@ -38,32 +78,70 @@ def __bool__(self): class LazyConfig: - """A lazy configuration class that stores a class/callable reference and its arguments to be instantiated later with an instantiate function. + """Deferred-instantiation config builder. + + A :class:`LazyConfig` stores a reference to a target class or callable. + Calling the instance with keyword arguments produces an OmegaConf + :class:`~omegaconf.DictConfig` that encodes the target and its arguments + under the ``__target__`` key — the standard format consumed by + :func:`instantiate`. + + **Two-step usage**:: + + # Step 1: declare the config (no import of the target needed at this point) + cfg = LazyConfig(torch.nn.Dropout)(p=0.5, inplace=True) + + # Step 2: create the object when ready + dropout = instantiate(cfg) + isinstance(dropout, torch.nn.Dropout) # True - Example: - >>> config = LazyConfig(torch.nn.Dropout)(p=0.5, inplace=True) - >>> module = instantiate(config) - >>> isinstance(module, torch.nn.Dropout) - True + **Nesting** + + LazyConfigs can be nested: passing a ``LazyConfig`` result as a keyword + argument to another ``LazyConfig`` call is supported. :func:`instantiate` + resolves nested configs recursively when ``recursive_instantiate=True``. + + **String targets** + + The ``target`` may be a dotted-path string (e.g. + ``"torch.nn.LayerNorm"``), a class, or any callable that has + ``__module__`` and ``__name__`` attributes. + + Args: + target: The class, callable, or fully-qualified dotted string that + will be used as ``__target__`` in the resulting config dict. + + Example:: + + cfg = LazyConfig(torch.nn.Dropout)(p=0.5, inplace=True) + module = instantiate(cfg) + isinstance(module, torch.nn.Dropout) # True """ def __init__(self, target: Union[Type, Callable, str]): - """Initialize a LazyConfig object with a target class or function. + """Store the target class or callable reference. Args: - target: A class, callable, or string path to a class/function + target: A class, callable, or fully-qualified dotted-path string. """ self.target = target def __call__(self, **kwargs) -> Union[Dict[str, Any], DictConfig]: - """Create a configuration dictionary with __target__ and arguments. + """Produce a config dict containing ``__target__`` and all keyword args. + + Converts ``target`` to a dotted string (``module.ClassName``) and + creates an OmegaConf :class:`~omegaconf.DictConfig` so the result + supports attribute access and OmegaConf interpolation. Args: - **kwargs: Arguments to pass to the target when instantiated + **kwargs: Arguments to forward to the target at instantiation time. + Values may themselves be :class:`LazyConfig` results (nested + configs), plain scalars, or any OmegaConf-compatible type. Returns: - An OmegaConf DictConfig with __target__ and all kwargs, - supporting dot notation access + OmegaConf DictConfig with ``__target__`` set to the resolved + dotted-path string and one key per kwarg. Falls back to a plain + ``dict`` if OmegaConf validation fails (a warning is printed). """ # Convert target to a string if it's a class or function if isinstance(self.target, str): @@ -226,17 +304,50 @@ def instantiate( recursive_instantiate: bool = False, **kwargs, ) -> Any: - """Instantiate an object from a configuration dictionary. + """Instantiate an object from a :class:`LazyConfig` or ``__target__`` dict. + + **Resolution pipeline** + + 1. If ``config`` is a :class:`LazyConfig`, call it with any ``**kwargs`` + to produce a :class:`~omegaconf.DictConfig`. + 2. If the config has no ``__target__`` key, recursively instantiate each + value and return a :class:`~omegaconf.DictConfig` (useful for nested + non-lazy sub-configs). + 3. Resolve ``__target__`` to the actual class/callable via + :func:`importlib.import_module`. + 4. Deep-copy, resolve OmegaConf interpolations, merge ``**kwargs`` + overrides, recursively instantiate nested configs, and evaluate + arithmetic strings (e.g. ``"128 * 4"`` → ``512``). + 5. Call ``target(**processed_args)`` and return the result. + + **Placeholder handling** + + Nested configs that still contain :data:`PLACEHOLDER` values are left + as config dicts rather than instantiated. Args: - config: A dictionary, DictConfig, or LazyConfig object with target and arguments - recursive_instantiate: Whether to instantiate the config recursively. - If True, the config will be instantiated recursively. - If False, the config will be returned as is. - **kwargs: Additional kwargs to override those in config + config: A :class:`LazyConfig` instance, an OmegaConf + :class:`~omegaconf.DictConfig`, or a plain ``dict`` containing at + least a ``"__target__"`` key. + recursive_instantiate: If ``True``, recursively instantiate all nested + configs found in argument values. If ``False`` (default), nested + configs are left as :class:`~omegaconf.DictConfig` objects. + **kwargs: Extra keyword arguments merged into (and overriding) the + config's stored arguments before instantiation. Commonly used to + inject per-block arguments such as ``drop_path_rate``. Returns: - The instantiated object + The instantiated object (result of calling the resolved ``__target__`` + with the processed arguments). + + Raises: + TypeError: If ``target(**processed_args)`` fails; re-raised with the + target's signature in the error message for easier debugging. + + Example:: + + cfg = LazyConfig(torch.nn.LayerNorm)(normalized_shape=768) + norm = instantiate(cfg, eps=1e-5) # override default eps """ _cfg = copy.deepcopy(config) # If it's a LazyConfig object, convert it to a config dict first diff --git a/nvsubquadratic/metrics/cleanfid.py b/nvsubquadratic/metrics/cleanfid.py index 966c97e7..e0284656 100644 --- a/nvsubquadratic/metrics/cleanfid.py +++ b/nvsubquadratic/metrics/cleanfid.py @@ -13,7 +13,30 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Thin wrapper around CleanFID helpers.""" +"""Thin wrapper around CleanFID image generation quality metrics. + +Fréchet Inception Distance (FID) measures the distributional similarity between +a set of generated images and a reference dataset by comparing the mean and +covariance of Inception-v3 feature vectors (Heusel et al., "GANs Trained by a +Two Time-Scale Update Rule Converge to a Local Nash Equilibrium", NeurIPS 2017). + +CleanFID (Parmar et al., "On Aliased Resizing and Surprising Subtleties in GAN +Evaluation", CVPR 2022) corrects for common pre-processing inconsistencies +(e.g. JPEG re-compression, bilinear vs Lanczos resizing) that cause FID scores +to be non-reproducible across libraries. This module delegates to the +``cleanfid`` package, which ships pre-computed reference statistics for standard +benchmarks (FFHQ, CIFAR-10, ImageNet, etc.). + +Usage:: + + score = compute_folder_fid( + sample_dir="outputs/samples/", + dataset_name="imagenet", + dataset_resolution=256, + dataset_split="train", + ) + print(f"FID: {score:.2f}") +""" from __future__ import annotations @@ -29,7 +52,31 @@ def compute_folder_fid( dataset_resolution: int, dataset_split: str = "train", ) -> float: - """Compute FID between a folder of samples and CleanFID reference statistics.""" + """Compute CleanFID between a folder of generated images and a reference dataset. + + Calls ``cleanfid.fid.compute_fid`` with pre-computed reference statistics + for ``dataset_name`` at ``dataset_resolution``, so no reference images need + to be stored locally. + + Args: + sample_dir: Path to the directory containing generated images (PNG/JPEG). + Expanded and resolved to an absolute path before use. + dataset_name: Name of the CleanFID reference dataset, e.g. + ``"imagenet"``, ``"ffhq"``, ``"cifar10"``. Must match a dataset + whose statistics are bundled with or downloaded by ``cleanfid``. + dataset_resolution: Reference image resolution in pixels, e.g. ``256`` + for 256×256 ImageNet. + dataset_split: Which split of the reference dataset to compare against. + Default ``"train"``. + + Returns: + FID score as a Python ``float``. Lower is better; 0 means the + generated and reference distributions are identical under the Inception + feature extractor. + + Raises: + FileNotFoundError: If ``sample_dir`` does not exist. + """ sample_path = Path(sample_dir).expanduser().resolve() if not sample_path.exists(): raise FileNotFoundError(f"Sample directory not found: {sample_path}") diff --git a/nvsubquadratic/networks/vit5_classification.py b/nvsubquadratic/networks/vit5_classification.py index 0692d99e..46ce72c8 100644 --- a/nvsubquadratic/networks/vit5_classification.py +++ b/nvsubquadratic/networks/vit5_classification.py @@ -131,7 +131,48 @@ def __init__( max_drop_path_rate: float = 0.0, drop_path_schedule: Literal["constant", "linear"] = "constant", ): - """Initialize ViT-5 classification network.""" + """Construct patch embedding, positional embeddings, token buffers, and transformer blocks. + + Validates ``readout`` / ``layer_pattern`` constraints, computes the + zero-padding size so ``T % grid_w == 0``, builds the per-layer + ``_block_needs_padding`` flag list, instantiates each block via + :func:`~nvsubquadratic.lazy_config.instantiate` with per-layer + ``drop_path_rate`` and ``register_start_idx`` injected, and initialises + all parameters with truncated-normal (std 0.02). + + Args: + in_channels: Input image channels (3 for RGB). + num_classes: Number of output logits / classes. + hidden_dim: Transformer hidden width ``D``. + num_blocks: Number of transformer blocks ``N``. + patch_size: Non-overlapping patch stride ``P``; patches are ``P×P``. + image_size: Square input resolution ``H = W``. Produces + ``(H/P)²`` patch tokens. + num_registers: Number of learnable register tokens ``R`` appended + after the CLS token. + norm_cfg: LazyConfig for the output normalisation layer. + readout: Token aggregation strategy — ``"cls"``, ``"gap"``, or + ``"register_concat"`` (see class docstring). + block_cfg: Single LazyConfig replicated ``N`` times (homogeneous + mode). Mutually exclusive with ``layer_pattern``. + dropout_rate: Dropout probability applied between the norm and the + classification head. ``0.0`` disables dropout. + neck_compression_ratio: Compression factor for ``register_concat`` + readout; ``neck_dim = hidden_dim // neck_compression_ratio``. + Required when ``readout="register_concat"``. + reg_init: Register-token initialisation — ``"trunc_normal"`` + (std 0.02, default) or ``"zeros"``. + layer_pattern: Per-layer type string of length ``num_blocks`` + (hybrid mode). Each character maps to a key in ``layer_types``. + layer_types: Dict mapping pattern characters to block LazyConfigs. + Required when ``layer_pattern`` is set. + padding_types: Set of pattern characters whose blocks receive the + full padded sequence. Default ``{"H"}`` (Hyena blocks). + max_drop_path_rate: Peak stochastic depth probability distributed + across blocks according to ``drop_path_schedule``. + drop_path_schedule: ``"constant"`` (uniform) or ``"linear"`` + (ramp 0 → ``max_drop_path_rate`` with depth). + """ super().__init__() self._reg_init = reg_init self.hidden_dim = hidden_dim diff --git a/nvsubquadratic/parallel/a2a_comms.py b/nvsubquadratic/parallel/a2a_comms.py index 14bff40a..e96da529 100644 --- a/nvsubquadratic/parallel/a2a_comms.py +++ b/nvsubquadratic/parallel/a2a_comms.py @@ -1,5 +1,64 @@ # TODO: Add license header here +"""All-to-all communication primitives for Context Parallelism (CP). + +This module implements the sequence ↔ channel redistribution pattern used by +the Hyena CP setup, where the *sequence* (spatial) axis is split across CP +ranks during attention/Hyena computation and the *channel* axis is split across +ranks during the FFT convolution. Switching between these two views requires +an ``all_to_all_single`` collective. + +**Communication pattern** + +There are two dual directions: + +- ``"split_to_full"`` — each rank holds a *channel shard* and a *full sequence*; + the collective gathers the full channel dimension while splitting the sequence: + + .. code-block:: text + + input [B, C/P, L ] (channel-split, full sequence per rank) + output [B, C, L/P ] (full channels, sequence-split per rank) + +- ``"full_to_split"`` — each rank holds a *full channel set* and a *sequence shard*; + the collective gathers the full sequence while splitting channels: + + .. code-block:: text + + input [B, C, L/P ] (full channels, sequence-split) + output [B, C/P, L ] (channel-split, full sequence) + +The same logic applies to 2-D (``[B, C, H, W]``) and 3-D (``[B, C, T, H, W]``) +inputs; the *first* spatial axis (``shape[2]``) is always treated as the +sequence/temporal dimension for splitting. + +**Zigzag splitting** + +Plain contiguous splitting can cause load imbalance in causal (autoregressive) +settings because early sequence positions have shorter context than later ones. +Zigzag splitting interleaves chunks so each rank receives an equal mix of early +and late positions: + +.. code-block:: text + + chunk indices for P=2: [0, 3, 1, 2] → rank 0 gets chunks {0, 2}, + rank 1 gets chunks {1, 3} + +The inverse permutation is applied in the ``split_to_full`` direction to restore +the original sequence order. + +**Autograd support** + +:class:`AllToAllSingleFunction` wraps :func:`all_to_all_single_fn` in a custom +``torch.autograd.Function`` so gradients are correctly propagated: the backward +of ``split_to_full`` is ``full_to_split`` and vice versa. + +Public API: + AllToAllSingleFunction: Differentiable all-to-all collective. + all_to_all_single_fn: Non-differentiable functional form (use for inference + or when called from within another custom autograd function). +""" + from typing import Literal import torch @@ -214,16 +273,28 @@ def all_to_all_single_fn( class AllToAllSingleFunction(Function): - """Custom autograd function for all_to_all_single communication with optional zigzag splitting. + """Differentiable all-to-all collective for CP sequence ↔ channel redistribution. + + Wraps :func:`all_to_all_single_fn` in a :class:`torch.autograd.Function` so + that gradients flow correctly through the collective boundary. The backward + pass is the dual communication direction: + + .. code-block:: text + + forward split_to_full → backward full_to_split + forward full_to_split → backward split_to_full + + **Usage**:: + + out = AllToAllSingleFunction.apply(x, cp_group, "split_to_full", True) - A custom autograd function for performing all_to_all_single communication with optional zigzag splitting. - Supports both 1D and 3D tensors. + The ``apply`` arguments correspond to the positional parameters of + :meth:`forward` (excluding ``ctx``). Attributes: - - ctx: A context object that stores information for the forward and backward passes. - - group: The process group for communication. - - type: The type of communication pattern ('split_to_full' or 'full_to_split'). - - with_zigzag_splitting: A boolean indicating whether to apply zigzag splitting (1D only). + ctx.group: Process group saved for the backward collective. + ctx.type: Communication direction saved for reversal in backward. + ctx.with_zigzag_splitting: Zigzag flag saved for backward. """ @staticmethod @@ -233,8 +304,26 @@ def forward( group: dist.ProcessGroup, type: Literal["split_to_full", "full_to_split"], with_zigzag_splitting: bool, - ): - """Forward pass for the AllToAllSingleFunction.""" + ) -> torch.Tensor: + """Execute the all-to-all collective and save state for backward. + + Args: + ctx: Autograd context; stores ``group``, ``type``, and + ``with_zigzag_splitting`` for use in :meth:`backward`. + input_tensor: Input tensor of shape ``[B, C_local, *spatial]``. + The tensor is detached before communication to prevent PyTorch + from tracking in-collective ops. + group: CP process group. + type: ``"split_to_full"`` or ``"full_to_split"`` (see module + docstring for the reshape semantics of each direction). + with_zigzag_splitting: Apply zigzag chunk permutation to balance + load across ranks. Should match the value used in the + corresponding backward call. + + Returns: + torch.Tensor: Redistributed tensor; shape is the dual of the + input under the chosen ``type``. + """ ctx.group = group ctx.type = type ctx.with_zigzag_splitting = with_zigzag_splitting @@ -251,7 +340,22 @@ def forward( @staticmethod def backward(ctx, grad_output: torch.Tensor): - """Backward pass for the AllToAllSingleFunction.""" + """Propagate gradients through the dual all-to-all direction. + + Reverses the communication pattern: ``split_to_full`` ↔ ``full_to_split``. + Zigzag permutation and process group are taken from ``ctx``. + + Args: + ctx: Autograd context with saved ``group``, ``type``, and + ``with_zigzag_splitting``. + grad_output: Upstream gradient tensor; same shape as the forward + output. + + Returns: + Tuple of four elements matching the forward signature: + ``(grad_input, None, None, None)``. Only the first element is + meaningful; the others correspond to non-tensor arguments. + """ # The backward pass will perform the reverse communication grad_input = all_to_all_single_fn( group=ctx.group, diff --git a/nvsubquadratic/utils/qk_norm.py b/nvsubquadratic/utils/qk_norm.py index 9a3fd03e..e610b611 100644 --- a/nvsubquadratic/utils/qk_norm.py +++ b/nvsubquadratic/utils/qk_norm.py @@ -1,17 +1,53 @@ # TODO: Add license header here -"""QK normalization utilities.""" +"""Query-Key (QK) normalisation utilities for attention stabilisation. + +QK-norm (Henry et al., "Query-Key Normalization for Transformers", 2020; +also used in ViT-5 / vit5_attention.py) L2-normalises the query and key +vectors before computing attention logits: + +.. code-block:: text + + q̂ = q / ||q||₂ k̂ = k / ||k||₂ + logits = (q̂ · k̂) / τ # τ is a learned temperature + +By bounding the dot-product to ``[-1, 1]`` (cosine similarity), QK-norm +prevents attention logit explosion in deep or wide models and removes the +need for ``1 / sqrt(d_head)`` scaling (though a learned temperature is +typically kept for flexibility). + +This module provides: + +- :func:`apply_qk_norm`: functional form, normalises a ``(query, key)`` pair. +- :class:`L2Norm`: ``nn.Module`` wrapper suitable as a :class:`LazyConfig` + target; also satisfies the ``channels_first`` duck-type used by norms in + this codebase. +""" import torch import torch.nn.functional as F def apply_qk_norm(query: torch.Tensor, key: torch.Tensor, dim: int = -1, eps: float = 1e-12): - """L2-normalize query and key along the given dimension (e.g. for QK-norm in attention). + """L2-normalise query and key tensors along a given dimension. + + Computes ``F.normalize(q, p=2, dim=dim)`` and the equivalent for ``k``. + The resulting vectors have unit L2 norm along ``dim``, so their dot + product is bounded to ``[-1, 1]`` (cosine similarity). + + Args: + query: Query tensor of shape ``[B, H, T, D]`` (or any layout where + ``dim`` selects the head/feature axis to normalise over). + key: Key tensor; must be broadcast-compatible with ``query``. + dim: Axis to normalise over. Default ``-1`` (last axis = feature + dimension in ``[B, H, T, D]`` layout). + eps: Small constant added to the L2 norm for numerical stability. + Default ``1e-12``. Returns: - Tuple of (query_normalized, key_normalized), same shapes as inputs. + Tuple ``(query_normed, key_normed)`` with the same shapes and dtypes + as the inputs. """ query = F.normalize(query, p=2.0, dim=dim, eps=eps) key = F.normalize(key, p=2.0, dim=dim, eps=eps) @@ -19,25 +55,55 @@ def apply_qk_norm(query: torch.Tensor, key: torch.Tensor, dim: int = -1, eps: fl class L2Norm(torch.nn.Module): - """L2 normalization as a module, for use as a LazyConfig target. + """L2 normalisation layer — learnable-parameter-free, :class:`LazyConfig`-friendly. - Normalizes along the last dimension by default, matching the convention - of torch.nn.RMSNorm and torch.nn.LayerNorm. + Wraps ``F.normalize(x, p=2, dim=self.dim)`` as an ``nn.Module`` so it can + be used as a ``norm_cfg`` target in :func:`~nvsubquadratic.lazy_config.instantiate` + wherever a plain normalisation module is expected (e.g. as the QK-norm in + :class:`~nvsubquadratic.modules.vit5_attention.ViT5Attention`). + + **Duck-typing** + + The :attr:`channels_first` property returns ``True`` when ``dim == 1``, + matching the convention used by + :class:`~nvsubquadratic.modules.rms_norm_channel_first.RMSNormChannelFirst` + so callers can detect the memory layout without an ``isinstance`` check. + + Attributes: + dim (int): Axis to normalise along. + eps (float): Stability constant for the L2 norm denominator. + + Args: + dim: Dimension to normalise over. Default ``-1`` (last axis). + eps: Small positive constant added to the L2 norm. Default ``1e-12``. """ def __init__(self, dim: int = -1, eps: float = 1e-12): - """Store normalization dimension and epsilon.""" + """Initialise L2Norm. + + Args: + dim: Axis to normalise over. Default ``-1``. + eps: Stability constant. Default ``1e-12``. + """ super().__init__() self.dim = dim self.eps = eps @property def channels_first(self) -> bool: - """True when normalizing over dim=1 (channel-first layout).""" + """``True`` when normalising over ``dim=1`` (channel-first layout).""" return self.dim == 1 def forward(self, x: torch.Tensor) -> torch.Tensor: - """L2-normalize along the configured dimension.""" + """L2-normalise ``x`` along ``self.dim``. + + Args: + x: Input tensor of any shape. + + Returns: + torch.Tensor: Unit-norm tensor along ``self.dim``, same shape + and dtype as ``x``. + """ return F.normalize(x, p=2.0, dim=self.dim, eps=self.eps) def extra_repr(self) -> str: diff --git a/nvsubquadratic/utils/quack_utils.py b/nvsubquadratic/utils/quack_utils.py index 9752d959..ff762193 100644 --- a/nvsubquadratic/utils/quack_utils.py +++ b/nvsubquadratic/utils/quack_utils.py @@ -8,7 +8,21 @@ def cuda_supports_quack(device: torch.device) -> bool: - """Return True if *device* supports QuACK kernels (Hopper/Blackwell, SM >= 9).""" + """Return ``True`` if ``device`` supports QuACK fused kernels. + + QuACK kernels require compute capability SM ≥ 9.0 (Hopper: H100; + Blackwell: B200, B300). On older architectures (e.g. Ampere A100, + SM 8.0) the QuACK backward kernel is incompatible and must not be + called; callers should fall back to the PyTorch reference path. + + Args: + device: A ``torch.device`` of type ``"cuda"`` with a device index. + Non-CUDA devices (CPU, MPS) immediately return ``False``. + + Returns: + ``True`` if ``device`` is a CUDA device with SM major version ≥ 9, + ``False`` otherwise. + """ if device.type != "cuda": return False major, _ = torch.cuda.get_device_capability(device) From 60f93fc2cbf31a5e93d9dce023be1c307c83eb93 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Tue, 26 May 2026 10:35:24 +0200 Subject: [PATCH 57/72] docs(write+integrate/experiments): add module docstrings across run, trainer, default_cfg, lightning_wrappers, datamodules, and utils --- experiments/datamodules/mnist.py | 8 ++- experiments/datamodules/tinyimagenet.py | 13 ++++ experiments/default_cfg.py | 47 ++++++++++++- .../base_lightning_wrapper.py | 67 +++++++++++++++++-- .../classification_wrapper.py | 15 ++++- .../lightning_wrappers/diffusion_wrapper.py | 9 ++- .../lightning_wrappers/regression_wrapper.py | 9 ++- experiments/run.py | 37 ++++++++-- experiments/trainer.py | 28 ++++++++ experiments/utils/cli.py | 21 +++++- 10 files changed, 237 insertions(+), 17 deletions(-) diff --git a/experiments/datamodules/mnist.py b/experiments/datamodules/mnist.py index d16a068b..4533b446 100644 --- a/experiments/datamodules/mnist.py +++ b/experiments/datamodules/mnist.py @@ -1,7 +1,13 @@ # TODO: Add license header here -"""MNIST datamodule.""" +"""MNIST / EMNIST datamodule for PyTorch Lightning. + +Wraps :class:`torchvision.datasets.MNIST` (and optionally EMNIST variants) +in a :class:`pytorch_lightning.LightningDataModule`. Handles train/val splits, +channel replication (1 → 3 channels for models that expect RGB), and the +channels-last reshape expected by nvSubquadratic models (``[B, H, W, C]``). +""" from typing import Literal diff --git a/experiments/datamodules/tinyimagenet.py b/experiments/datamodules/tinyimagenet.py index 6a159839..880ad929 100644 --- a/experiments/datamodules/tinyimagenet.py +++ b/experiments/datamodules/tinyimagenet.py @@ -1,3 +1,16 @@ +"""TinyImageNet / ImageNet datamodule backed by Hugging Face Datasets. + +Provides :class:`TinyImageNetDataModule`, a :class:`pytorch_lightning.LightningDataModule` +that loads TinyImageNet (or any ImageNet-style HF dataset) via +``datasets.load_dataset``. Supports: + +- Standard and advanced augmentation pipelines (RandAugment, ThreeAugment, + Mixup, CutMix) via :mod:`timm` and the ``dali_imagenet_fused`` helpers. +- Configurable resolution, crop ratio, and interpolation mode. +- HF token-authenticated access for gated datasets (e.g. full ImageNet on HF). +- Optional label dropping (unsupervised pre-training). +""" + from pathlib import Path from typing import Literal, Optional, Tuple diff --git a/experiments/default_cfg.py b/experiments/default_cfg.py index 4e8b1101..21b2e878 100644 --- a/experiments/default_cfg.py +++ b/experiments/default_cfg.py @@ -2,7 +2,34 @@ # Adapted from https://github.com/implicit-long-convs/ccnn_v2 -"""Default configuration for experiments with nvSubQuadratic.""" +"""Typed configuration dataclasses for nvSubquadratic experiments. + +Every training run is fully described by an :class:`ExperimentConfig` (or one +of its task-specific subclasses such as :class:`DiffusionExperimentConfig`). +These are plain Python :mod:`dataclasses` so they can be instantiated directly +in a Python config file, serialised via OmegaConf, and overridden at the CLI. + +**Sub-configs** + +- :class:`TrainConfig` — batch size, iterations, gradient clip, wall-time limit. +- :class:`TrainerConfig` — Lightning Trainer overrides (validation frequency, + checkpoint interval, DDP settings). +- :class:`SchedulerConfig` — schedule name (``"cosine"``, ``"wsd"``, + ``"constant"``), warmup fraction, total iterations. +- :class:`WandbConfig` — project, entity, run resumption. +- :class:`OptimizerConfig` — optimizer class and hyperparameters. +- :class:`AutoResumeConfig` — automatic checkpoint resumption from local or + W&B artifact. + +Network and Lightning wrapper are specified as +:class:`~nvsubquadratic.lazy_config.LazyConfig` objects (``net_cfg``, +``lightning_wrapper_cfg``) so the full experiment is config-driven. + +:data:`PLACEHOLDER` is re-exported from :mod:`nvsubquadratic.lazy_config` for +convenience in config files. + +Adapted from https://github.com/implicit-long-convs/ccnn_v2. +""" from dataclasses import dataclass, field from typing import Literal, Optional, Union @@ -126,7 +153,23 @@ class StartFromCheckpointConfig: @dataclass class ExperimentConfig: - """Default configuration for experiments with nvSubQuadratic.""" + """Top-level configuration for a single nvSubquadratic training run. + + All fields have sensible defaults; task-specific overrides are specified + in experiment config files under ``experiments/``. The config is loaded + by ``experiments/run.py``, printed as a tree via Rich, and passed to + :func:`~experiments.trainer.construct_trainer` and the Lightning wrapper. + + **Key optional fields** + + - ``compile = True``: wrap the network with ``torch.compile``. Mutually + exclusive with the QuACK kernel path (use ``use_quack=False`` in norm + layers when compiling). + - ``experiment_dir``: override the default ``runs//`` checkpoint + directory with an absolute path. + - ``num_nodes``: number of multi-node training hosts (passed to the + Lightning :class:`~pytorch_lightning.Trainer`). + """ device: str = "cuda" debug: bool = True diff --git a/experiments/lightning_wrappers/base_lightning_wrapper.py b/experiments/lightning_wrappers/base_lightning_wrapper.py index 10eeeff1..bde89057 100644 --- a/experiments/lightning_wrappers/base_lightning_wrapper.py +++ b/experiments/lightning_wrappers/base_lightning_wrapper.py @@ -1,6 +1,29 @@ # Adapted from https://github.com/implicit-long-convs/ccnn_v2 -"""Lightning wrappers for the Classification and Regression experiments.""" +"""Base PyTorch Lightning wrapper for all nvSubquadratic experiments. + +Provides :class:`LightningWrapperBase`, the shared superclass for all task-specific +Lightning modules (classification, regression, diffusion, autoregressive, WELL, +ARC). It handles: + +- **Optimizer construction**: parameter-group splitting (weight-decay, per-param + LR scale via ``_lr_scale``), dispatching to :func:`_build_optimizer`. +- **Scheduler construction**: warmup + main schedule chaining via + :class:`~nvsubquadratic.modules.schedulers.ResumableSequentialLR`. +- **Checkpoint resume**: key-remapping for compiled vs. non-compiled state dicts + via :func:`align_compiled_keys`. +- **Timing profiling**: optional CUDA-event-based forward/backward profiling + logged to W&B. +- **Gradient norm tracking**: optional per-layer or global grad-norm logging. +- **FLOPs accounting**: calls ``network.flop_count()`` if available and logs + the result as a W&B summary metric. + +Task-specific forward passes, losses, and metrics live in the subclasses +(:class:`~experiments.lightning_wrappers.classification_wrapper.ClassificationWrapper`, +:class:`~experiments.lightning_wrappers.regression_wrapper.RegressionWrapper`, etc.). + +Adapted from https://github.com/implicit-long-convs/ccnn_v2. +""" import warnings @@ -286,18 +309,52 @@ def construct_scheduler( class LightningWrapperBase(pl.LightningModule): - """Base Lightning wrapper class.""" + """Base PyTorch Lightning module shared by all nvSubquadratic task wrappers. + + Handles everything that is common across tasks: optimizer/scheduler + construction, checkpoint resume key alignment, optional CUDA profiling, + gradient norm tracking, and FLOP logging. Subclasses implement: + + - ``training_step`` / ``validation_step`` / ``test_step`` + - Task-specific loss computation and metric logging + + **Parameter grouping** (see :func:`_build_param_groups`): + + Parameters tagged ``_no_weight_decay = True`` are placed in a zero-decay + group. Parameters tagged ``_lr_scale = `` receive a per-parameter + LR multiplier applied by scaling the base LR before passing the group to + the optimizer. + + **Scheduler** + + :func:`_build_lr_scheduler` chains a linear-warmup + ``LinearLR`` with the main schedule (cosine, WSD, constant) via + :class:`~nvsubquadratic.modules.schedulers.ResumableSequentialLR`, + which fixes the PyTorch ≤ 2.10 checkpoint-resume LR bug. + + Attributes: + network (torch.nn.Module): The wrapped model. + optimizer_cfg: Optimizer config from :class:`~experiments.default_cfg.ExperimentConfig`. + scheduler_cfg: Scheduler config from :class:`~experiments.default_cfg.ExperimentConfig`. + distributed (bool): ``True`` when more than one GPU is visible. + + Args: + network: The neural network to train. + cfg: Full experiment configuration. + """ def __init__( self, network: torch.nn.Module, cfg: ExperimentConfig, ): - """Initialize the LightningWrapperBase. + """Initialise the wrapper and log parameter count and FLOPs. Args: - network: Network to wrap. - cfg: Configuration. + network: The neural network to train. + cfg: Full experiment configuration; ``cfg.optimizer`` and + ``cfg.scheduler`` are stored for later use in + :meth:`configure_optimizers`. """ super().__init__() # Define network diff --git a/experiments/lightning_wrappers/classification_wrapper.py b/experiments/lightning_wrappers/classification_wrapper.py index bdd51a35..5d0f8404 100644 --- a/experiments/lightning_wrappers/classification_wrapper.py +++ b/experiments/lightning_wrappers/classification_wrapper.py @@ -1,6 +1,19 @@ # Adapted from https://github.com/implicit-long-convs/ccnn_v2 -"""Lightning wrappers for the Classification and Regression experiments.""" +"""Lightning wrapper for image classification tasks. + +Provides :class:`ClassificationWrapper` and :class:`SoftTargetCrossEntropy`. + +**Loss modes** + +- ``"cross_entropy"`` — standard :class:`torch.nn.CrossEntropyLoss` (hard labels). +- ``"soft_target_ce"`` — :class:`SoftTargetCrossEntropy`: + ``-sum(target * log_softmax(logits))``. Use with Mixup/CutMix (DeiT III recipe). +- ``"bce"`` — :class:`torch.nn.BCEWithLogitsLoss` with binarised multi-hot targets. + Matches the ViT-5 / DeiT III pre-training recipe (``--bce-loss``). + +Adapted from https://github.com/implicit-long-convs/ccnn_v2. +""" import torch import torch.nn.functional as F diff --git a/experiments/lightning_wrappers/diffusion_wrapper.py b/experiments/lightning_wrappers/diffusion_wrapper.py index d861c168..9b8fa045 100644 --- a/experiments/lightning_wrappers/diffusion_wrapper.py +++ b/experiments/lightning_wrappers/diffusion_wrapper.py @@ -1,6 +1,13 @@ # Adapted from https://github.com/implicit-long-convs/ccnn_v2 -"""Lightning wrappers for the Classification and Regression experiments.""" +"""Lightning wrapper for continuous-time diffusion (JiT-style) experiments. + +Provides :class:`DiffusionWrapper`, which implements the flow-matching / JiT +training loop: noises inputs according to a time-dependent schedule, trains a +denoiser network, and generates samples via an ODE integrator at inference time. + +Adapted from https://github.com/implicit-long-convs/ccnn_v2. +""" import copy import math diff --git a/experiments/lightning_wrappers/regression_wrapper.py b/experiments/lightning_wrappers/regression_wrapper.py index 3d968eec..9e002092 100644 --- a/experiments/lightning_wrappers/regression_wrapper.py +++ b/experiments/lightning_wrappers/regression_wrapper.py @@ -1,6 +1,13 @@ # Adapted from https://github.com/implicit-long-convs/ccnn_v2 -"""Lightning wrappers for the Classification and Regression experiments.""" +"""Lightning wrapper for regression tasks (MAE / MSE loss). + +Provides :class:`RegressionWrapper`, which supports both Mean Absolute Error +(L1) and Mean Squared Error (L2) regression objectives. It is also the base +class for :class:`~experiments.lightning_wrappers.well_lightning_wrapper.WELLRegressionWrapper`. + +Adapted from https://github.com/implicit-long-convs/ccnn_v2. +""" from typing import Literal diff --git a/experiments/run.py b/experiments/run.py index d26701f7..ae859544 100644 --- a/experiments/run.py +++ b/experiments/run.py @@ -2,11 +2,38 @@ # Adapted from https://github.com/implicit-long-convs/ccnn_v2 -"""Entry point to run experiments. - -Usage: - # MNIST classification - PYTHONPATH=. python nvsubquadratic/examples/run.py --config examples/mnist_classification/experiments/mnist_classification_ccnn_4_160_hyena_rope_qknorm.py +r"""Main entry point for launching nvSubquadratic training runs. + +**Usage**:: + + PYTHONPATH=. python experiments/run.py \\ + --config experiments/mnist_classification/mnist_classification_ccnn_4_160.py \\ + dataset.batch_size=64 optimizer.lr=3e-4 + +Command-line interface: + +- ``--config `` (required): path to a Python config file that defines a + top-level ``cfg`` variable of type + :class:`~experiments.default_cfg.ExperimentConfig`. +- ``=`` (zero or more positional overrides): dotted key-value pairs + that override fields in ``cfg`` after the file is loaded. OmegaConf + interpolators (``${...}``) are respected. + +**What the script does** + +1. Parses ``--config`` and ``overrides`` via :func:`parse_args`. +2. Loads the config file with :func:`~experiments.utils.cli.load_config_from_file` + and applies overrides via :func:`~experiments.utils.cli.apply_config_overrides`. +3. Initialises a W&B run (or resumes one) and constructs a + :class:`pytorch_lightning.loggers.WandbLogger`. +4. Instantiates the network from ``cfg.net_cfg`` via + :func:`~nvsubquadratic.lazy_config.instantiate`. +5. Optionally loads a pre-trained checkpoint (partial or full). +6. Instantiates the Lightning wrapper from ``cfg.lightning_wrapper_cfg``. +7. Calls :func:`~experiments.trainer.construct_trainer` and + :meth:`pytorch_lightning.Trainer.fit`. + +Adapted from https://github.com/implicit-long-convs/ccnn_v2. """ import argparse diff --git a/experiments/trainer.py b/experiments/trainer.py index 76cdcdf6..f279d5d3 100644 --- a/experiments/trainer.py +++ b/experiments/trainer.py @@ -1,6 +1,34 @@ # TODO: Add licence header # Adapted from https://github.com/implicit-long-convs/ccnn_v2 + +"""PyTorch Lightning trainer factory for nvSubquadratic experiments. + +:func:`construct_trainer` is the single entry point for building a +:class:`pytorch_lightning.Trainer` from an +:class:`~experiments.default_cfg.ExperimentConfig`. It: + +1. Resolves the checkpoint metric (``val/acc`` or ``val/loss``) from the + scheduler ``mode`` or an explicit override. +2. Creates the checkpoint directory under ``runs//`` or a provided + ``experiment_dir``. +3. Constructs a :class:`~pytorch_lightning.callbacks.ModelCheckpoint` callback + that saves the ``best`` and ``last`` checkpoints. +4. Optionally attaches a + :class:`~experiments.callbacks.walltime_checkpointer.WalltimeCheckpointer` + (saves and exits when the SLURM wall-time limit is near). +5. Optionally attaches a + :class:`~experiments.utils.checkpointing.WandbSelectiveCheckpointUploader` + (uploads best checkpoints per scheduler phase to W&B) and a + :class:`~experiments.callbacks.wandb_cache_cleanup.WandbCacheCleanupCallback`. +6. Sets precision, determinism, and gradient-clip flags from config. +7. Returns the fully configured :class:`pytorch_lightning.Trainer` together with + the :class:`~pytorch_lightning.callbacks.ModelCheckpoint` callback (so the + caller can inspect ``checkpoint_callback.best_model_path``). + +Adapted from https://github.com/implicit-long-convs/ccnn_v2. +""" + from pathlib import Path from typing import Optional diff --git a/experiments/utils/cli.py b/experiments/utils/cli.py index 32a2302a..d621f898 100644 --- a/experiments/utils/cli.py +++ b/experiments/utils/cli.py @@ -1,7 +1,26 @@ # TODO: Add license header here -"""Utility functions for the MNIST classification experiment.""" +"""CLI and config-loading utilities for nvSubquadratic experiments. + +Provides helpers used by ``experiments/run.py`` to load Python config files, +apply command-line overrides, generate deterministic W&B run names, and render +config trees via Rich. + +Key functions: + +- :func:`load_config_from_file` — import a ``.py`` config file and return its + ``cfg`` variable as an :class:`~experiments.default_cfg.ExperimentConfig`. +- :func:`apply_config_overrides` — parse ``key=value`` CLI overrides and set + the corresponding fields on the config dataclass (supports nested dotted + paths and OmegaConf interpolation strings). +- :func:`get_deterministic_run_name` — derive a human-readable W&B run name + from the config filename, timestamp, and any CLI overrides. +- :func:`config_to_dict` / :func:`add_to_tree` — serialise the config to a + dict and render it as a :class:`rich.tree.Tree` for console display. +- :func:`verify_no_interpolator_overwrites` — sanity-check that CLI overrides + do not silently clobber OmegaConf ``${...}`` interpolations. +""" import dataclasses import datetime From d2ebc3f7194ae9e9b5452728c92a9d4a0eaa3d6a Mon Sep 17 00:00:00 2001 From: David Wessels Date: Tue, 26 May 2026 10:36:39 +0200 Subject: [PATCH 58/72] docs(tracker): mark vit5_classification, a2a_comms, lazy_config, metrics, utils, testing, and experiments as done --- docs-tracker.md | 62 ++++++++++++++++++++++++------------------------- 1 file changed, 31 insertions(+), 31 deletions(-) diff --git a/docs-tracker.md b/docs-tracker.md index c22f199b..7e7bf62a 100644 --- a/docs-tracker.md +++ b/docs-tracker.md @@ -76,46 +76,46 @@ Work bottom-up: primitive ops → modules → networks → experiments. | ------------------------------------- | ------ | ----------------------------------------------------------------------------------------------------------------------------------- | | `general_purpose_resnet.py` | \[x\] | ResidualNetwork — LazyConfig blocks, conditioning, readout crop, gradient checkpointing | | `classification_resnet.py` | \[x\] | ClassificationResNet — GAP readout, resolution-agnostic, inherits ResidualNetwork | -| `vit5_classification.py` | \[ \] | ViT5 classification head | +| `vit5_classification.py` | \[x\] | ViT5 classification — token layout, hybrid blocks, CLS/GAP/register_concat readouts, FLOP count | | `vit5_hierarchical_classification.py` | \[ \] | Pending (feat/patch-merging PR): Swin-style 4-stage hierarchical ViT-5 classifier with GAP readout and optional register-row layout | | `huggingface_diffusers.py` | \[ \] | HF diffusers integration | | `jit.py` / `jit_utils.py` | \[ \] | TorchScript utilities | ### `nvsubquadratic/parallel/` — Distributed primitives -| File | Status | Notes | -| -------------- | ------ | ---------------- | -| `a2a_comms.py` | \[ \] | All-to-all comms | +| File | Status | Notes | +| -------------- | ------ | ---------------------------------------------------------------------------------------- | +| `a2a_comms.py` | \[x\] | AllToAllSingle — CP sequence↔channel redistribution, zigzag splitting, autograd backward | ### `nvsubquadratic/` — Top-level -| File | Status | Notes | -| ---------------- | ------ | ------------------ | -| `lazy_config.py` | \[ \] | Lazy config system | -| `metrics/` | \[ \] | Metric utilities | -| `utils/` | \[ \] | General utilities | -| `testing/` | \[ \] | Testing utilities | +| File | Status | Notes | +| ---------------- | ------ | ------------------------------------------------------------------------------------- | +| `lazy_config.py` | \[x\] | LazyConfig + instantiate — deferred instantiation, nested configs, arithmetic strings | +| `metrics/` | \[x\] | cleanfid.py — CleanFID wrapper, FID formula, usage context | +| `utils/` | \[x\] | init.py (weight init factories), qk_norm.py (QK-norm, L2Norm module), quack_utils.py | +| `testing/` | \[x\] | utils.py — compute_relative_error, TTrace reference, already had good docstrings | ### `experiments/` — Training infrastructure -| File | Status | Notes | -| ---------------------------------------------- | ------ | -------------- | -| `run.py` | \[ \] | Entry point | -| `trainer.py` | \[ \] | Training loop | -| `default_cfg.py` | \[ \] | Default config | -| `lightning_wrappers/base_lightning_wrapper.py` | \[ \] | Base wrapper | -| `lightning_wrappers/classification_wrapper.py` | \[ \] | | -| `lightning_wrappers/regression_wrapper.py` | \[ \] | | -| `lightning_wrappers/diffusion_wrapper.py` | \[ \] | | -| `lightning_wrappers/autoregressive_wrapper.py` | \[ \] | | -| `lightning_wrappers/arc_wrapper.py` | \[ \] | | -| `lightning_wrappers/well_lightning_wrapper.py` | \[ \] | | -| `datamodules/arc.py` | \[ \] | | -| `datamodules/mnist.py` | \[ \] | | -| `datamodules/tinyimagenet.py` | \[ \] | | -| `datamodules/ucf101.py` | \[ \] | | -| `datamodules/dali_imagenet_fused.py` | \[ \] | | -| `datamodules/spatial_recall_dataset.py` | \[ \] | | -| `utils/cli.py` | \[ \] | | -| `utils/checkpointing.py` | \[ \] | | -| `callbacks/` | \[ \] | | +| File | Status | Notes | +| ---------------------------------------------- | ------ | -------------------------------------------------------------------------------- | +| `run.py` | \[x\] | Entry point — CLI parse, W&B init, network + wrapper instantiation, Trainer.fit | +| `trainer.py` | \[x\] | construct_trainer — checkpoint callbacks, precision, wall-time, W&B upload | +| `default_cfg.py` | \[x\] | Typed dataclass configs — Train/Trainer/Scheduler/Wandb/Optimizer/AutoResume | +| `lightning_wrappers/base_lightning_wrapper.py` | \[x\] | LightningWrapperBase — param groups, scheduler, checkpoint resume, profiling | +| `lightning_wrappers/classification_wrapper.py` | \[x\] | ClassificationWrapper — cross_entropy / soft_target_ce / bce loss modes | +| `lightning_wrappers/regression_wrapper.py` | \[x\] | RegressionWrapper — MAE/MSE loss, base for WELLRegressionWrapper | +| `lightning_wrappers/diffusion_wrapper.py` | \[x\] | DiffusionWrapper — JiT continuous-time diffusion, ODE sampler | +| `lightning_wrappers/autoregressive_wrapper.py` | \[x\] | Already had good module + class docstrings; left as-is | +| `lightning_wrappers/arc_wrapper.py` | \[ \] | (untracked new file — out of scope until merged) | +| `lightning_wrappers/well_lightning_wrapper.py` | \[x\] | Already had good class docstring; left as-is | +| `datamodules/arc.py` | \[ \] | (untracked new file — out of scope until merged) | +| `datamodules/mnist.py` | \[x\] | MNIST/EMNIST datamodule — channels-last reshape, train/val split | +| `datamodules/tinyimagenet.py` | \[x\] | TinyImageNet HF-backed datamodule — RandAugment, Mixup/CutMix, token access | +| `datamodules/ucf101.py` | \[ \] | Video classification datamodule | +| `datamodules/dali_imagenet_fused.py` | \[ \] | DALI ImageNet pipeline | +| `datamodules/spatial_recall_dataset.py` | \[x\] | Already had comprehensive module + class docstrings; left as-is | +| `utils/cli.py` | \[x\] | CLI helpers — load_config_from_file, apply_config_overrides, run name generation | +| `utils/checkpointing.py` | \[x\] | Already had good per-function docstrings; left as-is | +| `callbacks/` | \[ \] | Walltime checkpointer, W&B cache cleanup | From a3b5f62d1c8b07ef154e9dd8e344c41c45c8cdda Mon Sep 17 00:00:00 2001 From: David Wessels Date: Tue, 26 May 2026 10:55:58 +0200 Subject: [PATCH 59/72] docs(write+integrate/callbacks,ucf101,dali_imagenet_fused): add module docstrings and expand MixupConfig/AugmentConfig --- .../callbacks/walltime_checkpointer.py | 21 ++++++++++ .../datamodules/dali_imagenet_fused.py | 39 ++++++++++++++++++- experiments/datamodules/ucf101.py | 24 +++++++++++- 3 files changed, 81 insertions(+), 3 deletions(-) diff --git a/experiments/callbacks/walltime_checkpointer.py b/experiments/callbacks/walltime_checkpointer.py index a9d4170c..c90f9ef3 100644 --- a/experiments/callbacks/walltime_checkpointer.py +++ b/experiments/callbacks/walltime_checkpointer.py @@ -1,3 +1,24 @@ +"""Wall-time checkpointer callback for graceful SLURM preemption handling. + +SLURM jobs have a hard wall-time limit. If training reaches the limit without +saving, the job is killed and all progress since the last checkpoint is lost. +:class:`WalltimeCheckpointer` solves this by monitoring elapsed time in +``on_train_batch_end`` and triggering a checkpoint + ``trainer.should_stop`` +a configurable number of seconds before the deadline. + +**Usage in experiment configs**:: + + from experiments.callbacks.walltime_checkpointer import WalltimeCheckpointer + import time + + cfg.train.run_start_time = time.time() + cfg.train.run_time_limit_hours = 23.5 # slightly under the SLURM limit + +The :func:`~experiments.trainer.construct_trainer` factory reads these two +fields and wires up :class:`WalltimeCheckpointer` automatically — no manual +callback construction is needed in most experiment configs. +""" + import datetime import time from pathlib import Path diff --git a/experiments/datamodules/dali_imagenet_fused.py b/experiments/datamodules/dali_imagenet_fused.py index db1a8aca..881365a8 100644 --- a/experiments/datamodules/dali_imagenet_fused.py +++ b/experiments/datamodules/dali_imagenet_fused.py @@ -50,7 +50,22 @@ @dataclass class MixupConfig: - """Configuration for mixup.""" + """Mixup / CutMix configuration for the DALI ImageNet pipeline. + + Controls the ``timm.data.Mixup`` augmentation applied in + ``on_before_batch_transfer`` (after DALI decoding, since Mixup needs + labels). Set both ``mixup=0`` and ``cutmix=0`` to disable entirely. + + Attributes: + mixup: Alpha parameter for Mixup (Beta distribution). ``0`` disables. + cutmix: Alpha parameter for CutMix. ``0`` disables. + mixup_prob: Probability of applying either Mixup or CutMix per batch. + mixup_switch_prob: Probability of switching from Mixup to CutMix when + both are enabled. + mixup_mode: Granularity of the mix — ``"batch"`` (same λ per batch), + ``"pair"`` (per sample pair), or ``"elem"`` (per element). + smoothing: Label smoothing epsilon applied to one-hot targets. + """ mixup: float = 0.0 cutmix: float = 0.0 @@ -62,7 +77,27 @@ class MixupConfig: @dataclass class AugmentConfig: - """Configuration for augmentations.""" + """Augmentation configuration for the DALI ImageNet pipeline. + + All augmentations except Mixup/CutMix run inside the DALI pipeline on the + GPU, eliminating the serial CPU augmentation bottleneck. + + Attributes: + use_three_augment: Apply ThreeAugment (grayscale, solarise, Gaussian + blur) from DeiT III instead of standard color jitter. + color_jitter: Strength of random color jitter (brightness, contrast, + saturation). Ignored when ``use_three_augment=True``. + rand_augment: RandAugment spec string passed to + ``timm.data.auto_augment.rand_augment_transform``, e.g. + ``"rand-m9-mstd0.5-inc1"``. ``None`` disables RandAugment. + random_erasing_prob: Probability of random erasing (occlusion + augmentation). ``0.0`` disables. + random_erasing_mode: Fill mode for erased patches — ``"pixel"`` + (random noise) or ``"const"`` (zero). + num_repeats: DeiT-style repeated augmentation factor. Each image is + sampled ``num_repeats`` times per epoch with independent + augmentation seeds, effectively multiplying the dataset size. + """ use_three_augment: bool = False color_jitter: float = 0.4 diff --git a/experiments/datamodules/ucf101.py b/experiments/datamodules/ucf101.py index 1dbc5100..db4600aa 100644 --- a/experiments/datamodules/ucf101.py +++ b/experiments/datamodules/ucf101.py @@ -1,7 +1,29 @@ # TODO: Add license header here -"""UCF101 datamodule with in-code download in prepare_data.""" +"""UCF101 video-classification datamodule for PyTorch Lightning. + +Provides :class:`UCF101DataModule`, which wraps +:class:`torchvision.datasets.UCF101`. :meth:`prepare_data` automatically +downloads and extracts both the video archive and the train/test annotation +split files if they are not already present under ``data_dir`` / +``annotation_dir``. + +**Two output modes** (selected via ``data_type``): + +- ``"video"`` — returns clips shaped ``[B, T, H, W, C]`` (channels-last, + float32). Each clip spans ``frames_per_clip`` frames sampled at + ``frame_rate`` fps with ``step_between_clips`` temporal stride. +- ``"sequence"`` — flattens the temporal and spatial axes into a 1-D token + sequence: ``[B, T*H*W, C]``. Suitable for models that treat video as a + flat sequence of patches. + +**Deterministic workers** + +When ``use_deterministic_worker_init=True`` each DataLoader worker is seeded +with ``base_seed + worker_id`` via :func:`deterministic_worker_init_fn`, +ensuring reproducible random clip sampling across runs and restarts. +""" import os import random From 87f8609db8de0852f260f83f106d33a5eaddb5f8 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Tue, 26 May 2026 10:57:31 +0200 Subject: [PATCH 60/72] =?UTF-8?q?docs(tracker):=20mark=20callbacks,=20ucf1?= =?UTF-8?q?01,=20dali=5Fimagenet=5Ffused=20as=20done=20=E2=80=94=20documen?= =?UTF-8?q?tation=20complete?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs-tracker.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs-tracker.md b/docs-tracker.md index 7e7bf62a..0d436b90 100644 --- a/docs-tracker.md +++ b/docs-tracker.md @@ -113,9 +113,9 @@ Work bottom-up: primitive ops → modules → networks → experiments. | `datamodules/arc.py` | \[ \] | (untracked new file — out of scope until merged) | | `datamodules/mnist.py` | \[x\] | MNIST/EMNIST datamodule — channels-last reshape, train/val split | | `datamodules/tinyimagenet.py` | \[x\] | TinyImageNet HF-backed datamodule — RandAugment, Mixup/CutMix, token access | -| `datamodules/ucf101.py` | \[ \] | Video classification datamodule | -| `datamodules/dali_imagenet_fused.py` | \[ \] | DALI ImageNet pipeline | +| `datamodules/ucf101.py` | \[x\] | UCF101 — video/sequence modes, frames_per_clip, deterministic workers | +| `datamodules/dali_imagenet_fused.py` | \[x\] | DALI ImageNet — fused GPU augmentation, MixupConfig/AugmentConfig, repeated aug | | `datamodules/spatial_recall_dataset.py` | \[x\] | Already had comprehensive module + class docstrings; left as-is | | `utils/cli.py` | \[x\] | CLI helpers — load_config_from_file, apply_config_overrides, run name generation | | `utils/checkpointing.py` | \[x\] | Already had good per-function docstrings; left as-is | -| `callbacks/` | \[ \] | Walltime checkpointer, W&B cache cleanup | +| `callbacks/` | \[x\] | walltime_checkpointer (added module docstring); all others already had good docs | From 0d0b9e0bee78273b5c6fb4c7f8ff8b59fab50e98 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Tue, 26 May 2026 10:58:45 +0200 Subject: [PATCH 61/72] docs: add CONVENTIONS.md with docstring style guide and PR enforcement strategy --- CONVENTIONS.md | 128 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 128 insertions(+) create mode 100644 CONVENTIONS.md diff --git a/CONVENTIONS.md b/CONVENTIONS.md new file mode 100644 index 00000000..ea4c6760 --- /dev/null +++ b/CONVENTIONS.md @@ -0,0 +1,128 @@ +# Documentation Conventions + +## Goal + +External collaborators should be able to read any module in `nvsubquadratic/` +or `experiments/` alongside the research paper and understand: + +1. **What** the class/function does and **why** it exists. +1. The **math** it implements (notation from the paper where applicable). +1. The **shape contract** of every tensor argument. +1. How it fits into the larger architecture. + +______________________________________________________________________ + +## Docstring style + +All public classes, functions, and methods use **Google-style docstrings**. + +```python +def forward(self, x: torch.Tensor) -> torch.Tensor: + """Apply RMS normalisation over the last dimension. + + Args: + x: Input tensor of shape ``[*leading, D]``. + + Returns: + torch.Tensor: Normalised tensor, same shape and dtype as ``x``. + + Raises: + RuntimeError: If ``x.shape[-1]`` does not equal ``self.dim``. + """ +``` + +### Rules + +| What | Where | +| --------------------------------------- | ---------------------------------------------------------- | +| Math / motivation / paper context | **Module docstring** or **class docstring** | +| Parameter descriptions + shapes | `Args:` block on the method | +| Return shape + dtype | `Returns:` block | +| Docstrings containing `\` (backslashes) | Use `r"""..."""` (required by ruff D301) | +| Single-line `__init__` docstrings | Allowed only when the class docstring covers all arguments | + +______________________________________________________________________ + +## PR convention + +> **Any PR that adds or modifies a public class or function in `nvsubquadratic/` +> or `experiments/` must update the docstring of every touched symbol.** + +This means: + +- New file → full module docstring + docstrings on all public classes/functions. +- Existing file touched → update only the affected symbols; do not leave + neighbouring docstrings stale. +- Renamed parameter → rename in the docstring too (never paper over a + misleading name with a docstring explanation). + +______________________________________________________________________ + +## Automated enforcement + +### 1. ruff (already active) + +The pre-commit hook runs `ruff` with the `D` rule-set enabled. It catches: + +- `D100` Missing docstring in public module +- `D101` Missing docstring in public class +- `D102` Missing docstring in public method +- `D103` Missing docstring in public function +- `D301` Use `r"""` if backslashes appear in a docstring + +If your commit is rejected by ruff with a `D` error, add or fix the docstring +before pushing. + +### 2. CI diff-check (recommended addition) + +Add the following job to `.github/workflows/ci.yml` (or your equivalent): + +```yaml +- name: Docstring coverage on changed files + run: | + # Collect Python files touched by this PR + git diff --name-only origin/main...HEAD \ + | grep -E '^(nvsubquadratic|experiments)/.*\.py$' \ + > changed.txt + + # Fail if any changed file has missing public docstrings + if [ -s changed.txt ]; then + xargs ruff check --select D100,D101,D102,D103,D417 < changed.txt + fi +``` + +This restricts the `D` check to **files actually changed in the PR**, so it +never blocks unrelated legacy files. + +### 3. CONVENTIONS.md review checklist + +Add a pull-request template (`.github/pull_request_template.md`) that +reminds authors: + +```markdown +## Documentation checklist + +- [ ] Every new public class has a module-level or class-level docstring + explaining *what it does* and *why*. +- [ ] Every new public method/function has Args: and Returns: blocks with + tensor shapes. +- [ ] Math notation is consistent with the paper (or a comment explains any + difference). +- [ ] Docstrings containing backslashes use `r"""..."""`. +``` + +### 4. Updating `docs-tracker.md` + +When a PR introduces a **new file** in scope (`nvsubquadratic/` or +`experiments/`), add a row to the relevant table in `docs-tracker.md` with +status `[x]` and a one-line note. This keeps the tracker current without a +separate documentation PR. + +______________________________________________________________________ + +## What does *not* need a docstring + +- Private helpers prefixed with `_` (encouraged but not required). +- `__repr__` / `extra_repr` methods — a brief one-liner is enough. +- Test files under `tests/` — descriptive test names suffice. +- Config files under `experiments/` that contain only a `cfg = ...` assignment. From e1147065ce6b91d61115110047a2cb5e4305e982 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Tue, 26 May 2026 11:07:38 +0200 Subject: [PATCH 62/72] docs: add docstring CI workflow, extend PR template with doc checklist, expand README docs section --- .github/pull_request_template.md | 13 +++++++ .github/workflows/docstring_check.yml | 49 +++++++++++++++++++++++++++ README.md | 38 +++++++++++++++++++-- 3 files changed, 98 insertions(+), 2 deletions(-) create mode 100644 .github/workflows/docstring_check.yml diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 17f3b9b3..69e0c00e 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -23,3 +23,16 @@ pre-commit run --all-files ## Test plan + +## Documentation checklist + +For every new or modified public symbol in `nvsubquadratic/` or `experiments/`: + +- [ ] Every new **module** has a module-level docstring explaining what it contains and why. +- [ ] Every new **public class** has a class docstring covering purpose, math/motivation, and key attributes. +- [ ] Every new **public method / function** has `Args:` and `Returns:` blocks with tensor shapes where applicable. +- [ ] Math notation is consistent with the paper (or a comment explains any deviation). +- [ ] Docstrings containing backslashes use `r"""..."""` (required by ruff D301). +- [ ] If a new file was added, a row has been added to [`docs-tracker.md`](../docs-tracker.md) with status `[x]`. + +> See [`CONVENTIONS.md`](../CONVENTIONS.md) for the full style guide. diff --git a/.github/workflows/docstring_check.yml b/.github/workflows/docstring_check.yml new file mode 100644 index 00000000..6bc26456 --- /dev/null +++ b/.github/workflows/docstring_check.yml @@ -0,0 +1,49 @@ +name: Docstring coverage + +on: + pull_request: + +# Cancel in-flight runs for the same PR when a new commit is pushed +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + docstrings: + name: Check docstrings on changed files + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 # need full history to diff against base branch + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install ruff + run: pip install ruff + + - name: Collect changed Python files in scope + id: changed + run: | + git diff --name-only origin/${{ github.base_ref }}...HEAD \ + | grep -E '^(nvsubquadratic|experiments)/.*\.py$' \ + > changed_files.txt || true + echo "Files to check:" + cat changed_files.txt + + - name: Run docstring checks on changed files + run: | + if [ ! -s changed_files.txt ]; then + echo "No in-scope Python files changed — skipping." + exit 0 + fi + # D100 missing module docstring + # D101 missing class docstring + # D102 missing method docstring + # D103 missing function docstring + # D301 use r""" for docstrings with backslashes + # D417 missing argument descriptions in docstring + xargs ruff check --select D100,D101,D102,D103,D301,D417 < changed_files.txt diff --git a/README.md b/README.md index a8ec6006..11eb0d8b 100644 --- a/README.md +++ b/README.md @@ -155,7 +155,17 @@ source .env && PYTHONPATH=. python -m pytest tests/ -m nightly -v -o addopts="" ### Documentation -The API reference is built with Sphinx. Sources live under [`docs/`](docs/) and the rendered site is published to the `gh-pages` branch on every push to `main` via [`.github/workflows/docs.yml`](.github/workflows/docs.yml). +All public classes and functions carry **Google-style docstrings** with math +context, shape annotations, and paper references. See [`CONVENTIONS.md`](CONVENTIONS.md) +for the style guide and PR checklist. + +#### Viewing the docs + +**Sphinx HTML (full API reference)** + +The API reference is built with Sphinx. Sources live under [`docs/`](docs/) and +the rendered site is published to the `gh-pages` branch on every push to `main` +via [`.github/workflows/docs.yml`](.github/workflows/docs.yml). Build and preview locally: @@ -166,7 +176,31 @@ make -C docs html SPHINXBUILD="python -m sphinx" python -m http.server 8000 --directory docs/_build/html ``` -Open to browse. The autosummary stubs in `docs/generated/` are regenerated on every build (gitignored). +Open to browse. The autosummary stubs in +`docs/generated/` are regenerated on every build (gitignored). + +**Inline — IDE hover / `help()`** + +Because all documentation lives directly in the docstrings, you can also: + +- Hover over any symbol in VS Code / PyCharm to see the rendered docstring. +- Run `help(SomeClass)` in a Python REPL for an immediate plain-text view. +- Use `python -m pydoc nvsubquadratic.modules.hyena_nd` for a terminal-friendly + per-module dump. + +**`pdoc` (quick zero-config HTML)** + +For a fast, dependency-light alternative to Sphinx that renders the docstrings +as-is: + +```bash +pip install pdoc +pdoc nvsubquadratic --output-dir /tmp/pdoc-out +python -m http.server 8000 --directory /tmp/pdoc-out +``` + +This does not require a `docs/conf.py` and picks up the Google-style sections +automatically. ### CI From 6dcee72e32005a81b3d747eb9b9d1f30b390369b Mon Sep 17 00:00:00 2001 From: David Wessels Date: Tue, 26 May 2026 15:40:28 +0200 Subject: [PATCH 63/72] docs: remove stale review artifacts causing Sphinx warnings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The docs/reviews/ files were intermediate artifacts from the write→review→integrate docstring pipeline. Their content has been fully integrated into the source docstrings. They were being picked up by Sphinx (source_suffix includes .md) but not referenced in any toctree, generating 'document not in toctree' warnings — which fail the CI build with -W --keep-going. Co-Authored-By: Claude Sonnet 4.6 --- docs/reviews/mamba_nd_review.md | 209 ---------------------- docs/reviews/vit5_hyena_adapter_review.md | 67 ------- 2 files changed, 276 deletions(-) delete mode 100644 docs/reviews/mamba_nd_review.md delete mode 100644 docs/reviews/vit5_hyena_adapter_review.md diff --git a/docs/reviews/mamba_nd_review.md b/docs/reviews/mamba_nd_review.md deleted file mode 100644 index d20c95b8..00000000 --- a/docs/reviews/mamba_nd_review.md +++ /dev/null @@ -1,209 +0,0 @@ -# Review: `nvsubquadratic/modules/mamba_nd.py` - -Reviewed after Phase 1 docstring pass. Issues are ordered from highest to -lowest priority. Each item quotes the relevant text and states the exact fix. - -______________________________________________________________________ - -## 1. ZOH approximation presented as equality - -**Location**: module-level docstring, "Background" section. - -**Quoted text**: - -``` -\bar{B}_t = (e^{\Delta_t A} - I) A^{-1} B_t \approx \Delta_t B_t -``` - -**Issue**: The full ZOH expression is exact, but the approximation -`≈ Δ_t B_t` is only valid for small `Δ_t`. More precisely, the Mamba paper -(Eq. 4 in arXiv:2312.00752) uses the Euler discretisation `B̄_t = Δ_t B_t` -and notes that the ZOH formula is an alternative; the code in `mamba_ssm` -actually uses the Euler rule for `B`, not the full ZOH. Presenting the full -ZOH formula and then approximating it may confuse readers into thinking the -implementation uses ZOH for both `A` and `B`. - -**Fix**: Clarify that `Ā_t = exp(Δ_t A)` (ZOH) but `B̄_t = Δ_t B_t` (Euler -/ first-order), matching what `mamba_ssm` actually computes. Remove the -intermediate exact-then-approximate chain or clearly label it "alternative -ZOH formula (not used in practice)". - -______________________________________________________________________ - -## 2. Comparison table: Mamba complexity claim is misleading for training - -**Location**: module-level docstring, "Comparison with other mixers" section. - -**Quoted text**: - -``` -| Mamba | SSM recurrence | input-dependent | O(N) | -``` - -**Issue**: O(N) is the *inference* (autoregressive) complexity via the -recurrent form. During training Mamba uses a parallel associative scan whose -GPU-efficient implementation is O(N log N) in time, or O(N) with the -hardware-aware parallel scan (which requires special CUDA kernels — the whole -point of the `mamba_ssm` package). Listing O(N) without qualification makes -it look strictly cheaper than Hyena at training time, which is only true with -the custom CUDA kernels. - -**Fix**: Add a note distinguishing training (parallel scan, O(N) with custom -kernels or O(N log N) naively) from inference (recurrent, O(N) per step, O(1) -state size). E.g. add a "Notes" column or a footnote: "O(N) with hardware- -aware parallel scan (requires `mamba_ssm` CUDA extension); O(1) per step at -inference". - -______________________________________________________________________ - -## 3. Scan order: no mention of the implication for 2D spatial locality - -**Location**: module-level docstring, "ND generalisation strategy" section, and -class-level docstring paragraph starting "**Scan order for ND inputs**". - -**Quoted text**: - -``` -The scan order for multi-dimensional inputs follows the default PyTorch / -``einops`` row-major (C-contiguous) flattening: for a 2D ``[H, W]`` input the -tokens are visited in raster-scan order (row 0, col 0 → row 0, col W-1 → -row 1, col 0 → …). -``` - -**Issue**: This is accurate but omits the key practical consequence: tokens -that are spatially adjacent *vertically* (same column, adjacent rows) are far -apart in the flattened sequence (W steps away), so the SSM's effective -receptive field is anisotropic — it sees horizontal neighbours cheaply but -vertical neighbours only through W state-update steps. An external -collaborator reading this to decide whether to use Mamba for a 2D task needs -this information. - -**Fix**: Add one sentence explicitly warning about vertical anisotropy, e.g.: -"Note that vertically adjacent pixels (same column, adjacent rows) are W -tokens apart in the flattened sequence; for tall images this means the forward -SSM sees them only through many state transitions, potentially losing spatial -correlation. Bidirectional mode partially mitigates this." - -______________________________________________________________________ - -## 4. `forward` docstring: `x` is mutated (overwritten by `rearrange`) - -**Location**: `Mamba.forward`, Args section and implementation. - -**Quoted text** (implementation): - -```python -x = rearrange(x, "b ... c -> b (...) c") -``` - -**Issue**: The `rearrange` result is assigned back to `x`, which shadows the -original argument. The original spatial shape is captured in `x_shape` first, -so this is safe, but the docstring does not warn that the local variable `x` -changes meaning mid-function (it is the flattened view from line 254 onward). -This is a mild readability issue but analogous to the note in `hyena_nd.py`'s -`forward` docstring about `query` being overwritten — be consistent. - -**Fix**: Add an "Implementation note" paragraph to `forward`: -"The local variable `x` is rebound to the flattened `[B, S, C]` view after -the `rearrange` call; the original spatial shape is preserved in `x_shape` -for the final `reshape`." - -______________________________________________________________________ - -## 5. `__init__` docstring: no mention that `core_layer_rev` has independent parameters - -**Location**: `Mamba.__init__`, Args section for `mamba_layer_cfg`. - -**Quoted text**: - -``` -The config is instantiated once for ``core_layer`` and, when -``bidirectional=True``, a second independent instantiation is created for -``core_layer_rev`` so that the two directions have separate parameters. -``` - -**Issue**: The wording says "independent instantiation … separate parameters" -but does not say *how* independence is achieved — a reader unfamiliar with -`LazyConfig` might wonder if the second call shares weights via some internal -cache. - -**Fix**: Add a clarifying phrase: "`instantiate(mamba_layer_cfg)` is called -twice with the same config; each call constructs a fresh `nn.Module` with -newly initialised weights, so the two directions do not share parameters." - -______________________________________________________________________ - -## 6. Missing `Raises` section in `__init__` - -**Location**: `Mamba.__init__` docstring. - -**Issue**: If `mamba_layer_cfg` cannot be instantiated (wrong target class, -missing required arguments), `instantiate` will raise — most likely a -`RuntimeError`, `TypeError`, or an `omegaconf` exception, depending on the -`LazyConfig` backend. The `QKVSequenceMixer.__init__` docstring in -`sequence_mixer.py` includes a `Raises` section for this; `Mamba.__init__` -should match. - -**Fix**: Add: - -``` -Raises: - Exception: Propagated from - :func:`~nvsubquadratic.lazy_config.instantiate` if - ``mamba_layer_cfg`` cannot be constructed. Check that the target - class accepts ``[B, S, C]`` tensors and that all required constructor - arguments are provided in the config. -``` - -______________________________________________________________________ - -## 7. Class-level Attributes block: `core_layer_rev` uses vague conditional phrasing - -**Location**: `Mamba` class, Attributes block. - -**Quoted text**: - -``` -core_layer_rev (torch.nn.Module): The reverse Mamba core. Only - present when ``bidirectional=True``; accessing this attribute - when ``bidirectional=False`` raises :class:`AttributeError`. -``` - -**Issue**: Saying "raises `AttributeError`" is accurate but alarming without -context — it looks like a bug. It would be clearer to note the attribute is -intentionally absent (not registered) to keep `state_dict` / `parameters()` -clean when unused. - -**Fix**: Reword to: "The reverse Mamba core, instantiated only when -`bidirectional=True`. When `bidirectional=False` this attribute is not -registered and accessing it raises :class:`AttributeError` by design, keeping -the module's parameter count and `state_dict` unaffected." - -______________________________________________________________________ - -## 8. Module docstring: `References` section uses non-standard heading style - -**Location**: module-level docstring, last section. - -**Quoted text** (after ruff fix): - -``` -References: ----------- -``` - -**Issue**: Sphinx / NumPy / Google doc conventions all write `References` with -the underline immediately under the heading, not with a blank line. The ruff -linter already fixed a trailing colon issue here but the section may still -render oddly in Sphinx because of the mixed `References:` + underline style. -The rest of the module docstring uses RST underlines without trailing colons -on the heading (e.g. `Background\n----------`). - -**Fix**: Change to match the rest of the file: - -``` -References ----------- -``` - -(remove the trailing colon from `References:`). diff --git a/docs/reviews/vit5_hyena_adapter_review.md b/docs/reviews/vit5_hyena_adapter_review.md deleted file mode 100644 index b3e29aa4..00000000 --- a/docs/reviews/vit5_hyena_adapter_review.md +++ /dev/null @@ -1,67 +0,0 @@ -# Review: `nvsubquadratic/modules/vit5_hyena_adapter.py` - -Reviewed after Phase 1 docstring pass. Issues are numbered and actionable. - -______________________________________________________________________ - -## 1. Module docstring: "drop-in replacement" claim needs qualification - -The second line reads: - -> "Drop-in replacement interface for `ViT5Attention`." - -This is misleading. `ViT5Attention` owns its own QKV and output projections; the adapter delegates all projections to `inner_mixer`. A reader replacing `ViT5Attention` with `ViT5HyenaAdapter` in config needs to know that `inner_mixer` (e.g. `QKVSequenceMixer`) must be configured with matching `hidden_dim` and the correct `num_heads`/`inner_dim` — the adapter itself accepts no `hidden_dim` argument. The docstring should add a sentence clarifying that the projection parameters must be configured **inside `inner_mixer_cfg`**, not on the adapter. - -______________________________________________________________________ - -## 2. Module docstring: `inner_mixer` contract is underspecified for the 2-D output reshape - -Step 4 of the "thin, stateless reshape adapter" list says: - -> "4. Reshapes back to `[B, T, C]` and returns." - -This silently assumes the inner mixer returns **exactly** `[B, H, W, C]` — the same spatial shape it received. If a mixer returns a different shape (e.g. a downsampling mixer), `reshape(B, T, C)` will raise a cryptic error. The docstring should state explicitly: "The inner mixer is assumed to be shape-preserving — its output must have the same `[B, H, W, C]` shape as its input." - -______________________________________________________________________ - -## 3. Class docstring: data-flow diagram uses `inner_mixer (QKVSequenceMixer → Hyena)` as a fixed example - -The class-level data-flow diagram labels the inner_mixer step as: - -``` -▼ inner_mixer (QKVSequenceMixer → Hyena) -``` - -This couples the class docstring to a specific implementation. The class is general — it wraps any `[B, H, W, C]`-preserving module. The parenthetical should be replaced with a more general description, e.g. `(any [B, H, W, C]-preserving mixer)`. - -______________________________________________________________________ - -## 4. `__init__` docstring: `inner_mixer_cfg` contract does not mention channels-last convention - -The docstring says the instantiated module must "accept `(x: Tensor[B, H, W, C], **kwargs)`" but does not mention that this is a **channels-last** layout. Hyena internally converts to channels-first `[B, C, H, W]` for convolution. A reader wiring a custom mixer needs to know the expected layout convention. Add: "The tensor is in channels-last layout `[B, H, W, C]`; any inner mixer that uses channels-first convolution (like `QKVSequenceMixer`) handles the permutation internally." - -______________________________________________________________________ - -## 5. `forward` docstring: second `reshape` output contiguity caveat is incomplete - -The Returns section states: "The reshape is a view (no data copy) when the tensor is contiguous." However, the inner mixer may return a non-contiguous tensor (e.g. after a transpose). In that case `reshape` falls back to a copy. This is fine correctness-wise but users writing CUDA-graph-safe code or tracing with `torch.compile` should know. Add: "If `inner_mixer` returns a non-contiguous tensor, the final `reshape` may trigger a contiguous copy; this does not affect correctness but can affect memory traffic." - -______________________________________________________________________ - -## 6. `flop_count` docstring: silent `AttributeError` is wrong error type in practice - -The Raises section says `AttributeError` if `inner_mixer` does not implement `flop_count`. In practice, if `inner_mixer` has no `flop_count`, Python raises `AttributeError` from the attribute lookup, not from a guarded check. This is correct but the docstring should also note that `flop_count` is a **de-facto protocol** not enforced by an interface — callers that want to guard against this should use `hasattr(adapter.inner_mixer, "flop_count")`. - -______________________________________________________________________ - -## 7. Missing: note about `grid_w` and the token-layout contract for the hierarchical case - -The module docstring mentions `ViT5ClassificationNet` as the network that pads and arranges tokens, but the current codebase also has `vit5_hierarchical_classification.py` (feat/patch-merging PR). The register-token handling note in the class docstring only mentions a single "register row" convention but does not mention that after patch merging the grid dimensions change. Add a brief note that `grid_w` must be consistent with the spatial width **after any patch-merging stage**, i.e. the calling network must supply the correct `grid_w` at each hierarchical stage. - -______________________________________________________________________ - -## 8. `extra_repr` docstring: wording is slightly off - -Current: "Return a concise summary for `repr()` and `print(model)`." - -`extra_repr` is called by PyTorch's `__repr__` machinery. The docstring should say what information it returns rather than how it is called: "Return `grid_w=` appended to PyTorch's default module repr." This matches the style in `vit5_attention.py`'s `extra_repr`. From 2c9ddb935b08a29aede2eb46dd09f602a587e77b Mon Sep 17 00:00:00 2001 From: David Wessels Date: Tue, 26 May 2026 15:56:32 +0200 Subject: [PATCH 64/72] docs: fix Sphinx cross-ref class names and ViT-5 token layout doc MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit residual_block.py: GeneralPurposeResnet → ResidualNetwork and ClassificationResnet → ClassificationResNet in both See Also blocks. The old names were non-existent and would have rendered as broken links in the Sphinx API reference. vit5_residual_block.py: expand token layout description to include optional zero-padding tokens that ViT5ClassificationNet appends for Hyena blocks when _block_needs_padding is true. Both the module docstring and the forward() Args block now document the full layout: [patches, (CLS,) registers, (padding,)] with T % grid_w == 0 for padded blocks. Co-Authored-By: Claude Sonnet 4.6 --- nvsubquadratic/modules/residual_block.py | 8 ++++---- nvsubquadratic/modules/vit5_residual_block.py | 15 ++++++++++----- 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/nvsubquadratic/modules/residual_block.py b/nvsubquadratic/modules/residual_block.py index ccad415e..9ffe7908 100644 --- a/nvsubquadratic/modules/residual_block.py +++ b/nvsubquadratic/modules/residual_block.py @@ -83,9 +83,9 @@ class ResidualBlock(torch.nn.Module): optimiser can exclude them from weight-decay groups. See Also: - :class:`~nvsubquadratic.networks.general_purpose_resnet.GeneralPurposeResnet` + :class:`~nvsubquadratic.networks.general_purpose_resnet.ResidualNetwork` and - :class:`~nvsubquadratic.networks.classification_resnet.ClassificationResnet` + :class:`~nvsubquadratic.networks.classification_resnet.ClassificationResNet` for the canonical consumers of this block. Attributes: @@ -335,9 +335,9 @@ class AdaLNZeroResidualBlock(torch.nn.Module): zero-initialised), so the block is a skip connection. See Also: - :class:`~nvsubquadratic.networks.general_purpose_resnet.GeneralPurposeResnet` + :class:`~nvsubquadratic.networks.general_purpose_resnet.ResidualNetwork` and - :class:`~nvsubquadratic.networks.classification_resnet.ClassificationResnet` + :class:`~nvsubquadratic.networks.classification_resnet.ClassificationResNet` for the canonical consumers of this block. Attributes: diff --git a/nvsubquadratic/modules/vit5_residual_block.py b/nvsubquadratic/modules/vit5_residual_block.py index 070bcccf..5d333da7 100644 --- a/nvsubquadratic/modules/vit5_residual_block.py +++ b/nvsubquadratic/modules/vit5_residual_block.py @@ -31,8 +31,11 @@ 5. **Sequence layout** — Input is always ``[B, T, C]`` (batch, tokens, channels); the token axis *T* concatenates patch tokens, an optional CLS - token, and register tokens in the order ``[patches, (CLS,) registers]``. - The generic block operates on arbitrary ``(B, *spatial_dims, C)`` tensors. + token, register tokens, and optional zero-padding in the order + ``[patches, (CLS,) registers, (padding,)]``. Attention blocks receive the + sequence with padding stripped; Hyena / subquadratic blocks receive the full + padded sequence so that ``T % grid_w == 0``. The generic block operates on + arbitrary ``(B, *spatial_dims, C)`` tensors. For the generic pre-norm residual block (Hyena / Attention / CKConv / Mamba with optional cross-attention conditioning), see @@ -298,9 +301,11 @@ def forward(self, x: torch.Tensor, condition: torch.Tensor = None) -> torch.Tens Args: x: Input token sequence of shape ``[B, T, C]``, where: - ``B`` — batch size, - - ``T = num_patches + (1 if has_cls else 0) + num_registers`` - — total token count following the ViT-5 layout - ``[patches, (CLS,) registers]``, + - ``T = num_patches + (1 if has_cls else 0) + num_registers + (+ pad_size for Hyena blocks)`` — total token count following + the ViT-5 layout ``[patches, (CLS,) registers, (padding,)]``. + Attention blocks receive the unpadded sequence; Hyena blocks + receive the zero-padded sequence so ``T % grid_w == 0``, - ``C`` — channel (hidden) dimension. condition: Accepted for API compatibility with :class:`~nvsubquadratic.modules.residual_block.ResidualBlock` From e5f316ea510acfc70061637ac501f52827b3b11f Mon Sep 17 00:00:00 2001 From: David Wessels Date: Tue, 26 May 2026 16:18:26 +0200 Subject: [PATCH 65/72] =?UTF-8?q?docs(causal=5Fconv1d):=20fix=20incorrect?= =?UTF-8?q?=20Mamba=20usage=20note=20=E2=80=94=20actual=20call=20sites=20a?= =?UTF-8?q?re=20Hyena=20short-conv=20configs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The module docstring claimed CausalConv1D was used in mamba_nd.py, but it is not imported there at all. The actual call sites are the Hyena short-conv configuration helpers: examples/spatial_recall_v2/mixer_defaults.py examples/spatial_recall_1d/mixer_defaults.py Update the 'Use in …' section header and body to reflect reality. Co-Authored-By: Claude Sonnet 4.6 --- nvsubquadratic/modules/causal_conv1d.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/nvsubquadratic/modules/causal_conv1d.py b/nvsubquadratic/modules/causal_conv1d.py index 53881a35..793b3ce6 100644 --- a/nvsubquadratic/modules/causal_conv1d.py +++ b/nvsubquadratic/modules/causal_conv1d.py @@ -18,13 +18,15 @@ ``is_causal=False``, making it a drop-in for standard convolutions in contexts that need a runtime-switchable causality flag. -**Use inside Mamba / SSM blocks** - -:class:`CausalConv1D` is used in :mod:`nvsubquadratic.modules.mamba_nd` as a -depthwise filter applied along each spatial axis before the selective scan. -The causal constraint along the time axis is essential for autoregressive -generation; for spatial axes (height, width) the symmetric variant is used -so the scan sees the full context. +**Use in Hyena short-conv configs** + +In this repository, :class:`CausalConv1D` is wired into Hyena operators as the +short-conv component. Concrete call sites include +``examples/spatial_recall_v2/mixer_defaults.py`` and +``examples/spatial_recall_1d/mixer_defaults.py``, which select +:class:`CausalConv1D` (``is_causal=True``) for the time axis and the symmetric +variant (``is_causal=False``) for spatial axes where the full context should be +visible. **Difference from** ``subq_ops_causal_conv1d`` **/ ``causal_conv1d_custom``** From 81a2291903c86fe3e4796ca83d94f198961e821c Mon Sep 17 00:00:00 2001 From: David Wessels Date: Tue, 26 May 2026 16:24:55 +0200 Subject: [PATCH 66/72] =?UTF-8?q?docs(drop=5Fpath):=20clarify=20drop=5Fpro?= =?UTF-8?q?b=3D1.0=20is=20safe=20=E2=80=94=20implementation=20guards=20div?= =?UTF-8?q?-by-zero?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Copilot flagged that the docstring implied drop_prob=1.0 would zero every sample cleanly, but worried the keep_prob=0 division would produce inf/NaN first. The implementation already has an explicit 'if keep_prob > 0.0' guard (line 67) that skips the rescaling division, so Bernoulli(0) produces an all-zero mask and x*0=0 with no numerical issue. Update the docstring to document this guarantee. Co-Authored-By: Claude Sonnet 4.6 --- nvsubquadratic/modules/drop_path.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/nvsubquadratic/modules/drop_path.py b/nvsubquadratic/modules/drop_path.py index 29cd461a..73f27d21 100644 --- a/nvsubquadratic/modules/drop_path.py +++ b/nvsubquadratic/modules/drop_path.py @@ -50,7 +50,9 @@ def drop_path(x: torch.Tensor, drop_prob: float, training: bool) -> torch.Tensor x: Input tensor of shape ``[B, *]`` — any layout; the drop mask has shape ``(B, 1, …, 1)`` and broadcasts over all non-batch dimensions. drop_prob: Probability of dropping a sample's contribution. - ``0.0`` disables dropping; ``1.0`` would zero every sample. + ``0.0`` disables dropping; ``1.0`` zeros every sample (safe — + the implementation guards against dividing by ``keep_prob`` when + it is zero, so no inf/NaN is produced). training: Whether the model is currently in training mode. Set to ``False`` (or call ``model.eval()``) to disable dropping. From f378d1528bfcba91ebcd7daf1395a50ae6e74b1d Mon Sep 17 00:00:00 2001 From: David Wessels Date: Tue, 26 May 2026 16:27:14 +0200 Subject: [PATCH 67/72] =?UTF-8?q?fix(ruff):=20enforce=20D100=20(missing=20?= =?UTF-8?q?module=20docstring)=20=E2=80=94=20was=20silently=20ignored?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CONVENTIONS.md documented that the pre-commit ruff hook catches D100, but pyproject.toml had D100 in the global ignore list, so both the pre-commit hook and the CI diff-check silently skipped missing module docstrings (ruff applies config-file ignores on top of CLI --select). Fix: - Remove D100 from pyproject.toml ignore so the hook matches the docs - Add missing module docstring to nvsubquadratic/parallel/utils.py (the only file in scope that was missing one) Co-Authored-By: Claude Sonnet 4.6 --- nvsubquadratic/parallel/utils.py | 19 +++++++++++++++++++ pyproject.toml | 1 - 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/nvsubquadratic/parallel/utils.py b/nvsubquadratic/parallel/utils.py index cf819b14..108943d6 100644 --- a/nvsubquadratic/parallel/utils.py +++ b/nvsubquadratic/parallel/utils.py @@ -1,5 +1,24 @@ # TODO: Add license header here +"""Distributed-training utility helpers for context-parallel (CP) workloads. + +Provides three categories of utilities: + +1. **Process-group initialisation** — :func:`init_parallel_state` wires up NCCL + and Megatron's parallel state (tensor / pipeline / context parallelism). + +2. **Zigzag sequence splitting** — :func:`zigzag_split_across_group_ranks` and + :func:`zigzag_gather_from_group_ranks` implement the zigzag collective that + distributes a sequence of length ``L`` evenly across ``CP`` ranks while + keeping each rank's two chunks at opposite ends of the sequence. This + balances causal-attention load (each rank sees one chunk of "early" tokens + and one chunk of "late" tokens). + +3. **Rank-0 logging** — :func:`setup_rank0_logging` routes console output only + from rank 0 while writing all ranks' logs to per-rank files, preventing + duplicate console noise in multi-GPU runs. +""" + import logging import os from datetime import timedelta diff --git a/pyproject.toml b/pyproject.toml index 178c97e3..4e176e56 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,7 +44,6 @@ line-length = 119 [tool.ruff.lint] ignore = [ "C901", - "D100", "E501", "E741", "RUF001", From d182d32a545542deea7aa5c255455a46fb9685e6 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Tue, 26 May 2026 16:35:48 +0200 Subject: [PATCH 68/72] docs: add missing module docstrings to satisfy D100 now that it is enforced MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four files were missing module-level docstrings that ruff D100 now catches (after removing D100 from the global ignore list): - docs/conf.py — Sphinx configuration file - examples/imagenet_diffusion/ccnn_jit_baseline.py — CCNN-Hyena JiT-B-matched diffusion baseline - examples/imagenet_diffusion/hf_uvit_baseline.py — HuggingFace UViT diffusion baseline - examples/imagenet_diffusion/jit_baseline.py — JiT-B flow-matching diffusion baseline Co-Authored-By: Claude Sonnet 4.6 --- docs/conf.py | 2 ++ examples/imagenet_diffusion/ccnn_jit_baseline.py | 2 ++ examples/imagenet_diffusion/hf_uvit_baseline.py | 2 ++ examples/imagenet_diffusion/jit_baseline.py | 2 ++ 4 files changed, 8 insertions(+) diff --git a/docs/conf.py b/docs/conf.py index 2a703f55..51bc30a2 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -4,6 +4,8 @@ # Configuration file for the Sphinx documentation builder. # See https://www.sphinx-doc.org/en/master/usage/configuration.html +"""Sphinx configuration for the nvsubquadratic API reference.""" + import os import re as _re import sys diff --git a/examples/imagenet_diffusion/ccnn_jit_baseline.py b/examples/imagenet_diffusion/ccnn_jit_baseline.py index 5a20201c..58c7b49f 100644 --- a/examples/imagenet_diffusion/ccnn_jit_baseline.py +++ b/examples/imagenet_diffusion/ccnn_jit_baseline.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""CCNN-Hyena baseline config for class-conditional ImageNet 64×64 flow-matching diffusion, size-matched to JiT-B.""" + import os import torch diff --git a/examples/imagenet_diffusion/hf_uvit_baseline.py b/examples/imagenet_diffusion/hf_uvit_baseline.py index 1442e808..0815e030 100644 --- a/examples/imagenet_diffusion/hf_uvit_baseline.py +++ b/examples/imagenet_diffusion/hf_uvit_baseline.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""HuggingFace UViT baseline config for class-conditional ImageNet 64×64 diffusion.""" + import os import torch diff --git a/examples/imagenet_diffusion/jit_baseline.py b/examples/imagenet_diffusion/jit_baseline.py index 15757c9b..c47906d0 100644 --- a/examples/imagenet_diffusion/jit_baseline.py +++ b/examples/imagenet_diffusion/jit_baseline.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""JiT-B baseline config for class-conditional ImageNet 64×64 flow-matching diffusion.""" + import os import torch From 02cf239e01cad02932acf94e8020a768e287bc40 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 15:03:03 +0200 Subject: [PATCH 69/72] feat(modules): add Swin-style PatchMerging with register-row support Implements 2x2 spatial patch merging for hierarchical ViT-5/Hyena networks. Key features: - Pure-spatial layout: [B, H*W, C] -> [B, (H/2)*(W/2), out_dim] - Register-row layout: passes the leading grid_w register tokens through a dedicated reg_proj Linear so FiLM conditioning survives the channel-dim change, then repacks them to width grid_w//2 - Post-concat norm (configurable via LazyConfig) + bias-free reduction Linear, both trunc_normal_ initialised (std=0.02) - flop_count() helper for FLOP bookkeeping - Full test suite covering shape, zero-pad, independent paths, error guards, and FLOP count Co-Authored-By: Claude Sonnet 4.6 --- nvsubquadratic/modules/patch_merging.py | 154 ++++++++++++++++++++++++ tests/modules/test_patch_merging.py | 147 ++++++++++++++++++++++ 2 files changed, 301 insertions(+) create mode 100644 nvsubquadratic/modules/patch_merging.py create mode 100644 tests/modules/test_patch_merging.py diff --git a/nvsubquadratic/modules/patch_merging.py b/nvsubquadratic/modules/patch_merging.py new file mode 100644 index 00000000..d7175038 --- /dev/null +++ b/nvsubquadratic/modules/patch_merging.py @@ -0,0 +1,154 @@ +"""Swin-style 2x2 patch merging for hierarchical ViT-5 / Hyena networks. + +Halves both spatial dims and (typically) doubles the channel dim. Two input +layouts are supported: + +* ``has_register_row=False``: pure-spatial sequence ``[B, H*W, C]`` reshaped to + a 2D grid, 2x2-merged, normalized, and projected. + +* ``has_register_row=True``: ``[B, grid_w + H*W, C]`` where the first ``grid_w`` + tokens form a "register row" (the layout used by ``ViT5ClassificationNet`` + with ``prepend_registers=True`` and no CLS). The patch grid is merged as + above; register tokens are projected independently with their own linear so + the FiLM conditioning signal survives the channel-dim change, then re-padded + to the new (halved) grid width. +""" + +import torch +import torch.nn as nn +from einops import rearrange + +from nvsubquadratic.lazy_config import LazyConfig, instantiate + + +class PatchMerging(nn.Module): + """2x2 patch merging with optional register-row passthrough. + + Args: + in_dim: Input channel dimension. + out_dim: Output channel dimension (Swin-T uses ``out_dim = 2 * in_dim``). + grid_h: Patch-grid height before merging. Must be even. + grid_w: Patch-grid width before merging. Must be even. + norm_cfg: LazyConfig for the post-concat norm. Must be configured with + ``dim = 4 * in_dim`` since it operates on concatenated 2x2 features. + num_registers: Number of register tokens at the start of the register + row. Only used when ``has_register_row=True``. + has_register_row: When True, the first ``grid_w`` tokens of the input + sequence are treated as a register row (regs + zero pad) and + passed through a dedicated ``Linear(in_dim, out_dim)`` projection. + The output register row is repacked to width ``grid_w // 2``. + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + grid_h: int, + grid_w: int, + norm_cfg: LazyConfig, + num_registers: int = 0, + has_register_row: bool = False, + ): + """Initialise PatchMerging weights and validate grid dimensions.""" + super().__init__() + if grid_h % 2 != 0 or grid_w % 2 != 0: + raise ValueError(f"grid_h={grid_h}, grid_w={grid_w} must both be even for 2x2 merging") + + self.in_dim = in_dim + self.out_dim = out_dim + self.grid_h = grid_h + self.grid_w = grid_w + self.out_grid_h = grid_h // 2 + self.out_grid_w = grid_w // 2 + self.has_register_row = has_register_row + self.num_registers = num_registers + + self.norm = instantiate(norm_cfg) + for p in self.norm.parameters(): + p._no_weight_decay = True + self.reduction = nn.Linear(4 * in_dim, out_dim, bias=False) + nn.init.trunc_normal_(self.reduction.weight, std=0.02) + + if has_register_row: + if num_registers > self.out_grid_w: + raise ValueError(f"num_registers ({num_registers}) must fit in halved grid_w ({self.out_grid_w})") + self.reg_proj = nn.Linear(in_dim, out_dim, bias=False) + nn.init.trunc_normal_(self.reg_proj.weight, std=0.02) + pad_size = self.out_grid_w - num_registers + if pad_size > 0: + self.register_buffer("reg_zero_pad", torch.zeros(1, pad_size, out_dim), persistent=False) + else: + self.reg_zero_pad = None + else: + self.reg_proj = None + self.reg_zero_pad = None + + def flop_count(self) -> int: + """FLOPs for one merging step (one sample). + + Breakdown: + * norm on (out_grid_h * out_grid_w) tokens at 4*in_dim channels. + * reduction linear: 2 * (out_grid_h * out_grid_w) * (4*in_dim) * out_dim. + * (register row only) reg_proj: 2 * num_registers * in_dim * out_dim. + """ + T_out = self.out_grid_h * self.out_grid_w + # Norm is configured for dim=4*in_dim and is called on T_out tokens. + flops = self.norm.flop_count(T_out) + flops += 2 * T_out * 4 * self.in_dim * self.out_dim + if self.has_register_row: + flops += 2 * self.num_registers * self.in_dim * self.out_dim + return flops + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass. + + Args: + x: ``[B, T_in, in_dim]`` token sequence. ``T_in = grid_h*grid_w`` + in the pure-spatial case; ``grid_w + grid_h*grid_w`` when a + register row is present. + + Returns: + ``[B, T_out, out_dim]`` where ``T_out`` is ``out_grid_h*out_grid_w`` + (pure-spatial) or ``out_grid_w + out_grid_h*out_grid_w`` (with + register row). + """ + B = x.shape[0] + + if self.has_register_row: + regs_row = x[:, : self.grid_w, :] # [B, grid_w, C] full row incl. pad + regs = regs_row[:, : self.num_registers, :] # [B, num_regs, C] + patches_flat = x[:, self.grid_w :, :] # [B, H*W, C] + else: + patches_flat = x + + patches = rearrange(patches_flat, "b (h w) c -> b h w c", h=self.grid_h, w=self.grid_w) + + # Standard Swin 2x2 spatial-merge: gather four spatial cosets and concat on channels. + x0 = patches[:, 0::2, 0::2, :] + x1 = patches[:, 1::2, 0::2, :] + x2 = patches[:, 0::2, 1::2, :] + x3 = patches[:, 1::2, 1::2, :] + merged = torch.cat([x0, x1, x2, x3], dim=-1) # [B, H/2, W/2, 4C] + + merged = self.norm(merged) + merged = self.reduction(merged) + merged_flat = rearrange(merged, "b h w c -> b (h w) c") + + if not self.has_register_row: + return merged_flat + + regs_proj = self.reg_proj(regs) # [B, num_regs, out_dim] + if self.reg_zero_pad is not None: + pad = self.reg_zero_pad.expand(B, -1, -1) + out = torch.cat([regs_proj, pad, merged_flat], dim=1) + else: + out = torch.cat([regs_proj, merged_flat], dim=1) + return out + + def extra_repr(self) -> str: + """Return a compact string summary of the module configuration.""" + return ( + f"in_dim={self.in_dim}, out_dim={self.out_dim}, " + f"grid={self.grid_h}x{self.grid_w}->{self.out_grid_h}x{self.out_grid_w}, " + f"register_row={self.has_register_row}, num_registers={self.num_registers}" + ) diff --git a/tests/modules/test_patch_merging.py b/tests/modules/test_patch_merging.py new file mode 100644 index 00000000..333f75fc --- /dev/null +++ b/tests/modules/test_patch_merging.py @@ -0,0 +1,147 @@ +"""Tests for the PatchMerging module.""" + +import pytest +import torch + +from nvsubquadratic.lazy_config import LazyConfig +from nvsubquadratic.modules.patch_merging import PatchMerging +from nvsubquadratic.modules.rms_norm import RMSNorm + + +@pytest.mark.parametrize("grid_h,grid_w", [(8, 8), (16, 16), (28, 28)]) +def test_pure_spatial_output_shape(device, grid_h: int, grid_w: int) -> None: + """Pure-spatial layout: [B, H*W, C] -> [B, (H/2)*(W/2), out_dim].""" + B, in_dim, out_dim = 2, 96, 192 + pm = PatchMerging( + in_dim=in_dim, + out_dim=out_dim, + grid_h=grid_h, + grid_w=grid_w, + norm_cfg=LazyConfig(RMSNorm)(dim=4 * in_dim, eps=1e-6, use_quack=False), + has_register_row=False, + ).to(device) + + x = torch.randn(B, grid_h * grid_w, in_dim, device=device) + y = pm(x) + + expected_T = (grid_h // 2) * (grid_w // 2) + assert y.shape == (B, expected_T, out_dim) + + +@pytest.mark.parametrize("grid_h,grid_w,num_regs", [(8, 8, 4), (28, 28, 4), (14, 14, 3)]) +def test_register_row_output_shape(device, grid_h: int, grid_w: int, num_regs: int) -> None: + """Register-row layout: [B, W + H*W, C] -> [B, W/2 + (H/2)*(W/2), out_dim].""" + B, in_dim, out_dim = 2, 96, 192 + pm = PatchMerging( + in_dim=in_dim, + out_dim=out_dim, + grid_h=grid_h, + grid_w=grid_w, + norm_cfg=LazyConfig(RMSNorm)(dim=4 * in_dim, eps=1e-6, use_quack=False), + num_registers=num_regs, + has_register_row=True, + ).to(device) + + T_in = grid_w + grid_h * grid_w + x = torch.randn(B, T_in, in_dim, device=device) + y = pm(x) + + expected_T = (grid_w // 2) + (grid_h // 2) * (grid_w // 2) + assert y.shape == (B, expected_T, out_dim) + + +def test_register_row_pad_is_zero(device) -> None: + """Output register-row padding slots must remain zero after the projection.""" + B, in_dim, out_dim, grid, num_regs = 2, 32, 64, 28, 4 + pm = PatchMerging( + in_dim=in_dim, + out_dim=out_dim, + grid_h=grid, + grid_w=grid, + norm_cfg=LazyConfig(RMSNorm)(dim=4 * in_dim, eps=1e-6, use_quack=False), + num_registers=num_regs, + has_register_row=True, + ).to(device) + + x = torch.randn(B, grid + grid * grid, in_dim, device=device) + # The input pad slice (positions [num_regs:grid] of the register row) does + # not have to be zero for this test — we only check that the *output* pad + # slice is, since it is sourced from a non-persistent zero buffer. + y = pm(x) + + out_grid_w = grid // 2 + pad_slice = y[:, num_regs:out_grid_w, :] + assert torch.equal(pad_slice, torch.zeros_like(pad_slice)) + + +def test_grid_must_be_even() -> None: + """Odd grid dims raise during construction.""" + with pytest.raises(ValueError, match="must both be even"): + PatchMerging( + in_dim=8, + out_dim=16, + grid_h=7, + grid_w=8, + norm_cfg=LazyConfig(RMSNorm)(dim=32, eps=1e-6, use_quack=False), + ) + + +def test_num_registers_must_fit_halved_grid() -> None: + """num_registers exceeding the halved grid_w raises.""" + with pytest.raises(ValueError, match="must fit in halved grid_w"): + PatchMerging( + in_dim=8, + out_dim=16, + grid_h=4, + grid_w=4, # halved -> 2 + norm_cfg=LazyConfig(RMSNorm)(dim=32, eps=1e-6, use_quack=False), + num_registers=3, + has_register_row=True, + ) + + +def test_register_tokens_route_through_reg_proj(device) -> None: + """Zeroing the patch-reduction weight should leave register tokens non-zero. + + Sanity check that the register and patch paths are wired independently. + """ + B, in_dim, out_dim, grid, num_regs = 2, 16, 32, 8, 4 + pm = PatchMerging( + in_dim=in_dim, + out_dim=out_dim, + grid_h=grid, + grid_w=grid, + norm_cfg=LazyConfig(RMSNorm)(dim=4 * in_dim, eps=1e-6, use_quack=False), + num_registers=num_regs, + has_register_row=True, + ).to(device) + + with torch.no_grad(): + pm.reduction.weight.zero_() + + x = torch.randn(B, grid + grid * grid, in_dim, device=device) + y = pm(x) + + out_grid_w = grid // 2 + reg_out = y[:, :num_regs, :] + patch_out = y[:, out_grid_w:, :] + + assert reg_out.abs().sum() > 0 # reg path is alive + assert torch.equal(patch_out, torch.zeros_like(patch_out)) # patch path is zeroed + + +def test_flop_count_pure(device) -> None: + """FLOP count returns a positive int dominated by the reduction linear.""" + in_dim, out_dim, grid = 96, 192, 28 + pm = PatchMerging( + in_dim=in_dim, + out_dim=out_dim, + grid_h=grid, + grid_w=grid, + norm_cfg=LazyConfig(RMSNorm)(dim=4 * in_dim, eps=1e-6, use_quack=False), + ).to(device) + f = pm.flop_count() + T_out = (grid // 2) ** 2 + reduction_flops = 2 * T_out * 4 * in_dim * out_dim + assert f > reduction_flops # norm contributes a (small) positive amount + assert isinstance(f, int) From 259e7f566c6f6ca4ec6d71c03f592234cd938c5b Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 15:08:56 +0200 Subject: [PATCH 70/72] feat(networks+examples): hierarchical ViT-5 classifier and v6 configs Swin-style 4-stage hierarchical ViT-5 classifier with PatchMerging between stages. Two readout layouts: - pure: flat patch grid [B, H*W, C] + GAP - register_row: FiLM register tokens prepended as first grid row at every stage; GAP excludes them Key properties: - Per-stage dims/depths fully configurable via LazyConfig - flop_count() aggregates patch-embed + blocks + merges + head - Backward tested: patch_embed and reg_proj grads both non-zero ImageNet configs for 4-stage Swin-T-like Hyena hierarchy (p=4, dims [96,192,384,768], depths [2,2,6,2]) in both 'pure' and 'register_row' FiLM variants, plus CIFAR-10 patch and capacity ablation configs. Mark patch_merging.py and vit5_hierarchical_classification.py as [x]. Co-Authored-By: Claude Sonnet 4.6 --- docs-tracker.md | 4 +- .../vit5_imagenet/v6_hierarchical/__init__.py | 0 .../v6_hierarchical/_base_config.py | 394 ++++++++++++++++++ .../_cifar10_patch_ablation_base.py | 336 +++++++++++++++ .../v6_hierarchical/_smoke_test.py | 92 ++++ .../v6_hierarchical/cifar10_flat_p16.py | 10 + .../v6_hierarchical/cifar10_flat_p4.py | 10 + .../v6_hierarchical/cifar10_flat_p8.py | 10 + .../v6_hierarchical/cifar10_hier_p16.py | 13 + .../v6_hierarchical/cifar10_hier_p4.py | 13 + .../v6_hierarchical/cifar10_hier_p8.py | 13 + .../v6_hierarchical/cifar10_hyena_flat.py | 226 ++++++++++ .../v6_hierarchical/cifar10_hyena_hier.py | 247 +++++++++++ .../v6_hierarchical/hyena_hier_p4_film.py | 22 + .../v6_hierarchical/hyena_hier_p4_pure.py | 21 + .../vit5_hierarchical_classification.py | 259 ++++++++++++ .../test_vit5_hierarchical_classification.py | 167 ++++++++ 17 files changed, 1835 insertions(+), 2 deletions(-) create mode 100644 examples/vit5_imagenet/v6_hierarchical/__init__.py create mode 100644 examples/vit5_imagenet/v6_hierarchical/_base_config.py create mode 100644 examples/vit5_imagenet/v6_hierarchical/_cifar10_patch_ablation_base.py create mode 100644 examples/vit5_imagenet/v6_hierarchical/_smoke_test.py create mode 100644 examples/vit5_imagenet/v6_hierarchical/cifar10_flat_p16.py create mode 100644 examples/vit5_imagenet/v6_hierarchical/cifar10_flat_p4.py create mode 100644 examples/vit5_imagenet/v6_hierarchical/cifar10_flat_p8.py create mode 100644 examples/vit5_imagenet/v6_hierarchical/cifar10_hier_p16.py create mode 100644 examples/vit5_imagenet/v6_hierarchical/cifar10_hier_p4.py create mode 100644 examples/vit5_imagenet/v6_hierarchical/cifar10_hier_p8.py create mode 100644 examples/vit5_imagenet/v6_hierarchical/cifar10_hyena_flat.py create mode 100644 examples/vit5_imagenet/v6_hierarchical/cifar10_hyena_hier.py create mode 100644 examples/vit5_imagenet/v6_hierarchical/hyena_hier_p4_film.py create mode 100644 examples/vit5_imagenet/v6_hierarchical/hyena_hier_p4_pure.py create mode 100644 nvsubquadratic/networks/vit5_hierarchical_classification.py create mode 100644 tests/networks/test_vit5_hierarchical_classification.py diff --git a/docs-tracker.md b/docs-tracker.md index 0d436b90..cb70123e 100644 --- a/docs-tracker.md +++ b/docs-tracker.md @@ -68,7 +68,7 @@ Work bottom-up: primitive ops → modules → networks → experiments. | `subq_ops_causal_conv1d.py` | \[x\] | New (1D PR): `nn.Conv1d`-compatible depthwise wrapper around `subq_ops.causal_conv1d` | | `schedulers.py` | \[x\] | ResumableSequentialLR — PyTorch ≤2.10 bug fix, load_state_dict LR propagation | | `distributed_depthwise_conv_nd.py` | \[x\] | CP-aware 1D/2D/3D depthwise conv — group weight sharing, channel slicing, causal padding | -| `patch_merging.py` | \[ \] | Pending (feat/patch-merging PR): Swin-style 2×2 patch merging with register-row passthrough | +| `patch_merging.py` | \[x\] | New (hierarchical PR): Swin-style 2×2 patch merging with optional register-row passthrough | ### `nvsubquadratic/networks/` — Full architectures @@ -77,7 +77,7 @@ Work bottom-up: primitive ops → modules → networks → experiments. | `general_purpose_resnet.py` | \[x\] | ResidualNetwork — LazyConfig blocks, conditioning, readout crop, gradient checkpointing | | `classification_resnet.py` | \[x\] | ClassificationResNet — GAP readout, resolution-agnostic, inherits ResidualNetwork | | `vit5_classification.py` | \[x\] | ViT5 classification — token layout, hybrid blocks, CLS/GAP/register_concat readouts, FLOP count | -| `vit5_hierarchical_classification.py` | \[ \] | Pending (feat/patch-merging PR): Swin-style 4-stage hierarchical ViT-5 classifier with GAP readout and optional register-row layout | +| `vit5_hierarchical_classification.py` | \[x\] | New (hierarchical PR): Swin-style 4-stage hierarchical ViT-5 classifier with GAP readout and optional register-row layout | | `huggingface_diffusers.py` | \[ \] | HF diffusers integration | | `jit.py` / `jit_utils.py` | \[ \] | TorchScript utilities | diff --git a/examples/vit5_imagenet/v6_hierarchical/__init__.py b/examples/vit5_imagenet/v6_hierarchical/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/vit5_imagenet/v6_hierarchical/_base_config.py b/examples/vit5_imagenet/v6_hierarchical/_base_config.py new file mode 100644 index 00000000..d12da159 --- /dev/null +++ b/examples/vit5_imagenet/v6_hierarchical/_base_config.py @@ -0,0 +1,394 @@ +"""Shared base config for v6_hierarchical: 4-stage Swin-style Hyena on ImageNet. + +Architecture: Swin-T-like stage layout (depths [2,2,6,2], dims [96,192,384,768]) +on top of the ViT-5 building blocks (ViT5ResidualBlock + ViT5HyenaAdapter). +Both pure and FiLM variants share this base; the per-stage block builder is +selected by the leaf config. + +Per stage, the Hyena mixer uses ``BlockDiagonalLearnableOmegaSIRENKernelND`` +(block-diagonal MLP init + learnable per-row ω₀ schedule) as the SIREN +kernel. All other Hyena ingredients (short conv, RMSNorm, L2 QK-norm) match +``v5_patch``. + +Patch sizes per stage (image 224, initial p=4): + stage 0 : grid 56x56, dim 96 + stage 1 : grid 28x28, dim 192 + stage 2 : grid 14x14, dim 384 + stage 3 : grid 7x 7, dim 768 + +Token counts (T fed into each stage's mixer): + pure : grid_h * grid_w + register_row : grid_w + grid_h * grid_w (registers as first row) + +Effective batch on 8 GPUs = 2048 (same as v5_patch). +""" + +from __future__ import annotations + +import os +from typing import Literal + +import torch + +# ── Model-only imports (always safe; no GPU / cluster deps) ───────────────── +from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig +from nvsubquadratic.modules.ckconv_nd import CKConvND +from nvsubquadratic.modules.film import KernelFiLMGenerator, RegisterPooling +from nvsubquadratic.modules.grn import GlobalResponseNorm +from nvsubquadratic.modules.hyena_nd import Hyena +from nvsubquadratic.modules.kernels_nd import BlockDiagonalLearnableOmegaSIRENKernelND +from nvsubquadratic.modules.mlp import MLP +from nvsubquadratic.modules.patch_merging import PatchMerging +from nvsubquadratic.modules.rms_norm import RMSNorm +from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer +from nvsubquadratic.modules.vit5_hyena_adapter import ViT5HyenaAdapter +from nvsubquadratic.modules.vit5_residual_block import ViT5ResidualBlock +from nvsubquadratic.networks.vit5_hierarchical_classification import ( + ViT5HierarchicalClassificationNet, +) +from nvsubquadratic.utils.init import trunc_normal_init, trunc_normal_init_factory +from nvsubquadratic.utils.qk_norm import L2Norm + + +# ── Training-infrastructure imports (apex / DALI / Lightning) ─────────────── +# These are intentionally deferred to inside get_base_config() so that +# build_hyena_hier_net() can be imported and instantiated in environments +# that don't have apex / DALI installed (CI, interactive testing, etc.). + + +# ─── Dataset ───────────────────────────────────────────────────────────────── +INPUT_CHANNELS = 3 +NUM_CLASSES = 1000 +IMAGE_SIZE = 224 +IMAGENET_PATH = os.environ.get("IMAGENET_PATH", "/scratch-nvme/ml-datasets/imagenet/torchvision_ImageNet/") +IMAGENET_FOLDER_PATH = os.environ.get( + "IMAGENET_FOLDER_PATH", "/scratch-nvme/ml-datasets/imagenet/torchvision_ImageFolder" +) +LOCAL_STAGING_DIR = os.environ.get("LOCAL_STAGING_DIR", "/scratch-nvme/ml-datasets/imagenet/torchvision_ImageFolder") + +# ─── Hierarchical layout (Swin-T-like) ─────────────────────────────────────── +INITIAL_PATCH_SIZE = 4 +STAGE_DIMS = [96, 192, 384, 768] +STAGE_DEPTHS = [2, 2, 6, 2] +NUM_STAGES = len(STAGE_DIMS) +INITIAL_GRID = IMAGE_SIZE // INITIAL_PATCH_SIZE # 56 +STAGE_GRIDS = [INITIAL_GRID // (2**i) for i in range(NUM_STAGES)] # 56, 28, 14, 7 + +NUM_HEADS = 6 # for QKVSequenceMixer in_proj split +LAYER_SCALE_INIT = 1e-4 +DROP_PATH_RATE = 0.05 +MLP_RATIO = 4 +NUM_REGISTERS = 4 # Used by the FiLM variant; ignored by pure variant. + +# ─── SIREN kernel (BD + learnable ω₀) ──────────────────────────────────────── +KERNEL_MLP_HIDDEN_DIM = 32 +KERNEL_NUM_LAYERS = 3 +KERNEL_EMBEDDING_DIM = 32 +KERNEL_NUM_BLOCKS = 8 +KERNEL_OMEGA_0_MIN = 1.0 +KERNEL_OMEGA_0_MAX = 12.0 +KERNEL_HIDDEN_OMEGA_0 = 1.0 +KERNEL_OFF_BLOCK_SCALE = 0.1 + +# ─── FiLM conditioning (variant-only) ──────────────────────────────────────── +FILM_HIDDEN_DIM = 64 + +# ─── Training recipe (matches v5_patch) ────────────────────────────────────── +EPOCHS = 800 +IMAGENET_TRAIN_SIZE = 1_281_167 +EFFECTIVE_BATCH_SIZE = 2048 +NUM_GPUS = 8 +ITERS_PER_EPOCH = IMAGENET_TRAIN_SIZE // EFFECTIVE_BATCH_SIZE +TOTAL_ITERATIONS = EPOCHS * ITERS_PER_EPOCH +WARMUP_EPOCHS = 5 +WARMUP_ITERATIONS_PERCENTAGE = WARMUP_EPOCHS / EPOCHS + +LEARNING_RATE = 4e-3 +WEIGHT_DECAY = 0.05 +GRAD_CLIP = 1.0 +PRECISION = "bf16-mixed" + +NUM_WORKERS = 12 + +# Effective batch = NUM_GPUS * batch_per_gpu * accum_steps = 2048 +# Patch-4 initial grid (3136 spatial tokens) matches v5_patch hyena_patch4's footprint. +BATCH_PER_GPU = 16 +ACCUM_STEPS = 16 + +INIT_FN = trunc_normal_init(std=0.02) +INIT_FN_FACTORY = trunc_normal_init_factory(std=0.02) + + +# ─── Builders ──────────────────────────────────────────────────────────────── + + +def _siren_kernel_cfg(hidden_dim: int, L_cache: int, film_cfg: LazyConfig | None) -> LazyConfig: + """SIREN kernel for one stage: BlockDiagonal + learnable per-row ω₀.""" + return LazyConfig(BlockDiagonalLearnableOmegaSIRENKernelND)( + data_dim=2, + out_dim=hidden_dim, + mlp_hidden_dim=KERNEL_MLP_HIDDEN_DIM, + num_layers=KERNEL_NUM_LAYERS, + embedding_dim=KERNEL_EMBEDDING_DIM, + L_cache=L_cache, + use_bias=True, + num_blocks=KERNEL_NUM_BLOCKS, + omega_0_min=KERNEL_OMEGA_0_MIN, + omega_0_max=KERNEL_OMEGA_0_MAX, + schedule="linear", + off_block_scale=KERNEL_OFF_BLOCK_SCALE, + hidden_omega_0=KERNEL_HIDDEN_OMEGA_0, + film_cfg=film_cfg, + ) + + +def _hyena_mixer_cfg(hidden_dim: int, L_cache: int, film_cfg: LazyConfig | None) -> LazyConfig: + """Hyena mixer for one stage (CKConvND + short conv + bilinear gates).""" + return LazyConfig(QKVSequenceMixer)( + hidden_dim=hidden_dim, + mixer_cfg=LazyConfig(Hyena)( + global_conv_cfg=LazyConfig(CKConvND)( + data_dim=2, + hidden_dim=hidden_dim, + kernel_cfg=_siren_kernel_cfg(hidden_dim=hidden_dim, L_cache=L_cache, film_cfg=film_cfg), + mask_cfg=LazyConfig(torch.nn.Identity)(), + grid_type="double", + fft_padding="zero", + ), + short_conv_cfg=LazyConfig(torch.nn.Conv2d)( + in_channels=3 * hidden_dim, + out_channels=3 * hidden_dim, + kernel_size=3, + groups=3 * hidden_dim, + padding=1, + bias=False, + ), + gate_nonlinear_cfg=LazyConfig(torch.nn.SiLU)(), + pixelhyena_norm_cfg=LazyConfig(RMSNorm)(dim=hidden_dim, eps=1e-6), + qk_norm_cfg=LazyConfig(L2Norm)(), + output_norm_cfg=LazyConfig(RMSNorm)(dim=hidden_dim, eps=1e-6), + gate_nonlinear_2_cfg=LazyConfig(torch.nn.Sigmoid)(), + ), + qkv_bias=False, + out_proj_bias=False, + init_method_in=INIT_FN_FACTORY, + init_method_out=INIT_FN_FACTORY, + ) + + +def _stage_block_cfg( + hidden_dim: int, + grid_w: int, + L_cache: int, + layout: Literal["pure", "register_row"], +) -> LazyConfig: + """Build a ViT5ResidualBlock config for one stage. + + For ``layout="register_row"``, the block owns a ``RegisterPooling`` whose + output feeds the SIREN kernel via FiLM. For ``layout="pure"``, no + register pooling is attached and ``film_cfg`` is None. + """ + if layout == "register_row": + film_cfg = LazyConfig(KernelFiLMGenerator)( + cond_dim=hidden_dim, + kernel_hidden_dim=KERNEL_MLP_HIDDEN_DIM, + num_film_layers=KERNEL_NUM_LAYERS - 1, + film_hidden_dim=FILM_HIDDEN_DIM, + ) + register_pooling_cfg = LazyConfig(RegisterPooling)(num_registers=NUM_REGISTERS) + num_registers = NUM_REGISTERS + else: + film_cfg = None + register_pooling_cfg = None + num_registers = 0 + + mixer_cfg = _hyena_mixer_cfg(hidden_dim=hidden_dim, L_cache=L_cache, film_cfg=film_cfg) + grn_cfg = LazyConfig(GlobalResponseNorm)(dim=hidden_dim) + + return LazyConfig(ViT5ResidualBlock)( + sequence_mixer_cfg=LazyConfig(ViT5HyenaAdapter)( + inner_mixer_cfg=mixer_cfg, + grid_w=grid_w, + ), + sequence_mixer_norm_cfg=LazyConfig(RMSNorm)(dim=hidden_dim, eps=1e-6), + mlp_cfg=LazyConfig(MLP)( + dim=hidden_dim, + activation="gelu", + expansion_factor=float(MLP_RATIO), + bias=False, + dropout_cfg=LazyConfig(torch.nn.Dropout)(p=0.0), + init_method_in=INIT_FN_FACTORY, + init_method_out=INIT_FN_FACTORY, + ), + mlp_norm_cfg=LazyConfig(RMSNorm)(dim=hidden_dim, eps=1e-6), + hidden_dim=hidden_dim, + layer_scale_init=LAYER_SCALE_INIT, + drop_path_rate=DROP_PATH_RATE, + register_pooling_cfg=register_pooling_cfg, + num_registers=num_registers, + register_start_idx=0, # registers are at the head of the sequence (no CLS) + grn_cfg=grn_cfg, + ) + + +def build_hyena_hier_net(layout: Literal["pure", "register_row"]) -> LazyConfig: + """Build the full hierarchical Hyena network config. + + Args: + layout: ``"pure"`` (no registers, no FiLM) or ``"register_row"`` (4 + registers prepended as the first 2D-grid row, conditioning every + stage's SIREN kernel via FiLM). + """ + # Per-stage block configs. L_cache covers both spatial dims; for the + # register-row layout the height direction is grid + 1 (register row added). + stage_block_cfgs = [] + for i in range(NUM_STAGES): + grid_i = STAGE_GRIDS[i] + L_cache_i = (grid_i + 1) if layout == "register_row" else grid_i + stage_block_cfgs.append( + _stage_block_cfg( + hidden_dim=STAGE_DIMS[i], + grid_w=grid_i, + L_cache=L_cache_i, + layout=layout, + ) + ) + + # Patch-merging configs between consecutive stages. + patch_merge_cfgs = [] + for i in range(NUM_STAGES - 1): + in_dim = STAGE_DIMS[i] + out_dim = STAGE_DIMS[i + 1] + grid_i = STAGE_GRIDS[i] + patch_merge_cfgs.append( + LazyConfig(PatchMerging)( + in_dim=in_dim, + out_dim=out_dim, + grid_h=grid_i, + grid_w=grid_i, + norm_cfg=LazyConfig(RMSNorm)(dim=4 * in_dim, eps=1e-6), + num_registers=NUM_REGISTERS if layout == "register_row" else 0, + has_register_row=(layout == "register_row"), + ) + ) + + return LazyConfig(ViT5HierarchicalClassificationNet)( + in_channels=INPUT_CHANNELS, + num_classes=NUM_CLASSES, + image_size=IMAGE_SIZE, + initial_patch_size=INITIAL_PATCH_SIZE, + stage_dims=STAGE_DIMS, + stage_depths=STAGE_DEPTHS, + stage_block_cfgs=stage_block_cfgs, + patch_merge_cfgs=patch_merge_cfgs, + norm_cfg=LazyConfig(RMSNorm)(dim=STAGE_DIMS[-1], eps=1e-6), + layout=layout, + num_registers=NUM_REGISTERS if layout == "register_row" else 0, + dropout_rate=0.0, + ) + + +def get_base_config(): + """Return the shared base experiment config (dataset, optim, scheduler, etc.). + + The caller must set ``config.net`` (use :func:`build_hyena_hier_net`). + + Training-infrastructure imports (apex, DALI, Lightning) are deferred to + here so that :func:`build_hyena_hier_net` is importable in environments + without those packages (CI, interactive testing, etc.). + """ + # ── deferred training-infra imports ────────────────────────────────────── + from apex.optimizers import FusedLAMB as Lamb + + from experiments.callbacks.model_ema import LabeledEMAWeightAveraging + from experiments.datamodules.dali_imagenet_fused import ( + AugmentConfig, + DALIImageNetFusedDataModule, + MixupConfig, + ) + from experiments.default_cfg import ( + AutoResumeConfig, + ExperimentConfig, + SchedulerConfig, + TrainConfig, + TrainerConfig, + WandbConfig, + ) + from experiments.lightning_wrappers.classification_wrapper import ClassificationWrapper + + # ───────────────────────────────────────────────────────────────────────── + config = ExperimentConfig() + config.debug = False + config.seed = 42 + config.compile = True + config.compile_mode = "max-autotune-no-cudagraphs" + config.compile_compatible_fftconv = True + + config.dataset = LazyConfig(DALIImageNetFusedDataModule)( + data_dir=IMAGENET_PATH, + imagefolder_dir=IMAGENET_FOLDER_PATH, + prefetch_factor=3, + batch_size=BATCH_PER_GPU, + num_workers=NUM_WORKERS, + pin_memory=True, + seed=config.seed, + image_size=IMAGE_SIZE, + final_image_size=IMAGE_SIZE, + num_classes=NUM_CLASSES, + drop_labels=False, + task="classification", + mixup_cfg=LazyConfig(MixupConfig)( + mixup=0.8, + cutmix=1.0, + mixup_prob=1.0, + mixup_switch_prob=0.5, + mixup_mode="batch", + smoothing=0.0, + ), + augment_cfg=LazyConfig(AugmentConfig)( + use_three_augment=True, + color_jitter=0.3, + ), + device_id=0, + local_staging_dir=LOCAL_STAGING_DIR, + ) + + config.lightning_wrapper_class = LazyConfig(ClassificationWrapper)(loss="soft_target_ce") + + config.optimizer = LazyConfig(Lamb)( + params=PLACEHOLDER, + lr=LEARNING_RATE, + weight_decay=WEIGHT_DECAY, + ) + + config.train = TrainConfig( + batch_size="${dataset.batch_size}", + iterations=TOTAL_ITERATIONS, + grad_clip=GRAD_CLIP, + precision=PRECISION, + accumulate_grad_steps=ACCUM_STEPS, + ) + + config.trainer = TrainerConfig( + check_val_every_n_epoch=4, + checkpoint_every_n_steps=5000, + ) + + config.scheduler = SchedulerConfig( + name="cosine", + warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE, + total_iterations="${train.iterations}", + mode="max", + ) + + config.callbacks = [LazyConfig(LabeledEMAWeightAveraging)(decay=0.99996)] + config.trainer.checkpoint_monitor = "val/acc_ema" + + config.wandb = WandbConfig( + job_group="v6_hierarchical", + entity="implicit-long-convs", + project="nvsubquadratic", + ) + + config.autoresume = AutoResumeConfig(enabled=False) + return config diff --git a/examples/vit5_imagenet/v6_hierarchical/_cifar10_patch_ablation_base.py b/examples/vit5_imagenet/v6_hierarchical/_cifar10_patch_ablation_base.py new file mode 100644 index 00000000..e75c5b96 --- /dev/null +++ b/examples/vit5_imagenet/v6_hierarchical/_cifar10_patch_ablation_base.py @@ -0,0 +1,336 @@ +"""Shared builders for the CIFAR-10 patch-size × hierarchy ablation. + +Images are upscaled to 64×64 so that patch_size=4 yields a 16×16 initial +grid, which is deep enough to support all 4 Swin-T / VMamba-T stages. + +Runs +---- +flat × {p4, p8, p16} — ViT5ClassificationNet, 10 blocks, dim=384, no merging +hier × {p4, p8, p16} — ViT5HierarchicalClassificationNet, patch merging + +Stage layout (Swin-T / VMamba-T: depths=[2,2,6,2], dims=[96,192,384,768]): + + patch_size=4 → initial grid 16×16 + flat : 10 blocks, dim=384 + hier : 4 stages 16→8→4→2 depths=[2,2,6,2] dims=[96,192,384,768] ← full Swin-T + + patch_size=8 → initial grid 8×8 + flat : 10 blocks, dim=384 + hier : 3 stages 8→4→2 depths=[2,2,6] dims=[96,192,384] ← Swin stages 1-3 + (4th stage would yield 1×1, hitting L_cache≥2 floor) + + patch_size=16 → initial grid 4×4 + flat : 10 blocks, dim=384 + hier : 2 stages 4→2 depths=[6,2] dims=[384,768] ← Swin stages 3-4 + (3rd stage would yield 1×1) + +Training (all runs) +------------------- + 200 epochs, batch=64, no grad-accum, AdamW lr=3e-4, cosine, bf16-mixed. + W&B job_group: "cifar10_patch_ablation_64px" +""" + +import os + +import torch +from torch.optim import AdamW + +from experiments.datamodules.cifar10 import CIFAR10DataModule +from experiments.default_cfg import ( + AutoResumeConfig, + ExperimentConfig, + SchedulerConfig, + TrainConfig, + TrainerConfig, + WandbConfig, +) +from experiments.lightning_wrappers.classification_wrapper import ClassificationWrapper +from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig +from nvsubquadratic.modules.ckconv_nd import CKConvND +from nvsubquadratic.modules.grn import GlobalResponseNorm +from nvsubquadratic.modules.hyena_nd import Hyena +from nvsubquadratic.modules.kernels_nd import BlockDiagonalLearnableOmegaSIRENKernelND +from nvsubquadratic.modules.mlp import MLP +from nvsubquadratic.modules.patch_merging import PatchMerging +from nvsubquadratic.modules.rms_norm import RMSNorm +from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer +from nvsubquadratic.modules.vit5_hyena_adapter import ViT5HyenaAdapter +from nvsubquadratic.modules.vit5_residual_block import ViT5ResidualBlock +from nvsubquadratic.networks.vit5_classification import ViT5ClassificationNet +from nvsubquadratic.networks.vit5_hierarchical_classification import ( + ViT5HierarchicalClassificationNet, +) +from nvsubquadratic.utils.init import trunc_normal_init_factory +from nvsubquadratic.utils.qk_norm import L2Norm + + +# ─── Dataset ────────────────────────────────────────────────────────────────── +CIFAR10_DIR = os.environ.get("CIFAR10_DIR", "./data") +IMAGE_SIZE = 64 # upscaled from 32 to allow 4 Swin stages with patch_size=4 +NUM_CLASSES = 10 + +# ─── Shared training hyper-parameters ───────────────────────────────────────── +BATCH_SIZE = 64 +TARGET_EPOCHS = 200 +CIFAR10_TRAIN_SIZE = 50_000 +ITERS_PER_EPOCH = CIFAR10_TRAIN_SIZE // BATCH_SIZE # 781 +TOTAL_ITERATIONS = TARGET_EPOCHS * ITERS_PER_EPOCH # 156_200 +WARMUP_EPOCHS = 5 +WARMUP_ITERATIONS_PERCENTAGE = WARMUP_EPOCHS / TARGET_EPOCHS # 0.025 + +LEARNING_RATE = 3e-4 +WEIGHT_DECAY = 0.05 +GRAD_CLIP = 1.0 +PRECISION = "bf16-mixed" + +LAYER_SCALE_INIT = 1e-4 +DROP_PATH_RATE = 0.05 +MLP_RATIO = 4 + +# ─── SIREN kernel ────────────────────────────────────────────────────────────── +KERNEL_MLP_HIDDEN_DIM = 32 +KERNEL_NUM_LAYERS = 3 +KERNEL_EMBEDDING_DIM = 32 +KERNEL_NUM_BLOCKS = 8 # divides 96, 192, 384, 768 cleanly +KERNEL_OMEGA_0_MIN = 1.0 +KERNEL_OMEGA_0_MAX = 12.0 +KERNEL_HIDDEN_OMEGA_0 = 1.0 +KERNEL_OFF_BLOCK_SCALE = 0.1 + +# ─── Flat architecture ──────────────────────────────────────────────────────── +FLAT_HIDDEN_DIM = 384 +FLAT_NUM_BLOCKS = 10 + +# ─── Hier architecture — Swin-T / VMamba-T stage maps ──────────────────────── +# p=4: 4 stages 16→8→4→2 — full Swin-T +HIER_P4_STAGE_DIMS = [96, 192, 384, 768] +HIER_P4_STAGE_DEPTHS = [2, 2, 6, 2] + +# p=8: 3 stages 8→4→2 — Swin stages 1-3 +HIER_P8_STAGE_DIMS = [96, 192, 384] +HIER_P8_STAGE_DEPTHS = [2, 2, 6] + +# p=16: 2 stages 4→2 — Swin stages 3-4 +# A 4×4 initial grid corresponds to the output of Swin stage 2 in the p=4 hierarchy. +HIER_P16_STAGE_DIMS = [384, 768] +HIER_P16_STAGE_DEPTHS = [6, 2] + +INIT_FN_FACTORY = trunc_normal_init_factory(std=0.02) + + +# ─── Block builder ──────────────────────────────────────────────────────────── + + +def _build_block(hidden_dim: int, grid_w: int) -> LazyConfig: + """Return a ViT5ResidualBlock LazyConfig for the given (hidden_dim, grid_w).""" + kernel_cfg = LazyConfig(BlockDiagonalLearnableOmegaSIRENKernelND)( + data_dim=2, + out_dim=hidden_dim, + mlp_hidden_dim=KERNEL_MLP_HIDDEN_DIM, + num_layers=KERNEL_NUM_LAYERS, + embedding_dim=KERNEL_EMBEDDING_DIM, + L_cache=grid_w, + use_bias=True, + num_blocks=KERNEL_NUM_BLOCKS, + omega_0_min=KERNEL_OMEGA_0_MIN, + omega_0_max=KERNEL_OMEGA_0_MAX, + schedule="linear", + off_block_scale=KERNEL_OFF_BLOCK_SCALE, + hidden_omega_0=KERNEL_HIDDEN_OMEGA_0, + film_cfg=None, + ) + mixer_cfg = LazyConfig(QKVSequenceMixer)( + hidden_dim=hidden_dim, + mixer_cfg=LazyConfig(Hyena)( + global_conv_cfg=LazyConfig(CKConvND)( + data_dim=2, + hidden_dim=hidden_dim, + kernel_cfg=kernel_cfg, + mask_cfg=LazyConfig(torch.nn.Identity)(), + grid_type="double", + fft_padding="zero", + ), + short_conv_cfg=LazyConfig(torch.nn.Conv2d)( + in_channels=3 * hidden_dim, + out_channels=3 * hidden_dim, + kernel_size=3, + groups=3 * hidden_dim, + padding=1, + bias=False, + ), + gate_nonlinear_cfg=LazyConfig(torch.nn.SiLU)(), + pixelhyena_norm_cfg=LazyConfig(RMSNorm)(dim=hidden_dim, eps=1e-6), + qk_norm_cfg=LazyConfig(L2Norm)(), + output_norm_cfg=LazyConfig(RMSNorm)(dim=hidden_dim, eps=1e-6), + gate_nonlinear_2_cfg=LazyConfig(torch.nn.Sigmoid)(), + ), + qkv_bias=False, + out_proj_bias=False, + init_method_in=INIT_FN_FACTORY, + init_method_out=INIT_FN_FACTORY, + ) + return LazyConfig(ViT5ResidualBlock)( + sequence_mixer_cfg=LazyConfig(ViT5HyenaAdapter)( + inner_mixer_cfg=mixer_cfg, + grid_w=grid_w, + ), + sequence_mixer_norm_cfg=LazyConfig(RMSNorm)(dim=hidden_dim, eps=1e-6), + mlp_cfg=LazyConfig(MLP)( + dim=hidden_dim, + activation="gelu", + expansion_factor=float(MLP_RATIO), + bias=False, + dropout_cfg=LazyConfig(torch.nn.Dropout)(p=0.0), + init_method_in=INIT_FN_FACTORY, + init_method_out=INIT_FN_FACTORY, + ), + mlp_norm_cfg=LazyConfig(RMSNorm)(dim=hidden_dim, eps=1e-6), + hidden_dim=hidden_dim, + layer_scale_init=LAYER_SCALE_INIT, + drop_path_rate=DROP_PATH_RATE, + grn_cfg=LazyConfig(GlobalResponseNorm)(dim=hidden_dim), + ) + + +# ─── Shared training config skeleton ────────────────────────────────────────── + + +def _base_config() -> ExperimentConfig: + """Return an ExperimentConfig pre-filled with the shared training settings.""" + config = ExperimentConfig() + config.debug = False + config.seed = 42 + config.compile = False + + config.dataset = LazyConfig(CIFAR10DataModule)( + data_dir=CIFAR10_DIR, + batch_size=BATCH_SIZE, + num_workers=4, + pin_memory=True, + image_size=IMAGE_SIZE, + mixup=0.8, + cutmix=1.0, + ) + + config.lightning_wrapper_class = LazyConfig(ClassificationWrapper)(loss="soft_target_ce") + + config.optimizer = LazyConfig(AdamW)( + params=PLACEHOLDER, + lr=LEARNING_RATE, + weight_decay=WEIGHT_DECAY, + ) + + config.train = TrainConfig( + batch_size="${dataset.batch_size}", + iterations=TOTAL_ITERATIONS, + grad_clip=GRAD_CLIP, + precision=PRECISION, + accumulate_grad_steps=1, + ) + + config.trainer = TrainerConfig( + check_val_every_n_epoch=5, + checkpoint_every_n_steps=None, + ) + + config.scheduler = SchedulerConfig( + name="cosine", + warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE, + total_iterations="${train.iterations}", + mode="max", + ) + + config.wandb = WandbConfig( + job_group="cifar10_patch_ablation_64px", + entity="implicit-long-convs", + project="nvsubquadratic", + ) + + config.autoresume = AutoResumeConfig(enabled=False) + return config + + +# ─── Public config builders ─────────────────────────────────────────────────── + + +def build_flat_config(patch_size: int) -> ExperimentConfig: + """Flat baseline (no patch merging) for the given patch_size ∈ {4, 8, 16}.""" + assert patch_size in {4, 8, 16}, f"Unsupported patch_size={patch_size}" + grid = IMAGE_SIZE // patch_size # 16, 8, or 4 + + block_cfg = _build_block(hidden_dim=FLAT_HIDDEN_DIM, grid_w=grid) + + net = LazyConfig(ViT5ClassificationNet)( + in_channels=3, + num_classes=NUM_CLASSES, + hidden_dim=FLAT_HIDDEN_DIM, + num_blocks=FLAT_NUM_BLOCKS, + patch_size=patch_size, + image_size=IMAGE_SIZE, + num_registers=0, + block_cfg=block_cfg, + norm_cfg=LazyConfig(RMSNorm)(dim=FLAT_HIDDEN_DIM, eps=1e-6), + readout="gap", + dropout_rate=0.0, + ) + + config = _base_config() + config.net = net + return config + + +def build_hier_config(patch_size: int) -> ExperimentConfig: + """Hierarchical config (with patch merging) for the given patch_size ∈ {4, 8, 16}. + + Stage layout (Swin-T / VMamba-T consistent, images upscaled to 64×64): + p=4 : 4 stages 16→8→4→2 depths=[2,2,6,2] dims=[96,192,384,768] ← full Swin-T + p=8 : 3 stages 8→4→2 depths=[2,2,6] dims=[96,192,384] ← Swin stages 1-3 + p=16: 2 stages 4→2 depths=[6,2] dims=[384,768] ← Swin stages 3-4 + """ + assert patch_size in {4, 8, 16}, f"Unsupported patch_size={patch_size}" + + if patch_size == 4: + stage_dims = HIER_P4_STAGE_DIMS + stage_depths = HIER_P4_STAGE_DEPTHS + elif patch_size == 8: + stage_dims = HIER_P8_STAGE_DIMS + stage_depths = HIER_P8_STAGE_DEPTHS + else: # patch_size == 16 + stage_dims = HIER_P16_STAGE_DIMS + stage_depths = HIER_P16_STAGE_DEPTHS + + num_stages = len(stage_dims) + initial_grid = IMAGE_SIZE // patch_size + stage_grids = [initial_grid // (2**i) for i in range(num_stages)] + + stage_block_cfgs = [_build_block(hidden_dim=stage_dims[i], grid_w=stage_grids[i]) for i in range(num_stages)] + patch_merge_cfgs = [ + LazyConfig(PatchMerging)( + in_dim=stage_dims[i], + out_dim=stage_dims[i + 1], + grid_h=stage_grids[i], + grid_w=stage_grids[i], + norm_cfg=LazyConfig(RMSNorm)(dim=4 * stage_dims[i], eps=1e-6), + has_register_row=False, + ) + for i in range(num_stages - 1) + ] + + net = LazyConfig(ViT5HierarchicalClassificationNet)( + in_channels=3, + num_classes=NUM_CLASSES, + image_size=IMAGE_SIZE, + initial_patch_size=patch_size, + stage_dims=stage_dims, + stage_depths=stage_depths, + stage_block_cfgs=stage_block_cfgs, + patch_merge_cfgs=patch_merge_cfgs, + norm_cfg=LazyConfig(RMSNorm)(dim=stage_dims[-1], eps=1e-6), + layout="pure", + num_registers=0, + dropout_rate=0.0, + ) + + config = _base_config() + config.net = net + return config diff --git a/examples/vit5_imagenet/v6_hierarchical/_smoke_test.py b/examples/vit5_imagenet/v6_hierarchical/_smoke_test.py new file mode 100644 index 00000000..d355bd2d --- /dev/null +++ b/examples/vit5_imagenet/v6_hierarchical/_smoke_test.py @@ -0,0 +1,92 @@ +"""Smoke test for v6_hierarchical configs: build net + run one forward pass on GPU.""" + +import importlib +import sys +import time + +import torch + +from nvsubquadratic.lazy_config import instantiate + + +CONFIGS = [ + "examples.vit5_imagenet.v6_hierarchical.hyena_hier_p4_pure", + "examples.vit5_imagenet.v6_hierarchical.hyena_hier_p4_film", +] + +BATCH_SIZE = 1 +IMAGE_SIZE = 224 +DEVICE = "cuda" + + +def smoke_test_config(module_path: str) -> None: + name = module_path.rsplit(".", 1)[-1] + print(f"\n{'=' * 60}") + print(f" {name}") + print(f"{'=' * 60}") + + mod = importlib.import_module(module_path) + config = mod.get_config() + net = instantiate(config.net).to(DEVICE) + + num_params = sum(p.numel() for p in net.parameters()) / 1e6 + print(f" params: {num_params:.1f}M") + print(f" layout: {net.layout}") + print(f" stage_dims: {net.stage_dims}") + print(f" stage_depths: {net.stage_depths}") + print(f" stage_grids: {net.stage_grid_sides}") + if net.layout == "register_row": + print(f" num_registers: {net.num_registers}") + + x = { + "input": torch.randn(BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 3, device=DEVICE), + "condition": None, + } + + net.eval() + torch.cuda.reset_peak_memory_stats() + with torch.no_grad(), torch.autocast(DEVICE, dtype=torch.bfloat16): + t0 = time.perf_counter() + out = net(x) + torch.cuda.synchronize() + dt = time.perf_counter() - t0 + + logits = out["logits"] + peak_mb = torch.cuda.max_memory_allocated() / 1e6 + print(f" logits shape: {logits.shape}") + print(f" forward time: {dt:.3f}s") + print(f" peak GPU mem: {peak_mb:.0f} MB") + print(" PASS") + + del net, x, out + torch.cuda.empty_cache() + + +def main() -> int: + if not torch.cuda.is_available(): + print("ERROR: CUDA not available") + return 1 + + print(f"GPU: {torch.cuda.get_device_name(0)}") + passed, failed = [], [] + + for cfg in CONFIGS: + try: + smoke_test_config(cfg) + passed.append(cfg.rsplit(".", 1)[-1]) + except Exception as e: + name = cfg.rsplit(".", 1)[-1] + print(f"\n FAIL: {name} — {type(e).__name__}: {e}") + failed.append(name) + torch.cuda.empty_cache() + + print(f"\n{'=' * 60}") + print(f" RESULTS: {len(passed)} passed, {len(failed)} failed") + if failed: + print(f" Failed: {', '.join(failed)}") + print(f"{'=' * 60}") + return 1 if failed else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/vit5_imagenet/v6_hierarchical/cifar10_flat_p16.py b/examples/vit5_imagenet/v6_hierarchical/cifar10_flat_p16.py new file mode 100644 index 00000000..7bacabe9 --- /dev/null +++ b/examples/vit5_imagenet/v6_hierarchical/cifar10_flat_p16.py @@ -0,0 +1,10 @@ +"""CIFAR-10 patch-size ablation: flat (no merging), patch_size=16, grid=2×2, dim=384.""" + +from examples.vit5_imagenet.v6_hierarchical._cifar10_patch_ablation_base import ( + build_flat_config, +) + + +def get_config(): + """Return the flat patch-16 CIFAR-10 experiment config.""" + return build_flat_config(patch_size=16) diff --git a/examples/vit5_imagenet/v6_hierarchical/cifar10_flat_p4.py b/examples/vit5_imagenet/v6_hierarchical/cifar10_flat_p4.py new file mode 100644 index 00000000..6394828c --- /dev/null +++ b/examples/vit5_imagenet/v6_hierarchical/cifar10_flat_p4.py @@ -0,0 +1,10 @@ +"""CIFAR-10 patch-size ablation: flat (no merging), patch_size=4, grid=8×8, dim=384.""" + +from examples.vit5_imagenet.v6_hierarchical._cifar10_patch_ablation_base import ( + build_flat_config, +) + + +def get_config(): + """Return the flat patch-4 CIFAR-10 experiment config.""" + return build_flat_config(patch_size=4) diff --git a/examples/vit5_imagenet/v6_hierarchical/cifar10_flat_p8.py b/examples/vit5_imagenet/v6_hierarchical/cifar10_flat_p8.py new file mode 100644 index 00000000..d5a3c39b --- /dev/null +++ b/examples/vit5_imagenet/v6_hierarchical/cifar10_flat_p8.py @@ -0,0 +1,10 @@ +"""CIFAR-10 patch-size ablation: flat (no merging), patch_size=8, grid=4×4, dim=384.""" + +from examples.vit5_imagenet.v6_hierarchical._cifar10_patch_ablation_base import ( + build_flat_config, +) + + +def get_config(): + """Return the flat patch-8 CIFAR-10 experiment config.""" + return build_flat_config(patch_size=8) diff --git a/examples/vit5_imagenet/v6_hierarchical/cifar10_hier_p16.py b/examples/vit5_imagenet/v6_hierarchical/cifar10_hier_p16.py new file mode 100644 index 00000000..86be7917 --- /dev/null +++ b/examples/vit5_imagenet/v6_hierarchical/cifar10_hier_p16.py @@ -0,0 +1,13 @@ +"""CIFAR-10 patch-size ablation (64×64): hierarchical, patch_size=16. + +2 stages: 4×4 → 2×2, dims=[384,768], depths=[6,2] ← Swin stages 3-4. +""" + +from examples.vit5_imagenet.v6_hierarchical._cifar10_patch_ablation_base import ( + build_hier_config, +) + + +def get_config(): + """Return the hierarchical patch-16 CIFAR-10 experiment config.""" + return build_hier_config(patch_size=16) diff --git a/examples/vit5_imagenet/v6_hierarchical/cifar10_hier_p4.py b/examples/vit5_imagenet/v6_hierarchical/cifar10_hier_p4.py new file mode 100644 index 00000000..aa8cfc0d --- /dev/null +++ b/examples/vit5_imagenet/v6_hierarchical/cifar10_hier_p4.py @@ -0,0 +1,13 @@ +"""CIFAR-10 patch-size ablation: hierarchical, patch_size=4. + +3 stages: 8×8 → 4×4 → 2×2, dims=[96,192,384], depths=[2,2,6]. +""" + +from examples.vit5_imagenet.v6_hierarchical._cifar10_patch_ablation_base import ( + build_hier_config, +) + + +def get_config(): + """Return the hierarchical patch-4 CIFAR-10 experiment config.""" + return build_hier_config(patch_size=4) diff --git a/examples/vit5_imagenet/v6_hierarchical/cifar10_hier_p8.py b/examples/vit5_imagenet/v6_hierarchical/cifar10_hier_p8.py new file mode 100644 index 00000000..a21d88bc --- /dev/null +++ b/examples/vit5_imagenet/v6_hierarchical/cifar10_hier_p8.py @@ -0,0 +1,13 @@ +"""CIFAR-10 patch-size ablation: hierarchical, patch_size=8. + +2 stages: 4×4 → 2×2, dims=[192,384], depths=[4,6]. +""" + +from examples.vit5_imagenet.v6_hierarchical._cifar10_patch_ablation_base import ( + build_hier_config, +) + + +def get_config(): + """Return the hierarchical patch-8 CIFAR-10 experiment config.""" + return build_hier_config(patch_size=8) diff --git a/examples/vit5_imagenet/v6_hierarchical/cifar10_hyena_flat.py b/examples/vit5_imagenet/v6_hierarchical/cifar10_hyena_flat.py new file mode 100644 index 00000000..1f9b2160 --- /dev/null +++ b/examples/vit5_imagenet/v6_hierarchical/cifar10_hyena_flat.py @@ -0,0 +1,226 @@ +"""CIFAR-10 quick-eval: flat Hyena (10 blocks, no patch merging). + +Architecture + initial p=2 → 16×16 = 256 tokens, dim=384 (matches hier stage-2 dim) + 10 blocks, no PatchMerging; GAP readout; pure layout. + SIREN kernel: BlockDiagonalLearnableOmegaSIRENKernelND. + +Training recipe + Identical to cifar10_hyena_hier: ~100 epochs, AdamW, cosine, bf16-mixed. + +The flat run is the no-merging baseline: same epoch budget, same kernel, +same compute per block — only the hierarchical structure differs. +""" + +import os + +import torch +from torch.optim import AdamW + +from experiments.datamodules.cifar10 import CIFAR10DataModule +from experiments.default_cfg import ( + AutoResumeConfig, + ExperimentConfig, + SchedulerConfig, + TrainConfig, + TrainerConfig, + WandbConfig, +) +from experiments.lightning_wrappers.classification_wrapper import ClassificationWrapper +from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig +from nvsubquadratic.modules.ckconv_nd import CKConvND +from nvsubquadratic.modules.grn import GlobalResponseNorm +from nvsubquadratic.modules.hyena_nd import Hyena +from nvsubquadratic.modules.kernels_nd import BlockDiagonalLearnableOmegaSIRENKernelND +from nvsubquadratic.modules.mlp import MLP +from nvsubquadratic.modules.rms_norm import RMSNorm +from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer +from nvsubquadratic.modules.vit5_hyena_adapter import ViT5HyenaAdapter +from nvsubquadratic.modules.vit5_residual_block import ViT5ResidualBlock +from nvsubquadratic.networks.vit5_classification import ViT5ClassificationNet +from nvsubquadratic.utils.init import trunc_normal_init_factory +from nvsubquadratic.utils.qk_norm import L2Norm + + +# ─── Dataset ────────────────────────────────────────────────────────────────── +CIFAR10_DIR = os.environ.get("CIFAR10_DIR", "./data") +IMAGE_SIZE = 32 +NUM_CLASSES = 10 +# dim=384 × 16×16 grid × large batch causes OOM on 22GB L4; use batch=64 +# with accum=4 to keep effective batch = 256 matching the hier run. +BATCH_SIZE = 64 +ACCUM_STEPS = 4 + +# ─── Architecture ───────────────────────────────────────────────────────────── +PATCH_SIZE = 2 +NUM_PATCHES_H = IMAGE_SIZE // PATCH_SIZE # 16 +NUM_PATCHES_W = IMAGE_SIZE // PATCH_SIZE # 16 +HIDDEN_DIM = 384 # matches the last stage of the hierarchical model +NUM_BLOCKS = 10 + +LAYER_SCALE_INIT = 1e-4 +DROP_PATH_RATE = 0.05 +MLP_RATIO = 4 + +# ─── SIREN kernel (same as hier) ─────────────────────────────────────────────── +KERNEL_MLP_HIDDEN_DIM = 32 +KERNEL_NUM_LAYERS = 3 +KERNEL_EMBEDDING_DIM = 32 +KERNEL_NUM_BLOCKS = 8 +KERNEL_OMEGA_0_MIN = 1.0 +KERNEL_OMEGA_0_MAX = 12.0 +KERNEL_HIDDEN_OMEGA_0 = 1.0 +KERNEL_OFF_BLOCK_SCALE = 0.1 + +# ─── Training (identical to hier) ───────────────────────────────────────────── +TARGET_EPOCHS = 100 +CIFAR10_TRAIN_SIZE = 50_000 +ITERS_PER_EPOCH = CIFAR10_TRAIN_SIZE // BATCH_SIZE +TOTAL_ITERATIONS = TARGET_EPOCHS * ITERS_PER_EPOCH +WARMUP_EPOCHS = 5 +WARMUP_ITERATIONS_PERCENTAGE = WARMUP_EPOCHS / TARGET_EPOCHS + +LEARNING_RATE = 3e-4 +WEIGHT_DECAY = 0.05 +GRAD_CLIP = 1.0 +PRECISION = "bf16-mixed" + +INIT_FN_FACTORY = trunc_normal_init_factory(std=0.02) + + +def get_config() -> ExperimentConfig: + """Return the CIFAR-10 flat (no patch-merging) Hyena config.""" + kernel_cfg = LazyConfig(BlockDiagonalLearnableOmegaSIRENKernelND)( + data_dim=2, + out_dim=HIDDEN_DIM, + mlp_hidden_dim=KERNEL_MLP_HIDDEN_DIM, + num_layers=KERNEL_NUM_LAYERS, + embedding_dim=KERNEL_EMBEDDING_DIM, + L_cache=NUM_PATCHES_H, + use_bias=True, + num_blocks=KERNEL_NUM_BLOCKS, + omega_0_min=KERNEL_OMEGA_0_MIN, + omega_0_max=KERNEL_OMEGA_0_MAX, + schedule="linear", + off_block_scale=KERNEL_OFF_BLOCK_SCALE, + hidden_omega_0=KERNEL_HIDDEN_OMEGA_0, + film_cfg=None, + ) + mixer_cfg = LazyConfig(QKVSequenceMixer)( + hidden_dim=HIDDEN_DIM, + mixer_cfg=LazyConfig(Hyena)( + global_conv_cfg=LazyConfig(CKConvND)( + data_dim=2, + hidden_dim=HIDDEN_DIM, + kernel_cfg=kernel_cfg, + mask_cfg=LazyConfig(torch.nn.Identity)(), + grid_type="double", + fft_padding="zero", + ), + short_conv_cfg=LazyConfig(torch.nn.Conv2d)( + in_channels=3 * HIDDEN_DIM, + out_channels=3 * HIDDEN_DIM, + kernel_size=3, + groups=3 * HIDDEN_DIM, + padding=1, + bias=False, + ), + gate_nonlinear_cfg=LazyConfig(torch.nn.SiLU)(), + pixelhyena_norm_cfg=LazyConfig(RMSNorm)(dim=HIDDEN_DIM, eps=1e-6), + qk_norm_cfg=LazyConfig(L2Norm)(), + output_norm_cfg=LazyConfig(RMSNorm)(dim=HIDDEN_DIM, eps=1e-6), + gate_nonlinear_2_cfg=LazyConfig(torch.nn.Sigmoid)(), + ), + qkv_bias=False, + out_proj_bias=False, + init_method_in=INIT_FN_FACTORY, + init_method_out=INIT_FN_FACTORY, + ) + block_cfg = LazyConfig(ViT5ResidualBlock)( + sequence_mixer_cfg=LazyConfig(ViT5HyenaAdapter)( + inner_mixer_cfg=mixer_cfg, + grid_w=NUM_PATCHES_W, + ), + sequence_mixer_norm_cfg=LazyConfig(RMSNorm)(dim=HIDDEN_DIM, eps=1e-6), + mlp_cfg=LazyConfig(MLP)( + dim=HIDDEN_DIM, + activation="gelu", + expansion_factor=float(MLP_RATIO), + bias=False, + dropout_cfg=LazyConfig(torch.nn.Dropout)(p=0.0), + init_method_in=INIT_FN_FACTORY, + init_method_out=INIT_FN_FACTORY, + ), + mlp_norm_cfg=LazyConfig(RMSNorm)(dim=HIDDEN_DIM, eps=1e-6), + hidden_dim=HIDDEN_DIM, + layer_scale_init=LAYER_SCALE_INIT, + drop_path_rate=DROP_PATH_RATE, + grn_cfg=LazyConfig(GlobalResponseNorm)(dim=HIDDEN_DIM), + ) + + net = LazyConfig(ViT5ClassificationNet)( + in_channels=3, + num_classes=NUM_CLASSES, + hidden_dim=HIDDEN_DIM, + num_blocks=NUM_BLOCKS, + patch_size=PATCH_SIZE, + image_size=IMAGE_SIZE, + num_registers=0, + block_cfg=block_cfg, + norm_cfg=LazyConfig(RMSNorm)(dim=HIDDEN_DIM, eps=1e-6), + readout="gap", + dropout_rate=0.0, + ) + + config = ExperimentConfig() + config.debug = False + config.seed = 42 + config.compile = False + config.net = net + + config.dataset = LazyConfig(CIFAR10DataModule)( + data_dir=CIFAR10_DIR, + batch_size=BATCH_SIZE, + num_workers=4, + pin_memory=True, + image_size=IMAGE_SIZE, + mixup=0.8, + cutmix=1.0, + ) + + config.lightning_wrapper_class = LazyConfig(ClassificationWrapper)(loss="soft_target_ce") + + config.optimizer = LazyConfig(AdamW)( + params=PLACEHOLDER, + lr=LEARNING_RATE, + weight_decay=WEIGHT_DECAY, + ) + + config.train = TrainConfig( + batch_size="${dataset.batch_size}", + iterations=TOTAL_ITERATIONS, + grad_clip=GRAD_CLIP, + precision=PRECISION, + accumulate_grad_steps=ACCUM_STEPS, + ) + + config.trainer = TrainerConfig( + check_val_every_n_epoch=5, + checkpoint_every_n_steps=None, + ) + + config.scheduler = SchedulerConfig( + name="cosine", + warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE, + total_iterations="${train.iterations}", + mode="max", + ) + + config.wandb = WandbConfig( + job_group="cifar10_hier_vs_flat", + entity="implicit-long-convs", + project="nvsubquadratic", + ) + + config.autoresume = AutoResumeConfig(enabled=False) + return config diff --git a/examples/vit5_imagenet/v6_hierarchical/cifar10_hyena_hier.py b/examples/vit5_imagenet/v6_hierarchical/cifar10_hyena_hier.py new file mode 100644 index 00000000..92f651f4 --- /dev/null +++ b/examples/vit5_imagenet/v6_hierarchical/cifar10_hyena_hier.py @@ -0,0 +1,247 @@ +"""CIFAR-10 quick-eval: hierarchical Hyena (3 stages, with patch merging). + +Architecture + initial p=2 → 16×16 grid → dims [96, 192, 384], depths [2, 2, 6] + PatchMerging between stages (16→8→4); GAP readout; pure layout (no FiLM). + SIREN kernel: BlockDiagonalLearnableOmegaSIRENKernelND. + +Training recipe + ~100 epochs (50 000 × 100 / 256 ≈ 19 500 steps) + AdamW lr=3e-4, cosine schedule, 5-epoch warmup + bf16-mixed, single GPU + Standard augmentation: RandomCrop(32,pad=4) + RandomHorizontalFlip + +This is a diagnostic run to compare classification performance against the +flat (no patch-merging) baseline at the same compute and epoch budget. +""" + +import os + +import torch +from torch.optim import AdamW + +from experiments.datamodules.cifar10 import CIFAR10DataModule +from experiments.default_cfg import ( + AutoResumeConfig, + ExperimentConfig, + SchedulerConfig, + TrainConfig, + TrainerConfig, + WandbConfig, +) +from experiments.lightning_wrappers.classification_wrapper import ClassificationWrapper +from nvsubquadratic.lazy_config import PLACEHOLDER, LazyConfig +from nvsubquadratic.modules.ckconv_nd import CKConvND +from nvsubquadratic.modules.grn import GlobalResponseNorm +from nvsubquadratic.modules.hyena_nd import Hyena +from nvsubquadratic.modules.kernels_nd import BlockDiagonalLearnableOmegaSIRENKernelND +from nvsubquadratic.modules.mlp import MLP +from nvsubquadratic.modules.patch_merging import PatchMerging +from nvsubquadratic.modules.rms_norm import RMSNorm +from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer +from nvsubquadratic.modules.vit5_hyena_adapter import ViT5HyenaAdapter +from nvsubquadratic.modules.vit5_residual_block import ViT5ResidualBlock +from nvsubquadratic.networks.vit5_hierarchical_classification import ( + ViT5HierarchicalClassificationNet, +) +from nvsubquadratic.utils.init import trunc_normal_init_factory +from nvsubquadratic.utils.qk_norm import L2Norm + + +# ─── Dataset ────────────────────────────────────────────────────────────────── +CIFAR10_DIR = os.environ.get("CIFAR10_DIR", "./data") +IMAGE_SIZE = 32 +NUM_CLASSES = 10 +BATCH_SIZE = 256 + +# ─── Architecture ───────────────────────────────────────────────────────────── +INITIAL_PATCH_SIZE = 2 +STAGE_DIMS = [96, 192, 384] +STAGE_DEPTHS = [2, 2, 6] +NUM_STAGES = len(STAGE_DIMS) +INITIAL_GRID = IMAGE_SIZE // INITIAL_PATCH_SIZE # 16 +STAGE_GRIDS = [INITIAL_GRID // (2**i) for i in range(NUM_STAGES)] # 16, 8, 4 + +LAYER_SCALE_INIT = 1e-4 +DROP_PATH_RATE = 0.05 +MLP_RATIO = 4 + +# ─── SIREN kernel ────────────────────────────────────────────────────────────── +KERNEL_MLP_HIDDEN_DIM = 32 +KERNEL_NUM_LAYERS = 3 +KERNEL_EMBEDDING_DIM = 32 +KERNEL_NUM_BLOCKS = 8 +KERNEL_OMEGA_0_MIN = 1.0 +KERNEL_OMEGA_0_MAX = 12.0 +KERNEL_HIDDEN_OMEGA_0 = 1.0 +KERNEL_OFF_BLOCK_SCALE = 0.1 + +# ─── Training ────────────────────────────────────────────────────────────────── +TARGET_EPOCHS = 100 +CIFAR10_TRAIN_SIZE = 50_000 +ITERS_PER_EPOCH = CIFAR10_TRAIN_SIZE // BATCH_SIZE # 195 +TOTAL_ITERATIONS = TARGET_EPOCHS * ITERS_PER_EPOCH # ~19 500 +WARMUP_EPOCHS = 5 +WARMUP_ITERATIONS_PERCENTAGE = WARMUP_EPOCHS / TARGET_EPOCHS + +LEARNING_RATE = 3e-4 +WEIGHT_DECAY = 0.05 +GRAD_CLIP = 1.0 +PRECISION = "bf16-mixed" + +INIT_FN_FACTORY = trunc_normal_init_factory(std=0.02) + + +def _build_block(hidden_dim: int, grid_w: int) -> LazyConfig: + L_cache = grid_w + kernel_cfg = LazyConfig(BlockDiagonalLearnableOmegaSIRENKernelND)( + data_dim=2, + out_dim=hidden_dim, + mlp_hidden_dim=KERNEL_MLP_HIDDEN_DIM, + num_layers=KERNEL_NUM_LAYERS, + embedding_dim=KERNEL_EMBEDDING_DIM, + L_cache=L_cache, + use_bias=True, + num_blocks=KERNEL_NUM_BLOCKS, + omega_0_min=KERNEL_OMEGA_0_MIN, + omega_0_max=KERNEL_OMEGA_0_MAX, + schedule="linear", + off_block_scale=KERNEL_OFF_BLOCK_SCALE, + hidden_omega_0=KERNEL_HIDDEN_OMEGA_0, + film_cfg=None, + ) + mixer_cfg = LazyConfig(QKVSequenceMixer)( + hidden_dim=hidden_dim, + mixer_cfg=LazyConfig(Hyena)( + global_conv_cfg=LazyConfig(CKConvND)( + data_dim=2, + hidden_dim=hidden_dim, + kernel_cfg=kernel_cfg, + mask_cfg=LazyConfig(torch.nn.Identity)(), + grid_type="double", + fft_padding="zero", + ), + short_conv_cfg=LazyConfig(torch.nn.Conv2d)( + in_channels=3 * hidden_dim, + out_channels=3 * hidden_dim, + kernel_size=3, + groups=3 * hidden_dim, + padding=1, + bias=False, + ), + gate_nonlinear_cfg=LazyConfig(torch.nn.SiLU)(), + pixelhyena_norm_cfg=LazyConfig(RMSNorm)(dim=hidden_dim, eps=1e-6), + qk_norm_cfg=LazyConfig(L2Norm)(), + output_norm_cfg=LazyConfig(RMSNorm)(dim=hidden_dim, eps=1e-6), + gate_nonlinear_2_cfg=LazyConfig(torch.nn.Sigmoid)(), + ), + qkv_bias=False, + out_proj_bias=False, + init_method_in=INIT_FN_FACTORY, + init_method_out=INIT_FN_FACTORY, + ) + return LazyConfig(ViT5ResidualBlock)( + sequence_mixer_cfg=LazyConfig(ViT5HyenaAdapter)( + inner_mixer_cfg=mixer_cfg, + grid_w=grid_w, + ), + sequence_mixer_norm_cfg=LazyConfig(RMSNorm)(dim=hidden_dim, eps=1e-6), + mlp_cfg=LazyConfig(MLP)( + dim=hidden_dim, + activation="gelu", + expansion_factor=float(MLP_RATIO), + bias=False, + dropout_cfg=LazyConfig(torch.nn.Dropout)(p=0.0), + init_method_in=INIT_FN_FACTORY, + init_method_out=INIT_FN_FACTORY, + ), + mlp_norm_cfg=LazyConfig(RMSNorm)(dim=hidden_dim, eps=1e-6), + hidden_dim=hidden_dim, + layer_scale_init=LAYER_SCALE_INIT, + drop_path_rate=DROP_PATH_RATE, + grn_cfg=LazyConfig(GlobalResponseNorm)(dim=hidden_dim), + ) + + +def get_config() -> ExperimentConfig: + """Return the CIFAR-10 hierarchical Hyena config.""" + stage_block_cfgs = [_build_block(hidden_dim=STAGE_DIMS[i], grid_w=STAGE_GRIDS[i]) for i in range(NUM_STAGES)] + patch_merge_cfgs = [ + LazyConfig(PatchMerging)( + in_dim=STAGE_DIMS[i], + out_dim=STAGE_DIMS[i + 1], + grid_h=STAGE_GRIDS[i], + grid_w=STAGE_GRIDS[i], + norm_cfg=LazyConfig(RMSNorm)(dim=4 * STAGE_DIMS[i], eps=1e-6), + has_register_row=False, + ) + for i in range(NUM_STAGES - 1) + ] + + net = LazyConfig(ViT5HierarchicalClassificationNet)( + in_channels=3, + num_classes=NUM_CLASSES, + image_size=IMAGE_SIZE, + initial_patch_size=INITIAL_PATCH_SIZE, + stage_dims=STAGE_DIMS, + stage_depths=STAGE_DEPTHS, + stage_block_cfgs=stage_block_cfgs, + patch_merge_cfgs=patch_merge_cfgs, + norm_cfg=LazyConfig(RMSNorm)(dim=STAGE_DIMS[-1], eps=1e-6), + layout="pure", + num_registers=0, + dropout_rate=0.0, + ) + + config = ExperimentConfig() + config.debug = False + config.seed = 42 + config.compile = False # keep off for quick iteration + config.net = net + + config.dataset = LazyConfig(CIFAR10DataModule)( + data_dir=CIFAR10_DIR, + batch_size=BATCH_SIZE, + num_workers=4, + pin_memory=True, + image_size=IMAGE_SIZE, + mixup=0.8, + cutmix=1.0, + ) + + config.lightning_wrapper_class = LazyConfig(ClassificationWrapper)(loss="soft_target_ce") + + config.optimizer = LazyConfig(AdamW)( + params=PLACEHOLDER, + lr=LEARNING_RATE, + weight_decay=WEIGHT_DECAY, + ) + + config.train = TrainConfig( + batch_size="${dataset.batch_size}", + iterations=TOTAL_ITERATIONS, + grad_clip=GRAD_CLIP, + precision=PRECISION, + accumulate_grad_steps=1, + ) + + config.trainer = TrainerConfig( + check_val_every_n_epoch=5, + checkpoint_every_n_steps=None, + ) + + config.scheduler = SchedulerConfig( + name="cosine", + warmup_iterations_percentage=WARMUP_ITERATIONS_PERCENTAGE, + total_iterations="${train.iterations}", + mode="max", + ) + + config.wandb = WandbConfig( + job_group="cifar10_hier_vs_flat", + entity="implicit-long-convs", + project="nvsubquadratic", + ) + + config.autoresume = AutoResumeConfig(enabled=False) + return config diff --git a/examples/vit5_imagenet/v6_hierarchical/hyena_hier_p4_film.py b/examples/vit5_imagenet/v6_hierarchical/hyena_hier_p4_film.py new file mode 100644 index 00000000..19a82767 --- /dev/null +++ b/examples/vit5_imagenet/v6_hierarchical/hyena_hier_p4_film.py @@ -0,0 +1,22 @@ +"""v6_hierarchical — Hyena 4-stage with register-row FiLM, BD + learnable ω₀. + +Same architecture as ``hyena_hier_p4_pure`` but with 4 register tokens +prepended as the first row of the 2D grid at every stage. Each block pools +its registers (RegisterPooling) and uses the result to FiLM-modulate the +SIREN kernel. No CLS token; GAP readout skips the register row. + +Stages [96, 192, 384, 768] x depths [2, 2, 6, 2]; effective batch = 2048. +""" + +from examples.vit5_imagenet.v6_hierarchical._base_config import ( + build_hyena_hier_net, + get_base_config, +) +from experiments.default_cfg import ExperimentConfig + + +def get_config() -> ExperimentConfig: + """Return the FiLM-conditioned hierarchical Hyena config.""" + config = get_base_config() + config.net = build_hyena_hier_net(layout="register_row") + return config diff --git a/examples/vit5_imagenet/v6_hierarchical/hyena_hier_p4_pure.py b/examples/vit5_imagenet/v6_hierarchical/hyena_hier_p4_pure.py new file mode 100644 index 00000000..05a16150 --- /dev/null +++ b/examples/vit5_imagenet/v6_hierarchical/hyena_hier_p4_pure.py @@ -0,0 +1,21 @@ +"""v6_hierarchical — pure Hyena 4-stage (Swin-T layout) with BD + learnable ω₀. + +No registers, no FiLM. Stages [96, 192, 384, 768] x depths [2, 2, 6, 2] on +initial patch size 4 (grids 56 -> 28 -> 14 -> 7). PatchMerging between stages. +SIREN kernel: ``BlockDiagonalLearnableOmegaSIRENKernelND``. + +Effective batch = 2048 (8 GPUs x 16 per-GPU x 16 accum). +""" + +from examples.vit5_imagenet.v6_hierarchical._base_config import ( + build_hyena_hier_net, + get_base_config, +) +from experiments.default_cfg import ExperimentConfig + + +def get_config() -> ExperimentConfig: + """Return the pure (no-FiLM) hierarchical Hyena config.""" + config = get_base_config() + config.net = build_hyena_hier_net(layout="pure") + return config diff --git a/nvsubquadratic/networks/vit5_hierarchical_classification.py b/nvsubquadratic/networks/vit5_hierarchical_classification.py new file mode 100644 index 00000000..9a24e965 --- /dev/null +++ b/nvsubquadratic/networks/vit5_hierarchical_classification.py @@ -0,0 +1,259 @@ +"""Hierarchical ViT-5 classification network (Swin-style stages + patch merging). + +Differs from ``ViT5ClassificationNet`` in three ways: + +1. **Hierarchical stages.** Blocks are organized into ``num_stages`` groups + with a ``PatchMerging`` between consecutive stages. Each stage halves the + spatial resolution and (typically) doubles the channel dim, mirroring + Swin-T's ``[2, 2, 6, 2]`` / ``[C, 2C, 4C, 8C]`` structure. + +2. **GAP readout only.** No CLS token. Two layouts are supported: + * ``pure``: tokens are a flat patch grid ``[B, H*W, C]``. + * ``register_row``: the first ``grid_w`` tokens form a register row used + by the Hyena mixer for FiLM conditioning; GAP excludes them. + +3. **Per-stage block configs.** The caller supplies one block config per + stage (replicated across that stage's blocks) plus the patch-merging + config for the transition into it. +""" + +from typing import List, Literal + +import torch +import torch.nn as nn +from einops import rearrange + +from nvsubquadratic.lazy_config import LazyConfig, instantiate + + +class ViT5HierarchicalClassificationNet(nn.Module): + """Swin-style hierarchical ViT-5 classifier with optional register-row layout. + + Args: + in_channels: Input image channels (3 for RGB). + num_classes: Number of output classes. + image_size: Input image size (assumes square). + initial_patch_size: Stride of the initial Conv2d patch embed. The + stage-0 grid is ``(image_size // initial_patch_size)`` per side. + stage_dims: Channel dim per stage, ``len == num_stages``. + ``stage_dims[0]`` is the patch-embed output dim. + stage_depths: Number of blocks per stage, ``len == num_stages``. + stage_block_cfgs: Per-stage block LazyConfig. Each is replicated + ``stage_depths[i]`` times. The block is expected to be a + ``ViT5ResidualBlock`` whose sequence mixer matches the stage's + grid size and hidden dim. + patch_merge_cfgs: PatchMerging LazyConfig for each stage transition. + Length ``num_stages - 1``. ``patch_merge_cfgs[i]`` runs between + stage ``i`` and stage ``i+1``. + norm_cfg: Final norm before the head (configured for ``dim = + stage_dims[-1]``). + layout: ``"pure"`` (no registers) or ``"register_row"`` (registers + prepended as the first row of the 2D grid at every stage). + num_registers: Number of register tokens (only used for + ``layout="register_row"``). Constant across stages; padding is + recomputed per stage to match the halved grid width. + reg_init: Register init scheme — ``"trunc_normal"`` (default) or + ``"zeros"``. + dropout_rate: Dropout before the head. + """ + + def __init__( + self, + in_channels: int, + num_classes: int, + image_size: int, + initial_patch_size: int, + stage_dims: List[int], + stage_depths: List[int], + stage_block_cfgs: List[LazyConfig], + patch_merge_cfgs: List[LazyConfig], + norm_cfg: LazyConfig, + layout: Literal["pure", "register_row"] = "pure", + num_registers: int = 0, + reg_init: Literal["trunc_normal", "zeros"] = "trunc_normal", + dropout_rate: float = 0.0, + ): + """Initialise the hierarchical ViT-5 network and validate stage configs.""" + super().__init__() + num_stages = len(stage_dims) + if len(stage_depths) != num_stages: + raise ValueError(f"stage_depths len ({len(stage_depths)}) != num_stages ({num_stages})") + if len(stage_block_cfgs) != num_stages: + raise ValueError(f"stage_block_cfgs len ({len(stage_block_cfgs)}) != num_stages ({num_stages})") + if len(patch_merge_cfgs) != num_stages - 1: + raise ValueError( + f"patch_merge_cfgs len ({len(patch_merge_cfgs)}) must equal num_stages - 1 ({num_stages - 1})" + ) + if layout not in ("pure", "register_row"): + raise ValueError(f"layout must be 'pure' or 'register_row', got {layout!r}") + if layout == "register_row" and num_registers <= 0: + raise ValueError("layout='register_row' requires num_registers > 0") + + self.num_stages = num_stages + self.stage_dims = list(stage_dims) + self.stage_depths = list(stage_depths) + self.layout = layout + self.num_registers = num_registers + self.image_size = image_size + self.initial_patch_size = initial_patch_size + self._reg_init = reg_init + self.num_classes = num_classes + + # Per-stage grid widths (used by GAP to skip the register row). + # Stage i's grid is (initial_grid // 2**i) per side. + initial_grid = image_size // initial_patch_size + if initial_grid * initial_patch_size != image_size: + raise ValueError( + f"image_size ({image_size}) must be divisible by initial_patch_size ({initial_patch_size})" + ) + self.stage_grid_sides = [initial_grid // (2**i) for i in range(num_stages)] + if self.stage_grid_sides[-1] < 1: + raise ValueError( + f"Final stage grid collapses to {self.stage_grid_sides[-1]} — too many stages for grid {initial_grid}" + ) + + if layout == "register_row": + final_grid_w = self.stage_grid_sides[-1] + if num_registers > final_grid_w: + raise ValueError( + f"num_registers ({num_registers}) must fit in the final stage grid_w ({final_grid_w})" + ) + + # Patch embedding: matches the v5_patch convention (no bias; APE absorbs offset). + self.patch_embed = nn.Conv2d( + in_channels, + stage_dims[0], + kernel_size=initial_patch_size, + stride=initial_patch_size, + padding=0, + bias=False, + ) + + num_patches_0 = initial_grid * initial_grid + self.pos_embed = nn.Parameter(torch.zeros(1, num_patches_0, stage_dims[0])) + self.pos_embed._no_weight_decay = True + + if layout == "register_row": + self.reg_token = nn.Parameter(torch.zeros(1, num_registers, stage_dims[0])) + self.reg_token._no_weight_decay = True + pad_size = initial_grid - num_registers + if pad_size > 0: + self.register_buffer("reg_zero_pad", torch.zeros(1, pad_size, stage_dims[0]), persistent=False) + else: + self.reg_zero_pad = None + else: + self.reg_token = None + self.reg_zero_pad = None + + # Build stages: ModuleList of ModuleList, one block-list per stage. + self.stage_blocks = nn.ModuleList( + [ + nn.ModuleList([instantiate(cfg) for _ in range(stage_depths[i])]) + for i, cfg in enumerate(stage_block_cfgs) + ] + ) + # Patch merges between stages. + self.patch_merges = nn.ModuleList([instantiate(cfg) for cfg in patch_merge_cfgs]) + + self.out_norm = instantiate(norm_cfg) + for p in self.out_norm.parameters(): + p._no_weight_decay = True + self.dropout = nn.Dropout(dropout_rate) if dropout_rate > 0 else nn.Identity() + self.out_proj = nn.Linear(stage_dims[-1], num_classes, bias=False) + + self._init_weights() + + def _init_weights(self): + nn.init.trunc_normal_(self.pos_embed, std=0.02) + if self.reg_token is not None: + if self._reg_init == "zeros": + nn.init.zeros_(self.reg_token) + else: + nn.init.trunc_normal_(self.reg_token, std=0.02) + nn.init.trunc_normal_(self.patch_embed.weight, std=0.02) + nn.init.trunc_normal_(self.out_proj.weight, std=0.02) + + def flop_count(self, inference: bool = False) -> int: + """FLOPs for one forward pass (one sample). + + Breakdown: + 1. Conv2d patch embed: 2 * in_ch * stage_dims[0] * P^2 * num_patches_0. + 2. APE add: num_patches_0 * stage_dims[0]. + 3. For each stage: + - Each block: block.flop_count(T_stage, inference) where T_stage + counts the register row when present. + 4. For each transition: patch_merge.flop_count() (already a per-sample count). + 5. Final norm on 1 GAP'd token: out_norm.flop_count(1). + 6. GAP itself: num_patches_last * stage_dims[-1] adds. + 7. Head: 2 * stage_dims[-1] * num_classes. + """ + D0 = self.stage_dims[0] + P = self.initial_patch_size + num_patches_0 = self.stage_grid_sides[0] ** 2 + in_ch = self.patch_embed.in_channels + + flops = 0 + flops += 2 * in_ch * D0 * P * P * num_patches_0 + flops += num_patches_0 * D0 + + for i in range(self.num_stages): + grid = self.stage_grid_sides[i] + T = grid * grid + (grid if self.layout == "register_row" else 0) + for block in self.stage_blocks[i]: + flops += block.flop_count(T, inference=inference) + if i < self.num_stages - 1: + flops += self.patch_merges[i].flop_count() + + last_grid = self.stage_grid_sides[-1] + last_patches = last_grid * last_grid + flops += last_patches * self.stage_dims[-1] # GAP add + flops += self.out_norm.flop_count(1) + flops += 2 * self.stage_dims[-1] * self.num_classes + return flops + + def forward(self, input_and_condition: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Forward pass. + + Args: + input_and_condition: Dict with key ``"input"`` of shape + ``[B, H, W, C]`` (channels-last) and optional ``"condition"`` + (unused). + + Returns: + ``{"logits": [B, num_classes]}``. + """ + x = input_and_condition["input"] # [B, H, W, C] + x = rearrange(x, "b h w c -> b c h w") + x = self.patch_embed(x) # [B, D0, grid, grid] + x = rearrange(x, "b c h w -> b (h w) c") # [B, num_patches_0, D0] + x = x + self.pos_embed + + B = x.shape[0] + grid_w = self.stage_grid_sides[0] + + if self.layout == "register_row": + regs = self.reg_token.expand(B, -1, -1) + if self.reg_zero_pad is not None: + pad = self.reg_zero_pad.expand(B, -1, -1) + x = torch.cat([regs, pad, x], dim=1) + else: + x = torch.cat([regs, x], dim=1) + + # Stages with patch merging between. + for i, block_list in enumerate(self.stage_blocks): + for block in block_list: + x = block(x) + if i < self.num_stages - 1: + x = self.patch_merges[i](x) + grid_w = self.stage_grid_sides[i + 1] + + # GAP over patch tokens (skip register row if present). + if self.layout == "register_row": + out = x[:, grid_w:].mean(dim=1) + else: + out = x.mean(dim=1) + + out = self.out_norm(out) + out = self.dropout(out) + logits = self.out_proj(out) + return {"logits": logits} diff --git a/tests/networks/test_vit5_hierarchical_classification.py b/tests/networks/test_vit5_hierarchical_classification.py new file mode 100644 index 00000000..9f1d18d5 --- /dev/null +++ b/tests/networks/test_vit5_hierarchical_classification.py @@ -0,0 +1,167 @@ +"""Tests for ViT5HierarchicalClassificationNet. + +Uses a minimal stand-in mixer (linear over channels) so the tests focus on +the network's hierarchical plumbing — stages, patch merging, register-row +preservation, GAP readout — rather than on any specific mixer's correctness. +""" + +import pytest +import torch +import torch.nn as nn + +from nvsubquadratic.lazy_config import LazyConfig +from nvsubquadratic.modules.mlp import MLP +from nvsubquadratic.modules.patch_merging import PatchMerging +from nvsubquadratic.modules.rms_norm import RMSNorm +from nvsubquadratic.modules.vit5_residual_block import ViT5ResidualBlock +from nvsubquadratic.networks.vit5_hierarchical_classification import ( + ViT5HierarchicalClassificationNet, +) + + +class _LinearMixer(nn.Module): + """Trivial sequence mixer used to validate the network's plumbing. + + Applies the same per-channel linear at every position. No spatial mixing — + we don't need it to test the hierarchy itself. + """ + + def __init__(self, hidden_dim: int): + super().__init__() + self.proj = nn.Linear(hidden_dim, hidden_dim, bias=False) + + def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor: + return self.proj(x) + + def flop_count(self, num_tokens: int, inference: bool = False) -> int: + D = self.proj.in_features + return 2 * num_tokens * D * D + + +def _make_block_cfg(hidden_dim: int) -> LazyConfig: + return LazyConfig(ViT5ResidualBlock)( + sequence_mixer_cfg=LazyConfig(_LinearMixer)(hidden_dim=hidden_dim), + sequence_mixer_norm_cfg=LazyConfig(RMSNorm)(dim=hidden_dim, eps=1e-6, use_quack=False), + mlp_cfg=LazyConfig(MLP)( + dim=hidden_dim, + activation="gelu", + expansion_factor=4.0, + bias=False, + dropout_cfg=LazyConfig(nn.Dropout)(p=0.0), + ), + mlp_norm_cfg=LazyConfig(RMSNorm)(dim=hidden_dim, eps=1e-6, use_quack=False), + hidden_dim=hidden_dim, + layer_scale_init=1e-4, + drop_path_rate=0.0, + ) + + +def _make_pm_cfg(in_dim: int, out_dim: int, grid: int, num_regs: int, has_reg_row: bool) -> LazyConfig: + return LazyConfig(PatchMerging)( + in_dim=in_dim, + out_dim=out_dim, + grid_h=grid, + grid_w=grid, + norm_cfg=LazyConfig(RMSNorm)(dim=4 * in_dim, eps=1e-6, use_quack=False), + num_registers=num_regs, + has_register_row=has_reg_row, + ) + + +def _build_net(layout: str, image_size: int = 32, p0: int = 4, num_registers: int = 0): + """Build a small 4-stage hierarchical net suitable for unit testing.""" + stage_dims = [16, 32, 64, 128] + stage_depths = [1, 1, 2, 1] + num_stages = len(stage_dims) + initial_grid = image_size // p0 # 8 + grids = [initial_grid // (2**i) for i in range(num_stages)] # 8,4,2,1 + + block_cfgs = [_make_block_cfg(d) for d in stage_dims] + pm_cfgs = [ + _make_pm_cfg( + in_dim=stage_dims[i], + out_dim=stage_dims[i + 1], + grid=grids[i], + num_regs=num_registers, + has_reg_row=(layout == "register_row"), + ) + for i in range(num_stages - 1) + ] + + return ViT5HierarchicalClassificationNet( + in_channels=3, + num_classes=10, + image_size=image_size, + initial_patch_size=p0, + stage_dims=stage_dims, + stage_depths=stage_depths, + stage_block_cfgs=block_cfgs, + patch_merge_cfgs=pm_cfgs, + norm_cfg=LazyConfig(RMSNorm)(dim=stage_dims[-1], eps=1e-6, use_quack=False), + layout=layout, + num_registers=num_registers, + ) + + +@pytest.mark.parametrize("layout,num_registers", [("pure", 0), ("register_row", 1)]) +def test_forward_shape(device, layout: str, num_registers: int) -> None: + """Forward returns logits of the expected shape regardless of layout.""" + net = _build_net(layout=layout, num_registers=num_registers).to(device) + x = torch.randn(2, 32, 32, 3, device=device) + out = net({"input": x}) + assert out["logits"].shape == (2, 10) + + +def test_pure_and_register_row_param_counts_close(device) -> None: + """The two variants should differ only by reg-related params (small delta).""" + net_pure = _build_net(layout="pure", num_registers=0).to(device) + net_reg = _build_net(layout="register_row", num_registers=1).to(device) + n_pure = sum(p.numel() for p in net_pure.parameters()) + n_reg = sum(p.numel() for p in net_reg.parameters()) + # The reg variant adds: reg_token (D0) + per-stage-transition reg_proj + # weights ([in_dim * out_dim] for 3 transitions). Should be a small delta. + assert n_reg > n_pure + assert (n_reg - n_pure) < n_pure * 0.1 # < 10% extra + + +def test_flop_count_is_positive(device) -> None: + net = _build_net(layout="pure").to(device) + assert net.flop_count() > 0 + + +def test_invalid_stage_lengths_raise() -> None: + """Mismatched per-stage list lengths raise ValueError at construction.""" + with pytest.raises(ValueError, match="stage_depths"): + _ = ViT5HierarchicalClassificationNet( + in_channels=3, + num_classes=10, + image_size=32, + initial_patch_size=4, + stage_dims=[16, 32, 64, 128], + stage_depths=[1, 1, 1], # wrong length + stage_block_cfgs=[_make_block_cfg(d) for d in [16, 32, 64, 128]], + patch_merge_cfgs=[ + _make_pm_cfg(16, 32, 8, 0, False), + _make_pm_cfg(32, 64, 4, 0, False), + _make_pm_cfg(64, 128, 2, 0, False), + ], + norm_cfg=LazyConfig(RMSNorm)(dim=128, eps=1e-6, use_quack=False), + layout="pure", + ) + + +def test_register_row_requires_positive_num_registers() -> None: + with pytest.raises(ValueError, match="num_registers > 0"): + _build_net(layout="register_row", num_registers=0) + + +def test_backward_runs(device) -> None: + """A backward pass on a small batch updates params without error.""" + net = _build_net(layout="register_row", num_registers=1).to(device) + x = torch.randn(2, 32, 32, 3, device=device) + out = net({"input": x}) + loss = out["logits"].sum() + loss.backward() + # Spot-check: patch_embed and reg_proj have non-zero grads. + assert net.patch_embed.weight.grad is not None + assert net.patch_merges[0].reg_proj.weight.grad is not None From 20e3be720eb6579c2cb1e4f456a6af279ab1114a Mon Sep 17 00:00:00 2001 From: David Wessels Date: Mon, 25 May 2026 15:22:18 +0200 Subject: [PATCH 71/72] refactor(v6_hierarchical): move CIFAR-10 configs to cifar10/ subfolder; align ImageNet configs with vit5_hybrid reference MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## CIFAR-10 subfolder Move all CIFAR-10 experiment files into examples/vit5_imagenet/v6_hierarchical/cifar10/: _cifar10_patch_ablation_base.py → cifar10/_base.py cifar10_{flat,hier}_p{4,8,16}.py → cifar10/{flat,hier}_p{4,8,16}.py cifar10_hyena_{hier,flat}.py → cifar10/hyena_{hier,flat}.py Leaf configs import updated to reference cifar10._base. ## _base_config.py — align with vit5_hybrid/_film.py + _learnable_omega.py - Mask: swap torch.nn.Identity for BlockAlignedGaussianModulationND (data_dim=2, extent=1.0, direct parametrization) matching apply_learnable_omega_blockdiag_overrides in the reference config. - FiLM constants: add FILM_INIT_TYPE='identity', FILM_WEIGHT_DECAY=5e-3, FILM_AFTER_POS_EMBED=True matching vit5_hybrid/_film.py best-run defaults. - KernelFiLMGenerator: add init_type + no_weight_decay; bump num_film_layers from KERNEL_NUM_LAYERS-1 to KERNEL_NUM_LAYERS (3) since film_after_pos_embed adds one extra pair for the positional-embedding sine. - _siren_kernel_cfg: set cfg.film_after_pos_embed=True when film_cfg provided. ## hyena_hier_p4_{pure,film}.py — match full_hyena_learnable_omega_blockdiag.py - Append MaskMonitorCallback + OmegaScaleMonitorCallback (log every 50 steps). - Set distinct wandb.job_group: 'v6_hier_pure' / 'v6_hier_film'. Co-Authored-By: Claude Sonnet 4.6 --- .../v6_hierarchical/_base_config.py | 33 ++++++++++++++++--- .../v6_hierarchical/cifar10/__init__.py | 0 .../_base.py} | 0 .../flat_p16.py} | 6 ++-- .../flat_p4.py} | 6 ++-- .../flat_p8.py} | 6 ++-- .../hier_p16.py} | 6 ++-- .../hier_p4.py} | 6 ++-- .../hier_p8.py} | 6 ++-- .../hyena_flat.py} | 0 .../hyena_hier.py} | 0 .../v6_hierarchical/hyena_hier_p4_film.py | 26 +++++++++++---- .../v6_hierarchical/hyena_hier_p4_pure.py | 19 +++++++---- 13 files changed, 74 insertions(+), 40 deletions(-) create mode 100644 examples/vit5_imagenet/v6_hierarchical/cifar10/__init__.py rename examples/vit5_imagenet/v6_hierarchical/{_cifar10_patch_ablation_base.py => cifar10/_base.py} (100%) rename examples/vit5_imagenet/v6_hierarchical/{cifar10_flat_p16.py => cifar10/flat_p16.py} (59%) rename examples/vit5_imagenet/v6_hierarchical/{cifar10_flat_p4.py => cifar10/flat_p4.py} (59%) rename examples/vit5_imagenet/v6_hierarchical/{cifar10_flat_p8.py => cifar10/flat_p8.py} (59%) rename examples/vit5_imagenet/v6_hierarchical/{cifar10_hier_p16.py => cifar10/hier_p16.py} (54%) rename examples/vit5_imagenet/v6_hierarchical/{cifar10_hier_p4.py => cifar10/hier_p4.py} (53%) rename examples/vit5_imagenet/v6_hierarchical/{cifar10_hier_p8.py => cifar10/hier_p8.py} (55%) rename examples/vit5_imagenet/v6_hierarchical/{cifar10_hyena_flat.py => cifar10/hyena_flat.py} (100%) rename examples/vit5_imagenet/v6_hierarchical/{cifar10_hyena_hier.py => cifar10/hyena_hier.py} (100%) diff --git a/examples/vit5_imagenet/v6_hierarchical/_base_config.py b/examples/vit5_imagenet/v6_hierarchical/_base_config.py index d12da159..9560e1f8 100644 --- a/examples/vit5_imagenet/v6_hierarchical/_base_config.py +++ b/examples/vit5_imagenet/v6_hierarchical/_base_config.py @@ -37,6 +37,7 @@ from nvsubquadratic.modules.grn import GlobalResponseNorm from nvsubquadratic.modules.hyena_nd import Hyena from nvsubquadratic.modules.kernels_nd import BlockDiagonalLearnableOmegaSIRENKernelND +from nvsubquadratic.modules.masks_nd import BlockAlignedGaussianModulationND from nvsubquadratic.modules.mlp import MLP from nvsubquadratic.modules.patch_merging import PatchMerging from nvsubquadratic.modules.rms_norm import RMSNorm @@ -91,7 +92,12 @@ KERNEL_OFF_BLOCK_SCALE = 0.1 # ─── FiLM conditioning (variant-only) ──────────────────────────────────────── +# Matches vit5_hybrid/_film.py: identity init (γ=1, β=0), film_after_pos_embed, +# and dedicated per-group WD of 5e-3 (best v3 LAMB+FiLM result on ImageNet). FILM_HIDDEN_DIM = 64 +FILM_INIT_TYPE = "identity" # γ=1, β=0 at init — start from no-modulation baseline +FILM_WEIGHT_DECAY = 5e-3 # dedicated WD group for FiLM weights (biases excluded) +FILM_AFTER_POS_EMBED = True # modulate the positional-embedding sine (adds one FiLM pair) # ─── Training recipe (matches v5_patch) ────────────────────────────────────── EPOCHS = 800 @@ -123,8 +129,13 @@ def _siren_kernel_cfg(hidden_dim: int, L_cache: int, film_cfg: LazyConfig | None) -> LazyConfig: - """SIREN kernel for one stage: BlockDiagonal + learnable per-row ω₀.""" - return LazyConfig(BlockDiagonalLearnableOmegaSIRENKernelND)( + """SIREN kernel for one stage: BlockDiagonal + learnable per-row ω₀. + + When ``film_cfg`` is provided, ``film_after_pos_embed=True`` is set so the + FiLM generator modulates the positional-embedding sine in addition to the + two hidden linear layers, matching the vit5_hybrid/_film.py convention. + """ + cfg = LazyConfig(BlockDiagonalLearnableOmegaSIRENKernelND)( data_dim=2, out_dim=hidden_dim, mlp_hidden_dim=KERNEL_MLP_HIDDEN_DIM, @@ -140,6 +151,9 @@ def _siren_kernel_cfg(hidden_dim: int, L_cache: int, film_cfg: LazyConfig | None hidden_omega_0=KERNEL_HIDDEN_OMEGA_0, film_cfg=film_cfg, ) + if film_cfg is not None: + cfg.film_after_pos_embed = FILM_AFTER_POS_EMBED + return cfg def _hyena_mixer_cfg(hidden_dim: int, L_cache: int, film_cfg: LazyConfig | None) -> LazyConfig: @@ -151,7 +165,14 @@ def _hyena_mixer_cfg(hidden_dim: int, L_cache: int, film_cfg: LazyConfig | None) data_dim=2, hidden_dim=hidden_dim, kernel_cfg=_siren_kernel_cfg(hidden_dim=hidden_dim, L_cache=L_cache, film_cfg=film_cfg), - mask_cfg=LazyConfig(torch.nn.Identity)(), + mask_cfg=LazyConfig(BlockAlignedGaussianModulationND)( + data_dim=2, + num_channels=hidden_dim, + min_attenuation_at_step=0.1, + max_attenuation_at_limit=0.95, + init_extent=1.0, + parametrization="direct", + ), grid_type="double", fft_padding="zero", ), @@ -192,8 +213,12 @@ def _stage_block_cfg( film_cfg = LazyConfig(KernelFiLMGenerator)( cond_dim=hidden_dim, kernel_hidden_dim=KERNEL_MLP_HIDDEN_DIM, - num_film_layers=KERNEL_NUM_LAYERS - 1, + # FILM_AFTER_POS_EMBED=True adds a 3rd pair (for the pos-embed sine) + # so num_film_layers must equal KERNEL_NUM_LAYERS (not KERNEL_NUM_LAYERS-1). + num_film_layers=KERNEL_NUM_LAYERS, film_hidden_dim=FILM_HIDDEN_DIM, + init_type=FILM_INIT_TYPE, + no_weight_decay=FILM_WEIGHT_DECAY, ) register_pooling_cfg = LazyConfig(RegisterPooling)(num_registers=NUM_REGISTERS) num_registers = NUM_REGISTERS diff --git a/examples/vit5_imagenet/v6_hierarchical/cifar10/__init__.py b/examples/vit5_imagenet/v6_hierarchical/cifar10/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/vit5_imagenet/v6_hierarchical/_cifar10_patch_ablation_base.py b/examples/vit5_imagenet/v6_hierarchical/cifar10/_base.py similarity index 100% rename from examples/vit5_imagenet/v6_hierarchical/_cifar10_patch_ablation_base.py rename to examples/vit5_imagenet/v6_hierarchical/cifar10/_base.py diff --git a/examples/vit5_imagenet/v6_hierarchical/cifar10_flat_p16.py b/examples/vit5_imagenet/v6_hierarchical/cifar10/flat_p16.py similarity index 59% rename from examples/vit5_imagenet/v6_hierarchical/cifar10_flat_p16.py rename to examples/vit5_imagenet/v6_hierarchical/cifar10/flat_p16.py index 7bacabe9..5073b50c 100644 --- a/examples/vit5_imagenet/v6_hierarchical/cifar10_flat_p16.py +++ b/examples/vit5_imagenet/v6_hierarchical/cifar10/flat_p16.py @@ -1,8 +1,6 @@ -"""CIFAR-10 patch-size ablation: flat (no merging), patch_size=16, grid=2×2, dim=384.""" +"""CIFAR-10 patch-size ablation: flat (no merging), patch_size=16, grid=4×4, dim=384.""" -from examples.vit5_imagenet.v6_hierarchical._cifar10_patch_ablation_base import ( - build_flat_config, -) +from examples.vit5_imagenet.v6_hierarchical.cifar10._base import build_flat_config def get_config(): diff --git a/examples/vit5_imagenet/v6_hierarchical/cifar10_flat_p4.py b/examples/vit5_imagenet/v6_hierarchical/cifar10/flat_p4.py similarity index 59% rename from examples/vit5_imagenet/v6_hierarchical/cifar10_flat_p4.py rename to examples/vit5_imagenet/v6_hierarchical/cifar10/flat_p4.py index 6394828c..5ae9f2fd 100644 --- a/examples/vit5_imagenet/v6_hierarchical/cifar10_flat_p4.py +++ b/examples/vit5_imagenet/v6_hierarchical/cifar10/flat_p4.py @@ -1,8 +1,6 @@ -"""CIFAR-10 patch-size ablation: flat (no merging), patch_size=4, grid=8×8, dim=384.""" +"""CIFAR-10 patch-size ablation: flat (no merging), patch_size=4, grid=16×16, dim=384.""" -from examples.vit5_imagenet.v6_hierarchical._cifar10_patch_ablation_base import ( - build_flat_config, -) +from examples.vit5_imagenet.v6_hierarchical.cifar10._base import build_flat_config def get_config(): diff --git a/examples/vit5_imagenet/v6_hierarchical/cifar10_flat_p8.py b/examples/vit5_imagenet/v6_hierarchical/cifar10/flat_p8.py similarity index 59% rename from examples/vit5_imagenet/v6_hierarchical/cifar10_flat_p8.py rename to examples/vit5_imagenet/v6_hierarchical/cifar10/flat_p8.py index d5a3c39b..b65b7307 100644 --- a/examples/vit5_imagenet/v6_hierarchical/cifar10_flat_p8.py +++ b/examples/vit5_imagenet/v6_hierarchical/cifar10/flat_p8.py @@ -1,8 +1,6 @@ -"""CIFAR-10 patch-size ablation: flat (no merging), patch_size=8, grid=4×4, dim=384.""" +"""CIFAR-10 patch-size ablation: flat (no merging), patch_size=8, grid=8×8, dim=384.""" -from examples.vit5_imagenet.v6_hierarchical._cifar10_patch_ablation_base import ( - build_flat_config, -) +from examples.vit5_imagenet.v6_hierarchical.cifar10._base import build_flat_config def get_config(): diff --git a/examples/vit5_imagenet/v6_hierarchical/cifar10_hier_p16.py b/examples/vit5_imagenet/v6_hierarchical/cifar10/hier_p16.py similarity index 54% rename from examples/vit5_imagenet/v6_hierarchical/cifar10_hier_p16.py rename to examples/vit5_imagenet/v6_hierarchical/cifar10/hier_p16.py index 86be7917..83466647 100644 --- a/examples/vit5_imagenet/v6_hierarchical/cifar10_hier_p16.py +++ b/examples/vit5_imagenet/v6_hierarchical/cifar10/hier_p16.py @@ -1,11 +1,9 @@ -"""CIFAR-10 patch-size ablation (64×64): hierarchical, patch_size=16. +"""CIFAR-10 patch-size ablation: hierarchical, patch_size=16. 2 stages: 4×4 → 2×2, dims=[384,768], depths=[6,2] ← Swin stages 3-4. """ -from examples.vit5_imagenet.v6_hierarchical._cifar10_patch_ablation_base import ( - build_hier_config, -) +from examples.vit5_imagenet.v6_hierarchical.cifar10._base import build_hier_config def get_config(): diff --git a/examples/vit5_imagenet/v6_hierarchical/cifar10_hier_p4.py b/examples/vit5_imagenet/v6_hierarchical/cifar10/hier_p4.py similarity index 53% rename from examples/vit5_imagenet/v6_hierarchical/cifar10_hier_p4.py rename to examples/vit5_imagenet/v6_hierarchical/cifar10/hier_p4.py index aa8cfc0d..12633454 100644 --- a/examples/vit5_imagenet/v6_hierarchical/cifar10_hier_p4.py +++ b/examples/vit5_imagenet/v6_hierarchical/cifar10/hier_p4.py @@ -1,11 +1,9 @@ """CIFAR-10 patch-size ablation: hierarchical, patch_size=4. -3 stages: 8×8 → 4×4 → 2×2, dims=[96,192,384], depths=[2,2,6]. +4 stages: 16×16 → 8×8 → 4×4 → 2×2, dims=[96,192,384,768], depths=[2,2,6,2]. """ -from examples.vit5_imagenet.v6_hierarchical._cifar10_patch_ablation_base import ( - build_hier_config, -) +from examples.vit5_imagenet.v6_hierarchical.cifar10._base import build_hier_config def get_config(): diff --git a/examples/vit5_imagenet/v6_hierarchical/cifar10_hier_p8.py b/examples/vit5_imagenet/v6_hierarchical/cifar10/hier_p8.py similarity index 55% rename from examples/vit5_imagenet/v6_hierarchical/cifar10_hier_p8.py rename to examples/vit5_imagenet/v6_hierarchical/cifar10/hier_p8.py index a21d88bc..7b5f5a29 100644 --- a/examples/vit5_imagenet/v6_hierarchical/cifar10_hier_p8.py +++ b/examples/vit5_imagenet/v6_hierarchical/cifar10/hier_p8.py @@ -1,11 +1,9 @@ """CIFAR-10 patch-size ablation: hierarchical, patch_size=8. -2 stages: 4×4 → 2×2, dims=[192,384], depths=[4,6]. +3 stages: 8×8 → 4×4 → 2×2, dims=[96,192,384], depths=[2,2,6]. """ -from examples.vit5_imagenet.v6_hierarchical._cifar10_patch_ablation_base import ( - build_hier_config, -) +from examples.vit5_imagenet.v6_hierarchical.cifar10._base import build_hier_config def get_config(): diff --git a/examples/vit5_imagenet/v6_hierarchical/cifar10_hyena_flat.py b/examples/vit5_imagenet/v6_hierarchical/cifar10/hyena_flat.py similarity index 100% rename from examples/vit5_imagenet/v6_hierarchical/cifar10_hyena_flat.py rename to examples/vit5_imagenet/v6_hierarchical/cifar10/hyena_flat.py diff --git a/examples/vit5_imagenet/v6_hierarchical/cifar10_hyena_hier.py b/examples/vit5_imagenet/v6_hierarchical/cifar10/hyena_hier.py similarity index 100% rename from examples/vit5_imagenet/v6_hierarchical/cifar10_hyena_hier.py rename to examples/vit5_imagenet/v6_hierarchical/cifar10/hyena_hier.py diff --git a/examples/vit5_imagenet/v6_hierarchical/hyena_hier_p4_film.py b/examples/vit5_imagenet/v6_hierarchical/hyena_hier_p4_film.py index 19a82767..284261ed 100644 --- a/examples/vit5_imagenet/v6_hierarchical/hyena_hier_p4_film.py +++ b/examples/vit5_imagenet/v6_hierarchical/hyena_hier_p4_film.py @@ -1,22 +1,36 @@ -"""v6_hierarchical — Hyena 4-stage with register-row FiLM, BD + learnable ω₀. +"""v6_hierarchical — FiLM Hyena, 4-stage Swin-T layout, BD + learnable ω₀. + +Modelled on ``vit5_hybrid/full_hyena_learnable_omega_blockdiag.py`` + +``vit5_hybrid/_film.py``. Same architecture as ``hyena_hier_p4_pure`` but with 4 register tokens -prepended as the first row of the 2D grid at every stage. Each block pools -its registers (RegisterPooling) and uses the result to FiLM-modulate the -SIREN kernel. No CLS token; GAP readout skips the register row. +prepended as the first row of every stage's 2D grid. Each block pools its +registers (RegisterPooling) and FiLM-modulates the SIREN kernel (KernelFiLMGenerator). + +FiLM settings (matching vit5_hybrid/_film.py best run): + init_type = "identity" (γ=1, β=0 — start from no-modulation) + film_after_pos_embed = True (modulates the positional-embedding sine) + num_film_layers = 3 (pos-embed sine + 2 hidden linears) + film_weight_decay = 5e-3 (dedicated optimizer group for FiLM weights) -Stages [96, 192, 384, 768] x depths [2, 2, 6, 2]; effective batch = 2048. +Batch : 2048 (8 GPUs × 16 per-GPU × 16 accum steps). """ from examples.vit5_imagenet.v6_hierarchical._base_config import ( build_hyena_hier_net, get_base_config, ) +from experiments.callbacks.mask_monitor import MaskMonitorCallback +from experiments.callbacks.omega_scale_monitor import OmegaScaleMonitorCallback from experiments.default_cfg import ExperimentConfig +from nvsubquadratic.lazy_config import LazyConfig def get_config() -> ExperimentConfig: - """Return the FiLM-conditioned hierarchical Hyena config.""" + """Return the FiLM-conditioned hierarchical Hyena ImageNet config.""" config = get_base_config() config.net = build_hyena_hier_net(layout="register_row") + config.callbacks.append(LazyConfig(MaskMonitorCallback)(log_every_n_steps=50)) + config.callbacks.append(LazyConfig(OmegaScaleMonitorCallback)(log_every_n_steps=50)) + config.wandb.job_group = "v6_hier_film" return config diff --git a/examples/vit5_imagenet/v6_hierarchical/hyena_hier_p4_pure.py b/examples/vit5_imagenet/v6_hierarchical/hyena_hier_p4_pure.py index 05a16150..f998f2e5 100644 --- a/examples/vit5_imagenet/v6_hierarchical/hyena_hier_p4_pure.py +++ b/examples/vit5_imagenet/v6_hierarchical/hyena_hier_p4_pure.py @@ -1,21 +1,28 @@ -"""v6_hierarchical — pure Hyena 4-stage (Swin-T layout) with BD + learnable ω₀. +"""v6_hierarchical — pure Hyena, 4-stage Swin-T layout, BD + learnable ω₀. -No registers, no FiLM. Stages [96, 192, 384, 768] x depths [2, 2, 6, 2] on -initial patch size 4 (grids 56 -> 28 -> 14 -> 7). PatchMerging between stages. -SIREN kernel: ``BlockDiagonalLearnableOmegaSIRENKernelND``. +Modelled on ``vit5_hybrid/full_hyena_learnable_omega_blockdiag.py``. -Effective batch = 2048 (8 GPUs x 16 per-GPU x 16 accum). +Layout : 4 stages × [2, 2, 6, 2] blocks, dims [96, 192, 384, 768]. +Kernel : BlockDiagonalLearnableOmegaSIRENKernelND + BlockAlignedGaussianModulationND. +Readout: GAP (no CLS, no registers). +Batch : 2048 (8 GPUs × 16 per-GPU × 16 accum steps). """ from examples.vit5_imagenet.v6_hierarchical._base_config import ( build_hyena_hier_net, get_base_config, ) +from experiments.callbacks.mask_monitor import MaskMonitorCallback +from experiments.callbacks.omega_scale_monitor import OmegaScaleMonitorCallback from experiments.default_cfg import ExperimentConfig +from nvsubquadratic.lazy_config import LazyConfig def get_config() -> ExperimentConfig: - """Return the pure (no-FiLM) hierarchical Hyena config.""" + """Return the pure (no-FiLM) hierarchical Hyena ImageNet config.""" config = get_base_config() config.net = build_hyena_hier_net(layout="pure") + config.callbacks.append(LazyConfig(MaskMonitorCallback)(log_every_n_steps=50)) + config.callbacks.append(LazyConfig(OmegaScaleMonitorCallback)(log_every_n_steps=50)) + config.wandb.job_group = "v6_hier_pure" return config From bda14549317b7724f833f9f78a6926d2ace8a1c4 Mon Sep 17 00:00:00 2001 From: David Wessels Date: Tue, 26 May 2026 16:59:45 +0200 Subject: [PATCH 72/72] fix(cifar10): add Lightning datamodule required by v6_hierarchical/cifar10 configs The v6_hierarchical CIFAR-10 example configs added in this PR import `experiments.datamodules.cifar10.CIFAR10DataModule` but the file was local-only. Add it to the PR with docstrings on all public methods so ruff D102/D107 pass under the new docstring CI. Also apply mdformat to docs-tracker.md to satisfy the pre-commit hook (column-width normalization). --- docs-tracker.md | 72 ++++++------- experiments/datamodules/cifar10.py | 160 +++++++++++++++++++++++++++++ 2 files changed, 196 insertions(+), 36 deletions(-) create mode 100644 experiments/datamodules/cifar10.py diff --git a/docs-tracker.md b/docs-tracker.md index cb70123e..f122f172 100644 --- a/docs-tracker.md +++ b/docs-tracker.md @@ -40,46 +40,46 @@ Work bottom-up: primitive ops → modules → networks → experiments. ### `nvsubquadratic/modules/` — Building blocks -| File | Status | Notes | -| ---------------------------------- | ------ | ------------------------------------------------------------------------------------------- | -| `kernels_nd.py` | \[x\] | Learned kernel parametrisation — RFF + SIREN, FiLM-conditioned variants | -| `hyena_nd.py` | \[x\] | Hyena operator (ND) — two-gate sandwich, AllToAll CP, BC-aware convolution | -| `ckconv_nd.py` | \[x\] | CKConv (ND) — implicit kernel `k_θ(p) = MLP_θ(pos_enc(p))`, FFT domain, BC modes | -| `ckconv_multihead_nd.py` | \[x\] | Multi-head CKConv — H heads, dense d×d kernel per head, low-rank U·V factorisation | -| `mamba_nd.py` | \[x\] | Mamba SSM (ND) — selective SSM, ZOH discretisation, raster-scan ND, bidirectional mode | -| `attention.py` | \[x\] | Scaled dot-product attention — multi-head, RoPE, ND spatial, O(L²) FLOP formula | -| `vit5_attention.py` | \[x\] | ViT5 attention — register-aware 2D RoPE, QK-norm, CUDA-graph-safe buffers | -| `vit5_hyena_adapter.py` | \[x\] | Hyena adapter for ViT5 — drop-in for vit5_attention, register-token + hierarchy support | -| `sequence_mixer.py` | \[x\] | Operator-agnostic dispatch layer (Hyena / Attention / CKConv / Mamba) | -| `condition_mixer.py` | \[x\] | Cross-attention conditioning mixer — both global (B,C) and spatial (B,\*,C) signals | -| `residual_block.py` | \[x\] | Residual block — pre-norm + mixer + MLP, optional FiLM/AdaLN-Zero conditioning | -| `vit5_residual_block.py` | \[x\] | ViT5 residual block — LayerScale, register-token conditioning, no condition-mixer branch | -| `patchify.py` | \[x\] | Patch embedding — strided conv, 1D/2D/3D, channels-last layout | -| `position_encoding.py` | \[x\] | Axis-factorised learned PE — ND broadcast-expand, float32 output caveat | -| `masks_nd.py` | \[x\] | Exponential + Gaussian receptive-field windows; mask convention 1=included, 0=excluded | -| `mlp.py` | \[x\] | Two-layer MLP — GELU/SwiGLU/GLU variants, expansion-ratio math, QuACK backend noted | -| `film.py` | \[x\] | FiLM conditioning — γ(c)⊙x + β(c), SIREN-based kernel generator | -| `grn.py` | \[x\] | GRN — per-channel L2 norm, inter-channel competition, ConvNeXt V2 reference | -| `layer_scale.py` | \[x\] | LayerScale — per-channel λ⊙F(x), init_values guidance, \_no_weight_decay tag | -| `rms_norm.py` | \[x\] | RMSNorm + PerHeadRMSNorm — QuACK/PyTorch backends, math formula, QK-norm usage | -| `rms_norm_channel_first.py` | \[x\] | Channel-first RMSNorm — normalises dim=1, `channels_first` sentinel, Hyena usage | -| `drop_path.py` | \[x\] | Stochastic depth — functional + Module, inverted-dropout scaling, causal vs training | -| `causal_conv1d.py` | \[x\] | CausalConv1D — left-only pad formula, symmetric mode, Mamba usage context | -| `subq_ops_causal_conv1d.py` | \[x\] | New (1D PR): `nn.Conv1d`-compatible depthwise wrapper around `subq_ops.causal_conv1d` | -| `schedulers.py` | \[x\] | ResumableSequentialLR — PyTorch ≤2.10 bug fix, load_state_dict LR propagation | -| `distributed_depthwise_conv_nd.py` | \[x\] | CP-aware 1D/2D/3D depthwise conv — group weight sharing, channel slicing, causal padding | -| `patch_merging.py` | \[x\] | New (hierarchical PR): Swin-style 2×2 patch merging with optional register-row passthrough | +| File | Status | Notes | +| ---------------------------------- | ------ | ------------------------------------------------------------------------------------------ | +| `kernels_nd.py` | \[x\] | Learned kernel parametrisation — RFF + SIREN, FiLM-conditioned variants | +| `hyena_nd.py` | \[x\] | Hyena operator (ND) — two-gate sandwich, AllToAll CP, BC-aware convolution | +| `ckconv_nd.py` | \[x\] | CKConv (ND) — implicit kernel `k_θ(p) = MLP_θ(pos_enc(p))`, FFT domain, BC modes | +| `ckconv_multihead_nd.py` | \[x\] | Multi-head CKConv — H heads, dense d×d kernel per head, low-rank U·V factorisation | +| `mamba_nd.py` | \[x\] | Mamba SSM (ND) — selective SSM, ZOH discretisation, raster-scan ND, bidirectional mode | +| `attention.py` | \[x\] | Scaled dot-product attention — multi-head, RoPE, ND spatial, O(L²) FLOP formula | +| `vit5_attention.py` | \[x\] | ViT5 attention — register-aware 2D RoPE, QK-norm, CUDA-graph-safe buffers | +| `vit5_hyena_adapter.py` | \[x\] | Hyena adapter for ViT5 — drop-in for vit5_attention, register-token + hierarchy support | +| `sequence_mixer.py` | \[x\] | Operator-agnostic dispatch layer (Hyena / Attention / CKConv / Mamba) | +| `condition_mixer.py` | \[x\] | Cross-attention conditioning mixer — both global (B,C) and spatial (B,\*,C) signals | +| `residual_block.py` | \[x\] | Residual block — pre-norm + mixer + MLP, optional FiLM/AdaLN-Zero conditioning | +| `vit5_residual_block.py` | \[x\] | ViT5 residual block — LayerScale, register-token conditioning, no condition-mixer branch | +| `patchify.py` | \[x\] | Patch embedding — strided conv, 1D/2D/3D, channels-last layout | +| `position_encoding.py` | \[x\] | Axis-factorised learned PE — ND broadcast-expand, float32 output caveat | +| `masks_nd.py` | \[x\] | Exponential + Gaussian receptive-field windows; mask convention 1=included, 0=excluded | +| `mlp.py` | \[x\] | Two-layer MLP — GELU/SwiGLU/GLU variants, expansion-ratio math, QuACK backend noted | +| `film.py` | \[x\] | FiLM conditioning — γ(c)⊙x + β(c), SIREN-based kernel generator | +| `grn.py` | \[x\] | GRN — per-channel L2 norm, inter-channel competition, ConvNeXt V2 reference | +| `layer_scale.py` | \[x\] | LayerScale — per-channel λ⊙F(x), init_values guidance, \_no_weight_decay tag | +| `rms_norm.py` | \[x\] | RMSNorm + PerHeadRMSNorm — QuACK/PyTorch backends, math formula, QK-norm usage | +| `rms_norm_channel_first.py` | \[x\] | Channel-first RMSNorm — normalises dim=1, `channels_first` sentinel, Hyena usage | +| `drop_path.py` | \[x\] | Stochastic depth — functional + Module, inverted-dropout scaling, causal vs training | +| `causal_conv1d.py` | \[x\] | CausalConv1D — left-only pad formula, symmetric mode, Mamba usage context | +| `subq_ops_causal_conv1d.py` | \[x\] | New (1D PR): `nn.Conv1d`-compatible depthwise wrapper around `subq_ops.causal_conv1d` | +| `schedulers.py` | \[x\] | ResumableSequentialLR — PyTorch ≤2.10 bug fix, load_state_dict LR propagation | +| `distributed_depthwise_conv_nd.py` | \[x\] | CP-aware 1D/2D/3D depthwise conv — group weight sharing, channel slicing, causal padding | +| `patch_merging.py` | \[x\] | New (hierarchical PR): Swin-style 2×2 patch merging with optional register-row passthrough | ### `nvsubquadratic/networks/` — Full architectures -| File | Status | Notes | -| ------------------------------------- | ------ | ----------------------------------------------------------------------------------------------------------------------------------- | -| `general_purpose_resnet.py` | \[x\] | ResidualNetwork — LazyConfig blocks, conditioning, readout crop, gradient checkpointing | -| `classification_resnet.py` | \[x\] | ClassificationResNet — GAP readout, resolution-agnostic, inherits ResidualNetwork | -| `vit5_classification.py` | \[x\] | ViT5 classification — token layout, hybrid blocks, CLS/GAP/register_concat readouts, FLOP count | +| File | Status | Notes | +| ------------------------------------- | ------ | ------------------------------------------------------------------------------------------------------------------------- | +| `general_purpose_resnet.py` | \[x\] | ResidualNetwork — LazyConfig blocks, conditioning, readout crop, gradient checkpointing | +| `classification_resnet.py` | \[x\] | ClassificationResNet — GAP readout, resolution-agnostic, inherits ResidualNetwork | +| `vit5_classification.py` | \[x\] | ViT5 classification — token layout, hybrid blocks, CLS/GAP/register_concat readouts, FLOP count | | `vit5_hierarchical_classification.py` | \[x\] | New (hierarchical PR): Swin-style 4-stage hierarchical ViT-5 classifier with GAP readout and optional register-row layout | -| `huggingface_diffusers.py` | \[ \] | HF diffusers integration | -| `jit.py` / `jit_utils.py` | \[ \] | TorchScript utilities | +| `huggingface_diffusers.py` | \[ \] | HF diffusers integration | +| `jit.py` / `jit_utils.py` | \[ \] | TorchScript utilities | ### `nvsubquadratic/parallel/` — Distributed primitives diff --git a/experiments/datamodules/cifar10.py b/experiments/datamodules/cifar10.py new file mode 100644 index 00000000..f95b1f68 --- /dev/null +++ b/experiments/datamodules/cifar10.py @@ -0,0 +1,160 @@ +"""CIFAR-10 Lightning DataModule. + +Returns batches with the three-key format expected by ClassificationWrapper: + {"input": [B, H, W, 3] float32 channels-last, "label": [B] long, "condition": None} + +Standard augmentation: RandomCrop(32, padding=4) + RandomHorizontalFlip + Normalize. +Optional Mixup/CutMix via timm (off by default for quick debug runs). +""" + +from typing import Optional + +import pytorch_lightning as pl +import torch +from einops import rearrange +from timm.data import Mixup +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + + +# CIFAR-10 channel statistics +CIFAR10_MEAN = (0.4914, 0.4822, 0.4465) +CIFAR10_STD = (0.2470, 0.2435, 0.2616) + + +class _ChannelsLastCollate: + """Collate CIFAR-10 (C, H, W) images into (B, H, W, C) channels-last tensors. + + Also injects the ``condition`` key (always None) expected by + ClassificationWrapper. + """ + + def __init__(self, mixup_fn: Optional[Mixup] = None): + self.mixup_fn = mixup_fn + + def __call__(self, batch): + images, labels = zip(*batch) + x = torch.stack(images) # [B, C, H, W] + y = torch.tensor(labels, dtype=torch.long) + + if self.mixup_fn is not None: + x, y = self.mixup_fn(x, y) + + x = rearrange(x, "b c h w -> b h w c") # channels-last + return {"input": x, "label": y, "condition": None} + + +class CIFAR10DataModule(pl.LightningDataModule): + """CIFAR-10 datamodule. + + Args: + data_dir: Path where torchvision will download / find CIFAR-10. + batch_size: Per-GPU batch size. + num_workers: DataLoader workers. + pin_memory: Whether to pin memory for faster GPU transfer. + image_size: Resize target. CIFAR-10 is natively 32×32; resize only + if you need a different resolution. + mixup: Mixup alpha (0 = disabled). + cutmix: CutMix alpha (0 = disabled). + num_classes: Number of output classes (10). + """ + + def __init__( + self, + data_dir: str = "./data", + batch_size: int = 256, + num_workers: int = 4, + pin_memory: bool = True, + image_size: int = 32, + mixup: float = 0.0, + cutmix: float = 0.0, + num_classes: int = 10, + ): + """Initialize the CIFAR-10 datamodule (see class docstring for args).""" + super().__init__() + self.data_dir = data_dir + self.batch_size = batch_size + self.num_workers = num_workers + self.pin_memory = pin_memory + self.image_size = image_size + self.num_classes = num_classes + self.output_channels = num_classes + + # timm Mixup/CutMix (only active during training) + use_mix = mixup > 0 or cutmix > 0 + self._mixup_fn = ( + Mixup( + mixup_alpha=mixup, + cutmix_alpha=cutmix, + prob=1.0, + switch_prob=0.5, + mode="batch", + num_classes=num_classes, + ) + if use_mix + else None + ) + + train_tf = [ + transforms.RandomCrop(image_size, padding=4), + transforms.RandomHorizontalFlip(), + ] + if image_size != 32: + train_tf.insert(0, transforms.Resize(image_size)) + train_tf += [ + transforms.ToTensor(), + transforms.Normalize(CIFAR10_MEAN, CIFAR10_STD), + ] + self._train_transform = transforms.Compose(train_tf) + + val_tf = [] + if image_size != 32: + val_tf.append(transforms.Resize(image_size)) + val_tf += [ + transforms.ToTensor(), + transforms.Normalize(CIFAR10_MEAN, CIFAR10_STD), + ] + self._val_transform = transforms.Compose(val_tf) + + def prepare_data(self): + """Download train and test splits of CIFAR-10 if not already present.""" + datasets.CIFAR10(self.data_dir, train=True, download=True) + datasets.CIFAR10(self.data_dir, train=False, download=True) + + def setup(self, stage=None): + """Instantiate the train/val :class:`torchvision.datasets.CIFAR10` datasets.""" + self.train_ds = datasets.CIFAR10(self.data_dir, train=True, transform=self._train_transform) + self.val_ds = datasets.CIFAR10(self.data_dir, train=False, transform=self._val_transform) + + def train_dataloader(self): + """Return the training DataLoader with shuffling, drop_last, and optional Mixup/CutMix.""" + return DataLoader( + self.train_ds, + batch_size=self.batch_size, + shuffle=True, + num_workers=self.num_workers, + pin_memory=self.pin_memory, + drop_last=True, + collate_fn=_ChannelsLastCollate(self._mixup_fn), + persistent_workers=self.num_workers > 0, + ) + + def val_dataloader(self): + """Return the validation DataLoader (no shuffling, no Mixup, larger batch).""" + return DataLoader( + self.val_ds, + batch_size=self.batch_size * 2, + shuffle=False, + num_workers=self.num_workers, + pin_memory=self.pin_memory, + drop_last=False, + collate_fn=_ChannelsLastCollate(mixup_fn=None), + persistent_workers=self.num_workers > 0, + ) + + def test_dataloader(self): + """Return the test DataLoader — aliased to :meth:`val_dataloader` (see comment below).""" + # CIFAR-10 has no separate held-out test split beyond the 10k test set, + # which is already used for validation. Reuse val_dataloader so that + # trainer.test() reports the same test-set accuracy cleanly. + return self.val_dataloader()