diff --git a/datasets/vcfdataset.py b/datasets/vcfdataset.py index f35286e..554b8f3 100644 --- a/datasets/vcfdataset.py +++ b/datasets/vcfdataset.py @@ -216,14 +216,26 @@ def __adjust_length(self, token_ids): attention_mask = attention_mask[: self.max_length] return token_ids, attention_mask - def _get_cres(self, gene_id: str, gene_info: dict, vcf_path: str) -> pd.DataFrame: - """Get the cres for a given gene id + def _resolve_cres( + self, gene_id: str, gene_info: dict, vcf_path: str + ) -> pd.DataFrame: + """Run the CRE-resolution pipeline for ``gene_id``. + + Reads the per-gene CRE manifest, runs ``ExtractSeqFromBed.process_subject`` + (which may silently drop CREs whose sequence extraction fails), and + applies the minus-strand reversal. The resulting rows are in the exact + order the model consumes them, so they align 1-to-1 with the cross- + attention key axis. Args: - gene_id: Gene ID + gene_id: ENSEMBL gene id. + gene_info: gene metadata returned by ``_get_gene_info``. + vcf_path: VCF used for sequence extraction (or ``None`` for ref). Returns: - pd.DataFrame: A dataframe containing the cres for the given gene id + pd.DataFrame with the per-CRE columns produced by ``process_subject`` + (notably ``chrom``, ``start_cre``, ``end_cre``, ``cCRE``, + ``sequence``) in model order. """ gene_cre_map_path = self.gene_cre_manifest.get_file_path(gene_id) genes_cre_map = multi_try_load_csv(gene_cre_map_path) @@ -244,6 +256,47 @@ def _get_cres(self, gene_id: str, gene_info: dict, vcf_path: str) -> pd.DataFram cres = cres.iloc[ ::-1 ] # reverse the cres if the gene is on the minus strand + return cres + + def get_cre_positions( + self, gene_id: str, vcf_path: str | None = None + ) -> pd.DataFrame: + """Return CRE genomic positions in the order the model consumes them. + + Re-runs the same CRE-resolution pipeline used internally during a + forward pass (manifest → ``ExtractSeqFromBed.process_subject`` → + strand flip). The returned rows therefore align 1-to-1 with the CRE + token (key) axis of the model's cross-attention, including any drops + ``process_subject`` may apply mid-manifest. + + Args: + gene_id: ENSEMBL gene id present in the dataset. + vcf_path: VCF path; defaults to ``self.vcf_path``. + + Returns: + DataFrame with columns ``chromosome``, ``start_cre``, ``end_cre``, + ``cre_name`` indexed in model order. + """ + if vcf_path is None: + vcf_path = self.vcf_path + gene_info = self._get_gene_info(gene_id) + cres = self._resolve_cres(gene_id, gene_info, vcf_path) + return ( + cres[["chrom", "start_cre", "end_cre", "cCRE"]] + .rename(columns={"chrom": "chromosome", "cCRE": "cre_name"}) + .reset_index(drop=True) + ) + + def _get_cres(self, gene_id: str, gene_info: dict, vcf_path: str) -> pd.DataFrame: + """Get the cres for a given gene id + + Args: + gene_id: Gene ID + + Returns: + pd.DataFrame: A dataframe containing the cres for the given gene id + """ + cres = self._resolve_cres(gene_id, gene_info, vcf_path) X = [] attentions = [] diff --git a/notebooks/explore_attention.ipynb b/notebooks/explore_attention.ipynb new file mode 100644 index 0000000..1bac9f6 --- /dev/null +++ b/notebooks/explore_attention.ipynb @@ -0,0 +1,726 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "intro-md", + "metadata": {}, + "source": [ + "# Exploring CRE → gene-tissue CLS attention in VariantFormer\n", + "\n", + "This notebook demonstrates how to use the `LogAttention` callback (added in the\n", + "`feature/attention-callback` branch) to extract per-layer, per-head attention\n", + "matrices during a forward pass of VariantFormer, and how to visualize the\n", + "attention from the **gene-tissue CLS token** — a per-tissue learnable\n", + "embedding (`MultiRegistry.registry_tokens[tissue_id]`) prepended to the gene\n", + "tokens, which the model pools to predict gene expression — over the\n", + "**cis-regulatory elements (CREs)** surrounding a gene. It plays the role a\n", + "CLS token would in a vanilla transformer, but is conditioned on tissue\n", + "identity, so we call it the *gene-tissue CLS*.\n", + "\n", + "By reading the **gene-tissue CLS row** (query index 0) of\n", + "the gene cross-attention matrix (`gene_modulator_crossMHA_`), we can\n", + "see which CREs the model 'looks at' when summarizing the gene for a given\n", + "tissue.\n", + "\n", + "The notebook:\n", + "\n", + "1. Loads a real `VF` checkpoint (`v4_ag`,\n", + " the all-genes model).\n", + "2. Builds a single-gene single-tissue batch from the example VCF\n", + " (`HG00096.vcf.gz`) and the reference genome.\n", + "3. Runs one forward pass under `log_attn.record_attention(...)` to capture\n", + " the attention matrices from a few layers of the gene modulator.\n", + "4. Plots the gene-tissue CLS row of the cross-attention as a function of\n", + " CRE genomic position, plus per-head heatmaps and a layer comparison.\n", + "5. Repeats the forward pass on the **reference genome** (no VCF) and\n", + " compares the two attention profiles to highlight where the variants\n", + " in `HG00096.vcf.gz` shift the model's focus.\n", + "\n", + "> **Requirements**: a CUDA GPU and the artifacts produced by\n", + "> `python download_artifacts.py` (model checkpoints, reference genome, VCF)." + ] + }, + { + "cell_type": "markdown", + "id": "setup-md", + "metadata": {}, + "source": [ + "## 1. Setup" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "setup-code", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-03T22:06:03.816129Z", + "iopub.status.busy": "2026-05-03T22:06:03.816030Z", + "iopub.status.idle": "2026-05-03T22:06:06.725393Z", + "shell.execute_reply": "2026-05-03T22:06:06.724841Z" + } + }, + "outputs": [], + "source": [ + "import os, sys, logging\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "import torch\n", + "import matplotlib.pyplot as plt\n", + "\n", + "# Walk up from cwd until we find the variantformer repo root (it contains\n", + "# the `processors/vcfprocessor.py` file). This works under Jupyter, VS Code,\n", + "# and `jupyter nbconvert --execute` without depending on ipynbname.\n", + "def _find_repo_root(start: Path) -> Path:\n", + " for p in [start, *start.parents]:\n", + " if (p / 'processors' / 'vcfprocessor.py').exists():\n", + " return p\n", + " raise RuntimeError('Could not locate variantformer repo root')\n", + "\n", + "REPO_PATH = _find_repo_root(Path.cwd())\n", + "sys.path.insert(0, str(REPO_PATH))\n", + "\n", + "from processors.vcfprocessor import VCFProcessor\n", + "from seq2gene.attn_log_callback import LogAttention\n", + "\n", + "logging.basicConfig(level=logging.WARNING, format='%(asctime)s %(levelname)s %(message)s')\n", + "\n", + "assert torch.cuda.is_available(), 'This notebook requires a CUDA GPU.'\n", + "print(f'GPU: {torch.cuda.get_device_name(0)}')\n", + "print(f'CUDA memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB')\n", + "print(f'Repo path: {REPO_PATH}')" + ] + }, + { + "cell_type": "markdown", + "id": "config-md", + "metadata": {}, + "source": [ + "## 2. Pick a gene, a tissue, and a VCF\n", + "\n", + "We use **APOE** (`ENSG00000130203.9`) on chromosome 19 with the example HG00096\n", + "VCF. APOE is a strongly tissue-regulated gene with well-characterized\n", + "regulatory elements, which makes it a good qualitative test case.\n", + "\n", + "Feel free to change `gene_id`, `tissue`, and `vcf_path` to explore other\n", + "examples. The query format mirrors the one used in `notebooks/vcf2exp.ipynb`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "config-code", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-03T22:06:06.726975Z", + "iopub.status.busy": "2026-05-03T22:06:06.726768Z", + "iopub.status.idle": "2026-05-03T22:06:06.733445Z", + "shell.execute_reply": "2026-05-03T22:06:06.733068Z" + } + }, + "outputs": [], + "source": [ + "MODEL_CLASS = 'v4_ag' # all-genes checkpoint (v4_pcg is protein-coding only)\n", + "GENE_ID = 'ENSG00000130203.9' # APOE\n", + "TISSUE = 'whole blood'\n", + "VCF_PATH = str(REPO_PATH / '_artifacts' / 'HG00096.vcf.gz')\n", + "\n", + "LAYERS_TO_LOG = [0, 6, 12, 18, 24] # Sparse selection of the 25 layers\n", + "\n", + "query_df = pd.DataFrame({'gene_id': [GENE_ID], 'tissues': [TISSUE]})\n", + "query_df" + ] + }, + { + "cell_type": "markdown", + "id": "loaddata-md", + "metadata": {}, + "source": [ + "## 3. Build the dataset and load the model\n", + "\n", + "`VCFProcessor` handles tokenization of the gene window and the surrounding\n", + "CREs, including any variants from the VCF. We load the checkpoint manually\n", + "instead of using `Trainer.predict` so we can wrap a single forward pass in\n", + "the `record_attention` context manager." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "loaddata-code", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-03T22:06:06.734712Z", + "iopub.status.busy": "2026-05-03T22:06:06.734606Z", + "iopub.status.idle": "2026-05-03T22:06:25.259006Z", + "shell.execute_reply": "2026-05-03T22:06:25.258342Z" + } + }, + "outputs": [], + "source": [ + "processor = VCFProcessor(model_class=MODEL_CLASS)\n", + "vcf_dataset, dataloader = processor.create_data(VCF_PATH, query_df)\n", + "\n", + "model, ckpt_path, _trainer = processor.load_model()\n", + "state_dict = torch.load(ckpt_path, map_location='cpu', weights_only=False)\n", + "model.load_state_dict(state_dict.get('state_dict', state_dict), strict=False)\n", + "_ = model.cuda().eval()\n", + "\n", + "# The model reads `self.trainer.precision` to know whether to autocast.\n", + "# We use bf16-mixed (matches the training setup) and let torch.amp do the cast.\n", + "model.trainer = type('T', (), {'precision': 'bf16-mixed'})()\n", + "print('Model loaded.')" + ] + }, + { + "cell_type": "markdown", + "id": "forward-md", + "metadata": {}, + "source": [ + "## 4. Forward pass with attention logging\n", + "\n", + "`LogAttention.record_attention(model)` flips the `log_attn_matrix` flag on\n", + "every selected encoder layer of the gene modulator (and optionally the\n", + "epigenetics modulator). During the forward pass the layer recomputes the\n", + "attention matrix in PyTorch (matching FlashAttention's output up to fp16\n", + "tolerance — see `tests/test_attn_log_callback.py`) and stores it in\n", + "`layer.attn_matrix`. On context exit the callback collects, processes, and\n", + "moves the matrices to CPU and resets the flags." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "forward-code", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-03T22:06:25.260708Z", + "iopub.status.busy": "2026-05-03T22:06:25.260566Z", + "iopub.status.idle": "2026-05-03T22:06:44.365389Z", + "shell.execute_reply": "2026-05-03T22:06:44.364549Z" + } + }, + "outputs": [], + "source": [ + "# keep_heads=True keeps the per-head dimension. Each per-batch entry is\n", + "# a (H, Q, K) tensor instead of being averaged to (Q, K).\n", + "log_attn = LogAttention(layer_ids=LAYERS_TO_LOG, keep_heads=True)\n", + "\n", + "batch = next(iter(dataloader))\n", + "for k, v in batch.items():\n", + " if isinstance(v, torch.Tensor):\n", + " batch[k] = v.cuda()\n", + " elif isinstance(v, list) and v and isinstance(v[0], torch.Tensor):\n", + " batch[k] = [t.cuda() for t in v]\n", + "\n", + "with torch.no_grad(), torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16):\n", + " with log_attn.record_attention(model, log_epigenetics=False, log_gene=True):\n", + " preds = model.predict_step(batch, 0)\n", + "\n", + "print('Captured matrices:')\n", + "for name, mats in log_attn.attention_matrices.items():\n", + " print(f' {name}: list of {len(mats)} batches, first shape = {tuple(mats[0].shape)}')\n", + "\n", + "pred_value = preds['pred_gene_exp'][0]\n", + "print(f'\\nPredicted log-expression for {GENE_ID} in {TISSUE!r}: {float(pred_value.flatten()[0]):.3f}')" + ] + }, + { + "cell_type": "markdown", + "id": "cre-md", + "metadata": {}, + "source": [ + "## 5. Recover CRE genomic coordinates\n", + "\n", + "The dataloader doesn't pass CRE coordinates through the batch — they are\n", + "consumed only as token IDs. `VCFDataset.get_cre_positions` replays the\n", + "exact CRE-resolution pipeline used during the forward pass (manifest →\n", + "`ExtractSeqFromBed.process_subject` → strand flip), so the returned rows\n", + "align 1-to-1 with the CRE token (key) axis of the model's cross-attention\n", + "— including any rows `process_subject` drops mid-manifest when sequence\n", + "extraction fails." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cre-code", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-03T22:06:44.367285Z", + "iopub.status.busy": "2026-05-03T22:06:44.367104Z", + "iopub.status.idle": "2026-05-03T22:07:09.868803Z", + "shell.execute_reply": "2026-05-03T22:07:09.868283Z" + } + }, + "outputs": [], + "source": [ + "gene_info = vcf_dataset._get_gene_info(GENE_ID)\n", + "# Pull positions directly from the dataset so they are guaranteed to align\n", + "# with the CRE tokens the model consumed (drops anywhere in the manifest\n", + "# are honored, not just at the tail).\n", + "cre_df = vcf_dataset.get_cre_positions(GENE_ID, VCF_PATH)\n", + "\n", + "cre_df['midpoint'] = (cre_df['start_cre'] + cre_df['end_cre']) // 2\n", + "gene_start_site = gene_info['start'] if gene_info['strand'] == '+' else gene_info['end']\n", + "cre_df['distance_to_gene_start'] = cre_df['midpoint'] - gene_start_site\n", + "if gene_info['strand'] == '-':\n", + " cre_df['distance_to_gene_start'] = -cre_df['distance_to_gene_start']\n", + "\n", + "print(\n", + " f\"Gene {GENE_ID} on {gene_info['chromosome']} strand={gene_info['strand']} \"\n", + " f\"gene_start_site={gene_start_site}\"\n", + ")\n", + "print(f'Number of CREs the model sees: {len(cre_df)}')\n", + "cre_df.head()" + ] + }, + { + "cell_type": "markdown", + "id": "shape-md", + "metadata": {}, + "source": [ + "## 6. Inspect the captured matrices\n", + "\n", + "For `Seq2GenePredictorCombinedModulator` with `gene_pooling='multi_registry'`,\n", + "the gene cross-attention matrix has shape `(num_query_tokens, num_cre_tokens)`\n", + "where `num_query_tokens = 1 (gene-tissue CLS) + num_gene_windows` and\n", + "`num_cre_tokens` equals the number of CREs surrounding the gene.\n", + "\n", + "Because we passed `keep_heads=True` the callback retains the **head**\n", + "dimension, so each per-batch entry has shape\n", + "`(num_heads, num_query_tokens, num_cre_tokens)`. We average over heads here\n", + "to make the per-layer view; the per-head heatmap below uses the raw matrix\n", + "directly." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "shape-code", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-03T22:07:09.870568Z", + "iopub.status.busy": "2026-05-03T22:07:09.870450Z", + "iopub.status.idle": "2026-05-03T22:07:09.875427Z", + "shell.execute_reply": "2026-05-03T22:07:09.874965Z" + } + }, + "outputs": [], + "source": [ + "# Stack per-layer attention into a dict[layer_id] -> tensor[H, Q, K]\n", + "per_head_attention = {} # raw, per-head\n", + "gene_tissue_cls_attention = {} # head-averaged gene-tissue CLS row vs. CRE\n", + "full_attention = {} # head-averaged (Q, K), all gene tokens\n", + "K_attn = None # number of CRE keys produced by the model\n", + "for name, mats in log_attn.attention_matrices.items():\n", + " if not name.startswith('gene_modulator_crossMHA_'):\n", + " continue\n", + " layer_id = int(name.rsplit('_', 1)[-1])\n", + " A_heads = mats[0].cpu().to(torch.float32).numpy() # (H, Q, K)\n", + " per_head_attention[layer_id] = A_heads\n", + " A = A_heads.mean(axis=0) # head-average: (Q, K)\n", + " full_attention[layer_id] = A\n", + " # The gene-tissue CLS token is the first query token.\n", + " gene_tissue_cls_attention[layer_id] = A[0]\n", + " K_attn = A.shape[1]\n", + "\n", + "# `cre_df` already matches the model's CRE order 1-to-1 (drops included),\n", + "# and with a single-sample batch the attention K-axis is exactly the CRE\n", + "# token count — no batch-padding columns to strip. Assert this invariant\n", + "# so any future regression (e.g. running with batch_size > 1 across genes\n", + "# of differing CRE counts) fails loudly instead of silently misaligning.\n", + "assert K_attn == len(cre_df), (\n", + " f'attention has {K_attn} keys but dataset returned {len(cre_df)} CREs; '\n", + " f'alignment is broken.'\n", + ")\n", + "x_kb = cre_df['distance_to_gene_start'].values / 1000.0 # used by the plots below\n", + "\n", + "ordered_layers = sorted(gene_tissue_cls_attention)\n", + "print('Cross-attention shapes:')\n", + "for layer_id in ordered_layers:\n", + " print(\n", + " f' layer {layer_id:2d}: per_head={per_head_attention[layer_id].shape} '\n", + " f'full={full_attention[layer_id].shape} '\n", + " f'gene_tissue_cls_row={gene_tissue_cls_attention[layer_id].shape}'\n", + " )" + ] + }, + { + "cell_type": "markdown", + "id": "plot1-md", + "metadata": {}, + "source": [ + "## 7. Plot gene-tissue CLS → CRE attention vs. distance to gene start site\n", + "\n", + "Each panel shows, for one layer, how strongly the gene-tissue CLS token\n", + "attends to each surrounding CRE as a function of (signed) distance from\n", + "the gene start site. Bars increasing near 0 indicate the model relies more\n", + "on promoter-proximal CREs; bars at large distances indicate distal\n", + "enhancer-like contributions." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "plot1-code", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-03T22:07:09.884977Z", + "iopub.status.busy": "2026-05-03T22:07:09.884856Z", + "iopub.status.idle": "2026-05-03T22:07:13.197798Z", + "shell.execute_reply": "2026-05-03T22:07:13.197276Z" + } + }, + "outputs": [], + "source": [ + "fig, axes = plt.subplots(\n", + " len(ordered_layers), 1, figsize=(11, 2.0 * len(ordered_layers)), sharex=True\n", + ")\n", + "if len(ordered_layers) == 1:\n", + " axes = [axes]\n", + "\n", + "x_kb = cre_df['distance_to_gene_start'].values / 1000.0\n", + "for ax, layer_id in zip(axes, ordered_layers):\n", + " weights = gene_tissue_cls_attention[layer_id]\n", + " ax.bar(x_kb, weights, width=2.0, color='steelblue', edgecolor='none')\n", + " ax.axvline(0, color='red', lw=0.8, ls='--', alpha=0.7, label='gene start')\n", + " ax.set_ylabel(f'layer {layer_id}\\nattention')\n", + " ax.spines['top'].set_visible(False)\n", + " ax.spines['right'].set_visible(False)\n", + "axes[-1].set_xlabel('Distance from gene start site (kb, gene-strand orientation)')\n", + "axes[0].set_title(\n", + " f'Gene-tissue CLS → CRE cross-attention for {GENE_ID} in {TISSUE!r}'\n", + ")\n", + "fig.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "plot2-md", + "metadata": {}, + "source": [ + "## 8. Top attended CREs (last layer)\n", + "\n", + "Surfacing the actual cCRE names ranked by gene-tissue CLS attention weight is\n", + "often the\n", + "most directly interpretable view: each row corresponds to one ENCODE candidate\n", + "regulatory element." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "plot2-code", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-03T22:07:13.199027Z", + "iopub.status.busy": "2026-05-03T22:07:13.198810Z", + "iopub.status.idle": "2026-05-03T22:07:13.205738Z", + "shell.execute_reply": "2026-05-03T22:07:13.205294Z" + } + }, + "outputs": [], + "source": [ + "last_layer = ordered_layers[-1]\n", + "topk = 15\n", + "ranked = cre_df.copy()\n", + "ranked['attention'] = gene_tissue_cls_attention[last_layer]\n", + "ranked = ranked.sort_values('attention', ascending=False).head(topk).reset_index(drop=True)\n", + "ranked[['cre_name', 'chromosome', 'start_cre', 'end_cre', 'distance_to_gene_start', 'attention']]" + ] + }, + { + "cell_type": "markdown", + "id": "plot3-md", + "metadata": {}, + "source": [ + "## 9. Full attention heatmap (gene tokens × CREs, last layer)\n", + "\n", + "Beyond the gene-tissue CLS row, the cross-attention matrix has a row per\n", + "gene window\n", + "token. Plotting the full matrix shows whether different gene windows pick\n", + "out different CREs or share roughly the same set." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "plot3-code", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-03T22:07:13.206855Z", + "iopub.status.busy": "2026-05-03T22:07:13.206650Z", + "iopub.status.idle": "2026-05-03T22:07:13.312595Z", + "shell.execute_reply": "2026-05-03T22:07:13.312153Z" + } + }, + "outputs": [], + "source": [ + "A = full_attention[last_layer][:, : len(cre_df)]\n", + "fig, ax = plt.subplots(figsize=(12, max(2.5, 0.4 * A.shape[0])))\n", + "im = ax.imshow(\n", + " A,\n", + " aspect='auto',\n", + " cmap='magma',\n", + " extent=[x_kb.min(), x_kb.max(), A.shape[0] - 0.5, -0.5],\n", + " interpolation='nearest',\n", + ")\n", + "ax.axvline(0, color='cyan', lw=0.7, ls='--')\n", + "ax.set_yticks(np.arange(A.shape[0]))\n", + "y_labels = ['gene-tissue CLS'] + [f'gene window {i}' for i in range(1, A.shape[0])]\n", + "ax.set_yticklabels(y_labels)\n", + "ax.set_xlabel('Distance from gene start site (kb)')\n", + "ax.set_title(f'Cross-attention layer {last_layer} for {GENE_ID} — {TISSUE!r}')\n", + "fig.colorbar(im, ax=ax, label='attention weight')\n", + "fig.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "plot4-md", + "metadata": {}, + "source": [ + "## 10. Per-head attention at the gene-tissue CLS row (last layer)\n", + "\n", + "Because we recorded with `keep_heads=True`, the head-resolved matrices are\n", + "already in `per_head_attention`. No second forward pass is needed." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "plot4-code", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-03T22:07:13.313988Z", + "iopub.status.busy": "2026-05-03T22:07:13.313871Z", + "iopub.status.idle": "2026-05-03T22:07:13.449935Z", + "shell.execute_reply": "2026-05-03T22:07:13.449476Z" + } + }, + "outputs": [], + "source": [ + "head_mats = per_head_attention[ordered_layers[-1]] # (H, Q, K)\n", + "n_heads = head_mats.shape[0]\n", + "gene_tissue_cls_per_head = head_mats[:, 0, : len(cre_df)] # (H, K_real)\n", + "\n", + "fig, ax = plt.subplots(figsize=(12, max(3, 0.18 * n_heads)))\n", + "im = ax.imshow(\n", + " gene_tissue_cls_per_head,\n", + " aspect='auto',\n", + " cmap='magma',\n", + " extent=[x_kb.min(), x_kb.max(), n_heads - 0.5, -0.5],\n", + " interpolation='nearest',\n", + ")\n", + "ax.axvline(0, color='cyan', lw=0.7, ls='--')\n", + "ax.set_xlabel('Distance from gene start site (kb)')\n", + "ax.set_ylabel('attention head')\n", + "ax.set_title(f'Per-head gene-tissue CLS → CRE attention, layer {ordered_layers[-1]}')\n", + "fig.colorbar(im, ax=ax, label='attention weight')\n", + "fig.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "ref-md", + "metadata": {}, + "source": [ + "## 11. Repeat on the reference genome (no VCF)\n", + "\n", + "Setting `vcf_path=None` makes `VCFDataset` skip `bcftools consensus` and\n", + "feed the **reference** sequence to the model. We reuse the same gene,\n", + "tissue, model, and `LayersToLog` selection, swap in a fresh\n", + "`LogAttention` callback (so the captured matrices are kept separate from\n", + "the VCF run above), and then plot the gene-tissue CLS row alongside the\n", + "VCF version at the last layer to highlight where the variants in HG00096\n", + "shift the model's attention." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ref-forward-code", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-03T22:07:13.451482Z", + "iopub.status.busy": "2026-05-03T22:07:13.451332Z", + "iopub.status.idle": "2026-05-03T22:07:16.842138Z", + "shell.execute_reply": "2026-05-03T22:07:16.841373Z" + } + }, + "outputs": [], + "source": [ + "vcf_dataset_ref, dataloader_ref = processor.create_data(None, query_df)\n", + "log_attn_ref = LogAttention(layer_ids=LAYERS_TO_LOG, keep_heads=True)\n", + "\n", + "batch_ref = next(iter(dataloader_ref))\n", + "for k, v in batch_ref.items():\n", + " if isinstance(v, torch.Tensor):\n", + " batch_ref[k] = v.cuda()\n", + " elif isinstance(v, list) and v and isinstance(v[0], torch.Tensor):\n", + " batch_ref[k] = [t.cuda() for t in v]\n", + "\n", + "with torch.no_grad(), torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16):\n", + " with log_attn_ref.record_attention(model, log_epigenetics=False, log_gene=True):\n", + " preds_ref = model.predict_step(batch_ref, 0)\n", + "\n", + "pred_ref = float(preds_ref['pred_gene_exp'][0].flatten()[0])\n", + "pred_vcf = float(pred_value.flatten()[0])\n", + "print(f'Predicted log-expression for {GENE_ID} in {TISSUE!r}:')\n", + "print(f' reference genome : {pred_ref:.3f}')\n", + "print(f' HG00096 VCF : {pred_vcf:.3f}')\n", + "print(f' delta (vcf-ref) : {pred_vcf - pred_ref:+.3f}')" + ] + }, + { + "cell_type": "markdown", + "id": "ref-align-md", + "metadata": {}, + "source": [ + "Pull the reference-genome CRE positions from the new dataset (same\n", + "alignment guarantee as in section 5) and head-average the captured\n", + "matrices, then align the K-axis exactly as we did for the VCF run." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ref-align-code", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-03T22:07:16.843531Z", + "iopub.status.busy": "2026-05-03T22:07:16.843400Z", + "iopub.status.idle": "2026-05-03T22:07:26.676578Z", + "shell.execute_reply": "2026-05-03T22:07:26.675883Z" + } + }, + "outputs": [], + "source": [ + "cre_df_ref = vcf_dataset_ref.get_cre_positions(GENE_ID)\n", + "cre_df_ref['midpoint'] = (cre_df_ref['start_cre'] + cre_df_ref['end_cre']) // 2\n", + "cre_df_ref['distance_to_gene_start'] = cre_df_ref['midpoint'] - gene_start_site\n", + "if gene_info['strand'] == '-':\n", + " cre_df_ref['distance_to_gene_start'] = -cre_df_ref['distance_to_gene_start']\n", + "\n", + "gene_tissue_cls_attention_ref = {}\n", + "for name, mats in log_attn_ref.attention_matrices.items():\n", + " if not name.startswith('gene_modulator_crossMHA_'):\n", + " continue\n", + " layer_id = int(name.rsplit('_', 1)[-1])\n", + " A_heads = mats[0].cpu().to(torch.float32).numpy() # (H, Q, K)\n", + " A = A_heads.mean(axis=0) # (Q, K)\n", + " gene_tissue_cls_attention_ref[layer_id] = A[0, : len(cre_df_ref)]\n", + "\n", + "x_kb_ref = cre_df_ref['distance_to_gene_start'].values / 1000.0\n", + "print(f'reference CREs: {len(cre_df_ref)} (VCF run had {len(cre_df)})')" + ] + }, + { + "cell_type": "markdown", + "id": "ref-plot-md", + "metadata": {}, + "source": [ + "Compare the gene-tissue CLS → CRE attention from the reference and VCF\n", + "runs at the last logged layer. Differences indicate CREs whose tokenized\n", + "sequence (and therefore the model's attention) is sensitive to the\n", + "variants carried by HG00096." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ref-plot-code", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-03T22:07:26.678033Z", + "iopub.status.busy": "2026-05-03T22:07:26.677906Z", + "iopub.status.idle": "2026-05-03T22:07:27.840710Z", + "shell.execute_reply": "2026-05-03T22:07:27.840221Z" + } + }, + "outputs": [], + "source": [ + "last_layer = ordered_layers[-1]\n", + "attn_ref = gene_tissue_cls_attention_ref[last_layer]\n", + "attn_vcf = gene_tissue_cls_attention[last_layer]\n", + "\n", + "fig, axes = plt.subplots(2, 1, figsize=(14, 5), sharex=True)\n", + "for ax, (label, values, x_axis) in zip(\n", + " axes,\n", + " [\n", + " ('reference genome', attn_ref, x_kb_ref),\n", + " ('HG00096 VCF', attn_vcf, x_kb),\n", + " ],\n", + "):\n", + " ax.bar(x_axis, values, width=4.0, color='steelblue', alpha=0.85)\n", + " ax.axvline(0, color='red', lw=0.8, ls='--', alpha=0.7, label='gene start')\n", + " ax.set_ylabel(f'{label}\\nattention')\n", + " ax.legend(loc='upper right', fontsize=8)\n", + "axes[-1].set_xlabel('Distance from gene start site (kb, gene-strand orientation)')\n", + "fig.suptitle(\n", + " f'Gene-tissue CLS → CRE attention at layer {last_layer} for {GENE_ID} — {TISSUE!r}'\n", + ")\n", + "fig.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "summary-md", + "metadata": {}, + "source": [ + "## Recap\n", + "\n", + "- `LogAttention(layer_ids=[...]).record_attention(model)` is the only thing\n", + " needed to capture attention during a regular forward pass; it flips the\n", + " layer-level `log_attn_matrix` flag on the way in and resets it on the way\n", + " out.\n", + "- For `multi_registry` pooling, the **first row** of every\n", + " `gene_modulator_crossMHA_` matrix is the gene-tissue CLS → CRE\n", + " attention used\n", + " for the gene-regulation analysis.\n", + "- The recompute is mathematically equivalent to FlashAttention's internal\n", + " computation (verified in `tests/test_attn_log_callback.py::TestManualAttentionMatchesFlashAttention`).\n", + "- To go beyond a single gene, batch multiple `(gene_id, tissues)` rows in\n", + " the query DataFrame and aggregate `log_attn.attention_matrices` across\n", + " batches — each entry is a list of one tensor per batch.\n", + "- Passing `vcf_path=None` to `processor.create_data` runs the whole\n", + " pipeline on the reference genome, which is useful as a baseline for\n", + " variant-effect comparisons (section 11)." + ] + }, + { + "cell_type": "markdown", + "id": "8f9fa96d", + "metadata": {}, + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "my-venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.10" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/seq2gene/attn_log_callback.py b/seq2gene/attn_log_callback.py new file mode 100644 index 0000000..34238c1 --- /dev/null +++ b/seq2gene/attn_log_callback.py @@ -0,0 +1,353 @@ +"""Attention-matrix logging callback for VariantFormer. + +This module provides a Lightning callback (``LogAttention``) and a context +manager (``LogAttention.record_attention``) that capture the per-layer, +per-head attention matrices from the transformer layers in VariantFormer's +``epigenetics_modulator``/``gene_modulator`` (or the equivalent +``combined_modulator.cre_layers``/``combined_modulator.gene_layers`` in the +memory-optimized model) during a forward pass. + +Internally it works by toggling ``log_attn_matrix`` on the +:class:`~seq2gene.modules.layers.FlashAttLayer` instances of interest before +the forward pass and reading the resulting ``attn_matrix`` tensor afterwards. +The capture path bypasses FlashAttention to recompute a plain +softmax(Q K^T / sqrt(d)) attention matrix (with ALiBi if enabled), so it is +intentionally heavier than the production path; only enable it for the layers +you actually want to inspect. + +Typical usage with a one-off forward pass:: + + from seq2gene.attn_log_callback import LogAttention + + log_attn = LogAttention(layer_ids=[0, 4, 8]) + with log_attn.record_attention(model, log_epigenetics=True, log_gene=True): + _ = model(*inputs) + # log_attn.attention_matrices is now a dict of + # {f"{modulator}_{mha}_{layer_idx}": [tensor_per_batch_item, ...]} +""" + +from __future__ import annotations + +import dataclasses +from contextlib import contextmanager +from typing import Iterable, Optional + +import lightning.pytorch as pl +import matplotlib.pyplot as plt +import seaborn as sns +import torch +import torch.nn as nn + +from seq2gene.modules.layers import ( + ContextFlashAttentionEncoderLayer, + ContextFlashCrossAttentionEncoderLayer, + FlashAttentionEncoderLayer, +) + + +@dataclasses.dataclass +class MatrixLogMetadata: + """Metadata used to label a captured attention matrix for plotting.""" + + modulator_name: str # one of "epigenetics_modulator" / "gene_modulator" + mha_name: str # human-readable mha name (e.g. "epigenetics_modulator_mixer_3") + layer_id: int + step_number: int + cross_attn: bool = False + batch_idx: int = -1 # -1 means "first batch" / unspecified + aggregation_op: str = "mean" # one of "mean", "sum", "max" + + +def create_heatmap( + mat, + title: str, + figsize: tuple[int, int] = (12, 10), + cmap: str = "viridis", +) -> plt.Figure: + """Render an attention matrix as a Matplotlib heatmap. + + Args: + mat: 2D ``np.ndarray`` of attention weights. + title: figure title. + figsize: figure size in inches. + cmap: matplotlib colormap. + """ + fig, ax = plt.subplots(figsize=figsize) + sns.heatmap(mat, cmap=cmap, ax=ax) + ax.set_title(title) + ax.set_xlabel("Key position") + ax.set_ylabel("Query position") + plt.tight_layout() + return fig + + +def process_attention_matrix(matrix: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + """Take the first batch element, average over heads, return on CPU.""" + if matrix is None: + return None + processed = matrix[0] # first batch element + processed = torch.mean(processed, dim=0) # average over heads + return processed.detach().cpu() + + +def process_attention_matrix_all_batches( + matrix: Optional[torch.Tensor], + keep_heads: bool = False, +) -> Optional[list[torch.Tensor]]: + """Return one CPU tensor per batch element. + + If ``keep_heads`` is False (default) the head dimension is averaged out + and each entry has shape ``(Q, K)``. If True, each entry retains its + full ``(H, Q, K)`` shape. + """ + if matrix is None: + return None + processed: list[torch.Tensor] = [] + for batch_item in matrix: # batch_item: (H, Q, K) + if keep_heads: + processed.append(batch_item.detach().cpu()) + else: + processed.append(torch.mean(batch_item, dim=0).detach().cpu()) + return processed + + +def create_heatmap_from_matrix( + mlm: MatrixLogMetadata, matrix: Optional[torch.Tensor] +) -> Optional[plt.Figure]: + """Build a labeled heatmap for one already-processed (head-averaged) matrix.""" + if matrix is None: + return None + cmap = "RdBu_r" if mlm.cross_attn else "magma" + cmat = matrix.to(torch.float32).numpy() + step_str = str(mlm.step_number).zfill(6) + fig = create_heatmap( + cmat, + title=f"{mlm.modulator_name}_{mlm.mha_name}_L{mlm.layer_id} @ step {step_str}", + figsize=(14, 12), + cmap=cmap, + ) + return fig + + +_LAYER_TYPES = ( + ContextFlashAttentionEncoderLayer, + FlashAttentionEncoderLayer, + ContextFlashCrossAttentionEncoderLayer, +) + + +class LogAttention(pl.Callback): + """Lightning callback that captures attention matrices from selected layers. + + It can be used in two modes: + + 1. **As a Lightning callback** during ``trainer.predict``: it hooks + ``on_predict_batch_start``/``on_predict_batch_end`` and captures + attention every ``freq_steps`` batches. + 2. **As a context manager** via :meth:`record_attention` for a single + manual forward pass. + + Captured matrices are accumulated on ``self.attention_matrices`` keyed by + ``f"{modulator_name}_{mha_name}_{layer_idx}"``. Use the + :func:`create_heatmap` / :func:`create_heatmap_from_matrix` helpers + (matplotlib + seaborn) if you want to render them, or plot them yourself. + + Args: + freq_steps: capture every Nth predict batch (only relevant in + callback mode). + layer_ids: layer indices (within each modulator's ``ModuleList``) to + capture. Indices outside the list are silently skipped. + keep_heads: if True, keep the per-head dimension when accumulating + matrices (each entry has shape ``(H, Q, K)``); if False (default), + average over heads (each entry has shape ``(Q, K)``). + """ + + def __init__( + self, + freq_steps: int = 100, + layer_ids: Optional[Iterable[int]] = None, + keep_heads: bool = False, + ): + super().__init__() + if layer_ids is None: + layer_ids = [1, 4, 9, 13, 17, 23, 24] + self.freq_steps = freq_steps + self.layer_ids = list(layer_ids) + self.keep_heads = keep_heads + self.attention_matrices: dict[str, list[torch.Tensor]] = {} + + # ------------------------------------------------------------------ utils + def _should_process_batch(self, batch_idx: int) -> bool: + return batch_idx != 0 and batch_idx % self.freq_steps == 0 + + def _get_modulators( + self, pl_module: pl.LightningModule + ) -> tuple[nn.Module, nn.Module]: + """Return ``(epigenetics_layers, gene_layers)`` for either model variant. + + Supports both the dual-modulator model + (``Seq2GenePredictor.epigenetics_modulator`` / + ``Seq2GenePredictor.gene_modulator``) and the memory-optimized + ``Seq2GenePredictorCombinedModulator.combined_modulator``. + """ + if hasattr(pl_module, "combined_modulator"): + return ( + pl_module.combined_modulator.cre_layers, + pl_module.combined_modulator.gene_layers, + ) + if hasattr(pl_module, "epigenetics_modulator") and hasattr( + pl_module, "gene_modulator" + ): + return ( + pl_module.epigenetics_modulator.epigenetics_modulator, + pl_module.gene_modulator.gene_modulator, + ) + raise AttributeError( + "Model must have either 'combined_modulator' or both " + "'epigenetics_modulator' and 'gene_modulator' attributes" + ) + + def _set_log_flags( + self, + layer: nn.Module, + set_to: bool, + ) -> None: + """Toggle ``log_attn_matrix`` on every FlashAttLayer inside ``layer``.""" + if isinstance(layer, ContextFlashAttentionEncoderLayer): + layer.crossMHA.log_attn_matrix = set_to + layer.mixer.log_attn_matrix = set_to + elif isinstance(layer, FlashAttentionEncoderLayer): + layer.mixer.log_attn_matrix = set_to + elif isinstance(layer, ContextFlashCrossAttentionEncoderLayer): + layer.crossMHA.log_attn_matrix = set_to + else: + raise AttributeError(f"unknown layer type: {type(layer)}") + + def _process_attention_matrices( + self, + layer: nn.Module, + layer_idx: int, + modulator_name: str, + ) -> dict[str, list[torch.Tensor]]: + """Pull captured matrices off a layer and return them keyed by mha name.""" + matrices: dict[str, list[torch.Tensor]] = {} + mhas: list[tuple[nn.Module, str]] = [] + + if isinstance( + layer, + (ContextFlashAttentionEncoderLayer, ContextFlashCrossAttentionEncoderLayer), + ) and hasattr(layer, "crossMHA"): + mhas.append((layer.crossMHA, f"{modulator_name}_crossMHA_{layer_idx}")) + if isinstance( + layer, (ContextFlashAttentionEncoderLayer, FlashAttentionEncoderLayer) + ) and hasattr(layer, "mixer"): + mhas.append((layer.mixer, f"{modulator_name}_mixer_{layer_idx}")) + + for mha, mha_name in mhas: + if not hasattr(mha, "attn_matrix") or mha.attn_matrix is None: + continue + processed_matrices = process_attention_matrix_all_batches( + mha.attn_matrix, + keep_heads=self.keep_heads, + ) + # Free GPU memory from the captured raw tensor immediately. + mha.attn_matrix = None + if processed_matrices is not None: + matrices[mha_name] = processed_matrices + + return matrices + + def _process_modulator( + self, + modulator: nn.Module, + modulator_name: str, + step_number: int, + set_flags: bool, + ) -> None: + """Walk a modulator's ``ModuleList`` and toggle / extract on selected layers.""" + for layer_idx, layer in enumerate(modulator): + if layer_idx not in self.layer_ids: + continue + self._set_log_flags(layer, set_flags) + if not set_flags: # teardown phase: now read & clear + matrices = self._process_attention_matrices( + layer, layer_idx, modulator_name + ) + if matrices: + self.attention_matrices.update(matrices) + + # ------------------------------------------------------------- public API + def reset(self) -> None: + """Clear any previously recorded matrices.""" + self.attention_matrices = {} + + @contextmanager + def record_attention( + self, + pl_module: pl.LightningModule, + step_number: int = 0, + log_epigenetics: bool = True, + log_gene: bool = True, + ): + """Context manager that records attention for a single forward pass. + + Example:: + + log_attn = LogAttention(layer_ids=[0, 4]) + with log_attn.record_attention(model): + _ = model(*inputs) + log_attn.attention_matrices # populated + """ + epigenetics_modulator, gene_modulator = self._get_modulators(pl_module) + if log_epigenetics: + self._process_modulator( + epigenetics_modulator, "epigenetics_modulator", step_number, True + ) + if log_gene: + self._process_modulator( + gene_modulator, "gene_modulator", step_number, True + ) + + try: + yield self + finally: + if log_epigenetics: + self._process_modulator( + epigenetics_modulator, + "epigenetics_modulator", + step_number, + False, + ) + if log_gene: + self._process_modulator( + gene_modulator, "gene_modulator", step_number, False + ) + + # --------------------------------------------------------- lightning hooks + def on_predict_batch_start(self, trainer, pl_module, batch, batch_idx): # noqa: D401 + if not self._should_process_batch(batch_idx): + return + epigenetics_modulator, gene_modulator = self._get_modulators(pl_module) + self._process_modulator( + epigenetics_modulator, + "epigenetics_modulator", + trainer.global_step, + True, + ) + self._process_modulator( + gene_modulator, "gene_modulator", trainer.global_step, True + ) + + def on_predict_batch_end(self, trainer, pl_module, outputs, batch, batch_idx): # noqa: D401 + if not self._should_process_batch(batch_idx): + return + epigenetics_modulator, gene_modulator = self._get_modulators(pl_module) + self._process_modulator( + epigenetics_modulator, + "epigenetics_modulator", + trainer.global_step, + False, + ) + self._process_modulator( + gene_modulator, "gene_modulator", trainer.global_step, False + ) diff --git a/seq2gene/modules/layers.py b/seq2gene/modules/layers.py index c704948..b909a44 100644 --- a/seq2gene/modules/layers.py +++ b/seq2gene/modules/layers.py @@ -9,6 +9,7 @@ unpad_input, ) from typing import Union +from einops import rearrange # From https://github.com/ofirpress/attention_with_linear_biases/blob/4b92f28a005ead2567abe2359f633e73e08f3833/fairseq/models/transformer.py#L742 @@ -349,6 +350,14 @@ def __init__( use_alibi=use_alibi, cross_attn=cross_attn, ) + self.use_alibi = use_alibi + self.causal = False + # When True, the layer recomputes the full attention matrix during + # forward and stores it on `self.attn_matrix`. Off by default to keep + # the FlashAttention fast path unaffected; the LogAttention callback + # toggles it on for selected layers around the forward pass. + self.log_attn_matrix = False + self.attn_matrix = None def forward( self, @@ -485,8 +494,205 @@ def forward( else: x = self.MHA(src) + + if self.log_attn_matrix: + self.attn_matrix = self._compute_logged_attn_matrix( + src=src, + cntx=cntx, + unpad_info=unpad_info, + context_unpad_info=context_unpad_info, + ) return x + def _compute_logged_attn_matrix( + self, + src: torch.Tensor, + cntx: torch.Tensor | None, + unpad_info: dict | None, + context_unpad_info: dict | None, + ) -> torch.Tensor: + """Re-pad the unpadded inputs and compute the full attention matrix. + + This is only invoked when ``log_attn_matrix`` is True. It mirrors how + the model is actually run in this codebase: the modulators always pass + unpadded tensors through with companion ``unpad_info``/``context_unpad_info`` + dicts. + """ + assert ( + unpad_info is not None + ), "log_attn_matrix=True requires unpad_info; logging is only supported on the unpadded code path" + if self.cross_attn: + assert ( + context_unpad_info is not None + ), "log_attn_matrix=True with cross_attn=True requires context_unpad_info" + + max_seqlen_q = unpad_info["max_seqlen"] + max_seqlen_k = ( + context_unpad_info["max_seqlen"] + if context_unpad_info is not None + else unpad_info["max_seqlen"] + ) + + # Reconstruct the padded src and its boolean key-padding mask + src_mask = torch.ones(src.shape[0], src.shape[1], device=src.device) + original_src_key_padding_mask = ~pad_input( + src_mask, + unpad_info["indices"], + unpad_info["batch"], + unpad_info["seqlen"], + ).bool()[:, :, 0] + src_padded = pad_input( + src, unpad_info["indices"], unpad_info["batch"], unpad_info["seqlen"] + ) + + if self.cross_attn: + context_mask = torch.ones( + cntx.shape[0], cntx.shape[1], device=cntx.device + ) + original_context_key_padding_mask = ~pad_input( + context_mask, + context_unpad_info["indices"], + context_unpad_info["batch"], + context_unpad_info["seqlen"], + ).bool()[:, :, 0] + cntx_padded = pad_input( + cntx, + context_unpad_info["indices"], + context_unpad_info["batch"], + context_unpad_info["seqlen"], + ) + else: + cntx_padded = None + original_context_key_padding_mask = None + + return self.calculate_attention_matrix( + src_padded, + cntx_padded, + original_src_key_padding_mask, + original_context_key_padding_mask, + max_seqlen_q, + max_seqlen_k, + ) + + def calculate_attention_matrix( + self, + src: torch.Tensor, + cntx: torch.Tensor | None = None, + src_key_padding_mask: torch.Tensor | None = None, + context_key_padding_mask: torch.Tensor | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + cleanup: bool = True, + ) -> torch.Tensor: + """Compute the full attention matrix (with ALiBi) for self/cross attention. + + Args: + src: Source sequence, shape ``[B, S_q, D]``. + cntx: Context sequence for cross attention, shape ``[B, S_k, D]``. + For self-attention this is ignored and ``src`` is used. + src_key_padding_mask: Boolean padding mask for the query, shape ``[B, S_q]`` + (True = padded position). + context_key_padding_mask: Boolean padding mask for the key/context, + shape ``[B, S_k]`` (True = padded position). + max_seqlen_q / max_seqlen_k: maximum unpadded query/key sequence + lengths in the batch. Used to align ALiBi distances when + Q and K have different lengths. + cleanup: whether to free temporary tensors and call + ``torch.cuda.empty_cache()`` after the computation. + + Returns: + Attention probabilities, shape ``[B, H, S_q, S_k]``. For self-attention + ``S_q == S_k``. + """ + if self.cross_attn: + assert cntx is not None, "Context tensor must be provided for cross attention" + else: + cntx = src + + delta = ( + max_seqlen_k - max_seqlen_q + if max_seqlen_q is not None and max_seqlen_k is not None + else 0 + ) + + if self.cross_attn: + Q = self.MHA.Wq(src) + KV = self.MHA.Wkv(cntx) + Q = rearrange(Q, "b s (h d) -> b s h d", d=self.MHA.head_dim) + KV = rearrange( + KV, "b s (two h d) -> two b s h d", two=2, d=self.MHA.head_dim + ) + K, V = KV + else: + QKV = self.MHA.Wqkv(src) + QKV = rearrange( + QKV, "b s (three h d) -> three b s h d", three=3, d=self.MHA.head_dim + ) + Q, K, V = QKV + + B, S_q, H, D = Q.shape + S_k = K.shape[1] + + Q = rearrange(Q, "b s h d -> (b h) s d") + K = rearrange(K, "b s h d -> (b h) s d") + V = rearrange(V, "b s h d -> (b h) s d") + + softmax_scale = 1.0 / math.sqrt(D) + logits = torch.einsum("btd,bsd->bts", Q, K * softmax_scale) + + if self.use_alibi: + alibi_slopes = ( + get_alibi_slopes(H).to(src.device).repeat(B).to(src.dtype) + ) + alibi_slopes = alibi_slopes.unsqueeze(1).unsqueeze(2) + + q_idx = torch.arange(S_q, dtype=src.dtype, device=src.device).unsqueeze(1) + k_idx = torch.arange(S_k, dtype=src.dtype, device=src.device).unsqueeze(0) + distance = q_idx + delta - k_idx + + bias = alibi_slopes * torch.abs(distance) + logits = logits - bias + del distance, bias, q_idx, k_idx, alibi_slopes + + if self.cross_attn: + key_padding_mask = context_key_padding_mask + else: + key_padding_mask = src_key_padding_mask + + if key_padding_mask is not None: + logits = rearrange(logits, "(b h) sq sk -> b h sq sk", b=B, h=H) + mask = key_padding_mask.unsqueeze(1).unsqueeze(2) # (b, 1, 1, sk) + logits.masked_fill_(mask, float("-inf")) + logits = rearrange(logits, "b h sq sk -> (b h) sq sk") + + if self.causal: + causal_mask = torch.triu( + torch.ones(S_q, S_k, dtype=logits.dtype, device=logits.device), + diagonal=1, + ) + logits = logits.masked_fill(causal_mask.bool(), float("-inf")) + del causal_mask + + attn = torch.softmax(logits, dim=-1) + attn = attn.nan_to_num(0) + attn = rearrange(attn, "(b h) sq sk -> b h sq sk", b=B, h=H, sq=S_q, sk=S_k) + + if src_key_padding_mask is not None: + attn_mask = src_key_padding_mask.unsqueeze(1).unsqueeze(3) # (B, 1, S_q, 1) + attn.masked_fill_(attn_mask, 0.0) + + if cleanup: + del logits + del Q, K, V + if self.cross_attn: + del KV + else: + del QKV + if torch.cuda.is_available(): + torch.cuda.empty_cache() + attn = attn.detach() + return attn + class StartToken(nn.Module): def __init__(self, emb_dim): diff --git a/tests/test_attn_log_callback.py b/tests/test_attn_log_callback.py new file mode 100644 index 0000000..3d1af98 --- /dev/null +++ b/tests/test_attn_log_callback.py @@ -0,0 +1,669 @@ +"""Unit tests for the attention-logging callback and layer hooks. + +These tests run end-to-end on a tiny model with random weights so they do not +depend on the downloaded ``_artifacts`` (model checkpoints, reference genome, +or VCFs). They do require a CUDA GPU because FlashAttention only runs on GPU. +""" + +import unittest +from typing import Optional + +import numpy as np +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F +import lightning.pytorch as pl +from einops import rearrange + +from seq2gene.attn_log_callback import ( + LogAttention, + MatrixLogMetadata, + create_heatmap, + process_attention_matrix, + process_attention_matrix_all_batches, +) +from seq2gene.modules.layers import ( + ContextFlashAttentionEncoderLayer, + ContextFlashCrossAttentionEncoderLayer, + EpigeneticsModulator, + FlashAttentionEncoderLayer, + FlashAttLayer, + GeneModulator, + get_alibi_slopes, +) +from seq2gene.model_combined_modulator import CombinedModulator + + +CUDA_AVAILABLE = torch.cuda.is_available() +SKIP_CUDA_REASON = "FlashAttention requires CUDA" + + +# --------------------------------------------------------------------------- +# Tiny model harness + + +class _TinyModel(pl.LightningModule): + """Minimal LightningModule with the modulator attributes the callback + auto-detects. + + It exposes both an ``epigenetics_modulator`` (an ``EpigeneticsModulator`` + wrapping a ``ModuleList`` of self-attention encoder layers) and a + ``gene_modulator`` (a ``GeneModulator`` wrapping a ``ModuleList`` of + cross-attention encoder layers), matching the structure the callback + looks for in the production model. + """ + + def __init__( + self, + emb_dim: int = 64, + num_heads: int = 4, + num_layers: int = 3, + use_alibi: bool = True, + ): + super().__init__() + self.epigenetics_modulator = EpigeneticsModulator( + emb_dim=emb_dim, + num_heads=num_heads, + num_layers=num_layers, + use_alibi=use_alibi, + mlp_dout=0.0, + use_context=False, + ) + self.gene_modulator = GeneModulator( + emb_dim=emb_dim, + num_heads=num_heads, + num_layers=num_layers, + use_alibi=use_alibi, + mlp_dout=0.0, + only_cross_attention=True, + use_res=False, + cross_alibi=False, + ) + + def forward(self, cre, gene, cre_mask, gene_mask, precision=None): + modulator_outputs = self.epigenetics_modulator( + cre, + context=None, + src_key_padding_mask=cre_mask, + precision=precision, + keep_intermediates_unpadded=True, + ) + g = self.gene_modulator( + gene, + modulator_outputs, + res=None, + padding_mask=cre_mask, + src_key_padding_mask=gene_mask, + precision=precision, + ) + return g + + +def _make_padding_mask(batch: int, seqlen: int, n_unpadded: int) -> torch.Tensor: + """Build a [batch, seqlen] mask with True = padded.""" + mask = torch.ones(batch, seqlen, dtype=torch.bool) + mask[:, :n_unpadded] = False + return mask + + +def _build_inputs( + batch: int = 2, + cre_seq: int = 16, + gene_seq: int = 8, + cre_unpadded: int = 12, + gene_unpadded: int = 6, + emb_dim: int = 64, + device: str = "cuda", + dtype: torch.dtype = torch.float32, +): + """Build random inputs. + + The model weights remain in fp32; the ``precision`` argument passed to + the modulator's forward controls when tensors are cast to fp16 for the + FlashAttention call. So inputs default to fp32 here. + """ + cre = torch.randn(batch, cre_seq, emb_dim, device=device, dtype=dtype) + gene = torch.randn(batch, gene_seq, emb_dim, device=device, dtype=dtype) + cre_mask = _make_padding_mask(batch, cre_seq, cre_unpadded).to(device) + gene_mask = _make_padding_mask(batch, gene_seq, gene_unpadded).to(device) + return cre, gene, cre_mask, gene_mask + + +# --------------------------------------------------------------------------- +# Pure CPU tests (no GPU needed) + + +class TestLayerAttributes(unittest.TestCase): + def test_flash_att_layer_has_log_flags(self): + layer = FlashAttLayer(d_model=64, nhead=4, use_alibi=True) + self.assertFalse(layer.log_attn_matrix) + self.assertIsNone(layer.attn_matrix) + self.assertTrue(hasattr(layer, "calculate_attention_matrix")) + self.assertTrue(layer.use_alibi) + self.assertFalse(layer.causal) + + def test_flash_att_layer_default_use_alibi_false(self): + layer = FlashAttLayer(d_model=64, nhead=4, use_alibi=False) + self.assertFalse(layer.use_alibi) + + +class TestProcessAttentionHelpers(unittest.TestCase): + def test_process_attention_matrix_none(self): + self.assertIsNone(process_attention_matrix(None)) + self.assertIsNone(process_attention_matrix_all_batches(None)) + + def test_process_attention_matrix_first_batch_head_avg(self): + # [B=2, H=3, S_q=4, S_k=5] + mat = torch.randn(2, 3, 4, 5) + out = process_attention_matrix(mat) + self.assertEqual(out.shape, (4, 5)) + # Should equal the head-averaged first batch element + torch.testing.assert_close(out, mat[0].mean(dim=0)) + + def test_process_attention_matrix_all_batches(self): + mat = torch.randn(2, 3, 4, 5) + out = process_attention_matrix_all_batches(mat) + self.assertEqual(len(out), 2) + for i, item in enumerate(out): + self.assertEqual(item.shape, (4, 5)) + torch.testing.assert_close(item, mat[i].mean(dim=0)) + + def test_process_attention_matrix_all_batches_keep_heads(self): + """When ``keep_heads=True`` we must keep the (H, Q, K) shape.""" + mat = torch.randn(2, 3, 4, 5) + out = process_attention_matrix_all_batches(mat, keep_heads=True) + self.assertEqual(len(out), 2) + for i, item in enumerate(out): + self.assertEqual(item.shape, (3, 4, 5)) + torch.testing.assert_close(item, mat[i]) + + +class TestMatrixLogMetadata(unittest.TestCase): + def test_defaults(self): + m = MatrixLogMetadata( + modulator_name="epigenetics_modulator", + mha_name="epigenetics_modulator_mixer_3", + layer_id=3, + step_number=10, + ) + self.assertFalse(m.cross_attn) + self.assertEqual(m.batch_idx, -1) + self.assertEqual(m.aggregation_op, "mean") + + +class TestCreateHeatmap(unittest.TestCase): + def test_create_heatmap_returns_figure(self): + import matplotlib + + matplotlib.use("Agg") + mat = np.random.rand(6, 8) + fig = create_heatmap(mat, title="t", figsize=(4, 3)) + self.assertEqual(fig.axes[0].get_xlabel(), "Key position") + self.assertEqual(fig.axes[0].get_ylabel(), "Query position") + + +class TestSetLogFlagsOnEncoderLayers(unittest.TestCase): + """The flag-setting paths are CPU-only (just attribute toggles).""" + + def setUp(self): + self.callback = LogAttention(layer_ids=[0]) + + def test_context_flash_attention_encoder_layer(self): + layer = ContextFlashAttentionEncoderLayer(d_model=64, nhead=4) + self.callback._set_log_flags(layer, True) + self.assertTrue(layer.mixer.log_attn_matrix) + self.assertTrue(layer.crossMHA.log_attn_matrix) + self.callback._set_log_flags(layer, False) + self.assertFalse(layer.mixer.log_attn_matrix) + self.assertFalse(layer.crossMHA.log_attn_matrix) + + def test_flash_attention_encoder_layer(self): + layer = FlashAttentionEncoderLayer(d_model=64, nhead=4) + self.callback._set_log_flags(layer, True) + self.assertTrue(layer.mixer.log_attn_matrix) + # Self-attention layer has no crossMHA + self.assertFalse(hasattr(layer, "crossMHA")) + + def test_context_flash_cross_attention_encoder_layer(self): + layer = ContextFlashCrossAttentionEncoderLayer(d_model=64, nhead=4) + self.callback._set_log_flags(layer, True) + self.assertTrue(layer.crossMHA.log_attn_matrix) + + def test_unknown_layer_raises(self): + with self.assertRaises(AttributeError): + self.callback._set_log_flags(nn.Linear(2, 2), True) + + +class TestGetModulators(unittest.TestCase): + def test_dual_modulator_module(self): + callback = LogAttention() + + class _Dual(pl.LightningModule): + def __init__(self): + super().__init__() + self.epigenetics_modulator = EpigeneticsModulator( + emb_dim=32, num_heads=2, num_layers=2, + use_alibi=False, mlp_dout=0.0, use_context=False, + ) + self.gene_modulator = GeneModulator( + emb_dim=32, num_heads=2, num_layers=2, + use_alibi=False, mlp_dout=0.0, only_cross_attention=True, + ) + + epi, gene = callback._get_modulators(_Dual()) + self.assertIsInstance(epi, nn.ModuleList) + self.assertIsInstance(gene, nn.ModuleList) + + def test_combined_modulator_module(self): + callback = LogAttention() + + class _Combined(pl.LightningModule): + def __init__(self): + super().__init__() + self.combined_modulator = CombinedModulator( + emb_dim=32, num_heads=2, num_layers=2, + use_alibi=False, mlp_dout=0.0, use_context=False, + only_cross_attention=True, + ) + + epi, gene = callback._get_modulators(_Combined()) + self.assertIsInstance(epi, nn.ModuleList) + self.assertIsInstance(gene, nn.ModuleList) + + def test_missing_attributes_raises(self): + callback = LogAttention() + + class _Bad(pl.LightningModule): + pass + + with self.assertRaises(AttributeError): + callback._get_modulators(_Bad()) + + +# --------------------------------------------------------------------------- +# GPU-required tests: actually run a forward pass and verify capture. + + +@pytest.mark.skipif(not CUDA_AVAILABLE, reason=SKIP_CUDA_REASON) +class TestRecordAttentionEndToEnd(unittest.TestCase): + @classmethod + def setUpClass(cls): + torch.manual_seed(0) + cls.emb_dim = 64 + cls.num_heads = 4 + cls.num_layers = 3 + cls.model = _TinyModel( + emb_dim=cls.emb_dim, + num_heads=cls.num_heads, + num_layers=cls.num_layers, + use_alibi=True, + ).to("cuda").eval() + + def _forward(self): + cre, gene, cre_mask, gene_mask = _build_inputs( + batch=2, + cre_seq=16, + gene_seq=8, + cre_unpadded=12, + gene_unpadded=6, + emb_dim=self.emb_dim, + device="cuda", + dtype=torch.float32, + ) + with torch.no_grad(): + _ = self.model(cre, gene, cre_mask, gene_mask, precision=torch.float16) + + def test_record_attention_populates_dict(self): + log_attn = LogAttention(layer_ids=[0, 1]) + with log_attn.record_attention(self.model): + self._forward() + + self.assertGreater(len(log_attn.attention_matrices), 0) + # We expect one self-attn entry per epigenetics layer (mixer) and one + # cross-attn entry per gene layer (crossMHA), for layers 0 and 1. + expected_epi_keys = { + "epigenetics_modulator_mixer_0", + "epigenetics_modulator_mixer_1", + } + expected_gene_keys = { + "gene_modulator_crossMHA_0", + "gene_modulator_crossMHA_1", + } + for k in expected_epi_keys | expected_gene_keys: + self.assertIn(k, log_attn.attention_matrices, f"missing {k}") + mats = log_attn.attention_matrices[k] + self.assertEqual(len(mats), 2) # batch size = 2 + for m in mats: + self.assertEqual(m.dim(), 2) # head-averaged, so 2D + # softmax rows on unpadded query positions sum to ~1 + # (head-averaged so this is approximate but still close) + row_sums = m.sum(dim=-1) + # The first 6 (gene) or 12 (epi) rows should sum to ~1. + non_zero_rows = row_sums[row_sums > 0] + self.assertTrue( + torch.all(torch.abs(non_zero_rows - 1.0) < 0.05), + f"row sums for {k} were not ~1: {non_zero_rows}", + ) + + def test_record_attention_skips_layers_outside_layer_ids(self): + log_attn = LogAttention(layer_ids=[0]) + with log_attn.record_attention(self.model): + self._forward() + keys = list(log_attn.attention_matrices.keys()) + for k in keys: + self.assertTrue(k.endswith("_0"), f"unexpected key {k}") + + def test_log_flags_are_reset_after_context_exit(self): + log_attn = LogAttention(layer_ids=[0, 1]) + with log_attn.record_attention(self.model): + self._forward() + # After exit, no FlashAttLayer should still have log_attn_matrix=True + for layer in self.model.epigenetics_modulator.epigenetics_modulator: + self.assertFalse(layer.mixer.log_attn_matrix) + for layer in self.model.gene_modulator.gene_modulator: + self.assertFalse(layer.crossMHA.log_attn_matrix) + + def test_recorded_matrices_are_on_cpu(self): + log_attn = LogAttention(layer_ids=[0]) + with log_attn.record_attention(self.model): + self._forward() + for mats in log_attn.attention_matrices.values(): + for m in mats: + self.assertEqual(m.device.type, "cpu") + + def test_record_attention_keep_heads_preserves_head_dimension(self): + """``keep_heads=True`` must keep ``(H, Q, K)`` per-batch entries.""" + log_attn = LogAttention( + layer_ids=[0], keep_heads=True + ) + with log_attn.record_attention(self.model): + self._forward() + + self.assertGreater(len(log_attn.attention_matrices), 0) + for k, mats in log_attn.attention_matrices.items(): + for m in mats: + self.assertEqual(m.dim(), 3, f"{k}: expected (H, Q, K) tensor") + self.assertEqual(m.shape[0], self.num_heads) + self.assertEqual(m.device.type, "cpu") + + def test_record_attention_default_averages_heads(self): + """Default behaviour (``keep_heads=False``) must collapse heads.""" + log_attn = LogAttention(layer_ids=[0]) + with log_attn.record_attention(self.model): + self._forward() + for mats in log_attn.attention_matrices.values(): + for m in mats: + self.assertEqual(m.dim(), 2) # (Q, K), heads averaged out + + def test_record_attention_with_only_epigenetics(self): + log_attn = LogAttention(layer_ids=[0]) + with log_attn.record_attention(self.model, log_epigenetics=True, log_gene=False): + self._forward() + for k in log_attn.attention_matrices: + self.assertTrue(k.startswith("epigenetics_modulator")) + + def test_disabling_log_does_not_change_forward_output(self): + """Sanity: log_attn_matrix=False should leave the forward output + identical, since the recompute happens on a side path.""" + cre, gene, cre_mask, gene_mask = _build_inputs( + batch=2, emb_dim=self.emb_dim, device="cuda", dtype=torch.float32 + ) + with torch.no_grad(): + out_no_log = self.model(cre, gene, cre_mask, gene_mask, precision=torch.float16) + + log_attn = LogAttention(layer_ids=[0, 1]) + with log_attn.record_attention(self.model): + out_with_log = self.model( + cre, gene, cre_mask, gene_mask, precision=torch.float16 + ) + + # Forward outputs must match exactly: log path is read-only relative + # to the main computation. + torch.testing.assert_close(out_no_log, out_with_log) + + +@pytest.mark.skipif(not CUDA_AVAILABLE, reason=SKIP_CUDA_REASON) +class TestCalculateAttentionMatrixDirect(unittest.TestCase): + """Test ``FlashAttLayer.calculate_attention_matrix`` directly. + + These run in fp32 to keep the softmax row-sum tolerances tight; the + method does not depend on FlashAttention so it works on either dtype. + """ + + def test_self_attention_matrix_shape_and_softmax(self): + torch.manual_seed(0) + d_model, nhead, B, S = 64, 4, 2, 8 + layer = FlashAttLayer(d_model, nhead, use_alibi=True).cuda() + src = torch.randn(B, S, d_model, device="cuda") + attn = layer.calculate_attention_matrix( + src, max_seqlen_q=S, max_seqlen_k=S + ) + self.assertEqual(attn.shape, (B, nhead, S, S)) + # Each row sums to 1 (softmax over keys). + row_sums = attn.float().sum(dim=-1) + torch.testing.assert_close( + row_sums, torch.ones_like(row_sums), atol=1e-5, rtol=1e-5 + ) + + def test_cross_attention_matrix_shape(self): + torch.manual_seed(1) + d_model, nhead, B, S_q, S_k = 64, 4, 2, 8, 11 + layer = FlashAttLayer( + d_model, nhead, use_alibi=True, cross_attn=True + ).cuda() + src = torch.randn(B, S_q, d_model, device="cuda") + cntx = torch.randn(B, S_k, d_model, device="cuda") + attn = layer.calculate_attention_matrix( + src, cntx, max_seqlen_q=S_q, max_seqlen_k=S_k + ) + self.assertEqual(attn.shape, (B, nhead, S_q, S_k)) + row_sums = attn.float().sum(dim=-1) + torch.testing.assert_close( + row_sums, torch.ones_like(row_sums), atol=1e-5, rtol=1e-5 + ) + + def test_padding_mask_zeroes_attention(self): + torch.manual_seed(2) + d_model, nhead, B, S = 64, 4, 2, 8 + layer = FlashAttLayer(d_model, nhead, use_alibi=False).cuda() + src = torch.randn(B, S, d_model, device="cuda") + # Mark last 3 positions as padded for both batches. + mask = torch.zeros(B, S, dtype=torch.bool, device="cuda") + mask[:, -3:] = True + attn = layer.calculate_attention_matrix( + src, + src_key_padding_mask=mask, + context_key_padding_mask=mask, + max_seqlen_q=S - 3, + max_seqlen_k=S - 3, + ) + # Padded query rows should have all zeros. + padded_rows = attn[:, :, -3:, :] + self.assertTrue(torch.all(padded_rows == 0)) + # Padded key columns should have all zeros (because softmax + # of -inf is 0). + padded_cols = attn[:, :, :-3, -3:] + self.assertTrue(torch.all(padded_cols == 0)) + + +@pytest.mark.skipif(not CUDA_AVAILABLE, reason=SKIP_CUDA_REASON) +class TestManualAttentionMatchesFlashAttention(unittest.TestCase): + """Verify the *manually computed* attention matrix is consistent with + what FlashAttention actually does. + + The strategy: + + 1. Build a ``FlashAttLayer`` with random weights. + 2. Run ``self.MHA(src)`` – this is the real (FlashAttention) output and is + defined to be ``out_proj(attn @ V)`` where ``attn = softmax(QK^T/sqrt(d) + + alibi_bias)`` and ``Q,K,V`` come from the same ``Wqkv``/``Wq``/``Wkv`` + projections we use in ``calculate_attention_matrix``. + 3. Recompute QKV ourselves from the layer's projection modules, take our + captured ``attn`` matrix, multiply by ``V``, reshape, and apply the + same ``out_proj``. + 4. Assert the two outputs match within a tolerance dominated by fp16 + precision. + + If steps 1 and 3 produce the same tensor, then the attention matrix our + callback exposes really is the per-head attention probability matrix + used by FlashAttention. This is the key correctness property the + callback advertises. + """ + + def _qkv_from_layer( + self, layer: FlashAttLayer, src: torch.Tensor, cntx: Optional[torch.Tensor] + ): + """Replicate flash_attn's QKV projection so our V matches theirs.""" + if layer.cross_attn: + assert cntx is not None + Q = layer.MHA.Wq(src) + KV = layer.MHA.Wkv(cntx) + Q = rearrange(Q, "b s (h d) -> b s h d", d=layer.MHA.head_dim) + KV = rearrange( + KV, "b s (two h d) -> two b s h d", two=2, d=layer.MHA.head_dim + ) + K, V = KV + else: + QKV = layer.MHA.Wqkv(src) + QKV = rearrange( + QKV, "b s (three h d) -> three b s h d", three=3, d=layer.MHA.head_dim + ) + Q, K, V = QKV + return Q, K, V # each [B, S_*, H, D] + + def _reconstruct_mha_output( + self, + layer: FlashAttLayer, + attn: torch.Tensor, # [B, H, S_q, S_k] + V: torch.Tensor, # [B, S_k, H, D] + ) -> torch.Tensor: + """Apply attn @ V then the layer's out_proj to mimic MHA's full forward.""" + # attn @ V: [B, H, S_q, S_k] @ [B, H, S_k, D] -> [B, H, S_q, D] + V_h = rearrange(V, "b s h d -> b h s d") + out_h = torch.einsum("bhqk,bhkd->bhqd", attn, V_h) # [B, H, S_q, D] + # FlashAttention's MHA output is shape [B, S_q, D_model] after out_proj + out = rearrange(out_h, "b h s d -> b s (h d)") + out = layer.MHA.out_proj(out) + return out + + @staticmethod + def _build_layer(use_alibi: bool, cross_attn: bool, d_model=64, nhead=4): + # dropout=0.0 + .eval() so dropout never fires; with fp16 + dropout + # active, FlashAttention's randomness would make the comparison + # non-deterministic. This mirrors how the layer is used at inference. + layer = ( + FlashAttLayer( + d_model, + nhead, + dropout=0.0, + use_alibi=use_alibi, + cross_attn=cross_attn, + ) + .cuda() + .to(torch.float16) + ) + layer.eval() + return layer + + def test_self_attention_no_alibi_matches_real_mha_output(self): + """Self-attention without ALiBi: ``out_proj(my_attn @ V)`` must equal + the real ``MHA(src)`` (which is ``out_proj(FlashAttn(Wqkv(src)))``). + """ + torch.manual_seed(0) + d_model, nhead, B, S = 64, 4, 2, 16 + layer = self._build_layer(use_alibi=False, cross_attn=False, d_model=d_model, nhead=nhead) + src = torch.randn(B, S, d_model, device="cuda", dtype=torch.float16) + + with torch.no_grad(): + real_out = layer.MHA(src) # FlashAttention output + attn = layer.calculate_attention_matrix( + src, max_seqlen_q=S, max_seqlen_k=S + ) + _, _, V = self._qkv_from_layer(layer, src, None) + manual_out = self._reconstruct_mha_output(layer, attn, V) + + torch.testing.assert_close(real_out, manual_out, atol=1e-2, rtol=1e-2) + + def test_self_attention_with_alibi_matches_real_mha_output(self): + """Same equivalence with ALiBi turned on, exercising the ALiBi branch + of ``calculate_attention_matrix``.""" + torch.manual_seed(1) + d_model, nhead, B, S = 64, 4, 2, 16 + layer = self._build_layer(use_alibi=True, cross_attn=False, d_model=d_model, nhead=nhead) + src = torch.randn(B, S, d_model, device="cuda", dtype=torch.float16) + + with torch.no_grad(): + real_out = layer.MHA(src) + attn = layer.calculate_attention_matrix( + src, max_seqlen_q=S, max_seqlen_k=S + ) + _, _, V = self._qkv_from_layer(layer, src, None) + manual_out = self._reconstruct_mha_output(layer, attn, V) + + torch.testing.assert_close(real_out, manual_out, atol=1e-2, rtol=1e-2) + + def test_cross_attention_no_alibi_matches_real_mha_output(self): + """Cross-attention with different Q/K lengths.""" + torch.manual_seed(2) + d_model, nhead, B, S_q, S_k = 64, 4, 2, 8, 13 + layer = self._build_layer(use_alibi=False, cross_attn=True, d_model=d_model, nhead=nhead) + src = torch.randn(B, S_q, d_model, device="cuda", dtype=torch.float16) + cntx = torch.randn(B, S_k, d_model, device="cuda", dtype=torch.float16) + + with torch.no_grad(): + real_out = layer.MHA(src, cntx) # cross-attention output + attn = layer.calculate_attention_matrix( + src, cntx, max_seqlen_q=S_q, max_seqlen_k=S_k + ) + _, _, V = self._qkv_from_layer(layer, src, cntx) + manual_out = self._reconstruct_mha_output(layer, attn, V) + + torch.testing.assert_close(real_out, manual_out, atol=1e-2, rtol=1e-2) + + def test_cross_attention_with_alibi_matches_real_mha_output(self): + torch.manual_seed(3) + d_model, nhead, B, S_q, S_k = 64, 4, 2, 8, 13 + layer = self._build_layer(use_alibi=True, cross_attn=True, d_model=d_model, nhead=nhead) + src = torch.randn(B, S_q, d_model, device="cuda", dtype=torch.float16) + cntx = torch.randn(B, S_k, d_model, device="cuda", dtype=torch.float16) + + with torch.no_grad(): + real_out = layer.MHA(src, cntx) + attn = layer.calculate_attention_matrix( + src, cntx, max_seqlen_q=S_q, max_seqlen_k=S_k + ) + _, _, V = self._qkv_from_layer(layer, src, cntx) + manual_out = self._reconstruct_mha_output(layer, attn, V) + + torch.testing.assert_close(real_out, manual_out, atol=1e-2, rtol=1e-2) + + def test_self_attention_matches_torch_sdpa_no_alibi(self): + """Cross-check against ``torch.nn.functional.scaled_dot_product_attention`` + as an independent reference for self-attention with no ALiBi.""" + torch.manual_seed(4) + d_model, nhead, B, S = 64, 4, 2, 12 + layer = ( + FlashAttLayer(d_model, nhead, use_alibi=False, cross_attn=False) + .cuda() + ) + src = torch.randn(B, S, d_model, device="cuda") + + with torch.no_grad(): + attn = layer.calculate_attention_matrix( + src, max_seqlen_q=S, max_seqlen_k=S + ) + Q, K, V = self._qkv_from_layer(layer, src, None) + # SDPA expects [B, H, S, D] + Qh = rearrange(Q, "b s h d -> b h s d") + Kh = rearrange(K, "b s h d -> b h s d") + Vh = rearrange(V, "b s h d -> b h s d") + sdpa_out = F.scaled_dot_product_attention(Qh, Kh, Vh) # [B, H, S, D] + # attn @ V (manually) should give the same per-head outputs + manual_out = torch.einsum("bhqk,bhkd->bhqd", attn, Vh) + + torch.testing.assert_close(manual_out, sdpa_out, atol=1e-4, rtol=1e-4) + + +if __name__ == "__main__": + unittest.main()