File size: 3,679 Bytes
e29f88c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | """Differentiable fixed-prefix execution; one working query after the first pass."""
import torch
from torch import Tensor, nn
from .memory import PromptMemory
from torch.utils.checkpoint import checkpoint
from transformers.models.qwen2.modeling_qwen2 import apply_rotary_pos_emb, eager_attention_forward
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
def pick(x: Tensor, positions: list[int]) -> Tensor:
return torch.stack([x[row, position:position + 1] for row, position in enumerate(positions)])
def layer_compute(layer: nn.Module, hidden: Tensor, embeddings: tuple[Tensor, Tensor],
mask: Tensor | None, positions: list[int] | None = None,
cached: tuple[Tensor, Tensor] | None = None) -> tuple[Tensor, Tensor, Tensor]:
"""Pure operation so non-reentrant checkpointing preserves prefix gradients."""
attention = layer.self_attn
residual = hidden
x = layer.input_layernorm(hidden)
shape = (*x.shape[:-1], -1, attention.head_dim)
q = attention.q_proj(x).view(shape).transpose(1, 2)
k = attention.k_proj(x).view(shape).transpose(1, 2)
v = attention.v_proj(x).view(shape).transpose(1, 2)
q, k = apply_rotary_pos_emb(q, k, *embeddings)
if cached is not None:
keys, values = cached
active = (torch.arange(keys.shape[-2], device=x.device)[None, :] ==
torch.tensor(positions, device=x.device)[:, None])[:, None, :, None]
k = torch.where(active, k, keys)
v = torch.where(active, v, values)
interface = ALL_ATTENTION_FUNCTIONS.get_interface(attention.config._attn_implementation, eager_attention_forward)
output, _ = interface(attention, q, k, v, mask,
dropout=attention.attention_dropout if layer.training else 0.0,
scaling=attention.scaling, sliding_window=None)
output = output.reshape(*x.shape[:-1], -1).contiguous()
hidden = residual + attention.o_proj(output)
hidden = hidden + layer.mlp(layer.post_attention_layernorm(hidden))
return hidden, k, v
def fixed_prefix_layer(layer: nn.Module, hidden: Tensor, embeddings: tuple[Tensor, Tensor],
mask: Tensor | None, memory: PromptMemory, index: int,
use_checkpoint: bool) -> Tensor:
"""Reuse prefix K/V without detaching it or caching across model forwards."""
positions = memory.write_mask[..., 0].long().argmax(-1).tolist()
cached = None
first = memory.output is None
if not first:
hidden = pick(hidden, positions)
embeddings = tuple(pick(e.expand(memory.write_mask.shape[0], -1, -1), positions) for e in embeddings)
cached = memory.keys_values[index]
if mask is None:
keep = torch.arange(cached[0].shape[-2], device=hidden.device)[None, :] <= torch.tensor(positions, device=hidden.device)[:, None]
mask = torch.zeros_like(keep, dtype=hidden.dtype).masked_fill(~keep, torch.finfo(hidden.dtype).min)[:, None, None, :]
else:
mask = torch.stack([mask[row if mask.shape[0] > 1 else 0, :, pos:pos + 1, :] for row, pos in enumerate(positions)])
def compute(x, cos, sin, *kv):
return layer_compute(layer, x, (cos, sin), mask, positions, tuple(kv) if kv else None)
args = (hidden, *embeddings, *(cached or ()))
result, keys, values = (checkpoint(compute, *args, use_reentrant=False)
if use_checkpoint and torch.is_grad_enabled() else compute(*args))
if first:
memory.keys_values.append((keys, values))
memory.layer_outputs.append(result)
return result
return torch.where(memory.write_mask, result, memory.layer_outputs[index])
|