Repository: mlflow/mlflow

Issue:
[FR] Add support for using HF Accelerate and efficiently loading weights
### Willingness to contribute

Yes. I can contribute this feature independently.

### Proposal Summary

HuggingFace models are often large Transformer models which benefit from GPU acceleration. 

This [article](https://huggingface.co/docs/accelerate/usage_guides/big_modeling) explains why default behavior used in loading PyTorch models results in multiple copies of the model coming to exist in host RAM before the model is loaded to the GPU - this results in OOM errors even on instances where the model should be able to load.

This is because:
1. MLflow does not support sharded checkpoints during save time for models that are logged.
2. MLflow does not use Accelerate to load the model and fails to use the features that reduce memory usage
3. MLflow does not support using the device_map flag.

### Motivation

> #### What is the use case for this feature?
HuggingFace models are often large Transformer models which benefit from GPU acceleration.  Any user of LLMs with HF will benefit from this feature.

> #### Why is this use case valuable to support for MLflow users in general?

> #### Why is this use case valuable to support for your project(s) or organization?

> #### Why is it currently difficult to achieve this use case?


### Details

_No response_

### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [X] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:

Resolving patch:
diff --git a/mlflow/environment_variables.py b/mlflow/environment_variables.py
--- a/mlflow/environment_variables.py
+++ b/mlflow/environment_variables.py
@@ -209,6 +209,32 @@ def get(self):
     "MLFLOW_DEFAULT_PREDICTION_DEVICE", str, None
 )
 
+#: Specifies to Huggingface whether to use the automatic device placement logic of
+# HuggingFace accelerate. If it's set to false, the low_cpu_mem_usage flag will not be
+# set to True and device_map will not be set to "auto".
+MLFLOW_HUGGINGFACE_DISABLE_ACCELERATE_FEATURES = _BooleanEnvironmentVariable(
+    "MLFLOW_DISABLE_HUGGINGFACE_ACCELERATE_FEATURES", False
+)
+
+#: Specifies to Huggingface whether to use the automatic device placement logic of
+# HuggingFace accelerate. If it's set to false, the low_cpu_mem_usage flag will not be
+# set to True and device_map will not be set to "auto".
+MLFLOW_HUGGINGFACE_USE_DEVICE_MAP = _BooleanEnvironmentVariable(
+    "MLFLOW_HUGGINGFACE_USE_DEVICE_MAP", True
+)
+
+#: Specifies to Huggingface to use the automatic device placement logic of HuggingFace accelerate.
+#: This can be set to values supported by the version of HuggingFace Accelerate being installed.
+MLFLOW_HUGGINGFACE_DEVICE_MAP_STRATEGY = _EnvironmentVariable(
+    "MLFLOW_HUGGINGFACE_DEVICE_MAP_STRATEGY", str, "auto"
+)
+
+#: Specifies to Huggingface to use the low_cpu_mem_usage flag powered by HuggingFace accelerate.
+#: If it's set to false, the low_cpu_mem_usage flag will be set to False.
+MLFLOW_HUGGINGFACE_USE_LOW_CPU_MEM_USAGE = _BooleanEnvironmentVariable(
+    "MLFLOW_HUGGINGFACE_USE_LOW_CPU_MEM_USAGE", True
+)
+
 #: Specifies whether or not to allow using a file URI as a model version source.
 #: Please be aware that setting this environment variable to True is potentially risky
 #: because it can allow access to arbitrary files on the specified filesystem
diff --git a/mlflow/transformers.py b/mlflow/transformers.py
--- a/mlflow/transformers.py
+++ b/mlflow/transformers.py
@@ -32,7 +32,13 @@
     LOG_MODEL_PARAM_DOCS,
     docstring_version_compatibility_warning,
 )
-from mlflow.environment_variables import MLFLOW_DEFAULT_PREDICTION_DEVICE
+from mlflow.environment_variables import (
+    MLFLOW_DEFAULT_PREDICTION_DEVICE,
+    MLFLOW_HUGGINGFACE_DISABLE_ACCELERATE_FEATURES,
+    MLFLOW_HUGGINGFACE_USE_DEVICE_MAP,
+    MLFLOW_HUGGINGFACE_DEVICE_MAP_STRATEGY,
+    MLFLOW_HUGGINGFACE_USE_LOW_CPU_MEM_USAGE,
+)
 from mlflow.utils.environment import (
     _mlflow_conda_env,
     _validate_env_arguments,
@@ -100,7 +106,7 @@ def _model_packages(model) -> List[str]:
     """
     engine = _get_engine_type(model)
     if engine == "torch":
-        return ["torch", "torchvision"]
+        return ["torch", "torchvision", "accelerate"]
     else:
         return [engine]
 
@@ -825,12 +831,38 @@ def is_gpu_available():
     return is_gpu
 
 
+def _try_load_model_with_device(model_instance, model_path, device, conf):
+    load_model_conf = {}
+    # Assume if torch_dtype was specified in the conf, then it must be with a
+    # pipeline for which it's compatible.
+    if _TORCH_DTYPE_KEY in conf:
+        load_model_conf[_TORCH_DTYPE_KEY] = conf[_TORCH_DTYPE_KEY]
+
+    try:
+        load_model_conf["device"] = device
+        model = model_instance.from_pretrained(model_path, **load_model_conf)
+    except (ValueError, TypeError, NotImplementedError):
+        _logger.warning("Could not specify device parameter for this pipeline type")
+        load_model_conf.pop("device", None)
+        model = model_instance.from_pretrained(model_path, **load_model_conf)
+    return model
+
+
 def _load_model(path: str, flavor_config, return_type: str, device=None, **kwargs):
     """
     Loads components from a locally serialized ``Pipeline`` object.
     """
     import transformers
 
+    model_instance = getattr(transformers, flavor_config[_PIPELINE_MODEL_TYPE_KEY])
+    local_path = pathlib.Path(path)
+    pipeline_path = local_path.joinpath(
+        flavor_config.get(_PIPELINE_BINARY_KEY, _PIPELINE_BINARY_FILE_NAME)
+    )
+    conf = {
+        "task": flavor_config[_TASK_KEY],
+    }
+
     if device is None:
         if MLFLOW_DEFAULT_PREDICTION_DEVICE.get():
             try:
@@ -839,19 +871,39 @@ def _load_model(path: str, flavor_config, return_type: str, device=None, **kwarg
                 device = _TRANSFORMERS_DEFAULT_CPU_DEVICE_ID
         elif is_gpu_available():
             device = _TRANSFORMERS_DEFAULT_GPU_DEVICE_ID
+    # Note that we don't set the device in the conf yet because device is
+    # incompatible with device_map.
+    accelerate_model_conf = {}
+    if MLFLOW_HUGGINGFACE_USE_DEVICE_MAP.get():
+        device_map_strategy = MLFLOW_HUGGINGFACE_DEVICE_MAP_STRATEGY.get()
+        conf["device_map"] = device_map_strategy
+        accelerate_model_conf["device_map"] = device_map_strategy
+        # Cannot use device with device_map
+        device = None
 
-    local_path = pathlib.Path(path)
-    pipeline_path = local_path.joinpath(
-        flavor_config.get(_PIPELINE_BINARY_KEY, _PIPELINE_BINARY_FILE_NAME)
-    )
-
-    model_instance = getattr(transformers, flavor_config[_PIPELINE_MODEL_TYPE_KEY])
-    conf = {
-        "task": flavor_config[_TASK_KEY],
-        "model": model_instance.from_pretrained(pipeline_path),
-    }
     if device is not None:
         conf["device"] = device
+        accelerate_model_conf["device"] = device
+
+    if _TORCH_DTYPE_KEY in flavor_config or _TORCH_DTYPE_KEY in kwargs:
+        if _TORCH_DTYPE_KEY in kwargs:
+            dtype_val = kwargs[_TORCH_DTYPE_KEY]
+        else:
+            dtype_val = _deserialize_torch_dtype_if_exists(flavor_config)
+        conf[_TORCH_DTYPE_KEY] = dtype_val
+        accelerate_model_conf[_TORCH_DTYPE_KEY] = dtype_val
+
+    accelerate_model_conf["low_cpu_mem_usage"] = MLFLOW_HUGGINGFACE_USE_LOW_CPU_MEM_USAGE.get()
+
+    if not MLFLOW_HUGGINGFACE_DISABLE_ACCELERATE_FEATURES.get():
+        try:
+            model = model_instance.from_pretrained(pipeline_path, **accelerate_model_conf)
+        except (ValueError, TypeError, NotImplementedError):
+            model = _try_load_model_with_device(model_instance, pipeline_path, device, conf)
+    else:
+        model = _try_load_model_with_device(model_instance, pipeline_path, device, conf)
+
+    conf["model"] = model
 
     if _PROCESSOR_TYPE_KEY in flavor_config:
         conf[_PROCESSOR_KEY] = _load_component(
@@ -863,9 +915,6 @@ def _load_model(path: str, flavor_config, return_type: str, device=None, **kwarg
         component_type = flavor_config[component_type_key]
         conf[component_key] = _load_component(local_path, component_key, component_type)
 
-    if _TORCH_DTYPE_KEY in flavor_config:
-        conf[_TORCH_DTYPE_KEY] = _deserialize_torch_dtype_if_exists(flavor_config)
-
     for key in _METADATA_PIPELINE_SCALAR_CONFIG_KEYS:
         if key in flavor_config:
             conf[key] = flavor_config[key]


Test patch:
diff --git a/tests/transformers/test_transformers_model_export.py b/tests/transformers/test_transformers_model_export.py
--- a/tests/transformers/test_transformers_model_export.py
+++ b/tests/transformers/test_transformers_model_export.py
@@ -2702,17 +2702,19 @@ def test_load_as_pipeline_preserves_framework_and_dtype(model_path):
     base_loaded = mlflow.transformers.load_model(model_path)
     assert base_loaded.torch_dtype == torch.bfloat16
     assert base_loaded.framework == "pt"
+    assert base_loaded.model.dtype == torch.bfloat16
 
     loaded_pipeline = mlflow.transformers.load_model(model_path, torch_dtype=torch.float64)
 
     assert loaded_pipeline.torch_dtype == torch.float64
     assert loaded_pipeline.framework == "pt"
+    assert loaded_pipeline.model.dtype == torch.float64
 
     prediction = loaded_pipeline.predict("Hello there. How are you today?")
-    assert prediction == [{"translation_text": "Bonjour, comment êtes-vous aujourd'hui ?"}]
+    assert prediction[0]["translation_text"].startswith("Bonjour")
 
 
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float64, torch.int16])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float64])
 @pytest.mark.skipif(
     Version(transformers.__version__) < Version("4.26.1"), reason="Feature does not exist"
 )
@@ -2739,7 +2741,7 @@ def test_load_pyfunc_mutate_torch_dtype(model_path, dtype):
 
     prediction = loaded_pipeline.predict("Hello there. How are you today?")
 
-    assert prediction == "Bonjour, comment êtes-vous aujourd'hui ?"
+    assert prediction.startswith("Bonjour")
 
 
 @pytest.mark.skipif(