Merge branch 'main' into update_oneshot_tests

integrations/huggingface-transformers/tutorials/sentiment-analysis/docs-sentiment-analysis-python-custom-teacher-rottentomatoes.ipynb

@@ -364,13 +364,11 @@
     "model_kwargs = {\"config\": model_config}\n",
     "model_kwargs[\"state_dict\"], s_delayed = SparseAutoModel._loadable_state_dict(model_path)\n",
     "model = AutoModelForSequenceClassification.from_pretrained(model_path,**model_kwargs,)\n",
-    "SparseAutoModel.log_model_load(model, model_path, \"student\", s_delayed) # prints metrics on sparsity profile\n",
     "\n",
     "# initialize teacher using familiar HF AutoModel\n",
     "teacher_kwargs = {\"config\": teacher_config}\n",
     "teacher_kwargs[\"state_dict\"], t_delayed = SparseAutoModel._loadable_state_dict(teacher_path)\n",
-    "teacher = AutoModelForSequenceClassification.from_pretrained(teacher_path,**teacher_kwargs,)\n",
-    "SparseAutoModel.log_model_load(teacher, teacher_path, \"teacher\", t_delayed)"
+    "teacher = AutoModelForSequenceClassification.from_pretrained(teacher_path,**teacher_kwargs,)"
    ]
   },
   {

integrations/huggingface-transformers/tutorials/sentiment-analysis/docs-sentiment-analysis-python-sst2.ipynb

@@ -440,11 +440,7 @@
     "\n",
     "teacher_kwargs = {'config':teacher_config}\n",
     "teacher_kwargs[\"state_dict\"], t_delayed = SparseAutoModel._loadable_state_dict(teacher_path)\n",
-    "teacher = AutoModelForSequenceClassification.from_pretrained(teacher_path, **teacher_kwargs,)\n",
-    "\n",
-    "# optional - prints metrics about sparsity profiles of the models\n",
-    "SparseAutoModel.log_model_load(model, model_path, \"student\", s_delayed)\n",
-    "SparseAutoModel.log_model_load(teacher, teacher_path, \"teacher\", t_delayed)"
+    "teacher = AutoModelForSequenceClassification.from_pretrained(teacher_path, **teacher_kwargs,)"
    ]
   },
   {

integrations/huggingface-transformers/tutorials/text-classification/docs-text-classification-python-custom-teacher-tweeteval.ipynb

@@ -515,10 +515,7 @@
     "# initialize model using familiar HF AutoModel\n",
     "model_kwargs = {\"config\": config}\n",
     "model_kwargs[\"state_dict\"], s_delayed = SparseAutoModel._loadable_state_dict(model_path)\n",
-    "model = AutoModelForSequenceClassification.from_pretrained(model_path, **model_kwargs,)\n",
-    "\n",
-    "# prints metrics on sparsity profile\n",
-    "SparseAutoModel.log_model_load(model, model_path, \"student\", s_delayed)"
+    "model = AutoModelForSequenceClassification.from_pretrained(model_path, **model_kwargs,)"
    ]
   },
   {

integrations/huggingface-transformers/tutorials/text-classification/docs-text-classification-python-qqp.ipynb

@@ -375,13 +375,11 @@
     "model_kwargs = {\"config\": model_config}\n",
     "model_kwargs[\"state_dict\"], s_delayed = SparseAutoModel._loadable_state_dict(model_path)\n",
     "model = AutoModelForSequenceClassification.from_pretrained(model_path, **model_kwargs,)\n",
-    "SparseAutoModel.log_model_load(model, model_path, \"student\", s_delayed) # prints metrics on sparsity profile\n",
     "\n",
     "# initialize teacher using familiar HF AutoModel\n",
     "teacher_kwargs = {\"config\": teacher_config}\n",
     "teacher_kwargs[\"state_dict\"], t_delayed = SparseAutoModel._loadable_state_dict(teacher_path)\n",
-    "teacher = AutoModelForSequenceClassification.from_pretrained(teacher_path, **teacher_kwargs,)\n",
-    "SparseAutoModel.log_model_load(teacher, teacher_path, \"teacher\", t_delayed) # prints metrics on sparsity profile"
+    "teacher = AutoModelForSequenceClassification.from_pretrained(teacher_path, **teacher_kwargs,)\n"
    ]
   },
   {

integrations/huggingface-transformers/tutorials/text-classification/docs-text-classification-python-sick.ipynb

@@ -361,7 +361,6 @@
     "model_kwargs = {\"config\": config}\n",
     "model_kwargs[\"state_dict\"], s_delayed = SparseAutoModel._loadable_state_dict(model_path)\n",
     "model = AutoModelForSequenceClassification.from_pretrained(model_path,**model_kwargs,)\n",
-    "SparseAutoModel.log_model_load(model, model_path, \"student\", s_delayed) # prints metrics on sparsity profile\n",
     "\n",
     "# FYI: there is a factory function called SparseAutoModel that does the same as above\n",
     "# model, teacher = SparseAutoModel.text_classification_from_pretrained_distil(\n",

integrations/huggingface-transformers/tutorials/token-classification/docs-token-classification-python-custom-teacher-wnut.ipynb

@@ -587,8 +587,7 @@
     "# initialize model\n",
     "model_kwargs = {\"config\": config}\n",
     "model_kwargs[\"state_dict\"], s_delayed = SparseAutoModel._loadable_state_dict(model_path)\n",
-    "model = AutoModelForTokenClassification.from_pretrained(model_path, **model_kwargs,)\n",
-    "SparseAutoModel.log_model_load(model, model_path, \"student\", s_delayed) # prints metrics on sparsity profile"
+    "model = AutoModelForTokenClassification.from_pretrained(model_path, **model_kwargs,)"
    ]
   },
   {

setup.py

@@ -55,6 +55,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 = [

src/sparseml/pytorch/utils/logger.py

@@ -45,11 +45,21 @@
     wandb = None
     wandb_err = err
 
+
+try:
+    from clearml import Task
+
+    clearml_err = None
+except Exception as err:
+    clearml = None
+    clearml_err = err
+
 from sparseml.utils import ALL_TOKEN, create_dirs
 
 
 __all__ = [
     "BaseLogger",
+    "ClearMLLogger",
     "LambdaLogger",
     "PythonLogger",
     "TensorBoardLogger",
@@ -628,6 +638,101 @@ def save(
         return True
 
 
+class ClearMLLogger(LambdaLogger):
+    @staticmethod
+    def available() -> bool:
+        """
+        :return: True if wandb is available and installed, False, otherwise
+        """
+        return not clearml_err
+
+    def __init__(
+        self,
+        name: str = "clearml",
+        enabled: bool = True,
+        project_name: str = "sparseml",
+        task_name: str = "",
+    ):
+        if task_name == "":
+            now = datetime.now()
+            task_name = now.strftime("%d-%m-%Y_%H.%M.%S")
+
+        self.task = Task.init(project_name=project_name, task_name=task_name)
+
+        super().__init__(
+            lambda_func=self.log_scalar,
+            name=name,
+            enabled=enabled,
+        )
+
+    def log_hyperparams(
+        self,
+        params: Dict,
+        level: Optional[int] = None,
+    ) -> bool:
+        """
+        :param params: Each key-value pair in the dictionary is the name of the
+            hyper parameter and it's corresponding value.
+        :return: True if logged, False otherwise.
+        """
+        if not self.enabled:
+            return False
+
+        self.task.connect(params)
+        return True
+
+    def log_scalar(
+        self,
+        tag: str,
+        value: float,
+        step: Optional[int] = None,
+        wall_time: Optional[float] = None,
+        level: Optional[int] = None,
+    ) -> bool:
+        """
+        :param tag: identifying tag to log the value with
+        :param value: value to save
+        :param step: global step for when the value was taken
+        :param wall_time: global wall time for when the value was taken,
+            defaults to time.time()
+        :param kwargs: additional logging arguments to support Python and custom loggers
+        :return: True if logged, False otherwise.
+        """
+        logger = self.task.get_logger()
+        # each series is superimposed on the same plot on title
+        logger.report_scalar(
+            title=tag, series=str(level) or tag, value=value, iteration=step
+        )
+        return True
+
+    def log_scalars(
+        self,
+        tag: str,
+        values: Dict[str, float],
+        step: Optional[int] = None,
+        wall_time: Optional[float] = None,
+        level: Optional[int] = None,
+    ) -> bool:
+        """
+        :param tag: identifying tag to log the values with
+        :param values: values to save
+        :param step: global step for when the values were taken
+        :param wall_time: global wall time for when the values were taken,
+            defaults to time.time()
+        :param kwargs: additional logging arguments to support Python and custom loggers
+        :return: True if logged, False otherwise.
+        """
+        for k, v in values.items():
+            self.log_scalar(
+                tag=f"{tag}.{k}",
+                value=v,
+                step=step,
+                wall_time=wall_time,
+                level=level,
+            )
+        return True
+
+
 class SparsificationGroupLogger(BaseLogger):
     """
     Modifier logger that handles outputting values to other supported systems.

src/sparseml/transformers/sparsification/modification/base.py

@@ -24,7 +24,7 @@
 __all__ = ["check_transformers_version"]
 
 _TRANSFORMERS_MIN_VERSION = "4.39.0"
-_TRANSFORMERS_MAX_VERSION = "4.39.2"
+_TRANSFORMERS_MAX_VERSION = "4.39.3"
 
 
 def check_transformers_version(
@@ -56,7 +56,7 @@ def check_transformers_version(
         _LOGGER.warning(
             "Attempting to modify the transformers model to support "
             "the SparseML-specific functionalities. However, the detected "
-            f"transformers version ({current_version}) does not fall within the"
+            f"transformers version ({current_version}) does not fall within the "
             f"supported version range ({min_version} - {max_version}). "
             "This may lead to unexpected behavior. Please ensure that the "
             "correct transformers version is installed."

src/sparseml/transformers/sparsification/sparse_model.py

@@ -40,9 +40,7 @@
     modify_save_pretrained,
 )
 from sparseml.transformers.sparsification.modification import modify_model
-from sparseml.transformers.utils.helpers import resolve_recipe
-from sparseml.utils import download_zoo_training_dir
-from sparseml.utils.fsdp.context import main_process_first_context
+from sparseml.transformers.utils.helpers import download_model_directory, resolve_recipe
 
 
 __all__ = ["SparseAutoModel", "SparseAutoModelForCausalLM", "get_shared_tokenizer_src"]
@@ -101,15 +99,9 @@ def skip(*args, **kwargs):
             else pretrained_model_name_or_path
         )
 
-        if pretrained_model_name_or_path.startswith("zoo:"):
-            _LOGGER.debug(
-                "Passed zoo stub to SparseAutoModelForCausalLM object. "
-                "Loading model from SparseZoo training files..."
-            )
-            with main_process_first_context():
-                pretrained_model_name_or_path = download_zoo_training_dir(
-                    zoo_stub=pretrained_model_name_or_path
-                )
+        pretrained_model_name_or_path = download_model_directory(
+            pretrained_model_name_or_path, **kwargs
+        )
 
         # determine compression format, if any, from the model config
         compressor = infer_compressor_from_model_config(pretrained_model_name_or_path)

src/sparseml/transformers/utils/helpers.py

@@ -34,8 +34,10 @@
 from transformers.trainer_utils import get_last_checkpoint
 from transformers.utils import PaddingStrategy
 
-from huggingface_hub import HUGGINGFACE_CO_URL_HOME, hf_hub_download
+from huggingface_hub import HUGGINGFACE_CO_URL_HOME, HfFileSystem, hf_hub_download
 from sparseml.export.helpers import ONNX_MODEL_NAME
+from sparseml.utils import download_zoo_training_dir
+from sparseml.utils.fsdp.context import main_process_first_context
 from sparsezoo import Model, setup_model
 
 
@@ -52,6 +54,8 @@
     "ALL_TASK_NAMES",
     "create_fake_dataloader",
     "POSSIBLE_TOKENIZER_FILES",
+    "download_repo_from_huggingface_hub",
+    "download_model_directory",
 ]
 
 
@@ -92,6 +96,7 @@ class TaskNames(Enum):
     "special_tokens_map.json",
     "tokenizer_config.json",
 }
+RELEVANT_HF_SUFFIXES = ["json", "md", "bin", "safetensors", "yaml", "yml", "py"]
 
 
 def remove_past_key_value_support_from_config(config: AutoConfig) -> AutoConfig:
@@ -553,3 +558,94 @@ def fetch_recipe_path(target: str):
         recipe_path = hf_hub_download(repo_id=target, filename=DEFAULT_RECIPE_NAME)
 
     return recipe_path
+
+
+def download_repo_from_huggingface_hub(repo_id, **kwargs):
+    """
+    Download relevant model files from the Hugging Face Hub
+    using the huggingface_hub.hf_hub_download function
+
+    Note(s):
+    - Does not download the entire repo, only the relevant files
+    for the model, such as the model weights, tokenizer files, etc.
+    - Does not re-download files that already exist locally, unless
+    the force_download flag is set to True
+
+    :pre-condition: the repo_id must be a valid Hugging Face Hub repo id
+    :param repo_id: the repo id to download
+    :param kwargs: additional keyword arguments to pass to hf_hub_download
+    """
+    hf_filesystem = HfFileSystem()
+    files = hf_filesystem.ls(repo_id)
+
+    if not files:
+        raise ValueError(f"Could not find any files in HF repo {repo_id}")
+
+    # All file(s) from hf_filesystem have "name" key
+    # Extract the file names from the files
+    relevant_file_names = (
+        Path(file["name"]).name
+        for file in files
+        if any(file["name"].endswith(suffix) for suffix in RELEVANT_HF_SUFFIXES)
+    )
+
+    hub_kwargs_names = (
+        "subfolder",
+        "repo_type",
+        "revision",
+        "library_name",
+        "library_version",
+        "cache_dir",
+        "local_dir",
+        "local_dir_use_symlinks",
+        "user_agent",
+        "force_download",
+        "force_filename",
+        "proxies",
+        "etag_timeout",
+        "resume_download",
+        "token",
+        "local_files_only",
+        "headers",
+        "legacy_cache_layout",
+        "endpoint",
+    )
+    hub_kwargs = {name: kwargs[name] for name in hub_kwargs_names if name in kwargs}
+
+    for file_name in relevant_file_names:
+        last_file = hf_hub_download(repo_id=repo_id, filename=file_name, **hub_kwargs)
+
+    # parent directory of the last file is the model directory
+    return str(Path(last_file).parent.resolve().absolute())
+
+
+def download_model_directory(pretrained_model_name_or_path: str, **kwargs):
+    """
+    Download the model directory from the HF hub or SparseZoo if the model
+    is not found locally
+
+    :param pretrained_model_name_or_path: the name of or path to the model to load
+        can be a SparseZoo/HuggingFace model stub
+    :param kwargs: additional keyword arguments to pass to the download function
+    :return: the path to the downloaded model directory
+    """
+    pretrained_model_path: Path = Path(pretrained_model_name_or_path)
+
+    if pretrained_model_path.exists():
+        _LOGGER.debug(
+            "Model directory already exists locally.",
+        )
+        return pretrained_model_name_or_path
+
+    with main_process_first_context():
+        if pretrained_model_name_or_path.startswith("zoo:"):
+            _LOGGER.debug(
+                "Passed zoo stub to SparseAutoModelForCausalLM object. "
+                "Loading model from SparseZoo training files..."
+            )
+            return download_zoo_training_dir(zoo_stub=pretrained_model_name_or_path)
+
+        _LOGGER.debug("Downloading model from HuggingFace Hub.")
+        return download_repo_from_huggingface_hub(
+            repo_id=pretrained_model_name_or_path, **kwargs
+        )