diff --git a/src/gfn/estimators.py b/src/gfn/estimators.py index b816cd76..12ffeb9e 100644 --- a/src/gfn/estimators.py +++ b/src/gfn/estimators.py @@ -1239,12 +1239,20 @@ class RecurrentDiscretePolicyEstimator(RecurrentPolicyMixin, DiscretePolicyEstim per-step artifacts. Non-recurrent estimators should use the default PolicyMixin and the standard ``DiscretePolicyEstimator`` base class instead. + An optional, fixed-length-only performance mode (opt-in; the default leaves + behavior unchanged): + + - ``use_kv_cache=True`` decodes one token per step against a persistent, in-place + preallocated KV-cache, cutting rollout cost from O(L^3) to O(L^2). The in-place + buffers cannot carry gradients, so cached sampling runs under ``no_grad`` -- it is + a sampling-time optimization, ideal for inference/rollout. + Notes ----- - - Forward is intended for on-policy generation; off-policy evaluation over - entire trajectories typically requires different batching and masking. - ``init_carry`` is a hard requirement for compatibility with the recurrent PolicyMixin. + - The cached fast path assumes fixed-length trajectories; see the Gap-2 TODO in + ``gfn/utils/prob_calculations.py``. Attributes: module: The neural network module to use. @@ -1261,14 +1269,27 @@ def __init__( preprocessor: Preprocessor | None = None, is_backward: bool = False, debug: bool = False, + use_kv_cache: bool = False, + cache_max_len: int | None = None, ): """Initializes a RecurrentDiscretePolicyEstimator. + See the class docstring for what ``use_kv_cache`` does and why cached sampling + is ``no_grad``. + Args: module: The neural network module to use. n_actions: Total number of actions in the discrete environment. preprocessor: Preprocessor object that transforms states to tensors. + is_backward: Flag indicating whether this estimator is for backward policy. debug: If True, enables expensive validation checks. + use_kv_cache: Opt into the O(L^2) in-place KV-cache sampling fast path. + Cached sampling runs under ``no_grad``, so it is a sampling-time + optimization (ideal for inference/rollout). If False (default), sampling + uses the corrected grad-bearing full-prefix forward each step. + cache_max_len: KV-cache preallocation length; defaults to the module's + ``max_position_embeddings``. Rarely set by hand; must be at least the + env horizon + 1 (the +1 is the BOS token). """ if preprocessor is None: preprocessor = IdentityPreprocessor(output_dim=None) @@ -1280,49 +1301,71 @@ def __init__( is_backward=is_backward, debug=debug, ) + self.use_kv_cache = use_kv_cache + # A preallocated in-place cache is only available for modules that expose a + # bounded position range (e.g. TransformerDiscreteSequenceModel). Modules + # without one (e.g. RNNs, whose carry is already fixed-size) decode one token + # per step with their ordinary carry. + max_positions = getattr(module, "max_position_embeddings", None) + if cache_max_len is None: + cache_max_len = max_positions + # Validate the cache budget eagerly (against the module's position range) so a + # too-small value fails here, not deep inside a rollout with a confusing error. + if ( + cache_max_len is not None + and max_positions is not None + and cache_max_len > max_positions + ): + raise ValueError( + f"cache_max_len ({cache_max_len}) exceeds the module's " + f"max_position_embeddings ({max_positions})." + ) + self._cache_max_len = cache_max_len + + @property + def _bos_index(self) -> int: + """BOS token index (the sequence model reserves ``vocab_size`` for BOS).""" + return int(getattr(self.module, "vocab_size", self.n_actions - 1)) def forward( self, states: States, carry: dict[str, torch.Tensor], ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: - """Forward pass of the module. + """Per-step forward: next-action logits for the current rollout states. + + Both paths build the same canonical token stream ``[BOS, w_0, ..., w_{t-1}]`` + and take the logits at the last position, so they are numerically equivalent + to one teacher-forced forward over the full sequence: + + - ``use_kv_cache=False`` (default): re-embed the whole committed prefix with a + fresh carry each step (grad-bearing, O(L^3) over a rollout). This is the + corrected baseline. + - ``use_kv_cache=True``: feed only the newest committed token against the + persistent carry (in-place KV-cache for transformers), under ``no_grad``. Args: - states: The input states. - carry: The carry from the previous step. + states: The current (active) rollout states, ``(B, W)`` word indices with + ``-1`` padding. + carry: The recurrent carry threaded from the previous step. Returns: - The output of the module, as a tensor of shape (*batch_shape, output_dim). + ``(logits (B, n_actions), updated carry)``. """ - # Prepare integer token sequences without -1 padding and use a BOS index. - # We infer the active sequence length per row from (token != -1). tokens = states.tensor.long() + if tokens.dim() != 2: + raise ValueError( + "RecurrentDiscretePolicyEstimator expects 2D state tensors " + f"(batch, seq); got shape {tuple(tokens.shape)}." + ) + lengths = (states.tensor >= 0).sum(dim=-1) # committed words per row - # Replace padding (-1) with BOS index expected by the sequence model. - # RecurrentDiscreteSequenceModel reserves index == vocab_size for BOS. - bos_index = getattr(self.module, "vocab_size", self.n_actions - 1) - tokens = torch.where( - tokens < 0, torch.as_tensor(bos_index, device=tokens.device), tokens - ) - - # Determine a common prefix length across the (active) batch. - # Active rows in a rollout step share the same length; use max for safety. - # We still derive length from original states.tensor where -1 marks padding. - valid_mask = states.tensor >= 0 - if valid_mask.ndim == 1: - max_len = int(valid_mask.sum().item()) + if self.use_kv_cache: + logits, carry = self._forward_cached(tokens, lengths, carry) else: - max_len = int(valid_mask.sum(dim=-1).max().item()) - if max_len == 0: - max_len = 1 # Ensure at least BOS is processed - - # Trim to the common active prefix length and run the sequence model. - seq_input = tokens[..., :max_len] - logits, carry = self.module(seq_input, carry) + logits, carry = self._forward_full_prefix(tokens, lengths) - # Use the logits corresponding to the last processed token. - logits = logits[:, -1, :] # (b, n_actions) + logits = logits[:, -1, :] # (B, n_actions) if self.expected_output_dim is not None: assert logits.shape[-1] == self.expected_output_dim, ( @@ -1332,20 +1375,104 @@ def forward( return logits, carry - def init_carry( + def _canonical_sequence( + self, tokens: torch.Tensor, lengths: torch.Tensor + ) -> torch.Tensor: + """Build ``[BOS, w_0, ..., w_{L-1}]`` up to the longest committed length ``L``. + + Used by the (no-cache) full-prefix path and the teacher-forced loss recompute. + """ + bos = self._bos_index + max_len = int(lengths.max().item()) + # Replace any -1 padding with BOS; unlike the cached path this one has no + # fixed-length guard, so shorter rows can carry padding inside [:max_len]. + words = torch.where(tokens < 0, torch.full_like(tokens, bos), tokens)[ + :, :max_len + ] + bos_col = torch.full( + (tokens.shape[0], 1), bos, dtype=torch.long, device=tokens.device + ) + return torch.cat([bos_col, words], dim=1) # (B, L + 1) + + def _forward_full_prefix( self, - batch_size: int, - device: torch.device, + tokens: torch.Tensor, + lengths: torch.Tensor, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Grad-bearing baseline: full prefix + fresh carry each step.""" + seq = self._canonical_sequence(tokens, lengths) + fresh = self._module_init_carry(seq.shape[0], seq.device, preallocate=False) + logits, new_carry = self.module(seq, fresh) + return logits, new_carry + + def _forward_cached( + self, + tokens: torch.Tensor, + lengths: torch.Tensor, + carry: dict[str, torch.Tensor], + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Cached fast path: feed only the newest committed token, under ``no_grad``. + + Assumes fixed-length rollouts (all active rows share a committed length). The + always-on backstop is the module's batch-size carry check (a masked, unsliced + carry mismatches the active batch); the stronger equal-length / cursor-sync + checks are gated behind ``debug`` so the hot path avoids per-step device syncs. + """ + length0 = int(lengths[0].item()) if lengths.numel() > 0 else 0 + + if self.debug and lengths.numel() > 0: + if not bool(torch.all(lengths == length0)): + raise ValueError( + "use_kv_cache=True requires equal-length active trajectories " + "(fixed-length env). Variable-length carry slicing is out of scope " + "(see the Gap-2 TODO in gfn/utils/prob_calculations.py)." + ) + cursor = int(carry["cache_len"].item()) if "cache_len" in carry else 0 + if cursor != length0: + raise ValueError( + f"KV-cache cursor ({cursor}) is out of sync with the committed " + f"length ({length0}); the carry must be threaded unmodified " + "across steps." + ) + + # Feed the canonical stream [BOS, w_0, ..., w_{L-1}] one token per step: at a + # committed length of ``length0`` that is column ``length0 - 1``, or BOS when + # nothing is committed yet. + if length0 == 0: + step_tok = torch.full( + (tokens.shape[0], 1), + self._bos_index, + dtype=torch.long, + device=tokens.device, + ) + else: + step_tok = tokens[:, length0 - 1 : length0] + + with torch.no_grad(): + logits, new_carry = self.module(step_tok, carry) + return logits, new_carry + + def _module_init_carry( + self, batch_size: int, device: torch.device, preallocate: bool ) -> dict[str, torch.Tensor]: init_carry = getattr(self.module, "init_carry", None) if not callable(init_carry): raise NotImplementedError( "Module does not implement init_carry(batch_size, device)." ) - init_carry_fn = cast(Callable[[int, torch.device], Any], init_carry) - + init_carry_fn = cast(Callable[..., dict[str, torch.Tensor]], init_carry) + if preallocate and self._cache_max_len is not None: + return init_carry_fn(batch_size, device, max_len=self._cache_max_len) return init_carry_fn(batch_size, device) + def init_carry( + self, + batch_size: int, + device: torch.device, + ) -> dict[str, torch.Tensor]: + # Preallocate an in-place cache only for the cached fast path. + return self._module_init_carry(batch_size, device, preallocate=self.use_kv_cache) + class DiffusionPolicyEstimator(PolicyMixin, Estimator): """Base class for diffusion policy estimators.""" diff --git a/src/gfn/utils/modules.py b/src/gfn/utils/modules.py index 14e5fc61..82c68d4b 100644 --- a/src/gfn/utils/modules.py +++ b/src/gfn/utils/modules.py @@ -1284,7 +1284,27 @@ def forward( hidden: torch.Tensor, key_carry: torch.Tensor, value_carry: torch.Tensor, + cache_len: Optional[int] = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Self-attention over the K/V carry, in one of two carry modes. + + Args: + hidden: (batch, timesteps, embed_dim) block input. + key_carry: The running key cache. + value_carry: The running value cache. + cache_len: Selects the carry mode. When ``None`` (default) the block runs + in **grow-by-cat** mode: the new K/V are concatenated onto the carry + and freshly allocated buffers are returned. This is autograd-safe and + used by the full-sequence / no-cache paths. When an int, it is the + write cursor for **in-place** mode: ``key_carry`` / ``value_carry`` are + preallocated buffers, the new K/V are written at + ``[cache_len:cache_len+T]`` without reallocation, and the *same* buffer + objects are returned. In-place mutation is only version-counter-safe + under ``no_grad``. + + Returns: + (block output, updated key carry, updated value carry). + """ batch, timesteps, _ = hidden.size() normed_hidden = self.norm1(hidden) @@ -1297,30 +1317,45 @@ def forward( k = k.view(batch, timesteps, self.num_heads, self.head_dim).transpose(1, 2) v = v.view(batch, timesteps, self.num_heads, self.head_dim).transpose(1, 2) - carry_length = key_carry.size(2) - updated_key_carry = torch.cat((key_carry, k), dim=2) - updated_value_carry = torch.cat((value_carry, v), dim=2) - - attn_scores = torch.matmul(q, updated_key_carry.transpose(-2, -1)) / math.sqrt( + if cache_len is None: + # Grow-by-cat mode (grad-safe): O(L^2) reallocation across a rollout. + prefix_length = key_carry.size(2) + updated_key_carry = torch.cat((key_carry, k), dim=2) + updated_value_carry = torch.cat((value_carry, v), dim=2) + keys = updated_key_carry + values = updated_value_carry + else: + # In-place mode: write the new K/V into the preallocated buffer at the + # current cursor without reallocating (narrow+copy_ is TorchScript-safe). + prefix_length = cache_len + end = cache_len + timesteps + key_carry.narrow(2, cache_len, timesteps).copy_(k) + value_carry.narrow(2, cache_len, timesteps).copy_(v) + updated_key_carry = key_carry + updated_value_carry = value_carry + keys = key_carry.narrow(2, 0, end) + values = value_carry.narrow(2, 0, end) + + attn_scores = torch.matmul(q, keys.transpose(-2, -1)) / math.sqrt( float(self.head_dim) ) - if timesteps > 1 or carry_length > 0: - total_kv_length = carry_length + timesteps + if timesteps > 1 or prefix_length > 0: + total_kv_length = prefix_length + timesteps kv_positions = torch.arange( total_kv_length, device=hidden.device, dtype=torch.long ) query_positions = torch.arange( timesteps, device=hidden.device, dtype=torch.long ).unsqueeze(1) - causal_mask = kv_positions.unsqueeze(0) <= (query_positions + carry_length) + causal_mask = kv_positions.unsqueeze(0) <= (query_positions + prefix_length) attn_scores = attn_scores.masked_fill( ~causal_mask.unsqueeze(0).unsqueeze(0), float("-inf") ) attn_weights = torch.softmax(attn_scores, dim=-1) attn_weights = self.attn_dropout(attn_weights) - attn_output = torch.matmul(attn_weights, updated_value_carry) + attn_output = torch.matmul(attn_weights, values) attn_output = attn_output.transpose(1, 2).reshape( batch, timesteps, self.embed_dim ) @@ -1406,20 +1441,57 @@ def init_carry( self, batch_size: int, device: torch.device, + max_len: Optional[int] = None, ) -> dict[str, torch.Tensor]: + """Initialize the recurrent carry. + + Args: + batch_size: Number of parallel rollouts. + device: Device to allocate the carry on. + max_len: If ``None`` (default), allocates length-0 K/V buffers that are + grown by concatenation each step (grad-safe; O(L^2) reallocation). + If an int, preallocates fixed ``[batch, n_heads, max_len, head_dim]`` + K/V buffers plus a scalar ``cache_len`` write cursor, so the block + writes new K/V in place (O(L), ``no_grad`` only). ``max_len`` must be + at least the maximum number of tokens decoded in a rollout (the fixed + env horizon + 1 for the BOS token) and no larger than + ``max_position_embeddings``. + + Returns: + The initialized carry dict. + """ weight = self.token_embedding.weight carry: dict[str, torch.Tensor] = { "position": torch.zeros(batch_size, dtype=torch.long, device=device), } - empty_key = weight.new_empty(batch_size, self.num_heads, 0, self.head_dim).to( - device - ) - empty_value = weight.new_empty(batch_size, self.num_heads, 0, self.head_dim).to( - device - ) - for key_name, value_name in zip(self.key_names, self.value_names): - carry[key_name] = empty_key.clone() - carry[value_name] = empty_value.clone() + if max_len is None: + empty_key = weight.new_empty( + batch_size, self.num_heads, 0, self.head_dim + ).to(device) + empty_value = weight.new_empty( + batch_size, self.num_heads, 0, self.head_dim + ).to(device) + for key_name, value_name in zip(self.key_names, self.value_names): + carry[key_name] = empty_key.clone() + carry[value_name] = empty_value.clone() + else: + if max_len <= 0: + raise ValueError("max_len must be positive for a preallocated carry.") + if max_len > self.max_position_embeddings: + raise ValueError( + "max_len exceeds max_position_embeddings; increase the model's " + "max_position_embeddings to preallocate a carry of this length." + ) + # Scalar cursor into the preallocated buffers (shared across rows, which + # advance in lockstep for the fixed-length fast path). + carry["cache_len"] = torch.zeros((), dtype=torch.long, device=device) + for key_name, value_name in zip(self.key_names, self.value_names): + carry[key_name] = weight.new_empty( + batch_size, self.num_heads, max_len, self.head_dim + ).to(device) + carry[value_name] = weight.new_empty( + batch_size, self.num_heads, max_len, self.head_dim + ).to(device) return carry @@ -1462,6 +1534,13 @@ def forward( updated_carry: dict[str, torch.Tensor] = {} + # An int cursor selects the in-place preallocated carry path in each block; + # None keeps the grow-by-cat path. The cursor is shared by all layers since + # they decode the same number of tokens. + cache_len: Optional[int] = None + if "cache_len" in carry: + cache_len = int(carry["cache_len"].item()) + for idx, layer in enumerate(self.layers): key_name = self.key_names[idx] value_name = self.value_names[idx] @@ -1471,22 +1550,28 @@ def forward( ) key_carry = carry[key_name] value_carry = carry[value_name] - if key_carry.size(0) != batch or key_carry.size(1) != self.num_heads: + if ( + key_carry.size(0) != batch + or value_carry.size(0) != batch + or key_carry.size(1) != self.num_heads + or value_carry.size(1) != self.num_heads + or key_carry.size(-1) != self.head_dim + or value_carry.size(-1) != self.head_dim + or key_carry.device != device + or value_carry.device != device + ): raise ValueError( - "Key carry shape is incompatible with the provided tokens." + "Key/value carry is incompatible with the provided tokens " + "(batch, num_heads, head_dim, or device mismatch)." ) - if value_carry.size(0) != batch or value_carry.size(1) != self.num_heads: + # Actionable guard kept separate: overflow points the user at max_len. + if cache_len is not None and cache_len + timesteps > key_carry.size(2): raise ValueError( - "Value carry shape is incompatible with the provided tokens." + "Preallocated carry is too short for the decoded sequence; " + "increase max_len passed to init_carry." ) - if (key_carry.size(-1) != self.head_dim) or ( - value_carry.size(-1) != self.head_dim - ): - raise ValueError("Key/value carry head dimension mismatch detected.") - if key_carry.device != device or value_carry.device != device: - raise ValueError("Key/value carry tensors must share the input device.") hidden, updated_key_carry, updated_value_carry = layer( - hidden, key_carry, value_carry + hidden, key_carry, value_carry, cache_len ) updated_carry[key_name] = updated_key_carry updated_carry[value_name] = updated_value_carry @@ -1495,6 +1580,8 @@ def forward( logits = self.output_projection(hidden) updated_carry["position"] = positions + timesteps + if "cache_len" in carry: + updated_carry["cache_len"] = carry["cache_len"] + timesteps return logits, updated_carry @property diff --git a/testing/test_kv_cache.py b/testing/test_kv_cache.py new file mode 100644 index 00000000..21ffe9cb --- /dev/null +++ b/testing/test_kv_cache.py @@ -0,0 +1,238 @@ +"""Correctness gates for the fixed-length KV-cache sampling optimization. + +Covers two layers: + +- The transformer module's in-place preallocated carry (PR1a): incremental decode + equals a single full-sequence forward, the buffers are preallocated and mutated in + place, and the decode position is load-bearing. +- The recurrent estimator's ``use_kv_cache`` toggle (PR1b): the cached fast path and + the corrected full-prefix baseline produce identical per-step logits, both equal to + a teacher-forced forward, and the fixed-length guard fires on ragged batches. + +These are the gates the spec calls "load-bearing": they turn red on the previous +(grow-both) recurrent path and green on the corrected implementation. +""" + +from typing import Literal + +import pytest +import torch + +from gfn.estimators import RecurrentDiscretePolicyEstimator +from gfn.gym.bitSequence import BitSequence +from gfn.samplers import Sampler +from gfn.utils.modules import TransformerDiscreteSequenceModel + +ATOL = 1e-5 + + +def _make_model( + vocab_size: int = 7, + embedding_dim: int = 16, + num_heads: int = 2, + num_layers: int = 2, + max_position_embeddings: int = 32, + positional_embedding: Literal["learned", "sinusoidal"] = "learned", +) -> TransformerDiscreteSequenceModel: + model = TransformerDiscreteSequenceModel( + vocab_size=vocab_size, + embedding_dim=embedding_dim, + num_heads=num_heads, + ff_hidden_dim=2 * embedding_dim, + num_layers=num_layers, + max_position_embeddings=max_position_embeddings, + dropout=0.0, + positional_embedding=positional_embedding, + ) + model.eval() + return model + + +# -------------------------------------------------------------------------------- +# PR1a: module-level in-place preallocated carry +# -------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "batch, num_heads, num_layers, seq_len", + [(1, 1, 1, 4), (3, 2, 2, 6), (2, 4, 3, 9)], +) +@pytest.mark.parametrize("positional_embedding", ["learned", "sinusoidal"]) +def test_inplace_cache_matches_full_forward( + batch: int, + num_heads: int, + num_layers: int, + seq_len: int, + positional_embedding: Literal["learned", "sinusoidal"], +) -> None: + """Incremental in-place decode == a single full-sequence forward (dense gate).""" + vocab_size = 7 + model = _make_model( + vocab_size=vocab_size, + embedding_dim=4 * num_heads, + num_heads=num_heads, + num_layers=num_layers, + max_position_embeddings=seq_len + 2, + positional_embedding=positional_embedding, + ) + device = torch.device("cpu") + tokens = torch.randint(0, vocab_size, (batch, seq_len)) + + with torch.no_grad(): + carry = model.init_carry(batch, device) + logits_full, _ = model(tokens, carry) + + carry = model.init_carry(batch, device, max_len=seq_len) + step_logits = [] + for t in range(seq_len): + lo, carry = model(tokens[:, t : t + 1], carry) + step_logits.append(lo) + logits_inplace = torch.cat(step_logits, dim=1) + + assert torch.allclose(logits_full, logits_inplace, atol=ATOL) + + +def test_inplace_cache_is_preallocated_and_mutated_in_place() -> None: + """Buffers are horizon-length from step 0, the same object across steps, and the + cursor advances one slot per decoded token.""" + vocab_size, batch, seq_len = 7, 3, 6 + model = _make_model(vocab_size=vocab_size, max_position_embeddings=seq_len + 2) + device = torch.device("cpu") + tokens = torch.randint(0, vocab_size, (batch, seq_len)) + + with torch.no_grad(): + carry = model.init_carry(batch, device, max_len=seq_len) + key0_obj = carry["key_0"] + # Preallocated to the full horizon from the very first step. + assert carry["key_0"].size(2) == seq_len + assert carry["value_0"].size(2) == seq_len + assert int(carry["cache_len"].item()) == 0 + + for t in range(seq_len): + _, carry = model(tokens[:, t : t + 1], carry) + assert carry["key_0"] is key0_obj # same tensor object, mutated in place + assert carry["key_0"].size(2) == seq_len # never reallocated / grown + assert int(carry["cache_len"].item()) == t + 1 + + +def test_decode_position_is_load_bearing() -> None: + """Perturbing the decode position by +1 must break equivalence (guards abs-pos).""" + vocab_size, batch, seq_len = 7, 2, 5 + model = _make_model(vocab_size=vocab_size, max_position_embeddings=seq_len + 2) + device = torch.device("cpu") + tokens = torch.randint(0, vocab_size, (batch, seq_len)) + + with torch.no_grad(): + carry_ref = model.init_carry(batch, device, max_len=seq_len) + logits_ref, _ = model(tokens[:, 0:1], carry_ref) + + carry_bad = model.init_carry(batch, device, max_len=seq_len) + carry_bad["position"] = carry_bad["position"] + 1 # off-by-one + logits_bad, _ = model(tokens[:, 0:1], carry_bad) + + assert not torch.allclose(logits_ref, logits_bad, atol=ATOL) + + +# -------------------------------------------------------------------------------- +# PR1b: estimator-level use_kv_cache toggle +# -------------------------------------------------------------------------------- + + +def _env_and_model(word_size: int = 2, seq_size: int = 12): + env = BitSequence( + word_size=word_size, + seq_size=seq_size, + n_modes=2, + device_str="cpu", + seed=0, + debug=False, + ) + model = _make_model( + vocab_size=env.n_actions, + num_heads=2, + num_layers=2, + max_position_embeddings=env.words_per_seq + 4, + ) + return env, model + + +def _per_step_logits(env, est, words: torch.Tensor) -> torch.Tensor: + """Drive the estimator over reconstructed per-step states (0..L words committed).""" + batch = words.shape[0] + device = words.device + carry = est.init_carry(batch, device) + out = [] + with torch.no_grad(): + for t in range(env.words_per_seq + 1): + tensor = words.clone() + tensor[:, t:] = -1 + logits, carry = est(env.States(tensor), carry) + out.append(logits) + return torch.stack(out, dim=0) # (L+1, B, n_actions) + + +def test_estimator_cache_equivalence() -> None: + """use_kv_cache OFF == ON == a single teacher-forced forward, at the logit level.""" + env, model = _env_and_model() + batch, L = 4, env.words_per_seq + words = torch.randint(0, env.n_actions - 1, (batch, L)) + + est_off = RecurrentDiscretePolicyEstimator( + module=model, n_actions=env.n_actions, use_kv_cache=False + ) + est_on = RecurrentDiscretePolicyEstimator( + module=model, n_actions=env.n_actions, use_kv_cache=True + ) + logits_off = _per_step_logits(env, est_off, words) + logits_on = _per_step_logits(env, est_on, words) + + with torch.no_grad(): + bos = torch.full((batch, 1), env.n_actions, dtype=torch.long) + seq = torch.cat([bos, words], dim=1) + tf_logits, _ = model(seq, model.init_carry(batch, torch.device("cpu"))) + tf_logits = tf_logits.transpose(0, 1) # (L+1, B, n_actions) + + assert torch.allclose(logits_off, tf_logits, atol=ATOL) + assert torch.allclose(logits_on, tf_logits, atol=ATOL) + assert torch.allclose(logits_off, logits_on, atol=ATOL) + + +def test_estimator_fixed_length_guard() -> None: + """The cached path rejects ragged (variable-length) active batches (debug on).""" + env, model = _env_and_model() + # The equal-length check is gated behind debug (it forces a per-step device sync). + est = RecurrentDiscretePolicyEstimator( + module=model, n_actions=env.n_actions, use_kv_cache=True, debug=True + ) + carry = est.init_carry(3, torch.device("cpu")) + ragged = torch.randint(0, env.n_actions - 1, (3, env.words_per_seq)) + ragged[0, 1:] = -1 # row 0 has a different committed length + with pytest.raises(ValueError, match="equal-length"): + est(env.States(ragged), carry) + + +def test_cache_max_len_validated_eagerly() -> None: + """A cache_max_len beyond the model's position range fails at construction.""" + env, model = _env_and_model() + with pytest.raises(ValueError, match="max_position_embeddings"): + RecurrentDiscretePolicyEstimator( + module=model, + n_actions=env.n_actions, + use_kv_cache=True, + cache_max_len=model.max_position_embeddings + 10, + ) + + +def test_sampling_runs_in_both_modes() -> None: + """Both toggle settings produce valid fixed-length trajectories end-to-end.""" + env, model = _env_and_model() + for use_cache in (False, True): + est = RecurrentDiscretePolicyEstimator( + module=model, n_actions=env.n_actions, use_kv_cache=use_cache + ) + with torch.no_grad(): + trajs = Sampler(estimator=est).sample_trajectories( + env, n=5, save_logprobs=True, save_estimator_outputs=False + ) + # Fixed-length: every trajectory has the same number of steps. + assert trajs.max_length == env.words_per_seq + 1 diff --git a/tutorials/examples/benchmark_kv_cache.py b/tutorials/examples/benchmark_kv_cache.py new file mode 100644 index 00000000..ff5897cb --- /dev/null +++ b/tutorials/examples/benchmark_kv_cache.py @@ -0,0 +1,344 @@ +#!/usr/bin/env python +r"""Instructional benchmark: KV-caching for autoregressive GFlowNet sampling. + +Run me: + + python tutorials/examples/benchmark_kv_cache.py + python tutorials/examples/benchmark_kv_cache.py --seq_sizes 16 32 64 128 --plot out.png + +I sample the *same* transformer policy on ``BitSequence`` two ways and time both, so the +wall-clock difference is exactly the KV-cache. No GPU needed -- the win is a FLOP +reduction and shows up on CPU (and grows with sequence length). + +================================================================================ +WHY KV-CACHING HELPS +================================================================================ +A GFlowNet policy over sequences is *autoregressive*: to choose token ``t`` it runs a +transformer over the tokens chosen so far and reads the logits at the last position. + +Inside each attention layer, the query at position ``t`` attends to the keys/values +(K/V) of every earlier position ``0..t``. Crucially, the K/V of positions ``0..t-1`` +are a function of tokens that are *already fixed* -- they do not change as the sequence +grows. Yet naive sampling recomputes the whole prefix from scratch at every step: + + step t=3, NO cache: embed+attend [w0 w1 w2] <- recompute K/V for w0,w1 again + step t=3, cache: embed+attend [ w2] <- compute K/V for w2 only, + reuse cached K/V for w0,w1 + +So a length-``L`` rollout without a cache re-embeds an ever-growing prefix each step: +roughly ``sum_t O(t^2) = O(L^3)`` attention FLOPs. A cache stores each position's K/V +the first time it is computed and reuses it, so step ``t`` only computes the new +token's Q/K/V and attends it against the ``t`` cached keys: ``sum_t O(t) = O(L^2)``. +That is a ~``L``x speedup that *widens with sequence length* -- the effect this script +makes visible. + +================================================================================ +HOW torchgfn IMPLEMENTS IT +================================================================================ +The seam is the *recurrent carry* threaded through sampling by ``Sampler``: + + RecurrentDiscretePolicyEstimator.forward(states, carry) -> (logits, carry) + +The carry is a ``dict[str, Tensor]`` holding, per transformer layer, the running K/V +cache, plus a scalar write cursor and the absolute position. Two forward paths (both +build the same ``[BOS, w0, ..., w_{t-1}]`` token stream and are numerically identical +to one full-sequence forward): + + * ``use_kv_cache=False`` (default): re-embed the full committed prefix with a fresh + carry each step. Correct and grad-bearing, but O(L^3) -- the honest baseline. + + * ``use_kv_cache=True``: feed only the newest token; ``TransformerDiscreteSequenceModel`` + writes its K/V *in place* into a preallocated ``[B, n_heads, max_len, head_dim]`` + buffer at the cursor (no ``torch.cat``, no realloc) and attends against the filled + prefix. This is the O(L^2) fast path. + +The autograd catch: writing K/V in place mutates a buffer whose earlier slice was saved +for a previous step's backward, which invalidates the autograd graph. So the cached path +runs under ``torch.no_grad()`` -- which is fine, because *sampling only needs sampled +actions, not gradients*. This makes ``use_kv_cache=True`` ideal for inference/rollout; +to *train* with it, recompute the loss log-probs with a grad-bearing forward +(``recalculate_all_logprobs=True``). A companion "teacher-forced" parallel recompute +makes that efficient and turns this sampling win into an end-to-end training win. + +================================================================================ +WHAT THIS SCRIPT MEASURES +================================================================================ +For each sequence length it (1) asserts the two configs produce identical per-step +logits (``|Δ logits| ~ 1e-7`` -- proof it is the same model), then times pure rollout +under ``no_grad`` for each -- isolating the cache (O(L^3) vs O(L^2)). +""" + +from __future__ import annotations + +import argparse +import statistics +import time +from typing import Callable + +import torch + +from gfn.estimators import RecurrentDiscretePolicyEstimator +from gfn.gflownet import TBGFlowNet +from gfn.gym.bitSequence import BitSequence +from gfn.utils.common import set_seed +from gfn.utils.modules import TransformerDiscreteSequenceModel + + +def _sync(device: torch.device) -> None: + """Block until queued device work finishes so timings are accurate. + + GPU/MPS kernels launch asynchronously, so ``perf_counter`` would otherwise time + only the Python dispatch, not the compute. A no-op on CPU. + """ + if device.type == "cuda": + torch.cuda.synchronize() + elif device.type == "mps": + torch.mps.synchronize() + + +def build(args, seq_size: int, use_kv_cache: bool, device: torch.device): + """Construct ``(env, gflownet)`` for one benchmark configuration. + + This is the only place the KV-cache API appears -- everything else (sampler, loss, + optimizer) is the ordinary torchgfn stack, unchanged. + + Args: + args: Parsed CLI arguments (model/env hyperparameters). + seq_size: BitSequence length; with ``word_size=1`` this equals the number of + autoregressive decode steps per trajectory (the ``L`` in the docstring). + use_kv_cache: Selects the sampling path (see module docstring). + device: Torch device. + + Returns: + The environment and a ``TBGFlowNet`` wrapping a transformer forward policy. + """ + env = BitSequence( + word_size=args.word_size, + seq_size=seq_size, + n_modes=args.n_modes, + device_str=str(device), + seed=args.seed, + debug=False, + ) + # The KV-cache is preallocated to the model's position range, so it must cover the + # whole rollout: one slot per word plus the leading BOS token (+ a little headroom). + max_positions = env.words_per_seq + 2 + model = TransformerDiscreteSequenceModel( + vocab_size=env.n_actions, + embedding_dim=args.embedding_dim, + num_heads=args.num_heads, + ff_hidden_dim=args.ff_hidden_dim, + num_layers=args.num_layers, + max_position_embeddings=max_positions, + dropout=0.0, + ).to(device) + est = RecurrentDiscretePolicyEstimator( + module=model, + n_actions=env.n_actions, + is_backward=False, + # The single switch that turns on the in-place, feed-newest-token fast path. + use_kv_cache=use_kv_cache, + ) + # Tree-structured DAG (each string has a unique parent) => backward policy is + # constant, hence pb=None + constant_pb=True. Only the forward policy is timed. + gflownet = TBGFlowNet(pf=est, pb=None, init_logZ=0.0, constant_pb=True).to(device) + return env, gflownet + + +def _median_ms(fn: Callable[[], None], device: torch.device, warmup: int, iters: int): + """Return the median (min, max) wall-clock of ``fn`` in milliseconds. + + Warm-up iterations are discarded so one-time costs (lazy allocations, the first + preallocation of the cache buffers) do not pollute the measurement. + """ + for _ in range(warmup): + fn() + _sync(device) + samples = [] + for _ in range(iters): + _sync(device) + t0 = time.perf_counter() + fn() + _sync(device) + samples.append((time.perf_counter() - t0) * 1e3) + lo, hi = min(samples), max(samples) + return statistics.median(samples), lo, hi + + +def time_sampling(gflownet, env, batch: int, device, warmup, iters): + """Time pure rollout under ``no_grad`` -- the operation the cache directly speeds up. + + Both configs run under ``no_grad`` here, so this isolates the cache: identical model, + identical outputs, the only difference is O(L^3) re-forward vs O(L^2) cached decode. + """ + + def _step(): + with torch.no_grad(): + gflownet.sample_trajectories( + env, n=batch, save_logprobs=False, save_estimator_outputs=False + ) + + return _median_ms(_step, device, warmup, iters) + + +def check_equivalence(args, seq_size: int, device: torch.device) -> float: + """Prove the cache changes speed, not results, at the per-step logit level. + + We reconstruct the states at every decode step of a fixed batch and run the + estimator over them both ways (``use_kv_cache`` False and True). Because the cached + incremental decode equals the full-prefix forward, the logits must agree to + floating-point tolerance; a large gap would mean the cache computes the wrong thing. + Returns the max absolute logit difference. + """ + set_seed(args.seed) + env, gflownet = build(args, seq_size, use_kv_cache=False, device=device) + # Both estimators wrap the SAME module weights, so the only difference is the path. + model = gflownet.pf.module + est_off = RecurrentDiscretePolicyEstimator( + module=model, n_actions=env.n_actions, use_kv_cache=False + ) + est_on = RecurrentDiscretePolicyEstimator( + module=model, n_actions=env.n_actions, use_kv_cache=True + ) + + batch, L = args.batch_size, env.words_per_seq + words = torch.randint(0, env.n_actions - 1, (batch, L), device=device) + + def per_step_logits(est) -> torch.Tensor: + carry = est.init_carry(batch, device) + out = [] + with torch.no_grad(): + for t in range(L + 1): # steps 0..L (L+1 decode steps) + tensor = words.clone() + tensor[:, t:] = -1 # only first t words committed + logits, carry = est(env.States(tensor), carry) + out.append(logits) + return torch.stack(out, dim=0) + + return (per_step_logits(est_off) - per_step_logits(est_on)).abs().max().item() + + +def main(args) -> None: + device = torch.device(args.device) + print( + f"device={device} batch={args.batch_size} " + f"model=(dim={args.embedding_dim}, heads={args.num_heads}, " + f"layers={args.num_layers}) word_size={args.word_size}\n" + ) + + # Columns: sampling base/cache + speedup, and the equivalence residual proving both + # columns describe the same model. + header = ( + f"{'seq_len':>7} | {'sample base':>12} {'sample cache':>13} " + f"{'speedup':>8} | {'|Δlogits|':>10}" + ) + print(header) + print("-" * len(header)) + + rows = [] + for seq_size in args.seq_sizes: + set_seed(args.seed) + + # Correctness gate first: never report a speedup for a model that changed. + max_delta = check_equivalence(args, seq_size, device) + assert ( + max_delta < 1e-4 + ), f"baseline vs cached logits differ by {max_delta} at seq_size={seq_size}" + + # Fresh models per config so neither warms the other's caches/allocator. + env_b, gfn_b = build(args, seq_size, use_kv_cache=False, device=device) + env_c, gfn_c = build(args, seq_size, use_kv_cache=True, device=device) + + s_base = time_sampling( + gfn_b, env_b, args.batch_size, device, args.warmup, args.iters + ) + s_cache = time_sampling( + gfn_c, env_c, args.batch_size, device, args.warmup, args.iters + ) + speedup = s_base[0] / s_cache[0] + rows.append((seq_size, s_base[0], s_cache[0], speedup)) + print( + f"{seq_size:>7} | {s_base[0]:>10.2f}ms {s_cache[0]:>11.2f}ms " + f"{speedup:>7.1f}x | {max_delta:>10.1e}" + ) + + print( + "\nBoth configs are numerically identical (see |Δlogits|); the speedup is the " + "KV-cache.\nIt isolates the sampling win (O(L^3)->O(L^2)) and widens with " + "seq_len -- that is the point.\nTo turn this into an end-to-end training " + "speedup, pair it with the teacher-forced loss recompute." + ) + + if args.plot: + _save_plot(rows, args.plot) + + +def _save_plot(rows, path: str) -> None: + """Optionally save a sampling-time / speedup figure (requires matplotlib).""" + try: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + except ImportError: + print("matplotlib not available; skipping --plot") + return + + seq = [r[0] for r in rows] + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4)) + ax1.plot(seq, [r[1] for r in rows], "o-", label="baseline (no cache)") + ax1.plot(seq, [r[2] for r in rows], "s-", label="KV-cache") + ax1.set( + xlabel="sequence length (steps)", + ylabel="ms / rollout", + title="Pure sampling time", + ) + ax1.legend() + ax1.grid(alpha=0.3) + ax2.plot(seq, [r[3] for r in rows], "o-") + ax2.axhline(1.0, color="gray", ls="--", lw=0.8) + ax2.set( + xlabel="sequence length (steps)", + ylabel="speedup (x)", + title="KV-cache sampling speedup", + ) + ax2.grid(alpha=0.3) + fig.tight_layout() + fig.savefig(path, dpi=120) + print(f"saved plot to {path}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Benchmark KV-cached vs baseline autoregressive GFlowNet sampling.", + ) + parser.add_argument("--device", type=str, default="cpu", help="cpu | mps | cuda") + parser.add_argument( + "--seq_sizes", + type=int, + nargs="+", + default=[8, 16, 32, 64], + help="BitSequence lengths to sweep (== rollout lengths with word_size=1)", + ) + parser.add_argument( + "--word_size", + type=int, + default=1, + help="Bits per action; 1 => one decode step per sequence bit", + ) + parser.add_argument("--n_modes", type=int, default=4) + parser.add_argument("--batch_size", type=int, default=32) + parser.add_argument("--embedding_dim", type=int, default=64) + parser.add_argument("--num_heads", type=int, default=4) + parser.add_argument("--ff_hidden_dim", type=int, default=128) + parser.add_argument("--num_layers", type=int, default=2) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--iters", type=int, default=5) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--plot", + type=str, + default=None, + help="Optional path to save a speedup figure (PNG)", + ) + main(parser.parse_args())