amachine.am_transformers.am_set_trainable

  1"""
  2Sets requires_grad on a HuggingFace transformer model based on three
  3independent axes — decoder layers, input embeddings, output layer (lm_head +
  4final norm) — and prints a readable summary of what ended up trainable.
  5
  6Depends on model_introspection.py for all module-finding logic; this file is
  7purely about applying requires_grad and reporting the result.
  8"""
  9
 10from __future__ import annotations
 11
 12from typing import Literal
 13
 14import torch
 15
 16from .am_model_introspection import (
 17    ResolvedModel,
 18    TiedEmbeddingConflictError,
 19    resolve_model,
 20)
 21
 22def set_trainable_parameters(
 23    model: torch.nn.Module,
 24    layers_to_train: Literal["all"] | list[int],
 25    train_input_embeddings: bool,
 26    train_output_layer: bool,
 27    *,
 28    strict_tie_check: bool = False,
 29    verbose: bool = True,
 30) -> list[torch.nn.Parameter]:
 31    """Freeze the whole model, then selectively unfreeze:
 32      - the requested decoder layers (their own norm/attention/mlp submodules
 33        come along for free, since we unfreeze at the layer level)
 34      - the input embedding matrix, if requested
 35      - the lm_head + final pre-lm_head norm, if requested
 36
 37    Raises TiedEmbeddingConflictError if input/output embeddings are tied
 38    (share the same weight tensor) and train_input_embeddings != train_output_layer,
 39    since the architecture cannot satisfy training one without the other.
 40
 41    Args:
 42        layers_to_train: "all", or a list of decoder layer indices (negative
 43            indices allowed, e.g. [-1, -2] for the last two layers).
 44        train_input_embeddings: unfreeze the input embedding matrix.
 45        train_output_layer: unfreeze lm_head + the final pre-lm_head norm.
 46        strict_tie_check: passed through to detect_tied_embeddings — raise
 47            instead of warn on config/identity disagreement. Set True when you
 48            control model construction directly.
 49        verbose: print a per-component trainable/frozen summary.
 50
 51    Returns:
 52        List of parameters with requires_grad=True, suitable for passing
 53        directly to an optimizer.
 54    """
 55    resolved = resolve_model(model, strict_tie_check=strict_tie_check)
 56
 57    if resolved.embeddings_tied and train_input_embeddings != train_output_layer:
 58        raise TiedEmbeddingConflictError(
 59            "Input embeddings and lm_head are tied (share the same weight tensor "
 60            f"via config.tie_word_embeddings={resolved.tie_word_embeddings_config}), "
 61            f"but train_input_embeddings={train_input_embeddings} and "
 62            f"train_output_layer={train_output_layer} disagree. These cannot be "
 63            "set independently on this model — either set both to the same "
 64            "value, or untie the weights first, e.g.:\n"
 65            "    out = model.get_output_embeddings()\n"
 66            "    out.weight = torch.nn.Parameter(out.weight.clone())"
 67        )
 68
 69    # 1. Freeze everything.
 70    for p in model.parameters():
 71        p.requires_grad = False
 72
 73    # 2. Unfreeze requested decoder layers.
 74    num_layers = len(resolved.decoder_layers)
 75    layer_indices = (
 76        list(range(num_layers)) if layers_to_train == "all" else list(layers_to_train)
 77    )
 78    # Normalize negative indices for reporting purposes.
 79    normalized_indices = {i % num_layers for i in layer_indices}
 80    for idx in layer_indices:
 81        for p in resolved.decoder_layers[idx].parameters():
 82            p.requires_grad = True
 83
 84    # 3. Input embeddings.
 85    if train_input_embeddings and resolved.input_embeddings is not None:
 86        for p in resolved.input_embeddings.parameters():
 87            p.requires_grad = True
 88
 89    # 4. Output layer: lm_head + final norm.
 90    if train_output_layer and resolved.output_embeddings is not None:
 91        for p in resolved.output_embeddings.parameters():
 92            p.requires_grad = True
 93            
 94    if train_output_layer and resolved.final_norm is not None:
 95        for p in resolved.final_norm.parameters():
 96            p.requires_grad = True
 97
 98    trainable_parameters = [p for p in model.parameters() if p.requires_grad]
 99    if not trainable_parameters:
100        raise RuntimeError("No trainable parameters selected — check your training config.")
101
102    if verbose:
103        _print_summary(
104            model=model,
105            resolved=resolved,
106            num_layers=num_layers,
107            normalized_indices=normalized_indices,
108            train_input_embeddings=train_input_embeddings,
109            train_output_layer=train_output_layer,
110        )
111
112    return trainable_parameters
113
114
115# ---------------------------------------------------------------------------
116# Reporting
117# ---------------------------------------------------------------------------
118
119def _count_params(module: torch.nn.Module | None) -> tuple[int, int]:
120    """Returns (trainable_count, total_count) for a module's own parameters."""
121    if module is None:
122        return 0, 0
123    total = 0
124    trainable = 0
125    for p in module.parameters():
126        n = p.numel()
127        total += n
128        if p.requires_grad:
129            trainable += n
130    return trainable, total
131
132
133def _fmt(n: int) -> str:
134    if n >= 1_000_000:
135        return f"{n / 1_000_000:.2f}M"
136    if n >= 1_000:
137        return f"{n / 1_000:.1f}K"
138    return str(n)
139
140
141def _print_summary(
142    model: torch.nn.Module,
143    resolved: ResolvedModel,
144    num_layers: int,
145    normalized_indices: set[int],
146    train_input_embeddings: bool,
147    train_output_layer: bool,
148) -> None:
149    lines: list[str] = []
150    lines.append("=" * 72)
151    lines.append("Trainable parameter summary")
152    lines.append("=" * 72)
153
154    # Decoder layers, one line each so nothing is silently missed.
155    lines.append(f"\nDecoder layers ({num_layers} total):")
156    for idx in range(num_layers):
157        trainable_n, total_n = _count_params(resolved.decoder_layers[idx])
158        status = "TRAIN" if idx in normalized_indices else "frozen"
159        marker = "x" if idx in normalized_indices else " "
160        lines.append(
161            f"  [{marker}] layer {idx:>3}  {status:<6}  "
162            f"{_fmt(trainable_n):>8} / {_fmt(total_n):<8} trainable"
163        )
164
165    # Input embeddings.
166    in_trainable, in_total = _count_params(resolved.input_embeddings)
167    lines.append("\nInput embeddings:")
168    if resolved.input_embeddings is None:
169        lines.append("  not found on this model")
170    else:
171        status = "TRAIN" if train_input_embeddings else "frozen"
172        lines.append(
173            f"  [{'x' if train_input_embeddings else ' '}] {status:<6}  "
174            f"{_fmt(in_trainable):>8} / {_fmt(in_total):<8} trainable  "
175            f"({type(resolved.input_embeddings).__name__})"
176        )
177
178    # Output layer: lm_head + final norm.
179    out_trainable, out_total = _count_params(resolved.output_embeddings)
180    norm_trainable, norm_total = _count_params(resolved.final_norm)
181    lines.append("\nOutput layer:")
182    if resolved.output_embeddings is None:
183        lines.append("  lm_head: not found on this model")
184    else:
185        status = "TRAIN" if train_output_layer else "frozen"
186        lines.append(
187            f"  [{'x' if train_output_layer else ' '}] lm_head      {status:<6}  "
188            f"{_fmt(out_trainable):>8} / {_fmt(out_total):<8} trainable  "
189            f"({type(resolved.output_embeddings).__name__})"
190        )
191    if resolved.final_norm is None:
192        lines.append("  final norm:  not found on this model (skipped, non-fatal)")
193    else:
194        status = "TRAIN" if train_output_layer else "frozen"
195        lines.append(
196            f"  [{'x' if train_output_layer else ' '}] final norm   {status:<6}  "
197            f"{_fmt(norm_trainable):>8} / {_fmt(norm_total):<8} trainable  "
198            f"({type(resolved.final_norm).__name__})"
199        )
200
201    # Tied-embedding note.
202    lines.append("\nTied embeddings:")
203    lines.append(
204        f"  {resolved.embeddings_tied}  "
205        f"(config.tie_word_embeddings={resolved.tie_word_embeddings_config})"
206    )
207
208    # Anything outside decoder layers / embeddings / lm_head / final norm
209    # that ended up trainable (e.g. unexpected top-level modules) or that
210    # exists but wasn't touched at all — catches architectures with extra
211    # top-level submodules (adapters, pooler heads, etc.) so nothing is silently
212    # missed either way.
213    accounted_ids = {id(p) for p in resolved.decoder_layers.parameters()}
214    if resolved.input_embeddings is not None:
215        accounted_ids |= {id(p) for p in resolved.input_embeddings.parameters()}
216    if resolved.output_embeddings is not None:
217        accounted_ids |= {id(p) for p in resolved.output_embeddings.parameters()}
218    if resolved.final_norm is not None:
219        accounted_ids |= {id(p) for p in resolved.final_norm.parameters()}
220
221    other_trainable = 0
222    other_total = 0
223    other_names: list[str] = []
224    for name, p in model.named_parameters():
225        if id(p) in accounted_ids:
226            continue
227        other_total += p.numel()
228        if p.requires_grad:
229            other_trainable += p.numel()
230            other_names.append(name)
231
232    if other_total > 0:
233        lines.append("\nOther top-level parameters (outside layers/embeddings/lm_head/norm):")
234        lines.append(
235            f"  {_fmt(other_trainable)} / {_fmt(other_total)} trainable"
236        )
237        if other_names:
238            lines.append(f"  trainable: {', '.join(other_names[:10])}"
239                          + (" ..." if len(other_names) > 10 else ""))
240
241    # Grand total.
242    total_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
243    total_all = sum(p.numel() for p in model.parameters())
244    pct = 100 * total_trainable / total_all if total_all else 0.0
245    lines.append("\n" + "-" * 72)
246    lines.append(
247        f"TOTAL: {_fmt(total_trainable)} / {_fmt(total_all)} trainable ({pct:.2f}%)"
248    )
249    lines.append("=" * 72)
250
251    print("\n".join(lines))