Make Transformers more torch-exportable and dynamo-friendly by IlyasMoutawwakil · Pull Request #42317 · huggingface/transformers

IlyasMoutawwakil · 2025-11-21T08:03:56Z

What does this PR do?

First proposals include:

check_with(error_type, cond, lambda: msg) instead of if cond: raise error_type(msg), which also works with torch.export/torch.compile to hint to the compiler that the condition is expected to be true at export/compile time.
vectorization of some loops / comprehension lists into traceable, optimized and non-blocking versions.

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline,
Pull Request section?
Was this discussed/approved via a Github issue or the forum? Please add a link
to it if that's the case.
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

HuggingFaceDocBuilderDev · 2025-11-21T08:13:02Z

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.

IlyasMoutawwakil · 2025-11-21T11:35:07Z

src/transformers/models/colqwen2/modeling_colqwen2.py

-            offsets = image_grid_thw[:, 1] * image_grid_thw[:, 2]  # (num_patches_h, num_patches_w)
-            pixel_values = torch.cat(
-                [pixel_sequence[:offset] for pixel_sequence, offset in zip(pixel_values, offsets)],
-                dim=0,
-            )  # (num_patches_h * num_patches_w, pixel_values)
+            offsets = image_grid_thw[:, 1] * image_grid_thw[:, 2]  # (batch_size,)
+            arange = torch.arange(pixel_values.shape[1], device=offsets.device)  # (max_len,)
+            mask = arange.unsqueeze(0) < offsets.unsqueeze(1)  # (batch_size, max_len)
+            pixel_values = pixel_values[mask]  # (total_valid_patches, channels, height, width)


avoiding looping over tensor

very nic eindeed!

IlyasMoutawwakil · 2025-11-21T11:36:36Z

src/transformers/models/fuyu/modeling_fuyu.py

-            for patch in pixel_values
-        ]
-        return patch_embeddings
+        return self.vision_embed_tokens(pixel_values)


need opinion about this

cc @molbap maybe (looks like this was added in #27007)

Don't know why I'm seeing this only now 👴 from what I remember pixel_values for that model is a list of Tensors hence the weird list comp, if tests pass however it should be ~ok!

Copilot

Pull Request Overview

This PR makes Transformers more export-friendly by introducing torch_check for dynamic assertions and implementing various export-related optimizations.

Key Changes

Introduces a new torch_check utility function that wraps torch._check to enable export-friendly error checking
Replaces raise ValueError with torch_check across numerous models for runtime validation
Implements performance optimizations including vectorizing batch operations, simplifying list comprehensions, and fixing instance variable assignments
Corrects error messages (e.g., "Videos features and image tokens" → "Video features and video tokens")
Adds proper training guards for weight clamping operations

Reviewed Changes

Copilot reviewed 87 out of 87 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
src/transformers/utils/import_utils.py	Adds `torch_check` function wrapper around `torch._check`
src/transformers/utils/init.py	Exports the new `torch_check` function
src/transformers/models//modeling_.py	Replaces ValueError raises with torch_check calls (50+ files)
src/transformers/models/idefics3/modeling_idefics3.py	Vectorizes position embedding computation from loop to batched operations
src/transformers/models/llava_next_video/modeling_llava_next_video.py	Fixes bug where instance variables were set in forward method
src/transformers/models/timesfm/modeling_timesfm.py	Simplifies frequency handling from loop to slice operation
src/transformers/models/tapas/modeling_tapas.py	Fixes tensor shape construction bug
src/transformers/models/ctrl/modeling_ctrl.py	Converts pos_encoding to registered buffer
src/transformers/models/gemma3n/modeling_gemma3n.py	Guards weight clamping with training check
src/transformers/models/fuyu/modeling_fuyu.py	Simplifies get_image_features to remove unnecessary list comprehension
src/transformers/models/dac/modeling_dac.py	Adds explicit dtype to torch.full call
src/transformers/models/colqwen2/modeling_colqwen2.py	Vectorizes pixel value filtering with mask-based indexing
src/transformers/models/biogpt/modeling_biogpt.py	Simplifies position_ids computation

src/transformers/utils/import_utils.py

ArthurZucker

My main comment is to use good default when you define the checking function this way most of the cases were use it or gonna be very simple.

Otherwise would be nice to ducment the good practices that you expose here, and potentially add a test in make repo-fix for simple rules.

Great work!

ArthurZucker · 2026-01-09T15:04:54Z

examples/modular-transformers/modeling_new_task_model.py

-                f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
-            )
+        special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+        check_with(


I think this needs a better name! something that says "torch_compile_check" something explicit for users as to why we use this!

I will name it torch_compilable_check as it is compilable without being bound to torch.compile, tell me if it works for you

ArthurZucker · 2026-01-09T15:06:50Z

src/transformers/models/colqwen2/modeling_colqwen2.py

-            offsets = image_grid_thw[:, 1] * image_grid_thw[:, 2]  # (num_patches_h, num_patches_w)
-            pixel_values = torch.cat(
-                [pixel_sequence[:offset] for pixel_sequence, offset in zip(pixel_values, offsets)],
-                dim=0,
-            )  # (num_patches_h * num_patches_w, pixel_values)
+            offsets = image_grid_thw[:, 1] * image_grid_thw[:, 2]  # (batch_size,)
+            arange = torch.arange(pixel_values.shape[1], device=offsets.device)  # (max_len,)
+            mask = arange.unsqueeze(0) < offsets.unsqueeze(1)  # (batch_size, max_len)
+            pixel_values = pixel_values[mask]  # (total_valid_patches, channels, height, width)


very nic eindeed!

ArthurZucker · 2026-01-19T10:55:20Z

examples/modular-transformers/modeling_new_task_model.py

+            lambda: f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {n_image_features}",
+        )


Given that you defined the function check with I think we should not have to use lambda here

yes we can support both str and lambda returning a string (for when we want the message to only be evaluated if cond is false)

ArthurZucker · 2026-01-19T10:55:56Z

examples/modular-transformers/modeling_test_detr.py

-                "Make sure to align the spatial shapes with the sequence length of the encoder hidden states"
-            )
+        check_with(
+            ValueError,


We should put the value error as a default because it seems to be used everywhere this way the more common cases were checked with function is used will be simplified

makes sense !

ArthurZucker · 2026-01-19T10:58:45Z

src/transformers/models/mm_grounding_dino/modeling_mm_grounding_dino.py

+    position_ids = torch.clamp(position_ids, min=0).to(torch.long)

-    return attention_mask, position_ids.to(torch.long)
+    return attention_mask, position_ids


very nice work here!

ArthurZucker · 2026-01-19T11:03:54Z

src/transformers/models/zamba/modeling_zamba.py

+                if attention_mask is not None:
                    hidden_states = hidden_states * attention_mask[:, -hidden_states.shape[-1] :].unsqueeze(1)
                conv_state = nn.functional.pad(hidden_states, (self.conv_kernel_size - hidden_states.shape[-1], 0))
                cache_params.conv_states[self.layer_idx] = conv_state
                hidden_states = self.act(self.conv1d(hidden_states)[..., :seq_len])
-                if attention_mask is not None and not torch.all(attention_mask == 1):
+                if attention_mask is not None:


this change is weird same for the next one in this file

the data-dependency on not torch.all(attention_mask == 1) breaks graphs, I can revert the change and try to find better alternatives later (in another PR).

no I mean look at the two if else

ArthurZucker · 2026-01-19T11:05:32Z

src/transformers/utils/import_utils.py

+
+    if isinstance(cond, torch.Tensor):
+        cond = cond.item()
+    torch._check_with(error_type, cond, msg)


yeah... that does sound good actually but only if we can catch to give a good detailed error!

ArthurZucker · 2026-01-20T17:40:38Z

LGTM, now just the flagged change that looks a bit weird (check the if else)

github-actions · 2026-01-21T15:05:13Z

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

run-slow: aria, aya_vision, bart, bigbird_pegasus, biogpt, chameleon, cohere2_vision, colqwen2, ctrl, d_fine, dac, deepseek_vl, deepseek_vl_hybrid, deformable_detr, emu3, ernie4_5_vl_moe

github-actions · 2026-01-21T15:27:32Z

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=42317&sha=99be85

…ace#42317) * make vlms export friendly * seq2seq lms * biogpt * more vlms * colqwen2 * vision models * more vlms * more vlms * more vlms * vectorized vision embedding * fixup * more vlms * more vlms * generate_masks_with_special_tokens_and_transfer_map * custom torch_check * use custom torch_check * revert grounding dino changes * fixup * remove file * undo * undo * testing * fixes * standard error message * use torch._check_with to raise value error instead of torch._check's runtime error * fix recurrent gemma * only itemize tensors * use spatial shapes list instead of tensor * fix udop use_cache default value * use tracable condition for seq2seq lms * make smolvlm exportable * fix fastvlm and t5gemma2 * fix qwen2_audio and idefics * remove script * tbc * skip mra model * helper * style and document * fix * set experts impl to batched * make xmod exportable and efficient * make more ssms exportable * fix * revert recurrent gemma * skip models that use chunked attention or rope_index * qwen3_next * assert async * tensorize (mm) grounding dino mask generation * style * fix repo * address comments * fix qwen2 audio and vits checks * skip two models using kernels by default * skip granite moe hybrid using custom kernels * disable mamba kernels * vits splinter and videomae

make vlms export friendly

1ab6216

IlyasMoutawwakil marked this pull request as draft November 21, 2025 08:04

IlyasMoutawwakil added 20 commits November 21, 2025 09:17

seq2seq lms

25456f0

biogpt

6d4f744

more vlms

b8135b4

colqwen2

2d3f986

vision models

7e03b17

more vlms

a58b049

more vlms

e776c18

more vlms

5377e2c

vectorized vision embedding

6b4c37d

fixup

44e8e47

more vlms

df1559f

more vlms

bb6b8a3

generate_masks_with_special_tokens_and_transfer_map

b606772

custom torch_check

fdc9535

use custom torch_check

0184d20

revert grounding dino changes

a4d754e

fixup

75c0fd3

remove file

713c895

undo

2a4081b

undo

73cd1ba

IlyasMoutawwakil commented Nov 21, 2025

View reviewed changes

IlyasMoutawwakil requested a review from Copilot November 21, 2025 11:39

Copilot started reviewing on behalf of IlyasMoutawwakil November 21, 2025 11:40 View session

Copilot finished reviewing on behalf of IlyasMoutawwakil November 21, 2025 11:42

Copilot AI reviewed Nov 21, 2025

View reviewed changes

src/transformers/utils/import_utils.py Outdated Show resolved Hide resolved

testing

8aea2bb

IlyasMoutawwakil and others added 7 commits January 8, 2026 14:36

fix

4431633

revert recurrent gemma

faa4540

skip models that use chunked attention or rope_index

01d93a2

qwen3_next

23f4485

assert async

60050d8

Merge branch 'main' into export-friendly

526fb5c

tensorize (mm) grounding dino mask generation

799ad58

IlyasMoutawwakil changed the title ~~Make Transformers more export-friendly~~ Make Transformers more torch-exportable and dynamo-friendly Jan 8, 2026

IlyasMoutawwakil and others added 4 commits January 9, 2026 08:24

Merge branch 'main' into export-friendly

c551dd0

Merge branch 'main' into export-friendly

ae78766

style

a6d56f1

fix repo

04dee52

ArthurZucker approved these changes Jan 19, 2026

View reviewed changes

IlyasMoutawwakil added 5 commits January 19, 2026 16:01

address comments

d3e145d

Merge branch 'main' into export-friendly

ac7c941

fix qwen2 audio and vits checks

86df05d

skip two models using kernels by default

e98872d

skip granite moe hybrid using custom kernels

b9f20c2

IlyasMoutawwakil added 2 commits January 21, 2026 09:35

disable mamba kernels

d20b998

vits splinter and videomae

b2884ef

IlyasMoutawwakil added 2 commits January 21, 2026 16:05

Merge branch 'main' into export-friendly

d45e2cd

Merge branch 'main' into export-friendly

99be85b

ArthurZucker merged commit eff263c into main Jan 22, 2026
24 of 26 checks passed

ArthurZucker deleted the export-friendly branch January 22, 2026 09:07

vasqu mentioned this pull request Jan 22, 2026

gemma3n executorch export fails missing self training guard and erfinv not supported #43412

Open

		lambda: f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {n_image_features}",
		)

Conversation

IlyasMoutawwakil commented Nov 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Before submitting

Who can review?

Uh oh!

HuggingFaceDocBuilderDev commented Nov 21, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

IlyasMoutawwakil Nov 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Key Changes

Reviewed Changes

Uh oh!

Uh oh!

ArthurZucker left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

IlyasMoutawwakil Jan 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

IlyasMoutawwakil Jan 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ArthurZucker commented Jan 20, 2026

Uh oh!

github-actions bot commented Jan 21, 2026

Uh oh!

github-actions bot commented Jan 21, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

IlyasMoutawwakil commented Nov 21, 2025 •

edited

Loading

IlyasMoutawwakil Nov 21, 2025 •

edited

Loading

IlyasMoutawwakil Jan 19, 2026 •

edited

Loading

IlyasMoutawwakil Jan 19, 2026 •

edited

Loading