[Python] MLFlow hangs on pytorch log_model for large model file - any help to solve this?

When trying to upload a VLM model artifact to a remote tracking uri, the MLFlow run hangs and fails after a few hours. I am using Mlflow 3.1.1.

The model file is aroung 5GB, we've set the server timeout to be 30 minutes and my local timeout to be 30 minutes, and my env variables are set correctly. I can load and perform inference with the .pth file locally perfectly fine with torch.load - could this be a bug with MLFlow? It seems so simple yet I can't get it to work.

I saved the LoRa model locally after training it using Lightning with:

torch.save(module_model.model, "/workspace/saved_models/vlm_model.pth")

and am attempting to start a run and upload the model file to the remote server with:

model = torch.load("/workspace/saved_models/vlm_model.pth", weights_only=False).to('cuda')

mlflow.pytorch.log_model(
pytorch_model=model,
name="model",
)

Continue reading...