Skip to content

fix: preserve sampling penalties in stateful generation fallback - #1888

Merged
jlarson4 merged 1 commit into
TransformerLensOrg:devfrom
emerardd:fix/stateful-generation-fallback-contract
Oct 12, 2026
Merged

jlarson4 merged 1 commit into
TransformerLensOrg:devfrom
emerardd:fix/stateful-generation-fallback-contract

Conversation

@emerardd

Copy link
Copy Markdown
Contributor

Description

Fix silent loss of sampling penalties when TransformerBridge.generate() delegates a stateful/SSM model to hf_generate() with use_past_kv_cache=False.

The fallback builds an allowlist of HF arguments but omits repetition_penalty. This makes an explicit penalty silently revert to the HF generation default. Forward it unchanged, including the neutral value, so the public Bridge argument determines the requested behavior.

HF does not expose a directly equivalent argument for Bridge's additive, token-count-based freq_penalty. Reject nonzero values on this fallback with an actionable NotImplementedError directing callers to the hooked stateful cache path, rather than silently ignoring the request. The generate() docstring now documents this limitation. Zero frequency penalty and the normal hooked generation loop are unchanged.

This is deliberately limited to sampling penalties. It does not resolve other fallback return-format or Bridge-only output options, redesign caching, or modify adapters. No new dependency is required.

Type of change

  • Bug fix

Validation

  • Before the production fix, the five new forwarding/rejection checks failed; the four existing fallback tests passed.
  • After the fix: 39 passed across test_generate_fallback_flags.py, test_hf_generate_batch_padding.py, and utilities/test_logits_utils.py.
  • A deterministic tiny GPT-2 generated sequence matches direct HF generation with repetition_penalty=2.0 and differs from the unpenalized sequence.
  • Dispatch tests use a tiny config-built GPT-2 with the stateful flag enabled to exercise the fallback, not a downloaded or mocked model load. This validates argument dispatch and HF behavior, not real SSM kernel execution.
  • Scoped pinned pycln, isort, Black and git diff --check passed. Full-repository mypy passed for 391 source files.
  • The full unit suite was not run. Two SWIG import deprecation warnings were emitted by the focused tests; no new warning path was added.

Checklist

  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation (API docstring)
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective
  • New and existing unit tests pass locally with my changes (affected surface passed; full suite not run)
  • I have not rewritten tests relating to key interfaces which would affect backward compatibility

@jlarson4

Copy link
Copy Markdown
Collaborator

Great solution, thanks @emerardd

@jlarson4
jlarson4 merged commit 0973df0 into TransformerLensOrg:dev Oct 12, 2026
27 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.

2 participants