from __future__ import annotations import math import random from collections.abc import Callable, Iterable, Sequence from pathlib import Path from typing import Any import numpy as np import torch import torch.nn.functional as F from torch import Tensor def _require_matplotlib() -> Any: try: import matplotlib.pyplot as plt # type: ignore except ImportError as exc: # pragma: no cover - optional dependency raise ImportError("Install mini-transformer[viz] to enable plotting utilities.") from exc return plt # ----------------------------------------------------------------------------- # Global seeding utilities # ----------------------------------------------------------------------------- def set_global_seed(seed: int, *, deterministic: bool = True) -> None: """Seed Python, NumPy, and PyTorch (CPU and CUDA) for reproducible runs.""" random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) if deterministic: torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False def make_worker_init_fn(seed: int) -> Callable[[int], None]: """Return a DataLoader worker init function that derives unique seeds.""" def _init_fn(worker_id: int) -> None: worker_seed = seed + worker_id random.seed(worker_seed) np.random.seed(worker_seed) torch.manual_seed(worker_seed) return _init_fn # ----------------------------------------------------------------------------- # Attention building blocks # ----------------------------------------------------------------------------- def split_heads(x: torch.Tensor, num_heads: int) -> torch.Tensor: """Split the last dimension into ``(num_heads, d_head)`` and permute to (B, H, S, d_head).""" if not isinstance(x, torch.Tensor): raise TypeError(f"x must be a torch.Tensor, got {type(x)}") if not isinstance(num_heads, int): raise TypeError(f"num_heads must be an int, got {type(num_heads)}") if x.ndim != 3: raise ValueError( f"x must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(x.shape)}" ) if num_heads <= 0: raise ValueError(f"num_heads must be > 0; got {num_heads}") batch_size, seq_length, d_model = x.shape if d_model % num_heads != 0: raise ValueError( f"d_model ({d_model}) must be divisible by num_heads ({num_heads});" f" got remainder {d_model % num_heads}" ) d_head = d_model // num_heads return x.reshape(batch_size, seq_length, num_heads, d_head).permute(0, 2, 1, 3) def join_heads(x: torch.Tensor) -> torch.Tensor: """Merge ``(num_heads, d_head)`` back into the model dimension.""" if not isinstance(x, torch.Tensor): raise TypeError(f"Expected torch.Tensor, got {type(x)}") if x.ndim != 4: raise ValueError( f"Expected 4D torch.Tensor (batch, num_heads, seq_len, d_head), got {tuple(x.shape)}" ) batch_size, num_heads, seq_length, d_head = x.shape return x.permute(0, 2, 1, 3).contiguous().view(batch_size, seq_length, num_heads * d_head) def calculate_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor | None, *, dropout_p: float = 0.0, return_probs: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """Scaled dot-product attention with optional dropout and probability return. Args: query, key, value: ``[B, H, S, Dh]`` tensors on the same device/dtype. mask: Optional boolean tensor broadcastable to ``[B, H, Sq, Sk]`` where ``True`` entries are masked. ``None`` disables masking. dropout_p: Dropout probability applied to the attention weights. Callers should pass ``0.0`` when not training. return_probs: When ``True`` returns ``(context, probs)``; otherwise only the context tensor is returned. """ if not (query.dim() == key.dim() == value.dim() == 4): raise ValueError( "query, key, value must be 4D tensors shaped (B, H, S, Dh);" f" got q={tuple(query.shape)}, k={tuple(key.shape)}, v={tuple(value.shape)}" ) if query.device != key.device or query.device != value.device: raise RuntimeError("query, key, value must be on the same device") if mask is not None: if mask.dtype != torch.bool: raise TypeError("mask must be boolean when provided") if mask.device != query.device: mask = mask.to(query.device) target_shape = (query.size(0), query.size(1), query.size(2), key.size(2)) if mask.shape != target_shape: try: mask = mask.expand(target_shape) except RuntimeError as exc: raise ValueError( f"mask with shape {tuple(mask.shape)} not broadcastable to {target_shape}" ) from exc p = float(dropout_p) if p < 0 or p >= 1: raise ValueError(f"dropout_p must be in [0, 1), got {dropout_p}") return _attention_with_probs(query, key, value, mask, p, return_probs) def _attention_with_probs( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor | None, dropout_p: float, return_probs: bool, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """Manual attention path supporting optional probability return.""" head_dim = query.size(-1) work_dtype = torch.float32 if query.dtype in (torch.float16, torch.bfloat16) else query.dtype q = query.to(work_dtype) k = key.to(work_dtype) scores = torch.matmul(q, k.transpose(-2, -1)) scores.mul_(1.0 / math.sqrt(head_dim)) full_mask_rows = None if mask is not None: mask = mask.to(scores.device) fill_value = torch.finfo(scores.dtype).min scores = scores.masked_fill(mask, fill_value) full_mask_rows = mask.all(dim=-1, keepdim=True) if full_mask_rows.any(): scores = scores.masked_fill(full_mask_rows, 0.0) row_max = scores.max(dim=-1, keepdim=True).values row_max = torch.where(torch.isfinite(row_max), row_max, torch.zeros_like(row_max)) logits = scores - row_max probs = torch.softmax(logits, dim=-1) if full_mask_rows is not None: probs = torch.where(full_mask_rows, torch.zeros_like(probs), probs) if dropout_p > 0.0: probs = F.dropout(probs, p=dropout_p, training=True) context = torch.matmul(probs.to(value.dtype), value) if return_probs: probs_out = probs.to(work_dtype) return context, probs_out return context # ----------------------------------------------------------------------------- # Positional encodings # ----------------------------------------------------------------------------- def sinusoidal_positional_encoding(S: int, D: int) -> torch.Tensor: """Return the classic sinusoidal positional encoding table (shape ``[S, D]``).""" if S <= 0 or D <= 0: raise ValueError(f"S and D must be > 0, got S={S}, D={D}") positions = torch.arange(S, dtype=torch.float32).unsqueeze(1) div_terms = torch.exp(torch.arange(0, D, 2, dtype=torch.float32) * -(math.log(10000.0) / D)) pe = torch.zeros(S, D, dtype=torch.float32) pe[:, 0::2] = torch.sin(positions * div_terms) pe[:, 1::2] = torch.cos(positions * div_terms[: D // 2]) return pe # ----------------------------------------------------------------------------- # Sampling utilities # ----------------------------------------------------------------------------- def _ensure_2d_logits(logits: torch.Tensor) -> torch.Tensor: """Reshape logits so that the last dimension is vocabulary sized and the rest flatten.""" if logits.dim() == 2: return logits if logits.dim() >= 3: V = logits.size(-1) return logits.reshape(-1, V) raise ValueError(f"sample_from_logits: logits must be at least 2D, got {tuple(logits.shape)}") def _apply_allow_deny_mask( logits: torch.Tensor, *, allowed_tokens: Iterable[int] | None, disallowed_tokens: Iterable[int] | None, filter_value: float, ) -> torch.Tensor: if allowed_tokens is not None: mask = torch.zeros_like(logits, dtype=torch.bool) idx = torch.tensor(list(allowed_tokens), device=logits.device) idx = idx[(idx >= 0) & (idx < logits.size(-1))] if idx.numel() > 0: mask.index_fill_(-1, idx, True) logits = torch.where(mask, logits, torch.full_like(logits, filter_value)) if disallowed_tokens is not None: idx = torch.tensor(list(disallowed_tokens), device=logits.device) idx = idx[(idx >= 0) & (idx < logits.size(-1))] if idx.numel() > 0: logits.index_fill_(-1, idx, filter_value) return logits def _top_k_filtering( logits: torch.Tensor, top_k: int | None, *, min_tokens_to_keep: int, filter_value: float, ) -> torch.Tensor: if top_k is None or top_k <= 0: return logits k = min(max(top_k, min_tokens_to_keep), logits.size(-1)) values, _ = torch.topk(logits, k, dim=-1) threshold = values[..., -1, None] return torch.where(logits < threshold, torch.full_like(logits, filter_value), logits) def _top_p_filtering( logits_scaled: torch.Tensor, probs: torch.Tensor, top_p: float, *, min_tokens_to_keep: int, filter_value: float, ) -> torch.Tensor: if top_p is None or not (0.0 < top_p < 1.0): return logits_scaled sorted_probs, sorted_idx = torch.sort(probs, dim=-1, descending=True) cumulative = torch.cumsum(sorted_probs, dim=-1) to_remove = cumulative > top_p to_remove[..., :min_tokens_to_keep] = False scatter_mask = torch.zeros_like(to_remove, dtype=torch.bool).scatter(-1, sorted_idx, to_remove) return torch.where(scatter_mask, torch.full_like(logits_scaled, filter_value), logits_scaled) def sample_from_logits( logits: torch.Tensor, *, do_sample: bool = False, temperature: float = 1.0, top_k: int | None = None, top_p: float | None = None, min_tokens_to_keep: int = 1, allowed_tokens: Iterable[int] | None = None, disallowed_tokens: Iterable[int] | None = None, repetition_ctx: torch.Tensor | None = None, presence_penalty: float = 0.0, frequency_penalty: float = 0.0, no_repeat_ngram_size: int | None = None, filter_value: float = -float("inf"), rng: torch.Generator | None = None, return_probs: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """Warp logits (temperature, top-k, nucleus, penalties) and produce next-token ids.""" if temperature <= 0: raise ValueError("temperature must be > 0") logits2d = _ensure_2d_logits(logits).to(dtype=torch.float32) batch, vocab = logits2d.shape logits2d = _apply_allow_deny_mask( logits2d, allowed_tokens=allowed_tokens, disallowed_tokens=disallowed_tokens, filter_value=filter_value, ) if repetition_ctx is not None and (presence_penalty > 0.0 or frequency_penalty > 0.0): if repetition_ctx.dim() != 2 or repetition_ctx.size(0) != batch: raise ValueError( f"repetition_ctx must be [B, T]; got {tuple(repetition_ctx.shape)} with B={batch}" ) counts = torch.zeros((batch, vocab), device=logits2d.device, dtype=torch.float32) ctx = repetition_ctx valid = (ctx >= 0) & (ctx < vocab) if valid.any(): ids = ctx.masked_select(valid) bidx = ( torch.arange(batch, device=logits2d.device) .unsqueeze(1) .expand_as(ctx) .masked_select(valid) ) counts.index_put_( (bidx, ids), torch.ones_like(ids, dtype=torch.float32), accumulate=True ) if presence_penalty > 0.0: logits2d = logits2d - presence_penalty * (counts > 0).to(logits2d.dtype) if frequency_penalty > 0.0: logits2d = logits2d - frequency_penalty * counts if ( no_repeat_ngram_size is not None and no_repeat_ngram_size >= 2 and repetition_ctx is not None ): ctx_cpu = repetition_ctx.detach().to("cpu") for b, seq in enumerate(ctx_cpu.tolist()): tokens = [t for t in seq if 0 <= t < vocab] if len(tokens) < no_repeat_ngram_size: continue history: dict[tuple[int, ...], set[int]] = {} n = int(no_repeat_ngram_size) for i in range(len(tokens) - (n - 1)): key = tuple(tokens[i : i + n - 1]) nxt = tokens[i + n - 1] history.setdefault(key, set()).add(nxt) key = tuple(tokens[-(n - 1) :]) banned = history.get(key) if banned: idx = torch.tensor(list(banned), device=logits2d.device, dtype=torch.long) logits2d[b].index_fill_(0, idx, filter_value) logits_scaled = logits2d / float(temperature) logits_scaled = _top_k_filtering( logits_scaled, top_k=top_k, min_tokens_to_keep=min_tokens_to_keep, filter_value=filter_value ) probs = F.softmax(logits_scaled, dim=-1) if top_p is not None and 0.0 < top_p < 1.0: logits_scaled = _top_p_filtering( logits_scaled, probs, top_p=top_p, min_tokens_to_keep=min_tokens_to_keep, filter_value=filter_value, ) probs = F.softmax(logits_scaled, dim=-1) if do_sample: next_ids = torch.multinomial(probs, num_samples=1, replacement=True, generator=rng).squeeze( -1 ) else: next_ids = torch.argmax(probs, dim=-1) next_ids = next_ids.to(dtype=torch.long, device=logits.device) if return_probs: return next_ids, probs.to(device=logits.device, dtype=probs.dtype) return next_ids # ----------------------------------------------------------------------------- # Attention mask helpers # ----------------------------------------------------------------------------- def create_causal_mask(x: torch.Tensor, num_heads: int) -> torch.Tensor: """Return a boolean causal mask ``[B, num_heads, S, S]`` for decoder self-attention.""" if x.ndim != 2: raise ValueError(f"Expected input of shape (B, S), got {x.shape}") if num_heads <= 0: raise ValueError(f"num_heads must be > 0, got {num_heads}") batch, seq_len = x.shape if seq_len <= 0: raise ValueError(f"Sequence length must be > 0, got {seq_len}") base = torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=x.device), diagonal=1) return base.view(1, 1, seq_len, seq_len).expand(batch, num_heads, seq_len, seq_len).contiguous() def create_qk_padding_mask( query_attention_mask: torch.Tensor, key_attention_mask: torch.Tensor ) -> torch.Tensor: """Combine query and key padding masks (boolean) into a ``[B, H, Sq, Sk]`` mask.""" Bq, Hq, _, Sq = query_attention_mask.shape Bk, Hk, _, Sk = key_attention_mask.shape if Bq != Bk or Hq != Hk: raise ValueError( f"Padding mask batch/head mismatch: query {query_attention_mask.shape} vs key {key_attention_mask.shape}" ) q_mask = query_attention_mask.to(torch.bool).view(Bq, Hq, Sq, 1) k_mask = key_attention_mask.to(torch.bool).view(Bk, Hk, 1, Sk) return (q_mask | k_mask).expand(Bq, Hq, Sq, Sk) def broadcast_padding_mask(mask: torch.Tensor, num_heads: int) -> torch.Tensor: """Expand a ``[B, S]`` padding mask to ``[B, H, 1, S]``.""" if not isinstance(mask, torch.Tensor): raise TypeError(f"mask must be a torch.Tensor, got {type(mask)}") if mask.dim() != 2: raise ValueError(f"mask must be [B, S], got shape {tuple(mask.shape)}") if not isinstance(num_heads, int) or num_heads <= 0: raise ValueError(f"num_heads must be a positive int, got {num_heads}") batch, seq_len = mask.shape return mask.unsqueeze(1).unsqueeze(2).expand(batch, num_heads, 1, seq_len) def combine_masks( causal_mask: torch.Tensor | None, padding_mask: torch.Tensor | None ) -> torch.Tensor | None: """Combine causal and padding masks (boolean OR) while handling ``None`` values.""" if padding_mask is None: return causal_mask.to(torch.bool) if causal_mask is not None else None if causal_mask is None: return padding_mask.to(torch.bool) m1 = causal_mask.to(torch.bool) m2 = padding_mask.to(torch.bool).to(device=m1.device) try: return m1 | m2 except RuntimeError as exc: raise ValueError( f"Masks not broadcastable: causal {tuple(m1.shape)} vs padding {tuple(m2.shape)}" ) from exc # ----------------------------------------------------------------------------- # Attention introspection / plotting # ----------------------------------------------------------------------------- def extract_all_attention_maps( model, src_ids: torch.Tensor, tgt_ids: torch.Tensor, src_padding_2d: torch.Tensor, tgt_padding_2d: torch.Tensor, ): """Collect attention probability tensors for encoder only/decoder self/cross attention.""" was_training = model.training model.eval() device = src_ids.device num_heads = model.cfg.num_heads src_pad_b = broadcast_padding_mask(src_padding_2d.to(device), num_heads) tgt_pad_b = broadcast_padding_mask(tgt_padding_2d.to(device), num_heads) tgt_causal = create_causal_mask(tgt_ids.to(device), num_heads) x = model.embed(src_ids.to(device)) y = model.embed(tgt_ids.to(device)) enc_layers = model.encoder.layers dec_layers = model.decoder.layers enc_self_maps: list[Tensor] = [] dec_self_maps: list[Tensor] = [] dec_cross_maps: list[Tensor] = [] cur = x for layer in enc_layers: mha = layer.attention_layer attn_input = layer.norm1(cur) if getattr(layer, "pre_norm", False) else cur q = split_heads(mha.query_linear(attn_input), mha.num_heads) k = split_heads(mha.key_linear(attn_input), mha.num_heads) v = split_heads(mha.value_linear(attn_input), mha.num_heads) mask = create_qk_padding_mask(src_pad_b, src_pad_b) _, probs = calculate_attention(q, k, v, mask, dropout_p=0.0, return_probs=True) enc_self_maps.append(probs) cur = layer(cur, src_pad_b) mem_full = cur y_in = y for layer in dec_layers: pre_norm = getattr(layer, "pre_norm", False) self_mha = layer.self_attention_layer self_attn_input = layer.norm1(y_in) if pre_norm else y_in q = split_heads(self_mha.query_linear(self_attn_input), self_mha.num_heads) k = split_heads(self_mha.key_linear(self_attn_input), self_mha.num_heads) v = split_heads(self_mha.value_linear(self_attn_input), self_mha.num_heads) pad = create_qk_padding_mask(tgt_pad_b, tgt_pad_b) combined_mask = combine_masks(tgt_causal, pad) self_ctx, self_probs = calculate_attention( q, k, v, combined_mask, dropout_p=0.0, return_probs=True ) dec_self_maps.append(self_probs) self_ctx_merged = join_heads(self_ctx) self_projected = self_mha.output_linear(self_ctx_merged) self_residual = y_in + layer.dropout1(self_projected) if pre_norm: cross_query_input = layer.norm2(self_residual) else: cross_query_input = layer.norm1(self_residual) cross_mha = layer.cross_attention_layer cq = split_heads(cross_mha.query_linear(cross_query_input), cross_mha.num_heads) ck = split_heads(cross_mha.key_linear(mem_full), cross_mha.num_heads) cv = split_heads(cross_mha.value_linear(mem_full), cross_mha.num_heads) cmask = create_qk_padding_mask(tgt_pad_b, src_pad_b) _, cross_probs = calculate_attention(cq, ck, cv, cmask, dropout_p=0.0, return_probs=True) dec_cross_maps.append(cross_probs) y_in = layer(mem_full, y_in, src_pad_b, tgt_pad_b, tgt_causal) maps = {"enc_self": enc_self_maps, "dec_self": dec_self_maps, "dec_cross": dec_cross_maps} model.train(was_training) return maps def attach_tokens(maps_dict, tokenizer, src_ids, tgt_ids): """Attach decoded token strings to a maps dictionary produced by ``extract_all_attention_maps``.""" src_tok = [ tokenizer.convert_ids_to_tokens(row, skip_special_tokens=False) for row in src_ids.tolist() ] tgt_tok = [ tokenizer.convert_ids_to_tokens(row, skip_special_tokens=False) for row in tgt_ids.tolist() ] maps_dict["src_tokens"] = src_tok maps_dict["tgt_tokens"] = tgt_tok return maps_dict def _imshow(ax, array2d, title="", vmin=0.0, vmax=1.0): image = ax.imshow(array2d, aspect="auto", vmin=vmin, vmax=vmax) ax.set_title(title, fontsize=9) ax.set_xticks([]) ax.set_yticks([]) return image def plot_layer_heads_grid( maps_dict: dict, layer_idx: int, batch: int = 0, heads: list[int] | None = None, figsize_per_cell: tuple[float, float] = (2.2, 2.2), show_colorbar: bool = True, vmin: float = 0.0, vmax: float = 1.0, *, show: bool = True, save_path: str | Path | None = None, ): """Plot a single layer as a 3-row (enc self / dec self / dec cross) grid.""" plt = _require_matplotlib() enc_layers = maps_dict["enc_self"] dec_self_layers = maps_dict["dec_self"] dec_cross_layers = maps_dict["dec_cross"] enc_idx = min(layer_idx, len(enc_layers) - 1) dec_idx = min(layer_idx, len(dec_self_layers) - 1) enc = enc_layers[enc_idx][batch] dec_self = dec_self_layers[dec_idx][batch] dec_cross = dec_cross_layers[dec_idx][batch] num_heads = enc.size(0) heads = list(range(num_heads)) if heads is None else heads rows = 3 cols = len(heads) fig_width = max(1, int(round(figsize_per_cell[0] * cols))) fig_height = max(1, int(round(figsize_per_cell[1] * rows))) fig, axes = plt.subplots( rows, cols, figsize=(fig_width, fig_height), squeeze=False, constrained_layout=True ) images = [] for column, head in enumerate(heads): images.append( _imshow(axes[0, column], enc[head].cpu().float().numpy(), f"Enc h{head}", vmin, vmax) ) _imshow( axes[1, column], dec_self[head].cpu().float().numpy(), f"Dec self h{head}", vmin, vmax ) _imshow( axes[2, column], dec_cross[head].cpu().float().numpy(), f"Dec cross h{head}", vmin, vmax ) if show_colorbar and images: fig.colorbar(images[0], ax=axes, fraction=0.02, pad=0.01) fig.suptitle(f"Layer {layer_idx}", fontsize=12) if save_path is not None: save_path = Path(save_path) save_path.parent.mkdir(parents=True, exist_ok=True) fig.savefig(save_path, dpi=150, bbox_inches="tight") if show: plt.show() else: plt.close(fig) return fig def plot_all_layers_all_heads( maps_dict: dict, batch: int = 0, max_layers: int | None = None, heads: list[int] | None = None, figsize_per_cell: tuple[float, float] = (2.0, 2.0), vmin: float = 0.0, vmax: float = 1.0, save_pdf_path: str | None = None, *, show: bool = True, ): """Render attention grids for every layer (optionally saving a multi-page PDF).""" plt = _require_matplotlib() from matplotlib.backends.backend_pdf import PdfPages # type: ignore enc_layers = len(maps_dict["enc_self"]) dec_layers = len(maps_dict["dec_self"]) total_layers = max(enc_layers, dec_layers) if max_layers is not None: total_layers = min(total_layers, max_layers) def _render_layers(record_page): for layer in range(total_layers): fig = plot_layer_heads_grid( maps_dict, layer_idx=layer, batch=batch, heads=heads, figsize_per_cell=figsize_per_cell, show_colorbar=True, vmin=vmin, vmax=vmax, show=show, ) if record_page is not None: record_page(fig) plt.close(fig) if save_pdf_path: with PdfPages(save_pdf_path) as pdf: _render_layers(pdf.savefig) print(f"[saved] multi-page attention viewer -> {save_pdf_path}") else: _render_layers(None) def _format_token_sequence(tokens: list[str]) -> str: return " | ".join(tokens) def _resolve_attention_layers(available: int, requested: Sequence[int] | None) -> list[int]: if requested is None: return list(range(available)) return sorted({int(idx) for idx in requested if 0 <= int(idx) < available}) def _resolve_attention_heads(available: int, requested: Sequence[int] | None) -> list[int]: if requested is None: return list(range(available)) return sorted({int(idx) for idx in requested if 0 <= int(idx) < available}) # ----------------------------------------------------------------------------- # Debug helper # ----------------------------------------------------------------------------- def debug_transformer_forward( model, tokenizer, src_ids: torch.Tensor, tgt_ids: torch.Tensor, *, logits: torch.Tensor | None = None, pad_id: int, device: str | torch.device = "cpu", batch_index: int = 0, sample_index: int = 0, show_attention: bool = False, save_attention: bool = False, average_heads: bool = False, attention_types: Sequence[str] = ("enc_self", "dec_self", "dec_cross"), attention_layers: Sequence[int] | None = None, attention_heads: Sequence[int] | None = None, attention_figsize: tuple[float, float] = (4.0, 4.0), skip_special_tokens: bool = False, log_fn: Callable[[str], None] | None = print, return_maps: bool = False, save_dir: str | Path | None = None, run_dir: str | Path | None = None, ) -> dict[str, Any]: """Run a forward pass, compute diagnostics, and optionally render attention heatmaps.""" if log_fn is None: def log_fn(_: str) -> None: # type: ignore[redefinition] return device = torch.device(device) if src_ids.dim() != 2 or tgt_ids.dim() != 2: raise ValueError("src_ids and tgt_ids must be rank-2 tensors") batch_size = src_ids.size(0) if not (0 <= sample_index < batch_size): raise IndexError(f"sample_index {sample_index} out of range for batch size {batch_size}") src_ids = src_ids.to(device) tgt_ids = tgt_ids.to(device) if tgt_ids.size(1) < 2: raise ValueError("tgt_ids must contain at least BOS and one target token") decoder_in = tgt_ids[:, :-1] labels = tgt_ids[:, 1:] src_pad_mask = src_ids.eq(pad_id) tgt_pad_mask = decoder_in.eq(pad_id) tgt_pad_full = tgt_ids.eq(pad_id) was_training = model.training attention_requested = show_attention or save_attention or return_maps if logits is None: if not callable(model): raise TypeError("model must be callable") model.eval() with torch.no_grad(): logits = model(src_ids, decoder_in, src_pad_mask, tgt_pad_mask) model.train(was_training) else: logits = logits.to(device) if logits.dim() != 3 or logits.size(0) != batch_size: raise ValueError("logits must be of shape [batch, seq_len, vocab]") vocab_size = logits.size(-1) mask = labels.ne(pad_id) with torch.no_grad(): log_probs = F.log_softmax(logits, dim=-1) nll = F.nll_loss( log_probs.reshape(-1, vocab_size), labels.reshape(-1), reduction="sum", ignore_index=pad_id, ) tokens_total = int(mask.sum().item()) avg_nll = (nll / max(tokens_total, 1)).item() perplexity = math.exp(avg_nll) if tokens_total > 0 else float("nan") pred_ids = logits.argmax(dim=-1) correct_mask = pred_ids.eq(labels) & mask correct_tokens = int(correct_mask.sum().item()) token_accuracy = correct_tokens / max(tokens_total, 1) sample_labels = labels[sample_index] sample_preds = pred_ids[sample_index] sample_mask = mask[sample_index] sample_tokens_total = int(sample_mask.sum().item()) sample_correct_tokens = int((sample_preds.eq(sample_labels) & sample_mask).sum().item()) sample_token_accuracy = sample_correct_tokens / max(sample_tokens_total, 1) sample_exact_match = bool(torch.equal(sample_preds[sample_mask], sample_labels[sample_mask])) sample_decoder_in = decoder_in[sample_index] sample_pred_sequence = torch.cat([sample_decoder_in[:1], sample_preds], dim=0) with torch.no_grad(): sample_log_probs = F.log_softmax(logits[sample_index], dim=-1) sample_nll = F.nll_loss( sample_log_probs, sample_labels, reduction="sum", ignore_index=pad_id, ) sample_avg_nll = (sample_nll / max(sample_tokens_total, 1)).item() sample_perplexity = math.exp(sample_avg_nll) if sample_tokens_total > 0 else float("nan") sample_src_ids = src_ids[sample_index] sample_tgt_ids = tgt_ids[sample_index] sample_src_tokens = tokenizer.convert_ids_to_tokens( sample_src_ids.tolist(), skip_special_tokens=skip_special_tokens, ) sample_tgt_tokens = tokenizer.convert_ids_to_tokens( sample_tgt_ids.tolist(), skip_special_tokens=skip_special_tokens, ) sample_pred_tokens = tokenizer.convert_ids_to_tokens( sample_pred_sequence.tolist(), skip_special_tokens=skip_special_tokens, ) sample_src_text = tokenizer.decode(sample_src_ids.tolist(), skip_special_tokens=True) sample_tgt_text = tokenizer.decode(sample_tgt_ids.tolist(), skip_special_tokens=True) sample_pred_text = tokenizer.decode(sample_pred_sequence.tolist(), skip_special_tokens=True) log_fn("\n--- DEBUG TRANSFORMER FORWARD ---") log_fn(f"Batch index : {batch_index}") log_fn(f"Sample index: {sample_index}") log_fn(f"Batch token accuracy: {token_accuracy * 100:.2f}% ({correct_tokens}/{tokens_total})") log_fn( f"Batch NLL / ppl : {avg_nll:.4f} / {perplexity:.4f}" if tokens_total > 0 else "Batch NLL / ppl : n/a" ) log_fn( f"Sample token accuracy: {sample_token_accuracy * 100:.2f}% ({sample_correct_tokens}/{sample_tokens_total})" ) log_fn( f"Sample NLL / ppl : {sample_avg_nll:.4f} / {sample_perplexity:.4f}" if sample_tokens_total > 0 else "Sample NLL / ppl : n/a" ) log_fn(f"Sample exact match : {sample_exact_match}") log_fn("-- Source sequence --") log_fn(f"IDs : {sample_src_ids.tolist()}") log_fn(f"Tokens: {_format_token_sequence(sample_src_tokens)}") log_fn(f"Text : {sample_src_text}") log_fn("-- Target sequence --") log_fn(f"IDs : {sample_tgt_ids.tolist()}") log_fn(f"Tokens: {_format_token_sequence(sample_tgt_tokens)}") log_fn(f"Text : {sample_tgt_text}") log_fn("-- Predicted sequence --") log_fn(f"IDs : {sample_pred_sequence.tolist()}") log_fn(f"Tokens: {_format_token_sequence(sample_pred_tokens)}") log_fn(f"Text : {sample_pred_text}") result: dict[str, Any] = { "batch_index": batch_index, "token_accuracy": token_accuracy, "avg_negative_log_likelihood": avg_nll, "perplexity": perplexity, "tokens_total": tokens_total, "correct_tokens": correct_tokens, "sample": { "index": sample_index, "token_accuracy": sample_token_accuracy, "exact_match": sample_exact_match, "avg_negative_log_likelihood": sample_avg_nll, "perplexity": sample_perplexity, "source_ids": sample_src_ids.tolist(), "target_ids": sample_tgt_ids.tolist(), "predicted_ids": sample_pred_sequence.tolist(), "source_tokens": sample_src_tokens, "target_tokens": sample_tgt_tokens, "predicted_tokens": sample_pred_tokens, "source_text": sample_src_text, "target_text": sample_tgt_text, "predicted_text": sample_pred_text, }, } if attention_requested: plt = _require_matplotlib() if not hasattr(model, "embed"): raise AttributeError("model must expose an embed() method for debugging") with torch.no_grad(): attention_maps = extract_all_attention_maps( model, src_ids, tgt_ids, src_pad_mask, tgt_pad_full, ) target_types = [t for t in attention_types if t in attention_maps] if not target_types: target_types = [k for k in ("enc_self", "dec_self", "dec_cross") if k in attention_maps] target_dir: Path | None = None if save_attention: if save_dir is not None: target_dir = Path(save_dir) elif run_dir is not None: target_dir = Path(run_dir) / "debug" / "attention" else: log_fn( "[warn] save_attention requested but no save_dir/run_dir provided; skipping save." ) if target_dir is not None: target_dir.mkdir(parents=True, exist_ok=True) sample_label = f"batch{batch_index}_sample{sample_index}" figure_records: list[dict[str, Any]] = [] for att_type in target_types: layer_maps = attention_maps.get(att_type, []) if not layer_maps: continue layers_to_plot = _resolve_attention_layers(len(layer_maps), attention_layers) for layer_idx in layers_to_plot: layer_tensor = layer_maps[layer_idx][sample_index] head_candidates = _resolve_attention_heads(layer_tensor.size(0), attention_heads) if not head_candidates: continue if average_heads: averaged = layer_tensor[head_candidates].mean(dim=0).cpu().float().numpy() fig, ax = plt.subplots(figsize=attention_figsize) im = ax.imshow(averaged, aspect="auto", vmin=0.0, vmax=1.0) ax.set_title(f"{att_type} L{layer_idx} (avg heads)") ax.set_xlabel("Key index") ax.set_ylabel("Query index") fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04) save_path = None if target_dir is not None: save_path = target_dir / f"{sample_label}_{att_type}_L{layer_idx}_avg.png" fig.savefig(save_path, dpi=150, bbox_inches="tight") if show_attention: plt.show() else: plt.close(fig) figure_records.append( { "type": att_type, "layer": layer_idx, "heads": "average", "path": str(save_path) if save_path else None, } ) else: cols = len(head_candidates) fig, axes = plt.subplots( 1, cols, figsize=(attention_figsize[0] * cols, attention_figsize[1]), squeeze=False, constrained_layout=True, ) for col, head_idx in enumerate(head_candidates): ax = axes[0, col] head_map = layer_tensor[head_idx].cpu().float().numpy() im = ax.imshow(head_map, aspect="auto", vmin=0.0, vmax=1.0) ax.set_title(f"{att_type} L{layer_idx} H{head_idx}") ax.set_xlabel("Key index") if col == 0: ax.set_ylabel("Query index") fig.colorbar(im, ax=axes.ravel().tolist(), fraction=0.046, pad=0.04) save_path = None if target_dir is not None: head_tag = "-".join(str(h) for h in head_candidates) save_path = ( target_dir / f"{sample_label}_{att_type}_L{layer_idx}_H{head_tag}.png" ) fig.savefig(save_path, dpi=150, bbox_inches="tight") if show_attention: plt.show() else: plt.close(fig) figure_records.append( { "type": att_type, "layer": layer_idx, "heads": head_candidates, "path": str(save_path) if save_path else None, } ) if target_dir is not None: raw_path = target_dir / f"{sample_label}_attention.pt" torch.save( {k: [v.cpu() for v in tensors] for k, tensors in attention_maps.items()}, raw_path ) result.setdefault("attention", {})["raw_path"] = str(raw_path) summary_path = target_dir / f"{sample_label}_attention_all_layers.pdf" plot_all_layers_all_heads( attention_maps, batch=sample_index, figsize_per_cell=attention_figsize, save_pdf_path=str(summary_path), show=False, ) result.setdefault("attention", {})["summary_path"] = str(summary_path) if figure_records: result.setdefault("attention", {})["figures"] = figure_records if return_maps: result["attention_maps"] = attention_maps log_fn("--- END DEBUG ---\n") return result # ----------------------------------------------------------------------------- # Config compatibility checker # ----------------------------------------------------------------------------- def check_tokenizer_model_compatibility(model_cfg, tokenizer_cfg): """Ensure tokenizer and model configuration agree on core vocabulary settings.""" if model_cfg.vocab_size != tokenizer_cfg.vocab_size: raise ValueError( f"Vocab size mismatch: model={model_cfg.vocab_size}, tokenizer={tokenizer_cfg.vocab_size}" ) if model_cfg.max_seq_len != tokenizer_cfg.max_seq_len: raise ValueError( f"Max sequence length mismatch: model={model_cfg.max_seq_len}, tokenizer={tokenizer_cfg.max_seq_len}" ) if model_cfg.pad_id != tokenizer_cfg.pad_id: raise ValueError( f"pad_id mismatch: model={model_cfg.pad_id}, tokenizer={tokenizer_cfg.pad_id}" ) if model_cfg.bos_id != tokenizer_cfg.bos_id: raise ValueError( f"bos_id mismatch: model={model_cfg.bos_id}, tokenizer={tokenizer_cfg.bos_id}" ) if model_cfg.eos_id != tokenizer_cfg.eos_id: raise ValueError( f"eos_id mismatch: model={model_cfg.eos_id}, tokenizer={tokenizer_cfg.eos_id}" )