diff --git a/setup.py b/setup.py index ff56f27da6c..aa0d28d6fc8 100644 --- a/setup.py +++ b/setup.py @@ -56,6 +56,7 @@ "GPUtil>=1.4.0", "protobuf>=3.12.2,<=3.20.3", "click>=7.1.2,!=8.0.0", # latest version < 8.0 + blocked version with reported bug + "clearml==1.14.4", ] _nm_deps = [f"{'sparsezoo' if is_release else 'sparsezoo-nightly'}~={version_nm_deps}"] _deepsparse_deps = [ diff --git a/tests/sparseml/test_clear_ml.py b/tests/sparseml/test_clear_ml.py new file mode 100644 index 00000000000..78c70be2ca0 --- /dev/null +++ b/tests/sparseml/test_clear_ml.py @@ -0,0 +1,50 @@ +# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path + +import torch + +from clearml import Task +from sparseml.transformers import apply + + +def test_oneshot_and_finetune(tmp_path: Path): + recipe_str = "tests/sparseml/transformers/finetune/test_alternate_recipe.yaml" + model = "Xenova/llama2.c-stories15M" + device = "cuda:0" + if not torch.cuda.is_available(): + device = "cpu" + dataset = "wikitext" + dataset_config_name = "wikitext-2-raw-v1" + concatenate_data = True + run_stages = True + output_dir = tmp_path + max_steps = 50 + splits = {"train": "train[:50%]", "calibration": "train[50%:60%]"} + + Task.init(project_name="test", task_name="test_oneshot_and_finetune") + + apply( + model=model, + dataset=dataset, + dataset_config_name=dataset_config_name, + run_stages=run_stages, + output_dir=output_dir, + recipe=recipe_str, + max_steps=max_steps, + concatenate_data=concatenate_data, + splits=splits, + oneshot_device=device, + )