End-to-End Workflow with torchtune¶
In this tutorial, we’ll walk through an end-to-end example of how you can fine-tune, evaluate, optionally quantize and then run generation with your favorite LLM using torchtune. We’ll also go over how you can use some popular tools and libraries from the community seemlessly with torchtune.
Different type of recipes available in torchtune beyond fine-tuning
End-to-end example connecting all of these recipes
Different tools and libraries you can use with torchtune
Be familiar with the overview of torchtune
Make sure to install torchtune
Concepts such as configs and checkpoints
Overview¶
Fine-tuning an LLM is usually only one step in a larger workflow. An example workflow that you might have can look something like this:
Download a popular model from HF Hub
Fine-tune the model using a relevant fine-tuning technique. The exact technique used will depend on factors such as the model, amount and nature of training data, your hardware setup and the end task for which the model will be used
Evaluate the model on some benchmarks to validate model quality
Run some generations to make sure the model output looks reasonable
Quantize the model for efficient inference
[Optional] Export the model for specific environments such as inference on a mobile phone
In this tutorial, we’ll cover how you can use torchtune for all of the above, leveraging integrations with popular tools and libraries from the ecosystem.
We’ll use the Llama2 7B model for this tutorial. You can find a complete set of models supported by torchtune here.
Download Llama2 7B¶
In this tutorial, we’ll use the Hugging Face model weights for the Llama2 7B mode. For more information on checkpoint formats and how these are handled in torchtune, take a look at this tutorial on checkpoints.
To download the HF format Llama2 7B model, we’ll use the tune CLI.
tune download \
meta-llama/Llama-2-7b-hf \
--output-dir <checkpoint_dir> \
--hf-token <ACCESS TOKEN>
Make a note of <checkpoint_dir>
, we’ll use this many times in this tutorial.
Finetune the model using LoRA¶
For this tutorial, we’ll fine-tune the model using LoRA. LoRA is a parameter efficient fine-tuning technique which is especially helpful when you don’t have a lot of GPU memory to play with. LoRA freezes the base LLM and adds a very small percentage of learnable parameters. This helps keep memory associated with gradients and optimizer state low. Using torchtune, you should be able to fine-tune a Llama2 7B model with LoRA in less than 16GB of GPU memory using bfloat16 on a RTX 3090/4090. For more information on how to use LoRA, take a look at our LoRA Tutorial.
We’ll fine-tune using our single device LoRA recipe and use the standard settings from the default config.
This will fine-tune our model using a batch_size=2
and dtype=bfloat16
. With these settings the model
should have a peak memory usage of ~16GB and total training time of around two hours for each epoch.
We’ll need to make some changes to the config to make sure our recipe can access the
right checkpoints.
Let’s look for the right config for this use case by using the tune CLI.
tune ls
RECIPE CONFIG
full_finetune_single_device llama2/7B_full_low_memory
mistral/7B_full_low_memory
full_finetune_distributed llama2/7B_full
llama2/13B_full
mistral/7B_full
lora_finetune_single_device llama2/7B_lora_single_device
llama2/7B_qlora_single_device
mistral/7B_lora_single_device
...
For this tutorial we’ll use the llama2/7B_lora_single_device
config.
The config already points to the HF Checkpointer and the right checkpoint files. All we need to do is update the checkpoint directory for both the model and the tokenizer. Let’s do this using the overrides in the tune CLI while starting training!
tune run lora_finetune_single_device \
--config llama2/7B_lora_single_device \
checkpointer.checkpoint_dir=<checkpoint_dir> \
tokenizer.path=<checkpoint_dir>/tokenizer.model \
checkpointer.output_dir=<checkpoint_dir>
Once training is complete, you’ll see the following in the logs.
[_checkpointer.py:473] Model checkpoint of size 9.98 GB saved to <checkpoint_dir>/hf_model_0001_0.pt
[_checkpointer.py:473] Model checkpoint of size 3.50 GB saved to <checkpoint_dir>/hf_model_0002_0.pt
[_checkpointer.py:484] Adapter checkpoint of size 0.01 GB saved to <checkpoint_dir>/adapter_0.pt
The final trained weights are merged with the original model and split across two checkpoint files similar to the source checkpoints from the HF Hub (see the LoRA Tutorial for more details). In fact the keys will be identical between these checkpoints. We also have a third checkpoint file which is much smaller in size and contains the learnt LoRA adapter weights. For this tutorial, we’ll only use the model checkpoints and not the adapter weights.
Run Evaluation using EleutherAI’s Eval Harness¶
We’ve fine-tuned a model. But how well does this model really do? Let’s run some Evaluations!
torchtune integrates with
EleutherAI’s evaluation harness.
An example of this is available through the
eleuther_eval
recipe. In this tutorial, we’re going to directly use this recipe by
modifying its associated config eleuther_evaluation.yaml
.
Note
For this section of the tutorial, you should first run pip install lm_eval==0.4.*
to install the EleutherAI evaluation harness.
Since we plan to update all of the checkpoint files to point to our fine-tuned checkpoints, let’s first copy over the config to our local working directory so we can make changes. This will be easier than overriding all of these elements through the CLI.
tune cp eleuther_evaluation ./custom_eval_config.yaml \
For this tutorial we’ll use the truthfulqa_mc2
task from the harness.
This task measures a model’s propensity to be truthful when answering questions and
measures the model’s zero-shot accuracy on a question followed by one or more true
responses and one or more false responses. Let’s first run a baseline without fine-tuning.
tune run eleuther_eval --config ./custom_eval_config.yaml
checkpointer.checkpoint_dir=<checkpoint_dir> \
tokenizer.path=<checkpoint_dir>/tokenizer.model
[evaluator.py:324] Running loglikelihood requests
[eleuther_eval.py:195] Eval completed in 121.27 seconds.
[eleuther_eval.py:197] truthfulqa_mc2: {'acc,none': 0.388...
The model has an accuracy around 38.8%. Let’s compare this with the fine-tuned model.
First, we modify custom_eval_config.yaml
to include the fine-tuned checkpoints.
checkpointer:
_component_: torchtune.utils.FullModelHFCheckpointer
# directory with the checkpoint files
# this should match the output_dir specified during
# finetuning
checkpoint_dir: <checkpoint_dir>
# checkpoint files for the fine-tuned model. This should
# match what's shown in the logs above
checkpoint_files: [
hf_model_0001_0.pt,
hf_model_0002_0.pt,
]
output_dir: <checkpoint_dir>
model_type: LLAMA2
# Make sure to update the tokenizer path to the right
# checkpoint directory as well
tokenizer:
_component_: torchtune.models.llama2.llama2_tokenizer
path: <checkpoint_dir>/tokenizer.model
Now, let’s run the recipe.
tune run eleuther_eval --config ./custom_eval_config.yaml
The results should look something like this.
[evaluator.py:324] Running loglikelihood requests
[eleuther_eval.py:195] Eval completed in 121.27 seconds.
[eleuther_eval.py:197] truthfulqa_mc2: {'acc,none': 0.489 ...
Our fine-tuned model gets ~48% on this task, which is ~10 points better than the baseline. Great! Seems like our fine-tuning helped.
Generation¶
We’ve run some evaluations and the model seems to be doing well. But does it really generate meaningful text for the prompts you care about? Let’s find out!
For this, we’ll use the generate recipe and the associated config.
Let’s first copy over the config to our local working directory so we can make changes.
tune cp generation ./custom_generation_config.yaml
Let’s modify custom_generation_config.yaml
to include the following changes.
checkpointer:
_component_: torchtune.utils.FullModelHFCheckpointer
# directory with the checkpoint files
# this should match the output_dir specified during
# finetuning
checkpoint_dir: <checkpoint_dir>
# checkpoint files for the fine-tuned model. This should
# match what's shown in the logs above
checkpoint_files: [
hf_model_0001_0.pt,
hf_model_0002_0.pt,
]
output_dir: <checkpoint_dir>
model_type: LLAMA2
# Make sure to update the tokenizer path to the right
# checkpoint directory as well
tokenizer:
_component_: torchtune.models.llama2.llama2_tokenizer
path: <checkpoint_dir>/tokenizer.model
Once the config is updated, let’s kick off generation! We’ll use the
default settings for sampling with top_k=300
and a
temperature=0.8
. These parameters control how the probabilities for
sampling are computed. These are standard settings for Llama2 7B and
we recommend inspecting the model with these before playing around with
these parameters.
We’ll use a different prompt from the one in the config
tune run generate --config ./custom_generation_config.yaml \
prompt="What are some interesting sites to visit in the Bay Area?"
Once generation is complete, you’ll see the following in the logs.
[generate.py:92] Exploratorium in San Francisco has made the cover of Time Magazine,
and its awesome. And the bridge is pretty cool...
[generate.py:96] Time for inference: 11.61 sec total, 25.83 tokens/sec
[generate.py:99] Memory used: 15.72 GB
Indeed, the bridge is pretty cool! Seems like our LLM knows a little something about the Bay Area!
Speeding up Generation using Quantization¶
We saw that the generation recipe took around 11.6 seconds to generate 300 tokens. One technique commonly used to speed up inference is quantization. torchtune provides an integration with the TorchAO quantization APIs. Let’s first quantize the model using 4-bit weights-only quantization and see if this improves generation speed.
For this, we’ll use the quantization recipe.
Let’s first copy over the config to our local working directory so we can make changes.
tune cp quantization ./custom_quantization_config.yaml
Let’s modify custom_quantization_config.yaml
to include the following changes.
checkpointer:
_component_: torchtune.utils.FullModelHFCheckpointer
# directory with the checkpoint files
# this should match the output_dir specified during
# finetuning
checkpoint_dir: <checkpoint_dir>
# checkpoint files for the fine-tuned model. This should
# match what's shown in the logs above
checkpoint_files: [
hf_model_0001_0.pt,
hf_model_0002_0.pt,
]
output_dir: <checkpoint_dir>
model_type: LLAMA2
Once the config is updated, let’s kick off quantization! We’ll use the default quantization method from the config.
tune run quantize --config ./custom_quantization_config.yaml
Once quantization is complete, you’ll see the following in the logs.
[quantize.py:68] Time for quantization: 19.76 sec
[quantize.py:69] Memory used: 13.95 GB
[quantize.py:82] Model checkpoint of size 3.67 GB saved to <checkpoint_dir>/hf_model_0001_0-4w.pt
Note
Unlike the fine-tuned checkpoints, this outputs a single checkpoint file. This is because our quantization APIs currently don’t support any conversion across formats. As a result you won’t be able to use these quantized models outside of torchtune. But you should be able to use these with the generation and evaluation recipes within torchtune. These results will help inform which quantization methods you should use with your favorite inference engine.
Now that we have the quantized model, let’s re-run generation.
Modify custom_generation_config.yaml
to include the following changes.
checkpointer:
# we need to use the custom torchtune checkpointer
# instead of the HF checkpointer for loading
# quantized models
_component_: torchtune.utils.FullModelTorchTuneCheckpointer
# directory with the checkpoint files
# this should match the output_dir specified during
# finetuning
checkpoint_dir: <checkpoint_dir>
# checkpoint files point to the quantized model
checkpoint_files: [
hf_model_0001_0-4w.pt,
]
output_dir: <checkpoint_dir>
model_type: LLAMA2
# we also need to update the quantizer to what was used during
# quantization
quantizer:
_component_: torchtune.utils.quantization.Int4WeightOnlyQuantizer
groupsize: 256
Once the config is updated, let’s kick off generation! We’ll use the same sampling parameters as before. We’ll also use the same prompt we did with the unquantized model.
tune run generate --config ./custom_generation_config.yaml \
prompt="What are some interesting sites to visit in the Bay Area?"
Once generation is complete, you’ll see the following in the logs.
[generate.py:92] A park in San Francisco that sits at the top of a big hill.
There are lots of trees and a beautiful view of San Francisco...
[generate.py:96] Time for inference: 4.13 sec total, 72.62 tokens/sec
[generate.py:99] Memory used: 17.85 GB
With quantization (and torch compile under the hood), we’ve sped up generation by almost 3x!
Using torchtune checkpoints with other libraries¶
As we mentioned above, one of the benefits of handling of the checkpoint conversion is that you can directly work with standard formats. This helps with interoperability with other libraries since torchtune doesn’t add yet another format to the mix.
Let’s take a look at an example of how this would work with a popular codebase used for running performant inference with LLMs - gpt-fast. This section assumes that you’ve cloned that repository on your machine.
gpt-fast
makes some assumptions about the checkpoint and the availability of
the key-to-file mapping i.e. a file mapping parameter names to the files containing them.
Let’s satisfy these assumptions, by creating this mapping
file. Let’s assume we’ll be using <new_dir>/Llama-2-7B-hf
as the directory
for this. gpt-fast
assumes that the directory with checkpoints has the
same format at the HF repo-id.
import json
import torch
# create the output dictionary
output_dict = {"weight_map": {}}
# Load the checkpoints
sd_1 = torch.load('<checkpoint_dir>/hf_model_0001_0.pt', mmap=True, map_location='cpu')
sd_2 = torch.load('<checkpoint_dir>/hf_model_0002_0.pt', mmap=True, map_location='cpu')
# create the weight map
for key in sd_1.keys():
output_dict['weight_map'][key] = "hf_model_0001_0.pt"
for key in sd_2.keys():
output_dict['weight_map'][key] = "hf_model_0002_0.pt"
with open('<new_dir>/Llama-2-7B-hf/pytorch_model.bin.index.json', 'w') as f:
json.dump(output_dict, f)
Now that we’ve created the weight_map, let’s copy over our checkpoints.
cp <checkpoint_dir>/hf_model_0001_0.pt <new_dir>/Llama-2-7B-hf/
cp <checkpoint_dir>/hf_model_0002_0.pt <new_dir>/Llama-2-7B-hf/
cp <checkpoint_dir>/tokenizer.model <new_dir>/Llama-2-7B-hf/
Once the directory structure is setup, let’s convert the checkpoints and run inference!
cd gpt-fast/
# convert the checkpoints into a format readable by gpt-fast
python scripts/convert_hf_checkpoint.py \
--checkpoint_dir <new_dir>/Llama-2-7B-hf/ \
--model 7B
# run inference using the converted model
python generate.py \
--compile \
--checkpoint_path <new_dir>/Llama-2-7B-hf/model.pth \
--device cuda
The output should look something like this:
Hello, my name is Justin. I am a middle school math teacher
at WS Middle School ...
Time for inference 5: 1.94 sec total, 103.28 tokens/sec
Bandwidth achieved: 1391.84 GB/sec
And thats it! Try your own prompt!
Uploading your model to the Hugging Face Hub¶
Your new model is working great and you want to share it with the world. The easiest way to do this
is utilizing the huggingface-cli
command, which works seamlessly with torchtune. Simply point the CLI
to your finetuned model directory like so:
huggingface-cli upload <hf-repo-id> <checkpoint-dir>
The command should output a link to your repository on the Hub. If the repository doesn’t exist yet, it will be created automatically:
https://huggingface.co/<hf-repo-id>/tree/main/.
Note
Before uploading, make sure you are authenticated with Hugging Face by running huggingface-cli login
.
For more details on the huggingface-cli upload
feature check out the Hugging Face docs.
Hopefully this tutorial gave you some insights into how you can use torchtune for your own workflows. Happy Tuning!