Logging finetuned model using transformers mlflow flavor in azure

Intermediate
5

Gonna mark this as solved because I figured out the solution.

The issue seems to be that when working in an azure job it has issues when dealing with AutoPeftModelForCausalLM and by association I assume Peft models in general. It struggles to use the variable that you assign to the peft model with the error that I mentioned above. If you instead refer to the models location in the mlflow.transformers.log_model args you can solve the problem with some extra steps. Code here:

peft_model = AutoPeftModelForCausalLM.from_pretrained(
    'models/finetuned_llama3', 
    local_files_only=True,
    config=LoraConfig(
        r=lora_config_dict["r"],
        lora_alpha=lora_config_dict["lora_alpha"],
        target_modules=lora_config_dict["target_modules"],
        lora_dropout=lora_config_dict["lora_dropout"],
        bias=lora_config_dict["bias"],
        task_type=lora_config_dict["task_type"]
    ), 
    device_map="cuda"
)
peft_model.model.config.quantization_config.use_exllama = True
peft_model.model.config.quantization_config.exllama_config = {"version": 2}

with open("models/finetuned_llama3/config.json", "w") as f:
    json.dump(peft_model.config.to_dict(), f, indent=4)

mlflow.transformers.log_model(
    transformers_model='models/finetuned_llama3',
    artifact_path="models/finetuned_llama3",
    registered_model_name="huggingface-finetuned-model",
    task="text-generation",
    save_pretrained=True
)

The extra step you need to take is adding the config file from you peft model to the directory that your model is saved in. This is because the config file you need is an attribute of the peft mode but if not in the folder that your finetuned model is saved in. The log model statement complains about that so you need to add the config file to that folder (seen in my json.dump).

Hopefully if someone else has this issue I hope they find this thread.

1 Like