Skip to content

fix: handle attention masks without a padding token - #1909

Open
AaravMehta-07 wants to merge 2 commits into
TransformerLensOrg:devfrom
AaravMehta-07:codex/no-pad-attention-mask-20261016
Open

AaravMehta-07 wants to merge 2 commits into
TransformerLensOrg:devfrom
AaravMehta-07:codex/no-pad-attention-mask-20261016

Conversation

@AaravMehta-07

Copy link
Copy Markdown
Contributor

Description

get_attention_mask calls tokens.ne(tokenizer.pad_token_id) whenever a tokenizer is present. A tokenizer without a padding token has pad_token_id=None, so this raises TypeError instead of attending the supplied unpadded token tensor. This also reaches TransformerBridge.input_to_embed with a left-padding tokenizer.

Return the already-created all-ones mask when no padding ID exists. This does not mutate the tokenizer vocabulary or add model embeddings. Existing left/right padding and BOS-equals-pad handling are unchanged. The annotation now also reflects the existing None-tokenizer branch.

This is an independently reproduced bug; no existing issue is being closed.

Type Of Change

  • Bug fix (non-breaking change which fixes an issue)

Validation

Executed on Python 3.12, CPU, through uv run with the local validation dependency overlay and normal project jaxtyping/beartype pytest options:

  • New helper/native-consumer regressions against unmodified production code: 5 failed with Tensor.ne(None), as expected.
  • python -m pytest -q -s tests/unit/utilities/test_tokenize_utils.py tests/unit/model_bridge/test_boot_native.py: 89 passed after the fix.
  • python -m mypy .: no issues in 388 source files.
  • make format, make check-format, and git diff --check: passed. The upstream format recipe's literal isort format . reports a missing path; the separate actual isort --check-only . succeeds. Black skips untouched notebooks because its Jupyter extra is absent; Python formatting is checked.

Tests cover left/right tokenizer settings and int32/int64 inputs; real padding/BOS controls; and a real tiny native bridge with a local WordLevel tokenizer, comparing mask, residuals and logits with an explicit all-ones mask. No model or HF load is mocked.

The full project suite, GPU/MPS checks and hosted CI have not been run locally.

Checklist

  • Added failing-before regression tests.
  • New and existing tests in the two affected test modules pass locally.
  • Kept the production change focused; no interface tests were weakened.
  • Documentation: no new public feature or usage change; existing fallback behavior is extended to tokenizers without a padding sentinel.

AI Assistance

Codex assisted with investigation, implementation, regression tests, local test execution, technical review and this description. The results above were executed locally. No human code review or maintainer approval is claimed.

@jlarson4 jlarson4 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just one comment, more detail about how the tokenizer system is connected that is relevant to this bug fix. Let me know if you have any questions

def test_input_to_embed_without_pad_id_attends_every_token():
bridge = TransformerBridge.boot_native(_cfg(device="cpu"))
bridge.eval()
bridge.tokenizer = TokenizersBackend(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The tokenizer only lacks a pad id here because assigning bridge.tokenizer after a tokenizer-less boot skips configure_tokenizer, using boot_native(cfg, tokenizer=...) would have set one. The setter treats any first assignment as the __init__ one, so this path also leaves to_tokens unable to tokenize and keeps a stale tokenizer_prepends_bos (a Qwen tokenizer attached this way silently loses its BOS). Could you look at fixing the first-assignment check in the tokenizer setter in bridge_core.py, so a tokenizer attached after boot gets the same setup as one passed at boot?

@jlarson4

jlarson4 commented Nov 2, 2026

Copy link
Copy Markdown
Collaborator

@AaravMehta-07 Please let me know if you'll have a moment to wrap up this PR, or if I should find another contributor who'd be willing to finish it off. Thanks!

@AaravMehta-07

Copy link
Copy Markdown
Contributor Author

Thanks for explaining the lifecycle issue, and sorry for the delay. You're right: the previous presence/non-None test confused a first attachment after tokenizer-less boot with the preconfigured assignment inside init.

The correction initializes the property even for a tokenizer-less boot and checks whether _tokenizer has already been initialized, rather than whether its current value is non-None. Thus first attachment after boot and attachment after clearing now both run configure_tokenizer; the initial preconfigured boot assignment still avoids duplicate setup.

Added real local WordLevel tokenizer regressions comparing late attachment with boot_native(..., tokenizer=...): padding ID, BOS/EOS detection, vocabulary config and actual to_tokens output (including BOS). Both a tokenizer-less boot and a clear/reattach cycle are covered. The native input-to-embed test now exercises the correctly normalized tokenizer rather than depending on the broken setter. The low-level no-pad attention-mask guard remains defensive for direct utility callers.

Validation on Python 3.12/CPU: three late-attachment regression cases failed before the production fix. The three affected tokenizer/native-bridge test modules now pass 97 tests, including GPT-2/NeoX reassignment. Scoped Black/isort and git diff --check pass. No GPU, full-suite or hosted-CI result is claimed.

Full mypy .: no issues in 388 source files with mypy 1.18.2 and Transformers 5.13.0 (within the project's supported dependency range). The initial newest-dependency check found two unrelated-file errors with mypy 2.4.1/Transformers 5.22.2; I have not patched or suppressed them. Both Transformers versions passed the same 97 tests. Formatting was checked with Black 25.1.0 and isort 5.8.0; repository-wide formatter churn was excluded.

AI-assisted implementation, reproduction, tests and response.

This branch has not been deployed

No deployments
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