Blockwise mask fn as opt arg in all masking functions by zucchini-nlp · Pull Request #45477 · huggingface/transformers

zucchini-nlp · 2026-04-16T14:12:03Z

What does this PR do?

As per title, I think this pattern is used quite often and deserves to be a public mask-fn. Used currently in gemma/paligemma family, GIT, PI0 and will be used in two upcoming models (deepseekOcr and Molmo2)

This PR allows all these models ton go non-vmap path, which iirc is more preferable for us.

I opted for adding 'blockwise_mask' as an arg in existing mask fn, it doesn't seems like a new mask type of itself. We might have to create otherwise - create_blockwise_causal_mask, create_blockwise_sliding_causal_mask, etc.

HuggingFaceDocBuilderDev · 2026-04-16T14:27:14Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

zucchini-nlp · 2026-04-17T13:17:21Z

run-slow: gemma3, gemma4, git, paligemma, pi0

github-actions · 2026-04-17T13:18:49Z

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/gemma3", "models/gemma4", "models/git", "models/paligemma", "models/pi0"]
quantizations: []

github-actions · 2026-04-17T14:29:31Z

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	36c679a5	workflow commit (merge commit)
PR	e8f06b29	branch commit (from PR)
main	77de8dd8	base commit (on `main`)

Model CI Report

❌ 17 new failed tests from this PR 😭

gemma3:
tests/models/gemma3/test_modeling_gemma3.py::Gemma3IntegrationTest::test_model_4b_batch (❌ ⟹ ❌)
tests/models/gemma3/test_modeling_gemma3.py::Gemma3IntegrationTest::test_model_4b_batch_crops (❌ ⟹ ❌)
tests/models/gemma3/test_modeling_gemma3.py::Gemma3IntegrationTest::test_model_4b_bf16 (❌ ⟹ ❌)
tests/models/gemma3/test_modeling_gemma3.py::Gemma3IntegrationTest::test_model_4b_crops (❌ ⟹ ❌)
tests/models/gemma3/test_modeling_gemma3.py::Gemma3IntegrationTest::test_model_4b_multiimage (❌ ⟹ ❌)
gemma4:
tests/models/gemma4/test_modeling_gemma4.py::Gemma4IntegrationTest::test_export_text_only (❌ ⟹ ❌)
paligemma:
tests/models/paligemma/test_modeling_paligemma.py::PaliGemmaForConditionalGenerationIntegrationTest::test_integration_detection_bug (❌ ⟹ ❌)
tests/models/paligemma/test_modeling_paligemma.py::PaliGemmaForConditionalGenerationIntegrationTest::test_small_model_integration_test (❌ ⟹ ❌)
tests/models/paligemma/test_modeling_paligemma.py::PaliGemmaForConditionalGenerationIntegrationTest::test_small_model_integration_test_multiimage (❌ ⟹ ❌)
tests/models/paligemma/test_modeling_paligemma.py::PaliGemmaForConditionalGenerationIntegrationTest::test_small_model_integration_test_paligemma_VQA (❌ ⟹ ❌)
tests/models/paligemma/test_modeling_paligemma.py::PaliGemmaForConditionalGenerationIntegrationTest::test_small_model_integration_test_paligemma_batched (❌ ⟹ ❌)
tests/models/paligemma/test_modeling_paligemma.py::PaliGemmaForConditionalGenerationIntegrationTest::test_small_model_integration_test_paligemma_batched_bf16 (❌ ⟹ ❌)
tests/models/paligemma/test_modeling_paligemma.py::PaliGemmaForConditionalGenerationIntegrationTest::test_small_model_integration_test_paligemma_batched_f16 (❌ ⟹ ❌)
tests/models/paligemma/test_modeling_paligemma.py::PaliGemmaForConditionalGenerationIntegrationTest::test_small_model_integration_test_paligemma_empty_prompt (❌ ⟹ ❌)
pi0:
tests/models/pi0/test_modeling_pi0.py::PI0ModelIntegrationTest::test_pi0_base_libero (❌ ⟹ ❌)
tests/models/pi0/test_modeling_pi0.py::PI0ModelIntegrationTest::test_pi0_base_reference_values (❌ ⟹ ❌)
tests/models/pi0/test_modeling_pi0.py::PI0ModelIntegrationTest::test_train_pi0_base_libero (❌ ⟹ ❌)

zucchini-nlp · 2026-04-17T14:33:10Z

oh, no, I forgot about conversion. Will wait until merged and trigger test again

zucchini-nlp · 2026-04-20T12:34:05Z

run-slow: gemma3, gemma4, git, paligemma, pi0

github-actions · 2026-04-20T12:35:24Z

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/gemma3", "models/gemma4", "models/git", "models/paligemma", "models/pi0"]
quantizations: []

github-actions · 2026-04-20T13:21:00Z

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	43e0b2ad	workflow commit (merge commit)
PR	0e17f9a4	branch commit (from PR)
main	c3ec5ff4	base commit (on `main`)

Model CI Report

❌ 6 new failed tests from this PR 😭

gemma3:
tests/models/gemma3/test_modeling_gemma3.py::Gemma3IntegrationTest::test_model_4b_multiimage (✅ ⟹ ❌)
gemma4:
tests/models/gemma4/test_modeling_gemma4.py::Gemma4IntegrationTest::test_export_text_only (❌ ⟹ ❌)
paligemma:
tests/models/paligemma/test_modeling_paligemma.py::PaliGemmaForConditionalGenerationIntegrationTest::test_integration_detection_bug (✅ ⟹ ❌)
pi0:
tests/models/pi0/test_modeling_pi0.py::PI0ModelIntegrationTest::test_pi0_base_libero (✅ ⟹ ❌)
tests/models/pi0/test_modeling_pi0.py::PI0ModelIntegrationTest::test_pi0_base_reference_values (✅ ⟹ ❌)
tests/models/pi0/test_modeling_pi0.py::PI0ModelIntegrationTest::test_train_pi0_base_libero (❌ ⟹ ❌)

zucchini-nlp · 2026-04-20T14:09:34Z

run-slow: gemma3, paligemma, pi0

github-actions · 2026-04-20T14:10:55Z

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/gemma3", "models/paligemma", "models/pi0"]
quantizations: []

github-actions · 2026-04-20T14:58:52Z

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	e73db476	workflow commit (merge commit)
PR	3617c364	branch commit (from PR)
main	0db33792	base commit (on `main`)

Model CI Report

❌ 5 new failed tests from this PR 😭

gemma3:
tests/models/gemma3/test_modeling_gemma3.py::Gemma3IntegrationTest::test_model_4b_multiimage (✅ ⟹ ❌)
paligemma:
tests/models/paligemma/test_modeling_paligemma.py::PaliGemmaForConditionalGenerationIntegrationTest::test_integration_detection_bug (✅ ⟹ ❌)
pi0:
tests/models/pi0/test_modeling_pi0.py::PI0ModelIntegrationTest::test_pi0_base_libero (✅ ⟹ ❌)
tests/models/pi0/test_modeling_pi0.py::PI0ModelIntegrationTest::test_pi0_base_reference_values (✅ ⟹ ❌)
tests/models/pi0/test_modeling_pi0.py::PI0ModelIntegrationTest::test_train_pi0_base_libero (❌ ⟹ ❌)

…le-mask

zucchini-nlp · 2026-04-23T14:58:40Z

run-slow: gemma3, paligemma, pi0

zucchini-nlp · 2026-04-23T15:03:58Z

run-slow: gemma3, paligemma, pi0

github-actions · 2026-04-23T15:05:35Z

Workflow Run ⚙️💔 This comment contains run-slow, but unknown error occurred and the workflow run aborted!

Cyrilvallez

Nice, I like having this natively!
Just a few things, most particularly I think it's best to pad the mask directly once at the beginning, similar to https://github.com/huggingface/transformers/blob/main/src/transformers/masking_utils.py#L185-L194

zucchini-nlp · 2026-04-27T10:08:49Z

run-slow: gemma3, gemma4, git, paligemma, pi0

github-actions · 2026-04-27T10:10:03Z

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/gemma3", "models/gemma4", "models/git", "models/paligemma", "models/pi0"]
quantizations: []

github-actions · 2026-04-27T10:53:38Z

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	1033cdbc	workflow commit (merge commit)
PR	9b9706ce	branch commit (from PR)
main	bbb51c83	base commit (on `main`)

Model CI Report

❌ 3 new failed tests from this PR 😭

gemma4:
tests/models/gemma4/test_modeling_gemma4.py::Gemma4IntegrationTest::test_export_text_only (❌ ⟹ ❌)
pi0:
tests/models/pi0/test_modeling_pi0.py::PI0ModelIntegrationTest::test_pi0_base_reference_values (✅ ⟹ ❌)
tests/models/pi0/test_modeling_pi0.py::PI0ModelIntegrationTest::test_train_pi0_base_libero (❌ ⟹ ❌)

zucchini-nlp · 2026-04-27T12:25:19Z

Addressed the comments @Cyrilvallez

vasqu · 2026-04-27T13:33:09Z

I think we wait on Cyril for a final pass? I can also take a look if you want but I don't think it's rushing

Cyrilvallez

I believe we still need to take the offset into account (see comment), but otherwise looks good to me! cc @vasqu here for last check and approval as you checked it as well!

github-actions · 2026-04-28T10:10:41Z

[For maintainers] Suggested jobs to run (before merge)

run-slow: gemma3, gemma4, git, paligemma, pi0

zucchini-nlp · 2026-04-28T10:11:16Z

final pass with slow CI on rebased branch

zucchini-nlp · 2026-04-28T10:11:38Z

run-slow: gemma3, gemma4, git, paligemma, pi0

github-actions · 2026-04-28T10:12:57Z

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/gemma3", "models/gemma4", "models/git", "models/paligemma", "models/pi0"]
quantizations: []

github-actions · 2026-04-28T13:20:45Z

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=45477&sha=c4c5ce

zucchini-nlp · 2026-04-29T09:35:14Z

run-slow: gemma3, gemma4, git, paligemma, pi0

github-actions · 2026-04-29T09:36:45Z

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/gemma3", "models/gemma4", "models/git", "models/paligemma", "models/pi0"]
quantizations: []

github-actions · 2026-04-29T10:18:30Z

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	4e759e21	workflow commit (merge commit)
PR	c4c5ce1a	branch commit (from PR)
main	9120f5e4	base commit (on `main`)

Model CI Report

❌ 3 new failed tests from this PR 😭

gemma4:
tests/models/gemma4/test_modeling_gemma4.py::Gemma4IntegrationTest::test_export_text_only (❌ ⟹ ❌)
pi0:
tests/models/pi0/test_modeling_pi0.py::PI0ModelIntegrationTest::test_pi0_base_libero (✅ ⟹ ❌)
tests/models/pi0/test_modeling_pi0.py::PI0ModelIntegrationTest::test_train_pi0_base_libero (❌ ⟹ ❌)

zucchini-nlp added 5 commits April 7, 2026 19:54

works i think

9368ece

more models

513035c

like this

0bb610a

pi0

d527a15

Merge branch 'main' into split-out-gemma-style-mask

fe59952

zucchini-nlp added 3 commits April 17, 2026 13:57

Merge branch 'main' into split-out-gemma-style-mask

1099438

fix some

6fa84ea

fix more

19c3066

zucchini-nlp changed the title ~~[WIP] Add blockwise mask fn as opt arg for all masking functions~~ Blockwise mask fn as opt arg for all masking functions Apr 17, 2026

zucchini-nlp changed the title ~~Blockwise mask fn as opt arg for all masking functions~~ Blockwise mask fn as opt arg in all masking functions Apr 17, 2026

.

e8f06b2

Merge branch 'main' into split-out-gemma-style-mask

0e17f9a

for now

3617c36

zucchini-nlp added 3 commits April 23, 2026 16:38

assumes always prefix, unless type ids are passed

96ef824

Merge remote-tracking branch 'upstream/main' into split-out-gemma-sty…

add5818

…le-mask

ig it is sdpa choosing a different backend

7286713

Cyrilvallez reviewed Apr 27, 2026

View reviewed changes

Comment thread src/transformers/masking_utils.py Outdated

Comment thread src/transformers/masking_utils.py Outdated

Comment thread src/transformers/masking_utils.py

Comment thread src/transformers/masking_utils.py

zucchini-nlp added 3 commits April 27, 2026 11:56

comments: pad and move around

9555c59

forward in generation

ef1c767

fix repo

9b9706c

prepare token types when sampling as well

c1178aa

vasqu reviewed Apr 27, 2026

View reviewed changes

Comment thread src/transformers/masking_utils.py Outdated

consume as unused kwags

d46668a

Cyrilvallez reviewed Apr 28, 2026

View reviewed changes

Comment thread src/transformers/masking_utils.py Outdated

zucchini-nlp added 2 commits April 28, 2026 12:09

mirror padding mask and add kv-offset

c8fb8fd

oops

be84c3a

Merge branch 'main' into split-out-gemma-style-mask

89ababc

vasqu mentioned this pull request Apr 28, 2026

Add Deepseek-OCR-2 model #45075

Open

wait no

82cf4d0

huggingface deleted a comment from github-actions Bot Apr 28, 2026

Merge branch 'main' into split-out-gemma-style-mask

c4c5ce1

evalstate mentioned this pull request Apr 28, 2026

Cumulative defect fixes from recent Transformers PRs evalstate/transformers#41

Open

zucchini-nlp requested a review from vasqu April 29, 2026 09:35

Conversation

zucchini-nlp commented Apr 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Uh oh!

HuggingFaceDocBuilderDev commented Apr 16, 2026

Uh oh!

zucchini-nlp commented Apr 17, 2026

Uh oh!

github-actions Bot commented Apr 17, 2026

Uh oh!

github-actions Bot commented Apr 17, 2026

CI Results

Commit Info

Model CI Report

Uh oh!

zucchini-nlp commented Apr 17, 2026

Uh oh!

zucchini-nlp commented Apr 20, 2026

Uh oh!

github-actions Bot commented Apr 20, 2026

Uh oh!

github-actions Bot commented Apr 20, 2026

CI Results

Commit Info

Model CI Report

Uh oh!

zucchini-nlp commented Apr 20, 2026

Uh oh!

github-actions Bot commented Apr 20, 2026

Uh oh!

github-actions Bot commented Apr 20, 2026

CI Results

Commit Info

Model CI Report

Uh oh!

zucchini-nlp commented Apr 23, 2026

Uh oh!

zucchini-nlp commented Apr 23, 2026

Uh oh!

github-actions Bot commented Apr 23, 2026

Uh oh!

Cyrilvallez left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zucchini-nlp commented Apr 27, 2026

Uh oh!

github-actions Bot commented Apr 27, 2026

Uh oh!

github-actions Bot commented Apr 27, 2026

CI Results

Commit Info

Model CI Report

Uh oh!

zucchini-nlp commented Apr 27, 2026

Uh oh!

Uh oh!

vasqu commented Apr 27, 2026

Uh oh!

Cyrilvallez left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

github-actions Bot commented Apr 28, 2026

Uh oh!

zucchini-nlp commented Apr 28, 2026

Uh oh!

zucchini-nlp commented Apr 28, 2026

Uh oh!

github-actions Bot commented Apr 28, 2026

Uh oh!

github-actions Bot commented Apr 28, 2026

Uh oh!

zucchini-nlp commented Apr 29, 2026

Uh oh!

github-actions Bot commented Apr 29, 2026

zucchini-nlp commented Apr 16, 2026 •

edited

Loading