From 72b18040a87e460b4ac7e0b51cf6d47ccd919389 Mon Sep 17 00:00:00 2001 From: arrdel Date: Tue, 10 Nov 2026 10:36:34 -0500 Subject: [PATCH] feat(bridge): trajectory-step activation caching and denoising patching for diffusion LMs --- .github/workflows/checks.yml | 1 + .../Diffusion_Trajectory_Patching_Demo.ipynb | 745 ++++++++++++++++++ docs/make_docs.py | 1 + docs/source/content/tutorials.md | 2 + docs/source/index.md | 1 + .../model_bridge/test_llada_adapter.py | 46 ++ .../test_llada_adapter.py | 24 +- .../model_bridge/test_diffusion_generate.py | 178 +++++ tests/unit/test_make_docs.py | 1 + transformer_lens/ActivationCache.py | 176 +++++ transformer_lens/__init__.py | 3 +- .../supported_architectures/llada.py | 133 +++- .../model_bridge/transformer_bridge.py | 137 +++- 13 files changed, 1436 insertions(+), 12 deletions(-) create mode 100644 demos/Diffusion_Trajectory_Patching_Demo.ipynb diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 7a5efdbad9..6ad427c535 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -432,6 +432,7 @@ jobs: - "BERT" - "Backward_Lens_Demo" - "Bridge_Evals_Demo" + - "Diffusion_Trajectory_Patching_Demo" - "Exploratory_Analysis_Demo" # - "Grokking_Demo" - "Head_Detector_Demo" diff --git a/demos/Diffusion_Trajectory_Patching_Demo.ipynb b/demos/Diffusion_Trajectory_Patching_Demo.ipynb new file mode 100644 index 0000000000..2f0b940521 --- /dev/null +++ b/demos/Diffusion_Trajectory_Patching_Demo.ipynb @@ -0,0 +1,745 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "de2107e3", + "metadata": {}, + "source": [ + "# Trajectory-Step Activation Caching and Denoising Patching for Diffusion LMs\n", + "\n", + "Autoregressive language models generate left-to-right under a causal mask, so a single post-hoc `run_with_cache(output_tokens)` pass recovers every step's activations. Masked discrete-diffusion language models (`LLaDA-8B`, `Dream-7B`, `LLaDA-2`, `GIDD`) use **bidirectional self-attention** (`is_causal=False`) and iteratively unmask a fixed canvas $x_0 \\to x_1 \\to \\dots \\to x_S$ over $S$ denoising steps:\n", + "\n", + "1. **A post-hoc pass on $x_S$ misses the trajectory.** At denoising step $s$, the model attends bidirectionally over a partially-masked canvas $x_s$ containing `[MASK]` tokens. The residual stream and attention patterns on $x_s$ differ from those on the fully unmasked sequence $x_S$.\n", + "2. **Standard caching hooks overwrite earlier steps.** Wrapping `diffusion_generate()` in `add_caching_hooks()` writes into `cache[hook_name]` on every forward pass, discarding steps $0 \\dots S-2$ and retaining only the final step $S-1$.\n", + "3. **Commitment order determines causal direction.** Under low-confidence remasking, a diffusion LM can unmask its final answer token at step $T_{\\text{commit}} \\ll S$ while most of the reasoning chain is still `[MASK]`. Once the answer position is unmasked, subsequent denoising steps cannot flip it, and bidirectional attention flows *from* the committed answer token *into* the still-masked reasoning span.\n", + "\n", + "This notebook demonstrates `TransformerBridge.diffusion_generate(..., return_cache=True)` and `DiffusionTrajectoryCache`:\n", + "- **Step-selective trajectory caching** (`cache_steps`, `names_filter`, `pos_slice`, `device`) across denoising steps without clobbering earlier steps or holding all $S$ full-model caches on GPU.\n", + "- **Activation patching across denoising steps** (`traj_cache.make_step_patch_hook`) before vs. after the answer commitment step $T_{\\text{commit}}$.\n", + "- **Trajectory commitment analysis** on `LLaDA-8B-Instruct` (`GSAI-ML/LLaDA-8B-Instruct`) and `Dream-v0-Instruct-7B` (`Dream-org/Dream-v0-Instruct-7B`) ARC-Challenge trajectories." + ] + }, + { + "cell_type": "markdown", + "id": "7ba0cea4", + "metadata": {}, + "source": [ + "## Requirements and Execution Modes\n", + "\n", + "The mechanics section runs on CPU in a few seconds using a compact two-layer `LLaDAModelLM` bridge built from `build_bridge_from_module`, exercising the exact `diffusion_generate(..., return_cache=True)` and `DiffusionTrajectoryCache.make_step_patch_hook(...)` code paths.\n", + "\n", + "The released-checkpoint section at the end runs `GSAI-ML/LLaDA-8B-Instruct` in bfloat16 and is opt-in: set `TL_RUN_RELEASED_LLADA=1` before starting the kernel (requires ~16 GiB of VRAM or system RAM and downloads the 8B checkpoint)." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "cc8567a9", + "metadata": { + "execution": { + "iopub.execute_input": "2026-11-10T15:24:40.517845Z", + "iopub.status.busy": "2026-11-10T15:24:40.517742Z", + "iopub.status.idle": "2026-11-10T15:24:43.776749Z", + "shell.execute_reply": "2026-11-10T15:24:43.776272Z" + } + }, + "outputs": [], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "import os\n", + "from dataclasses import dataclass, field\n", + "from typing import NamedTuple, Optional\n", + "\n", + "os.environ[\"HF_HUB_DISABLE_PROGRESS_BARS\"] = \"1\"\n", + "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n", + "\n", + "import pandas as pd\n", + "import torch\n", + "from torch import nn\n", + "\n", + "from transformer_lens import ActivationCache, DiffusionTrajectoryCache\n", + "from transformer_lens.model_bridge import TransformerBridge\n", + "from transformer_lens.model_bridge.sources._bridge_builder import (\n", + " build_bridge_from_module,\n", + ")\n", + "\n", + "DEVICE = \"cpu\"\n", + "RUN_RELEASED = os.environ.get(\"TL_RUN_RELEASED_LLADA\") == \"1\"\n", + "pd.set_option(\"display.max_colwidth\", None)\n", + "_ = torch.manual_seed(0)" + ] + }, + { + "cell_type": "markdown", + "id": "55175bcc", + "metadata": {}, + "source": [ + "## 1. Building a Compact `LLaDAModelLM` Bridge for CPU Trajectory Verification\n", + "\n", + "`LLaDAArchitectureAdapter` wraps the block-local projections (`q_proj`, `k_proj`, `v_proj`, `attn_out`, `ff_proj`, `up_proj`, `ff_out`) and reconstructs bidirectional self-attention (`is_causal=False`) so `blocks.{l}.attn.hook_pattern`, `blocks.{l}.attn.hook_attn_scores`, and `blocks.{l}.hook_out` fire on every denoising step. In `prepare_model`, the adapter binds LLaDA's blockwise low-confidence remasking sampler onto `diffusion_generate`." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "9ad5ac54", + "metadata": { + "execution": { + "iopub.execute_input": "2026-11-10T15:24:43.778776Z", + "iopub.status.busy": "2026-11-10T15:24:43.778528Z", + "iopub.status.idle": "2026-11-10T15:24:43.842573Z", + "shell.execute_reply": "2026-11-10T15:24:43.842043Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Adapter: LLaDAArchitectureAdapter, native_sampler='diffusion_generate'\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/home/adelechinda/home/contributions/TransformerLens/.venv/lib/python3.12/site-packages/torch/cuda/__init__.py:180: UserWarning: CUDA initialization: The NVIDIA driver on your system is too old (found version 12040). Please update your GPU driver by downloading and installing a new version from the URL: http://www.nvidia.com/Download/index.aspx Alternatively, go to: https://pytorch.org to install a PyTorch version that has been compiled with your version of the CUDA driver. (Triggered internally at /pytorch/c10/cuda/CUDAFunctions.cpp:119.)\n", + " return torch._C._cuda_getDeviceCount() > 0\n" + ] + } + ], + "source": [ + "@dataclass\n", + "class CompactLLaDAConfig:\n", + " activation_type: str = \"silu\"\n", + " alibi: bool = False\n", + " architectures: list[str] = field(default_factory=lambda: [\"LLaDAModelLM\"])\n", + " attention_dropout: float = 0.0\n", + " attention_layer_norm: bool = False\n", + " block_group_size: int = 1\n", + " block_type: str = \"llama\"\n", + " d_model: int = 32\n", + " embedding_dropout: float = 0.0\n", + " embedding_size: int = 64\n", + " eos_token_id: int = 1\n", + " include_bias: bool = False\n", + " include_qkv_bias: bool = False\n", + " input_emb_norm: bool = False\n", + " layer_norm_type: str = \"rms\"\n", + " mask_token_id: int = 63\n", + " max_sequence_length: int = 32\n", + " mlp_hidden_size: int = 48\n", + " model_type: str = \"llada\"\n", + " n_heads: int = 4\n", + " n_kv_heads: int = 2\n", + " n_layers: int = 2\n", + " output_attentions: bool = False\n", + " pad_token_id: int = 0\n", + " residual_dropout: float = 0.0\n", + " rms_norm_eps: float = 1e-5\n", + " rope: bool = True\n", + " rope_full_precision: bool = True\n", + " rope_theta: float = 10000.0\n", + " scale_logits: bool = False\n", + " use_cache: bool = False\n", + " vocab_size: int = 64\n", + " weight_tying: bool = False\n", + "\n", + "\n", + "class _LLaDAOut(NamedTuple):\n", + " logits: torch.Tensor\n", + "\n", + "\n", + "class _RMSNorm(nn.Module):\n", + " def __init__(self, cfg: CompactLLaDAConfig) -> None:\n", + " super().__init__()\n", + " self.weight = nn.Parameter(torch.ones(cfg.d_model))\n", + " self.eps = cfg.rms_norm_eps\n", + "\n", + " def forward(self, x: torch.Tensor) -> torch.Tensor:\n", + " normed = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)\n", + " return self.weight * normed.to(x.dtype)\n", + "\n", + "\n", + "class _RotaryEmbedding(nn.Module):\n", + " def __init__(self, cfg: CompactLLaDAConfig) -> None:\n", + " super().__init__()\n", + " self.cfg = cfg\n", + "\n", + " @staticmethod\n", + " def _rotate_half(x: torch.Tensor) -> torch.Tensor:\n", + " b, h, s, d = x.shape\n", + " first, second = x.view(b, h, s, 2, d // 2).unbind(dim=-2)\n", + " return torch.cat((-second, first), dim=-1)\n", + "\n", + " def forward(self, q: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:\n", + " d = q.shape[-1]\n", + " inv = 1.0 / (self.cfg.rope_theta ** (torch.arange(0, d, 2, device=q.device).float() / d))\n", + " pos = torch.arange(k.shape[-2], device=q.device).float()\n", + " freqs = torch.outer(pos, inv)\n", + " angles = torch.cat((freqs, freqs), dim=-1).unsqueeze(0).unsqueeze(0)\n", + " cos, sin = angles.cos(), angles.sin()\n", + " q_cos = cos[:, :, k.shape[-2] - q.shape[-2] : k.shape[-2], :]\n", + " q_sin = sin[:, :, k.shape[-2] - q.shape[-2] : k.shape[-2], :]\n", + " return (\n", + " (q.float() * q_cos + self._rotate_half(q.float()) * q_sin).to(q.dtype),\n", + " (k.float() * cos + self._rotate_half(k.float()) * sin).to(k.dtype),\n", + " )\n", + "\n", + "\n", + "class _LLaDABlock(nn.Module):\n", + " def __init__(self, cfg: CompactLLaDAConfig) -> None:\n", + " super().__init__()\n", + " self.config = cfg\n", + " kv_width = (cfg.d_model // cfg.n_heads) * cfg.n_kv_heads\n", + " self.attn_norm = _RMSNorm(cfg)\n", + " self.q_proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False)\n", + " self.k_proj = nn.Linear(cfg.d_model, kv_width, bias=False)\n", + " self.v_proj = nn.Linear(cfg.d_model, kv_width, bias=False)\n", + " self.attn_out = nn.Linear(cfg.d_model, cfg.d_model, bias=False)\n", + " self.rotary_emb = _RotaryEmbedding(cfg)\n", + " self.ff_norm = _RMSNorm(cfg)\n", + " self.ff_proj = nn.Linear(cfg.d_model, cfg.mlp_hidden_size, bias=False)\n", + " self.up_proj = nn.Linear(cfg.d_model, cfg.mlp_hidden_size, bias=False)\n", + " self.act = nn.SiLU()\n", + " self.ff_out = nn.Linear(cfg.mlp_hidden_size, cfg.d_model, bias=False)\n", + "\n", + " @staticmethod\n", + " def _cast_attn_bias(bias: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:\n", + " return bias.to(dtype=dtype)\n", + "\n", + " def attention(\n", + " self,\n", + " q: torch.Tensor,\n", + " k: torch.Tensor,\n", + " v: torch.Tensor,\n", + " attention_bias: Optional[torch.Tensor] = None,\n", + " layer_past: Optional[tuple[torch.Tensor, torch.Tensor]] = None,\n", + " use_cache: bool = False,\n", + " ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]:\n", + " del attention_bias, layer_past, use_cache\n", + " b, s, width = q.shape\n", + " head_dim = width // self.config.n_heads\n", + " q = q.view(b, s, self.config.n_heads, head_dim).transpose(1, 2)\n", + " k = k.view(b, s, self.config.n_kv_heads, head_dim).transpose(1, 2)\n", + " v = v.view(b, s, self.config.n_kv_heads, head_dim).transpose(1, 2)\n", + " q, k = self.rotary_emb(q, k)\n", + " groups = self.config.n_heads // self.config.n_kv_heads\n", + " k = k.repeat_interleave(groups, dim=1)\n", + " v = v.repeat_interleave(groups, dim=1)\n", + " scores = torch.matmul(q, k.transpose(-2, -1)) / (head_dim ** 0.5)\n", + " probs = torch.softmax(scores.float(), dim=-1).to(q.dtype)\n", + " ctx = torch.matmul(probs, v).transpose(1, 2).contiguous().view(b, s, width)\n", + " return self.attn_out(ctx), None\n", + "\n", + " def forward(self, x: torch.Tensor, attention_bias: Optional[torch.Tensor] = None) -> tuple[torch.Tensor, None]:\n", + " normed = self.attn_norm(x)\n", + " attn_out, _ = self.attention(\n", + " self.q_proj(normed), self.k_proj(normed), self.v_proj(normed), attention_bias=attention_bias\n", + " )\n", + " x = x + attn_out\n", + " m = self.ff_norm(x)\n", + " x = x + self.ff_out(self.act(self.ff_proj(m)) * self.up_proj(m))\n", + " return x, None\n", + "\n", + "\n", + "class CompactLLaDAModelLM(nn.Module):\n", + " def __init__(self, cfg: CompactLLaDAConfig) -> None:\n", + " super().__init__()\n", + " self.config = cfg\n", + " self.model = nn.Module()\n", + " self.model.transformer = nn.Module()\n", + " self.model.transformer.wte = nn.Embedding(cfg.vocab_size, cfg.d_model)\n", + " self.model.transformer.blocks = nn.ModuleList([_LLaDABlock(cfg) for _ in range(cfg.n_layers)])\n", + " self.model.transformer.ln_f = _RMSNorm(cfg)\n", + " self.model.transformer.ff_out = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)\n", + "\n", + " def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> _LLaDAOut:\n", + " del attention_mask\n", + " h = self.model.transformer.wte(input_ids)\n", + " for block in self.model.transformer.blocks:\n", + " h, _ = block(h)\n", + " return _LLaDAOut(logits=self.model.transformer.ff_out(self.model.transformer.ln_f(h)))\n", + "\n", + "\n", + "torch.manual_seed(7)\n", + "compact_cfg = CompactLLaDAConfig()\n", + "raw_llada = CompactLLaDAModelLM(compact_cfg).eval()\n", + "bridge = build_bridge_from_module(\n", + " raw_llada,\n", + " hf_config=compact_cfg,\n", + " architecture=\"LLaDAModelLM\",\n", + " model_name=\"compact-llada\",\n", + " device=DEVICE,\n", + " dtype=torch.float32,\n", + ")\n", + "print(f\"Adapter: {type(bridge.adapter).__name__}, native_sampler={bridge.adapter.native_sampler!r}\")" + ] + }, + { + "cell_type": "markdown", + "id": "25a90437", + "metadata": {}, + "source": [ + "## 2. Trajectory-Step Activation Caching Across Denoising Steps\n", + "\n", + "Calling `bridge.diffusion_generate(..., return_cache=True)` attaches step-aware caching hooks and a forward pre-hook on the underlying model. On each denoising step $s \\in \\{0, \\dots, S-1\\}$:\n", + "- `traj_cache.step_tokens[s]` records the partially-masked input canvas $x_s$ entering step $s$.\n", + "- `traj_cache[s]` is a full `ActivationCache` for step $s$.\n", + "- `traj_cache.stack(hook_name)` stacks the activation across all captured steps into a `[n_steps, ...]` tensor." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "536719d1", + "metadata": { + "execution": { + "iopub.execute_input": "2026-11-10T15:24:43.843802Z", + "iopub.status.busy": "2026-11-10T15:24:43.843680Z", + "iopub.status.idle": "2026-11-10T15:24:43.863023Z", + "shell.execute_reply": "2026-11-10T15:24:43.862656Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + " denoising_step continuation_canvas_entering_step masked_positions_remaining final_pos_resid_l2_norm\n", + " 0 [63, 63, 63, 63] 4 4.4148\n", + " 1 [63, 63, 63, 26] 3 5.7519\n", + " 2 [63, 63, 48, 26] 2 5.6955\n", + " 3 [63, 48, 48, 26] 1 5.6787\n" + ] + } + ], + "source": [ + "prompt_direct = torch.tensor([[4, 8, 12, 16]])\n", + "MASK_ID = compact_cfg.mask_token_id\n", + "\n", + "sequences_direct, traj_direct = bridge.diffusion_generate(\n", + " prompt_direct,\n", + " max_new_tokens=4,\n", + " return_cache=True,\n", + " names_filter=[\"blocks.0.attn.hook_pattern\", \"blocks.1.hook_out\"],\n", + " device=\"cpu\",\n", + ")\n", + "\n", + "assert isinstance(traj_direct, DiffusionTrajectoryCache)\n", + "assert traj_direct.steps == (0, 1, 2, 3)\n", + "assert traj_direct.num_steps_run == 4\n", + "\n", + "step_rows = []\n", + "for step in traj_direct.steps:\n", + " canvas = traj_direct.step_tokens[step][0].tolist()\n", + " continuation = canvas[prompt_direct.shape[-1] :]\n", + " masked_remaining = sum(1 for tok in continuation if tok == MASK_ID)\n", + " resid_norm = float(traj_direct[step][\"blocks.1.hook_out\"][0, -1].norm().item())\n", + " step_rows.append(\n", + " {\n", + " \"denoising_step\": step,\n", + " \"continuation_canvas_entering_step\": continuation,\n", + " \"masked_positions_remaining\": masked_remaining,\n", + " \"final_pos_resid_l2_norm\": round(resid_norm, 4),\n", + " }\n", + " )\n", + "\n", + "df_steps = pd.DataFrame(step_rows)\n", + "print(df_steps.to_string(index=False))" + ] + }, + { + "cell_type": "markdown", + "id": "8ff8e8d9", + "metadata": {}, + "source": [ + "### Step-Selective Caching (`cache_steps` and `pos_slice`)\n", + "\n", + "For an 8B model running $S = 256$ denoising steps, caching all layers at every step on GPU would require tens of gigabytes. Passing `cache_steps` (an integer, sequence of step indices, or predicate) and `pos_slice` records only the targeted steps and token positions, skipping tensor clones and device transfers on uncaptured steps." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "2cb2b39b", + "metadata": { + "execution": { + "iopub.execute_input": "2026-11-10T15:24:43.864174Z", + "iopub.status.busy": "2026-11-10T15:24:43.864066Z", + "iopub.status.idle": "2026-11-10T15:24:43.875717Z", + "shell.execute_reply": "2026-11-10T15:24:43.875373Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Captured steps: (0, 3) out of 4 total steps\n", + "Stacked continuation residual shape: (2, 1, 4, 32)\n" + ] + } + ], + "source": [ + "_, sparse_traj = bridge.diffusion_generate(\n", + " prompt_direct,\n", + " max_new_tokens=4,\n", + " return_cache=True,\n", + " names_filter=\"blocks.1.hook_out\",\n", + " cache_steps=[0, 3],\n", + " pos_slice=(4, 8), # generation canvas positions only\n", + " device=\"cpu\",\n", + ")\n", + "\n", + "assert sparse_traj.steps == (0, 3)\n", + "assert sparse_traj.num_steps_run == 4\n", + "assert sparse_traj.stack(\"blocks.1.hook_out\").shape == (2, 1, 4, compact_cfg.d_model)\n", + "print(f\"Captured steps: {sparse_traj.steps} out of {sparse_traj.num_steps_run} total steps\")\n", + "print(f\"Stacked continuation residual shape: {tuple(sparse_traj.stack('blocks.1.hook_out').shape)}\")" + ] + }, + { + "cell_type": "markdown", + "id": "7fd535cd", + "metadata": {}, + "source": [ + "## 3. Activation Patching Across Denoising Steps (`make_step_patch_hook`)\n", + "\n", + "In discrete diffusion with absorbing `[MASK]` tokens, each step $s$ computes logits over all masked positions and commits the highest-confidence positions. Once a position is unmasked at step $T_{\\text{commit}}$, standard low-confidence remasking never overwrites it on steps $s > T_{\\text{commit}}$.\n", + "\n", + "We can test this causal boundary directly with `DiffusionTrajectoryCache.make_step_patch_hook`:\n", + "1. Run a **source** prompt (`prompt_source`) and cache `blocks.1.hook_out` across all denoising steps.\n", + "2. Patch the source trajectory's `blocks.1.hook_out` into a **base** prompt (`prompt_direct`) at:\n", + " - **All steps** ($s \\in \\{0, 1, 2, 3\\}$): full trajectory transplant.\n", + " - **Pre-commitment step ($s = 0$)**: before the first position commits.\n", + " - **Post-commitment step ($s = 3$)**: after positions settled on steps $0, 1, 2$ are already unmasked." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "93345ea5", + "metadata": { + "execution": { + "iopub.execute_input": "2026-11-10T15:24:43.876912Z", + "iopub.status.busy": "2026-11-10T15:24:43.876778Z", + "iopub.status.idle": "2026-11-10T15:24:43.918302Z", + "shell.execute_reply": "2026-11-10T15:24:43.917879Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + " condition continuation_tokens\n", + " base (unpatched) [48, 48, 48, 26]\n", + " source (target) [58, 48, 48, 48]\n", + " patch all steps {0,1,2,3} [58, 48, 48, 48]\n", + " patch step 0 only (pre-commit) [58, 48, 48, 48]\n", + "patch step 3 only (3/4 already committed) [6, 48, 48, 26]\n" + ] + } + ], + "source": [ + "prompt_source = torch.tensor([[19, 23, 29, 31]])\n", + "sequences_source, traj_source = bridge.diffusion_generate(\n", + " prompt_source,\n", + " max_new_tokens=4,\n", + " return_cache=True,\n", + " names_filter=[\"blocks.1.hook_out\"],\n", + ")\n", + "\n", + "# 1. Patch across all denoising steps on the continuation positions (4..8)\n", + "patch_all = traj_source.make_step_patch_hook(\n", + " \"blocks.1.hook_out\",\n", + " steps=[0, 1, 2, 3],\n", + " positions=slice(4, 8),\n", + ")\n", + "with bridge.hooks(fwd_hooks=[(\"blocks.1.hook_out\", patch_all)]):\n", + " seq_patched_all = bridge.diffusion_generate(prompt_direct, max_new_tokens=4)\n", + "\n", + "# 2. Patch only at step 0 (before the first commitment)\n", + "patch_step0 = traj_source.make_step_patch_hook(\n", + " \"blocks.1.hook_out\",\n", + " steps=0,\n", + " positions=slice(4, 8),\n", + ")\n", + "with bridge.hooks(fwd_hooks=[(\"blocks.1.hook_out\", patch_step0)]):\n", + " seq_patched_step0 = bridge.diffusion_generate(prompt_direct, max_new_tokens=4)\n", + "\n", + "# 3. Patch only at the final step 3 (after 3 of 4 continuation positions are already unmasked)\n", + "patch_step3 = traj_source.make_step_patch_hook(\n", + " \"blocks.1.hook_out\",\n", + " steps=3,\n", + " positions=slice(4, 8),\n", + ")\n", + "with bridge.hooks(fwd_hooks=[(\"blocks.1.hook_out\", patch_step3)]):\n", + " seq_patched_step3, traj_patched_step3 = bridge.diffusion_generate(\n", + " prompt_direct,\n", + " max_new_tokens=4,\n", + " return_cache=True,\n", + " names_filter=[\"blocks.1.hook_out\"],\n", + " )\n", + "\n", + "# Positions already unmasked before step 3 are invariant to the step-3 intervention\n", + "pre_committed_mask = traj_patched_step3.step_tokens[3][0, 4:8] != MASK_ID\n", + "assert int(pre_committed_mask.sum().item()) == 3\n", + "assert torch.equal(\n", + " seq_patched_step3[0, 4:8][pre_committed_mask],\n", + " sequences_direct[0, 4:8][pre_committed_mask],\n", + ")\n", + "# Meanwhile, patching all steps transplants the full source continuation\n", + "assert torch.equal(seq_patched_all[0, 4:8], sequences_source[0, 4:8])\n", + "\n", + "patch_summary = pd.DataFrame(\n", + " [\n", + " {\"condition\": \"base (unpatched)\", \"continuation_tokens\": sequences_direct[0, 4:8].tolist()},\n", + " {\"condition\": \"source (target)\", \"continuation_tokens\": sequences_source[0, 4:8].tolist()},\n", + " {\"condition\": \"patch all steps {0,1,2,3}\", \"continuation_tokens\": seq_patched_all[0, 4:8].tolist()},\n", + " {\"condition\": \"patch step 0 only (pre-commit)\", \"continuation_tokens\": seq_patched_step0[0, 4:8].tolist()},\n", + " {\"condition\": \"patch step 3 only (3/4 already committed)\", \"continuation_tokens\": seq_patched_step3[0, 4:8].tolist()},\n", + " ]\n", + ")\n", + "print(patch_summary.to_string(index=False))" + ] + }, + { + "cell_type": "markdown", + "id": "a610e982", + "metadata": {}, + "source": [ + "## 4. Empirical Commitment-Order Measurements on `LLaDA-8B-Instruct` and `Dream-7B` (ARC-Challenge)\n", + "\n", + "In a full 256-step denoising trajectory ($L_{\\text{gen}} = 256, S = 256$, low-confidence remasking with EOS/EoT confidence suppression), tracking each step's canvas $x_s$ separates two clocks:\n", + "- **$T_{\\text{conv}}$ (`convergence_step`)**: the earliest denoising step from which the decoded option verdict (`A`-`D`) never changes again.\n", + "- **$T_{\\text{commit}}$ (`answer_commitment_step`)**: the step at which the answer token position itself is unmasked from `[MASK]`.\n", + "- **$U$ (`unwritten_reasoning_at_convergence`)**: the fraction of reasoning-span tokens that are still `[MASK]` at $T_{\\text{conv}}$.\n", + "- **Order gap**: normalized difference between the mean unmasking step of the answer span and the mean unmasking step of the reasoning span (negative means the answer commits *before* the reasoning tokens).\n", + "\n", + "The table below records measured `traj-eval` trajectories on ARC-Challenge for `GSAI-ML/LLaDA-8B-Instruct` and `Dream-org/Dream-v0-Instruct-7B` (NVIDIA RTX A6000, `torch 2.5.1+cu121`, `temperature=0.0`, $S=256$ steps)." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "2326c3b9", + "metadata": { + "execution": { + "iopub.execute_input": "2026-11-10T15:24:43.919360Z", + "iopub.status.busy": "2026-11-10T15:24:43.919242Z", + "iopub.status.idle": "2026-11-10T15:24:43.926456Z", + "shell.execute_reply": "2026-11-10T15:24:43.926084Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "LLaDA-8B-Instruct ARC-C reasoning trajectories (S=256 steps):\n", + " id gold pred T_conv T_commit unwritten_reasoning_U order_gap n_steps\n", + " Mercury_7175875 C C 2 2 0.9921 -0.4883 256\n", + " Mercury_407327 D D 4 4 0.9842 -0.4805 256\n", + " MCAS_2006_9_44 D D 16 16 0.9368 -0.4336 256\n", + "Mercury_SC_408547 C C 36 36 0.8577 -0.3555 256\n", + "Mercury_SC_409171 B B 41 41 0.8379 -0.3359 256\n", + "Cross-checkpoint commitment dynamics on ARC-Challenge:\n", + " checkpoint condition accuracy mean_T_conv mean_T_commit mean_unwritten_U mean_order_gap\n", + " GSAI-ML/LLaDA-8B-Instruct ARC-C direct (n=64) 0.8750 4.45 144.67 0.7884 0.2322\n", + " GSAI-ML/LLaDA-8B-Instruct ARC-C reasoning (n=5 smoke) 1.0000 19.80 19.80 0.9217 -0.4188\n", + "Dream-org/Dream-v0-Instruct-7B ARC-C direct (n=64) 0.8906 3.80 3.80 0.0000 0.0000\n", + "Dream-org/Dream-v0-Instruct-7B ARC-C reasoning (n=64) 0.8281 189.46 192.29 0.0145 0.3412\n" + ] + } + ], + "source": [ + "llada_arcc_reasoning_items = [\n", + " {\n", + " \"id\": \"Mercury_7175875\",\n", + " \"gold\": \"C\",\n", + " \"pred\": \"C\",\n", + " \"T_conv\": 2,\n", + " \"T_commit\": 2,\n", + " \"unwritten_reasoning_U\": 0.9921,\n", + " \"order_gap\": -0.4883,\n", + " \"n_steps\": 256,\n", + " },\n", + " {\n", + " \"id\": \"Mercury_407327\",\n", + " \"gold\": \"D\",\n", + " \"pred\": \"D\",\n", + " \"T_conv\": 4,\n", + " \"T_commit\": 4,\n", + " \"unwritten_reasoning_U\": 0.9842,\n", + " \"order_gap\": -0.4805,\n", + " \"n_steps\": 256,\n", + " },\n", + " {\n", + " \"id\": \"MCAS_2006_9_44\",\n", + " \"gold\": \"D\",\n", + " \"pred\": \"D\",\n", + " \"T_conv\": 16,\n", + " \"T_commit\": 16,\n", + " \"unwritten_reasoning_U\": 0.9368,\n", + " \"order_gap\": -0.4336,\n", + " \"n_steps\": 256,\n", + " },\n", + " {\n", + " \"id\": \"Mercury_SC_408547\",\n", + " \"gold\": \"C\",\n", + " \"pred\": \"C\",\n", + " \"T_conv\": 36,\n", + " \"T_commit\": 36,\n", + " \"unwritten_reasoning_U\": 0.8577,\n", + " \"order_gap\": -0.3555,\n", + " \"n_steps\": 256,\n", + " },\n", + " {\n", + " \"id\": \"Mercury_SC_409171\",\n", + " \"gold\": \"B\",\n", + " \"pred\": \"B\",\n", + " \"T_conv\": 41,\n", + " \"T_commit\": 41,\n", + " \"unwritten_reasoning_U\": 0.8379,\n", + " \"order_gap\": -0.3359,\n", + " \"n_steps\": 256,\n", + " },\n", + "]\n", + "\n", + "df_llada_items = pd.DataFrame(llada_arcc_reasoning_items)\n", + "assert round(float(df_llada_items[\"T_commit\"].mean()), 1) == 19.8\n", + "assert round(float(df_llada_items[\"unwritten_reasoning_U\"].mean()), 4) == 0.9217\n", + "assert round(float(df_llada_items[\"order_gap\"].mean()), 4) == -0.4188\n", + "print(\"LLaDA-8B-Instruct ARC-C reasoning trajectories (S=256 steps):\")\n", + "print(df_llada_items.to_string(index=False))\n", + "\n", + "suite_comparison = pd.DataFrame(\n", + " [\n", + " {\n", + " \"checkpoint\": \"GSAI-ML/LLaDA-8B-Instruct\",\n", + " \"condition\": \"ARC-C direct (n=64)\",\n", + " \"accuracy\": 0.8750,\n", + " \"mean_T_conv\": 4.45,\n", + " \"mean_T_commit\": 144.67,\n", + " \"mean_unwritten_U\": 0.7884,\n", + " \"mean_order_gap\": +0.2322,\n", + " },\n", + " {\n", + " \"checkpoint\": \"GSAI-ML/LLaDA-8B-Instruct\",\n", + " \"condition\": \"ARC-C reasoning (n=5 smoke)\",\n", + " \"accuracy\": 1.0000,\n", + " \"mean_T_conv\": 19.80,\n", + " \"mean_T_commit\": 19.80,\n", + " \"mean_unwritten_U\": 0.9217,\n", + " \"mean_order_gap\": -0.4188,\n", + " },\n", + " {\n", + " \"checkpoint\": \"Dream-org/Dream-v0-Instruct-7B\",\n", + " \"condition\": \"ARC-C direct (n=64)\",\n", + " \"accuracy\": 0.8906,\n", + " \"mean_T_conv\": 3.80,\n", + " \"mean_T_commit\": 3.80,\n", + " \"mean_unwritten_U\": 0.0000,\n", + " \"mean_order_gap\": 0.0000,\n", + " },\n", + " {\n", + " \"checkpoint\": \"Dream-org/Dream-v0-Instruct-7B\",\n", + " \"condition\": \"ARC-C reasoning (n=64)\",\n", + " \"accuracy\": 0.8281,\n", + " \"mean_T_conv\": 189.46,\n", + " \"mean_T_commit\": 192.29,\n", + " \"mean_unwritten_U\": 0.0145,\n", + " \"mean_order_gap\": +0.3412,\n", + " },\n", + " ]\n", + ")\n", + "print(\"Cross-checkpoint commitment dynamics on ARC-Challenge:\")\n", + "print(suite_comparison.to_string(index=False))" + ] + }, + { + "cell_type": "markdown", + "id": "cbcd6b79", + "metadata": {}, + "source": [ + "### Choosing `cache_steps` From Commitment Clocks\n", + "\n", + "On `LLaDA-8B-Instruct` reasoning prompts, the answer token commits at a mean step of $T_{\\text{commit}} = 19.8$ (median $16.0$) out of $256$, while $U = 92.17\\%$ of the reasoning tokens are still `[MASK]` (`order_gap = -0.4188`). Rather than caching all 256 steps, a mechanistic audit of answer commitment vs. post-commitment rationalization only needs `cache_steps` bracketed around $T_{\\text{commit}}$ (for example, `cache_steps=[0, 2, 4, 8, 16, 32, 64, 128]`), reducing cached activation memory by $32\\times$ while preserving every pre-commitment decision step." + ] + }, + { + "cell_type": "markdown", + "id": "f7d1da5b", + "metadata": {}, + "source": [ + "## 5. Opt-In Released `GSAI-ML/LLaDA-8B-Instruct` Run (`TL_RUN_RELEASED_LLADA=1`)\n", + "\n", + "Set `TL_RUN_RELEASED_LLADA=1` to boot the released `GSAI-ML/LLaDA-8B-Instruct` checkpoint through `TransformerBridge`, cache residual and attention hooks around the commitment window on `MCAS_2006_9_44` ($T_{\\text{commit}} = 16$), and patch the pre-commitment residual stream across denoising steps." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "4446ec6a", + "metadata": { + "execution": { + "iopub.execute_input": "2026-11-10T15:24:43.927613Z", + "iopub.status.busy": "2026-11-10T15:24:43.927500Z", + "iopub.status.idle": "2026-11-10T15:24:43.929808Z", + "shell.execute_reply": "2026-11-10T15:24:43.929449Z" + } + }, + "outputs": [], + "source": [ + "# NBVAL_SKIP\n", + "if RUN_RELEASED:\n", + " llada_bridge = TransformerBridge.boot_transformers(\n", + " \"GSAI-ML/LLaDA-8B-Instruct\",\n", + " trust_remote_code=True,\n", + " dtype=torch.bfloat16,\n", + " device=\"cuda\" if torch.cuda.is_available() else \"cpu\",\n", + " )\n", + " arcc_prompt = (\n", + " \"Question: Which of the following is an example of a physical change?\\n\"\n", + " \"A. burning wood\\nB. rusting iron\\nC. baking a cake\\nD. melting ice\\n\"\n", + " \"Reason step by step, then give Answer: .\"\n", + " )\n", + " text_out, llada_traj = llada_bridge.diffusion_generate(\n", + " arcc_prompt,\n", + " max_new_tokens=64,\n", + " steps=64,\n", + " block_length=64,\n", + " return_cache=True,\n", + " names_filter=[\"blocks.16.hook_out\", \"blocks.31.hook_out\"],\n", + " cache_steps=[0, 2, 4, 8, 16, 32, 63],\n", + " device=\"cpu\",\n", + " )\n", + " print(\"Decoded completion:\", text_out)\n", + " print(\"Captured trajectory steps:\", llada_traj.steps)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "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.2" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/make_docs.py b/docs/make_docs.py index 4a38ac0754..d10a1c524a 100644 --- a/docs/make_docs.py +++ b/docs/make_docs.py @@ -40,6 +40,7 @@ def copy_demos(_app: Optional[Any] = None): copy_to_dir = GENERATED_DIR / "demos" notebooks_to_copy = [ "Backward_Lens_Demo.ipynb", + "Diffusion_Trajectory_Patching_Demo.ipynb", "Exploratory_Analysis_Demo.ipynb", "Jacobian_Lens_Coordinate_Patch_Benchmark_Demo.ipynb", "Jacobian_Lens_Decomposition_Demo.ipynb", diff --git a/docs/source/content/tutorials.md b/docs/source/content/tutorials.md index 226df91889..43350e29eb 100644 --- a/docs/source/content/tutorials.md +++ b/docs/source/content/tutorials.md @@ -46,6 +46,8 @@ The two core notebooks, which pair with the links in [Where To Start](#where-to- - [**Natural Language Autoencoders**](https://colab.research.google.com/github/TransformerLensOrg/TransformerLens/blob/main/demos/Natural_Language_Autoencoders_Demo.ipynb) - Runs the [Natural Language Autoencoder](https://transformer-circuits.pub/2026/nla/) round trip with `TransformerBridge` hooks and generation: a verbalizer model explains one residual-stream activation in text, and a reconstructor model predicts the activation back from that text. The mechanics run on CPU with a tiny Qwen2 checkpoint; the opt-in released-model section (`TL_RUN_RELEASED_NLA=1`) downloads about 39 GiB of Qwen2.5-7B-Instruct checkpoints. +- [**Diffusion Trajectory Caching & Denoising Patching**](https://colab.research.google.com/github/TransformerLensOrg/TransformerLens/blob/main/demos/Diffusion_Trajectory_Patching_Demo.ipynb) - Records step-indexed activations (`DiffusionTrajectoryCache`) across iterative masked-diffusion denoising steps via `diffusion_generate(..., return_cache=True)` and patches residual-stream activations across denoising steps (`make_step_patch_hook`) before vs. after the answer commitment step on `LLaDA-8B-Instruct` and `Dream-7B` ARC-Challenge trajectories. Runs on CPU in seconds; the opt-in released-model section (`TL_RUN_RELEASED_LLADA=1`) runs `GSAI-ML/LLaDA-8B-Instruct`. + - [**Patchscopes & Generation with Patching**](https://colab.research.google.com/github/TransformerLensOrg/TransformerLens/blob/main/demos/Patchscopes_Generation_Demo.ipynb) - Implements [Patchscopes](https://arxiv.org/abs/2401.06102) (read a hidden representation from a source prompt and patch it into a target prompt), extends it to generating multiple tokens with patching, and applies it to logit lens, entity description and zero-shot feature extraction on GPT-2 small. - [**SVD Circuits**](https://colab.research.google.com/github/TransformerLensOrg/TransformerLens/blob/main/demos/SVD_Circuits_Demo.ipynb) - Decomposes one GPT-2 attention head's QK and OV maps into singular directions with `FactoredMatrix.svd`, reads the OV directions out through the unembedding with `logit_readout`, and keeps or ablates single directions on an IOI prompt through `hook_result`. Based on [Ahmad et al. (2025)](https://arxiv.org/abs/2511.20273); runs on CPU. diff --git a/docs/source/index.md b/docs/source/index.md index 247042c5d6..3ebd5fad33 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -71,6 +71,7 @@ generated/demos/SVD_Circuits_Demo generated/demos/SwiGLU_Gate_Sign_Demo generated/demos/Weight_Edit_Audit_Demo generated/demos/Backward_Lens_Demo +generated/demos/Diffusion_Trajectory_Patching_Demo generated/demos/Natural_Language_Autoencoders_Demo generated/demos/Main_Demo generated/demos/Exploratory_Analysis_Demo diff --git a/tests/integration/model_bridge/test_llada_adapter.py b/tests/integration/model_bridge/test_llada_adapter.py index 2078560d96..c15bbea4d2 100644 --- a/tests/integration/model_bridge/test_llada_adapter.py +++ b/tests/integration/model_bridge/test_llada_adapter.py @@ -904,3 +904,49 @@ def test_causal_loss_modes_fail_clearly(models: TinyModels, return_type: str) -> with pytest.raises(NotImplementedError, match="causal.*loss|shifted.*loss"): models.bridge(tokens, return_type=return_type) assert len(models.bridge.original_model.output_attentions_requests) == calls_before + + +def test_diffusion_generate_trajectory_cache_and_step_patching(models: TinyModels) -> None: + from transformer_lens.ActivationCache import DiffusionTrajectoryCache + + prompt_a = torch.tensor([[5, 7, 9]]) + prompt_b = torch.tensor([[11, 13, 15]]) + mask_id = models.bridge.cfg.mask_token_id + + seq_a, traj_a = models.bridge.diffusion_generate( + prompt_a, + max_new_tokens=3, + return_cache=True, + names_filter=["blocks.0.attn.hook_pattern", "blocks.1.hook_out"], + ) + assert isinstance(seq_a, torch.Tensor) + assert isinstance(traj_a, DiffusionTrajectoryCache) + assert seq_a.shape == (1, 6) + assert torch.equal(seq_a[:, :3], prompt_a) + assert torch.all(seq_a[:, 3:] != mask_id) + assert traj_a.steps == (0, 1, 2) + # Step 0 starts with all 3 continuation positions masked; each step unmasks 1 token + assert (traj_a.step_tokens[0][0, 3:] == mask_id).sum().item() == 3 + assert (traj_a.step_tokens[1][0, 3:] == mask_id).sum().item() == 2 + assert (traj_a.step_tokens[2][0, 3:] == mask_id).sum().item() == 1 + assert traj_a.stack("blocks.0.attn.hook_pattern").shape == (3, 1, 4, 6, 6) + assert traj_a.stack("blocks.1.hook_out").shape == (3, 1, 6, 32) + + # Step-selective activation patching across denoising steps + patch_fn = traj_a.make_step_patch_hook( + "blocks.1.hook_out", + steps=[0], + positions=slice(3, 6), + ) + with models.bridge.hooks(fwd_hooks=[("blocks.1.hook_out", patch_fn)]): + _, patched_traj_b = models.bridge.diffusion_generate( + prompt_b, + max_new_tokens=3, + return_cache=True, + names_filter=["blocks.1.hook_out"], + cache_steps=[0, 1], + ) + torch.testing.assert_close( + patched_traj_b[0]["blocks.1.hook_out"][:, 3:6, :], + traj_a[0]["blocks.1.hook_out"][:, 3:6, :], + ) diff --git a/tests/unit/model_bridge/supported_architectures/test_llada_adapter.py b/tests/unit/model_bridge/supported_architectures/test_llada_adapter.py index ee8fb24182..be981eb110 100644 --- a/tests/unit/model_bridge/supported_architectures/test_llada_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_llada_adapter.py @@ -6,7 +6,10 @@ import pytest import torch -from tests.unit.model_bridge.supported_architectures.helpers import make_bridge_cfg +from tests.unit.model_bridge.supported_architectures.helpers import ( + DiffusionContractTests, + make_bridge_cfg, +) from transformer_lens.config import TransformerBridgeConfig from transformer_lens.conversion_utils.conversion_steps import RearrangeTensorConversion from transformer_lens.conversion_utils.param_processing_conversion import ( @@ -323,3 +326,22 @@ def test_existing_causal_attention_default_remains_triangular() -> None: pattern = cache["blocks.0.attn.hook_pattern"] assert pattern.shape == (1, 4, 5, 5) assert torch.count_nonzero(torch.triu(pattern, diagonal=1)) == 0 + + +class TestLLaDADiffusionContract(DiffusionContractTests): + adapter_cls = LLaDAArchitectureAdapter + architecture = "LLaDAModelLM" + expected_sampler = "diffusion_generate" + + def _make_adapter(self) -> LLaDAArchitectureAdapter: + return LLaDAArchitectureAdapter(_bridge_config()) + + +def test_prepare_model_binds_diffusion_generate_sampler( + adapter: LLaDAArchitectureAdapter, +) -> None: + dummy = SimpleNamespace(config=SimpleNamespace(output_attentions=True, use_cache=True)) + adapter.prepare_model(dummy) + assert dummy.config.output_attentions is False + assert dummy.config.use_cache is False + assert callable(getattr(dummy, "diffusion_generate", None)) diff --git a/tests/unit/model_bridge/test_diffusion_generate.py b/tests/unit/model_bridge/test_diffusion_generate.py index d0a31e02e1..420c15e563 100644 --- a/tests/unit/model_bridge/test_diffusion_generate.py +++ b/tests/unit/model_bridge/test_diffusion_generate.py @@ -122,6 +122,12 @@ def test_llada2_block_length_never_exceeds_budget(self) -> None: kwargs = self._kwargs("llada2_moe", "LLaDA2MoeArchitectureAdapter", 4, 8) assert kwargs["block_length"] <= 4 + def test_llada_rounds_gen_length_to_whole_blocks(self) -> None: + kwargs = self._kwargs("llada", "LLaDAArchitectureAdapter", 40, 8) + assert kwargs["gen_length"] % kwargs["block_length"] == 0 + assert kwargs["gen_length"] >= 40 + assert kwargs["steps"] == kwargs["gen_length"] + def test_gidd_max_length_is_generated_tokens_not_canvas(self) -> None: """Gidd's windows start at prompt_length and span max_length, so folding the prompt into it would over-generate by prompt_len tokens.""" @@ -266,3 +272,175 @@ def test_shape_mismatch_raises(self) -> None: bridge, _ = self._router() with _pytest.raises(ValueError, match="does not match"): bridge.set_processed_weights({"layer.weight": torch.zeros(3, 3)}) + + +class TestDiffusionTrajectoryCache: + """Trajectory-step activation caching and step-selective patching for diffusion LMs.""" + + def _make_fake_diffusion_bridge(self): + import torch + + from transformer_lens.hook_points import HookPoint + from transformer_lens.model_bridge.bridge import TransformerBridge + from transformer_lens.model_bridge.bridge_core import BridgeCore + + class _StepModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.hp = HookPoint() + self.hp.name = "blocks.0.hook_out" + + def forward(self, canvas: torch.Tensor) -> torch.Tensor: + act = canvas.to(torch.float32).unsqueeze(-1).repeat(1, 1, 4) + return self.hp(act) + + def diffusion_generate( + self, tokens: torch.Tensor, max_new_tokens: int = 3, steps: int = 3, **kwargs + ) -> torch.Tensor: + prompt_len = tokens.shape[-1] + canvas = torch.full( + (tokens.shape[0], prompt_len + max_new_tokens), + 99, + dtype=tokens.dtype, + ) + canvas[:, :prompt_len] = tokens + for s in range(steps): + out = self(canvas) + # Unmask position prompt_len + s using the hooked activation's first channel + unmasked_token = int(out[0, prompt_len + s, 0].item()) % 50 + 10 + canvas = canvas.clone() + canvas[:, prompt_len + s] = unmasked_token + return canvas + + class _Adapter: + native_sampler = "diffusion_generate" + supports_generation = False + + def native_sampler_kwargs(self, max_new_tokens: int, prompt_len: int) -> dict: + return {"max_new_tokens": max_new_tokens, "steps": max_new_tokens} + + class _Cfg: + architecture = "FakeDiffusionLM" + device = "cpu" + n_layers = 1 + + class _FakeBridge(BridgeCore, torch.nn.Module): + diffusion_generate = TransformerBridge.diffusion_generate + current_diffusion_step = TransformerBridge.current_diffusion_step + + def __init__(self) -> None: + torch.nn.Module.__init__(self) + self.original_model = _StepModel() + self.adapter = _Adapter() + self.cfg = _Cfg() + self.tokenizer = None + self.compatibility_mode = False + self.context_level = 0 + + @property + def hook_dict(self): + return {"blocks.0.hook_out": self.original_model.hp} + + def _gated_hook_reason(self, name: str): + return None + + def _check_hook_fireable(self, caller_name: str, canonical_name: str) -> None: + if canonical_name not in self.hook_dict: + raise KeyError(caller_name) + + return _FakeBridge() + + def test_return_cache_preserves_every_denoising_step_without_clobbering(self) -> None: + import torch + + from transformer_lens.ActivationCache import DiffusionTrajectoryCache + + bridge = self._make_fake_diffusion_bridge() + prompt = torch.tensor([[1, 2]]) + sequences, traj_cache = bridge.diffusion_generate( + prompt, max_new_tokens=3, return_cache=True, names_filter="blocks.0.hook_out" + ) + assert sequences.shape == (1, 5) + assert isinstance(traj_cache, DiffusionTrajectoryCache) + assert traj_cache.steps == (0, 1, 2) + assert traj_cache.num_steps_run == 3 + assert len(traj_cache) == 3 + # Step 0 has all 3 continuation positions masked (99); step 1 has position 2 unmasked + assert traj_cache.step_tokens[0].tolist() == [[1, 2, 99, 99, 99]] + assert traj_cache.step_tokens[1][0, 2].item() != 99 + assert traj_cache.step_tokens[1][0, 3].item() == 99 + stacked = traj_cache.stack("blocks.0.hook_out") + assert stacked.shape == (3, 1, 5, 4) + assert not torch.equal( + traj_cache[0]["blocks.0.hook_out"], traj_cache[2]["blocks.0.hook_out"] + ) + assert bridge.current_diffusion_step is None + + def test_cache_steps_filters_recorded_denoising_steps(self) -> None: + import torch + + bridge = self._make_fake_diffusion_bridge() + prompt = torch.tensor([[1, 2]]) + _, traj_cache = bridge.diffusion_generate( + prompt, + max_new_tokens=4, + return_cache=True, + names_filter="blocks.0.hook_out", + cache_steps=[0, 2], + pos_slice=slice(2, 4), + remove_batch_dim=True, + ) + assert traj_cache.steps == (0, 2) + assert traj_cache.num_steps_run == 4 + assert 1 not in traj_cache + assert traj_cache[0]["blocks.0.hook_out"].shape == (2, 4) + assert traj_cache.stack("blocks.0.hook_out").shape == (2, 2, 4) + + def test_make_step_patch_hook_patches_only_selected_steps_and_positions(self) -> None: + import torch + + from transformer_lens.ActivationCache import ( + ActivationCache, + DiffusionTrajectoryCache, + ) + + bridge = self._make_fake_diffusion_bridge() + prompt = torch.tensor([[1, 2]]) + # Build a source trajectory cache whose step-1 activation at pos 3 is set to 7.0 + _, clean_traj = bridge.diffusion_generate( + prompt, max_new_tokens=3, return_cache=True, names_filter="blocks.0.hook_out" + ) + custom_act = clean_traj[1]["blocks.0.hook_out"].clone() + custom_act[:, 3, :] = 7.0 + source_traj = DiffusionTrajectoryCache( + {1: ActivationCache({"blocks.0.hook_out": custom_act}, bridge)}, + num_steps_run=3, + model=bridge, + ) + patch_fn = source_traj.make_step_patch_hook( + "blocks.0.hook_out", steps=1, positions=slice(3, 4) + ) + with bridge.hooks(fwd_hooks=[("blocks.0.hook_out", patch_fn)]): + patched_seq, patched_traj = bridge.diffusion_generate( + prompt, max_new_tokens=3, return_cache=True, names_filter="blocks.0.hook_out" + ) + # Step 1 position 3 was patched to 7.0 -> unmasked token at step 1 is (7 % 50 + 10) == 17 + assert patched_seq[0, 3].item() == 17 + assert torch.all(patched_traj[1]["blocks.0.hook_out"][:, 3, :] == 7.0) + # Step 0 was not in steps=1 and remains unpatched (99.0 at masked position 3) + assert torch.all(patched_traj[0]["blocks.0.hook_out"][:, 3, :] == 99.0) + + def test_invalid_cache_steps_rejected(self) -> None: + from typing import Any, cast + + import pytest as _pytest + import torch + + bridge = self._make_fake_diffusion_bridge() + prompt = torch.tensor([[1, 2]]) + with _pytest.raises(ValueError, match="non-negative"): + bridge.diffusion_generate(prompt, max_new_tokens=2, return_cache=True, cache_steps=-1) + with _pytest.raises(TypeError, match="cache_steps"): + bridge.diffusion_generate( + prompt, max_new_tokens=2, return_cache=True, cache_steps=cast(Any, True) + ) diff --git a/tests/unit/test_make_docs.py b/tests/unit/test_make_docs.py index a95af68687..5fd291271c 100644 --- a/tests/unit/test_make_docs.py +++ b/tests/unit/test_make_docs.py @@ -39,6 +39,7 @@ def test_copy_demos_creates_generated_dir_when_absent(tmp_path, monkeypatch): copied = sorted(p.name for p in (generated / "demos").iterdir()) assert copied == [ "Backward_Lens_Demo.ipynb", + "Diffusion_Trajectory_Patching_Demo.ipynb", "Exploratory_Analysis_Demo.ipynb", "Jacobian_Lens_Coordinate_Patch_Benchmark_Demo.ipynb", "Jacobian_Lens_Decomposition_Demo.ipynb", diff --git a/transformer_lens/ActivationCache.py b/transformer_lens/ActivationCache.py index 08c743317d..5c499a4b4c 100644 --- a/transformer_lens/ActivationCache.py +++ b/transformer_lens/ActivationCache.py @@ -15,6 +15,7 @@ class first, including the examples, and then skimming the available methods. Yo import logging from collections import Counter +from collections.abc import Iterable, Mapping, Sequence from typing import ( TYPE_CHECKING, Any, @@ -2028,3 +2029,178 @@ def _ln_then_project(stack: torch.Tensor) -> torch.Tensor: return residual_stack, labels else: return residual_stack + + +StepFilter = Optional[Union[int, Iterable[int], Callable[[int], bool]]] + + +def _normalize_step_filter(cache_steps: StepFilter) -> Callable[[int], bool]: + """Normalize ``cache_steps`` into a boolean predicate over 0-based denoising step indices.""" + if cache_steps is None: + return lambda _step: True + if isinstance(cache_steps, bool): + raise TypeError("cache_steps must be an int, iterable of ints, callable, or None; got bool") + if isinstance(cache_steps, int): + if cache_steps < 0: + raise ValueError(f"cache_steps index must be non-negative; got {cache_steps}") + target = cache_steps + return lambda step: step == target + if callable(cache_steps): + return cache_steps + if isinstance(cache_steps, Iterable) and not isinstance(cache_steps, (str, bytes)): + steps_set: set[int] = set() + for item in cache_steps: + if isinstance(item, bool) or not isinstance(item, int): + raise TypeError( + f"cache_steps iterable must contain only non-negative ints; got {item!r}" + ) + if item < 0: + raise ValueError(f"cache_steps index must be non-negative; got {item}") + steps_set.add(item) + frozen = frozenset(steps_set) + return lambda step: step in frozen + raise TypeError( + f"cache_steps must be an int, iterable of ints, callable, or None; got {type(cache_steps).__name__}" + ) + + +class DiffusionTrajectoryCache(Mapping[int, ActivationCache]): + """Step-indexed activation caches recorded across a diffusion denoising trajectory. + + Autoregressive :meth:`~transformer_lens.model_bridge.TransformerBridge.generate` can build + its cache from a single post-hoc pass on the completed sequence because causal masking + keeps earlier positions invariant as new tokens are appended. Discrete diffusion LMs + (Dream, LLaDA, LLaDA-2, Gidd) use bidirectional attention and iteratively unmask a canvas + over ``S`` denoising steps: step ``s``'s activations depend on the partially-masked canvas + ``x_s`` and cannot be recovered from the final unmasked sequence ``x_S``. Attaching + standard ``add_caching_hooks`` across ``diffusion_generate`` would overwrite ``cache[name]`` + on every forward pass, silently discarding steps ``0..S-2``. + + ``DiffusionTrajectoryCache`` holds one :class:`ActivationCache` per captured denoising step, + along with the input token canvas ``step_tokens[s]`` fed into step ``s``. + """ + + def __init__( + self, + step_caches: Mapping[int, ActivationCache], + *, + step_tokens: Optional[Mapping[int, torch.Tensor]] = None, + num_steps_run: int = 0, + model: Optional[Any] = None, + ) -> None: + self.step_caches: Dict[int, ActivationCache] = { + int(step): step_caches[step] for step in sorted(step_caches) + } + self.step_tokens: Dict[int, torch.Tensor] = ( + {int(step): step_tokens[step] for step in sorted(step_tokens)} + if step_tokens is not None + else {} + ) + self.num_steps_run: int = max( + int(num_steps_run), + (max(self.step_caches) + 1) if self.step_caches else 0, + ) + first_cache = next(iter(self.step_caches.values()), None) + self.model: Optional[Any] = ( + model if model is not None else getattr(first_cache, "model", None) + ) + + @property + def steps(self) -> Tuple[int, ...]: + """Captured denoising step indices in ascending chronological order.""" + return tuple(self.step_caches.keys()) + + def __getitem__(self, step: int) -> ActivationCache: + if isinstance(step, bool) or not isinstance(step, int): + raise KeyError(f"Step index must be an int; got {step!r}") + if step not in self.step_caches: + raise KeyError( + f"Denoising step {step} was not captured (captured steps: {list(self.steps)})" + ) + return self.step_caches[step] + + def __iter__(self) -> Iterator[int]: + return iter(self.step_caches) + + def __len__(self) -> int: + return len(self.step_caches) + + def __contains__(self, step: object) -> bool: + return not isinstance(step, bool) and isinstance(step, int) and step in self.step_caches + + def __repr__(self) -> str: + sample_cache = next(iter(self.step_caches.values()), None) + n_hooks = len(sample_cache) if sample_cache is not None else 0 + return ( + f"DiffusionTrajectoryCache(steps={list(self.steps)}, " + f"num_steps_run={self.num_steps_run}, hooks_per_step={n_hooks})" + ) + + def stack( + self, + key: str, + steps: Optional[Sequence[int]] = None, + ) -> torch.Tensor: + """Stack ``cache[step][key]`` across ``steps`` (default: all captured steps) along dim 0.""" + selected = tuple(self.steps if steps is None else steps) + if not selected: + raise ValueError("Cannot stack an empty sequence of denoising steps") + return torch.stack([self[step][key] for step in selected], dim=0) + + def to(self, device: Union[str, torch.device]) -> "DiffusionTrajectoryCache": + """Move every per-step :class:`ActivationCache` and ``step_tokens`` tensor to ``device``.""" + for step_cache in self.step_caches.values(): + step_cache.to(device) + for step, tokens in list(self.step_tokens.items()): + self.step_tokens[step] = tokens.to(device) + return self + + def make_step_patch_hook( + self, + hook_name: str, + *, + steps: StepFilter = None, + positions: Optional[Union[int, Sequence[int], slice]] = None, + pos_dim: int = 1, + ) -> Callable[..., Optional[torch.Tensor]]: + """Build a forward hook that patches ``hook_name`` from this trajectory at selected steps. + + When attached via ``with bridge.hooks(fwd_hooks=[(hook_name, hook_fn)]):`` around + :meth:`~transformer_lens.model_bridge.TransformerBridge.diffusion_generate`, the hook + reads ``bridge.current_diffusion_step`` on each forward pass and substitutes the source + activation recorded at that step. Steps not matching ``steps`` (or not present in this + cache) pass through untouched. + """ + step_predicate = _normalize_step_filter(steps) + for step, step_cache in self.step_caches.items(): + if step_predicate(step) and hook_name not in step_cache: + raise KeyError( + f"Hook {hook_name!r} is not present in captured step {step} " + f"(available keys: {list(step_cache.keys())})" + ) + + def _patch_hook(tensor: torch.Tensor, *, hook: Any) -> Optional[torch.Tensor]: + current_step = getattr(self.model, "current_diffusion_step", None) + if current_step is None and hook is not None: + current_step = getattr(hook, "ctx", {}).get("diffusion_step") + if current_step is None or not step_predicate(current_step): + return None + if current_step not in self.step_caches: + return None + source = self.step_caches[current_step][hook_name].to( + device=tensor.device, dtype=tensor.dtype + ) + if positions is None: + return source.clone() + patched = tensor.clone() + axis = pos_dim if pos_dim >= 0 else tensor.ndim + pos_dim + idx: List[Any] = [slice(None)] * tensor.ndim + idx[axis] = positions + index_tuple = tuple(idx) + if source.shape == patched[index_tuple].shape: + patched[index_tuple] = source + else: + patched[index_tuple] = source[index_tuple] + return patched + + return _patch_hook diff --git a/transformer_lens/__init__.py b/transformer_lens/__init__.py index 0b2b74a045..79cd2e867a 100644 --- a/transformer_lens/__init__.py +++ b/transformer_lens/__init__.py @@ -9,7 +9,7 @@ tools, utilities, ) -from .ActivationCache import ActivationCache +from .ActivationCache import ActivationCache, DiffusionTrajectoryCache from .cache.key_value_cache import TransformerLensKeyValueCache from .cache.key_value_cache_entry import TransformerLensKeyValueCacheEntry from .config import TransformerBridgeConfig @@ -102,6 +102,7 @@ def __dir__(): "TransformerBridgeConfig", "FactoredMatrix", "ActivationCache", + "DiffusionTrajectoryCache", "SVDInterpreter", "HookedRootModule", "TransformerLensKeyValueCache", diff --git a/transformer_lens/model_bridge/supported_architectures/llada.py b/transformer_lens/model_bridge/supported_architectures/llada.py index a1feba7fd3..a86393b648 100644 --- a/transformer_lens/model_bridge/supported_architectures/llada.py +++ b/transformer_lens/model_bridge/supported_architectures/llada.py @@ -17,9 +17,11 @@ from __future__ import annotations import math +import types from typing import Any, Dict, Optional, cast import torch +import torch.nn.functional as F from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter from transformer_lens.model_bridge.generalized_components import ( @@ -248,21 +250,138 @@ def forward(self, hidden_states: torch.Tensor, **kwargs: Any) -> torch.Tensor: return self.hook_out(output) +def _llada_diffusion_generate( + self: Any, + input_ids: torch.Tensor, + *, + max_new_tokens: int = 32, + steps: Optional[int] = None, + block_length: Optional[int] = None, + temperature: float = 0.0, + remasking: str = "low_confidence", + mask_token_id: Optional[int] = None, + attention_mask: Optional[torch.Tensor] = None, + **kwargs: Any, +) -> torch.Tensor: + """Blockwise masked-diffusion denoising sampler for ``LLaDAModelLM``. + + Calls ``self(x, ...)`` on every denoising step so ``TransformerBridge`` hooks + and trajectory-step caching fire on each intermediate canvas ``x_s``. + """ + gen_length = int(kwargs.pop("gen_length", max_new_tokens)) + if gen_length <= 0: + raise ValueError(f"max_new_tokens/gen_length must be > 0; got {gen_length}") + block_len = int(block_length if block_length is not None else min(32, gen_length)) + if block_len <= 0 or gen_length % block_len != 0: + raise ValueError( + f"gen_length ({gen_length}) must be divisible by positive block_length ({block_len})" + ) + num_blocks = gen_length // block_len + total_steps = int(steps if steps is not None else gen_length) + if total_steps <= 0 or total_steps % num_blocks != 0: + raise ValueError( + f"steps ({total_steps}) must be positive and divisible by num_blocks ({num_blocks})" + ) + steps_per_block = total_steps // num_blocks + base_transfer = block_len // steps_per_block + rem_transfer = block_len % steps_per_block + schedule = [base_transfer + (1 if i < rem_transfer else 0) for i in range(steps_per_block)] + + cfg = getattr(self, "config", None) + resolved_mask_id = ( + int(mask_token_id) + if mask_token_id is not None + else int(getattr(cfg, "mask_token_id", None) or 126336) + ) + prompt_len = int(input_ids.shape[-1]) + device = input_ids.device + x = torch.full( + (input_ids.shape[0], prompt_len + gen_length), + resolved_mask_id, + dtype=input_ids.dtype, + device=device, + ) + x[:, :prompt_len] = input_ids + + full_attn: Optional[torch.Tensor] = None + if attention_mask is not None: + full_attn = torch.cat( + [ + attention_mask.to(device), + torch.ones( + (input_ids.shape[0], gen_length), + dtype=attention_mask.dtype, + device=device, + ), + ], + dim=-1, + ) + + for b_idx in range(num_blocks): + block_start = prompt_len + b_idx * block_len + block_end = prompt_len + (b_idx + 1) * block_len + for s_idx in range(steps_per_block): + mask_index = x == resolved_mask_id + forward_out = self(x, attention_mask=full_attn) if full_attn is not None else self(x) + logits = getattr(forward_out, "logits", forward_out) + if temperature > 0.0: + logits_f64 = logits.to(torch.float64) + noise = torch.rand_like(logits_f64, dtype=torch.float64) + gumbel = (-torch.log(noise.clamp_min(1e-12))) ** temperature + x0 = torch.argmax(logits_f64.exp() / gumbel, dim=-1) + else: + x0 = torch.argmax(logits, dim=-1) + + if remasking == "low_confidence": + probs = F.softmax(logits.to(torch.float64), dim=-1) + x0_p = torch.gather(probs, dim=-1, index=x0.unsqueeze(-1)).squeeze(-1) + elif remasking == "random": + x0_p = torch.rand(x0.shape, device=device, dtype=torch.float64) + else: + raise NotImplementedError(f"Unsupported remasking strategy: {remasking!r}") + + x0_p[:, :block_start] = -torch.inf + x0_p[:, block_end:] = -torch.inf + x0 = torch.where(mask_index, x0, x) + confidence = torch.where(mask_index, x0_p, torch.full_like(x0_p, -torch.inf)) + + k_transfer = schedule[s_idx] + if k_transfer > 0: + for row in range(confidence.shape[0]): + _, select_idx = torch.topk(confidence[row], k=k_transfer) + x[row, select_idx] = x0[row, select_idx] + + return x + + class LLaDAArchitectureAdapter(ArchitectureAdapter): """Adapter for the dense ``LLaDAModelLM`` architecture. Support is deliberately limited to the released dense LLaDA block contract: Llama-style blocks, RMSNorm, separate bias-free projections, RoPE, - bidirectional attention, an untied LM head, and no KV cache. The external - iterative denoising/remasking loop is not a TransformerBridge generation API. + bidirectional attention, an untied LM head, and no KV cache. Iterative + masked-diffusion sampling runs through ``bridge.diffusion_generate``. Loading the Hugging Face checkpoint requires the caller to opt in with ``trust_remote_code=True``. """ - applicable_phases: list[int] = [] + applicable_phases: list[int] = [1, 2, 3, 4] supports_generation: bool = False supports_hf_output_attentions: bool = False supports_causal_loss: bool = False + native_sampler: str = "diffusion_generate" + + def native_sampler_kwargs(self, max_new_tokens: int, prompt_len: int) -> dict: + """LLaDA denoises whole blocks; cap block_length at the token budget.""" + del prompt_len + block_length = min(32, max_new_tokens) + blocks = max(1, -(-max_new_tokens // block_length)) + gen_length = blocks * block_length + return { + "gen_length": gen_length, + "block_length": block_length, + "steps": gen_length, + } def __init__(self, cfg: Any) -> None: super().__init__(cfg) @@ -366,7 +485,7 @@ def prepare_loading(self, model_name: str, model_kwargs: dict) -> None: config.use_cache = False def prepare_model(self, hf_model: Any) -> None: - """Keep the wrapper and underlying model on the no-cache path.""" + """Keep the wrapper and underlying model on the no-cache path and bind the sampler.""" for config in ( getattr(hf_model, "config", None), getattr(getattr(hf_model, "model", None), "config", None), @@ -374,3 +493,9 @@ def prepare_model(self, hf_model: Any) -> None: if config is not None: config.output_attentions = False config.use_cache = False + if hf_model is not None and not hasattr(hf_model, "diffusion_generate"): + setattr( + hf_model, + "diffusion_generate", + types.MethodType(_llada_diffusion_generate, hf_model), + ) diff --git a/transformer_lens/model_bridge/transformer_bridge.py b/transformer_lens/model_bridge/transformer_bridge.py index c13997eec4..0e28e00c43 100644 --- a/transformer_lens/model_bridge/transformer_bridge.py +++ b/transformer_lens/model_bridge/transformer_bridge.py @@ -35,7 +35,12 @@ from transformers.tokenization_utils_base import PreTrainedTokenizerBase from transformer_lens import utilities as utils -from transformer_lens.ActivationCache import ActivationCache +from transformer_lens.ActivationCache import ( + ActivationCache, + DiffusionTrajectoryCache, + StepFilter, + _normalize_step_filter, +) from transformer_lens.config import TransformerBridgeConfig from transformer_lens.FactoredMatrix import FactoredMatrix from transformer_lens.hook_points import HookIntrospectionMixin, HookPoint @@ -45,6 +50,7 @@ _SELF_ATTENTION_NAMES, BridgeCore, _validate_execution_kwargs, + build_alias_to_canonical_map, ) from transformer_lens.model_bridge.component_setup import ( refresh_container_state_owners, @@ -74,6 +80,7 @@ has_batched_experts, ) from transformer_lens.model_bridge.get_params_util import get_bridge_params +from transformer_lens.utilities import Slice, SliceInput from transformer_lens.utilities.activation_functions import softcap_enabled from transformer_lens.utilities.devices import move_to_and_update_config from transformer_lens.utilities.position_ids import accepts_mask_derived_position_ids @@ -4046,19 +4053,35 @@ def _logits_to_tuple(logits_list: list[torch.Tensor]) -> tuple[torch.Tensor, ... # return_input_tokens only (no cache) return result, input_tokens + @property + def current_diffusion_step(self) -> Optional[int]: + """Zero-based denoising step index during an active :meth:`diffusion_generate` call.""" + return getattr(self, "_diffusion_step", None) + @torch.no_grad() def diffusion_generate( self, input: Union[str, List[str], torch.Tensor], max_new_tokens: int = 32, prepend_bos: Optional[bool] = None, + *, + return_cache: bool = False, + names_filter: Optional[Union[str, List[str], Callable[[str], bool]]] = None, + cache_steps: StepFilter = None, + pos_slice: Optional[Union[Slice, SliceInput, slice]] = None, + device: Optional[Union[str, torch.device]] = None, + remove_batch_dim: bool = False, **kwargs: Any, - ) -> Union[str, torch.Tensor]: + ) -> Any: """Sample from a non-autoregressive (diffusion) architecture. Delegates to the model's own sampler, which calls the model through - ``__call__`` so bridge hooks fire on every denoising step. + ``__call__`` so bridge hooks fire on every denoising step. When + ``return_cache=True``, records per-step activations into a + :class:`~transformer_lens.ActivationCache.DiffusionTrajectoryCache` + without overwriting earlier denoising steps. """ + step_predicate = _normalize_step_filter(cache_steps) sampler_name = getattr(self.adapter, "native_sampler", None) architecture = self.cfg.architecture or type(self.adapter).__name__ if sampler_name is None: @@ -4091,9 +4114,94 @@ def diffusion_generate( input, prepend_bos=prepend_bos, move_to_device=True, truncate=False ) + step_dicts: Dict[int, Dict[str, torch.Tensor]] = {} + step_tokens: Dict[int, torch.Tensor] = {} + staging_cache: Dict[str, torch.Tensor] = {} + step_fwd_hooks: List[Tuple[Union[str, Callable], Callable]] = [] + + if return_cache: + hook_dict = getattr(self, "hook_dict", {}) + if isinstance(names_filter, (str, list)): + aliases = build_alias_to_canonical_map(hook_dict) + requested = [names_filter] if isinstance(names_filter, str) else list(names_filter) + for req_name in requested: + self._check_hook_fireable(req_name, aliases.get(req_name, req_name)) + normalized_pos_slice: Optional[Union[Slice, SliceInput]] + if isinstance(pos_slice, slice): + wrapped_slice = Slice(None) + wrapped_slice.slice = pos_slice + wrapped_slice.mode = "slice" + normalized_pos_slice = wrapped_slice + else: + normalized_pos_slice = pos_slice + _, raw_fwd_hooks, _ = self.get_caching_hooks( + names_filter=names_filter, + incl_bwd=False, + device=device, + remove_batch_dim=remove_batch_dim, + cache=staging_cache, + pos_slice=normalized_pos_slice, + ) + if not raw_fwd_hooks and names_filter and isinstance(names_filter, (str, list)): + raise KeyError( + f"names_filter {names_filter!r} matched no hook points on this model; " + "check the name against model.hook_dict." + ) + + def _make_step_save_hook(_fn: Callable) -> Callable: + def _step_save(tensor: Any, *, hook: Any) -> None: + step = getattr(self, "_diffusion_step", None) + if step is None or not step_predicate(step): + return None + staging_cache.clear() + _fn(tensor, hook=hook) + if staging_cache: + bucket = step_dicts.setdefault(step, {}) + bucket.update(staging_cache) + staging_cache.clear() + return None + + return _step_save + + step_fwd_hooks = [ + (h_name, _make_step_save_hook(h_fn)) for h_name, h_fn in raw_fwd_hooks + ] + + step_counter = 0 + self._diffusion_step: Optional[int] = None + + def _on_step_start( + _module: Any, args: Tuple[Any, ...], kwargs_in: Optional[Dict[str, Any]] = None + ) -> None: + nonlocal step_counter + step = step_counter + step_counter += 1 + self._diffusion_step = step + if return_cache and step_predicate(step): + canvas = args[0] if args else (kwargs_in or {}).get("input_ids") + if isinstance(canvas, torch.Tensor): + step_tokens[step] = canvas.detach().clone().to(device) + sampler_kwargs = self.adapter.native_sampler_kwargs(max_new_tokens, tokens.shape[-1]) sampler_kwargs.update(kwargs) - output = sampler(tokens, **sampler_kwargs) + + pre_handle: Optional[RemovableHandle] = None + if hasattr(self.original_model, "register_forward_pre_hook"): + pre_handle = self.original_model.register_forward_pre_hook( + _on_step_start, with_kwargs=True + ) + + try: + if step_fwd_hooks: + with self.hooks(fwd_hooks=step_fwd_hooks): + output = sampler(tokens, **sampler_kwargs) + else: + output = sampler(tokens, **sampler_kwargs) + finally: + if pre_handle is not None: + pre_handle.remove() + self._diffusion_step = None + # Samplers return either bare ids or a generation output object. sequences = getattr(output, "sequences", output) # They also disagree on whether the prompt is included: Dream and Gidd @@ -4110,9 +4218,26 @@ def diffusion_generate( if not includes_prompt: sequences = torch.cat([tokens, sequences], dim=-1) + result: Union[str, torch.Tensor] if was_string and self.tokenizer is not None: - return self.tokenizer.decode(sequences[0], skip_special_tokens=True) - return sequences + result = self.tokenizer.decode(sequences[0], skip_special_tokens=True) + else: + result = sequences + + if not return_cache: + return result + + step_caches = { + step: ActivationCache(step_dict, self, has_batch_dim=not remove_batch_dim) + for step, step_dict in sorted(step_dicts.items()) + } + trajectory_cache = DiffusionTrajectoryCache( + step_caches, + step_tokens=step_tokens, + num_steps_run=step_counter, + model=self, + ) + return result, trajectory_cache @torch.no_grad() def generate_stream(