SmolVLM encoder not getting delegated to XNNPack

[jackzhxng]

🐛 Describe the bug

SmolVLM from Optimum (PR - huggingface/optimum-executorch#163) is not getting fully delegated to XNNPack in the encoder subgraph.

The pte dump shows this - P1983888427, you can see undelegated addmms still in the encoder.

Printed EP right after torch.export: P1984977202
Printed EP right after to_edge_transform_and_lower: P1985017274

Versions

N/A

cc @GregoryComer @digantdesai @cbilgin @mergennachin

[digantdesai]

Thanks @jackzhxng for this. First, does this happen on v1.0 or only on main?

Looking at all three pastes, seems like the addmm nodes do indeed come from linear decomposition. Spot checked (transformers/models/idefics3/modeling_idefics3.py:231) for example.

This makes me question (without digging too deep) - to_edge_transform_and_lower do-no-decompose feature. Can you dig into that a bit more or @GregoryComer or I can take over.

[jackzhxng]

@digantdesai happening on this fairly recent ET nightly pin.

This is where to_edge_transform_and_lower is called - https://github.com/huggingface/optimum-executorch/blob/main/optimum/exporters/executorch/recipes/xnnpack.py#L84. It is called for both decoder and encoder, but the encoder is not getting fully delegated.

[GregoryComer]

I enabled whyNoPartition logging and added a bit more granular logging locally, and I'm seeing the following. (As an aside, I intend to make this logging easier to enable without making code changes soon).

ERROR:executorch.backends.xnnpack.partition.config.xnnpack_config:arg tensor has free unbacked symbols or numel == 0
ERROR:executorch.backends.xnnpack.partition.config.xnnpack_config:inputs_valid: False, outputs_valid: True
ERROR:executorch.backends.xnnpack.partition.config.xnnpack_config:WhyNoPartition: Node linear was not partitioned because invalid dtype.
ERROR:executorch.backends.xnnpack.partition.config.xnnpack_config:arg tensor has free unbacked symbols or numel == 0
ERROR:executorch.backends.xnnpack.partition.config.xnnpack_config:inputs_valid: False, outputs_valid: True
ERROR:executorch.backends.xnnpack.partition.config.xnnpack_config:WhyNoPartition: Node linear_1 was not partitioned because invalid dtype.
... (and many more)

This check is failing:

executorch/backends/xnnpack/partition/config/xnnpack_config.py

Lines 177 to 181 in 926312e

	# XNNPACK does not support empty tensors. But we can't get numel()
	# for unbacked symints, so we conservatively bail out here if any
	# dimension of the tensor is unbacked symint.
	if has_free_unbacked_symbols(arg_val) or arg_val.numel() == 0:
	return False

I'm continuing to dig...

[GregoryComer]

If I'm understanding the torch logic correctly, the problem is that an unbacked symint is introduced by the data-dependent view here:

             # File: /home/jackzhxng/.conda/envs/tr3/lib/python3.10/site-packages/transformers/models/idefics3/modeling_idefics3.py:701 in get_image_features, code: real_images_inds = (pixel_values == 0.0).sum(dim=(-1, -2, -3)) != nb_values_per_image
            eq: "b8[13, 3, 384, 384]" = torch.ops.aten.eq.Scalar(view, 0.0)
            sum_1: "i64[13]" = torch.ops.aten.sum.dim_IntList(eq, [-1, -2, -3]);  eq = None
            ne: "b8[13]" = torch.ops.aten.ne.Scalar(sum_1, 442368);  sum_1 = None
             # File: /home/jackzhxng/.conda/envs/tr3/lib/python3.10/site-packages/transformers/models/idefics3/modeling_idefics3.py:702 in get_image_features, code: pixel_values = pixel_values[real_images_inds].contiguous()
            index: "f32[u0, 3, 384, 384]" = torch.ops.aten.index.Tensor(view, [ne]);  view = ne = None

The specific check XNNPACK is failing on could be re-written to be more flexible, but I'm guessing it will fail downstream as it won't be able to memory plan since the tensor size is not directly derived from the input shape bounds? My initial guess is that this logic needs to be re-written to give statically traceable shapes to work with ET's current design.

[GregoryComer]

Here's the data-dependent code in transformers:
https://github.com/huggingface/transformers/blob/94df0e65602922be2831b3faa457a2bde78b936b/src/transformers/models/idefics3/modeling_idefics3.py#L636-L637

From a quick read, it looks like it's removing all-zero images from the batched tensor, leading to a bsz <= input bsz. The output bsz is data dependent.

[GregoryComer]

I tried locally removing the check in XNNPACK for unbacked symints and get the following error:

File "/home/gjcomer/.conda/envs/executorch/lib/python3.12/site-packages/torch/fx/experimental/symbolic_shapes.py", line 7574, in _evaluate_expr
    raise self._make_data_dependent_error(
torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode: Could not guard on data-dependent expression Ne(Mod(u0, 729*u0), 0) (unhinted: Ne(Mod(u0, 729*u0), 0)).  (Size-like symbols: u0)

I don't fully understand why it's lower-able when not delegating this node. The data dependence error only shows up when this node is delegated? In the above post-to-edge paste, there is still an unbacked symint, but export is okay with it.

 # File: /home/jackzhxng/.conda/envs/tr3/lib/python3.10/site-packages/transformers/models/idefics3/modeling_idefics3.py:702 in get_image_features, code: pixel_values = pixel_values[real_images_inds].contiguous()
aten_index_tensor: "f32[u15, 3, 384, 384]" = executorch_exir_dialects_edge__ops_aten_index_Tensor(aten_view_copy_default, [aten_ne_scalar]);  aten_view_copy_default = aten_ne_scalar = None

 # File: /home/jackzhxng/.conda/envs/tr3/lib/python3.10/site-packages/optimum/exporters/executorch/integrations.py:70 in export_friendly_forward, code: pos_ids = (bucket_coords_h[:, None] * N + bucket_coords_w[None, :]).reshape(-1)
aten_mul_tensor_2: "i64[27, 1]" = executorch_exir_dialects_edge__ops_aten_mul_Tensor(aten_unsqueeze_copy_default, _lifted_tensor_constant4);  aten_unsqueeze_copy_default = _lifted_tensor_constant4 = None

 #
sym_size: "Sym(u15)" = torch.ops.aten.sym_size.int(aten_index_tensor, 0)

 # File: /home/jackzhxng/.conda/envs/tr3/lib/python3.10/site-packages/transformers/models/idefics3/modeling_idefics3.py:284 in forward, code: return self.proj(x)
mul_1: "Sym(81*u15)" = sym_size * 81

[jackzhxng]

That is strange how export goes through cc @tugsbayasgalan @pianpwk

I think the way out here would be to just rewrite this function to assume there are no padding images, and handle any removal of padding images outside of the model in the C++ runner. I'll give it a shot.

[digantdesai]

executorch/backends/xnnpack/partition/config/xnnpack_config.py

@GregoryComer can we update return False to record Why reason here for future?