amachine.am_transformers.am_model_introspection

  1"""
  2Architecture-agnostic module resolution for HuggingFace transformer models.
  3
  4This module is purely about *finding* things on a model — decoder layers,
  5the final pre-lm_head norm, and whether input/output embeddings are tied.
  6It does not mutate the model in any way (no requires_grad changes here).
  7
  8Designed to work across common architecture families (Llama, GPT-2, GPT-NeoX,
  9T5-style, custom MoE/shared-MLP variants, etc.) as well as common wrappers
 10(DDP, torch.compile, FSDP) by walking through wrapper attributes.
 11"""
 12
 13from __future__ import annotations
 14
 15import warnings
 16from dataclasses import dataclass
 17
 18import torch
 19
 20
 21# ---------------------------------------------------------------------------
 22# Known attribute paths for common architecture families.
 23# These are checked in order; the first match wins. Add new architectures
 24# here as needed rather than special-casing them elsewhere.
 25# ---------------------------------------------------------------------------
 26
 27DECODER_LAYER_PATHS: list[tuple[str, ...]] = [
 28    ("model", "layers"),              # Llama, Mistral, Gemma, Qwen, etc.
 29    ("transformer", "h"),              # GPT-2 family
 30    ("gpt_neox", "layers"),            # GPT-NeoX / Pythia
 31    ("model", "decoder", "layers"),    # OPT, some encoder-decoder decoders
 32    ("layers",),                       # bare / already-unwrapped models
 33]
 34
 35FINAL_NORM_PATHS: list[tuple[str, ...]] = [
 36    ("model", "norm"),                       # Llama, Mistral, Gemma, Qwen, etc.
 37    ("transformer", "ln_f"),                  # GPT-2 family
 38    ("gpt_neox", "final_layer_norm"),         # GPT-NeoX / Pythia
 39    ("model", "decoder", "final_layer_norm"), # OPT
 40]
 41
 42# Attributes that indicate "the real model is nested inside this wrapper".
 43WRAPPER_ATTRS: tuple[str, ...] = (
 44    "module",                 # DDP / DataParallel
 45    "_orig_mod",               # torch.compile
 46    "_fsdp_wrapped_module",    # FSDP fallback (best-effort; see caveats below)
 47)
 48
 49
 50class ModuleResolutionError(RuntimeError):
 51    """Raised when a required module (e.g. decoder layers) cannot be located."""
 52
 53
 54class TiedEmbeddingConflictError(RuntimeError):
 55    """Raised when tied input/output embeddings conflict with the caller's
 56    stated intent to train them independently."""
 57
 58
 59@dataclass(frozen=True)
 60class ResolvedModel:
 61    """Bundle of everything resolved from a model, so callers only pay the
 62    wrapper-walking / path-guessing cost once."""
 63
 64    root: torch.nn.Module                    # the model, after unwrapping to the level where layers were found
 65    decoder_layers: torch.nn.ModuleList
 66    final_norm: torch.nn.Module | None
 67    input_embeddings: torch.nn.Module | None
 68    output_embeddings: torch.nn.Module | None
 69    embeddings_tied: bool
 70    tie_word_embeddings_config: bool | None   # raw config value, for reference/debugging
 71
 72
 73# ---------------------------------------------------------------------------
 74# Internals
 75# ---------------------------------------------------------------------------
 76
 77def _resolve_path(module: torch.nn.Module, path: tuple[str, ...]) -> object | None:
 78    obj: object = module
 79    for name in path:
 80        obj = getattr(obj, name, None)
 81        if obj is None:
 82            return None
 83    return obj
 84
 85
 86def _walk_wrappers(model: torch.nn.Module):
 87    """Yield model, then anything reachable through known wrapper attrs,
 88    breadth-first-ish, without revisiting the same object twice."""
 89    pending, visited = [model], set()
 90    while pending:
 91        current = pending.pop()
 92        if id(current) in visited:
 93            continue
 94        visited.add(id(current))
 95        yield current
 96        for attr in WRAPPER_ATTRS:
 97            wrapped = getattr(current, attr, None)
 98            if isinstance(wrapped, torch.nn.Module) and wrapped is not current:
 99                pending.append(wrapped)
100
101
102# ---------------------------------------------------------------------------
103# Public resolution functions
104# ---------------------------------------------------------------------------
105
106def find_decoder_layers(model: torch.nn.Module) -> torch.nn.ModuleList:
107    """Find the decoder layer stack, descending through common wrappers if needed.
108
109    Raises ModuleResolutionError if no known path matches on this model or any
110    of its unwrapped inner modules. If your architecture isn't found, add its
111    attribute path to DECODER_LAYER_PATHS above.
112    """
113    for current in _walk_wrappers(model):
114        for path in DECODER_LAYER_PATHS:
115            obj = _resolve_path(current, path)
116            if obj is not None and hasattr(obj, "__len__"):
117                return obj  # type: ignore[return-value]
118    raise ModuleResolutionError(
119        "Could not locate decoder layers. Checked paths: "
120        f"{DECODER_LAYER_PATHS}. Add this architecture's attribute path to "
121        "DECODER_LAYER_PATHS if it uses a different naming convention."
122    )
123
124
125def find_final_norm(model: torch.nn.Module) -> torch.nn.Module | None:
126    """Best-effort lookup of the final pre-lm_head norm.
127
128    Returns None (rather than raising) if not found — some architectures may
129    genuinely lack one, or use an attribute path not yet in FINAL_NORM_PATHS.
130    Callers should decide whether a missing final norm is fatal for their use case.
131    """
132    for current in _walk_wrappers(model):
133        for path in FINAL_NORM_PATHS:
134            obj = _resolve_path(current, path)
135            if isinstance(obj, torch.nn.Module):
136                return obj
137    return None
138
139
140def detect_tied_embeddings(model: torch.nn.Module, strict: bool = False) -> bool:
141    """Detect whether input embeddings and lm_head share the same weight tensor.
142
143    Uses `is` identity on the .weight attribute as the primary signal, since
144    this is exact (not heuristic) for a plainly-constructed, unwrapped model.
145    Cross-checks against config.tie_word_embeddings when available.
146
147    IMPORTANT: call this as early as possible — right after model construction
148    /from_pretrained(), before FSDP, DeepSpeed ZeRO-3, quantization, or PEFT
149    wrapping. Those can partition, flatten, or replace weight tensors such that
150    identity no longer reflects the true tying relationship.
151
152    Args:
153        strict: if True, raise TiedEmbeddingConflictError on a config/identity
154            disagreement instead of warning. Appropriate when you control model
155            construction directly (e.g. generating configs from scratch) and
156            expect config.tie_word_embeddings to be authoritative.
157    """
158    input_emb = model.get_input_embeddings()
159    output_emb = model.get_output_embeddings()
160    if input_emb is None or output_emb is None:
161        return False
162
163    identity_tied = getattr(input_emb, "weight", None) is getattr(output_emb, "weight", None)
164    config_tied = getattr(getattr(model, "config", None), "tie_word_embeddings", None)
165
166    if config_tied is not None and bool(config_tied) != identity_tied:
167        message = (
168            f"Tied-embedding detection disagreement: config.tie_word_embeddings="
169            f"{config_tied} but weight identity check says {identity_tied}."
170        )
171        if strict:
172            raise TiedEmbeddingConflictError(
173                message + " This model was likely constructed without calling "
174                "tie_weights(), or was wrapped (FSDP/DeepSpeed/quantization/PEFT) "
175                "before this check ran."
176            )
177        warnings.warn(
178            message + " Trusting the identity check. If this model went through "
179            "FSDP/DeepSpeed/quantization/PEFT before this check, re-run detection "
180            "earlier in your pipeline.",
181            stacklevel=2,
182        )
183
184    return identity_tied
185
186
187def resolve_model(model: torch.nn.Module, strict_tie_check: bool = False) -> ResolvedModel:
188    """Run all resolution steps once and bundle the results.
189
190    Call this once at setup time (post model-construction, pre distributed/
191    quantization wrapping) and pass the result to downstream consumers (e.g.
192    trainable_parameters.set_trainable_parameters) instead of re-resolving.
193    """
194    decoder_layers = find_decoder_layers(model)
195    final_norm = find_final_norm(model)
196    input_embeddings = model.get_input_embeddings()
197    output_embeddings = model.get_output_embeddings()
198    embeddings_tied = detect_tied_embeddings(model, strict=strict_tie_check)
199    tie_config = getattr(getattr(model, "config", None), "tie_word_embeddings", None)
200
201    return ResolvedModel(
202        root=model,
203        decoder_layers=decoder_layers,
204        final_norm=final_norm,
205        input_embeddings=input_embeddings,
206        output_embeddings=output_embeddings,
207        embeddings_tied=embeddings_tied,
208        tie_word_embeddings_config=tie_config,
209    )