Skip to content

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

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

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

Conversation

@AaravMehta-07

Copy link
Copy Markdown

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?

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