Skip to content

Fix run_with_cache BOS handling for string inputs - #1625

Merged
jlarson4 merged 1 commit into
TransformerLensOrg:dev-4.xfrom
emerardd:fix/run-with-cache-prepend-bos
Aug 10, 2026
Merged

Fix run_with_cache BOS handling for string inputs#1625
jlarson4 merged 1 commit into
TransformerLensOrg:dev-4.xfrom
emerardd:fix/run-with-cache-prepend-bos

Conversation

@emerardd

@emerardd emerardd commented Aug 8, 2026

Copy link
Copy Markdown
Contributor

Description

Summary

  • Preserve an explicit prepend_bos value when BridgeCore.run_with_cache() eagerly tokenizes a scalar string input.
  • Consume the tokenization option before forwarding the resulting token tensor, rather than leaving an ineffective prepend_bos argument on the tensor forward.
  • Add offline native-Bridge regression coverage for direct-vs-cached parity with both prepend_bos=True and False.

Root cause and impact

run_with_cache() converts scalar string inputs with self.to_tokens(input) before calling forward(), but did not pass the caller's prepend_bos value. to_tokens() therefore fell back to cfg.default_prepend_bos, typically True. Although the original argument was subsequently forwarded, the input was already a token tensor and could no longer be retokenized.

As a result, bridge(text, prepend_bos=False) and bridge.run_with_cache(text, prepend_bos=False) could run on different token sequences. The cached path silently included an extra BOS position, shifting logits and position-indexed activations used for attention analysis, attribution, and activation patching.

The fix lives in BridgeCore, so it preserves the shared local/remote cache contract while leaving tensor/list inputs and driver-aware device placement unchanged. This PR intentionally does not broaden into padding_side, whose tokenization support differs between TransformerBridge and RemoteBridge and is not required to fix the demonstrated bug.

No documentation change is required because the existing prepend_bos contract is already documented; this change makes run_with_cache() conform to it.

Fixes #1624

Type of change

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • This change requires a documentation update

Screenshots

N/A — code-only behavior fix.

Validation

  • Target regression: 2 passed.
  • Full affected files (test_boot_native.py and test_remote_bridge.py): 62 passed.
  • Related model-bridge unit surface selected by run_with_cache, tokenizer, and prepend_bos: 68 passed, 6 skipped.
  • Offline tiny native Bridge parity: direct logits, cached logits, and hook_embed all have shape (1, 1, 8) with prepend_bos=False; maximum direct-vs-cached logit difference is 0.0.
  • mypy .: no issues in 424 source files.
  • pycln, isort, Black, and git diff --check: passed for the changed files.

Checklist:

  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation (no change required; existing docs already specify this behavior)
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes (the affected unit-test surface passed; the complete suite was not rerun locally)
  • I have not rewritten tests relating to key interfaces which would affect backward compatibility

@emerardd
emerardd marked this pull request as ready for review August 8, 2026 11:26
@emerardd

emerardd commented Aug 8, 2026

Copy link
Copy Markdown
Contributor Author

The only failing check is Full Code Coverage, and both failures reproduce identically on the exact base commit (7db5f8dc) before this PR changes:

  • test_one_left_padded_row_does_not_perturb_its_neighbours
    • maximum absolute difference: 9.1552734375e-05 with atol=1e-6
  • test_bridge_hooked_parity_multi_step_optimization
    • step 1 mean logits difference: 0.0010319872526451945 with a 0.001 threshold

Base coverage run:
https://github.com/TransformerLensOrg/TransformerLens/actions/runs/31233076384/job/93040556397

PR coverage run:
https://github.com/TransformerLensOrg/TransformerLens/actions/runs/31254874477/job/93096680161

The values and failure locations are identical in both runs. This PR only changes scalar-string tokenization in run_with_cache(), while both failing tests use direct tensor forward paths.

I have therefore kept this PR focused and have not changed the unrelated numerical tolerances. Please let me know whether you would prefer the coverage job to be retried, or these base-branch failures to be handled separately.

@jlarson4
jlarson4 merged commit 956c989 into TransformerLensOrg:dev-4.x Aug 10, 2026
49 of 50 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug Report] TransformerBridge.run_with_cache ignores prepend_bos=False for string inputs

2 participants