runtime/quantlib/wrap_model.py

"""In-place model conversion utilities for Sherry and RTN QAT."""

from __future__ import annotations

import fnmatch
import re
from collections.abc import Sequence

import torch.nn as nn

from .sherry_quant import SherryLinear
from .rtn_quant import RTNEmbedding, RTNLinear


DEFAULT_TARGET_PATTERNS = (
    "q_proj",
    "k_proj",
    "v_proj",
    "o_proj",
    "gate_proj",
    "up_proj",
    "down_proj",
)
DEFAULT_EXCLUDE_PATTERNS = (
    "audio",
    "encoder",
    "embed",
    "lm_head",
    "projector",
)


def _matches(path: str, leaf: str, patterns: Sequence[str]) -> bool:
    """Match exact leaves, path substrings, shell globs, or explicit regexes."""

    for pattern in patterns:
        if pattern.startswith("re:"):
            if re.search(pattern[3:], path):
                return True
        elif any(character in pattern for character in "*?["):
            if fnmatch.fnmatchcase(path, pattern) or fnmatch.fnmatchcase(leaf, pattern):
                return True
        elif pattern == leaf or pattern in path:
            return True
    return False


def wrap_sherry(
    model: nn.Module,
    target_patterns: Sequence[str] | None = None,
    exclude_patterns: Sequence[str] | None = None,
    *,
    group_size: int = 128,
) -> list[str]:
    """Replace matching decoder ``nn.Linear`` children in place.

    The default targets are Qwen3's seven attention/MLP projections and require
    a ``layers.<integer>`` decoder path, while audio towers, encoders, embeddings,
    projectors, and ``lm_head`` are excluded.  Plain strings match path substrings
    or exact leaf names; glob patterns and ``re:...`` regexes are also accepted.

    Differences from AngelSlim Sherry: this avoids a version-pinned fork of the
    Transformers modeling file, preserves weight/bias key paths, copies weights
    into fp32 masters, and returns an auditable replacement list.
    """

    use_default_targets = target_patterns is None
    targets = tuple(DEFAULT_TARGET_PATTERNS if target_patterns is None else target_patterns)
    excludes = tuple(DEFAULT_EXCLUDE_PATTERNS if exclude_patterns is None else exclude_patterns)
    replaced: list[str] = []

    # Materialize before mutation: named_modules otherwise observes replacements.
    for parent_path, parent in list(model.named_modules()):
        for leaf, child in list(parent.named_children()):
            if not isinstance(child, nn.Linear):
                continue
            path = f"{parent_path}.{leaf}" if parent_path else leaf
            if use_default_targets and re.search(r"(?:^|\.)layers\.\d+\.", path) is None:
                continue
            if excludes and _matches(path, leaf, excludes):
                continue
            if not _matches(path, leaf, targets):
                continue
            if isinstance(child, RTNLinear):
                raise ValueError(
                    f"quantization wrapper conflict at {path!r}: already wrapped by RTN"
                )
            if isinstance(child, SherryLinear):
                continue
            converted = SherryLinear.from_linear(child, group_size=group_size)
            setattr(parent, leaf, converted)
            replaced.append(path)
    return replaced


def wrap_rtn(
    model: nn.Module,
    target_patterns: Sequence[str],
    exclude_patterns: Sequence[str] | None = None,
    *,
    bits: int = 8,
    granularity: str = "per_channel",
    group_size: int = 128,
) -> list[str]:
    """Replace matching linear/embedding children with fp32-master RTN layers.

    Patterns have the same exact-leaf, substring, glob, and ``re:`` semantics as
    :func:`wrap_sherry`.  Calling this over a path already wrapped by Sherry is
    also an explicit conflict.
    """

    targets = tuple(target_patterns)
    excludes = tuple(() if exclude_patterns is None else exclude_patterns)
    replaced: list[str] = []

    for parent_path, parent in list(model.named_modules()):
        for leaf, child in list(parent.named_children()):
            if not isinstance(child, (nn.Linear, nn.Embedding)):
                continue
            path = f"{parent_path}.{leaf}" if parent_path else leaf
            if excludes and _matches(path, leaf, excludes):
                continue
            if not _matches(path, leaf, targets):
                continue
            if isinstance(child, SherryLinear):
                raise ValueError(
                    f"quantization wrapper conflict at {path!r}: already wrapped by Sherry"
                )
            if isinstance(child, (RTNLinear, RTNEmbedding)):
                continue
            if isinstance(child, nn.Embedding):
                converted = RTNEmbedding.from_embedding(
                    child,
                    bits=bits,
                    granularity=granularity,
                    group_size=group_size,
                )
            else:
                converted = RTNLinear.from_linear(
                    child,
                    bits=bits,
                    granularity=granularity,
                    group_size=group_size,
                )
            setattr(parent, leaf, converted)
            replaced.append(path)
    return replaced


def freeze_non_quantized(model: nn.Module) -> list[str]:
    """Freeze parameters outside Sherry layers and return trainable key names.

    AngelSlim constructs a fully quantized modeling fork and leaves global
    freezing to its trainer.  The wrapper approach needs an explicit helper;
    parameters owned by each Sherry module (fp32 master and optional bias) remain
    trainable while embeddings, norms, heads, and skipped towers are frozen.
    """

    for parameter in model.parameters():
        parameter.requires_grad_(False)
    trainable: list[str] = []
    for module_name, module in model.named_modules():
        if not isinstance(module, SherryLinear):
            continue
        for parameter_name, parameter in module.named_parameters(recurse=False):
            parameter.requires_grad_(True)
            path = f"{module_name}.{parameter_name}" if module_name else parameter_name
            trainable.append(path)
    return trainable