From 0942a2d2a5523ddb0b63f9eaac415cab242580dc Mon Sep 17 00:00:00 2001
From: shaahji <96227573+shaahji@users.noreply.github.com>
Date: Thu, 30 Jan 2025 18:02:29 -0800
Subject: [PATCH] Fix issues with CI builds (#1589)

## Fix issues with CI builds

* Relax version of bitsandbytes
* Add triton to requirements
* Few fixes for using newer version of torch

## Checklist before requesting a review
- [ ] Add unit tests for this change.
- [ ] Make sure all tests can pass.
- [ ] Update documents if necessary.
- [x] 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.
- [ ] Is this PR including examples changes? If yes, please remember to
update [example
documentation](https://github.com/microsoft/Olive/blob/main/docs/source/examples.md)
in a follow-up PR.

## (Optional) Issue link
---
 olive/model/handler/pytorch.py             | 2 +-
 olive/olive_config.json                    | 4 ++--
 test/requirements-test-gpu.txt             | 4 ++--
 test/unit_test/model/test_pytorch_model.py | 2 +-
 4 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/olive/model/handler/pytorch.py b/olive/model/handler/pytorch.py
index e87c9d39b..cb4b53c4b 100644
--- a/olive/model/handler/pytorch.py
+++ b/olive/model/handler/pytorch.py
@@ -151,7 +151,7 @@ def load_model(self, rank: int = None, cache_model: bool = True) -> "torch.nn.Mo
             elif self.model_file_format == ModelFileFormat.PYTORCH_TORCH_SCRIPT:
                 model = torch.jit.load(self.model_path)
             elif self.model_file_format == ModelFileFormat.PYTORCH_ENTIRE_MODEL:
-                model = torch.load(self.model_path)
+                model = torch.load(self.model_path, weights_only=False)
             elif self.model_file_format == ModelFileFormat.PYTORCH_SLICE_GPT_MODEL:
                 model = self._load_slicegpt_model()
             elif self.model_file_format == ModelFileFormat.PYTORCH_STATE_DICT:
diff --git a/olive/olive_config.json b/olive/olive_config.json
index 0027ef697..1bee9abfe 100644
--- a/olive/olive_config.json
+++ b/olive/olive_config.json
@@ -363,13 +363,13 @@
     "extra_dependencies": {
         "auto-opt": [ "optimum" ],
         "azureml": [ "azure-ai-ml>=1.11.1", "azure-keyvault-secrets", "azure-identity", "azureml-fsspec" ],
-        "bnb": [ "bitsandbytes" ],
+        "bnb": [ "bitsandbytes", "triton" ],
         "capture-onnx-graph": [ "onnxruntime-genai", "optimum" ],
         "cpu": [ "onnxruntime" ],
         "directml": [ "onnxruntime-directml" ],
         "docker": [ "docker" ],
         "shared-cache": [ "azure-identity", "azure-storage-blob" ],
-        "finetune": [ "onnxruntime-genai", "optimum", "accelerate>=0.30.0", "peft", "scipy", "bitsandbytes" ],
+        "finetune": [ "onnxruntime-genai", "optimum", "accelerate>=0.30.0", "peft", "scipy", "bitsandbytes", "triton" ],
         "flash-attn": [ "flash_attn" ],
         "gpu": [ "onnxruntime-gpu" ],
         "inc": [ "neural-compressor" ],
diff --git a/test/requirements-test-gpu.txt b/test/requirements-test-gpu.txt
index e8b9592f6..65674e827 100644
--- a/test/requirements-test-gpu.txt
+++ b/test/requirements-test-gpu.txt
@@ -1,5 +1,5 @@
 -r requirements-test.txt
 auto-gptq
 autoawq
-# only available on Linux currently
-bitsandbytes==0.43.3
+bitsandbytes
+triton
diff --git a/test/unit_test/model/test_pytorch_model.py b/test/unit_test/model/test_pytorch_model.py
index ac26454e1..bee910886 100644
--- a/test/unit_test/model/test_pytorch_model.py
+++ b/test/unit_test/model/test_pytorch_model.py
@@ -58,7 +58,7 @@ def test_load_from_path(torch_load):
     model = PyTorchModelHandler(model_path="test_path")
 
     assert model.load_model() == "dummy_pytorch_model"
-    torch_load.assert_called_once_with("test_path")
+    torch_load.assert_called_once_with("test_path", weights_only=False)
 
 
 @patch("olive.model.handler.pytorch.UserModuleLoader")