asr_layers.py

"""Module-surgery helpers for ASRModel: MPS-safe embeddings and encoder unfreezing.

Two self-contained groups of `nn.Module` utilities that ASRModel applies to its
loaded sub-models:

- `mps_unsafe_parameters` / `ChunkedEmbedding` / `chunk_oversized_embeddings`
  work around PyTorch's 32-bit indexing limit in its Metal kernels.
- `find_encoder_layer_stack` / `unfreeze_encoder_top_layers` locate an
  encoder's transformer-block list and switch its top blocks back to trainable.
"""

import math

import torch
import torch.nn as nn
from torch.nn import functional

# PyTorch's Metal kernels index tensor storage with 32-bit offsets, so a gather
# into a tensor holding more than INT32_MAX elements wraps around and silently
# returns the wrong rows -- no error, no warning, just corrupt values.
MPS_MAX_TENSOR_ELEMENTS = 2**31 - 1


def mps_unsafe_parameters(model: nn.Module) -> list[tuple[str, int]]:
    """Parameters too large for MPS to index correctly, as (name, numel) pairs.

    Gemma 4 E2B trips this on `embed_tokens_per_layer`: its per-layer embedding
    table is 262144 x (35*256) = 2,348,810,240 elements, 201M past INT32_MAX.
    On MPS the lookup returns garbage, which Gemma then adds into all 35 layers
    at every position. The failure is entirely silent -- the model loads, runs,
    and emits fluent, confident, wrong transcripts. Measured on this stack:
    teacher-forced loss 0.14 on CPU versus 3.80 on MPS from identical weights.

    Returns an empty list for models that are safe, so callers can treat a
    non-empty result as "this model needs `chunk_oversized_embeddings`".
    """
    return [
        (name, param.numel())
        for name, param in model.named_parameters()
        if param.numel() > MPS_MAX_TENSOR_ELEMENTS
    ]


class ChunkedEmbedding(nn.Module):
    """Drop-in `nn.Embedding` that splits its table to stay within 32-bit indexing.

    The overflow lives in the kernel's arithmetic over the weight's *storage*,
    so slicing is not enough -- views share that storage and stay broken. Each
    chunk therefore has to be its own contiguous allocation. Gathering per
    chunk and concatenating on the feature axis reproduces the full gather
    exactly: verified bit-identical to CPU on MPS, where the single-tensor
    lookup was off by up to 3.16.

    Wraps rather than replaces the semantics of `Gemma4TextScaledWordEmbedding`,
    whose forward multiplies the lookup by `embed_scale` -- dropping that would
    shrink every per-layer embedding by sqrt(hidden_size_per_layer_input).
    """

    def __init__(self, embedding: nn.Embedding) -> None:
        """Split `embedding`'s table into MPS-safe column chunks, keeping its metadata."""
        super().__init__()
        self.num_embeddings = embedding.num_embeddings
        self.embedding_dim = embedding.embedding_dim
        self.padding_idx = embedding.padding_idx
        # Gemma 4's scaled embedding carries this; a plain nn.Embedding does not.
        self.embed_scale = getattr(embedding, "embed_scale", None)

        num_chunks = math.ceil(embedding.weight.numel() / MPS_MAX_TENSOR_ELEMENTS)
        device = embedding.weight.device
        # Chunk from a CPU copy: carving contiguous pieces out of an
        # already-oversized MPS tensor would run the same overflowing kernel
        # this class exists to avoid.
        source = embedding.weight.detach().cpu()
        self.chunks = nn.ParameterList(
            [
                nn.Parameter(
                    chunk.contiguous().to(device), requires_grad=embedding.weight.requires_grad
                )
                for chunk in torch.chunk(source, num_chunks, dim=1)
            ]
        )
        del source

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        """Look up `input_ids` in every chunk and concatenate along the feature dim."""
        out = torch.cat([functional.embedding(input_ids, chunk) for chunk in self.chunks], dim=-1)
        if self.embed_scale is not None:
            out = out * self.embed_scale.to(out.dtype)
        return out

    @property
    def weight(self) -> torch.Tensor:
        """Reassembled table, for callers that introspect `.weight`.

        Materializes the full tensor, so this is a debugging/serialization
        convenience rather than something to call on a hot path.
        """
        return torch.cat(list(self.chunks), dim=1)

    def extra_repr(self) -> str:
        """Describe the table like `nn.Embedding` does, plus the chunk count."""
        return (
            f"{self.num_embeddings}, {self.embedding_dim}, "
            f"chunks={len(self.chunks)} (MPS int32-indexing workaround)"
        )


def find_encoder_layer_stack(encoder: nn.Module) -> tuple[str, nn.ModuleList] | None:
    """Locate the encoder's ordered transformer-block list.

    Returns `(attribute_path, ModuleList)` or None. Encoders disagree on the
    attribute name -- Granite Speech uses `layers`, Whisper `encoder.layers`,
    others `blocks` or `encoder_layers` -- so probe the known names rather
    than hardcoding one and silently unfreezing nothing on a new encoder.
    """
    candidates = ("layers", "blocks", "encoder_layers", "encoder.layers", "encoder.blocks")
    for path in candidates:
        try:
            obj = encoder.get_submodule(path)
        except AttributeError:
            continue
        if isinstance(obj, nn.ModuleList) and len(obj) > 0:
            return path, obj
    return None


def unfreeze_encoder_top_layers(
    encoder: nn.Module, top_n: int, include_post_projections: bool = True
) -> list[str]:
    """Unfreeze the top `top_n` encoder blocks, optionally with the output projections.

    Assumes the caller has already frozen the whole encoder. Returns the sorted
    names of every parameter switched back on, so callers (and tests) can
    assert on exactly what moved rather than trusting a count.

    The post-stack projections (Granite's `out` / `out_mid`) are included by
    default: they sit between the top block and the projector, so leaving them
    frozen forces the newly-trainable blocks to adapt through a fixed output
    map. That is a reasonable prior, but on this stack it is not what the
    gradient says. Probing the step-2000 granite_qwen_lora checkpoint on real
    audio, `out`/`out_mid` carry 33.57M params -- 23% of a top-4 budget -- for
    an RMS-per-param gradient of 1.14e-05 against 1.4-1.7e-04 in the top
    blocks, i.e. **0.27% of the selection's squared gradient**. Dropping them
    moved the aggregate norm from 1.5840 to 1.5827. Pass False to reclaim the
    parameters; the default stays True so existing recipes are unchanged.

    The pre-stack `input_linear` stays frozen either way -- it is the feature
    front-end, the part most expensive to damage and least specialised to the
    encoder's own CTC head.

    Raises if the layer stack cannot be found or `top_n` exceeds its depth:
    silently unfreezing nothing would produce a run that looks like a
    partial-unfreeze experiment and is actually the frozen baseline.
    """
    if top_n <= 0:
        return []
    found = find_encoder_layer_stack(encoder)
    if found is None:
        msg = (
            f"Could not locate a transformer-block ModuleList on "
            f"{type(encoder).__name__}; cannot unfreeze its top {top_n} layers. "
            f"Children: {[n for n, _ in encoder.named_children()]}"
        )
        raise ValueError(msg)
    path, stack = found
    depth = len(stack)
    if top_n > depth:
        msg = f"encoder_trainable_top_layers={top_n} exceeds the encoder's {depth} blocks"
        raise ValueError(msg)

    unfrozen: list[str] = []
    for idx in range(depth - top_n, depth):
        for name, param in stack[idx].named_parameters():
            param.requires_grad_(True)
            unfrozen.append(f"{path}.{idx}.{name}")
    # Post-stack projections, identified positionally: direct children that are
    # not the stack itself and not the pre-stack input projection.
    for name, child in encoder.named_children() if include_post_projections else ():
        if name == path.split(".")[0] or name.startswith("input"):
            continue
        for pname, param in child.named_parameters():
            param.requires_grad_(True)
            unfrozen.append(f"{name}.{pname}")
    return sorted(unfrozen)


def chunk_oversized_embeddings(root: nn.Module) -> list[str]:
    """Replace every embedding MPS cannot index with a `ChunkedEmbedding`.

    Returns the qualified names that were replaced, empty when nothing needed
    it. Idempotent: `ChunkedEmbedding` is not an `nn.Embedding`, so a second
    pass finds nothing.
    """
    replaced: list[str] = []
    for module_name, module in list(root.named_modules()):
        for child_name, child in list(module.named_children()):
            if isinstance(child, nn.Embedding) and child.weight.numel() > MPS_MAX_TENSOR_ELEMENTS:
                setattr(module, child_name, ChunkedEmbedding(child))
                replaced.append(f"{module_name}.{child_name}" if module_name else child_name)
    return replaced