-
Notifications
You must be signed in to change notification settings - Fork 4k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[BUG] Azure Databricks disk_offload error #11907
Comments
@vitaliy-sharandin Thanks for reporting this. Could you share your model logging code? |
I ran the following code but could not reproduce the error: %pip install -U git+https://github.com/huggingface/transformers torch accelerate==0.29.3 mlflow
dbutils.library.restartPython()
########
import transformers
import torch
model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
pipeline = transformers.pipeline(
"text-generation",
model=model_id,
model_kwargs={"torch_dtype": torch.bfloat16},
token="...",
)
import mlflow
import uuid
mlflow.set_registry_uri("databricks-uc")
with mlflow.start_run() as run:
mlflow.transformers.log_model(pipeline, "model")
mlflow.register_model(
model_uri=f"runs:/{run.info.run_id}/model",
name=f"..."
) |
The main difference between our code is that I am first fine-tuning adapters with peft and trying to register the run which has only adapters saved and base model reference without model weights. I have also read MLFLow Transformers guide which specifies that you don't need to use Here is my notebook: |
Thanks for the notebook! Let me run the notebook and see If I can reproduce the issue. |
@vitaliy-sharandin Can you try inserting this code before loading the model to see if it can fix the error? def get_model_with_peft_adapter(base_model, peft_adapter_path):
from peft import PeftModel
return PeftModel.from_pretrained(base_model, peft_adapter_path, offload_folder="offload")
mlflow.transformers.get_model_with_peft_adapter = get_model_with_peft_adapter Not sure if |
It doesn't quite make sense, as I don't have adapters to load pre-model-tuning, so I don't have value for |
@mlflow/mlflow-team Please assign a maintainer and start triaging this issue. |
@vitaliy-sharandin the traceback says |
@harupy Sorry, I have misunderstood your code at first. I did what you've proposed and it led to new error, please check out the notebook. |
@harupy Any updates? |
Issues Policy acknowledgement
Where did you encounter this bug?
Azure Databricks
Willingness to contribute
No. I cannot contribute a bug fix at this time.
MLflow version
mlflow==2.12.1
System information
Describe the problem
I encounter a
disk_offload
error whenever I try to register model in Unity Catalogue.Tracking information
Code to reproduce issue
Stack trace
Other info / logs
What component(s) does this bug affect?
area/artifacts
: Artifact stores and artifact loggingarea/build
: Build and test infrastructure for MLflowarea/deployments
: MLflow Deployments client APIs, server, and third-party Deployments integrationsarea/docs
: MLflow documentation pagesarea/examples
: Example codearea/model-registry
: Model Registry service, APIs, and the fluent client calls for Model Registryarea/models
: MLmodel format, model serialization/deserialization, flavorsarea/recipes
: Recipes, Recipe APIs, Recipe configs, Recipe Templatesarea/projects
: MLproject format, project running backendsarea/scoring
: MLflow Model server, model deployment tools, Spark UDFsarea/server-infra
: MLflow Tracking server backendarea/tracking
: Tracking Service, tracking client APIs, autologgingWhat interface(s) does this bug affect?
area/uiux
: Front-end, user experience, plotting, JavaScript, JavaScript dev serverarea/docker
: Docker use across MLflow's components, such as MLflow Projects and MLflow Modelsarea/sqlalchemy
: Use of SQLAlchemy in the Tracking Service or Model Registryarea/windows
: Windows supportWhat language(s) does this bug affect?
language/r
: R APIs and clientslanguage/java
: Java APIs and clientslanguage/new
: Proposals for new client languagesWhat integration(s) does this bug affect?
integrations/azure
: Azure and Azure ML integrationsintegrations/sagemaker
: SageMaker integrationsintegrations/databricks
: Databricks integrationsThe text was updated successfully, but these errors were encountered: