How to save finetuned HuggingFace models in AzureML job

matsuo_basho 30

I launch an AzureML job that finetunes a HuggingFace model through the CLI. The model is a text generation model, so I cannot use the finetune pipeline from the model registry. I can't figure out how to actually save the finetuned model to a blobstore.

Below is the part of the code where the model is saved. This works locally, but I'm unable to find the output anywhere after the job runs if run like this in AzureML.

    trainer = Trainer(
        model,
        training_args,
        train_dataset=tokenized_dataset['train'],
        eval_dataset=tokenized_dataset['test'],
        data_collator=data_collator,
        tokenizer=tokenizer)

    trainer.train()

    trainer.save_model()

I also tried passing the model_dir argument to save_model, but that didn't work. For reference, the base model I'm using: Salesforce/codegen-350M-mono

Ramr-msft 17,736 Reputation points

2023-11-20T07:07:38.1366667+00:00
Thanks for the question, To save your model to the output directory, you can use the AZUREML_DATAREFERENCE_output environment variable to get the path to the output directory. Here’s how you can modify your code:

import os

# ... trainer.train() output_dir = os.environ['AZUREML_DATAREFERENCE_output'] trainer.save_model(output_dir)

This will save your model to the output directory of the job After the job is finished, you can find the model in the default datastore in the azureml/JobRunId directory.

Remember to replace JobRunId with the actual ID of your job. You can find this ID in the Azure Machine Learning Studio.

Share via

How to save finetuned HuggingFace models in AzureML job

Your answer