# Save, load and do inference with fine-tuned model

**URL:** <https://discuss.huggingface.co/t/save-load-and-do-inference-with-fine-tuned-model/76291>\
**Category:** Beginners\
**Created:** [March 6, 2024, 11:28pm UTC](https://discuss.huggingface.co/t/save-load-and-do-inference-with-fine-tuned-model/76291 "2024-03-06T23:28:43Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![Iamexperimenting](https://avatars.discourse-cdn.com/v4/letter/i/839c29/32.png) [@Iamexperimenting](https://discuss.huggingface.co/u/Iamexperimenting)\
**Post date:** [March 6, 2024, 11:28pm UTC](https://discuss.huggingface.co/t/save-load-and-do-inference-with-fine-tuned-model/76291/1 "2024-03-06T23:28:43Z")

</div>

Hi team,

I’m using huggingface framework to fine-tune LLMs. Currently, I’m using mistral model. I wanted to save the fine-tuned model and load it later and do inference with it.

Since, I’m new to Huggingface framework I would like to get your guidance on saving, loading, and inferencing.

I remember in PyTorch we need to use `with torch.no_grad():` context manager to do inference. But, I’m not seeing such thing in HuggingFace.

@nielsr Could you please guide me here?

```auto

import pandas as pd
import torch
from datasets import Dataset, load_dataset
from random import randrange
from peft import LoraConfig, get_peft_model, AutoPeftModelForCausalLM, prepare_model_for_int8_training
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TrainingArguments
from trl import SFTTrainer
import warnings
warnings.filterwarnings("ignore")

df = pd.read_csv("train.csv")
train = Dataset.from_pandas(df)
model_id = "mistralai/Mistral-7B-Instruct-v0.1"

tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True,device_map="auto")
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

model = AutoModelForCausalLM.from_pretrained(model_id,
                                             load_in_4bit = True,
                                             torch_dtype = torch.float16,
                                             device_map="auto")
model.resize_token_embeddings(len(tokenizer))
model = prepare_model_for_int8_training(model)
peft_config = LoraConfig(
                          lora_alpha=16,
                          lora_dropout=0.1,
                          r=64,
                          bias="none",
                          task_type="CAUSAL_LM"
                        )
model = get_peft_model(model, peft_config)

args = TrainingArguments(
    output_dir='custom_domain',
    num_train_epochs=2, # adjust based on the data size
    per_device_train_batch_size=8, # use 4 if you have more GPU RAM
    optim = "adamw_torch",
    logging_steps = 100,
    save_total_limit = 2,
    save_strategy = "no",
    load_best_model_at_end=False,
    learning_rate=2e-4,
    fp16=not torch.cuda.is_bf16_supported(),
    bf16 = torch.cuda.is_bf16_supported(),
    evaluation_strategy="epoch",
    seed=42,
    warmup_ratio = 0.1,
    lr_scheduler_type = "linear",
    report_to="none",
    torch_compile = True
    #dataloader_num_workers = 4
)

# Create the trainer
trainer = SFTTrainer(
    model=model,
    train_dataset=train,
    # eval_dataset=test,
    dataset_text_field='text',
    peft_config=peft_config,
    max_seq_length=512,
    tokenizer=tokenizer,
    args=args,
    packing=False,
)

trainer.train()

```

I couldn’t find any complete example in the HuggingFace page. Can you please provide me relevant example?

Additionally, I noticed mistral model is not able to compile using HuggingFace framework during fine-tuning.

---

<div class="post-metadata">

**Author:** ![CKeibel](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/ckeibel/32/55416_2.png) [@CKeibel](https://discuss.huggingface.co/u/CKeibel)\
**Post date:** [March 7, 2024, 8:40am UTC](https://discuss.huggingface.co/t/save-load-and-do-inference-with-fine-tuned-model/76291/2 "2024-03-07T08:40:37Z")

</div>

For inference, causal models have the [generate()](https://github.com/huggingface/transformers/blob/31ec2cb2badfbdd4c1ac9c6c9b8a74e974984206/src/transformers/generation_utils.py#L831) function which is wrapped in torch.no\_grad.

```auto
tokenized_prompt = tokenizer(prompt, return_tensors="pt").to("cuda")
outputs = model.generate(**tokenized_prompt)
tokenizer.decode(outputs[0])

```

PEFT models requires explicit positional arguments (input\_ids)

```auto
input_ids= tokenizer(prompt, return_tensors="pt").input_ids.to("cuda")
outputs = model.generate(input_ids=input_ids)
tokenizer.decode(outputs[0])

```

You can load a model with the from\_pretained function. Instead of the huggingface model\_id, enter the path to your saved model.

```auto
model = AutoModelForCausalLM.from_pretrained("path/to/model.pt")

```

Saving works via the `save_pretrained()` function.

```auto
model.save_pretrained("path/to/model.pt")

```

Since you have trained the model with PEFT, you can also only save and load the adapter. [Here](https://huggingface.co/docs/peft/developer_guides/lora) is a good guide on how to do this.

**EDIT:**  
You specified a `output_dir` in your `TrainingArguments` which should help you saving your model. (Sorry, I never actually use the trainer and fine-tune in plain pytorch.)

---

<div class="post-metadata">

**Author:** ![Iamexperimenting](https://avatars.discourse-cdn.com/v4/letter/i/839c29/32.png) [@Iamexperimenting](https://discuss.huggingface.co/u/Iamexperimenting)\
**Post date:** [March 7, 2024, 2:13pm UTC](https://discuss.huggingface.co/t/save-load-and-do-inference-with-fine-tuned-model/76291/3 "2024-03-07T14:13:57Z")

</div>

I’m seeing different methods to save the fine-tuned model. That confuses me.

```auto
Example1 : model.save_pretrained('./output/')
Example1 : trainer.save_model('./output/')
Example1 : trainer.model.save_pretrained('./output/')

```

and some example with `merge and unload`.

@nielsr can you provide some example for fine-tuned model?

---

<div class="post-metadata">

**Author:** ![nielsr](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/nielsr/32/39617_2.png) [@nielsr](https://discuss.huggingface.co/u/nielsr)\
**Post date:** [March 8, 2024, 6:13pm UTC](https://discuss.huggingface.co/t/save-load-and-do-inference-with-fine-tuned-model/76291/4 "2024-03-08T18:13:02Z")

</div>

Hi,

@CKeibel explained it well. If you’re using the Trainer API, you can specify an output\_dir to which it will automatically save the model. You can specify the saving frequency in the `TrainingArguments` (like every epoch, every x steps, etc.).

Afterwards, you can load the model using the `from_pretrained` method, by specifying the path to the folder.

See [this demo notebook](https://github.com/NielsRogge/Transformers-Tutorials/blob/master/Mistral/Supervised_fine_tuning_(SFT)_of_an_LLM_using_Hugging_Face_tooling.ipynb) which showcases fine-tuning Mistral-7B which includes an inference section.