apply review comments

Author: eaidova

optimum/exporters/openvino/model_configs.py

@@ -764,43 +764,6 @@ def patch_model_for_export(
         return CodeGenModelPatcher(self, model, model_kwargs=model_kwargs)
 
 
-class DBRXDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator):
-    def __init__(
-        self,
-        task: str,
-        normalized_config: NormalizedTextConfig,
-        batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"],
-        sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"],
-        random_batch_size_range: Optional[Tuple[int, int]] = None,
-        random_sequence_length_range: Optional[Tuple[int, int]] = None,
-        **kwargs,
-    ):
-        super().__init__(
-            task=task,
-            normalized_config=normalized_config,
-            batch_size=batch_size,
-            sequence_length=sequence_length,
-            random_batch_size_range=random_batch_size_range,
-            random_sequence_length_range=random_sequence_length_range,
-        )
-        self.num_key_value_heads = normalized_config.num_key_value_heads
-
-    def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
-        shape = (
-            self.batch_size,
-            self.num_key_value_heads,
-            self.sequence_length,
-            self.hidden_size // self.num_attention_heads,
-        )
-        return [
-            (
-                self.random_float_tensor(shape, framework=framework, dtype=float_dtype),
-                self.random_float_tensor(shape, framework=framework, dtype=float_dtype),
-            )
-            for _ in range(self.num_layers)
-        ]
-
-
 @register_in_tasks_manager(
     "dbrx",
     *["text-generation", "text-generation-with-past"],
@@ -815,8 +778,8 @@ class DBRXOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig):
         num_key_value_heads="attn_config.kv_n_heads",
         allow_new=True,
     )
-    DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, DBRXDummyPastKeyValuesGenerator)
-    DUMMY_PKV_GENERATOR_CLASS = DBRXDummyPastKeyValuesGenerator
+    DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator)
+    DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator
 
     def patch_model_for_export(
         self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None

optimum/exporters/openvino/model_patcher.py

@@ -1468,6 +1468,7 @@ def _dbrx_experts_forward(
     return out
 
 
+# adapted from https://github.com/huggingface/transformers/blob/v4.40.2/src/transformers/models/dbrx/modeling_dbrx.py#L1228
 def _dbrx_update_causal_mask_legacy(
     self, attention_mask: Optional[torch.Tensor], input_tensor: torch.Tensor, cache_position: torch.Tensor
 ) -> Optional[torch.Tensor]:
@@ -1479,6 +1480,9 @@ def _dbrx_update_causal_mask_legacy(
         return None
 
     dtype, device = input_tensor.dtype, input_tensor.device
+    # difference with original modeling
+    # using minimum from dtype with larger bandwith (floa32) may lead to overflow
+    # during execution on platforms with default lower precision (bfloat16, float16)
     min_dtype = torch.finfo(torch.float16).min
     sequence_length = input_tensor.shape[1]
     if hasattr(self.blocks[0].norm_attn_norm.attn, "past_key_value"):  # static cache
@@ -1487,7 +1491,9 @@ def _dbrx_update_causal_mask_legacy(
         target_length = (
             attention_mask.shape[-1] if isinstance(attention_mask, torch.Tensor) else cache_position[-1] + 1
         )
-
+    # difference with original modeling
+    # removed target_length = int(target_length).
+    # Casting to int leads to constant folding during tracing that makes impossible to use model for sequence of different length
     causal_mask = torch.full((sequence_length, target_length), fill_value=1, dtype=dtype, device=device) * min_dtype
     if sequence_length != 1:
         causal_mask = torch.triu(causal_mask, diagonal=1)
@@ -1535,6 +1541,7 @@ def _dbrx_update_causal_mask_legacy(
     return causal_mask
 
 
+# adopted from https://github.com/huggingface/transformers/blob/1b3dba9417eebe16b7c206d1dfca6a4c7f11dbec/src/transformers/models/dbrx/modeling_dbrx.py#L1204
 def _dbrx_update_causal_mask_latest(
     self,
     attention_mask: torch.Tensor,
@@ -1631,18 +1638,22 @@ def _dbrx_update_causal_mask_latest(
 class DBRXModelPatcher(DecoderModelPatcher):
     def __enter__(self):
         super().__enter__()
+        # dbrx has some accuracy issues with bf16 with transformers >= 4.40
+        # fill causal mask in slightly different way for avoid overflow on some platforms
         self._model.transformer._orig_update_causal_mask = self._model.transformer._update_causal_mask
         self._model.transformer._update_causal_mask = types.MethodType(
             _dbrx_update_causal_mask, self._model.transformer
         )
 
         for block in self._model.transformer.blocks:
             rotary_emb = block.norm_attn_norm.attn.rotary_emb
+            # initialize inv_freq for torchscript tracing
             if rotary_emb.inv_freq is None:
                 inv_freq = 1.0 / (
                     rotary_emb.base ** (torch.arange(0, rotary_emb.dim, 2, dtype=torch.int64).float() / rotary_emb.dim)
                 )
                 rotary_emb.inv_freq = inv_freq
+            # remove continue-operator from iteration loop over experts
             block.ffn.experts._orig_forward = block.ffn.experts.forward
             block.ffn.experts.forward = types.MethodType(_dbrx_experts_forward, block.ffn.experts)