Move WORKDIR to storage disk in GPU CI & Use dynamo export for tests (#2313)

## Describe your changes

1. Remove optimum-intel version pin from tests, so we can test the
latest transformers.
2. Use dynamo export by default in our tests, since torchscript is no
longer been maintained.
3. Fix a couple of tests for latest transformers and dynamo export
change. Skip a couple of tests that need further investigations.
4. Move GPU agent work dir and docker work dir to storage disk,
otherwise the main disk will be out of space.
5. Add `register_dynamic_cache_export_support` to fix dynamic cache
registration issue.
6. Remove default 13 target_opset in tests.

## Checklist before requesting a review
- [ ] Add unit tests for this change.
- [ ] Make sure all tests can pass.
- [ ] Update documents if necessary.
- [ ] Lint and apply fixes to your code by running `lintrunner -a`
- [ ] Is this a user-facing change? If yes, give a description of this
change to be included in the release notes.

## (Optional) Issue link

Author: xiaoyu-work

.azure_pipelines/job_templates/olive-test-cpu-template.yaml

@@ -38,7 +38,7 @@ jobs:
   - script: |
         python -m pip install pytest
         python -m pip install -r $(Build.SourcesDirectory)/test/$(requirements_file)
-
+        python -m pip list
         coverage run --source=$(Build.SourcesDirectory)/olive -m pytest -v -s -p no:warnings --disable-warnings --log-cli-level=WARNING --junitxml=$(Build.SourcesDirectory)/logs/test-TestOlive.xml $(Build.SourcesDirectory)/test --basetemp $(PYTEST_BASETEMP)
         coverage xml
     displayName: Test Olive

.azure_pipelines/job_templates/olive-test-linux-gpu-template.yaml

@@ -20,10 +20,34 @@ jobs:
   pool:
     name: ${{ parameters.pool}}
   variables:
-    PIP_CACHE_DIR: $(Pipeline.Workspace)/.cache/pip
-    HF_HOME: $(Pipeline.Workspace)/.cache/huggingface
+    PIP_CACHE_DIR: /mnt/storage/.cache/pip
+    HF_HOME: /mnt/storage/.cache/huggingface
 
   steps:
+  - script: |
+      set -euxo pipefail
+
+      # Move agent work directory to /mnt/storage via symlink
+      AGENT_ROOT=$(dirname "$(Agent.BuildDirectory)")
+      sudo mkdir -p /mnt/storage/vss_work
+      sudo chown -R $USER:$USER /mnt/storage/vss_work
+      sudo cp -a "$AGENT_ROOT"/* /mnt/storage/vss_work/ 2>/dev/null || true
+      sudo rm -rf "$AGENT_ROOT"
+      sudo ln -sf /mnt/storage/vss_work "$AGENT_ROOT"
+
+      # Move Docker and containerd to /mnt/storage
+      sudo systemctl stop docker containerd
+      sudo mkdir -p /mnt/storage/docker /mnt/storage/containerd /etc/containerd
+      echo '{"data-root": "/mnt/storage/docker"}' | sudo tee /etc/docker/daemon.json
+      containerd config default | sed 's|/var/lib/containerd|/mnt/storage/containerd|g' | sudo tee /etc/containerd/config.toml > /dev/null
+      sudo systemctl start containerd docker
+
+      # Move /tmp to /mnt/storage
+      sudo mkdir -p /mnt/storage/tmp
+      sudo chmod 1777 /mnt/storage/tmp
+      sudo mount --bind /mnt/storage/tmp /tmp
+    displayName: Move pipeline to /mnt/storage
+
   - template: build-docker-image-template.yaml
     parameters:
       python_version: ${{ parameters.python_version }}

.azure_pipelines/scripts/run_test.sh

@@ -33,6 +33,8 @@ pip install -r "$4"
 pip install huggingface-hub
 hf auth login --token "$7"
 
+pip list
+
 # Step 4: Run tests with or without coverage tracking
 XML_PATH="/logs/TestOlive.xml"
 if [ "$6" = "true" ]; then

olive/cli/optimize.py

@@ -797,7 +797,7 @@ def _get_onnx_io_datatype_converter_pass_config(self) -> dict[str, Any]:
     {
         "name": "wikitext2_train",
         "type": "HuggingfaceContainer",
-        "load_dataset_config": {"data_name": "wikitext", "subset": "wikitext-2-raw-v1", "split": "train"},
+        "load_dataset_config": {"data_name": "Salesforce/wikitext", "subset": "wikitext-2-raw-v1", "split": "train"},
         "pre_process_data_config": {
             "strategy": "line-by-line",
             "add_special_tokens": False,

olive/common/quant/hf_utils.py

@@ -146,6 +146,7 @@ class OliveHfQuantizer(HfQuantizer):
 
     # only support load and inference, no on-the-fly quantization
     requires_calibration = True
+    modules_to_not_convert: list[str] | None = None
 
     def _process_model_before_weight_loading(
         self, model: PreTrainedModel, keep_in_fp32_modules: list[str] | None = None, **kwargs

olive/data/component/sd_lora/dataloader.py

@@ -35,7 +35,7 @@ def __init__(
             seed: Random seed for reproducibility.
 
         """
-        super().__init__(dataset)
+        super().__init__()
         self.dataset = dataset
         self.batch_size = batch_size
         self.drop_last = drop_last

olive/data/template.py

@@ -39,14 +39,14 @@ def huggingface_data_config_template(model_name, task, **kwargs) -> DataConfig:
     **kwargs: dict
         Additional arguments:
         - olive.data.component.load_dataset_config.huggingface_dataset
-            - `data_name`: str, data name in huggingface dataset, e.g.: "glue", "squad"
+            - `data_name`: str, data name in huggingface dataset, e.g.: "nyu-mll/glue", "squad"
             - `subset`: str, subset of data, e.g.: "train", "validation", "test"
             - `split`: str, split of data, e.g.: "train", "validation", "test"
             - `data_files`: str | list | dict, path to source data file(s).
             e.g.
                 load_dataset_config={
                     "params": {
-                        "data_name": "glue",
+                        "data_name": "nyu-mll/glue",
                         "subset": "train",
                         "split": "train",
                         "data_files": "whatever.pt"

olive/passes/onnx/conversion.py

@@ -212,6 +212,11 @@ def _export_pytorch_model(
                     "Please upgrade PyTorch to 2.6.0 or above."
                 )
 
+            # Register DynamicCache export support
+            from transformers.integrations.executorch import register_dynamic_cache_export_support
+
+            register_dynamic_cache_export_support()
+
             if isinstance(dummy_inputs, dict):
                 dummy_kwargs = dummy_inputs
                 dummy_inputs = ()
@@ -236,7 +241,7 @@ def _export_pytorch_model(
                 dynamic_axes=io_config.dynamic_axes,
                 dynamic_shapes=io_config.dynamic_shapes,
                 dynamo=True,
-                fallback=True,
+                fallback=False,
                 optimize=config.optimize,
                 report=logger.isEnabledFor(logging.DEBUG),
             )

olive/passes/pytorch/sparsegpt_utils.py

@@ -90,12 +90,14 @@ def __init__(self, module):
             super().__init__()
             self.module = module
 
-        def forward(self, inputs, **kwargs):
+        def forward(self, *args, **kwargs):
+            # First positional argument is the hidden states (inputs)
+            layer_inputs = args[0] if args else kwargs.get("hidden_states")
             # handle batch dimension
-            for batch in range(inputs.shape[0]):
+            for batch in range(layer_inputs.shape[0]):
                 if cache["i"] >= num_samples:
                     break
-                inputs[cache["i"]] = inputs[batch]
+                inputs[cache["i"]] = layer_inputs[batch]
                 cache["i"] += 1
             cache["attention_mask"] = kwargs.get("attention_mask")
             for input_name in additional_input:

olive/passes/pytorch/train_utils.py

@@ -324,7 +324,7 @@ def get_calibration_data_config(
         model_name=model_name_or_path,
         task="text-generation",
         load_dataset_config={
-            "data_name": "wikitext",
+            "data_name": "Salesforce/wikitext",
             "subset": "wikitext-2-raw-v1",
             "split": split,
             "trust_remote_code": trust_remote_code,