Merge remote-tracking branch 'origin/transformers_future' into upstream-accelerate

IlyasMoutawwakil · IlyasMoutawwakil · commit 1dd9e469f77e · 2025-03-16T16:39:09.000Z
diff --git a/.github/workflows/fast_tests.yml b/.github/workflows/fast_tests.yml
@@ -1,12 +1,11 @@
 name: Unit and integration tests
 
-
 on:
   workflow_dispatch:
   pull_request:
-    branches: [ main ]
+    branches: [main]
   push:
-    branches: [ main ]
+    branches: [main]
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
@@ -16,50 +15,69 @@ jobs:
   transformers:
     name: Run tests for optimum.habana.transformers
     runs-on: [self-hosted, linux, x64, gaudi2, fast]
+
+    container:
+      image: docker://vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest
+      options: --runtime=habana --shm-size=64G --env HABANA_VISIBLE_DEVICES
+      env:
+        OMPI_MCA_btl_vader_single_copy_mechanism: none
+
     steps:
-      - name: Checkout
-        uses: actions/checkout@v2
-      - name: Pull image
+      - name: HL-SMI (1)
+        run: |
+          hl-smi
+          echo "HABANA_VISIBLE_DEVICES=${HABANA_VISIBLE_DEVICES}"
+          echo "HABANA_VISIBLE_MODULES=${HABANA_VISIBLE_MODULES}"
+
+      - name: Extract HPU visible modules
         run: |
-            docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest
+          export HABANA_VISIBLE_MODULES=$(hl-smi -Q module_id -f csv,noheader | tr '\n' ',' | sed 's/,$//')
+          echo "HABANA_VISIBLE_MODULES=${HABANA_VISIBLE_MODULES}" >> $GITHUB_ENV
+
+      - name: HL-SMI (2)
+        run: |
+          hl-smi
+          echo "HABANA_VISIBLE_DEVICES=${HABANA_VISIBLE_DEVICES}"
+          echo "HABANA_VISIBLE_MODULES=${HABANA_VISIBLE_MODULES}"
+
+      - name: Checkout
+        uses: actions/checkout@v4
+
       - name: Run tests
         run: |
-            docker run \
-            --rm \
-            -v $PWD:/root/workspace \
-            -v /scratch-1:/data \
-            --workdir=/root/workspace \
-            --runtime=habana \
-            -e HABANA_VISIBLE_DEVICES=$DOCKER_HABANA_VISIBLE_DEVICES \
-            -e OMPI_MCA_btl_vader_single_copy_mechanism=none \
-            -e HF_HOME=/data \
-            --cap-add=sys_nice \
-            --net=host \
-            --ipc=host \
-            vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest \
-            /bin/bash tests/ci/fast_tests.sh
+          /bin/bash tests/ci/fast_tests.sh
+
   diffusers:
     name: Run tests for optimum.habana.diffusers
     runs-on: [self-hosted, linux, x64, gaudi2, fast]
+ 
+    container:
+      image: docker://vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest
+      options: --runtime=habana --shm-size=64G --env HABANA_VISIBLE_DEVICES
+      env:
+        OMPI_MCA_btl_vader_single_copy_mechanism: none
+
     steps:
-      - name: Checkout
-        uses: actions/checkout@v2
-      - name: Pull image
+      - name: HL-SMI (1)
+        run: |
+          hl-smi
+          echo "HABANA_VISIBLE_DEVICES=${HABANA_VISIBLE_DEVICES}"
+          echo "HABANA_VISIBLE_MODULES=${HABANA_VISIBLE_MODULES}"
+
+      - name: Extract HPU visible modules
         run: |
-            docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest
+          export HABANA_VISIBLE_MODULES=$(hl-smi -Q module_id -f csv,noheader | tr '\n' ',' | sed 's/,$//')
+          echo "HABANA_VISIBLE_MODULES=${HABANA_VISIBLE_MODULES}" >> $GITHUB_ENV
+
+      - name: HL-SMI (2)
+        run: |
+          hl-smi
+          echo "HABANA_VISIBLE_DEVICES=${HABANA_VISIBLE_DEVICES}"
+          echo "HABANA_VISIBLE_MODULES=${HABANA_VISIBLE_MODULES}"
+
+      - name: Checkout
+        uses: actions/checkout@v4
+
       - name: Run tests
         run: |
-            docker run \
-            --rm \
-            -v $PWD:/root/workspace \
-            -v /scratch-1:/data \
-            --workdir=/root/workspace \
-            --runtime=habana \
-            -e HABANA_VISIBLE_DEVICES=$DOCKER_HABANA_VISIBLE_DEVICES \
-            -e OMPI_MCA_btl_vader_single_copy_mechanism=none \
-            -e HF_HOME=/data \
-            --cap-add=sys_nice \
-            --net=host \
-            --ipc=host \
-            vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest \
-            /bin/bash tests/ci/fast_tests_diffusers.sh
+          /bin/bash tests/ci/fast_tests_diffusers.sh
diff --git a/.github/workflows/upstream.yml b/.github/workflows/upstream.yml
@@ -0,0 +1,146 @@
+name: Upstream Integrations
+
+on:
+  workflow_dispatch:
+  schedule:
+    # every monday at 00:00 UTC
+    - cron: "0 0 * * 1"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+jobs:
+  transformers:
+    name: Upstream Transformers
+    runs-on: [self-hosted, linux, x64, gaudi2, fast]
+
+    container:
+      image: docker://vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest
+      options: --runtime=habana --shm-size=64G --env HABANA_VISIBLE_DEVICES
+      env:
+        OMPI_MCA_btl_vader_single_copy_mechanism: none
+        RUN_THIRD_PARTY_DEVICE_TESTS: 1
+        TRANSFORMERS_TEST_DEVICE: hpu
+        PT_ENABLE_INT64_SUPPORT: 1
+        PT_HPU_LAZY_MODE: 0
+        RUN_SLOW: 1
+
+    steps:
+      - name: HL-SMI (1)
+        run: |
+          hl-smi
+          echo "HABANA_VISIBLE_DEVICES=${HABANA_VISIBLE_DEVICES}"
+          echo "HABANA_VISIBLE_MODULES=${HABANA_VISIBLE_MODULES}"
+
+      - name: Extract HPU visible modules
+        run: |
+          export HABANA_VISIBLE_MODULES=$(hl-smi -Q module_id -f csv,noheader | tr '\n' ',' | sed 's/,$//')
+          echo "HABANA_VISIBLE_MODULES=${HABANA_VISIBLE_MODULES}" >> $GITHUB_ENV
+
+      - name: HL-SMI (2)
+        run: |
+          hl-smi
+          echo "HABANA_VISIBLE_DEVICES=${HABANA_VISIBLE_DEVICES}"
+          echo "HABANA_VISIBLE_MODULES=${HABANA_VISIBLE_MODULES}"
+
+      - name: Checkout to Transformers
+        uses: actions/checkout@v4
+        with:
+          repository: huggingface/transformers
+
+      - name: Install Transformers with Accelerate & DeepSpeed
+        run: |
+          pip install -e .[testing] "numpy<2.0.0" scipy scikit-learn \
+            git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0 \
+            git+https://github.com/huggingface/accelerate.git
+
+      - name: Run Trainer tests
+        run: |
+          pytest tests/trainer/test_trainer.py -s -vvvv
+
+      - name: Run Trainer Utils tests
+        run: |
+          pytest tests/trainer/test_trainer_utils.py -s -vvvv
+
+      - name: Run Trainer Seq2Seq tests
+        run: |
+          pytest tests/trainer/test_trainer_seq2seq.py -s -vvvv
+
+      - name: Run Trainer Distributed tests
+        run: |
+          pytest tests/trainer/test_trainer_distributed.py -s -vvvv
+
+      - name: Run FSDP Integration tests
+        run: |
+          pytest tests/fsdp/test_fsdp.py tests/trainer/test_trainer_fsdp.py -s -vvvv
+
+      - name: Run DeepSpeed Integration tests
+        run: |
+          pytest tests/deepspeed/test_deepspeed.py -s -vvvv
+
+  accelerate:
+    name: Upstream Accelerate
+    runs-on: [self-hosted, linux, x64, gaudi2, fast]
+
+    container:
+      image: docker://vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest
+      options: --runtime=habana --shm-size=64G --cap-add=sys_nice --env HABANA_VISIBLE_DEVICES
+      env:
+        OMPI_MCA_btl_vader_single_copy_mechanism: none
+        PT_ENABLE_INT64_SUPPORT: 1
+        PT_HPU_LAZY_MODE: 0
+        RUN_SLOW: 1
+
+    steps:
+      - name: HL-SMI (1)
+        run: |
+          hl-smi
+          echo "HABANA_VISIBLE_DEVICES=${HABANA_VISIBLE_DEVICES}"
+          echo "HABANA_VISIBLE_MODULES=${HABANA_VISIBLE_MODULES}"
+
+      - name: Extract HPU visible modules
+        run: |
+          export HABANA_VISIBLE_MODULES=$(hl-smi -Q module_id -f csv,noheader | tr '\n' ',' | sed 's/,$//')
+          echo "HABANA_VISIBLE_MODULES=${HABANA_VISIBLE_MODULES}" >> $GITHUB_ENV
+
+      - name: HL-SMI (2)
+        run: |
+          hl-smi
+          echo "HABANA_VISIBLE_DEVICES=${HABANA_VISIBLE_DEVICES}"
+          echo "HABANA_VISIBLE_MODULES=${HABANA_VISIBLE_MODULES}"
+
+      - name: Checkout to Accelerate
+        uses: actions/checkout@v4
+        with:
+          repository: huggingface/accelerate
+
+      - name: Install Accelerate with Transformers & DeepSpeed
+        run: |
+          pip install -e .[testing] \
+            git+https://github.com/huggingface/transformers.git \
+            git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0
+
+      - name: Run CLI tests
+        run: |
+          make test_cli
+
+      - name: Run Core tests
+        run: |
+          make test_core
+
+      - name: Run Big Modeling tests
+        run: |
+          make test_big_modeling
+
+      - name: Run FSDP integration tests
+        run: |
+          make test_fsdp
+
+      - name: Run DeepSpeed integration tests
+        run: |
+          make test_deepspeed
+
+      - name: Run Examples tests
+        run: |
+          make test_examples
diff --git a/examples/language-modeling/run_clm.py b/examples/language-modeling/run_clm.py
@@ -459,7 +459,7 @@ def main():
 
     # Note that chatglm2/3 has float16 dtype from config.json, and on Gaudi we need to use bfloat16.
     if config.model_type == "chatglm":
-        config.dtype = "torch.bfloat16"
+        config.torch_dtype = torch.bfloat16
 
     tokenizer_kwargs = {
         "cache_dir": model_args.cache_dir,
@@ -484,6 +484,11 @@ def main():
             if model_args.torch_dtype in ["auto", None]
             else getattr(torch, model_args.torch_dtype)
         )
+        # workaraund for https://github.com/huggingface/transformers/issues/36258
+        # TODO: remove after fix is avalible in a release version of `transformers``
+        if torch_dtype is None:
+            torch_dtype = getattr(config, "torch_dtype", None)
+
         model = AutoModelForCausalLM.from_pretrained(
             model_args.model_name_or_path,
             from_tf=bool(".ckpt" in model_args.model_name_or_path),
diff --git a/optimum/habana/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/optimum/habana/transformers/models/deepseek_v3/modeling_deepseek_v3.py
@@ -1692,7 +1692,6 @@ def forward(
 
         hidden_states = outputs[0]
         logits = self.lm_head(hidden_states)
-        logits = logits.float()
 
         loss = None
         if labels is not None:
diff --git a/optimum/habana/transformers/models/gpt2/modeling_gpt2.py b/optimum/habana/transformers/models/gpt2/modeling_gpt2.py
@@ -70,6 +70,7 @@ def _upcast_and_reordered_attn(self, query, key, value, attention_mask=None, hea
             attn_weights = attn_weights * head_mask
 
         attn_output = torch.matmul(attn_weights, value)
+        attn_output = attn_output.transpose(1, 2)
 
         return attn_output, attn_weights
 
diff --git a/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py b/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py
@@ -269,6 +269,7 @@ def gaudi_gpt_neox_model_forward(
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
     token_idx: Optional[torch.Tensor] = None,
+    **kwargs,
 ) -> Union[Tuple, BaseModelOutputWithPast]:
     """
     Copied from GPTNeoxModel.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neox/modeling_gpt_neox.py
diff --git a/optimum/habana/transformers/models/mllama/modeling_mllama.py b/optimum/habana/transformers/models/mllama/modeling_mllama.py
@@ -1152,8 +1152,11 @@ def _update_model_kwargs_for_generation(self, outputs, model_kwargs, is_encoder_
         # add cross-attn mask for new token
         if cross_attention_mask_prev is not None:
             token_idx = model_kwargs.get("token_idx", None)
+            token_idx_cpu = model_kwargs.get(
+                "token_idx_cpu", None
+            )  # returns an integer so following slicing ops happen using int instead of tensor
             if token_idx is not None:
-                mask = cross_attention_mask_prev[:, token_idx - 2 : token_idx - 1, ...]
+                mask = cross_attention_mask_prev[:, token_idx_cpu - 2 : token_idx_cpu - 1, ...]
                 cross_attention_mask_prev.index_copy_(1, token_idx - 1, mask)
                 model_kwargs["cross_attention_mask"] = cross_attention_mask_prev
             else:
diff --git a/optimum/habana/transformers/models/qwen2_vl/modeling_qwen2_vl.py b/optimum/habana/transformers/models/qwen2_vl/modeling_qwen2_vl.py
@@ -517,8 +517,9 @@ def forward(
 
 # from: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py#L1420
 class GaudiQwen2VLForConditionalGeneration(Qwen2VLForConditionalGeneration):
-    # todo: change when the following gets fixed https://github.com/huggingface/transformers/blame/66f29aaaf55c8fe0c3dbcd24beede2ca4effac56/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py#L390C5-L390C27 
+    # todo: change when the following gets fixed https://github.com/huggingface/transformers/blame/66f29aaaf55c8fe0c3dbcd24beede2ca4effac56/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py#L390C5-L390C27
     _supports_static_cache = True
+
     def forward(
         self,
         input_ids: torch.LongTensor = None,
diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py
@@ -935,6 +935,7 @@ def _gradient_checkpointing_wrap(func, *args, **kwargs):
             train_dataloader,
             len_dataloader,
             num_examples,
+            steps_trained_in_current_epoch,
         )
 
         hb_profiler = HabanaProfile(
@@ -1584,7 +1585,6 @@ def _prepare_input(self, data: Union[torch.Tensor, Any]) -> Union[torch.Tensor,
                 return data.to(**kwargs)
         return data
 
-
     # handled by accelerate now (in model preparation)
     # def autocast_smart_context_manager(self, cache_enabled: Optional[bool] = True):
     #     """
@@ -2643,7 +2643,14 @@ def _zero_model_grad(self, model):
                 model._zero_grad_kwargs = {}
 
     def get_num_items_in_batches(
-        self, args, epochs_trained, num_train_epochs, train_dataloader, len_dataloader, num_examples
+        self,
+        args,
+        epochs_trained,
+        num_train_epochs,
+        train_dataloader,
+        len_dataloader,
+        num_examples,
+        steps_trained_in_current_epoch,
     ):
         """
         Calculate the number of items in each batch for all epochs during training.
@@ -2659,10 +2666,15 @@ def get_num_items_in_batches(
         total_updates = steps_in_epoch // args.gradient_accumulation_steps + 1
         if args.gradient_accumulation_steps == 1:
             total_updates -= 1
+        global_step = 0
 
         num_items_in_batches = []
         for epoch in range(epochs_trained, num_train_epochs):
-            epoch_dataloader = train_dataloader
+            if epoch == epochs_trained and steps_trained_in_current_epoch > 0:
+                epoch_dataloader = skip_first_batches(train_dataloader, steps_trained_in_current_epoch)
+            else:
+                epoch_dataloader = train_dataloader
+
             if hasattr(epoch_dataloader, "set_epoch"):
                 epoch_dataloader.set_epoch(epoch)
 
@@ -2702,6 +2714,11 @@ def get_num_items_in_batches(
                     num_items_in_batch = None
 
                 num_items_in_batches[epoch].append(num_items_in_batch)
+                global_step += 1
+
+            # For iterable datasets, don't do more than max_steps steps
+            if len_dataloader is None and global_step >= args.max_steps:
+                break
 
         return num_items_in_batches
 
diff --git a/tests/baselines/fixture/tests/test_examples.json b/tests/baselines/fixture/tests/test_examples.json
