Synapse 1.15 addition

optimum/habana/transformers/models/wav2vec2/modeling_wav2vec2.py

@@ -190,6 +190,63 @@ def _gaudi_wav2vec2_sample_negative_indices(
     return sampled_negative_indices
 
 
+def gaudi_wav2vec2_forward(
+    self,
+    input_values: Optional[torch.Tensor],
+    attention_mask: Optional[torch.Tensor] = None,
+    mask_time_indices: Optional[torch.FloatTensor] = None,
+    output_attentions: Optional[bool] = None,
+    output_hidden_states: Optional[bool] = None,
+    return_dict: Optional[bool] = None,
+) -> Union[Tuple, Wav2Vec2BaseModelOutput]:
+    """
+    Copied from Transformers: https://github.com/huggingface/transformers/blob/bd469c40659ce76c81f69c7726759d249b4aef49/src/transformers/models/wav2vec2/modeling_wav2vec2.py#L1282
+    The only difference is that a clone of `hidden_states` is given to _mask_hidden_states to avoid an error.
+    """
+    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+    output_hidden_states = (
+        output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+    )
+    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+    extract_features = self.feature_extractor(input_values)
+    extract_features = extract_features.transpose(1, 2)
+
+    if attention_mask is not None:
+        # compute reduced attention_mask corresponding to feature vectors
+        attention_mask = self._get_feature_vector_attention_mask(
+            extract_features.shape[1], attention_mask, add_adapter=False
+        )
+
+    hidden_states, extract_features = self.feature_projection(extract_features)
+    hidden_states = self._mask_hidden_states(
+        hidden_states.clone(), mask_time_indices=mask_time_indices, attention_mask=attention_mask
+    )
+
+    encoder_outputs = self.encoder(
+        hidden_states,
+        attention_mask=attention_mask,
+        output_attentions=output_attentions,
+        output_hidden_states=output_hidden_states,
+        return_dict=return_dict,
+    )
+
+    hidden_states = encoder_outputs[0]
+
+    if self.adapter is not None:
+        hidden_states = self.adapter(hidden_states)
+
+    if not return_dict:
+        return (hidden_states, extract_features) + encoder_outputs[1:]
+
+    return Wav2Vec2BaseModelOutput(
+        last_hidden_state=hidden_states,
+        extract_features=extract_features,
+        hidden_states=encoder_outputs.hidden_states,
+        attentions=encoder_outputs.attentions,
+    )
+
+
 def _gaudi_wav2vec2_mask_hidden_states(
     self,
     hidden_states: torch.FloatTensor,
@@ -318,63 +375,6 @@ def gaudi_wav2vec2_encoder_forward(
     )
 
 
-def gaudi_wav2vec2_forward(
-    self,
-    input_values: Optional[torch.Tensor],
-    attention_mask: Optional[torch.Tensor] = None,
-    mask_time_indices: Optional[torch.FloatTensor] = None,
-    output_attentions: Optional[bool] = None,
-    output_hidden_states: Optional[bool] = None,
-    return_dict: Optional[bool] = None,
-) -> Union[Tuple, Wav2Vec2BaseModelOutput]:
-    """
-    Copied from Transformers: https://github.com/huggingface/transformers/blob/bd469c40659ce76c81f69c7726759d249b4aef49/src/transformers/models/wav2vec2/modeling_wav2vec2.py#L1282
-    The only difference is that a clone of `hidden_states` is given to _mask_hidden_states to avoid an error.
-    """
-    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-    output_hidden_states = (
-        output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-    )
-    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-    extract_features = self.feature_extractor(input_values)
-    extract_features = extract_features.transpose(1, 2)
-
-    if attention_mask is not None:
-        # compute reduced attention_mask corresponding to feature vectors
-        attention_mask = self._get_feature_vector_attention_mask(
-            extract_features.shape[1], attention_mask, add_adapter=False
-        )
-
-    hidden_states, extract_features = self.feature_projection(extract_features)
-    hidden_states = self._mask_hidden_states(
-        hidden_states.clone(), mask_time_indices=mask_time_indices, attention_mask=attention_mask
-    )
-
-    encoder_outputs = self.encoder(
-        hidden_states,
-        attention_mask=attention_mask,
-        output_attentions=output_attentions,
-        output_hidden_states=output_hidden_states,
-        return_dict=return_dict,
-    )
-
-    hidden_states = encoder_outputs[0]
-
-    if self.adapter is not None:
-        hidden_states = self.adapter(hidden_states)
-
-    if not return_dict:
-        return (hidden_states, extract_features) + encoder_outputs[1:]
-
-    return Wav2Vec2BaseModelOutput(
-        last_hidden_state=hidden_states,
-        extract_features=extract_features,
-        hidden_states=encoder_outputs.hidden_states,
-        attentions=encoder_outputs.attentions,
-    )
-
-
 def gaudi_wav2vec2_tdnnlayer_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
     """
     Copied from Transformers: https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/wav2vec2/modeling_wav2vec2.py#L2290

tests/baselines/llama_7b.json

@@ -67,9 +67,9 @@
                 "multi_card": {
                     "learning_rate": 3e-4,
                     "train_batch_size": 8,
-                    "perplexity": 2.4259,
-                    "train_runtime": 199.94,
-                    "train_samples_per_second": 88.664,
+                    "perplexity": 2.4502,
+                    "train_runtime": 210.305,
+                    "train_samples_per_second": 85.0801,
                     "extra_arguments": [
                         "--bf16 True",
                         "--gradient_accumulation_steps 2",
@@ -89,7 +89,7 @@
                         "--adam_epsilon 1e-08",
                         "--ddp_bucket_cap_mb 50",
                         "--validation_split_percentage 10",
-                        "--attn_softmax_bf16 True",
+                        "--attn_softmax_bf16",
                         "--pipelining_fwd_bwd False",
                         "--fsdp auto_wrap",
                         "--torch_compile_backend hpu_backend",
@@ -100,5 +100,4 @@
             }
         }
     }
-}
-
+}

tests/test_fsdp_examples.py

@@ -10,34 +10,37 @@
 from .test_examples import ACCURACY_PERF_FACTOR, TIME_PERF_FACTOR
 
 
-# Gaudi2 CI baselines
-# FSDP is not supported on Gaudi1
-MODELS_TO_TEST = {
-    "bf16": [
-        (
-            "bert-base-uncased",
-            "Habana/bert-base-uncased",
-            2807,
-            85.4688,
-            "question-answering",
-            24,
-            8,
-            "run_qa.py",
-            "full_shard",
-        ),
-        (
-            "meta-llama/Llama-2-7b-hf",
-            "",
-            54,
-            0.92,
-            "language-modeling",
-            8,
-            8,
-            "run_lora_clm.py",
-            "auto_wrap",
-        ),
-    ],
-}
+if os.environ.get("GAUDI2_CI", "0") == "1":
+    # Gaudi2 CI baselines
+    MODELS_TO_TEST = {
+        "bf16": [
+            (
+                "bert-base-uncased",
+                "Habana/bert-base-uncased",
+                2807,
+                85.4688,
+                "question-answering",
+                24,
+                8,
+                "run_qa.py",
+                "full_shard",
+            ),
+            (
+                "meta-llama/Llama-2-7b-hf",
+                "",
+                54,
+                0.92,
+                "language-modeling",
+                8,
+                8,
+                "run_lora_clm.py",
+                "auto_wrap",
+            ),
+        ],
+    }
+else:
+    # FSDP is not supported on Gaudi1
+    MODELS_TO_TEST = {"bf16": []}
 
 
 def _test_fsdp(

tests/test_text_generation_example.py

@@ -59,6 +59,7 @@
             ("mistralai/Mistral-7B-v0.1", 40.00435417311187),
             ("microsoft/phi-2", 90.10751623430603),
         ],
+        "fp8": [],
         "deepspeed": [
             ("bigscience/bloomz-7b1", 31.044523676681507),
         ],