fix issues

Author: xadupre

_unittests/ut_torch_export_patches/test_patch_torch.py

@@ -4,6 +4,9 @@
 from torch._dynamo._trace_wrapped_higher_order_op import TransformGetItemToIndex
 from onnx_diagnostic.ext_test_case import ExtTestCase, requires_torch
 from onnx_diagnostic.torch_export_patches.patches.patch_torch import patched_vmap
+from onnx_diagnostic.torch_export_patches.patches.patch_transformers import (
+    patched__vmap_for_bhqkv as _vmap_for_bhqkv2,
+)
 
 
 class TestPatchPatchTorch(ExtTestCase):
@@ -62,7 +65,7 @@ def test_vmap_tuple(self):
         got = patched_vmap(torch.dot, in_dims=(0, None))(x, y)
         self.assertEqualArray(expected, got)
 
-    def test_vmap_transformers_scenario(self):
+    def test_vmap_transformers_scenario_vmap(self):
         def padding_mask_function(padding_mask: torch.Tensor) -> Callable:
             def inner_mask(batch_idx, head_idx, q_idx, kv_idx):
                 return padding_mask[batch_idx, kv_idx]
@@ -127,12 +130,77 @@ def forward(self, batch_arange, head_arange, cache_position, kv_arange):
         got = Model()(*inputs)
         self.assertEqualArray(causal_mask, got)
 
+        if not requires_torch("4.10"):
+            DYN = torch.export.Dim.DYNAMIC
+            ds1 = {0: DYN}
+            ds2 = {0: DYN, 1: DYN}
+            ds = (ds2, ds1, ds1, ds1)
+            ep = torch.export.export(Model(), inputs, dynamic_shapes=ds)
+            self.assertEqualArray(causal_mask, ep.moule(*inputs))
+
+    def test_vmap_transformers_scenario_novmap(self):
+        def padding_mask_function(padding_mask: torch.Tensor) -> Callable:
+            def inner_mask(batch_idx, head_idx, q_idx, kv_idx):
+                return padding_mask[batch_idx, kv_idx]
+
+            return inner_mask
+
+        def and_masks(*mask_functions: list[Callable]) -> Callable:
+            def and_mask(batch_idx, head_idx, q_idx, kv_idx):
+                result = q_idx.new_ones((), dtype=torch.bool)
+                for mask in mask_functions:
+                    result = result & mask(batch_idx, head_idx, q_idx, kv_idx)
+                return result
+
+            return and_mask
+
+        def causal_mask_function(
+            batch_idx: int, head_idx: int, q_idx: int, kv_idx: int
+        ) -> bool:
+            return kv_idx <= q_idx
+
+        def _vmap_for_bhqkv(mask_function: Callable, bh_indices: bool = True) -> Callable:
+            dimensions = [(None, None, None, 0), (None, None, 0, None)]
+            if bh_indices:
+                dimensions.extend([(None, 0, None, None), (0, None, None, None)])
+            for dims in dimensions:
+                mask_function = torch.vmap(mask_function, in_dims=dims, out_dims=0)
+            return mask_function
+
+        padding_mask = torch.ones((2, 33)).to(torch.bool)
+        batch_arange = torch.tensor([0, 1], dtype=torch.int64)
+        head_arange = torch.tensor([0, 1], dtype=torch.int64)
+        cache_position = torch.tensor([30, 31, 32], dtype=torch.int64)
+        kv_arange = torch.arange(33, dtype=torch.int64)
+        mask_function = and_masks(causal_mask_function, padding_mask_function(padding_mask))
+        with TransformGetItemToIndex():
+            causal_mask = _vmap_for_bhqkv(mask_function)(
+                batch_arange, head_arange, cache_position, kv_arange
+            )
+        with TransformGetItemToIndex():
+            causal_mask2 = _vmap_for_bhqkv2(mask_function)(
+                batch_arange, head_arange, cache_position, kv_arange
+            )
+        self.assertEqualArray(causal_mask, causal_mask2)
+
+        class Model(torch.nn.Module):
+            def forward(self, batch_arange, head_arange, cache_position, kv_arange):
+                with TransformGetItemToIndex():
+                    causal_mask2 = _vmap_for_bhqkv2(mask_function)(
+                        batch_arange, head_arange, cache_position, kv_arange
+                    )
+                return causal_mask2
+
+        inputs = batch_arange, head_arange, cache_position, kv_arange
+        got = Model()(*inputs)
+        self.assertEqualArray(causal_mask, got)
+
         DYN = torch.export.Dim.DYNAMIC
         ds1 = {0: DYN}
         ds2 = {0: DYN, 1: DYN}
         ds = (ds2, ds1, ds1, ds1)
         ep = torch.export.export(Model(), inputs, dynamic_shapes=ds)
-        self.assertEqualArray(causal_mask, ep.moule(*inputs))
+        self.assertEqualArray(causal_mask, ep.module()(*inputs))
 
 
 if __name__ == "__main__":

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -334,6 +334,8 @@ def torch_export_patches(
         ####################
 
         if patch_transformers:
+            import transformers.masking_utils
+
             if verbose:
                 import transformers
 
@@ -345,6 +347,16 @@ def torch_export_patches(
                 patch_transformers_list, verbose=verbose
             )
 
+        if hasattr(transformers.masking_utils, "_vmap_for_bhqkv"):
+            if verbose:
+                print(
+                    "[torch_export_patches] patches transformers.masking_utils._vmap_for_bhqkv"
+                )
+            f_transformers__vmap_for_bhqkv = transformers.masking_utils._vmap_for_bhqkv
+            transformers.masking_utils._vmap_for_bhqkv = (
+                patch_transformers_list.patched__vmap_for_bhqkv
+            )
+
         if custom_patches:
             if verbose:
                 print("[torch_export_patches] applies custom patches")
@@ -443,6 +455,14 @@ def torch_export_patches(
                     patch_transformers_list, revert_patches_info, verbose=verbose
                 )
 
+                if hasattr(transformers.masking_utils, "_vmap_for_bhqkv"):
+                    if verbose:
+                        print(
+                            "[torch_export_patches] unpatch "
+                            "transformers.masking_utils._vmap_for_bhqkv"
+                        )
+                    transformers.masking_utils._vmap_for_bhqkv = f_transformers__vmap_for_bhqkv
+
             ########
             # caches
             ########

onnx_diagnostic/torch_export_patches/patches/patch_transformers.py

@@ -1,6 +1,6 @@
 import inspect
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple
 import torch
 import transformers
 from transformers.modeling_attn_mask_utils import AttentionMaskConverter
@@ -9,6 +9,31 @@
 from ...helpers.torch_helper import is_torchdynamo_exporting
 
 
+def patched__vmap_for_bhqkv(mask_function: Callable, bh_indices: bool = True) -> Callable:
+    """Patch for function ``transformers.masking_utils._vmap_for_bhqkv``."""
+    from ...helpers import string_type
+
+    dimensions = [(None, None, None, 0), (None, None, 0, None)]
+    if bh_indices:
+        dimensions.extend([(None, 0, None, None), (0, None, None, None)])
+    dimensions = [tuple(1 if d is None else -1 for d in shape) for shape in dimensions]
+    dimensions = tuple(reversed(dimensions))
+    indices = tuple(shape.index(-1) for shape in dimensions)
+
+    def vector_mask_function(
+        *args, mask_function=mask_function, dimensions=dimensions, indices=indices
+    ):
+        assert len(args) == len(
+            dimensions
+        ), f"Mismatch between args={string_type(args)} and dimensions={dimensions}"
+        new_args = [a.reshape(shape) for a, shape in zip(args, dimensions)]
+        max_shape = tuple(args[i].shape[0] for i in indices)
+        expanded_args = [a.expand(max_shape) for a in new_args]
+        return mask_function(*expanded_args)
+
+    return vector_mask_function
+
+
 def _patch_make_causal_mask(
     input_ids_shape: torch.Size,
     dtype: torch.dtype,