Technical Documentation: meta-pytorch/torchtune
ℹ️ Provenance: Multi-Tier Ingestion · CodeWiki Reference
Overview (`README.md`)
⚠️ Torchtune is no longer actively maintained: torchtune development wound down in 2025 — see The future of torchtune. Huge thanks to the 150+ contributors who made this library what it was.
torchtune
Overview | Installation | Get Started | Documentation | Community | Citing torchtune | License
📣 Recent updates 📣
- May 2025: torchtune has added support for Qwen3 models! Check out all the configs here
- April 2025: Llama4 is now available in torchtune! Try out our full and LoRA finetuning configs here
- February 2025: Multi-node training is officially open for business in torchtune! Full finetune on multiple nodes to take advantage of larger batch sizes and models.
- December 2024: torchtune now supports Llama 3.3 70B! Try it out by following our installation instructions here, then run any of the configs here.
- November 2024: torchtune has released v0.4.0 which includes stable support for exciting features like activation offloading and multimodal QLoRA
- November 2024: torchtune has added Gemma2 to its models!
- October 2024: torchtune added support for Qwen2.5 models - find the configs here
- September 2024: torchtune has support for Llama 3.2 11B Vision, Llama 3.2 3B, and Llama 3.2 1B models! Try them out by following our installation instructions here, then run any of the text configs here or vision configs here.
Overview 📚
torchtune is a PyTorch library for easily authoring, post-training, and experimenting with LLMs. It provides:
- Hackable training recipes for SFT, knowledge distillation, DPO, PPO, GRPO, and quantization-aware training
- Simple PyTorch implementations of popular LLMs like Llama, Gemma, Mistral, Phi, Qwen, and more
- Best-in-class memory efficiency, performance improvements, and scaling, utilizing the latest PyTorch APIs
- YAML configs for easily configuring training, evaluation, quantization or inference recipes
Post-training recipes
torchtune supports the entire post-training lifecycle. A successful post-trained model will likely utilize several of the below methods.
Supervised Finetuning (SFT)
| Type of Weight Update | 1 Device | >1 Device | >1 Node |
|---|---|---|---|
| Full | ✅ | ✅ | ✅ |
| LoRA/QLoRA | ✅ | ✅ | ✅ |
Example: tune run lora_finetune_single_device --config llama3_2/3B_lora_single_device
You can also run e.g. tune ls lora_finetune_single_device for a full list of available configs.
[Knowledge Distillation (KD)](https://pytorch.org/torchtune/0.4/tutorials/llama_kd_tutorial.html)
| Type of Weight Update | 1 Device | >1 Device | >1 Node |
|---|---|---|---|
| Full | ❌ | ❌ | ❌ |
| LoRA/QLoRA | ✅ | ✅ | ❌ |
Example: tune run knowledge_distillation_distributed --config qwen2/1.5B_to_0.5B_KD_lora_distributed
You can also run e.g. tune ls knowledge_distillation_distributed for a full list of available configs.
Reinforcement Learning / Reinforcement Learning from Human Feedback (RLHF)
| Method | Type of Weight Update | 1 Device | >1 Device | >1 Node |
|---|---|---|---|---|
| DPO | Full | ❌ | ✅ | ❌ |
| LoRA/QLoRA | ✅ | ✅ | ❌ | |
| PPO | Full | ✅ | ❌ | ❌ |
| LoRA/QLoRA | ❌ | ❌ | ❌ | |
| GRPO | Full | 🚧 | ✅ | ✅ |
| LoRA/QLoRA | ❌ | ❌ | ❌ |
Example: tune run lora_dpo_single_device --config llama3_1/8B_dpo_single_device
You can also run e.g. tune ls full_dpo_distributed for a full list of available configs.
[Quantization-Aware Training (QAT)](https://pytorch.org/torchtune/main/tutorials/qat_finetune.html)
| Type of Weight Update | 1 Device | >1 Device | >1 Node |
|---|---|---|---|
| Full | ✅ | ✅ | ❌ |
| LoRA/QLoRA | ❌ | ✅ | ❌ |
Example: tune run qat_distributed --config llama3_1/8B_qat_lora
You can also run e.g. tune ls qat_distributed or tune ls qat_single_device for a full list of available configs.
The above configs are just examples to get you started. The full list of recipes can be found here. If you'd like to work on one of the gaps you see, please submit a PR! If there's a entirely new post-training method you'd like to see implemented in torchtune, feel free to open an Issue.
Models
For the above recipes, torchtune supports many state-of-the-art models available on the Hugging Face Hub or Kaggle Hub. Some of our supported models:
| Model | Sizes |
|---|---|
| Llama4 | Scout (17B x 16E) [models, configs] |
| Llama3.3 | 70B [models, configs] |
| Llama3.2-Vision | 11B, 90B [models, configs] |
| Llama3.2 | 1B, 3B [models, configs] |
| Llama3.1 | 8B, 70B, 405B [models, configs] |
| Mistral | 7B [models, configs] |
| Gemma2 | 2B, 9B, 27B [models, configs] |
| Microsoft Phi4 | 14B [models, configs] |
| Microsoft Phi3 | Mini [models, configs] |
| Qwen3 | 0.6B, 1.7B, 4B, 8B, 14B, 32B [models, configs] |
| Qwen2.5 | 0.5B, 1.5B, 3B, 7B, 14B, 32B, 72B [models, configs] |
| Qwen2 | 0.5B, 1.5B, 7B [models, configs] |
We're always adding new models, but feel free to file an issue if there's a new one you would like to see in torchtune.
Memory and training speed
Below is an example of the memory requirements and training speed for different Llama 3.1 models.
[!NOTE]
For ease of comparison, all the below numbers are provided for batch size 2 (without gradient accumulation), a dataset packed to sequence length 2048, and torch compile enabled.
If you are interested in running on different hardware or with different models, check out our documentation on memory optimizations here to find the right setup for you.
| Model | Finetuning Method | Runnable On | Peak Memory per GPU | Tokens/sec * |
|---|---|---|---|---|
| Llama 3.1 8B | Full finetune | 1x 4090 | 18.9 GiB | 1650 |
| Llama 3.1 8B | Full finetune | 1x A6000 | 37.4 GiB | 2579 |
| Llama 3.1 8B | LoRA | 1x 4090 | 16.2 GiB | 3083 |
| Llama 3.1 8B | LoRA | 1x A6000 | 30.3 GiB | 4699 |
| Llama 3.1 8B | QLoRA | 1x 4090 | 7.4 GiB | 2413 |
| Llama 3.1 70B | Full finetune | 8x A100 | 13.9 GiB ** | 1568 |
| Llama 3.1 70B | LoRA | 8x A100 | 27.6 GiB | 3497 |
| Llama 3.1 405B | QLoRA | 8x A100 | 44.8 GB | 653 |
*= Measured over one full training epoch
**= Uses CPU offload with fused optimizer
Optimization flags
torchtune exposes a number of levers for memory efficiency and performance. The table below demonstrates the effects of applying some of these techniques sequentially to the Llama 3.2 3B model. Each technique is added on top of the previous one, except for LoRA and QLoRA, which do not use optimizer_in_bwd or AdamW8bit optimizer.
Baseline uses Recipe=full_finetune_single_device, Model=Llama 3.2 3B, Batch size=2, Max sequence length=4096, Precision=bf16, Hardware=A100
| Technique | Peak Memory Active (GiB) | % Change Memory vs Previous | Tokens Per Second | % Change Tokens/sec vs Previous |
|---|---|---|---|---|
| Baseline | 25.5 | - | 2091 | - |
| + Packed Dataset | 60.0 | +135.16% | 7075 | +238.40% |
| + Compile | 51.0 | -14.93% | 8998 | +27.18% |
| + Chunked Cross Entropy | 42.9 | -15.83% | 9174 | +1.96% |
| + Activation Checkpointing | 24.9 | -41.93% | 7210 | -21.41% |
| + Fuse optimizer step into backward | 23.1 | -7.29% | 7309 | +1.38% |
| + Activation Offloading | 21.8 | -5.48% | 7301 | -0.11% |
| + 8-bit AdamW | 17.6 | -19.63% | 6960 | -4.67% |
| LoRA | 8.5 | -51.61% | 8210 | +17.96% |
| QLoRA | 4.6 | -45.71% | 8035 | -2.13% |
The final row in the table vs baseline + Packed Dataset uses 81.9% less memory with a 284.3% increase in tokens per second.
Command to reproduce final row.
bash
tune run lora_finetune_single_device --config llama3_2/3B_qlora_single_device \
dataset.packed=True \
compile=True \
loss=torchtune.modules.loss.CEWithChunkedOutputLoss \
enable_activation_checkpointing=True \
optimizer_in_bwd=False \
enable_activation_offloading=True \
optimizer=torch.optim.AdamW \
tokenizer.max_seq_len=4096 \
gradient_accumulation_steps=1 \
epochs=1 \
batch_size=2
Installation 🛠️
torchtune is only tested with the latest stable PyTorch release (currently 2.6.0) as well as the preview nightly version, and leverages
torchvision for finetuning multimodal LLMs and torchao for the latest in quantization techniques; you should install these as well.
Install stable release
bash
# Install stable PyTorch, torchvision, torchao stable releases
pip install torch torchvision torchao
pip install torchtuneInstall nightly release
bash
# Install PyTorch, torchvision, torchao nightlies.
pip install --pre --upgrade torch torchvision torchao --index-url https://download.pytorch.org/whl/nightly/cu126 # full options are cpu/cu118/cu124/cu126/xpu/rocm6.2/rocm6.3/rocm6.4
pip install --pre --upgrade torchtune --extra-index-url https://download.pytorch.org/whl/nightly/cpuYou can also check out our install documentation for more information, including installing torchtune from source.
To confirm that the package is installed correctly, you can run the following command:
bash
tune --helpAnd should see the following output:
bash
usage: tune [-h] {ls,cp,download,run,validate} ...
Welcome to the torchtune CLI!
options:
-h, --help show this help message and exit
...
Get Started 🚀
To get started with torchtune, see our First Finetune Tutorial. Our End-to-End Workflow Tutorial will show you how to evaluate, quantize, and run inference with a Llama model. The rest of this section will provide a quick overview of these steps with Llama3.1.
Downloading a model
Follow the instructions on the official meta-llama repository to ensure you have access to the official Llama model weights. Once you have confirmed access, you can run the following command to download the weights to your local machine. This will also download the tokenizer model and a responsible use guide.
To download Llama3.1, you can run:
bash
tune download meta-llama/Meta-Llama-3.1-8B-Instruct \
--output-dir /tmp/Meta-Llama-3.1-8B-Instruct \
--ignore-patterns "original/consolidated.00.pth" \
--hf-token <HF_TOKEN> \[!Tip]
Set your environment variableHF_TOKENor pass in--hf-tokento the command in order to validate your access. You can find your token at https://huggingface.co/settings/tokens
Running finetuning recipes
You can finetune Llama3.1 8B with LoRA on a single GPU using the following command:
bash
tune run lora_finetune_single_device --config llama3_1/8B_lora_single_deviceFor distributed training, tune CLI integrates with torchrun.
To run a full finetune of Llama3.1 8B on two GPUs:
bash
tune run --nproc_per_node 2 full_finetune_distributed --config llama3_1/8B_full[!Tip]
Make sure to place any torchrun commands before the recipe specification. Any CLI args after this will override the config and not impact distributed training.
Modify Configs
There are two ways in which you can modify configs:
Config Overrides
You can directly overwrite config fields from the command line:
bash
tune run lora_finetune_single_device \
--config llama2/7B_lora_single_device \
batch_size=8 \
enable_activation_checkpointing=True \
max_steps_per_epoch=128Update a Local Copy
You can also copy the config to your local directory and modify the contents directly:
bash
tune cp llama3_1/8B_full ./my_custom_config.yaml
Copied to ./my_custom_config.yamlThen, you can run your custom recipe by directing the tune run command to your local files:
bash
tune run full_finetune_distributed --config ./my_custom_config.yamlCheck out tune --help for all possible CLI commands and options. For more information on using and updating configs, take a look at our config deep-dive.
Custom Datasets
torchtune supports finetuning on a variety of different datasets, including instruct-style, chat-style, preference datasets, and more. If you want to learn more about how to apply these components to finetune on your own custom dataset, please check out the provided links along with our API docs.
Custom Devices
torchtune supports finetuning on a variety of devices, including NVIDIA GPU, Intel XPU, AMD ROCm, Apple MPS, and Ascend NPU. If you're interested in running recipes on a custom device, such as Intel XPU, follow the steps below.
Step 1: Refer to the Getting Started on Intel GPU guide to configure your environment.
Step 2: Update device information via either CLI override or config changes.
You can directly overwrite config fields from the command line:
bash
tune run lora_finetune_single_device --config llama3_1/8B_lora_single_device device=xpuOr edit your local copy of configuration files and replace device: cuda with device: xpu
Community 🌍
torchtune focuses on integrating with popular tools and libraries from the ecosystem. These are just a few examples, with more under development:
- Hugging Face Hub for accessing model weights
- EleutherAI's LM Eval Harness for evaluating trained models
- Hugging Face Datasets for access to training and evaluation datasets
- PyTorch FSDP2 for distributed training
- torchao for lower precision dtypes and post-training quantization techniques
- Weights & Biases for logging metrics and checkpoints, and tracking training progress
- Comet as another option for logging
- ExecuTorch for on-device inference using finetuned models
- bitsandbytes for low memory optimizers for our single-device recipes
- PEFT for continued finetuning or inference with torchtune models in the Hugging Face ecosystem
Community Contributions
We really value our community and the contributions made by our wonderful users. We'll use this section to call out some of these contributions. If you'd like to help out as well, please see the CONTRIBUTING guide.
- @SalmanMohammadi for adding a comprehensive end-to-end recipe for Reinforcement Learning from Human Feedback (RLHF) finetuning with PPO to torchtune
- @fyabc for adding Qwen2 models, tokenizer, and recipe integration to torchtune
- @solitude-alive for adding the Gemma 2B model to torchtune, including recipe changes, numeric validations of the models and recipe correctness
- @yechenzhi for adding Direct Preference Optimization (DPO) to torchtune, including the recipe and config along with correctness checks
- @Optimox for adding all the Gemma2 variants to torchtune!
Acknowledgements 🙏
The transformer code in this repository is inspired by the original Llama2 code. We also want to give a huge shout-out to EleutherAI, Hugging Face and
Weights & Biases for being wonderful collaborators and for working with us on some of these integrations within torchtune. In addition, we want to acknowledge some other awesome libraries and tools from the ecosystem:
- gpt-fast for performant LLM inference techniques which we've adopted out-of-the-box
- llama recipes for spring-boarding the llama2 community
- bitsandbytes for bringing several memory and performance based techniques to the PyTorch ecosystem
- @winglian and axolotl for early feedback and brainstorming on torchtune's design and feature set.
- lit-gpt for pushing the LLM finetuning community forward.
- HF TRL for making reward modeling more accessible to the PyTorch community.
Citing torchtune 📝
If you find the torchtune library useful, please cite it in your work as below.
bibtex
@software{torchtune,
title = {torchtune: PyTorch's finetuning library},
author = {torchtune maintainers and contributors},
url = {https//github.com/pytorch/torchtune},
license = {BSD-3-Clause},
month = apr,
year = {2024}
}
License
torchtune is released under the BSD 3 license. However you may have other legal obligations that govern your use of other content, such as the terms of service for third-party models.
Overview (`tests/assets/README.md`)
Details on the assets in this folder
`m.model`
Description:
Creation:
Usage:
`tiny_fair_checkpoint.pt`
Description:
Creation:
Usage:
`tiny_llama2_checkpoint.pt`
Description:
Creation:
Usage:
`tiny_state_dict_with_one_key.pt`
Description:
Creation:
Usage:
Overview (`torchtune/dev/README.md`)
Torchtune dev components
This directory houses experimental components.
The purpose of torchtune/dev is to support bleeding-edge APIs
with less stringent requirements on testing, documentation, or stability.
The APIs in this folder are not public and are not guaranteed to adhere to backwards compatibility.
Overview (`torchtune/modules/_export/README.md`)
Export
This directory provides exportable variants of torchtune modules.
Modules in this directory:
- Take the same arguments to
__init__()andforward()as the corresponding reference modules in torchtune. - Give the output as the reference module in torchtune (unless stated otherwise in the docstring).
- Are guaranteed to work out of the box with torch.export.export().
- Should work out of the box with torch.aot_compile().
All modules should be covered by unit tests (under tests/torchtune/modules/_export/) that runs daily and on PRs touching this directory.
These modules are subject to change so proceed with caution.
Contributors: @larryliu0820, @Jack-Khuu, @dvorjackz
Overview (`tests/torchtune/models/llama2/scripts/README.md`)
Verifying Correctness against Reference Implementations
This repository puts a high bar on correctness and testing. To make sure our model and
module implementations are correct, we compare our implementation against reference implementations
where possible. This folder contains scripts used for these comparisons.
Running the scripts
You can run the scripts using the following command as an example.
Each script should print out the value being used in the associated unit tests.
text
python3 -m tests.llm.llama2.scripts.compare_attentionOverview (`tests/torchtune/models/mistral/scripts/README.md`)
Verifying correctness
This directory compares the current implementation of mistral to the reference implementation at https://github.com/mistralai/mistral-src/blob/main/one_file_ref.py. Additionally, torchtune.models.mistral._component_builders.mistral_mlp is compared in tests.torchtune.models.mistral.scripts.compare_feed_forward.py
Since torchtune.models.mistral shares nearly all components with torchtune.models.llama2, please see tests.torchtune.models.llama2.scripts for comparison scripts for individual components.
Running the scripts
You can run the scripts using the following command as an example.
Each script should print out the value being used in the associated unit tests.
text
python3 -m tests.torchtune.models.mistral.scripts.compare_mistralLicense Header (`docs/license_header.txt`)
Copyright (c) Meta Platforms, Inc. and affiliates.
All rights reserved.
This source code is licensed under the BSD-style license found in the
LICENSE file in the root directory of this source tree.
Requirements (`docs/requirements.txt`)
sphinx-gallery>0.11
sphinx==5.0.0
sphinx_design
sphinx_copybutton
sphinx-tabs
matplotlib
-e git+https://github.com/pytorch/pytorch_sphinx_theme.git#egg=pytorch_sphinx_theme
Source Api Ref Config (`docs/source/api_ref_config.rst`)
.. _config:
================ torchtune.config
.. currentmodule:: torchtune.config
.. autosummary::
:toctree: generated/
:nosignatures:
text
instantiate
parse
validate
log_configSource Api Ref Data (`docs/source/api_ref_data.rst`)
.. _data:
============== torchtune.data
.. currentmodule:: torchtune.data
Text templates
Templates for instruct prompts and chat prompts. Includes some specific formatting for difference datasets
and models.
.. autosummary::
:toctree: generated/
:nosignatures:
text
GrammarErrorCorrectionTemplate
SummarizeTemplate
QuestionAnswerTemplate
PromptTemplate
PromptTemplateInterface
ChatMLTemplateTypes
.. autosummary::
:toctree: generated/
:nosignatures:
text
Message
Role.. _message_transforms_ref:
Message transforms
Converts data from common schema and conversation JSON formats into a list of torchtune :class:Message.
.. autosummary::
:toctree: generated/
:nosignatures:
text
InputOutputToMessages
ShareGPTToMessages
OpenAIToMessages
ChosenRejectedToMessages
AlpacaToMessagesCollaters
Collaters used to collect samples into batches and handle any padding.
.. autosummary::
:toctree: generated/
:nosignatures:
text
padded_collate
padded_collate_dpo
padded_collate_packed
padded_collate_sft
padded_collate_tiled_images_and_mask
left_pad_sequenceHelper functions
Miscellaneous helper functions used in modifying data.
.. autosummary::
:toctree: generated/
:nosignatures:
text
validate_messages
truncate
load_image
format_content_with_images
mask_messagesSource Api Ref Datasets (`docs/source/api_ref_datasets.rst`)
.. _datasets:
================== torchtune.datasets
.. currentmodule:: torchtune.datasets
For a detailed general usage guide, please see :ref:datasets_overview.
Text datasets
torchtune supports several widely used text-only datasets to help quickly bootstrap your fine-tuning.
.. autosummary::
:toctree: generated/
:nosignatures:
text
alpaca_dataset
alpaca_cleaned_dataset
grammar_dataset
hh_rlhf_helpful_dataset
samsum_dataset
slimorca_dataset
stack_exchange_paired_dataset
cnn_dailymail_articles_dataset
wikitext_datasetImage + Text datasets
.. autosummary::
:toctree: generated/
:nosignatures:
text
multimodal.llava_instruct_dataset
multimodal.the_cauldron_dataset
multimodal.vqa_dataset.. _dataset_builders:
Generic dataset builders
torchtune also supports generic dataset builders for common formats like chat models and instruct models.
These are especially useful for specifying from a YAML config.
.. autosummary::
:toctree: generated/
:nosignatures:
text
instruct_dataset
chat_dataset
preference_dataset
text_completion_datasetGeneric dataset classes
Class representations for the above dataset builders.
.. autosummary::
:toctree: generated/
:nosignatures:
text
TextCompletionDataset
ConcatDataset
PackedDataset
PreferenceDataset
SFTDatasetSource Api Ref Generation (`docs/source/api_ref_generation.rst`)
.. _generation:
==================== torchtune.generation
.. currentmodule:: torchtune.generation
.. autosummary::
:toctree: generated/
:nosignatures:
text
generate
generate_next_token
sample
get_causal_mask_from_padding_mask
get_position_ids_from_padding_maskSource Api Ref Models (`docs/source/api_ref_models.rst`)
.. _models:
================ torchtune.models
.. currentmodule:: torchtune.models
llama4
Multimodal models from the Llama4 family that support text and image input.
Important: You need to request access on Hugging Face before downloading it.
To download the Llama-4-Maverick-17B-16E-Instruct model:
.. code-block:: bash
text
tune download meta-llama/Llama-4-Scout-17B-16E-Instruct --hf-token <HF_TOKEN>To download the Llama-4-Maverick-17B-128E-Instruct model:
.. code-block:: bash
text
tune download meta-llama/Llama-4-Maverick-17B-128E-Instruct --hf-token <HF_TOKEN>.. autosummary::
:toctree: generated/
:nosignatures:
text
llama4.llama4_scout_17b_16e
llama4.llama4_maverick_17b_128e
llama4.llama4_vision_encoder
llama4.llama4_vision_projection_head
llama4.llama4_decoder
llama4.Llama4VisionEncoder
llama4.Llama4VisionProjectionHead
llama4.Llama4Tokenizer
llama4.Llama4Transformllama3.3
Text-only models from the 3.3 version of Llama3 family <https://llama.meta.com/llama3/>_.
Important: You need to request access on Hugging Face <https://huggingface.co/meta-llama/Llama-3.3-70B-Instruct>__ before downloading it.
To download the Llama-3.3-70B-Instruct model:
.. code-block:: bash
text
tune download meta-llama/Llama-3.3-70B-Instruct --ignore-patterns "original/consolidated.00.pth" --hf-token <HF_TOKEN>.. autosummary::
:toctree: generated/
:nosignatures:
text
llama3_3.llama3_3_70b
llama3_3.lora_llama3_3_70b
llama3_3.qlora_llama3_3_70b.. note::
text
The Llama3.3 tokenizer reuses the :class:`~torchtune.models.llama3.llama3_tokenizer` class.llama3.2
Text-only models from the 3.2 version of Llama3 family <https://llama.meta.com/llama3/>_.
Important: You need to request access on Hugging Face <https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct>__ before downloading it.
To download the Llama-3.2-1B-Instruct model:
.. code-block:: bash
text
tune download meta-llama/Llama-3.2-1B-Instruct --output-dir /tmp/Llama-3.2-1B-Instruct --ignore-patterns "original/consolidated.00.pth" --hf-token <HF_TOKEN>To download the Llama-3.2-3B-Instruct model:
.. code-block:: bash
text
tune download meta-llama/Llama-3.2-3B-Instruct --output-dir /tmp/Llama-3.2-3B-Instruct --ignore-patterns "original/consolidated*" --hf-token <HF_TOKEN>.. autosummary::
:toctree: generated/
:nosignatures:
text
llama3_2.llama3_2_1b
llama3_2.llama3_2_3b
llama3_2.lora_llama3_2_1b
llama3_2.lora_llama3_2_3b
llama3_2.qlora_llama3_2_1b
llama3_2.qlora_llama3_2_3b.. note::
text
The Llama3.2 tokenizer reuses the :class:`~torchtune.models.llama3.llama3_tokenizer` class.llama3.2 vision
Vision-Language Models from the 3.2 version of Llama3 family <https://llama.meta.com/llama3/>_.
Important: You need to request access on Hugging Face <https://huggingface.co/meta-llama/Llama-3.2-11B-Vision-Instruct>__ before downloading it.
To download the Llama-3.2-11B-Instruct model:
.. code-block:: bash
text
tune download meta-llama/Llama-3.2-11B-Vision-Instruct --output-dir /tmp/Llama-3.2-11B-Vision-Instruct --hf-token <HF_TOKEN>.. autosummary::
:toctree: generated/
:nosignatures:
text
llama3_2_vision.llama3_2_vision_11b
llama3_2_vision.llama3_2_vision_transform
llama3_2_vision.lora_llama3_2_vision_11b
llama3_2_vision.qlora_llama3_2_vision_11b
llama3_2_vision.llama3_2_vision_decoder
llama3_2_vision.llama3_2_vision_encoder
llama3_2_vision.lora_llama3_2_vision_decoder
llama3_2_vision.lora_llama3_2_vision_encoder
llama3_2_vision.Llama3VisionEncoder
llama3_2_vision.Llama3VisionProjectionHead
llama3_2_vision.Llama3VisionTransform.. note::
text
The Llama3.2 tokenizer reuses the :class:`~torchtune.models.llama3.llama3_tokenizer` class.llama3 & llama3.1
Models 3 and 3.1 from the Llama3 family <https://llama.meta.com/llama3/>_.
Important: You need to request access on Hugging Face <https://huggingface.co/meta-llama/Meta-Llama-3.1-8B-Instruct>__ before downloading it.
To download the Llama3.1-8B-Instruct model:
.. code-block:: bash
text
tune download meta-llama/Meta-Llama-3.1-8B-Instruct --output-dir /tmp/Meta-Llama-3.1-8B-Instruct --ignore-patterns "original/consolidated.00.pth" --hf-token <HF_TOKEN>To download the Llama3.1-70B-Instruct model:
.. code-block:: bash
text
tune download meta-llama/Meta-Llama-3.1-70B-Instruct --output-dir /tmp/Meta-Llama-3.1-70B-Instruct --ignore-patterns "original/consolidated*" --hf-token <HF_TOKEN>To download the Llama3.1-405B-Instruct model:
.. code-block:: bash
text
tune download meta-llama/Meta-Llama-3.1-405B-Instruct --ignore-patterns "original/consolidated*" --hf-token <HF_TOKEN>To download the Llama3 weights of the above models, you can instead download from Meta-Llama-3-8B-Instruct andMeta-Llama-3-70B-Instruct, and remove the ignore patterns flag.
.. autosummary::
:toctree: generated/
:nosignatures:
text
llama3.llama3
llama3.lora_llama3
llama3.llama3_8b
llama3.lora_llama3_8b
llama3.qlora_llama3_8b
llama3.llama3_70b
llama3.lora_llama3_70b
llama3.qlora_llama3_70b
llama3.llama3_tokenizer
|
llama3_1.llama3_1
llama3_1.lora_llama3_1
llama3_1.llama3_1_8b
llama3_1.lora_llama3_1_8b
llama3_1.qlora_llama3_1_8b
llama3_1.llama3_1_70b
llama3_1.lora_llama3_1_70b
llama3_1.qlora_llama3_1_70b
llama3_1.llama3_1_405b
llama3_1.lora_llama3_1_405b
llama3_1.qlora_llama3_1_405b.. note::
text
The Llama3.1 tokenizer reuses the `llama3.llama3_tokenizer` builder class.llama2
All models from the Llama2 family <https://llama.meta.com/llama2/>_.
Important: You need to request access on Hugging Face <https://huggingface.co/meta-llama/Llama-2-7b-hf>__ before downloading it.
To download the Llama2-7B model:
.. code-block:: bash
tune download meta-llama/Llama-2-7b-hf --output-dir /tmp/Llama-2-7b-hf --hf-token
To download the Llama2-13B model:
.. code-block:: bash
text
tune download meta-llama/Llama-2-13b-hf --output-dir /tmp/Llama-2-13b-hf --hf-token <HF_TOKEN>To download the Llama2-70B model:
.. code-block:: bash
text
tune download meta-llama/Llama-2-70b-hf --output-dir /tmp/Llama-2-70b-hf --hf-token <HF_TOKEN>.. autosummary::
:toctree: generated/
:nosignatures:
text
llama2.llama2
llama2.lora_llama2
llama2.llama2_7b
llama2.lora_llama2_7b
llama2.qlora_llama2_7b
llama2.llama2_13b
llama2.lora_llama2_13b
llama2.qlora_llama2_13b
llama2.llama2_70b
llama2.lora_llama2_70b
llama2.qlora_llama2_70b
llama2.llama2_tokenizer
llama2.llama2_reward_7b
llama2.lora_llama2_reward_7b
llama2.qlora_llama2_reward_7b
llama2.Llama2ChatTemplateqwen-2.5
Models of size 0.5B, 1.5B, 3B, 7B, 14B, 32B, 72B from the Qwen2.5 family <https://huggingface.co/collections/Qwen/qwen25-66e81a666513e518adb90d9e>_.
To download the Qwen2.5 1.5B model, for example:
.. code-block:: bash
text
tune download Qwen/Qwen2.5-1.5B-Instruct --output-dir /tmp/Qwen2_5-1_5B-Instruct.. autosummary::
:toctree: generated/
:nosignatures:
text
qwen2_5.qwen2_5_0_5b
qwen2_5.lora_qwen2_5_0_5b
qwen2_5.qwen2_5_1_5b_base
qwen2_5.qwen2_5_1_5b_instruct
qwen2_5.lora_qwen2_5_1_5b_base
qwen2_5.lora_qwen2_5_1_5b_instruct
qwen2_5.qwen2_5_3b
qwen2_5.lora_qwen2_5_3b
qwen2_5.qwen2_5_7b_base
qwen2_5.qwen2_5_7b_instruct
qwen2_5.lora_qwen2_5_7b_base
qwen2_5.lora_qwen2_5_7b_instruct
qwen2_5.qwen2_5_14b_base
qwen2_5.qwen2_5_14b_instruct
qwen2_5.lora_qwen2_5_14b_base
qwen2_5.lora_qwen2_5_14b_instruct
qwen2_5.qwen2_5_32b_base
qwen2_5.qwen2_5_32b_instruct
qwen2_5.lora_qwen2_5_32b_base
qwen2_5.lora_qwen2_5_32b_instruct
qwen2_5.qwen2_5_72b_base
qwen2_5.qwen2_5_72b_instruct
qwen2_5.lora_qwen2_5_72b_base
qwen2_5.lora_qwen2_5_72b_instruct
qwen2_5.qwen2_5_tokenizerqwen-2
Models of size 0.5B, 1.5B, and 7B from the Qwen2 family <https://huggingface.co/collections/Qwen/qwen2-6659360b33528ced941e557f>_.
To download the Qwen2 1.5B model, for example:
.. code-block:: bash
text
tune download Qwen/Qwen2-1.5B-Instruct --output-dir /tmp/Qwen2-1.5B-Instruct.. autosummary::
:toctree: generated/
:nosignatures:
text
qwen2.qwen2
qwen2.lora_qwen2
qwen2.qwen2_0_5b
qwen2.lora_qwen2_0_5b
qwen2.qwen2_1_5b
qwen2.lora_qwen2_1_5b
qwen2.qwen2_7b
qwen2.lora_qwen2_7b
qwen2.qwen2_tokenizerphi-4
Models from the Phi-4 family <https://arxiv.org/abs/2412.08905>_.
To download the Phi-4 instruct model:
.. code-block:: bash
text
tune download microsoft/phi-4 --hf-token <HF_TOKEN>.. autosummary::
:toctree: generated/
:nosignatures:
text
phi4.phi4_14b
phi4.lora_phi4_14b
phi4.qlora_phi4_14b
phi4.phi4_tokenizerphi-3
Models from the Phi-3 mini family <https://news.microsoft.com/source/features/ai/the-phi-3-small-language-models-with-big-potential/>_.
To download the Phi-3 Mini 4k instruct model:
.. code-block:: bash
text
tune download microsoft/Phi-3-mini-4k-instruct --output-dir /tmp/Phi-3-mini-4k-instruct --hf-token <HF_TOKEN>.. autosummary::
:toctree: generated/
:nosignatures:
text
phi3.phi3
phi3.lora_phi3
phi3.phi3_mini
phi3.lora_phi3_mini
phi3.qlora_phi3_mini
phi3.phi3_mini_tokenizermistral
All models from Mistral AI family <https://mistral.ai/technology/#models>_.
Important: You need to request access on Hugging Face <https://huggingface.co/mistralai/Mistral-7B-v0.1>__ to download this model.
To download the Mistral 7B v0.1 model:
.. code-block:: bash
text
tune download mistralai/Mistral-7B-v0.1 --output-dir /tmp/Mistral-7B-v0.1 --ignore-patterns "*.safetensors" --hf-token <HF_TOKEN>.. autosummary::
:toctree: generated/
:nosignatures:
text
mistral.mistral
mistral.lora_mistral
mistral.mistral_classifier
mistral.lora_mistral_classifier
mistral.mistral_7b
mistral.lora_mistral_7b
mistral.qlora_mistral_7b
mistral.mistral_reward_7b
mistral.lora_mistral_reward_7b
mistral.qlora_mistral_reward_7b
mistral.mistral_tokenizer
mistral.MistralChatTemplategemma
Models of size 2B and 7B from the Gemma family <https://blog.google/technology/developers/gemma-open-models/>_.
Important: You need to request access on Hugging Face <https://huggingface.co/google/gemma-2b>__ to use this model.
To download the Gemma 2B model (not Gemma2):
.. code-block:: bash
text
tune download google/gemma-2b --ignore-patterns "gemma-2b.gguf" --hf-token <HF_TOKEN>To download the Gemma 7B model:
.. code-block:: bash
text
tune download google/gemma-7b --ignore-patterns "gemma-7b.gguf" --hf-token <HF_TOKEN>.. autosummary::
:toctree: generated/
:nosignatures:
text
gemma.gemma
gemma.lora_gemma
gemma.gemma_2b
gemma.lora_gemma_2b
gemma.qlora_gemma_2b
gemma.gemma_7b
gemma.lora_gemma_7b
gemma.qlora_gemma_7b
gemma.gemma_tokenizergemma2
Models of size 2B, 9B, 27B from the Gemma family <https://blog.google/technology/developers/gemma-open-models/>_.
Important: You need to request access on Hugging Face <https://huggingface.co/google/gemma-2-2b>__ to use this model.
To download the Gemma2 2B, 9B, 27B models :
.. code-block:: bash
text
tune download google/gemma-2-<MODEL_SIZE>b --ignore-patterns "gemma-2-<MODEL_SIZE>b.gguf" --hf-token <HF_TOKEN>.. autosummary::
:toctree: generated/
:nosignatures:
text
gemma2.gemma2
gemma2.lora_gemma2
gemma2.gemma2_2b
gemma2.lora_gemma2_2b
gemma2.qlora_gemma2_2b
gemma2.gemma2_9b
gemma2.lora_gemma2_9b
gemma2.qlora_gemma2_9b
gemma2.gemma2_27b
gemma2.lora_gemma2_27b
gemma2.qlora_gemma2_27b
gemma.gemma_tokenizerclip
Vision components to support multimodality using CLIP encoder <https://arxiv.org/abs/2103.00020>_.
.. autosummary::
:toctree: generated/
:nosignatures:
text
clip.clip_vision_encoder
clip.TokenPositionalEmbedding
clip.TiledTokenPositionalEmbedding
clip.TilePositionalEmbeddingSource Api Ref Modules (`docs/source/api_ref_modules.rst`)
================= torchtune.modules
.. currentmodule:: torchtune.modules
Modeling Components and Building Blocks
.. autosummary::
:toctree: generated/
:nosignatures:
text
MultiHeadAttention
FeedForward
KVCache
RotaryPositionalEmbeddings
RMSNorm
Fp32LayerNorm
TanhGate
TiedLinear
TransformerSelfAttentionLayer
TransformerCrossAttentionLayer
TransformerDecoder
VisionTransformer
LayerDropout
prepare_layer_dropout
classifier_modelLosses
.. autosummary::
:toctree: generated/
:nosignatures:
text
loss.LinearCrossEntropyLoss
loss.ForwardKLLoss
loss.ForwardKLWithChunkedOutputLoss
loss.ReverseKLWithChunkedOutputLoss
loss.SymmetricKLWithChunkedOutputLossOptimizers
.. autosummary::
:toctree: generated/
:nosignatures:
text
optim.OptimizerInBackwardBase Tokenizers
Base tokenizers are tokenizer models that perform the direct encoding of text
into token IDs and decoding of token IDs into text. These are typically byte pair encodings <https://en.wikipedia.org/wiki/Byte_pair_encoding>_ that underlie the
model specific tokenizers.
.. autosummary::
:toctree: generated/
:nosignatures:
text
transforms.tokenizers.SentencePieceBaseTokenizer
transforms.tokenizers.TikTokenBaseTokenizer
transforms.tokenizers.HuggingFaceBaseTokenizer
transforms.tokenizers.ModelTokenizer
transforms.tokenizers.BaseTokenizerTokenizer Utilities
These are helper methods that can be used by any tokenizer.
.. autosummary::
:toctree: generated/
:nosignatures:
text
transforms.tokenizers.tokenize_messages_no_special_tokens
transforms.tokenizers.parse_hf_tokenizer_jsonPEFT Components
.. autosummary::
:toctree: generated/
:nosignatures:
text
peft.LoRALinear
peft.DoRALinear
peft.AdapterModule
peft.get_adapter_params
peft.set_trainable_params
peft.get_adapter_state_dict
peft.validate_missing_and_unexpected_for_lora
peft.disable_adapterFusion Components
Components for building models that are a fusion of two+ pre-trained models.
.. autosummary::
:toctree: generated/
:nosignatures:
text
model_fusion.DeepFusionModel
model_fusion.FusionLayer
model_fusion.FusionEmbedding
model_fusion.register_fusion_module
model_fusion.get_fusion_paramsModule Utilities
These are utilities that are common to and can be used by all modules.
.. autosummary::
:toctree: generated/
:nosignatures:
common_utils.reparametrize_as_dtype_state_dict_post_hook
common_utils.local_kv_cache
common_utils.disable_kv_cache
common_utils.delete_kv_caches
embedding_utils.resize_token_embeddings
Vision Transforms
Functions used for preprocessing images.
.. autosummary::
:toctree: generated/
:nosignatures:
text
transforms.Transform
transforms.VisionCrossAttentionMaskSource Api Ref Rlhf (`docs/source/api_ref_rlhf.rst`)
=============== torchtune.rlhf
.. currentmodule:: torchtune.rlhf
Components and losses for RLHF algorithms like PPO and DPO.
.. autosummary::
:toctree: generated/
:nosignatures:
text
estimate_advantages
get_rewards_ppo
truncate_sequence_at_first_stop_token
loss.PPOLoss
loss.DPOLoss
loss.RSOLossSource Api Ref Training (`docs/source/api_ref_training.rst`)
================== torchtune.training
.. currentmodule:: torchtune.training
.. _checkpointing_label:
Checkpointing
torchtune offers checkpointers to allow seamless transitioning between checkpoint formats for training and interoperability with the rest of the ecosystem. For a comprehensive overview of
checkpointing, please see the :ref:checkpointing deep-dive .
.. autosummary::
:toctree: generated/
:nosignatures:
text
FullModelHFCheckpointer
FullModelMetaCheckpointer
FullModelTorchTuneCheckpointer
ModelType
FormattedCheckpointFiles
update_state_dict_for_classifier.. _mp_label:
Reduced Precision
Utilities for working in a reduced precision setting.
.. autosummary::
:toctree: generated/
:nosignatures:
text
get_dtype
set_default_dtype
validate_expected_param_dtype
get_quantizer_mode.. _dist_label:
Distributed
Utilities for enabling and working with distributed training.
.. autosummary::
:toctree: generated/
:nosignatures:
text
init_distributed
is_distributed
gather_cpu_state_dict
get_distributed_backend.. _ac_label:
Memory Management
Utilities to reduce memory consumption during training.
.. autosummary::
:toctree: generated/
:nosignatures:
text
apply_selective_activation_checkpointing
set_activation_checkpointing
OptimizerInBackwardWrapper
create_optim_in_bwd_wrapper
register_optim_in_bwd_hooks.. _lr_scheduler_label:
Schedulers
Utilities to control lr during the training process.
.. autosummary::
:toctree: generated/
:nosignatures:
text
get_cosine_schedule_with_warmup
get_lr.. _metric_logging_label:
Metric Logging
Various logging utilities.
.. autosummary::
:toctree: generated/
:nosignatures:
text
metric_logging.CometLogger
metric_logging.WandBLogger
metric_logging.TensorBoardLogger
metric_logging.StdoutLogger
metric_logging.DiskLogger
metric_logging.MLFlowLogger.. _perf_profiling_label:
Performance and Profiling
torchtune provides utilities to profile and debug the memory and performance
of your finetuning job.
.. autosummary::
:toctree: generated/
:nosignatures:
text
get_memory_stats
log_memory_stats
setup_torch_profilerMiscellaneous
.. autosummary::
:toctree: generated/
:nosignatures:
text
get_unmasked_sequence_lengths
disable_dropout
set_seedSource Api Ref Utilities (`docs/source/api_ref_utilities.rst`)
=============== torchtune.utils
.. currentmodule:: torchtune.utils
.. _gen_label:
Miscellaneous
.. autosummary::
:toctree: generated/
:nosignatures:
text
batch_to_device
get_device
get_logger
torch_version_ge
get_world_size_and_rankSource Index (`docs/source/index.rst`)
Welcome to the torchtune Documentation
torchtune is a Native-PyTorch library for LLM fine-tuning.
Getting Started
text
Topics in this section will help you get started with torchtune.
.. grid:: 3
.. grid-item-card:: :octicon:`file-code;1em`
What is torchtune?
:img-top: _static/img/card-background.svg
:link: overview.html
:link-type: url
A gentle introduction to torchtune and how you can
use the library in your projects.
.. grid-item-card:: :octicon:`file-code;1em`
Installation instructions
:img-top: _static/img/card-background.svg
:link: install.html
:link-type: url
A step-by-step tutorial on how to install torchtune.
.. grid-item-card:: :octicon:`file-code;1em`
Finetune your first model
:img-top: _static/img/card-background.svg
:link: tutorials/first_finetune_tutorial.html
:link-type: url
Follow a simple tutorial to finetune Llama2 with torchtune.
Tutorials
~~~~~~~~~
Ready to experiment? Check out some of the interactive
torchtune tutorials.
.. customcardstart::
.. customcarditem::
:header: Llama3 in torchtune
:card_description:
:image: _static/img/generic-pytorch-logo.png
:link: tutorials/llama3.html
:tags: finetuning,llama3
.. customcarditem::
:header: Finetuning with LoRA in torchtune
:card_description: Parameter-efficient finetuning of Llama2 using LoRA
:image: _static/img/generic-pytorch-logo.png
:link: tutorials/lora_finetune.html
:tags: finetuning,llama2,lora
.. customcarditem::
:header: Understanding QLoRA in torchtune
:card_description: Using QLoRA to quantize base model weights and maximize memory savings
:image: _static/img/generic-pytorch-logo.png
:link: tutorials/qlora_finetune.html
:tags: finetuning,llama2,qlora
.. customcarditem::
:header: Finetuning with QAT in torchtune
:card_description: Finetuning of Llama3 using QAT
:image: _static/img/generic-pytorch-logo.png
:link: tutorials/qat_finetune.html
:tags: finetuning,llama3,qat,quantization,evals
.. customcarditem::
:header: End-to-End Workflow with torchtune
:card_description: Train, Evaluate, Quantize and then Generate with your LLM.
:image: _static/img/generic-pytorch-logo.png
:link: tutorials/e2e_flow.html
:tags: finetuning,quantization,inference,evals,llama2
.. customcarditem::
:header: Distilling Llama3 8B into 1B
:card_description: Distilling Llama3 8B into 1B using Knowledge Distillation
:image: _static/img/generic-pytorch-logo.png
:link: tutorials/llama_kd_tutorial.html
:tags: finetuning,llama3,kd
.. customcarditem::
:header: Multi-node training w/ Llama3.3 70B
:card_description: Fine-tuning a large model on 2+ nodes
:image: _static/img/generic-pytorch-logo.png
:link: tutorials/multinode.html
:tags: multinode,llama3,slurm
.. customcardend::
.. ----------------------------------------------------------------------
.. Below is the toctree i.e. it defines the content of the left sidebar.
.. Each of the entry below corresponds to a file.rst in docs/source/.
.. ----------------------------------------------------------------------
.. toctree::
:glob:
:maxdepth: 1
:caption: Getting Started
:hidden:
overview
install
tutorials/first_finetune_tutorial
tune_cli
.. toctree::
:glob:
:maxdepth: 1
:caption: Finetuning Recipes
:hidden:
recipes/recipes_overview
recipes/lora_finetune_single_device
recipes/qat_distributed
recipes/dpo
.. toctree::
:glob:
:maxdepth: 1
:caption: Basics
:hidden:
basics/datasets_overview
basics/custom_datasets
basics/chat_datasets
basics/instruct_datasets
basics/multimodal_datasets
basics/preference_datasets
basics/text_completion_datasets
basics/model_transforms
basics/messages
basics/message_transforms
basics/tokenizers
basics/prompt_templates
basics/packing
basics/custom_components
.. toctree::
:glob:
:maxdepth: 1
:caption: Tutorials
:hidden:
tutorials/llama3
tutorials/chat
tutorials/lora_finetune
tutorials/qlora_finetune
tutorials/qat_finetune
tutorials/e2e_flow
tutorials/llama_kd_tutorial
tutorials/memory_optimizations
tutorials/multinode
.. toctree::
:glob:
:maxdepth: 1
:caption: Deep-Dives
:hidden:
deep_dives/checkpointer
deep_dives/configs
deep_dives/recipe_deepdive
deep_dives/comet_logging
deep_dives/wandb_logging
.. toctree::
:glob:
:maxdepth: 1
:caption: API Reference
:hidden:
api_ref_config
api_ref_data
api_ref_datasets
api_ref_generation
api_ref_models
api_ref_modules
api_ref_rlhf
api_ref_training
api_ref_utilities
---
### Source Install (`docs/source/install.rst`)
.. _install_label:
====================
Install Instructions
====================
Pre-requisites
--------------
torchtune requires PyTorch, so please install for your proper host and environment
using the `"Start Locally" <https://pytorch.org/get-started/locally/>`_ page. You should also install
torchvision (for multimodal LLMs) and torchao (for quantization APIs). You can install either stable or
nightly versions with the following commands:
.. code-block:: bash
# Install stable version of PyTorch libraries using pip
pip install torch torchvision torchao
# Or nightly install for latest features
pip install --pre torch torchvision torchao --index-url https://download.pytorch.org/whl/nightly/cu126 # full options are cpu/cu118/cu124/cu126/xpu/rocm6.2/rocm6.3/rocm6.4
Install via PyPI
----------------
The latest stable version of torchtune is hosted on PyPI and can be downloaded
with the following command:
.. code-block:: bash
pip install torchtune
To confirm that the package is installed correctly, you can run the following command:
.. code-block:: bash
tune
And should see the following output:
::
usage: tune [-h] {download,ls,cp,run,validate} ...
Welcome to the torchtune CLI!
options:
-h, --help show this help message and exit
...
|
Install via ``git clone``
-------------------------
If you want the latest and greatest features from torchtune or if you want to `become a contributor <https://github.com/pytorch/torchtune/blob/main/CONTRIBUTING.md>`_,
you can also install the package locally with the following command.
.. code-block:: bash
git clone https://github.com/pytorch/torchtune.git
cd torchtune
pip install -e .
# or for a developer installation
pip install -e ".[dev]"
|
Install nightly build
---------------------
torchtune gets built every evening with the latest commits to ``main`` branch. If you want the latest updates
to the package *without* installing via ``git clone``, you can install with the following command:
.. code-block:: bash
pip install --pre torchtune --extra-index-url https://download.pytorch.org/whl/nightly/cpu --no-cache-dir
.. note::
``--no-cache-dir`` will direct ``pip`` to not look for a cached version of torchtune, thereby overwriting
your existing torchtune installation.
If you already have PyTorch installed, torchtune will default to using that version. However, if you want to
use the nightly version of PyTorch, you can append the ``--force-reinstall`` option to the above command. If you
opt for this install method, you will likely need to change the "cpu" suffix in the index url to match your hardware backend.
For example, if you are running CUDA 12, your index url would be "https://download.pytorch.org/whl/nightly/cu126", and if
you are running on Intel XPU, your index url would be "https://download.pytorch.org/whl/nightly/xpu".
---
### Source Overview (`docs/source/overview.rst`)
.. _overview_label:
==================
torchtune Overview
==================
On this page, we'll walk through an overview of torchtune, including features, key concepts and additional pointers.
What is torchtune?
------------------
torchtune is a PyTorch library for easily authoring, fine-tuning and experimenting with LLMs. The library emphasizes 4 key aspects:
- **Simplicity and Extensibility**. Native-PyTorch, componentized design and easy-to-reuse abstractions
- **Correctness**. High bar on proving the correctness of components and recipes
- **Stability**. PyTorch just works. So should torchtune
- **Democratizing LLM fine-tuning**. Works out-of-the-box on different hardware
torchtune provides:
- Modular native-PyTorch implementations of popular LLMs
- Interoperability with popular model zoos through checkpoint-conversion utilities
- Training recipes for a variety of fine-tuning techniques
- Integration with `Hugging Face Datasets <https://huggingface.co/docs/datasets/en/index>`_ for training and `EleutherAI's Eval Harness <https://github.com/EleutherAI/lm-evaluation-harness>`_ for evaluation
- Support for distributed training using `FSDP2 <https://github.com/pytorch/torchtitan/blob/main/docs/fsdp.md>`_
- YAML configs for easily configuring training runs
Excited? To get started, checkout some of our tutorials, including:
- Our :ref:`quickstart guide <finetune_llama_label>` to finetune your first LLM using torchtune.
- Our :ref:`LoRA tutorial <lora_finetune_label>` to learn about parameter-efficient finetuning with torchtune.
- Our :ref:`QLoRA tutorial <qlora_finetune_label>` to attain maximal memory efficiency with torchtune.
You can check out our :ref:`recipes overview<recipes_overview_label>` to see all the fine-tuning techniques we support.
Key Concepts
------------
As you go through the tutorials and code, there are two concepts which will help you better understand and use torchtune.
**Configs.** YAML files which help you configure training settings (dataset, model, checkpoint) and
hyperparameters (batch size, learning rate) without modifying code.
See the ":ref:`All About Configs<config_tutorial_label>`" deep-dive for more information.
**Recipes.** Recipes can be thought of
as targeted end-to-end pipelines for training and optionally evaluating LLMs.
Each recipe implements a training method (eg: full fine-tuning) with a set of meaningful
features (eg: FSDP2 + Activation Checkpointing + Gradient Accumulation + Reduced Precision training)
applied to a given model family (eg: Llama3.1). See the ":ref:`What Are Recipes?<recipe_deepdive>`" deep-dive for more information.
.. _design_principles_label:
Design Principles
-----------------
torchtune embodies `PyTorch’s design philosophy <https://pytorch.org/docs/stable/community/design.html>`_, especially "usability over everything else".
**Native PyTorch**
torchtune is a native-PyTorch library. While we provide integrations with the surrounding ecosystem (eg: `Hugging Face Datasets <https://huggingface.co/docs/datasets/en/index>`_,
`EleutherAI's Eval Harness <https://github.com/EleutherAI/lm-evaluation-harness>`_), all of the core functionality is written in PyTorch.
**Simplicity and Extensibility**
torchtune is designed to be easy to understand, use and extend.
- Composition over implementation inheritance - layers of inheritance for code re-use makes the code hard to read and extend
- No training frameworks - explicitly outlining the training logic makes it easy to extend for custom use cases
- Code duplication is prefered over unecessary abstractions
- Modular building blocks over monolithic components
**Correctness**
torchtune provides well-tested components with a high bar on correctness. The library will never be the first to provide a feature, but available features will be thoroughly tested. We provide
- Extensive unit tests to ensure component-level numerical parity with reference implementations
- Checkpoint tests to ensure model-level numerical parity with reference implementations
- Integration tests to ensure recipe-level performance parity with reference implementations on standard benchmarks
---
### Source Tune Cli (`docs/source/tune_cli.rst`)
.. _cli_label:
=============
torchtune CLI
=============
This page is the documentation for using the torchtune CLI - a convenient way to
download models, find and copy relevant recipes/configs, and run recipes. It is automatically
available when you install torchtune.
Getting started
---------------
The ``--help`` option will show all the possible commands available through the torchtune CLI,
with a short description of each.
.. code-block:: bash
$ tune --help
usage: tune [-h] {download,ls,cp,run,validate,cat} ...
Welcome to the torchtune CLI!
options:
-h, --help show this help message and exit
subcommands:
{download,ls,cp,run,validate,cat}
download Download a model from the Hugging Face Hub.
ls List all built-in recipes and configs
...
The ``--help`` option is convenient for getting more details about any command. You can use it anytime to list all
available options and their details. For example, ``tune download --help`` provides more information on how
to download files using the CLI.
.. _tune_download_label:
Download a model
----------------
The ``tune download `` command downloads any model from the Hugging Face or Kaggle Model Hub.
.. list-table::
:widths: 30 60
* - \--output-dir
- Directory in which to save the model. Note: this is option not yet supported when `--source` is set to `kaggle`.
* - \--output-dir-use-symlinks
- To be used with `output-dir`. If set to 'auto', the cache directory will be used and the file will be either duplicated or symlinked to the local directory depending on its size. It set to `True`, a symlink will be created, no matter the file size. If set to `False`, the file will either be duplicated from cache (if already exists) or downloaded from the Hub and not cached.
* - \--hf-token
- Hugging Face API token. Needed for gated models like Llama.
* - \--ignore-patterns
- If provided, files matching any of the patterns are not downloaded. Defaults to ignoring safetensors files to avoid downloading duplicate weights.
* - \--source {huggingface,kaggle}
- If provided, downloads model weights from the provided on the designated source hub.
* - \--kaggle-username
- Kaggle username for authentication. Needed for private models or gated models like Llama2.
* - \--kaggle-api-key
- Kaggle API key. Needed for private models or gated models like Llama2. You can find your API key at https://kaggle.com/settings.
.. code-block:: bash
$ tune download meta-llama/Meta-Llama-3-8B-Instruct
Successfully downloaded model repo and wrote to the following locations:
./model/config.json
./model/README.md
./model/model-00001-of-00002.bin
...
.. code-block:: bash
$ tune download metaresearch/llama-3.2/pytorch/1b --source kaggle
Successfully downloaded model repo and wrote to the following locations:
/tmp/llama-3.2/pytorch/1b/tokenizer.model
/tmp/llama-3.2/pytorch/1b/params.json
/tmp/llama-3.2/pytorch/1b/consolidated.00.pth
**Download a gated model**
A lot of recent large pretrained models released from organizations like Meta or MistralAI require you to agree
to the usage terms and conditions before you are allowed to download their model. If this is the case, you can specify
a Hugging Face access token.
You can find the access token `here <https://huggingface.co/docs/hub/en/security-tokens>`_.
.. code-block:: bash
$ tune download meta-llama/Meta-Llama-3-8B-Instruct --hf-token <TOKEN>
Successfully downloaded model repo and wrote to the following locations:
./model/config.json
./model/README.md
./model/model-00001-of-00002.bin
...
.. note::
If you'd prefer, you can also use ``huggingface-cli login`` to permanently login to the Hugging Face Hub on your machine.
The ``tune download`` command will pull the access token from your environment.
**Specify model files you don't want to download**
Some checkpoint directories can be very large and it can eat up a lot of bandwith and local storage to download the all of the files every time, even if you might
not need a lot of them. This is especially common when the same checkpoint exists in different formats. You can specify patterns to ignore to prevent downloading files
with matching names. By default we ignore safetensor files, but if you want to include all files you can pass in an empty string.
.. code-block:: bash
$ tune download meta-llama/Meta-Llama-3-8B-Instruct --hf-token <TOKEN> --ignore-patterns None
Successfully downloaded model repo and wrote to the following locations:
./model/config.json
./model/README.md
./model/model-00001-of-00030.safetensors
...
.. note::
Just because a model can be downloaded does not mean that it will work OOTB with torchtune's
built-in recipes or configs. For a list of supported model families and architectures, see :ref:`models<models>`.
.. _tune_ls_label:
List built-in recipes and configs
---------------------------------
The ``tune ls`` command lists out all the built-in recipes and configs within torchtune.
.. code-block:: bash
$ tune ls
RECIPE CONFIG
full_finetune_single_device llama2/7B_full_low_memory
llama3/8B_full_single_device
mistral/7B_full_low_memory
phi3/mini_full_low_memory
full_finetune_distributed llama2/7B_full
llama2/13B_full
llama3/8B_full
llama3/70B_full
...
.. _tune_cp_cli_label:
Copy a built-in recipe or config
--------------------------------
The ``tune cp <recipe|config> `` command copies built-in recipes and configs to a provided location. This allows you to make a local copy of a library
recipe or config to edit directly for yourself. See :ref:`here <tune_cp_label>` for an example of how to use this command.
.. list-table::
:widths: 30 60
* - \-n, \--no-clobber
- Do not overwrite destination if it already exists
* - \--make-parents
- Create parent directories for destination if they do not exist. If not set to True, will error if parent directories do not exist
.. code-block:: bash
$ tune cp lora_finetune_distributed .
Copied file to ./lora_finetune_distributed.py
Run a recipe
------------
The ``tune run <recipe> --config <config>`` is a wrapper around `torchrun <https://pytorch.org/docs/stable/elastic/run.html>`_. ``tune run`` allows you to specify
a built-in recipe or config by name, or by path to use your local recipes/configs.
To run a tune recipe
.. code-block:: bash
tune run lora_finetune_single_device --config llama3/8B_lora_single_device
**Specifying distributed (torchrun) arguments**
``tune run`` supports launching distributed runs by passing through arguments preceding the recipe directly to torchrun. This follows the pattern used by torchrun
of specifying distributed and host machine flags before the script (recipe). For a full list of available flags for distributed setup, see the `torchrun docs <https://pytorch.org/docs/stable/elastic/run.html>`_.
Some common flags:
.. list-table::
:widths: 30 60
* - \--nproc-per-node
- Number of workers per node; supported values: [auto, cpu, gpu, int].
* - \--nnodes
- Number of nodes, or the range of nodes in form <minimum_nodes>:<maximum_nodes>.
* - \--max-restarts
- Maximum number of worker group restarts before failing.
* - \--rdzv-backend
- Rendezvous backend.
* - \--rdzv-endpoint
- Rendezvous backend endpoint; usually in form <host>:.
.. code-block:: bash
tune run --nnodes=1 --nproc-per-node=4 lora_finetune_distributed --config llama3/8B_lora
.. note::
If no arguments are provided before the recipe, tune will bypass torchrun and launch directly with ``python``. This can simplify running and debugging recipes
when distributed isn't needed. If you want to launch with torchrun, but use only a single device, you can specify ``tune run --nnodes=1 --nproc-per-node=1 <recipe> --config <config>``.
**Running a custom (local) recipe and config**
To use ``tune run`` with your own local recipes and configs, simply pass in a file path instead of a name to the run command. You can mix and match a custom recipe with a
torchtune config or vice versa or you can use both custom configs and recipes.
.. code-block:: bash
tune run my/fancy_lora.py --config my/configs/8B_fancy_lora.yaml
**Overriding the config**
You can override existing parameters from the command line using a key=value format. Let’s say you want to set the number of training epochs to 1.
Further information on config overrides can be found :ref:`here <cli_override>`.
.. code-block:: bash
tune run <RECIPE> --config <CONFIG> epochs=1
.. _validate_cli_label:
Validate a config
-----------------
The ``tune validate <config>`` command will validate that your config is formatted properly.
.. code-block:: bash
# If you've copied over a built-in config and want to validate custom changes
$ tune validate my_configs/llama3/8B_full.yaml
Config is well-formed!
.. _tune_cat_cli_label:
Inspect a config
---------------------
The ``tune cat <config>`` command pretty prints a configuration file, making it easy to use ``tune run`` with confidence. This command is useful for inspecting the structure and contents of a config file before running a recipe, ensuring that all parameters are correctly set.
You can also use the ``--sort`` option to print the config in sorted order, which can help in quickly locating specific keys.
.. list-table::
:widths: 30 60
* - \--sort
- Print the config in sorted order.
**Workflow Example**
1. **List all available configs:**
Use the ``tune ls`` command to list all the built-in recipes and configs within torchtune.
.. code-block:: bash
$ tune ls
RECIPE CONFIG
full_finetune_single_device llama2/7B_full_low_memory
llama3/8B_full_single_device
mistral/7B_full_low_memory
phi3/mini_full_low_memory
full_finetune_distributed llama2/7B_full
llama2/13B_full
llama3/8B_full
llama3/70B_full
...
2. **Inspect the contents of a config:**
Use the ``tune cat`` command to pretty print the contents of a specific config. This helps you understand the structure and parameters of the config.
.. code-block:: bash
$ tune cat llama2/7B_full
output_dir: /tmp/torchtune/llama2_7B/full
tokenizer:
_component_: torchtune.models.llama2.llama2_tokenizer
path: /tmp/Llama-2-7b-hf/tokenizer.model
max_seq_len: null
...
You can also print the config in sorted order:
.. code-block:: bash
$ tune cat llama2/7B_full --sort
3. **Run a recipe with parameter override:**
After inspecting the config, you can use the ``tune run`` command to run a recipe with the config. You can also override specific parameters directly from the command line. For example, to override the `output_dir` parameter:
.. code-block:: bash
$ tune run full_finetune_distributed --config llama2/7B_full output_dir=./
Learn more about config overrides :ref:`here <cli_override>`.
.. note::
You can find all the cat-able configs via the ``tune ls`` command.
---
### Source Basics Chat Datasets (`docs/source/basics/chat_datasets.rst`)
.. _chat_dataset_usage_label:
=============
Chat Datasets
=============
Chat datasets involve multi-turn conversations (multiple back-and-forths) between user and assistant.
.. code-block:: python
[
{"role": "user", "content": "What is the answer to the ultimate question of life?"},
{"role": "assistant", "content": "The answer is 42."},
{"role": "user", "content": "That's ridiculous"},
{"role": "assistant", "content": "Oh I know."},
]
This is more structured than freeform text association that models are typically pre-trained with,
where they learn to simply predict the next token instead of responding accurately to the user.
The primary entry point for fine-tuning with chat datasets in torchtune is the :func:`~torchtune.datasets.chat_dataset`
builder. This lets you specify a local or Hugging Face dataset that follows the chat data format
directly from the config and train your LLM on it.
.. _example_chat:
Example chat dataset
--------------------
.. code-block:: python
# data/my_data.json
[
{
"conversations": [
{
"from": "human",
"value": "What is the answer to life?"
},
{
"from": "gpt",
"value": "The answer is 42."
},
{
"from": "human",
"value": "That's ridiculous"
},
{
"from": "gpt",
"value": "Oh I know."
}
]
}
]
.. code-block:: python
from torchtune.models.mistral import mistral_tokenizer
from torchtune.datasets import chat_dataset
m_tokenizer = mistral_tokenizer(
path="/tmp/Mistral-7B-v0.1/tokenizer.model",
prompt_template="torchtune.models.mistral.MistralChatTemplate",
max_seq_len=8192,
)
ds = chat_dataset(
tokenizer=m_tokenizer,
source="json",
data_files="data/my_data.json",
split="train",
conversation_column="conversations",
conversation_style="sharegpt",
# By default, user prompt is ignored in loss. Set to True to include it
train_on_input=True,
new_system_prompt=None,
)
tokenized_dict = ds[0]
tokens, labels = tokenized_dict["tokens"], tokenized_dict["labels"]
print(m_tokenizer.decode(tokens))
# [INST] What is the answer to life? [/INST] The answer is 42. [INST] That's ridiculous [/INST] Oh I know.
print(labels)
# [1, 733, 16289, 28793, 1824, 349, 272, 4372, ...]
.. code-block:: yaml
# In config
tokenizer:
_component_: torchtune.models.mistral.mistral_tokenizer
path: /tmp/Mistral-7B-v0.1/tokenizer.model
prompt_template: torchtune.models.mistral.MistralChatTemplate
max_seq_len: 8192
dataset:
_component_: torchtune.datasets.chat_dataset
source: json
data_files: data/my_data.json
split: train
conversation_column: conversations
conversation_style: sharegpt
train_on_input: True
new_system_prompt: null
Chat dataset format
-------------------
Chat datasets typically have a single column named "conversations" or "messages" that contains a list of messages on a single topic
per sample. The list of messages could include a system prompt, multiple turns between user and assistant, and tool calls/returns.
.. code-block:: text
| conversations |
|--------------------------------------------------------------|
| [{"role": "user", "content": "What day is today?"}, |
| {"role": "assistant", "content": "It is Tuesday."}] |
| [{"role": "user", "content": "What about tomorrow?"}, |
| {"role": "assistant", "content": "Tomorrow is Wednesday."}] |
As an example, you can see the schema of the `SlimOrca dataset <https://huggingface.co/datasets/Open-Orca/SlimOrca-Dedup>`_.
Loading chat datasets from Hugging Face
---------------------------------------
You need to pass in the dataset repo name to ``source``, select one of the conversation styles in ``conversation_style``, and specify the ``conversation_column``.
For most HF datasets, you will also need to specify the ``split``.
.. code-block:: python
from torchtune.models.gemma import gemma_tokenizer
from torchtune.datasets import chat_dataset
g_tokenizer = gemma_tokenizer("/tmp/gemma-7b/tokenizer.model")
ds = chat_dataset(
tokenizer=g_tokenizer,
source="Open-Orca/SlimOrca-Dedup",
conversation_column="conversations",
conversation_style="sharegpt",
split="train",
)
.. code-block:: yaml
# Tokenizer is passed into the dataset in the recipe
dataset:
_component_: torchtune.datasets.chat_dataset
source: Open-Orca/SlimOrca-Dedup
conversation_column: conversations
conversation_style: sharegpt
split: train
Loading local and remote chat datasets
--------------------------------------
To load in a local or remote dataset via https that has conversational data, you need to additionally specify the ``data_files`` and ``split``
arguments. See Hugging Face's ``load_dataset`` `documentation <https://huggingface.co/docs/datasets/main/en/loading#local-and-remote-files>`_
for more details on loading local or remote files.
.. code-block:: python
from torchtune.models.gemma import gemma_tokenizer
from torchtune.datasets import chat_dataset
g_tokenizer = gemma_tokenizer("/tmp/gemma-7b/tokenizer.model")
ds = chat_dataset(
tokenizer=g_tokenizer,
source="json",
conversation_column="conversations",
conversation_style="sharegpt",
data_files="data/my_data.json",
split="train",
)
.. code-block:: yaml
# Tokenizer is passed into the dataset in the recipe
dataset:
_component_: torchtune.datasets.chat_dataset
source: json
conversation_column: conversations
conversation_style: sharegpt
data_files: data/my_data.json
split: train
Specifying conversation style
-----------------------------
The structure of the conversation in the raw dataset can vary widely with different role names and different fields
indicating the message content name. There are a few standardized formats that are common across many datasets.
We have built-in converters to convert these standardized formats into a list of torchtune :class:`~torchtune.data.Message`
that follows this format:
.. code-block:: python
[
{
"role": "system" | "user" | "assistant" | "ipython",
"content": <message>,
},
...
]
.. _sharegpt:
``"sharegpt"``
^^^^^^^^^^^^^^
The associated message transform is :class:`~torchtune.data.ShareGPTToMessages`. The expected format is:
.. code-block:: python
{
"conversations": [
{
"from": "system" | "human" | "gpt",
"value": <message>,
},
...
]
}
You can specify ``conversation_style=sharegpt`` in code or config:
.. code-block:: python
from torchtune.models.gemma import gemma_tokenizer
from torchtune.datasets import chat_dataset
g_tokenizer = gemma_tokenizer("/tmp/gemma-7b/tokenizer.model")
ds = chat_dataset(
tokenizer=g_tokenizer,
source="json",
conversation_column="conversations",
conversation_style="sharegpt",
data_files="data/my_data.json",
split="train",
)
.. code-block:: yaml
# Tokenizer is passed into the dataset in the recipe
dataset:
_component_: torchtune.datasets.chat_dataset
source: json
conversation_column: conversations
conversation_style: sharegpt
data_files: data/my_data.json
split: train
``"openai"``
^^^^^^^^^^^^
The associated message transform is :class:`~torchtune.data.OpenAIToMessages`. The expected format is:
.. code-block:: python
{
"messages": [
{
"role": "system" | "user" | "assistant",
"content": <message>,
},
...
]
}
You can specify ``conversation_style=openai`` in code or config:
.. code-block:: python
from torchtune.models.gemma import gemma_tokenizer
from torchtune.datasets import chat_dataset
g_tokenizer = gemma_tokenizer("/tmp/gemma-7b/tokenizer.model")
ds = chat_dataset(
tokenizer=g_tokenizer,
source="json",
conversation_column="conversations",
conversation_style="openai",
data_files="data/my_data.json",
split="train",
)
.. code-block:: yaml
# Tokenizer is passed into the dataset in the recipe
dataset:
_component_: torchtune.datasets.chat_dataset
source: json
conversation_column: conversations
conversation_style: openai
data_files: data/my_data.json
split: train
If your dataset does not fit one of the above conversation styles, then you will need to create a custom message transform.
Renaming columns
----------------
To specify the column that contains your conversation data, use ``conversation_column``.
.. code-block:: python
# data/my_data.json
[
{
"dialogue": [
{
"from": "human",
"value": "What is the answer to life?"
},
{
"from": "gpt",
"value": "The answer is 42."
},
{
"from": "human",
"value": "That's ridiculous"
},
{
"from": "gpt",
"value": "Oh I know."
}
]
}
]
.. code-block:: python
from torchtune.models.gemma import gemma_tokenizer
from torchtune.datasets import chat_dataset
g_tokenizer = gemma_tokenizer("/tmp/gemma-7b/tokenizer.model")
ds = chat_dataset(
tokenizer=g_tokenizer,
source="json",
conversation_column="dialogue",
conversation_style="sharegpt",
data_files="data/my_data.json",
split="train",
)
.. code-block:: yaml
# Tokenizer is passed into the dataset in the recipe
dataset:
_component_: torchtune.datasets.chat_dataset
source: json
conversation_column: dialogue
conversation_style: sharegpt
data_files: data/my_data.json
split: train
Chat templates
--------------
Chat templates are defined the same way as instruct templates in :func:`~torchtune.datasets.instruct_dataset`. See :ref:`instruct_template` for more info.
Built-in chat datasets
----------------------
- :class:`~torchtune.datasets.slimorca_dataset`
---
### Source Basics Custom Components (`docs/source/basics/custom_components.rst`)
.. _custom_components_label:
=============================
Custom Components and Recipes
=============================
torchtune lets you launch fine-tuning jobs directly from the command-line using both built-in and custom components,
such as datasets, models, recipes, and configs. This is done with the ``tune run`` command (see :ref:`cli_label`),
which can also be used from your project folder.
Setting up your torchtune project
---------------------------------
First, ensure that you have torchtune installed - see :ref:`install_label`. This will install the ``tune`` command
in your environment, so you can launch ``tune run`` from any directory. Let's create a new project directory and ensure
we can launch a built-in library recipe with a library config from that folder.
.. code-block:: bash
mkdir ~/my_project
cd ~/my_project
# This downloads the Llama 3.2 1B Instruct model
tune download meta-llama/Llama-3.2-1B-Instruct --output-dir /tmp/Llama-3.2-1B-Instruct --ignore-patterns "original/consolidated.00.pth"
# This launches a lora finetuning run with the default single device config
tune run lora_finetune_single_device --config llama3_2/1B_lora_single_device
Launching custom configs
------------------------
Often, you'll want to start with one of our default configs for a particular model and adjust a few training hyperparameters.
You can use the ``tune cp`` command to create a copy of a default config in your project directory so you can make modifications.
.. code-block:: bash
# Show all the default model configs for each recipe
tune ls
# This makes a copy of a Qwen2 full finetune config
tune cp qwen2/0.5B_full_single_device ~/my_project/config/qwen_config.yaml
Now, you can make modifications to the config directly in your project folder and launch the custom config. Make sure you are using
the correct recipe associated with the config and that you've downloaded the model. Even if you didn't start with copying a library
recipe, you can launch a completely custom config using the same command. Note that for custom configs, you must specify the file extension.
.. code-block:: bash
mkdir ~/my_project/config
tune run full_finetune_single_device --config ~/my_project/config/qwen_config.yaml
# Or launch directly from the project directory with a relative path
tune run full_finetune_single_device --config config/qwen_config.yaml
For a more detailed discussion on downloading models and modifying library configs, see :ref:`finetune_llama_label`.
Launching custom recipes
------------------------
torchtune's built-in recipes provide starting points for your fine-tuning workflows, but you can write your own training loop
with customized logic for your use case and launch training with ``tune run``. Similar to modifying library configs, you can
also copy one of our recipes as a starting point and modify, or write one completely from scratch. Note that for launching
custom recipes, you must specify the file extension.
.. code-block:: bash
mkdir ~/my_project/recipes
# Show all the default recipes
tune ls
# This makes a copy of the full finetune single device recipe locally
tune cp full_finetune_single_device ~/my_project/recipes/single_device.py
# Launch custom recipe with custom config from project directory
tune run recipes/single_device.py --config config/qwen_config.yaml
If you are writing a new recipe from scratch, we recommend following the Python convention of defining a ``main()`` function
in your script and decorating it with the :func:`~torchtune.config.parse` decorator. This will enable you to launch the recipe
with ``tune run`` and pass in a yaml file for the ``--config`` argument.
.. code-block:: python
from torchtune import config
from omegaconf import DictConfig
@config.parse
def main(cfg: DictConfig):
# Add all your recipe logic here, access config fields as attributes
if __name__ == "__main__":
# Config will be parsed from CLI, don't need to pass in here
main()
Launching with custom components
--------------------------------
torchtune supports full experimentation with custom models, datasets, optimizers, or any fine-tuning component. You can define
these locally in your repo and use them in your recipes and configs that you can launch with ``tune run``.
We recommend following the "builder" pattern when making your components. This means creating "builder" functions that set up
the classes you need with a few high level parameters that can be modified easily from the config. For example, we can define custom
model and dataset builders in our project directory:
.. code-block:: python
#
# In models/custom_decoder.py
#
class CustomTransformerDecoder(nn.Module):
# A custom architecture not present in torchtune
# Builder function for the custom model
def custom_model(num_layers: int, classification_head: bool = False):
# Any setup for defining the class
...
# Return the module you want to train
return CustomTransformerDecoder(...)
This allows us to expose our custom model in a config friendly manner - rather than having to define every argument needed to
construct our custom model in our config, we only expose the arguments which we care about modifying. This is how we implement
our models in torchtune - see :func:`~torchtune.models.llama3_2_vision.llama3_2_vision_11b` as an example.
.. code-block:: python
#
# In datasets/custom_dataset.py
#
from torchtune.datasets import SFTDataset, PackedDataset
from torchtune.data import InputOutputToMessages
from torchtune.modules.transforms.tokenizers import ModelTokenizer
# Example builder function for a custom code instruct dataset not in torchtune, but using
# different dataset building blocks from torchtune
def tiny_codes(tokenizer: ModelTokenizer, packed: bool = True):
"""
Python subset of nampdn-ai/tiny-codes. Instruct and code response pairs.
"""
ds = SFTDataset(
model_transform=tokenizer,
source="nampdn-ai/tiny-codes",
message_transform=InputOutputToMessages(
column_map={"input": "prompt", "output": "response"},
),
filter_fn=lambda x: x["language"] == "python",
split="train",
)
if packed:
return PackedDataset(ds, max_seq_len=tokenizer.max_seq_len, split_across_pack=False)
else:
return ds
.. note::
If you are using a default torchtune recipe with a custom dataset, you must define the first
positional argument to be the tokenizer or model transform. These are automatically passed into
dataset during instantiation and are defined separately in the config, not under the dataset field.
You can define the custom model and custom dataset in the config using the relative import path from where
you are launching with ``tune run``. It is best to define the path relative to your project root directory
and launch from there.
.. code-block:: yaml
# In YAML file config/custom_finetune.yaml
model:
_component_: models.custom_decoder.custom_model
num_layers: 32
# this is an optional param, so you can also omit this from the config
classification_head: False
dataset:
_component_: datasets.custom_dataset.tiny_codes
# we don't need to define a tokenizer here as it's automatically passed in
packed: True
.. code-block:: bash
cd ~/my_project/
tune run recipes/single_device.py --config config/custom_finetune.yaml
If your custom components are not being found or imported correctly, you can try to launch with ``tune run`` after
modifying the ``PYTHONPATH`` to ensure the files in your project directory are importable.
.. code-block:: bash
PYTHONPATH=${pwd}:PYTHONPATH tune run recipes/single_device.py --config config/custom_finetune.yaml
---
### Source Basics Custom Datasets (`docs/source/basics/custom_datasets.rst`)
.. _custom_dataset_usage_label:
===============
Custom Datasets
===============
If your dataset schema does not fit torchtune's built-in dataset builders, you can create
an end-to-end custom dataset pipeline by combining:
1. A custom message transform (raw sample -> ``messages``)
2. :class:`~torchtune.datasets.SFTDataset` (messages -> tokenized training samples)
This page shows the full flow in one place.
Create a custom message transform
---------------------------------
Start by converting your raw sample into torchtune :class:`~torchtune.data.Message` objects.
.. code-block:: python
from typing import Any, Mapping
from torchtune.data import Message
from torchtune.modules.transforms import Transform
class MyMessageTransform(Transform):
def __call__(self, sample: Mapping[str, Any]) -> Mapping[str, Any]:
return {
"messages": [
Message(role="user", content=sample["input"], masked=True, eot=True),
Message(role="assistant", content=sample["output"], masked=False, eot=True),
]
}
Create a custom dataset builder with ``SFTDataset``
---------------------------------------------------
Wrap the transform in a small dataset builder function.
.. code-block:: python
# data/dataset.py
from torchtune.datasets import SFTDataset
from data.message_transform import MyMessageTransform
def custom_dataset(tokenizer, **load_dataset_kwargs) -> SFTDataset:
return SFTDataset(
source="json",
data_files="data/my_data.json",
split="train",
message_transform=MyMessageTransform(),
model_transform=tokenizer,
**load_dataset_kwargs,
)
Use your custom dataset in a config
-----------------------------------
Point the recipe ``dataset`` component to your builder.
.. code-block:: yaml
dataset:
_component_: data.dataset.custom_dataset
For deeper details on message construction and custom component registration, see
:ref:`message_transform_usage_label` and :ref:`custom_components_label`.
---
### Source Basics Datasets Overview (`docs/source/basics/datasets_overview.rst`)
.. _datasets_overview:
=================
Datasets Overview
=================
torchtune lets you fine-tune LLMs and VLMs using any dataset found on Hugging Face Hub, downloaded locally,
or on a remote url. We provide built-in dataset builders to help you quickly bootstrap your fine-tuning project
for workflows including instruct tuning, preference alignment, continued pretraining, and more. Beyond those, torchtune
enables full customizability on your dataset pipeline, letting you train on any data format or schema.
The following tasks are supported:
- Text supervised fine-tuning
- :ref:`instruct_dataset_usage_label`
- :ref:`chat_dataset_usage_label`
- Fully custom datasets
- :ref:`custom_dataset_usage_label`
- Multimodal supervised fine-tuning
- :ref:`multimodal_dataset_usage_label`
- RLHF
- :ref:`preference_dataset_usage_label`
- Continued pre-training
- :ref:`text_completion_dataset_usage_label`
Data pipeline
-------------
.. image:: /_static/img/torchtune_datasets.svg
From raw data samples to the model inputs in the training recipe, all torchtune datasets follow
the same pipeline:
1. Raw data is queried one sample at a time from a Hugging Face dataset, local file, or remote file
2. :ref:`message_transform_usage_label` convert the raw sample which can take any format into a list of torchtune
:ref:`messages_usage_label`. Images are contained in the message object they are associated with.
3. :ref:`model_transform_usage_label` applies model-specific transforms to the messages, including tokenization (see :ref:`tokenizers_usage_label`),
prompt templating (see :ref:`prompt_templates_usage_label`), image transforms, and anything else required for that particular model.
4. The collater packages the processed samples together in a batch and the batch is passed into the model during training.
---
### Source Basics Instruct Datasets (`docs/source/basics/instruct_datasets.rst`)
.. _instruct_dataset_usage_label:
=================
Instruct Datasets
=================
Instruction tuning involves training an LLM to perform specific task(s). This typically takes the form
of a user command or prompt and the assistant's response, along with an optional system prompt that
describes the task at hand. This is more structured than freeform text association that models are
typically pre-trained with, where they learn to specifically predict the next token instead of completing
the task.
The primary entry point for fine-tuning with instruct datasets in torchtune is the :func:`~torchtune.datasets.instruct_dataset`
builder. This lets you specify a local or Hugging Face dataset that follows the instruct data format
directly from the config and train your LLM on it.
.. _example_instruct:
Example instruct dataset
------------------------
Here is an example of an instruct dataset to fine-tune for a grammar correction task.
.. code-block:: bash
head data/my_data.csv
# incorrect,correct
# This are a cat,This is a cat.
.. code-block:: python
from torchtune.models.gemma import gemma_tokenizer
from torchtune.datasets import instruct_dataset
g_tokenizer = gemma_tokenizer(
path="/tmp/gemma-7b/tokenizer.model",
prompt_template="torchtune.data.GrammarErrorCorrectionTemplate",
max_seq_len=8192,
)
ds = instruct_dataset(
tokenizer=g_tokenizer,
source="csv",
data_files="data/my_data.csv",
split="train",
# By default, user prompt is ignored in loss. Set to True to include it
train_on_input=True,
# Prepend a system message to every sample
new_system_prompt="You are an AI assistant. ",
# Use columns in our dataset instead of default
column_map={"input": "incorrect", "output": "correct"},
)
tokenized_dict = ds[0]
tokens, labels = tokenized_dict["tokens"], tokenized_dict["labels"]
print(g_tokenizer.decode(tokens))
# You are an AI assistant. Correct this to standard English:This are a cat---\nCorrected:This is a cat.
print(labels) # System message is masked out, but not user message
# [-100, -100, -100, -100, -100, -100, 27957, 736, 577, ...]
.. code-block:: yaml
# In config
tokenizer:
_component_: torchtune.models.gemma.gemma_tokenizer
path: /tmp/gemma-7b/tokenizer.model
prompt_template: torchtune.data.GrammarErrorCorrectionTemplate
max_seq_len: 8192
dataset:
source: csv
data_files: data/my_data.csv
split: train
train_on_input: True
new_system_prompt: You are an AI assistant.
column_map:
input: incorrect
output: correct
Instruct dataset format
-----------------------
Instruct datasets are expected to follow an input-output format, where the user prompt is in one column
and the assistant prompt is in another column.
.. code-block:: text
| input | output |
|-----------------|------------------|
| "user prompt" | "model response" |
As an example, you can see the schema of the `C4 200M dataset <https://huggingface.co/datasets/liweili/c4_200m>`_.
Loading instruct datasets from Hugging Face
-------------------------------------------
You simply need to pass in the dataset repo name to ``source``, which is then passed into Hugging Face's ``load_dataset``.
For most datasets, you will also need to specify the ``split``.
.. code-block:: python
# In code
from torchtune.models.gemma import gemma_tokenizer
from torchtune.datasets import instruct_dataset
g_tokenizer = gemma_tokenizer("/tmp/gemma-7b/tokenizer.model")
ds = instruct_dataset(
tokenizer=g_tokenizer,
source="liweili/c4_200m",
split="train"
)
.. code-block:: yaml
# In config
tokenizer:
_component_: torchtune.models.gemma.gemma_tokenizer
path: /tmp/gemma-7b/tokenizer.model
# Tokenizer is passed into the dataset in the recipe
dataset:
_component_: torchtune.datasets.instruct_dataset
source: liweili/c4_200m
split: train
This will use the default column names "input" and "output". To change the column names, use the ``column_map`` argument (see :ref:`column_map`).
Loading local and remote instruct datasets
------------------------------------------
To load in a local or remote dataset via https that follows the instruct format, you need to specify the ``source``, ``data_files`` and ``split``
arguments. See Hugging Face's ``load_dataset`` `documentation <https://huggingface.co/docs/datasets/main/en/loading#local-and-remote-files>`_
for more details on loading local or remote files.
.. code-block:: python
# In code
from torchtune.models.gemma import gemma_tokenizer
from torchtune.datasets import instruct_dataset
g_tokenizer = gemma_tokenizer("/tmp/gemma-7b/tokenizer.model")
ds = instruct_dataset(
tokenizer=g_tokenizer,
source="json",
data_files="data/my_data.json",
split="train",
)
.. code-block:: yaml
# In config
tokenizer:
_component_: torchtune.models.gemma.gemma_tokenizer
path: /tmp/gemma-7b/tokenizer.model
# Tokenizer is passed into the dataset in the recipe
dataset:
_component_: torchtune.datasets.instruct_dataset
source: json
data_files: data/my_data.json
split: train
.. _column_map:
Renaming columns
----------------
You can remap the default column names to the column names in your dataset by specifying
``column_map`` as ``{"<default column>": "<column in your dataset>"}``. The default column names
are detailed in each of the dataset builders (see :func:`~torchtune.datasets.instruct_dataset` and
:func:`~torchtune.datasets.chat_dataset` as examples).
For example, if the default column names are "input", "output" and you need to change them to something else,
such as "prompt", "response", then ``column_map = {"input": "prompt", "output": "response"}``.
.. code-block:: python
# data/my_data.json
[
{"prompt": "hello world", "response": "bye world"},
{"prompt": "are you a robot", "response": "no, I am an AI assistant"},
...
]
.. code-block:: python
from torchtune.models.gemma import gemma_tokenizer
from torchtune.datasets import instruct_dataset
g_tokenizer = gemma_tokenizer("/tmp/gemma-7b/tokenizer.model")
ds = instruct_dataset(
tokenizer=g_tokenizer,
source="json",
data_files="data/my_data.json",
split="train",
column_map={"input": "prompt", "output": "response"},
)
.. code-block:: yaml
# Tokenizer is passed into the dataset in the recipe
dataset:
_component_: torchtune.datasets.instruct_dataset
source: json
data_files: data/my_data.json
split: train
column_map:
input: prompt
output: response
.. _instruct_template:
Instruct templates
------------------
Typically for instruct datasets, you will want to add a :class:`~torchtune.data.PromptTemplate` to provide task-relevant
information. For example, for a grammar correction task, we may want to use a prompt template like :class:`~torchtune.data.GrammarErrorCorrectionTemplate`
to structure each of our samples. Prompt templates are passed into the tokenizer and automatically applied to the dataset
you are fine-tuning on. See :ref:`using_prompt_templates` for more details.
Built-in instruct datasets
--------------------------
- :class:`~torchtune.datasets.alpaca_dataset`
- :class:`~torchtune.datasets.grammar_dataset`
- :class:`~torchtune.datasets.samsum_dataset`
---
### Source Basics Message Transforms (`docs/source/basics/message_transforms.rst`)
.. _message_transform_usage_label:
==================
Message Transforms
==================
Message transforms perform the conversion of raw sample dictionaries from your dataset into torchtune's
:class:`~torchtune.data.Message` structure. Once you data is represented as Messages, torchtune will handle
tokenization and preparing it for the model.
.. TODO (rafiayub): place an image here to depict overall pipeline
Configuring message transforms
------------------------------
Most of our built-in message transforms contain parameters for controlling input masking (``masking_strategy``),
adding a system prompt (``new_system_prompt``), and changing the expected column names (``column_map``).
These are exposed in our dataset builders :func:`~torchtune.datasets.instruct_dataset` and :func:`~torchtune.datasets.chat_dataset`
so you don't have to worry about the message transform itself and can configure this directly from the config.
You can see :ref:`example_instruct` or :ref:`example_chat` for more details.
.. _custom_message_transform:
Custom message transforms
-------------------------
If our built-in message transforms do not configure for your particular dataset well,
you can create your own class with full flexibility. Simply inherit from the :class:`~torchtune.modules.transforms.Transform`
class and add your code in the ``__call__`` method.
A simple contrived example would be to take one column from the dataset as the user message and another
column as the model response. Indeed, this is quite similar to :class:`~torchtune.data.InputOutputToMessages`.
.. code-block:: python
from torchtune.modules.transforms import Transform
from torchtune.data import Message
from typing import Any, Mapping
from pprint import pprint
class MessageTransform(Transform):
def __call__(self, sample: Mapping[str, Any]) -> Mapping[str, Any]:
messages = [
Message(
role="user",
content=sample["input"],
masked=True,
eot=True,
),
Message(
role="assistant",
content=sample["output"],
masked=False,
eot=True,
),
]
return {"messages": messages}
input_sample = {"input": "hello world", "output": "bye world"}
transform = MessageTransform()
output_sample = transform(input_sample)
pprint(output_sample)
# {'messages': [Message(role='user', content=['hello world']),
# Message(role='assistant', content=['bye world'])]}
See :ref:`creating_messages` for more details on how to manipulate :class:`~torchtune.data.Message` objects.
To use this for your dataset, you must create a custom dataset builder that uses the underlying
dataset class, :class:`~torchtune.datasets.SFTDataset`.
.. code-block:: python
# In data/dataset.py
from torchtune.datasets import SFTDataset
def custom_dataset(tokenizer, **load_dataset_kwargs) -> SFTDataset:
message_transform = MyMessageTransform()
return SFTDataset(
source="json",
data_files="data/my_data.json",
split="train",
message_transform=message_transform,
model_transform=tokenizer,
**load_dataset_kwargs,
)
This can be used directly from the config.
.. code-block:: yaml
dataset:
_component_: data.dataset.custom_dataset
For a full end-to-end walkthrough (custom transform + ``SFTDataset`` + config wiring), see
:ref:`custom_dataset_usage_label`.
Example message transforms
--------------------------
- Instruct
- :class:`~torchtune.data.InputOutputToMessages`
- :class:`~torchtune.data.AlpacaToMessages`
- Chat
- :class:`~torchtune.data.ShareGPTToMessages`
- :class:`~torchtune.data.OpenAIToMessages`
- Preference
- :class:`~torchtune.data.ChosenRejectedToMessages`
---
### Source Basics Messages (`docs/source/basics/messages.rst`)
.. _messages_usage_label:
========
Messages
========
Messages are a core component in torchtune that govern how text and multimodal content is tokenized. It serves as the common interface
for all tokenizer and datasets APIs to operate on. Messages contain information about the text content, which role is sending the text
content, and other information relevant for special tokens in model tokenizers. For more information about the individual parameters
for Messages, see the API ref for :class:`~torchtune.data.Message`.
.. _creating_messages:
Creating Messages
-----------------
Messages can be created via the standard class constructor or directly from a dictionary.
.. code-block:: python
from torchtune.data import Message
msg = Message(
role="user",
content="Hello world!",
masked=True,
eot=True,
ipython=False,
)
# This is identical
msg = Message.from_dict(
{
"role": "user",
"content": "Hello world!",
"masked": True,
"eot": True,
"ipython": False,
},
)
print(msg.content)
# [{'type': 'text', 'content': 'Hello world!'}]
Content is formatted as a list of dictionaries. This is because Messages can also contain multimodal content, such as images.
Images in Messages
^^^^^^^^^^^^^^^^^^
For multimodal datasets, you need to add the image as a :class:`~PIL.Image.Image` to the corresponding :class:`~torchtune.data.Message`.
To add it to the beginning of the message, simply prepend it to the content list.
.. code-block:: python
import PIL
from torchtune.data import Message
img_msg = Message(
role="user",
content=[
{
"type": "image",
# Place your image here
"content": PIL.Image.new(mode="RGB", size=(4, 4)),
},
{"type": "text", "content": "What's in this image?"},
],
)
This will indicate to the model tokenizers where to add the image special token and will be processed by the model transform
appropriately.
In many cases, you will have an image path instead of a raw :class:`~PIL.Image.Image`. You can use the :func:`~torchtune.data.load_image`
utility for both local paths and remote paths.
.. code-block:: python
import PIL
from torchtune.data import Message, load_image
image_path = "path/to/image.jpg"
img_msg = Message(
role="user",
content=[
{
"type": "image",
# Place your image here
"content": load_image(image_path),
},
{"type": "text", "content": "What's in this image?"},
],
)
If your dataset contain image tags, or placeholder text to indicate where in the text the image should be inserted,
you can use the :func:`~torchtune.data.format_content_with_images` to split the text into the correct content list
that you can pass into the content field of Message.
.. code-block:: python
import PIL
from torchtune.data import format_content_with_images
content = format_content_with_images(
"<|image|>hello <|image|>world",
image_tag="<|image|>",
images=[PIL.Image.new(mode="RGB", size=(4, 4)), PIL.Image.new(mode="RGB", size=(4, 4))]
)
print(content)
# [
# {"type": "image", "content": },
# {"type": "text", "content": "hello "},
# {"type": "image", "content": },
# {"type": "text", "content": "world"}
# ]
Message transforms
^^^^^^^^^^^^^^^^^^
Message transforms are convenient utilities to format raw data into a list of torchtune :class:`~torchtune.data.Message`
objects.
.. code-block:: python
from torchtune.data import InputOutputToMessages
sample = {
"input": "What is your name?",
"output": "I am an AI assistant, I don't have a name."
}
transform = InputOutputToMessages()
output = transform(sample)
for message in output["messages"]:
print(message.role, message.text_content)
# user What is your name?
# assistant I am an AI assistant, I don't have a name.
See :ref:`message_transform_usage_label` for more discussion.
Formatting messages with prompt templates
-----------------------------------------
Prompt templates provide a way to format messages into a structured text template. You can simply call any class that inherits
from :class:`~torchtune.data.PromptTemplateInterface` on a list of Messages and it will add the appropriate text to the content
list.
.. code-block:: python
from torchtune.models.mistral import MistralChatTemplate
from torchtune.data import Message
msg = Message(
role="user",
content="Hello world!",
masked=True,
eot=True,
ipython=False,
)
template = MistralChatTemplate()
templated_msg = template([msg])
print(templated_msg[0].content)
# [{'type': 'text', 'content': '[INST] '},
# {'type': 'text', 'content': 'Hello world!'},
# {'type': 'text', 'content': ' [/INST] '}]
Accessing text content in messages
----------------------------------
.. code-block:: python
from torchtune.models.mistral import MistralChatTemplate
from torchtune.data import Message
msg = Message(
role="user",
content="Hello world!",
masked=True,
eot=True,
ipython=False,
)
template = MistralChatTemplate()
templated_msg = template([msg])
print(templated_msg[0].text_content)
# [INST] Hello world! [/INST]
Accessing images in messages
----------------------------
.. code-block:: python
from torchtune.data import Message
import PIL
msg = Message(
role="user",
content=[
{
"type": "image",
# Place your image here
"content": PIL.Image.new(mode="RGB", size=(4, 4)),
},
{"type": "text", "content": "What's in this image?"},
],
)
if msg.contains_media:
print(msg.get_media())
# []
Tokenizing messages
-------------------
All model tokenizers have a ``tokenize_messsages`` method that converts a list of
:class:`~torchtune.data.Message` objects into token IDs and a loss mask.
.. code-block:: python
from torchtune.models.mistral import mistral_tokenizer
from torchtune.data import Message
m_tokenizer = mistral_tokenizer(
path="/tmp/Mistral-7B-v0.1/tokenizer.model",
prompt_template="torchtune.models.mistral.MistralChatTemplate",
max_seq_len=8192,
)
msgs = [
Message(
role="user",
content="Hello world!",
masked=True,
eot=True,
ipython=False,
),
Message(
role="assistant",
content="Hi, I am an AI assistant.",
masked=False,
eot=True,
ipython=False,
)
]
tokens, mask = m_tokenizer.tokenize_messages(msgs)
print(tokens)
# [1, 733, 16289, 28793, 22557, 1526, 28808, 28705, 733, 28748, 16289, 28793, 15359, 28725, 315, 837, 396, 16107, 13892, 28723, 2]
print(mask) # User message is masked from the loss
# [True, True, True, True, True, True, True, True, True, True, True, True, False, False, False, False, False, False, False, False, False]
print(m_tokenizer.decode(tokens))
# [INST] Hello world! [/INST] Hi, I am an AI assistant.
---
### Source Basics Model Transforms (`docs/source/basics/model_transforms.rst`)
.. _model_transform_usage_label:
=====================
Multimodal Transforms
=====================
Multimodal model transforms apply model-specific data transforms to each modality and prepares :class:`~torchtune.data.Message`
objects to be input into the model. torchtune currently supports text + image model transforms.
These are intended to be drop-in replacements for tokenizers in multimodal datasets and support the standard
``encode``, ``decode``, and ``tokenize_messages``.
.. code-block:: python
# torchtune.models.llama3_2_vision.Llama3VisionTransform
class Llama3VisionTransform(ModelTokenizer, Transform):
def __init__(...):
# Text transform - standard tokenization
self.tokenizer = llama3_tokenizer(...)
# Image transforms
self.transform_image = CLIPImageTransform(...)
self.xattn_mask = VisionCrossAttentionMask(...)
.. code-block:: python
from torchtune.models.llama3_2_vision import Llama3VisionTransform
from torchtune.data import Message
from PIL import Image
sample = {
"messages": [
Message(
role="user",
content=[
{"type": "image", "content": Image.new(mode="RGB", size=(560, 560))},
{"type": "image", "content": Image.new(mode="RGB", size=(560, 560))},
{"type": "text", "content": "What is common in these two images?"},
],
),
Message(
role="assistant",
content="A robot is in both images.",
),
],
}
transform = Llama3VisionTransform(
path="/tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model",
tile_size=224,
patch_size=14,
)
tokenized_dict = transform(sample)
print(transform.decode(tokenized_dict["tokens"], skip_special_tokens=False))
# '<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n<|image|><|image|>What is common in these two images?<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\nA robot is in both images.<|eot_id|>'
print(tokenized_dict["encoder_input"]["images"][0].shape) # (num_tiles, num_channels, tile_height, tile_width)
# torch.Size([4, 3, 224, 224])
Using model transforms
----------------------
You can pass them into any multimodal dataset builder just as you would a model tokenizer.
.. code-block:: python
from torchtune.datasets.multimodal import the_cauldron_dataset
from torchtune.models.llama3_2_vision import Llama3VisionTransform
transform = Llama3VisionTransform(
path="/tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model",
tile_size=224,
patch_size=14,
)
ds = the_cauldron_dataset(
model_transform=transform,
subset="ai2d",
)
tokenized_dict = ds[0]
print(transform.decode(tokenized_dict["tokens"], skip_special_tokens=False))
# <|begin_of_text|><|start_header_id|>user<|end_header_id|>
#
# <|image|>Question: What do respiration and combustion give out
# Choices:
# A. Oxygen
# B. Carbon dioxide
# C. Nitrogen
# D. Heat
# Answer with the letter.<|eot_id|><|start_header_id|>assistant<|end_header_id|>
#
# Answer: B<|eot_id|>
print(tokenized_dict["encoder_input"]["images"][0].shape) # (num_tiles, num_channels, tile_height, tile_width)
# torch.Size([4, 3, 224, 224])
Creating model transforms
-------------------------
Model transforms are expected to process both text and images in the sample dictionary.
Both should be contained in the ``"messages"`` field of the sample.
The following methods are required on the model transform:
- ``tokenize_messages``
- ``__call__``
.. code-block:: python
from torchtune.modules.transforms.tokenizers import ModelTokenizer
from torchtune.modules.transforms import Transform
class MyMultimodalTransform(ModelTokenizer, Transform):
def __init__(...):
self.tokenizer = my_tokenizer_builder(...)
self.transform_image = MyImageTransform(...)
def tokenize_messages(
self,
messages: list[Message],
add_eos: bool = True,
) -> tuple[list[int], list[bool]]:
# Any other custom logic here
...
return self.tokenizer.tokenize_messages(
messages=messages,
add_eos=add_eos,
)
def __call__(
self, sample: Mapping[str, Any], inference: bool = False
) -> Mapping[str, Any]:
# Expected input parameters for vision encoder
encoder_input = {"images": [], "aspect_ratio": []}
messages = sample["messages"]
# Transform all images in sample
for message in messages:
for image in message.get_media():
out = self.transform_image({"image": image}, inference=inference)
encoder_input["images"].append(out["image"])
encoder_input["aspect_ratio"].append(out["aspect_ratio"])
sample["encoder_input"] = encoder_input
# Transform all text - returns same dictionary with additional keys "tokens" and "mask"
sample = self.tokenizer(sample, inference=inference)
return sample
transform = MyMultimodalTransform(...)
sample = {
"messages": [
Message(
role="user",
content=[
{"type": "image", "content": Image.new(mode="RGB", size=(224, 224))},
{"type": "image", "content": Image.new(mode="RGB", size=(224, 224))},
{"type": "text", "content": "What is common in these two images?"},
],
),
Message(
role="assistant",
content="A robot is in both images.",
),
],
}
tokenized_dict = transform(sample)
print(tokenized_dict)
# {'encoder_input': {'images': ..., 'aspect_ratio': ...}, 'tokens': ..., 'mask': ...}
Example model transforms
------------------------
- Llama 3.2 Vision
- :class:`~torchtune.models.llama3_2_vision.Llama3VisionTransform`
---
### Source Basics Multimodal Datasets (`docs/source/basics/multimodal_datasets.rst`)
.. _multimodal_dataset_usage_label:
===================
Multimodal Datasets
===================
Multimodal datasets include more than one data modality, e.g. text + image, and can be used to train transformer-based models.
torchtune currently only supports multimodal text+image chat datasets for Vision-Language Models (VLMs).
The primary entry point for fine-tuning with multimodal datasets in torchtune is the :func:`~torchtune.datasets.multimodal.multimodal_chat_dataset`
builder. This lets you specify a local or Hugging Face dataset that follows the multimodal chat data format
directly from the config and train your VLM on it.
.. _example_multimodal:
Example multimodal dataset
--------------------------
Here is an example of a multimodal chat dataset for a visual question-answering task. Note that there is a placeholder
in the text, ``""`` for where to place the image tokens. This will get replaced by the image special token
``<|image|>`` in the example below.
.. code-block:: python
# data/my_data.json
[
{
"dialogue": [
{
"from": "human",
"value": "What time is it on the clock?",
},
{
"from": "gpt",
"value": "It is 10:00 AM.",
},
],
"image_path": "images/clock.jpg",
},
...,
]
.. code-block:: python
from torchtune.models.llama3_2_vision import llama3_2_vision_transform
from torchtune.datasets.multimodal import multimodal_chat_dataset
model_transform = llama3_2_vision_transform(
path="/tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model",
prompt_template="torchtune.data.QuestionAnswerTemplate",
max_seq_len=8192,
image_size=560,
)
ds = multimodal_chat_dataset(
model_transform=model_transform,
source="json",
data_files="data/my_data.json",
column_map={
"dialogue": "conversations",
"image_path": "image",
},
image_dir="/home/user/dataset/", # /home/user/dataset/images/clock.jpg
image_tag="",
split="train",
)
tokenized_dict = ds[0]
print(model_transform.decode(tokenized_dict["tokens"], skip_special_tokens=False))
# '<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\nQuestion:<|image|>What time is it on the clock?Answer:<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\nIt is 10:00AM.<|eot_id|>'
print(tokenized_dict["encoder_input"]["images"][0].shape) # (num_tiles, num_channels, tile_height, tile_width)
# torch.Size([4, 3, 224, 224])
.. code-block:: yaml
tokenizer:
_component_: torchtune.models.llama3_2_vision_transform
path: /tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model
prompt_template: torchtune.data.QuestionAnswerTemplate
max_seq_len: 8192
image_size: 560
dataset:
_component_: torchtune.datasets.multimodal.multimodal_chat_dataset
source: json
data_files: data/my_data.json
split: train
column_map:
dialogue: conversations
image_path: image
image_dir: /home/user/dataset/
image_tag: ""
split: train
Multimodal dataset format
-------------------------
Multimodal datasets are currently expected to follow the :ref:`sharegpt` chat format, where the image paths are in one column
and the user-assistant conversations are in another column.
.. code-block:: text
| conversations | image |
|------------------------------------|--------------|
| [{"from": "human", "value": "Q1"}, | images/1.jpg |
| {"from": "gpt", "value": "A1"}] | |
As an example, you can see the schema of the `ShareGPT4V dataset <https://huggingface.co/datasets/Lin-Chen/ShareGPT4V>`_.
Currently, :func:`~torchtune.datasets.multimodal.multimodal_chat_dataset` only supports a single image path per conversation sample.
Loading multimodal datasets from Hugging Face
---------------------------------------------
You simply need to pass in the dataset repo name to ``source``, which is then passed into Hugging Face's ``load_dataset``.
For most datasets, you will also need to specify the ``split`` and/or the subset via ``name``.
.. code-block:: python
# In code
from torchtune.models.llama3_2_vision import llama3_2_vision_transform
from torchtune.datasets.multimodal import multimodal_chat_dataset
model_transform = llama3_2_vision_transform(
path="/tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model",
max_seq_len=8192,
image_size=560,
)
ds = multimodal_chat_dataset(
model_transform=model_transform,
source="Lin-Chen/ShareGPT4V",
split="train",
name="ShareGPT4V",
image_dir="/home/user/dataset/",
image_tag="",
)
.. code-block:: yaml
# In config
tokenizer:
_component_: torchtune.models.llama3_2_vision.llama3_2_vision_transform
path: /tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model
max_seq_len: 8192
image_size: 560
# Tokenizer is passed into the dataset in the recipe
dataset:
_component_: torchtune.datasets.multimodal.multimodal_chat_dataset
source: Lin-Chen/ShareGPT4V
split: train
name: ShareGPT4V
image_dir: /home/user/dataset/
image_tag: ""
This will use the default column names "conversations" and "image". To change the column names, use the ``column_map`` argument (see :ref:`column_map`).
Loading local and remote multimodal datasets
--------------------------------------------
To load in a local or remote dataset via https that follows the instruct format, you need to specify the ``source``, ``data_files`` and ``split``
arguments. See Hugging Face's ``load_dataset`` `documentation <https://huggingface.co/docs/datasets/main/en/loading#local-and-remote-files>`_
for more details on loading local or remote files. See :ref:`example_multimodal` above.
Loading images
--------------
In many cases, your dataset will contain paths to the images instead of the raw images themselves. :func:`~torchtune.datasets.multimodal.multimodal_chat_dataset`
will automatically handle this for you, but if you are writing a custom message transform for a custom multimodal dataset
(see :ref:`custom_message_transform`), you can use the :func:`~torchtune.data.load_image` utility directly.
.. code-block:: python
from torchtune.data import load_image
from pathlib import Path
sample = {
"conversations": [
{
"from": "human",
"value": "What time is it on the clock?",
},
{
"from": "gpt",
"value": "It is 10:00 AM.",
},
],
"image": "images/clock.jpg",
}
image_dir = "/home/user/dataset/"
pil_image = load_image(Path(image_dir) / Path(sample["image"]))
print(pil_image)
#
Then, you can add the PIL image directly to the content of the related message. Only PIL images are supported as image content
in :class:`~torchtune.data.Message`, not image paths or urls.
.. code-block:: python
from torchtune.data import Message
user_message = None
for msg in sample["conversations"]:
if msg["from"] == "human":
user_message = Message(
role="user",
content=[
{"type": "image", "content": pil_image},
{"type": "text", "content": msg["value"]},
]
)
print(user_message.contains_media)
# True
print(user_message.get_media())
# []
print(user_message.text_content)
# What time is it on the clock?
If the image paths in your dataset are relative paths, you can use the ``image_dir`` parameter in :func:`~torchtune.datasets.multimodal.multimodal_chat_dataset`
to prepend the full path where your images are downloaded locally.
Interleaving images in text
---------------------------
torchtune supports adding multiple images in any locations in the text, as long as your model supports it.
.. code-block:: python
import PIL
from torchtune.data import Message
image_dog = PIL.Image.new(mode="RGB", size=(4, 4))
image_cat = PIL.Image.new(mode="RGB", size=(4, 4))
image_bird = PIL.Image.new(mode="RGB", size=(4, 4))
user_message = Message(
role="user",
content=[
{"type": "image", "content": image_dog},
{"type": "text", "content": "This is an image of a dog. "},
{"type": "image", "content": image_cat},
{"type": "text", "content": "This is an image of a cat. "},
{"type": "image", "content": image_bird},
{"type": "text", "content": "This is a bird, the best pet of the three."},
]
)
print(user_message.contains_media)
# True
print(user_message.get_media())
# [, , ]
print(user_message.text_content)
# This is an image of a dog. This is an image of a cat. This is a bird, the best pet of the three.
Your dataset may contain image placeholder tags which indicate where in the text the image should be referenced
As an example, see `ShareGPT4V <https://huggingface.co/datasets/Lin-Chen/ShareGPT4V>`, which uses ``""``.
You can easily create the interleaved message content similar to above with the utility :func:`~torchtune.data.format_content_with_images`,
which replaces the image placeholder tags with the passed in images.
.. code-block:: python
import PIL
from torchtune.data import Message, format_content_with_images
image_dog = PIL.Image.new(mode="RGB", size=(4, 4))
image_cat = PIL.Image.new(mode="RGB", size=(4, 4))
image_bird = PIL.Image.new(mode="RGB", size=(4, 4))
text = "[img]This is an image of a dog. [img]This is an image of a cat. [img]This is a bird, the best pet of the three."
user_message = Message(
role="user",
content=format_content_with_images(
content=text,
image_tag="[img]",
images=[image_dog, image_cat, image_bird],
),
)
print(user_message.contains_media)
# True
print(user_message.get_media())
# [,, ]
print(user_message.text_content)
# This is an image of a dog. This is an image of a cat. This is a bird, the best pet of the three.
This is handled automatically for you in :func:`~torchtune.datasets.multimodal.multimodal_chat_dataset` when you pass in
``image_tag``.
Built-in multimodal datasets
----------------------------
- :class:`~torchtune.datasets.multimodal.the_cauldron_dataset`
- :class:`~torchtune.datasets.multimodal.llava_instruct_dataset`
---
### Source Basics Packing (`docs/source/basics/packing.rst`)
.. _packing_usage_label:
==============
Sample packing
==============
Sample packing involves concatenating multiple samples from your dataset into a single sequence, upto a maximum
sequence length. This requires some pre-processing of the dataset which may
slow down time-to-first-batch, but can introduce significant training speedups
depending on the dataset. In torchtune, sample packing is done by iterating through your dataset and performing
greedy packing upon dataset initialization. You can use sample packing with any of the single dataset builders by passing in
:code:`packed=True`.
To set the max sequence length to pack to, make sure to define ``max_seq_len`` on your tokenizer.
.. code-block:: python
from torchtune.datasets import alpaca_dataset, PackedDataset
from torchtune.models.llama3 import llama3_tokenizer
# Load in tokenizer
tokenizer = llama3_tokenizer(
path="/tmp/Llama-3.2-1B-Instruct/original/tokenizer.model",
max_seq_len=8192,
)
dataset = alpaca_dataset(
tokenizer=tokenizer,
packed=True,
)
print(isinstance(dataset, PackedDataset)) # True
.. code-block:: yaml
# YAML config
tokenizer:
_component_: torchtune.models.llama3.llama3_tokenizer
path: /tmp/Llama-3.2-1B-Instruct/original/tokenizer.model
max_seq_len: 8192
dataset:
_component_: torchtune.datasets.alpaca_dataset
packed: True
.. code-block:: bash
# Command line
tune run full_finetune_single_device --config llama3_2/1B_full_single_device \
dataset.packed=True tokenizer.max_seq_len=8192
torchtune will automatically handle document masking and relative position IDs when sample packing is enabled
to prevent different irrelevant samples from cross-attending. This is done via PyTorch's `Flex Attention <https://pytorch.org/blog/flexattention/#document-maskingjagged-sequences>`_,
which enables the use of flash attention with non-causal masks. If your hardware does not support Flex Attention
(for CUDA devices, it must be Turing or above), standard SDPA with memory-efficient attention will be used as a fallback,
while retaining the document masking and relative position IDs.
---
### Source Basics Preference Datasets (`docs/source/basics/preference_datasets.rst`)
.. _preference_dataset_usage_label:
===================
Preference Datasets
===================
Preference datasets are used for reward modelling, where the downstream task is to fine-tune a base model
to capture some underlying human preferences. Currently, these datasets are used in torchtune with the
Direct Preference Optimization (DPO) `recipe <https://github.com/pytorch/torchtune/blob/main/recipes/lora_dpo_single_device.py>`_.
The ground-truth in preference datasets is usually the outcome of a binary comparison between two completions for the same prompt,
and where a human annotator has indicated that one completion is more preferable than the other, according to some pre-set criterion.
These prompt-completion pairs could be instruct style (single-turn, optionally with a single prompt), chat style (multi-turn), or
some other set of interactions between a user and model (e.g. free-form text completion).
The primary entry point for fine-tuning with preference datasets in torchtune with the DPO recipe is :func:`~torchtune.datasets.preference_dataset`.
Example local preference dataset
--------------------------------
.. code-block:: bash
# my_preference_dataset.json
[
{
"chosen_conversations": [
{
"content": "What do I do when I have a hole in my trousers?",
"role": "user"
},
{ "content": "Fix the hole.", "role": "assistant" }
],
"rejected_conversations": [
{
"content": "What do I do when I have a hole in my trousers?",
"role": "user"
},
{ "content": "Take them off.", "role": "assistant" }
]
}
]
.. code-block:: python
from torchtune.models.mistral import mistral_tokenizer
from torchtune.datasets import preference_dataset
m_tokenizer = mistral_tokenizer(
path="/tmp/Mistral-7B-v0.1/tokenizer.model",
prompt_template="torchtune.models.mistral.MistralChatTemplate",
max_seq_len=8192,
)
column_map = {
"chosen": "chosen_conversations",
"rejected": "rejected_conversations"
}
ds = preference_dataset(
tokenizer=tokenizer,
source="json",
column_map=column_map,
data_files="my_preference_dataset.json",
train_on_input=False,
split="train",
)
tokenized_dict = ds[0]
print(m_tokenizer.decode(tokenized_dict["rejected_input_ids"]))
# user\n\nWhat do I do when I have a hole in my trousers?assistant\n\nTake them off.
print(tokenized_dict["rejected_labels"])
# [-100,-100,-100,-100,-100,-100,-100,-100,-100,-100,-100,-100, -100,-100,\
# -100,-100,-100,-100,-100,128006,78191,128007,271,18293,1124,1022,13,128009,-100]
This can also be accomplished via the yaml config:
.. code-block:: yaml
# In config
tokenizer:
_component_: torchtune.models.mistral.mistral_tokenizer
path: /tmp/Mistral-7B-v0.1/tokenizer.model
prompt_template: torchtune.models.mistral.MistralChatTemplate
max_seq_len: 8192
dataset:
_component_: torchtune.datasets.preference_dataset
source: json
data_files: my_preference_dataset.json
column_map:
chosen: chosen_conversations
rejected: rejected_conversations
train_on_input: False
split: train
In this example, we've also shown how `column_map` can be used when the "chosen" and/or "rejected" column names differ from the corresponding columns in your dataset.
Preference dataset format
-------------------------
Preference datasets are expected to have two columns: *"chosen"*, which indicates the human annotator's preferred response, and *"rejected"*, indicating
the human annotator's dis-preferred response. Each of these columns should contain a list of messages with an identical prompt.
The list of messages could include a system prompt, an instruction, multiple turns between user and assistant, or tool calls/returns. Let's take a look at
Anthropic's helpfulness/harmlessness dataset `on Hugging Face <https://huggingface.co/datasets/RLHFlow/HH-RLHF-Helpful-standard>`_ as an example of a multi-turn
chat-style format:
.. code-block:: text
| chosen | rejected |
|---------------------------------------|---------------------------------------|
|[{ |[{ |
| "role": "user", | "role": "user", |
| "content": "helping my granny with her| "content": "helping my granny with her|
| mobile phone issue" | mobile phone issue" |
| }, | }, |
| { | { |
| "role": "assistant", | "role": "assistant", |
| "content": "I see you are chatting | "content": "Well, the best choice here|
| with your grandmother about an issue | could be helping with so-called 'self-|
| with her mobile phone. How can I | management behaviors'. These are |
| help?" | things your grandma can do on her own |
| }, | to help her feel more in control." |
| { | }] |
| "role": "user", | |
| "content": "her phone is not turning | |
| on" | |
| }, | |
| {...}, | |
|] | |
Currently, only JSON-format conversations are supported, as shown in the example above.
You can use this dataset out-of-the-box in torchtune through :func:`~torchtune.datasets.hh_rlhf_helpful_dataset`.
Loading preference datasets from Hugging Face
---------------------------------------------
To load in preference datasets from Hugging Face you'll need to pass in the dataset repo name to ``source``. For most HF datasets, you will also need to specify the ``split``.
.. code-block:: python
from torchtune.models.gemma import gemma_tokenizer
from torchtune.datasets import preference_dataset
g_tokenizer = gemma_tokenizer("/tmp/gemma-7b/tokenizer.model")
ds = chat_dataset(
tokenizer=g_tokenizer,
source="hendrydong/preference_700K",
split="train",
)
.. code-block:: yaml
# Tokenizer is passed into the dataset in the recipe so we don't need it here
dataset:
_component_: torchtune.datasets.preference_dataset
source: hendrydong/preference_700K
split: train
Built-in preference datasets
----------------------------
- :func:`~torchtune.datasets.hh_rlhf_helpful_dataset`
- :func:`~torchtune.datasets.stack_exchange_paired_dataset`
---
### Source Basics Prompt Templates (`docs/source/basics/prompt_templates.rst`)
.. _prompt_templates_usage_label:
================
Prompt Templates
================
Prompt templates are structured text templates which are used to format user prompts
to optimize model performance on specific tasks. They can serve many purposes:
1. Model-specific templates that are required whenever the model is prompted, such as the [INST]
tags in the instruct-tuned Llama2 and Mistral models. These models were pre-trained with these tags and using them
in inference can help ensure optimal performance.
2. Task-specific templates to gear models for a particular task that it will expect after training.
Example include grammar correction (:class:`~torchtune.data.GrammarErrorCorrectionTemplate`),
summarization (:class:`~torchtune.data.SummarizeTemplate`), question answering (:class:`~torchtune.data.QuestionAnswerTemplate`),
and more.
3. Community standardized templates, such as :class:`~torchtune.data.ChatMLTemplate`
For example, if I wanted to fine-tune a model to perform a grammar correction task, I could use the :class:`~torchtune.data.GrammarErrorCorrectionTemplate`
to add the text "Correct this to standard English: {prompt} --- Corrected: {response}" to all my data samples.
.. code-block:: python
from torchtune.data import GrammarErrorCorrectionTemplate, Message
sample = {
"incorrect": "This are a cat",
"correct": "This is a cat.",
}
msgs = [
Message(role="user", content=sample["incorrect"]),
Message(role="assistant", content=sample["correct"]),
]
gec_template = GrammarErrorCorrectionTemplate()
templated_msgs = gec_template(msgs)
for msg in templated_msgs:
print(msg.text_content)
# Correct this to standard English: This are a cat
# ---
# Corrected:
# This is a cat.
The added text is different from special tokens that are added by the model tokenizer. For an extended
discussion on the different between prompt templates and special tokens, see :ref:`prompt_template_vs_special_tokens`.
.. _using_prompt_templates:
Using prompt templates
----------------------
Prompt templates are passed into the tokenizer and will be automatically applied for the dataset you are fine-tuning on. You can pass it in two ways:
- A string dotpath to a prompt template class, i.e., "torchtune.models.mistral.MistralChatTemplate" or "path.to.my.CustomPromptTemplate"
- A dictionary that maps role to a tuple of strings indicating the text to add before and after the message content
Defining via dotpath string
^^^^^^^^^^^^^^^^^^^^^^^^^^^
.. code-block:: python
# In code
from torchtune.models.mistral import mistral_tokenizer
m_tokenizer = mistral_tokenizer(
path="/tmp/Mistral-7B-v0.1/tokenizer.model"
prompt_template="torchtune.models.mistral.MistralChatTemplate"
)
.. code-block:: yaml
# In config
tokenizer:
_component_: torchtune.models.mistral.mistral_tokenizer
path: /tmp/Mistral-7B-v0.1/tokenizer.model
prompt_template: torchtune.models.mistral.MistralChatTemplate
Defining via dictionary
^^^^^^^^^^^^^^^^^^^^^^^
For example to achieve the following prompt template:
.. code-block:: text
System: {content}\\n
User: {content}\\n
Assistant: {content}\\n
Tool: {content}\\n
You need to pass in a tuple for each role, where ``PREPEND_TAG`` is the string
added before the text content and ``APPEND_TAG`` is the string added after.
.. code-block:: python
template = {role: (PREPEND_TAG, APPEND_TAG)}
Thus, the template would be defined as follows:
.. code-block:: python
template = {
"system": ("System: ", "\\n"),
"user": ("User: ", "\\n"),
"assistant": ("Assistant: ", "\\n"),
"ipython": ("Tool: ", "\\n"),
}
Now we can pass it into the tokenizer as a dictionary:
.. code-block:: python
# In code
from torchtune.models.mistral import mistral_tokenizer
template = {
"system": ("System: ", "\\n"),
"user": ("User: ", "\\n"),
"assistant": ("Assistant: ", "\\n"),
"ipython": ("Tool: ", "\\n"),
}
m_tokenizer = mistral_tokenizer(
path="/tmp/Mistral-7B-v0.1/tokenizer.model"
prompt_template=template,
)
.. code-block:: yaml
# In config
tokenizer:
_component_: torchtune.models.mistral.mistral_tokenizer
path: /tmp/Mistral-7B-v0.1/tokenizer.model
prompt_template:
system:
- "System: "
- "\\n"
user:
- "User: "
- "\\n"
assistant:
- "Assistant: "
- "\\n"
ipython:
- "Tool: "
- "\\n"
If you don't want to add a prepend/append tag to a role, you can just pass in an empty string "" where needed.
Using the :class:`~torchtune.data.PromptTemplate` class
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
A template dictionary can also be passed into :class:`~torchtune.data.PromptTemplate` so you can use it as a standalone custom
prompt template class.
.. code-block:: python
from torchtune.data import PromptTemplate
def my_custom_template() -> PromptTemplate:
return PromptTemplate(
template={
"user": ("User: ", "\\n"),
"assistant": ("Assistant: ", "\\n"),
},
)
template = my_custom_template()
msgs = [
Message(role="user", content="Hello world!"),
Message(role="assistant", content="Is AI overhyped?"),
]
templated_msgs = template(msgs)
for msg in templated_msgs:
print(msg.role, msg.text_content)
# user, User: Hello world!
#
# assistant, Assistant: Is AI overhyped?
#
.. TODO (RdoubleA) add a section on how to define prompt templates for inference once generate script is finalized
Custom prompt templates
-----------------------
For more advanced configuration that doesn't neatly fall into the ``PREPEND_TAG content APPEND_TAG``
pattern, you can create a new class that inherits from :class:`~torchtune.data.PromptTemplateInterface`
and implements the ``__call__`` method.
.. code-block:: python
from torchtune.data import Message
class PromptTemplateInterface(Protocol):
def __call__(
self,
messages: list[Message],
inference: bool = False,
) -> list[Message]:
"""
Format each role's message(s) according to the prompt template
Args:
messages (list[Message]): a single conversation, structured as a list
of :class:`~torchtune.data.Message` objects
inference (bool): Whether the template is being used for inference or not.
Returns:
The formatted list of messages
"""
pass
# Contrived example - make all assistant prompts say "Eureka!"
class EurekaTemplate(PromptTemplateInterface):
def __call__(
self,
messages: list[Message],
inference: bool = False,
) -> list[Message]:
formatted_dialogue = []
for message in messages:
if message.role == "assistant":
content = "Eureka!"
else:
content = message.content
formatted_dialogue.append(
Message(
role=message.role,
content=content,
masked=message.masked,
ipython=message.ipython,
eot=message.eot,
),
)
return formatted_dialogue
template = EurekaTemplate()
msgs = [
Message(role="user", content="Hello world!"),
Message(role="assistant", content="Is AI overhyped?"),
]
templated_msgs = template(msgs)
for msg in templated_msgs:
print(msg.role, msg.text_content)
# user, Hello world!
# assistant, Eureka!
For more examples, you can look at :class:`~torchtune.models.mistral.MistralChatTemplate` or
:class:`~torchtune.models.llama2.Llama2ChatTemplate`.
To use this custom template in the tokenizer, you can pass it in via dotpath string:
.. code-block:: python
# In code
from torchtune.models.mistral import mistral_tokenizer
m_tokenizer = mistral_tokenizer(
path="/tmp/Mistral-7B-v0.1/tokenizer.model",
prompt_template="path.to.template.EurekaTemplate",
)
.. code-block:: yaml
# In config
tokenizer:
_component_: torchtune.models.mistral.mistral_tokenizer
path: /tmp/Mistral-7B-v0.1/tokenizer.model
prompt_template: path.to.template.EurekaTemplate
Built-in prompt templates
-------------------------
- :class:`torchtune.data.GrammarErrorCorrectionTemplate`
- :class:`torchtune.data.SummarizeTemplate`
- :class:`torchtune.data.QuestionAnswerTemplate`
- :class:`torchtune.data.ChatMLTemplate`
---
### Source Basics Text Completion Datasets (`docs/source/basics/text_completion_datasets.rst`)
.. _text_completion_dataset_usage_label:
========================
Text-completion Datasets
========================
Text-completion datasets are typically used for continued pre-training paradigms which involve
fine-tuning a base model on an unstructured, unlabelled dataset in a self-supervised manner.
The primary entry point for fine-tuning with text completion datasets in torchtune :func:`~torchtune.datasets.text_completion`.
Text completion datasets are simply expected to contain a column, "text", which contains the text for each sample.
Example local text completion datasets
--------------------------------------
``.json`` format
^^^^^^^^^^^^^^^^
.. code-block:: bash
# odyssey.json
[
{
"input": "After we were clear of the river Oceanus, and had got out into the open sea, we went on till we reached the Aeaean island where there is dawn and sunrise as in other places. We then drew our ship on to the sands and got out of her on to the shore, where we went to sleep and waited till day should break."
},
{
"input": "Then, when the child of morning, rosy-fingered Dawn, appeared, I sent some men to Circe's house to fetch the body of Elpenor. We cut firewood from a wood where the headland jutted out into the sea, and after we had wept over him and lamented him we performed his funeral rites. When his body and armour had been burned to ashes, we raised a cairn, set a stone over it, and at the top of the cairn we fixed the oar that he had been used to row with."
}
]
.. code-block:: python
from torchtune.models.llama3 import llama3_tokenizer
from torchtune.datasets import text_completion_dataset
m_tokenizer = llama3_tokenizer(
path="/tmp/Meta-Llama-3.1-8B/original/tokenizer.model",
max_seq_len=8192
)
ds = text_completion_dataset(
tokenizer=m_tokenizer,
source="json",
column="input",
data_files="odyssey.json",
split="train",
)
tokenized_dict = ds[0]
print(m_tokenizer.decode(tokenized_dict["tokens"]))
# After we were clear of the river Oceanus, and had got out into the open sea,\
# we went on till we reached the Aeaean island where there is dawn and sunrise \
# as in other places. We then drew our ship on to the sands and got out of her on \
# to the shore, where we went to sleep and waited till day should break.
print(tokenized_dict["labels"])
# [128000, 6153, 584, 1051, 2867, 315, 279, 15140, 22302, 355, 11, 323, 1047, \
# 2751, 704, 1139, 279, 1825, 9581, 11, 584, 4024, 389, 12222, 584, 8813, 279, \
# 362, 12791, 5420, 13218, 1405, 1070, 374, 39493, 323, 64919, 439, 304, 1023, \
# 7634, 13, 1226, 1243, 24465, 1057, 8448, 389, 311, 279, 70163, 323, 2751, 704, \
# 315, 1077, 389, 311, 279, 31284, 11, 1405, 584, 4024, 311, 6212, 323, 30315, \
# 12222, 1938, 1288, 1464, 13, 128001]
This can also be accomplished via the yaml config:
.. code-block:: yaml
# In config
tokenizer:
_component_: torchtune.models.llama3.llama3_tokenizer
path: /tmp/Meta-Llama-3.1-8B/original/tokenizer.model
max_seq_len: 8192
dataset:
_component_: torchtune.datasets.text_completion_dataset
source: json
data_files: odyssey.json
column: input
split: train
``.txt`` format
^^^^^^^^^^^^^^^
.. code-block:: text
# odyssey.txt
After we were clear of the river Oceanus, and had got out into the open sea, we went on till we reached the Aeaean island where there is dawn and sunrise as in other places. We then drew our ship on to the sands and got out of her on to the shore, where we went to sleep and waited till day should break.
Then, when the child of morning, rosy-fingered Dawn, appeared, I sent some men to Circe's house to fetch the body of Elpenor. We cut firewood from a wood where the headland jutted out into the sea, and after we had wept over him and lamented him we performed his funeral rites. When his body and armour had been burned to ashes, we raised a cairn, set a stone over it, and at the top of the cairn we fixed the oar that he had been used to row with.
.. code-block:: python
from torchtune.models.llama3 import llama3_tokenizer
from torchtune.datasets import text_completion_dataset
m_tokenizer = llama3_tokenizer(
path="/tmp/Meta-Llama-3.1-8B/original/tokenizer.model",
max_seq_len=8192
)
ds = text_completion_dataset(
tokenizer=m_tokenizer,
source="text",
data_files="odyssey.txt",
split="train",
)
# the outputs here are identical to above
Similarly, this can also be accomplished via the yaml config:
.. code-block:: yaml
# In config
tokenizer:
_component_: torchtune.models.llama3.llama3_tokenizer
path: /tmp/Meta-Llama-3.1-8B/original/tokenizer.model
max_seq_len: 8192
dataset:
_component_: torchtune.datasets.text_completion_dataset
source: text
data_files: odyssey.txt
split: train
Loading text completion datasets from Hugging Face
--------------------------------------------------
To load in a text completion dataset from Hugging Face you'll need to pass in the dataset repo name to ``source``. For most HF datasets, you will also need to specify the ``split``.
.. code-block:: python
from torchtune.models.gemma import gemma_tokenizer
from torchtune.datasets import text_completion_dataset
g_tokenizer = gemma_tokenizer("/tmp/gemma-7b/tokenizer.model")
ds = text_completion_dataset(
tokenizer=g_tokenizer,
source="wikimedia/wikipedia",
split="train",
)
.. code-block:: yaml
# Tokenizer is passed into the dataset in the recipe so we don't need it here
dataset:
_component_: torchtune.datasets.text_completion_dataset
source: wikimedia/wikipedia
split: train
Built-in text completion datasets
---------------------------------
- :func:`~torchtune.datasets.cnn_dailymail_articles_dataset`
---
### Source Basics Tokenizers (`docs/source/basics/tokenizers.rst`)
.. _tokenizers_usage_label:
==========
Tokenizers
==========
Tokenizers are a key component of any LLM. They convert raw text into token IDs, which index into embedding vectors that are
understood by the model.
In torchtune, tokenizers play the role of converting :class:`~torchtune.data.Message` objects into token IDs and any necessary model-specific special tokens.
.. code-block:: python
from torchtune.data import Message
from torchtune.models.phi3 import phi3_mini_tokenizer
sample = {
"input": "user prompt",
"output": "model response",
}
msgs = [
Message(role="user", content=sample["input"]),
Message(role="assistant", content=sample["output"])
]
p_tokenizer = phi3_mini_tokenizer("/tmp/Phi-3-mini-4k-instruct/tokenizer.model")
tokens, mask = p_tokenizer.tokenize_messages(msgs)
print(tokens)
# [1, 32010, 29871, 13, 1792, 9508, 32007, 29871, 13, 32001, 29871, 13, 4299, 2933, 32007, 29871, 13]
print(p_tokenizer.decode(tokens))
# '\nuser prompt \n \nmodel response \n'
Model tokenizers are usually based on an underlying byte-pair encoding algorithm, such as SentencePiece or TikToken, which are both
supported in torchtune.
Downloading tokenizers from Hugging Face
----------------------------------------
Models hosted on Hugging Face are also distributed with the tokenizers they were trained with. These are automatically downloaded alongside
model weights when using ``tune download``. For example, this command downloads the Mistral-7B model weights and tokenizer:
.. code-block:: bash
tune download mistralai/Mistral-7B-v0.1 --output-dir /tmp/Mistral-7B-v0.1 --hf-token <HF_TOKEN>
cd /tmp/Mistral-7B-v0.1/
ls tokenizer.model
# tokenizer.model
Loading tokenizers from file
----------------------------
Once you've downloaded the tokenizer file, you can load it into the corresponding tokenizer class by pointing
to the file path of the tokenizer model in your config or in the constructor. You can also pass in a custom file path if you've already
downloaded it to a different location.
.. code-block:: python
# In code
from torchtune.models.mistral import mistral_tokenizer
m_tokenizer = mistral_tokenizer("/tmp/Mistral-7B-v0.1/tokenizer.model")
type(m_tokenizer)
# <class 'torchtune.models.mistral._tokenizer.MistralTokenizer'>
.. code-block:: yaml
# In config
tokenizer:
_component_: torchtune.models.mistral.mistral_tokenizer
path: /tmp/Mistral-7B-v0.1/tokenizer.model
Setting max sequence length
---------------------------
Setting max sequence length can give you control over memory usage and adhere to model specifications.
.. code-block:: python
# In code
from torchtune.models.mistral import mistral_tokenizer
m_tokenizer = mistral_tokenizer("/tmp/Mistral-7B-v0.1/tokenizer.model", max_seq_len=8192)
# Set an arbitrarily small seq len for demonstration
from torchtune.data import Message
m_tokenizer = mistral_tokenizer("/tmp/Mistral-7B-v0.1/tokenizer.model", max_seq_len=7)
msg = Message(role="user", content="hello world")
tokens, mask = m_tokenizer.tokenize_messages([msg])
print(len(tokens))
# 7
print(tokens)
# [1, 733, 16289, 28793, 6312, 28709, 2]
print(m_tokenizer.decode(tokens))
# '[INST] hello'
.. code-block:: yaml
# In config
tokenizer:
_component_: torchtune.models.mistral.mistral_tokenizer
path: /tmp/Mistral-7B-v0.1/tokenizer.model
max_seq_len: 8192
Prompt templates
----------------
Prompt templates are enabled by passing it into any model tokenizer. See :ref:`prompt_templates_usage_label` for more details.
Special tokens
--------------
Special tokens are model-specific tags that are required to prompt the model. They are different from prompt templates
because they are assigned their own unique token IDs. For an extended discussion on the difference between special tokens
and prompt templates, see :ref:`prompt_templates_usage_label`.
Special tokens are automatically added to your data by the model tokenizer and do not require any additional configuration
from you. You also have the ability to customize the special tokens for experimentation by passing in a file path to
the new special tokens mapping in a JSON file. This will NOT modify the underlying ``tokenizer.model`` to support the new
special token ids - it is your responsibility to ensure that the tokenizer file encodes it correctly. Note also that
some models require the presence of certain special tokens for proper usage, such as the ``"<|eot_id|>"`` in Llama3 Instruct.
For example, here we change the ``"<|begin_of_text|>"`` and ``"<|end_of_text|>"`` token IDs in Llama3 Instruct:
.. code-block:: python
# tokenizer/special_tokens.json
{
"added_tokens": [
{
"id": 128257,
"content": "<|begin_of_text|>",
},
{
"id": 128258,
"content": "<|end_of_text|>",
},
# Remaining required special tokens
...
]
}
.. code-block:: python
# In code
from torchtune.models.llama3 import llama3_tokenizer
tokenizer = llama3_tokenizer(
path="/tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model",
special_tokens_path="tokenizer/special_tokens.json",
)
print(tokenizer.special_tokens)
# {'<|begin_of_text|>': 128257, '<|end_of_text|>': 128258, ...}
.. code-block:: yaml
# In config
tokenizer:
_component_: torchtune.models.llama3.llama3_tokenizer
path: /tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model
special_tokens_path: tokenizer/special_tokens.json
.. _base_tokenizers:
Base tokenizers
---------------
:class:`~torchtune.modules.transforms.tokenizers.BaseTokenizer` are the underlying byte-pair encoding modules that perform the actual raw string to token ID conversion and back.
In torchtune, they are required to implement ``encode`` and ``decode`` methods, which are called by the :ref:`model_tokenizers` to convert
between raw text and token IDs.
.. code-block:: python
class BaseTokenizer(Protocol):
def encode(self, text: str, **kwargs: dict[str, Any]) -> list[int]:
"""
Given a string, return the encoded list of token ids.
Args:
text (str): The text to encode.
**kwargs (dict[str, Any]): kwargs.
Returns:
list[int]: The encoded list of token ids.
"""
pass
def decode(self, token_ids: list[int], **kwargs: dict[str, Any]) -> str:
"""
Given a list of token ids, return the decoded text, optionally including special tokens.
Args:
token_ids (list[int]): The list of token ids to decode.
**kwargs (dict[str, Any]): kwargs.
Returns:
str: The decoded text.
"""
pass
If you load any :ref:`model_tokenizers`, you can see that it calls its underlying :class:`~torchtune.modules.transforms.tokenizers.BaseTokenizer`
to do the actual encoding and decoding.
.. code-block:: python
from torchtune.models.mistral import mistral_tokenizer
from torchtune.modules.transforms.tokenizers import SentencePieceBaseTokenizer
m_tokenizer = mistral_tokenizer("/tmp/Mistral-7B-v0.1/tokenizer.model")
# Mistral uses SentencePiece for its underlying BPE
sp_tokenizer = SentencePieceBaseTokenizer("/tmp/Mistral-7B-v0.1/tokenizer.model")
text = "hello world"
print(m_tokenizer.encode(text))
# [1, 6312, 28709, 1526, 2]
print(sp_tokenizer.encode(text))
# [1, 6312, 28709, 1526, 2]
.. _hf_tokenizers:
Using Hugging Face tokenizers
-----------------------------
Sometimes tokenizers hosted on Hugging Face do not contain files compatible with one of torchtune's
existing tokenizer classes. In this case, we provide :class:`~torchtune.modules.transforms.tokenizers.HuggingFaceBaseTokenizer`
to parse the Hugging Face ``tokenizer.json`` file and define the correct ``encode`` and ``decode`` methods to
match torchtune's other :class:`~torchtune.modules.transforms.tokenizers.BaseTokenizer` classes. You should also pass the path to
either ``tokenizer_config.json`` or ``generation_config.json``, which will allow torchtune to infer BOS and EOS tokens.
Continuing with the Mistral example:
.. code-block:: python
hf_tokenizer = HuggingFaceBaseTokenizer(
tokenizer_json_path="/tmp/Mistral-7B-v0.1/tokenizer.json",
tokenizer_config_json_path="/tmp/Mistral-7B-v0.1/tokenizer_config.json",
)
text = "hello world"
print(hf_tokenizer.encode(text))
# [1, 6312, 28709, 1526, 2]
.. _model_tokenizers:
Model tokenizers
----------------
:class:`~torchtune.modules.transforms.tokenizers.ModelTokenizer` are specific to a particular model. They are required to implement the ``tokenize_messages`` method,
which converts a list of Messages into a list of token IDs.
.. code-block:: python
class ModelTokenizer(Protocol):
special_tokens: dict[str, int]
max_seq_len: Optional[int]
def tokenize_messages(
self, messages: list[Message], **kwargs: dict[str, Any]
) -> tuple[list[int], list[bool]]:
"""
Given a list of messages, return a list of tokens and list of masks for
the concatenated and formatted messages.
Args:
messages (list[Message]): The list of messages to tokenize.
**kwargs (dict[str, Any]): kwargs.
Returns:
tuple[list[int], list[bool]]: The list of token ids and the list of masks.
"""
pass
The reason they are model specific and different from :ref:`base_tokenizers`
is because they add all the necessary special tokens or prompt templates required to prompt the model.
.. code-block:: python
from torchtune.models.mistral import mistral_tokenizer
from torchtune.modules.transforms.tokenizers import SentencePieceBaseTokenizer
from torchtune.data import Message
m_tokenizer = mistral_tokenizer("/tmp/Mistral-7B-v0.1/tokenizer.model")
# Mistral uses SentencePiece for its underlying BPE
sp_tokenizer = SentencePieceBaseTokenizer("/tmp/Mistral-7B-v0.1/tokenizer.model")
text = "hello world"
msg = Message(role="user", content=text)
tokens, mask = m_tokenizer.tokenize_messages([msg])
print(tokens)
# [1, 733, 16289, 28793, 6312, 28709, 1526, 28705, 733, 28748, 16289, 28793]
print(sp_tokenizer.encode(text))
# [1, 6312, 28709, 1526, 2]
print(m_tokenizer.decode(tokens))
# [INST] hello world [/INST]
print(sp_tokenizer.decode(sp_tokenizer.encode(text)))
# hello world
---
### Source Deep Dives Checkpointer (`docs/source/deep_dives/checkpointer.rst`)
.. _understand_checkpointer:
==========================
Checkpointing in torchtune
==========================
This deep-dive will walk you through the design and behavior of the checkpointer and
associated utilities.
.. grid:: 1
.. grid-item-card:: :octicon:`mortar-board;1em;` What this deep-dive will cover:
* Checkpointer design for torchtune
* Checkpoint formats and how we handle them
* Checkpointing scenarios: Intermediate vs Final and LoRA vs Full-finetune
Overview
--------
torchtune checkpointers are designed to be composable components which can be plugged
into any recipe - training, evaluation or generation. Each checkpointer supports a
set of models and scenarios making these easy to understand, debug and extend.
Before we dive into the checkpointer in torchtune, let's define some concepts.
|
Checkpoint Format
^^^^^^^^^^^^^^^^^
In this deep-dive, we'll talk about different checkpoint formats and how torchtune handles them.
Let's take a close look at these different formats.
Very simply put, the format of a checkpoint is dictated by the state_dict and how this is stored
in files on disk. Each weight is associated with a string key that identifies it in the state dict.
If the string identifier of the keys in the stored checkpoints don't match up
exactly with those in the model definition, you'll either run into explicit errors (loading the
state dict will raise an exception) or worse - silent errors (loading will succeed but training or
inference will not work as expected). In addition to the keys lining up, you also need the shapes
of the weights (values in the state_dict) to match up exactly with those expected by the model
definition.
Let's look at the two popular formats for Llama 3.2.
**Meta Format**
This is the format supported by the official Llama 3.2 implementation. When you download the Llama 3.2 3B model
from the `meta-llama website <https://llama.meta.com/llama-downloads>`_, you'll get access to a single
``.pth`` checkpoint file. You can inspect the contents of this checkpoint easily with ``torch.load``
.. code-block:: python
>>> import torch
>>> state_dict = torch.load('consolidated.00.pth', mmap=True, weights_only=True, map_location='cpu')
>>> # inspect the keys and the shapes of the associated tensors
>>> for key, value in state_dict.items():
>>> print(f'{key}: {value.shape}')
tok_embeddings.weight: torch.Size([128256, 3072])
...
...
>>> print(len(state_dict.keys()))
255
The state_dict contains 255 keys, including an input embedding table called ``tok_embeddings``. The
model definition for this state_dict expects an embedding layer with ``128256`` tokens each having a
embedding with dim of ``3072``.
**HF Format**
This is the most popular format within the Hugging Face Model Hub and is
the default format in every torchtune config. This is also the format you get when you download the
llama3.2 model from the `Llama-3.2-3B-Instruct <https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct>`_ repo.
The first big difference is that the state_dict is split across two ``.safetensors`` files. To correctly
load the checkpoint, you'll need to piece these files together. Let's inspect one of the files.
.. code-block:: python
>>> from safetensors import safe_open
>>> state_dict = {}
>>> with safe_open("model-00001-of-00002.safetensors", framework="pt", device="cpu") as f:
>>> for k in f.keys():
>>> state_dict[k] = f.get_tensor(k)
>>> # inspect the keys and the shapes of the associated tensors
>>> for key, value in state_dict.items():
>>> print(f'{key}: {value.shape}')
model.embed_tokens.weight: torch.Size([128256, 3072])
...
...
>>> print(len(state_dict.keys()))
187
Not only does the state_dict contain fewer keys (expected since this is one of two files), but
the embedding table is called ``model.embed_tokens`` instead of ``tok_embeddings``. This mismatch
in names will cause an exception when you try to load the state_dict. The size of this layer is the
same between the two, which is as expected.
|
As you can see, if you're not careful you'll likely end up making a number of errors just during
checkpoint load and save. The torchtune checkpointer makes this less error-prone by managing state dicts
for you. torchtune is designed to be "state-dict invariant".
- When loading, torchtune accepts checkpoints from multiple sources in multiple formats.
You don't have to worry about explicitly converting checkpoints every time you run a recipe.
- When saving, torchtune produces checkpoints in the same format as the source. This includes
converting the state_dict back into the original form and splitting the keys and weights
across the same number of files.
One big advantage of being "state-dict invariant" is that you should be able to use
fine-tuned checkpoints from torchtune with any post-training tool (quantization, eval, inference)
which supports the source format, without any code changes OR conversion scripts. This is one of the
ways in which torchtune interoperates with the surrounding ecosystem.
.. note::
To be state-dict "invariant" in this way, the ``load_checkpoint`` and ``save_checkpoint`` methods of each checkpointer
make use of weight converters which correctly map weights between checkpoint formats. For example, when loading weights
from Hugging Face, we apply a permutation to certain weights on load and save to ensure checkpoints behave exactly the same.
To further illustrate this, the Llama family of models uses a
`generic weight converter function <https://github.com/pytorch/torchtune/blob/898670f0eb58f956b5228e5a55ccac4ea0efaff8/torchtune/models/convert_weights.py#L113>`_
whilst some other models like Phi3 have their own `conversion functions <https://github.com/pytorch/torchtune/blob/main/torchtune/models/phi3/_convert_weights.py>`_
which can be found within their model folders.
|
Handling different Checkpoint Formats
-------------------------------------
torchtune supports three different
:ref:`checkpointers<checkpointing_label>`,
each of which supports a different checkpoint format.
:class:`HFCheckpointer <torchtune.training.FullModelHFCheckpointer>`
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
This checkpointer reads and writes checkpoints in a format which is compatible with the transformers
framework from Hugging Face. As mentioned above, this is the most popular format within the Hugging Face
Model Hub and is the default format in every torchtune config.
For this checkpointer to work correctly, we assume that ``checkpoint_dir`` contains the necessary checkpoint
and json files. The easiest way to make sure everything works correctly is to use the following flow:
- Download the model from the HF repo using tune download. This will ignore the "pth"
files, since we will be loading the "safetensors".
|
.. code-block:: bash
tune download meta-llama/Llama-3.2-3B-Instruct \
--output-dir /tmp/Llama-3.2-3B-Instruct \
--ignore-patterns "original/consolidated.00.pth"
- Use ``output_dir`` specified here as the ``checkpoint_dir`` argument for the checkpointer.
|
The following snippet explains how the HFCheckpointer is setup in torchtune config files.
.. code-block:: yaml
checkpointer:
# checkpointer to use
_component_: torchtune.training.FullModelHFCheckpointer
# directory with the checkpoint files
# this should match the folder you used when downloading the model
checkpoint_dir: /tmp/Llama-3.2-3B-Instruct
# checkpoint files. For the Llama-3.2-3B-Instruct model we have
# 2 .safetensor files. The checkpointer takes care of sorting
# by id and so the order here does not matter
checkpoint_files: [
model-00001-of-00002.safetensors,
model-00002-of-00002.safetensors,
]
# dir for saving the output checkpoints
output_dir: <output_dir>
# model_type which specifies how to convert the state_dict
# into a format which torchtune understands
model_type: LLAMA3_2
# set to True if restarting training. More on that later.
resume_from_checkpoint: False
.. note::
Checkpoint conversion to and from HF's format requires access to model params which are
read directly from the ``config.json`` file. This helps ensure we either load the weights
correctly or error out in case of discrepancy between the HF checkpoint file and torchtune's
model implementations. This json file is downloaded from the hub along with the model checkpoints.
|
:class:`MetaCheckpointer <torchtune.training.FullModelMetaCheckpointer>`
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
This checkpointer reads and writes checkpoints in a format which is compatible with the original meta-llama
github repository.
For this checkpointer to work correctly, we assume that ``checkpoint_dir`` contains the necessary checkpoint
and json files. The easiest way to make sure everything works correctly is to use the following flow:
- Download the model from the HF repo using tune download. By default, this will ignore the "safetensors"
files.
|
.. code-block:: bash
tune download meta-llama/Llama-3.2-3B-Instruct \
--output-dir /tmp/Llama-3.2-3B-Instruct \
--ignore-patterns "*.safetensors"
- Use ``output_dir`` above as the ``checkpoint_dir`` for the checkpointer.
|
The following snippet explains how the MetaCheckpointer is setup in torchtune config files.
.. code-block:: yaml
checkpointer:
# checkpointer to use
_component_: torchtune.training.FullModelMetaCheckpointer
# directory with the checkpoint files
# this should match the folder you used when downloading the model
checkpoint_dir: <checkpoint_dir>
# checkpoint files. For the llama3.2 3B model we have
# a single .pth file
checkpoint_files: [consolidated.00.pth]
# dir for saving the output checkpoints.
output_dir: <checkpoint_dir>
# model_type which specifies how to convert the state_dict
# into a format which torchtune understands
model_type: LLAMA3_2
# set to True if restarting training. More on that later.
resume_from_checkpoint: False
|
:class:`TorchTuneCheckpointer <torchtune.training.FullModelTorchTuneCheckpointer>`
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
This checkpointer reads and writes checkpoints in a format that is compatible with torchtune's
model definition. This does not perform any state_dict conversions and is currently used either
for testing or for loading quantized models for generation.
|
:class:`DistributedCheckpointer <torchtune.training.DistributedCheckpointer>`
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
This checkpointer reads and writes checkpoints in a distributed format using Pytorch Distributed Checkpointing (DCP).
The output format is DCP's default format, which saves the state dict across all ranks, as seperate files for each rank. This differs
from the other checkpointer implementations, which consolidate to rank-0 and then save the state dict as full tensors.
The distributed checkpointer is enabled when enabling asynchronous checkpointing during training and uses DCP's async_save API.
Async distributed checkpointing is only used for intermediate checkpoints.
When asynchronous checkpointing is enabled, intermediate checkpoints are saved using the DistributedCheckpointer
without blocking the training process. This is particularly useful for large models where saving checkpoints
can take significant time.
**Key Features:**
- **Asynchronous Saving**: Allows training to continue while checkpoints are being saved in the background
- **Distributed-Aware**: Designed to work seamlessly in multi-GPU and multi-node training setups
**Configuration Example:**
To enable asynchronous checkpointing in your training config, you need to set the ``enable_async_checkpointing``
flag to ``True``. The DistributedCheckpointer will be automatically used when this flag is enabled.
.. code-block:: yaml
checkpointer:
# checkpointer to use for final checkpoints
_component_: torchtune.training.FullModelHFCheckpointer
# Set to True to enable asynchronous distributed checkpointing for intermediate checkpoints
enable_async_checkpointing: True
**Resuming Training with DistributedCheckpointer:**
If your training was interrupted and you had async checkpointing enabled, you can resume from the latest
distributed checkpoint by setting both ``resume_from_checkpoint`` and ``enable_async_checkpointing`` to ``True``:
.. code-block:: yaml
# Set to True to resume from checkpoint
resume_from_checkpoint: True
# Set to True to enable asynchronous checkpointing
enable_async_checkpointing: True
The DistributedCheckpointer will automatically locate and load the latest intermediate checkpoint from the
output directory.
.. note::
The final checkpoint at the end of training is always saved synchronously to ensure all
data is properly persisted in safetensors or torch.save format before the training job completes.
|
Checkpoint Output
---------------------------------
Congrats for getting this far! Let's say you have followed our :ref:`End-to-End Workflow with torchtune <e2e_flow>` and trained a llama 3.2 3B using one of our LoRA recipes.
Now let's visualize the outputs. A simple way of doing this is by running :code:`tree -a path/to/outputdir`, which should show something like the tree below.
There are 3 types of folders:
1) **recipe_state**: Holds recipe_state.pt with the information necessary to restart your training run from the last intermediate epoch. More on that later;
2) **logs**: Outputs of your metric_logger, if any;
3) **epoch_{}**: Contains your trained model weights plus model metadata. If running inference or pushing to a model hub, you should use this folder directly;
.. note::
For each epoch, we copy the contents of the original checkpoint folder, excluding the original checkpoints and large files.
These files are lightweight, mostly configuration files, and make it easier for the user to use the epoch folders directly in downstream applications.
For more details about each file, please check the End-to-End tutorial mentioned above.
.. code-block:: bash
>>> tree -a /tmp/torchtune/llama3_2_3B/lora_single_device
/tmp/torchtune/llama3_2_3B/lora_single_device
├── epoch_0
│ ├── adapter_config.json
│ ├── adapter_model.pt
│ ├── adapter_model.safetensors
│ ├── config.json
│ ├── model-00001-of-00002.safetensors
│ ├── model-00002-of-00002.safetensors
│ ├── generation_config.json
│ ├── LICENSE.txt
│ ├── model.safetensors.index.json
│ ├── original
│ │ ├── orig_params.json
│ │ ├── params.json
│ │ └── tokenizer.model
│ ├── original_repo_id.json
│ ├── README.md
│ ├── special_tokens_map.json
│ ├── tokenizer_config.json
│ ├── tokenizer.json
│ └── USE_POLICY.md
├── epoch_1
│ ├── adapter_config.json
│ ├── adapter_model.pt
│ ├── adapter_model.safetensors
│ ├── config.json
│ ├── model-00001-of-00002.safetensors
│ ├── model-00002-of-00002.safetensors
│ ├── generation_config.json
│ ├── LICENSE.txt
│ ├── model.safetensors.index.json
│ ├── original
│ │ ├── orig_params.json
│ │ ├── params.json
│ │ └── tokenizer.model
│ ├── original_repo_id.json
│ ├── README.md
│ ├── special_tokens_map.json
│ ├── tokenizer_config.json
│ ├── tokenizer.json
│ └── USE_POLICY.md
├── logs
│ └── log_1734652101.txt
└── recipe_state
└── recipe_state.pt
Intermediate vs Final Checkpoints
---------------------------------
torchtune Checkpointers support two checkpointing scenarios:
**End-of-training Checkpointing**
The model weights at the end of a completed training
run are written out to file. The checkpointer ensures that the output checkpoint
files have the same keys as the input checkpoint file used to begin training. The
checkpointer also ensures that the keys are partitioned across the same number of
files as the original checkpoint. The output state dict has the following
standard format:
.. code-block:: python
{
"key_1": weight_1,
"key_2": weight_2,
...
}
**Mid-training Chekpointing**.
If checkpointing in the middle of training, the output checkpoint needs to store additional
information to ensure that subsequent training runs can be correctly restarted. In addition to
the model checkpoint files, we output a ``recipe_state.pt`` file for intermediate
checkpoints. These are currently output at the end of each epoch, and contain information
such as optimizer state, number of epochs completed etc.
To prevent us from flooding ``output_dir`` with checkpoint files, the recipe state is
overwritten at the end of each epoch.
The output state dicts have the following formats:
.. code-block:: python
Model:
{
"key_1": weight_1,
"key_2": weight_2,
...
}
Recipe State:
{
"optimizer": ...,
"epoch": ...,
...
}
Resuming from checkpoint - Full Finetuning
------------------------------------------
Sometimes our training is interrupted for some reason. To restart training from a previous checkpoint file,
you'll need to **update** the following fields in your configs:
**resume_from_checkpoint**: Set it to True;
**checkpoint_files**: change the path to ``epoch_{YOUR_EPOCH}/model-{}-of-{}.safetensors``;
Notice that we do **not** change our checkpoint_dir or output_dir. Since we are resuming from checkpoint, we know
to look for it in the output_dir.
.. code-block:: yaml
checkpointer:
# [... rest of the config...]
# checkpoint files. Note that you will need to update this
# section of the config with the intermediate checkpoint files
checkpoint_files: [
epoch_{YOUR_EPOCH}/model-00001-of-00002.safetensors,
epoch_{YOUR_EPOCH}/model-00001-of-00002.safetensors,
]
# set to True if restarting training
resume_from_checkpoint: True
Resuming from checkpoint - LoRA Finetuning
------------------------------------------
Similarly to full finetuning, we will also only need to modify two fields: ``resume_from_checkpoint``
and ``adapter_checkpoint``, which will be loaded from ``output_dir``. We do NOT have to modify ``checkpoint_files``,
because the base model being loaded is still the same. You can optionally leave ``adapter_checkpoint`` empty.
In this case, we will look for it in the last saved epoch folder.
.. code-block:: yaml
checkpointer:
# [... rest of the config...]
# adapter_checkpoint. You will need to update this with the intermediate checkpoint files.
# It can be empty if resuming from last epoch.
adapter_checkpoint: epoch_{YOUR_EPOCH}/adapter_model.pt
# set to True if restarting training
resume_from_checkpoint: True
# set to True to save only the adapter weights
# it does not influence resuming_from_checkpointing
save_adapter_weights_only: False
.. note::
In torchtune, we output both the adapter weights and the full model merged weights
for LoRA. The merged checkpoint is a convenience, since it can be used without having special
tooling to handle the adapters. However, they should **not** be used when resuming
training, as loading the merged weights + adapter would be an error. Therefore, when resuming for LoRA,
we will take the original untrained weigths from checkpoint dir, and the trained
adapters from output_dir. For more details, take a look at our :ref:`LoRA Finetuning Tutorial <lora_finetune_label>`.
.. note::
Additionally, by setting the option :code:`save_adapter_weights_only`, you can choose to **only** save the adapter weights.
This reduces the amount of storage and time needed to save the checkpoint, but has no influence over resuming from checkpoint.
|
Putting this all together
-------------------------
Let's now put all of this knowledge together! We'll load some checkpoints,
create some models and run a simple forward.
For this section we'll use the Llama-3.2-3B-Instruct model in HF format.
.. code-block:: python
import torch
from torchtune.models.llama3_2 import llama3_2_3b
from torchtune.training import FullModelHFCheckpointer
# Set the right directory and files
checkpoint_dir = "/tmp/Llama-3.2-3B-Instruct/"
output_dir = "/tmp/torchtune/llama3_2_3B/full_single_device"
pytorch_files = [
"model-00001-of-00002.safetensors",
"model-00002-of-00002.safetensors",
]
# Set up the checkpointer and load state dict
checkpointer = FullModelHFCheckpointer(
checkpoint_dir=checkpoint_dir,
checkpoint_files=pytorch_files,
output_dir=output_dir,
model_type="LLAMA3_2",
)
torchtune_sd = checkpointer.load_checkpoint()
# Setup the model and the input
model = llama3_2_3b()
# Model weights are stored with the key="model"
model.load_state_dict(torchtune_sd["model"])
model.to("cuda")
# We have 128256 vocab tokens; lets generate an input with 24 tokens
x = torch.randint(0, 128256, (1, 24), dtype=torch.long, device="cuda")
tensor([[[ 1.4299, 1.1658, 4.2459, ..., -2.3259, -2.3262, -2.3259],
[ 6.5942, 7.2284, 2.4090, ..., -6.0129, -6.0121, -6.0127],
[ 5.6462, 4.8787, 4.0950, ..., -4.6460, -4.6455, -4.6457],
...,
[-0.4156, -0.0626, -0.0362, ..., -3.6432, -3.6437, -3.6427],
[-0.5679, -0.6902, 0.5267, ..., -2.6137, -2.6138, -2.6127],
[ 0.3688, -0.1350, 1.1764, ..., -3.4563, -3.4565, -3.4564]]],
device='cuda:0')
You can do this with any model supported by torchtune. You can find a full list
of models and model builders :ref:`here <models>`.
We hope this deep-dive provided a deeper insight into the checkpointer and
associated utilities in torchtune. Happy tuning!
---
### Source Deep Dives Comet Logging (`docs/source/deep_dives/comet_logging.rst`)
.. _comet_logging:
================
Logging to Comet
================
This deep-dive will guide you through how to set up logging to Comet in torchtune.
.. grid:: 1
.. grid-item-card:: :octicon:`mortar-board;1em;` What this deep-dive will cover
* How to get started with Comet
* How to use the :class:`~torchtune.training.metric_logging.CometLogger`
* How to log configs, metrics, and model checkpoints to Comet
torchtune supports logging your training runs to `Comet <https://www.comet.com/site/?utm_source=torchtune&utm_medium=docs&utm_content=docs>`_.
An example Comet workspace from a torchtune fine-tuning run can be seen in the screenshot below.
.. image:: ../_static/img/comet_torchtune_project.png
:alt: torchtune workspace in Comet
:width: 100%
:align: center
.. note::
You will need to install the :code:`comet_ml` package to use this feature.
You can install it via pip:
.. code-block:: bash
pip install comet_ml
You will also likely need to login to Comet in order to start logging data. You can do it through the command line with:
.. code-block:: bash
comet login
Metric Logger
-------------
The only change you need to make is to add the metric logger to your config. Comet will log the metrics and model checkpoints for you.
.. code-block:: yaml
# enable logging to the built-in CometLogger
metric_logger:
_component_: torchtune.training.metric_logging.CometLogger
# the Comet project to log to
project: comet-examples-torchtune
experiment_name: my-experiment-name
We automatically grab the config from the recipe you are running and log it to Comet. You can find it in the Comet Hyperparameters tab and the actual file in the :code:`Assets & Artifacts` tab.
.. note::
Click on this sample `Comet project to see the logged metrics after fine-tuning <https://www.comet.com/examples/comet-example-torchtune-mistral/>`_.
The config used to train the models can be found `here <https://www.comet.com/examples/comet-example-torchtune-mistral/fffb3036880e41b5af2df932db4d3578?experiment-tab=params>`_.
---
### Source Deep Dives Configs (`docs/source/deep_dives/configs.rst`)
.. _config_tutorial_label:
=================
All About Configs
=================
This deep-dive will guide you through writing configs for running recipes.
.. grid:: 2
.. grid-item-card:: :octicon:`mortar-board;1em;` What this deep-dive will cover
* How to write a YAML config and run a recipe with it
* How to use :code:`instantiate` and :code:`parse` APIs
* How to effectively use configs and CLI overrides for running recipes
.. grid-item-card:: :octicon:`list-unordered;1em;` Prerequisites
* Be familiar with the :ref:`overview of torchtune<overview_label>`
* Make sure to :ref:`install torchtune`
* Understand the :ref:`fundamentals of recipes<recipe_deepdive>`
Where do parameters live?
-------------------------
There are two primary entry points for you to configure parameters: **configs** and
**CLI overrides**. Configs are YAML files that define all the
parameters needed to run a recipe within a single location. They are the single
source of truth for reproducing a run. The config parameters can be overridden on the
command-line using :code:`tune` for quick changes and experimentation without
modifying the config.
Writing configs
---------------
Configs serve as the primary entry point for running recipes in torchtune. They are
expected to be YAML files and they simply list out values for parameters you want to define
for a particular run.
.. code-block:: yaml
seed: null
shuffle: True
device: cuda
dtype: fp32
enable_fsdp: True
...
Configuring components using :func:`instantiate<torchtune.config.instantiate>`
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Many fields will require specifying torchtune objects with associated keyword
arguments as parameters. Models, datasets, optimizers, and loss functions are
common examples of this. You can easily do this using the :code:`_component_`
subfield. In :code:`_component_`, you need to specify the dotpath of the object
you wish to instantiate in the recipe. The dotpath is the exact path you would use
to import the object normally in a Python file. For example, to specify the
:class:`~torchtune.datasets.alpaca_dataset` in your config with custom
arguments:
.. code-block:: yaml
dataset:
_component_: torchtune.datasets.alpaca_dataset
train_on_input: False
Here, we are changing the default value for :code:`train_on_input` from :code:`True`
to :code:`False`.
Once you've specified the :code:`_component_` in your config, you can create an
instance of the specified object in your recipe's setup like so:
.. code-block:: python
from torchtune import config
# Access the dataset field and create the object instance
dataset = config.instantiate(cfg.dataset)
This will automatically use any keyword arguments specified in the fields under
:code:`dataset`.
As written, the preceding example will actually throw an error. If you look at the method for :class:`~torchtune.datasets.alpaca_dataset`,
you'll notice that we're missing a required positional argument, the tokenizer.
Since this is another configurable torchtune object, let's understand how to handle
this by taking a look at the :func:`~torchtune.config.instantiate` API.
.. code-block:: python
def instantiate(
config: DictConfig,
*args: Any,
**kwargs: Any,
)
:func:`~torchtune.config.instantiate` also accepts positional arguments
and keyword arguments and automatically uses that with the config when creating
the object. This means we can not only pass in the tokenizer, but also add additional
keyword arguments not specified in the config if we'd like:
.. code-block:: yaml
# Tokenizer is needed for the dataset, configure it first
tokenizer:
_component_: torchtune.models.llama2.llama2_tokenizer
path: /tmp/tokenizer.model
dataset:
_component_: torchtune.datasets.alpaca_dataset
.. code-block:: python
# Note the API of the tokenizer we specified - we need to pass in a path
def llama2_tokenizer(path: str) -> Llama2Tokenizer:
# Note the API of the dataset we specified - we need to pass in a model tokenizer
# and any optional keyword arguments
def alpaca_dataset(
tokenizer: ModelTokenizer,
train_on_input: bool = True,
max_seq_len: int = 512,
) -> SFTDataset:
from torchtune import config
# Since we've already specified the path in the config, we don't need to pass
# it in
tokenizer = config.instantiate(cfg.tokenizer)
# We pass in the instantiated tokenizer as the first required argument, then
# we change an optional keyword argument
dataset = config.instantiate(
cfg.dataset,
tokenizer,
train_on_input=False,
)
Note that additional keyword arguments will overwrite any duplicated keys in the
config.
Referencing other config fields with interpolations
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Sometimes you need to use the same value more than once for multiple fields. You
can use *interpolations* to reference another field, and :func:`~torchtune.config.instantiate`
will automatically resolve it for you.
.. code-block:: yaml
output_dir: /tmp/alpaca-llama2-finetune
metric_logger:
_component_: torchtune.training.metric_logging.DiskLogger
log_dir: ${output_dir}
Validating your config
^^^^^^^^^^^^^^^^^^^^^^
We provide a convenient CLI utility, :ref:`tune validate<validate_cli_label>`, to quickly verify that
your config is well-formed and all components can be instantiated properly. You
can also pass in overrides if you want to test out the exact commands you will run
your experiments with. If any parameters are not well-formed, :ref:`tune validate<validate_cli_label>`
will list out all the locations where an error was found.
.. code-block:: bash
tune cp llama2/7B_lora_single_device ./my_config.yaml
tune validate ./my_config.yaml
Best practices for writing configs
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Let's discuss some guidelines for writing configs to get the most out of them.
Airtight configs
""""""""""""""""
While it may be tempting to put as much as you can in the config to give you
maximum flexibility in switching parameters for your experiments, we encourage
you to only include fields in the config that will be used or instantiated in the
recipe. This ensures full clarity on the options a recipe was run with and will
make it significantly easier to debug.
.. code-block:: yaml
# dont do this
alpaca_dataset:
_component_: torchtune.datasets.alpaca_dataset
slimorca_dataset:
...
# do this
dataset:
# change this in config or override when needed
_component_: torchtune.datasets.alpaca_dataset
Use public APIs only
""""""""""""""""""""
If a component you wish to specify in a config is located in a private file, use
the public dotpath in your config. These components are typically exposed in their
parent module's :code:`__init__.py` file. This way, you can guarantee the stability
of the API you are using in your config. There should be no underscores in your
component dotpath.
.. code-block:: yaml
# don't do this
dataset:
_component_: torchtune.datasets._alpaca.alpaca_dataset
# do this
dataset:
_component_: torchtune.datasets.alpaca_dataset
.. _cli_override:
Command-line overrides
----------------------
Configs are the primary location to collect all your parameters to run a recipe,
but sometimes you may want to quickly try different values without having to update
the config itself. To enable quick experimentation, you can specify override values
to parameters in your config via the :code:`tune` command. These should be specified
as key-value pairs :code:`k1=v1 k2=v2 ...`
For example, to run the :ref:`LoRA single-device finetuning <lora_finetune_recipe_label>` recipe with custom model and tokenizer directories, you can provide overrides:
.. code-block:: bash
tune run lora_finetune_single_device \
--config llama2/7B_lora_single_device \
checkpointer.checkpoint_dir=/home/my_model_checkpoint \
checkpointer.checkpoint_files=['file_1','file_2'] \
tokenizer.path=/home/my_tokenizer_path
Overriding components
^^^^^^^^^^^^^^^^^^^^^
If you would like to override a class or function in the config that is instantiated
via the :code:`_component_` field, you can do so by assigning to the parameter
name directly. Any nested fields in the components can be overridden with dot notation.
.. code-block:: yaml
dataset:
_component_: torchtune.datasets.alpaca_dataset
.. code-block:: bash
# Change to slimorca_dataset and set train_on_input to True
tune run lora_finetune_single_device --config my_config.yaml \
dataset=torchtune.datasets.slimorca_dataset dataset.train_on_input=True
Removing config fields
^^^^^^^^^^^^^^^^^^^^^^
You may need to remove certain parameters from the config when changing components
through overrides that require different keyword arguments. You can do so by using
the `~` flag and specify the dotpath of the config field you would like to remove.
For example, if you want to override a built-in config and use the
`bitsandbytes.optim.PagedAdamW8bit <https://huggingface.co/docs/bitsandbytes/main/en/reference/optim/adamw#bitsandbytes.optim.PagedAdamW8bit>`_
optimizer, you may need to delete parameters like ``foreach`` which are
specific to PyTorch optimizers. Note that this example requires that you have `bitsandbytes <https://github.com/bitsandbytes-foundation/bitsandbytes>`_
installed.
.. code-block:: yaml
# In configs/llama3/8B_full.yaml
optimizer:
_component_: torch.optim.AdamW
lr: 2e-5
foreach: False
.. code-block:: bash
# Change to PagedAdamW8bit and remove fused, foreach
tune run --nproc_per_node 4 full_finetune_distributed --config llama3/8B_full \
optimizer=bitsandbytes.optim.PagedAdamW8bit ~optimizer.foreach
---
### Source Deep Dives Recipe Deepdive (`docs/source/deep_dives/recipe_deepdive.rst`)
.. _recipe_deepdive:
=================
What Are Recipes?
=================
This deep-dive will walk you through the design of training-recipes in torchtune.
.. grid:: 1
.. grid-item-card:: :octicon:`mortar-board;1em;` What this deep-dive will cover
* What are recipes?
* What are the core components that make up a recipe?
* How should I structure a new recipe?
Recipes are the primary entry points for torchtune users. These can be thought of
as "targeted" end-to-end pipelines for training and optionally evaluating LLMs.
Each recipe implements a training method (eg: full fine-tuning) with a set of meaningful
features (eg: FSDP + Activation Checkpointing + Gradient Accumulation + Mixed Precision
training) applied to a given model family (eg: Llama2).
As model training gets more and more complex, it becomes harder to anticipate new model
architectures and training methodologies while also reasoning about every possible trade-off
(eg: memory vs model quality). We believe a) users are best suited to make trade-offs
specific to their use cases and b) there's no one-size-fits-all solution. As a result, recipes
are meant to be easy to understand, extend and debug, *and not* generalized entry points for
all possible settings.
Depending on your use case and level of expertise, you will routinely find yourself modifying
existing recipes (eg: adding new features) or writing new ones. torchtune makes writing recipes
easy by providing well-tested modular components/building-blocks and general utilities
(eg: :ref:`WandB Logging<metric_logging_label>` and :ref:`Checkpointing <checkpointing_label>`).
|
**Recipe Design**
Recipes in torchtune are designed to be:
- **Simple**. Written fully in native-PyTorch.
- **Correct**. Numerical parity verification for every component and extensive comparisons with
reference implementations and benchmarks.
- **Easy to Understand**. Each recipe provides a limited set of meaningful features, instead of
every possible feature hidden behind 100s of flags. Code duplication is preferred over unnecessary
abstractions.
- **Easy to Extend**. No dependency on training frameworks and no implementation inheritance. Users
don't need to go through layers-upon-layers of abstractions to figure out how to extend core
functionality.
- **Accessible to a spectrum of Users**. Users can decide how they want to interact with torchtune recipes:
- Start training models by modifying existing configs
- Modify existing recipes for custom cases
- Directly use available building blocks to write completely new recipes/training paradigms
Each recipe consists of three components:
- **Configurable parameters**, specified through yaml configs and command-line overrides
- **Recipe Script**, entry-point which puts everything together including parsing and validating
configs, setting up the environment, and correctly using the recipe class
- **Recipe Class**, core logic needed for training, exposed to users through a set of APIs
In the following sections, we'll take a closer look at each of these components.
For a complete working example, refer to the
`full finetuning recipe <https://github.com/pytorch/torchtune/blob/main/recipes/full_finetune_distributed.py>`_
in torchtune and the associated
`config <https://github.com/pytorch/torchtune/blob/main/recipes/configs/7B_full.yaml>`_.
.. TODO (SalmanMohammadi) ref to full finetune recipe doc
|
What Recipes are not?
---------------------
- **Monolithic Trainers.** A recipe is **not** a monolithic trainer meant to support every
possible feature through 100s of flags.
- **Generalized entry-points.** A recipe is **not** meant to support every possible model
architecture or fine-tuning method.
- **Wrappers around external frameworks.** A recipe is **not** meant to be a wrapper around
external frameworks. These are fully written in native-PyTorch using torchtune building blocks.
Dependencies are primarily in the form of additional utilities or interoperability with the
surrounding ecosystem (eg: EleutherAI's evaluation harness).
|
Recipe Script
-------------
This is the primary entry point for each recipe and provides the user with control over how
the recipe is set up, how models are trained and how the subsequent checkpoints are used.
This includes:
- Setting up of the environment
- Parsing and validating configs
- Training the model
- Setting up multi-stage training (eg: Distillation) using multiple recipe classes
Scripts should generally structure operations in the following order:
- Initialize the recipe class which in-turn initializes recipe state
- Load and Validate checkpoint to update recipe state if resuming training
- Initialize recipe components (model, tokenizer, optimizer, loss and dataloader)
from checkpoint (if applicable)
- Train the model
- Clean up recipe state after training is complete
An example script looks something like this:
.. code-block:: python
# Initialize the process group
init_process_group(backend="gloo" if cfg.device == "cpu" else "nccl")
# Setup the recipe and train the model
recipe = FullFinetuneRecipeDistributed(cfg=cfg)
recipe.setup(cfg=cfg)
recipe.train()
recipe.cleanup()
# Other stuff to do after training is complete
...
Recipe Class
------------
The recipe class carries the core logic for training a model. Each class implements a relevant
interface and exposes a set of APIs. For fine-tuning, the structure of this class is as follows:
Initialize recipe state including seed, device, dtype, metric loggers, relevant flags etc:
.. code-block:: python
def __init__(...):
self._device = utils.get_device(device=params.device)
self._dtype = training.get_dtype(dtype=params.dtype, device=self._device)
...
Load checkpoint, update recipe state from checkpoint, initialize components and load state dicts from checkpoint
.. code-block:: python
def setup(self, cfg: DictConfig):
ckpt_dict = self.load_checkpoint(cfg.checkpointer)
# Setup the model, including FSDP wrapping, setting up activation checkpointing and
# loading the state dict
self._model = self._setup_model(...)
self._tokenizer = self._setup_tokenizer(...)
# Setup Optimizer, including transforming for FSDP when resuming training
self._optimizer = self._setup_optimizer(...)
self._loss_fn = self._setup_loss(...)
self._sampler, self._dataloader = self._setup_data(...)
Run forward and backward across all epochs and save checkpoint at end of each epoch
.. code-block:: python
def train(...):
self._optimizer.zero_grad()
for curr_epoch in range(self.epochs_run, self.total_epochs):
for idx, batch in enumerate(self._dataloader):
...
with self._autocast:
logits = self._model(...)
...
loss = self._loss_fn(logits, labels)
if self.global_step % self._log_every_n_steps == 0:
self._metric_logger.log_dict(...)
loss.backward()
self._optimizer.step()
self._optimizer.zero_grad()
# Update the number of steps when the weights are updated
self.global_step += 1
self.save_checkpoint(epoch=curr_epoch)
Cleanup recipe state
.. code-block:: python
def cleanup(...)
self.metric_loggers.close()
...
Running Recipes with Configs
----------------------------
To run a recipe with a set of user-defined parameters, you will need to write a config file.
You can learn all about configs in our :ref:`config deep-dive<config_tutorial_label>`.
Config and CLI parsing using :code:`parse`
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
We provide a convenient decorator :func:`~torchtune.config.parse` that wraps
your recipe to enable running from the command-line with :ref:`tune <cli_label>` with config
and CLI override parsing.
.. code-block:: python
@config.parse
def recipe_main(cfg: DictConfig) -> None:
recipe = FullFinetuneRecipe(cfg=cfg)
recipe.setup(cfg=cfg)
recipe.train()
recipe.cleanup()
Running your recipe
^^^^^^^^^^^^^^^^^^^
You should be able to run your recipe by providing the direct paths to your custom
recipe and custom config using the :ref:`tune <cli_label>` command with any CLI overrides:
.. code-block:: bash
tune run --config k1=v1 k2=v2 ...
---
### Source Deep Dives Wandb Logging (`docs/source/deep_dives/wandb_logging.rst`)
.. _wandb_logging:
===========================
Logging to Weights & Biases
===========================
This deep-dive will guide you through how to set up logging to Weights & Biases
(W&B) in torchtune.
.. grid:: 1
.. grid-item-card:: :octicon:`mortar-board;1em;` What this deep-dive will cover
* How to get started with W&B
* How to use the :class:`~torchtune.training.metric_logging.WandBLogger`
* How to log configs, metrics, and model checkpoints to W&B
torchtune supports logging your training runs to `Weights & Biases <https://wandb.ai)>`_.
An example W&B workspace from a torchtune fine-tuning run can be seen in the screenshot below.
.. image:: ../_static/img/torchtune_workspace.png
:alt: torchtune workspace in W&B
:width: 100%
:align: center
.. note::
You will need to install the :code:`wandb` package to use this feature.
You can install it via pip:
.. code-block:: bash
pip install wandb
Then you need to login with your API key using the W&B CLI:
.. code-block:: bash
wandb login
Metric Logger
-------------
The only change you need to make is to add the metric logger to your config. Weights & Biases will log the metrics and model checkpoints for you.
.. code-block:: yaml
# enable logging to the built-in WandBLogger
metric_logger:
_component_: torchtune.training.metric_logging.WandBLogger
# the W&B project to log to
project: torchtune
We automatically grab the config from the recipe you are running and log it to W&B. You can find it in the W&B overview tab and the actual file in the :code:`Files` tab.
As a tip, you may see straggler `wandb` processes running in the background if your job crashes or otherwise exits without cleaning up resources. To kill these straggler processes, a command like ``ps
-aux | grep wandb | awk '{ print $2 }' | xargs kill`` can be used.
.. note::
Click on this sample `project to see the W&B workspace <https://wandb.ai/capecape/torchtune>`_.
The config used to train the models can be found `here <https://wandb.ai/capecape/torchtune/runs/6053ofw0/files/torchtune_config_j67sb73v.yaml>`_.
Logging Model Checkpoints to W&B
--------------------------------
You can also log the model checkpoints to W&B by modifying the desired script :code:`save_checkpoint` method.
A suggested approach would be something like this:
.. code-block:: python
def save_checkpoint(self, epoch: int) -> None:
...
## Let's save the checkpoint to W&B
## depending on the Checkpointer Class the file will be named differently
## Here is an example for the full_finetune case
checkpoint_file = Path.joinpath(
self._checkpointer._output_dir, f"torchtune_model_{epoch}"
).with_suffix(".pt")
wandb_at = wandb.Artifact(
name=f"torchtune_model_{epoch}",
type="model",
# description of the model checkpoint
description="Model checkpoint",
# you can add whatever metadata you want as a dict
metadata={
training.SEED_KEY: self.seed,
training.EPOCHS_KEY: self.epochs_run,
training.TOTAL_EPOCHS_KEY: self.total_epochs,
training.MAX_STEPS_KEY: self.max_steps_per_epoch,
}
)
wandb_at.add_file(checkpoint_file)
wandb.log_artifact(wandb_at)
---
### Source Recipes Dpo (`docs/source/recipes/dpo.rst`)
.. _dpo_recipe_label:
====================================
Direct Preference Optimization
====================================
This recipe supports several `Direct Preference Optimization <https://arxiv.org/abs/2305.18290>`_ (DPO)-style fine-tuning techniques.
These techniques aim to steer (or `align <https://en.wikipedia.org/wiki/AI_alignment>`_) a model towards some desirable behaviours.
For example, a common goal is to train language models to produce safe and honest outputs,
or to be `helpful and harmless <https://arxiv.org/abs/2204.05862>`_.
To see the best results when using this recipe, it may be helpful to first fine-tune your model with using supervised fine-tuning to ensure your model is
on-distribution for the domain you're interested in. To do this, check out our other fine-tuning recipes in the :ref:`recipe overview <recipes_overview_label>` which
support a variety of SFT paradigms.
After supervised fine-tuning, here is an example of using either LoRA-based finetuning, or full-finetuning Llama 3.1 8B with DPO:
.. note::
You may need to be granted access to the Llama model you're interested in. See
:ref:`here <download_llama_label>` for details on accessing gated repositories.
.. code-block:: bash
tune download meta-llama/Meta-Llama-3.1-8B-Instruct \
--ignore-patterns "original/consolidated.00.pth"
--HF_TOKEN <HF_TOKEN>
# run lora dpo on a single device
tune run lora_dpo_single_device --config llama3_1/8B_lora_dpo_single_device
# run lora dpo on two gpus
tune run --nproc_per_node 2 lora_dpo_distributed --config llama3_1/8B_lora_dpo
# run full dpo on four gpus
tune run --nproc_per_node 4 full_dpo_distributed --config llama3_1/8B_full_dpo
It's easy to get started with this recipe with your dataset of choice, including custom local datasets,
and datasets from Hugging Face. Check out our primer on :ref:`preference datasets ` to
see how to do this.
For this recipe we include different DPO-style losses:
* :class:`Direct Preference Optimization <torchtune.rlhf.loss.DPOLoss>` (DPO) loss [#]_. The DPO loss function
increases the relative log-probabilities of preferred to un-preferred responses, whilst using log probabilities
from a reference model to prevent policy degradation during training. Alongside RLHF, this is the most commonly used
alignment technique and is used to train a growing number of state-of-the-art LLMs e.g. Llama3.1, Gemma 2, Qwen2, etc.
This is a good starting point for alignment fine-tuning.
* :class:`Statistical Rejection Sampling Optimization <torchtune.rlhf.loss.RSOLoss>` (RSO) or "hinge" loss [#]_.
RSO builds on concepts from support vector machines and DPO, applying a margin-based approach that penalizes
low-quality responses while ensuring a significant gap between chosen and un-chosen log probabilities.
To use any of these, simply use the ``loss`` config entry or flag through the :ref:`cli_label`:
.. code-block:: bash
tune run lora_dpo_single_device --config llama2/7B_lora_dpo_single_device \
loss=torchtune.modules.loss.RSOLoss \
gamma=0.5
Also, you can pass your custom loss in our recipe. Note that its `forward` method should align with the following signature:
.. code-block:: python
def forward(self, policy_inputs: ChosenRejectedOutputs, reference_inputs: ChosenRejectedOutputs) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
...
Here, `ChosenRejectedOutputs` is a dataclass obtained from `concatenated_forward``:
.. code-block:: python
@dataclass
class ChosenRejectedOutputs:
chosen_logps: torch.Tensor
rejected_logps: torch.Tensor
chosen_logits: torch.Tensor
rejected_logits: torch.Tensor
If this is not sufficient and you need to compute additional values from the logits, you can modify `concatenated_forward` directly. To do this, use `tune cp` to copy the desired recipe, and don’t forget to use your own dataclass!
Refer to the TRL library for reference implementations of the desired losses. In particular, you may find useful loss calculations in trainers.
In case, if you don't want to calculate reference values, you can add `is_reference_free: True` to the loss definition in the recipe.
For a deeper understanding of the different levers you can pull when using this recipe,
see our documentation for the different PEFT training paradigms we support:
* :ref:`glossary_lora`
* :ref:`glossary_qlora`
* :ref:`glossary_dora`
Many of our other memory optimization features can be used in this recipe. You can learn more about all of our memory optimization features in our :ref:`memory optimization overview<memory_optimization_overview_label>`.
.. rubric:: References:
.. [#] Rafailov, R., Sharma, A., Mitchell, E., Manning, C.D., Ermon, S. and Finn, C., 2024.
Direct preference optimization: Your language model is secretly a reward model. Advances in Neural Information Processing Systems, 36.
.. [#] Liu, T., Zhao, Y., Joshi, R., Khalman, M., Saleh, M., Liu, P.J. and Liu, J., 2023.
Statistical rejection sampling improves preference optimization. arXiv preprint arXiv:2309.06657.
---
### Source Recipes Lora Finetune Single Device (`docs/source/recipes/lora_finetune_single_device.rst`)
.. _lora_finetune_recipe_label:
=============================
LoRA Single Device Finetuning
=============================
This recipe supports finetuning on next-token prediction tasks using parameter efficient fine-tuning techniques (PEFT)
such as :ref:`glossary_lora` and :ref:`glossary_qlora`. These techniques
significantly reduce memory consumption during training whilst still maintaining competitive performance.
We provide configs which you can get up and running quickly. Here is an example with Llama 3.1 8B:
.. note::
You may need to be granted access to the Llama model you're interested in. See
:ref:`here <download_llama_label>` for details on accessing gated repositories.
.. code-block:: bash
# download the model
tune download meta-llama/Meta-Llama-3.1-8B-Instruct \
--output-dir /tmp/Meta-Llama-3.1-8B-Instruct \
--ignore-patterns "original/consolidated.00.pth"
# run the recipe
tune run lora_finetune_single_device \
--config llama3_1/8B_lora_single_device
You can customize this recipe through the :ref:`cli_label`. For example, when fine-tuning with LoRA, you can adjust the layers which LoRA are applied to:
.. code-block:: bash
tune run lora_finetune_single_device \
--config llama3_1/8B_lora_single_device \
model.lora_attn_modules=“[q_proj,k_proj,v_proj]” \
model.apply_lora_to_mlp=True \
model.lora_rank=64 \
model.lora_alpha=128
For a deeper understanding of the different levers you can pull when using this recipe,
see our documentation for the different PEFT training paradigms we support:
* :ref:`glossary_lora`
* :ref:`glossary_qlora`
* :ref:`glossary_dora`
Many of our other memory optimization features can be used in this recipe. You can learn more about all of our memory optimization features in our :ref:`memory optimization overview<memory_optimization_overview_label>`.
Interested in seeing this recipe in action? Check out some of our tutorials to show off how it can be used:
* :ref:`Finetuning Llama2 with LoRA<lora_finetune_label>`
* :ref:`Finetuning Llama2 with QLoRA<qlora_finetune_label>`
* :ref:`Fine-tuning Llama3 with Chat Data<chat_tutorial_label>`
* :ref:`Meta Llama3 in torchtune<llama3_label>`
* :ref:`Fine-Tune Your First LLM<finetune_llama_label>`
---
### Source Recipes Qat Distributed (`docs/source/recipes/qat_distributed.rst`)
.. _qat_distributed_recipe_label:
=============================================
Distributed Quantization-Aware Training (QAT)
=============================================
QAT allows for taking advantage of memory-saving optimizations from quantization at inference time, without significantly
degrading model performance. In torchtune, we use `torchao <https://github.com/pytorch/ao>`_ to implement QAT.
This works by :ref:`simulating quantization numerics during fine-tuning <what_is_qat_label>`. While this may introduce memory and
compute overheads during training, our tests found that QAT significantly reduced performance degradation in evaluations of
quantized model, without compromising on model size reduction gains. Please see the `PyTorch blogpost <https://pytorch.org/blog/quantization-aware-training/>`_
on QAT for a deeper dive on how the technique works.
We provide pre-tested out-of-the-box configs which you can get up and running with the latest `Llama models <https://llama.meta.com/>`_
in just two steps:
.. code-block:: bash
tune download meta-llama/Meta-Llama-3-8B-Instruct \
--output-dir /tmp/Meta-Llama-3-8B-Instruct \
--ignore-patterns "original/consolidated.00.pth" \
--HF_TOKEN <HF_TOKEN>
tune run --nproc_per_node 6 qat_distributed \
--config llama3/8B_qat_full
.. note::
You may need to be granted access to the Llama model you're interested in. See
:ref:`here <download_llama_label>` for details on accessing gated repositories.
Also, this workload requires at least 6 GPUs, each with VRAM of at least 80GB e.g. A100s or H100s.
Currently, the main lever you can pull for QAT is by using *delayed fake quantization*.
Delayed fake quantization allows for control over the step after which fake quantization occurs.
Empirically, allowing the model to finetune without fake quantization initially allows the
weight and activation values to stabilize before fake quantizing them, potentially leading
to improved quantized accuracy. This can be specified through ``fake_quant_after_n_steps``. To
provide you with an idea of how to roughly configure this parameter, we've achieved best results with
``fake_quant_after_n_steps ~= total_steps // 2``.
In the future we plan to support different quantization strategies. For now, note that you'll need at least
``torch>=2.4.0`` to use the `Int8DynActInt4WeightQATQuantizer <https://github.com/pytorch/ao/blob/08024c686fdd3f3dc2817094f817f54be7d3c4ac/torchao/quantization/prototype/qat/api.py#L35>`_
strategy. Generally, the pipeline for training, quantizing, and evaluating a model using QAT is:
#. Run the ``qat_distributed`` recipe using the above command, or by following the tutorial. By default, this will use ``Int8DynActInt4WeightQATQuantizer``.
#. This produces an un-quantized model in the original data type. To get an actual quantized model, follow this with
``tune run quantize`` while specifying the same quantizer in the config, e.g.
.. code-block:: yaml
# QAT specific args
quantizer:
_component_: torchtune.training.quantization.Int8DynActInt4WeightQATQuantizer
groupsize: 256
#. :ref:`Evaluate<qat_eval_label>` or `run inference <https://github.com/pytorch/torchtune/blob/main/recipes/quantization.md#generate>`_
using your your quantized model by specifying the corresponding post-training quantizer:
.. code-block:: yaml
quantizer:
_component_: torchtune.training.quantization.Int8DynActInt4WeightQuantizer
groupsize: 256
.. note::
We're using config files to show how to customize the recipe in these examples. Check out the
:ref:`configs tutorial <config_tutorial_label>` to learn more.
Many of our other memory optimization features can be used in this recipe, too:
* Adjust :ref:`model precision <glossary_precision>`.
* Use :ref:`activation checkpointing <glossary_act_ckpt>`.
* Enable :ref:`gradient accumulation <glossary_grad_accm>`.
* Use :ref:`lower precision optimizers <glossary_low_precision_opt>`.
You can learn more about all of our memory optimization features in our :ref:`memory optimization overview<memory_optimization_overview_label>`.
Interested in seeing this recipe in action? Check out some of our tutorials to show off how it can be used:
* :ref:`qat_finetune_label`
---
### Source Recipes Recipes Overview (`docs/source/recipes/recipes_overview.rst`)
.. _recipes_overview_label:
================
Recipes Overview
================
Recipes are the primary entry points for torchtune users.
These can be thought of as **hackable, singularly-focused scripts for interacting with LLMs** including fine-tuning,
inference, evaluation, and quantization.
Each recipe consists of three components:
* **Configurable parameters**, specified through yaml configs and command-line overrides
* **Recipe script**, entry-point which puts everything together including parsing and validating configs, setting up the environment, and correctly using the recipe class
* **Recipe class**, core logic needed for fine-tuning, exposed through a set of APIs
.. note::
To learn more about the concept of "recipes", check out our technical deep-dive: :ref:`recipe_deepdive`.
Finetuning
----------
Our recipes include:
* :ref:`Single-device LoRA fine-tuning <lora_finetune_recipe_label>`.
* Single-device full fine-tuning
* Distributed full fine-tuning
* Distributed LoRA fine-tuning
* :ref:`Direct Preference Optimization (DPO) <dpo_recipe_label>`
* Proximal Policy Optimization (PPO)
* :ref:`Distributed Quantization-Aware Training (QAT)<qat_distributed_recipe_label>`.
For a full list, please run:
.. code-block:: bash
tune ls
.. Alignment finetuning
.. --------------------
.. Interested in alignment fine-tuning? You've come to the right place! We support the following alignment techniques:
.. Direct Preference Optimixation (DPO) Fine-Tuning
.. ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.. `Direct Preference Optimixation <https://arxiv.org/abs/2305.18290>`_ (DPO) stype techniques allow for aligning language models with respect
.. to a reward model objective function without the use of reinforcement learning. We support DPO preference fine-tuning with:
.. * :ref:`Single-device <lora_finetune_recipe_label>` and :ref:`multi-device <lora_finetune_recipe_label>` LoRA finetuning.
.. note::
Our recipe documentation is currently in construction. Please feel free to follow the progress in our tracker
issue `here <https://github.com/pytorch/torchtune/issues/1408>`_.
---
### Source Tutorials Chat (`docs/source/tutorials/chat.rst`)
.. _chat_tutorial_label:
=================================
Fine-Tuning Llama3 with Chat Data
=================================
Llama3 Instruct introduced a new prompt template for fine-tuning with chat data. In this tutorial,
we'll cover what you need to know to get you quickly started on preparing your own
custom chat dataset for fine-tuning Llama3 Instruct.
.. grid:: 2
.. grid-item-card:: :octicon:`mortar-board;1em;` You will learn:
* How the Llama3 Instruct format differs from Llama2
* All about prompt templates and special tokens
* How to use your own chat dataset to fine-tune Llama3 Instruct
.. grid-item-card:: :octicon:`list-unordered;1em;` Prerequisites
* Be familiar with :ref:`configuring datasets<chat_dataset_usage_label>`
* Know how to :ref:`download Llama3 Instruct weights <llama3_label>`
Template changes from Llama2 to Llama3
--------------------------------------
The Llama2 chat model requires a specific template when prompting the pre-trained
model. Since the chat model was pretrained with this prompt template, if you want to run
inference on the model, you'll need to use the same template for optimal performance
on chat data. Otherwise, the model will just perform standard text completion, which
may or may not align with your intended use case.
From the `official Llama2 prompt
template guide <https://llama.meta.com/docs/model-cards-and-prompt-formats/meta-llama-2>`_
for the Llama2 chat model, we can see that special tags are added:
.. code-block:: text
<s>[INST] <<SYS>>
You are a helpful, respectful, and honest assistant.
<</SYS>>
Hi! I am a human. [/INST] Hello there! Nice to meet you! I'm Meta AI, your friendly AI assistant </s>
Llama3 Instruct `overhauled <https://llama.meta.com/docs/model-cards-and-prompt-formats/meta-llama-3>`_
the template from Llama2 to better support multiturn conversations. The same text
in the Llama3 Instruct format would look like this:
.. code-block:: text
<|begin_of_text|><|start_header_id|>system<|end_header_id|>
You are a helpful, respectful, and honest assistant.<|eot_id|><|start_header_id|>user<|end_header_id|>
Hi! I am a human.<|eot_id|><|start_header_id|>assistant<|end_header_id|>
Hello there! Nice to meet you! I'm Meta AI, your friendly AI assistant<|eot_id|>
The tags are entirely different, and they are actually encoded differently than in
Llama2. Let's walk through tokenizing an example with the Llama2 template and the
Llama3 template to understand how.
.. note::
The Llama3 Base model uses a `different prompt template
<https://llama.meta.com/docs/model-cards-and-prompt-formats/meta-llama-3>`_ than Llama3 Instruct
because it has not yet been instruct tuned and the extra special tokens are untrained. If you
are running inference on the Llama3 Base model without fine-tuning we recommend the base
template for optimal performance. Generally, for instruct and chat data, we recommend using
Llama3 Instruct with its prompt template. The rest of this tutorial assumes you are using
Llama3 Instruct.
.. _prompt_template_vs_special_tokens:
Tokenizing prompt templates & special tokens
--------------------------------------------
Let's say I have a sample of a single user-assistant turn accompanied with a system
prompt:
.. code-block:: python
sample = [
{
"role": "system",
"content": "You are a helpful, respectful, and honest assistant.",
},
{
"role": "user",
"content": "Who are the most influential hip-hop artists of all time?",
},
{
"role": "assistant",
"content": "Here is a list of some of the most influential hip-hop "
"artists of all time: 2Pac, Rakim, N.W.A., Run-D.M.C., and Nas.",
},
]
Now, let's format this with the :class:`~torchtune.models.llama2.Llama2ChatTemplate` class and
see how it gets tokenized. The Llama2ChatTemplate is an example of a **prompt template**,
which simply structures a prompt with flavor text to indicate a certain task.
.. code-block:: python
from torchtune.data import Llama2ChatTemplate, Message
messages = [Message.from_dict(msg) for msg in sample]
formatted_messages = Llama2ChatTemplate.format(messages)
print(formatted_messages)
# [
# Message(
# role='user',
# content='[INST] <<SYS>>\nYou are a helpful, respectful, and honest assistant.\n<</SYS>>\n\nWho are the most influential hip-hop artists of all time? [/INST] ',
# ...,
# ),
# Message(
# role='assistant',
# content='Here is a list of some of the most influential hip-hop artists of all time: 2Pac, Rakim, N.W.A., Run-D.M.C., and Nas.',
# ...,
# ),
# ]
There are also special tokens used by Llama2, which are not in the prompt template.
If you look at our :class:`~torchtune.models.llama2.Llama2ChatTemplate` class, you'll notice that
we don't include the :code:`<s>` and :code:`</s>` tokens. These are the beginning-of-sequence
(BOS) and end-of-sequence (EOS) tokens that are represented differently in the tokenizer
than the rest of the prompt template. Let's tokenize this example with the
:func:`~torchtune.models.llama2.llama2_tokenizer` used by Llama2 to see
why.
.. code-block:: python
from torchtune.models.llama2 import llama2_tokenizer
tokenizer = llama2_tokenizer("/tmp/Llama-2-7b-hf/tokenizer.model")
user_message = formatted_messages[0].text_content
tokens = tokenizer.encode(user_message, add_bos=True, add_eos=True)
print(tokens)
# [1, 518, 25580, 29962, 3532, 14816, 29903, 6778, ..., 2]
We've added the BOS and EOS tokens when encoding our example text. This shows up
as IDs 1 and 2. We can verify that these are our BOS and EOS tokens.
.. code-block:: python
print(tokenizer._spm_model.spm_model.piece_to_id("<s>"))
# 1
print(tokenizer._spm_model.spm_model.piece_to_id("</s>"))
# 2
The BOS and EOS tokens are what we call special tokens, because they have their own
reserved token IDs. This means that they will index to their own individual vectors in
the model's learnt embedding table. The rest of the prompt template tags, :code:`[INST]`
and :code:`<<SYS>>` are tokenized as normal text and not their own IDs.
.. code-block:: python
print(tokenizer.decode(518))
# '['
print(tokenizer.decode(25580))
# 'INST'
print(tokenizer.decode(29962))
# ']'
print(tokenizer.decode([3532, 14816, 29903, 6778]))
# '<<SYS>>'
It's important to note that you should not place the special reserved tokens in your
input prompts manually, as it will be treated as normal text and not as a special
token.
.. code-block:: python
print(tokenizer.encode("<s>", add_bos=False, add_eos=False))
# [529, 29879, 29958]
Now let's take a look at Llama3's formatting to see how it's tokenized differently
than Llama2.
.. code-block:: python
from torchtune.models.llama3 import llama3_tokenizer
tokenizer = llama3_tokenizer("/tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model")
messages = [Message.from_dict(msg) for msg in sample]
tokens, mask = tokenizer.tokenize_messages(messages)
print(tokenizer.decode(tokens))
# '<|start_header_id|>system<|end_header_id|>\n\nYou are a helpful, respectful,
# and honest assistant.<|eot_id|><|start_header_id|>user<|end_header_id|>\n\nWho
# are the most influential hip-hop artists of all time?<|eot_id|><|start_header_id|>
# assistant<|end_header_id|>\n\nHere is a list of some of the most influential hip-hop
# artists of all time: 2Pac, Rakim, N.W.A., Run-D.M.C., and Nas.<|eot_id|>'
.. note::
We used the ``tokenize_messages`` API for Llama3, which is different than
encode. It simply manages adding all the special tokens in the correct
places after encoding the individual messages.
We can see that the tokenizer handled all the formatting without us specifying a prompt
template. It turns out that all of the additional tags are special tokens, and we don't require
a separate prompt template. We can verify this by checking if the tags get encoded
as their own token IDs.
.. code-block:: python
print(tokenizer.special_tokens["<|begin_of_text|>"])
# 128000
print(tokenizer.special_tokens["<|eot_id|>"])
# 128009
The best part is - all these special tokens are handled purely by the tokenizer.
That means you won't have to worry about messing up any required prompt templates!
When should I use a prompt template?
------------------------------------
Whether or not to use a prompt template is governed by what your desired inference
behavior is. You should use a prompt template if you are running inference on the
base model and it was pre-trained with a prompt template, or you want to prime a
fine-tuned model to expect a certain prompt structure on inference for a specific task.
It is not strictly necessary to fine-tune with a prompt template, but generally
specific tasks will require specific templates. For example, the :class:`~torchtune.data.SummarizeTemplate`
provides a lightweight structure to prime your fine-tuned model for prompts asking to summarize text.
This would wrap around the user message, with the assistant message untouched.
.. code-block:: python
f"Summarize this dialogue:\n{dialogue}\n---\nSummary:\n"
You can fine-tune Llama2 with this template even though the model was originally pre-trained
with the :class:`~torchtune.models.llama2.Llama2ChatTemplate`, as long as this is what the model
sees during inference. The model should be robust enough to adapt to a new template.
Fine-tuning on a custom chat dataset
------------------------------------
Let's test our understanding by trying to fine-tune the Llama3-8B instruct model with a custom
chat dataset. We'll walk through how to set up our data so that it can be tokenized
correctly and fed into our model.
Let's say we have a local dataset saved as a JSON file that contains conversations
with an AI model. How can we get something like this into a format
Llama3 understands and tokenizes correctly?
.. code-block:: python
# data/my_data.json
[
{
"dialogue": [
{
"from": "human",
"value": "What is your name?"
},
{
"from": "gpt",
"value": "I am an AI assistant, I don't have a name."
},
{
"from": "human",
"value": "Pretend you have a name."
},
{
"from": "gpt",
"value": "My name is Mark Zuckerberg."
}
]
},
]
Let's first take a look at the :ref:`dataset_builders` and see which fits our use case. Since we
have conversational data, :func:`~torchtune.datasets.chat_dataset` seems to be a good fit. For any
custom local dataset we always need to specify ``source``, ``data_files``, and ``split`` for any dataset
builder in torchtune. For :func:`~torchtune.datasets.chat_dataset`, we additionally need to specify
``conversation_column`` and ``conversation_style``. Our data follows the ``"sharegpt"`` format, so
we can specify that here. Altogether, our :func:`~torchtune.datasets.chat_dataset` call should
look like so:
.. code-block:: python
from torchtune.datasets import chat_dataset
from torchtune.models.llama3 import llama3_tokenizer
tokenizer = llama3_tokenizer("/tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model")
ds = chat_dataset(
tokenizer=tokenizer,
source="json",
data_files="data/my_data.json",
split="train",
conversation_column="dialogue",
conversation_style="sharegpt",
)
.. code-block:: yaml
# In config
tokenizer:
_component_: torchtune.models.llama3.llama3_tokenizer
path: /tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model
dataset:
_component_: torchtune.datasets.chat_dataset
source: json
data_files: data/my_data.json
split: train
conversation_column: dialogue
conversation_style: sharegpt
.. note::
You can pass in any keyword argument for `load_dataset <https://huggingface.co/docs/datasets/v2.20.0/en/package_reference/loading_methods#datasets.load_dataset>`_ into all our
Dataset classes and they will honor them. This is useful for common parameters
such as specifying the data split with :code:`split` or configuration with
:code:`name`
If you needed to add a prompt template, you would simply pass it into the tokenizer.
Since we're fine-tuning Llama3, the tokenizer will handle all formatting for
us and prompt templates are optional. Other models such as Mistral's :class:`~torchtune.models.mistral._tokenizer.MistralTokenizer`,
use a chat template by default (:class:`~torchtune.models.mistral.MistralChatTemplate`) to format
all messages according to their `recommendations <https://docs.mistral.ai/getting-started/open_weight_models/#chat-template>`_.
Now we're ready to start fine-tuning! We'll use the built-in LoRA single device recipe.
Use the :ref:`tune cp <tune_cp_cli_label>` command to get a copy of the :code:`8B_lora_single_device.yaml`
config and update it with your dataset configuration.
Launch the fine-tune!
.. code-block:: bash
$ tune run lora_finetune_single_device --config custom_8B_lora_single_device.yaml epochs=15
---
### Source Tutorials E2e Flow (`docs/source/tutorials/e2e_flow.rst`)
.. _e2e_flow:
==================================
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 seamlessly with torchtune.
.. grid:: 2
.. grid-item-card:: :octicon:`mortar-board;1em;` What this tutorial will cover:
* 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
.. grid-item-card:: :octicon:`list-unordered;1em;` Prerequisites
* Be familiar with the :ref:`overview of torchtune<overview_label>`
* Make sure to :ref:`install torchtune`
* Concepts such as :ref:`configs <config_tutorial_label>` and
:ref:`checkpoints `
Finetune your model
-------------------
First, let's download a model using the tune CLI. The following command will download the `Llama3.2 3B Instruct <https://ai.meta.com/blog/llama-3-2-connect-2024-vision-edge-mobile-devices/>`_
model from the Hugging Face Hub and save it to the local filesystem. Hugging Face uploaded the original
weights (``consolidated.00.pth``) and the weights compatible with the `from_pretrained() <https://huggingface.co/docs/huggingface_hub/main/en/guides/integrations#frompretrained>`_ API (``*.safetensors``).
We don't need both so we'll ignore the original weights when downloading.
.. code-block:: text
$ tune download meta-llama/Llama-3.2-3B-Instruct --ignore-patterns "original/consolidated.00.pth"
Successfully downloaded model repo and wrote to the following locations:
/tmp/Llama-3.2-3B-Instruct/.cache
/tmp/Llama-3.2-3B-Instruct/.gitattributes
/tmp/Llama-3.2-3B-Instruct/LICENSE.txt
/tmp/Llama-3.2-3B-Instruct/README.md
/tmp/Llama-3.2-3B-Instruct/USE_POLICY.md
/tmp/Llama-3.2-3B-Instruct/config.json
/tmp/Llama-3.2-3B-Instruct/generation_config.json
/tmp/Llama-3.2-3B-Instruct/model-00001-of-00002.safetensors
...
.. note::
For a list of all other models you can finetune out-of-the-box with torchtune, check out
our :ref:`models page<models>`.
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 Llama-3.2-3B-Instruct 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
:ref:`LoRA Tutorial <lora_finetune_label>`.
Let's look for the right config for this use case by using the tune CLI.
.. code-block:: text
$ tune ls
RECIPE CONFIG
full_finetune_single_device llama2/7B_full_low_memory
llama3/8B_full_single_device
llama3_1/8B_full_single_device
llama3_2/1B_full_single_device
llama3_2/3B_full_single_device
mistral/7B_full_low_memory
phi3/mini_full_low_memory
qwen2/7B_full_single_device
...
full_finetune_distributed llama2/7B_full
llama2/13B_full
llama3/8B_full
llama3_1/8B_full
llama3_2/1B_full
llama3_2/3B_full
mistral/7B_full
gemma2/9B_full
gemma2/27B_full
phi3/mini_full
qwen2/7B_full
...
lora_finetune_single_device llama2/7B_lora_single_device
llama2/7B_qlora_single_device
llama3/8B_lora_single_device
...
We'll fine-tune using our
:ref:`single device LoRA recipe <lora_finetune_recipe_label>`
and use the standard settings from the
`default config <https://github.com/pytorch/torchtune/blob/main/recipes/configs/llama3_2/3B_lora_single_device.yaml>`_.
This will fine-tune our model using a ``batch_size=4`` and ``dtype=bfloat16``. With these settings the model
should have a peak memory usage of ~16GB and total training time of around 2-3 hours for each epoch.
.. code-block:: text
$ tune run lora_finetune_single_device --config llama3_2/3B_lora_single_device
Setting manual seed to local seed 3977464327. Local seed is seed + rank = 3977464327 + 0
Hint: enable_activation_checkpointing is True, but enable_activation_offloading isn't. Enabling activation offloading should reduce memory further.
Writing logs to /tmp/torchtune/llama3_2_3B/lora_single_device/logs/log_1734708879.txt
Model is initialized with precision torch.bfloat16.
Memory stats after model init:
GPU peak memory allocation: 6.21 GiB
GPU peak memory reserved: 6.27 GiB
GPU peak memory active: 6.21 GiB
Tokenizer is initialized from file.
Optimizer and loss are initialized.
Loss is initialized.
Dataset and Sampler are initialized.
Learning rate scheduler is initialized.
Profiling disabled.
Profiler config after instantiation: {'enabled': False}
1|3|Loss: 1.943998098373413: 0%| | 3/1617 [00:21<3:04:47, 6.87s/it]
Congrats on training your model! Let's take a look at the artifacts produced by torchtune. A simple way of doing this is by running :code:`tree -a path/to/outputdir`, which should show something like the tree below.
There are 3 types of folders:
1) **recipe_state**: Holds recipe_state.pt with the information necessary to restart the last intermediate epoch. For more information, please check our deep-dive :ref:`Checkpointing in torchtune `.;
2) **logs**: Contains all the logging output from your training run: loss, memory, exceptions, etc.
3) **epoch_{}**: Contains your trained model weights plus model metadata. If running inference or pushing to a model hub, you should use this folder directly.
.. code-block:: text
$ tree -a /tmp/torchtune/llama3_2_3B/lora_single_device
/tmp/torchtune/llama3_2_3B/lora_single_device
├── epoch_0
│ ├── adapter_config.json
│ ├── adapter_model.pt
│ ├── adapter_model.safetensors
│ ├── config.json
│ ├── model-00001-of-00002.safetensors
│ ├── model-00002-of-00002.safetensors
│ ├── generation_config.json
│ ├── LICENSE.txt
│ ├── model.safetensors.index.json
│ ├── original
│ │ ├── orig_params.json
│ │ ├── params.json
│ │ └── tokenizer.model
│ ├── original_repo_id.json
│ ├── README.md
│ ├── special_tokens_map.json
│ ├── tokenizer_config.json
│ ├── tokenizer.json
│ └── USE_POLICY.md
├── epoch_1
│ ├── adapter_config.json
│ ...
├── logs
│ └── log_1734652101.txt
└── recipe_state
└── recipe_state.pt
Let's understand the files:
- ``adapter_model.safetensors`` and ``adapter_model.pt`` are your LoRA trained adapter weights. We save a duplicated .pt version of it to facilitate resuming from checkpoint.
- ``model-{}-of-{}.safetensors`` are your trained full model weights (not adapters). When LoRA finetuning, these are only present if we set ``save_adapter_weights_only=False``. In that case, we merge the base model with trained adapters, making inference easier.
- ``adapter_config.json`` is used by Huggingface PEFT when loading an adapter (more on that later);
- ``model.safetensors.index.json`` is used by Hugging Face ``from_pretrained()`` when loading the model weights (more on that later)
- All other files were originally in the checkpoint_dir. They are automatically copied during training. Files over 100MiB and ending in .safetensors, .pth, .pt, .bin are ignored, making it lightweight.
Evaluate your model
-------------------
We've fine-tuned a model. But how well does this model really do? Let's determine this through structured evaluation and playing with it.
.. _eval_harness_label:
Run evals using EleutherAI's Eval Harness.. TODO (SalmanMohammadi) ref eval recipe docs
torchtune integrates withEleutherAI's evaluation harness <https://github.com/EleutherAI/lm-evaluation-harness>.
An example of this is available through theeleuther_eval <https://github.com/pytorch/torchtune/blob/main/recipes/eleuther_eval.py> recipe. In this tutorial, we're going to directly use this recipe by
modifying its associated config eleuther_evaluation.yaml <https://github.com/pytorch/torchtune/blob/main/recipes/configs/eleuther_evaluation.yaml>_.
.. note::
For this section of the tutorial, you should first run :code:pip install lm_eval>=0.4.5
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.
.. code-block:: bash
text
$ tune cp eleuther_evaluation ./custom_eval_config.yaml
Copied file to custom_eval_config.yamlNotice that we are using the merged weights, and not the LoRA adapters.
.. code-block:: yaml
text
# TODO: update to your desired epoch
output_dir: /tmp/torchtune/llama3_2_3B/lora_single_device/epoch_0
# Tokenizer
tokenizer:
_component_: torchtune.models.llama3.llama3_tokenizer
path: ${output_dir}/original/tokenizer.model
model:
# Notice that we don't pass the lora model. We are using the merged weights,
_component_: torchtune.models.llama3_2.llama3_2_3b
checkpointer:
_component_: torchtune.training.FullModelHFCheckpointer
checkpoint_dir: ${output_dir}
checkpoint_files: [
model-00001-of-00002.safetensors,
model-00002-of-00002.safetensors,
]
output_dir: ${output_dir}
model_type: LLAMA3_2
### OTHER PARAMETERS -- NOT RELATED TO THIS CHECKPOINT
# Environment
device: cuda
dtype: bf16
seed: 1234 # It is not recommended to change this seed, b/c it matches EleutherAI's default seed
# EleutherAI specific eval args
tasks: ["truthfulqa_mc2"]
limit: null
max_seq_length: 4096
batch_size: 8
enable_kv_cache: True
# Quantization specific args
quantizer: nullFor this tutorial we'll use the truthfulqa_mc2 <https://github.com/sylinrl/TruthfulQA>_ 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.
.. code-block:: text
text
$ tune run eleuther_eval --config ./custom_eval_config.yaml
[evaluator.py:324] Running loglikelihood requests
...Generate some output
text
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 <https://github.com/pytorch/torchtune/blob/main/recipes/generate.py>`_
and the associated
`config <https://github.com/pytorch/torchtune/blob/main/recipes/configs/generation.yaml>`_.
Let's first copy over the config to our local working directory so we can make changes.
.. code-block:: text
$ tune cp generation ./custom_generation_config.yaml
Copied file to custom_generation_config.yaml
$ mkdir /tmp/torchtune/llama3_2_3B/lora_single_device/out
Let's modify ``custom_generation_config.yaml`` to include the following changes. Again, you only need
to replace two fields: ``output_dir`` and ``checkpoint_files``
.. code-block:: yaml
checkpoint_dir: /tmp/torchtune/llama3_2_3B/lora_single_device/epoch_0
output_dir: /tmp/torchtune/llama3_2_3B/lora_single_device/out
# Tokenizer
tokenizer:
_component_: torchtune.models.llama3.llama3_tokenizer
path: ${checkpoint_dir}/original/tokenizer.model
prompt_template: null
model:
# Notice that we don't pass the lora model. We are using the merged weights,
_component_: torchtune.models.llama3_2.llama3_2_3b
checkpointer:
_component_: torchtune.training.FullModelHFCheckpointer
checkpoint_dir: ${checkpoint_dir}
checkpoint_files: [
model-00001-of-00002.safetensors,
model-00002-of-00002.safetensors,
]
output_dir: ${output_dir}
model_type: LLAMA3_2
### OTHER PARAMETERS -- NOT RELATED TO THIS CHECKPOINT
device: cuda
dtype: bf16
seed: 1234
# Generation arguments; defaults taken from gpt-fast
prompt:
system: null
user: "Tell me a joke. "
max_new_tokens: 300
temperature: 0.6 # 0.8 and 0.6 are popular values to try
top_k: 300
enable_kv_cache: True
quantizer: null
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. We recommend inspecting the model with these before playing around with
these parameters.
.. code-block:: text
$ tune run generate --config ./custom_generation_config.yaml prompt.user="Tell me a joke. "
Tell me a joke. Here's a joke for you:
What do you call a fake noodle?
An impasta!
Introduce some quantizationWe rely on torchao <https://github.com/pytorch/ao>_ for post-training quantization <https://github.com/pytorch/ao/tree/main/torchao/quantization#quantization>_.
To quantize the fine-tuned model after installing torchao we can run the following command::
we also support `int8_weight_only()` and `int8_dynamic_activation_int8_weight()`, see
https://github.com/pytorch/ao/tree/main/torchao/quantization#other-available-quantization-techniques
for a full list of techniques that we support
from torchao.quantization.quant_api import quantize_, int4_weight_only
quantize_(model, int4_weight_only())
After quantization, we rely on torch.compile for speedups. For more details, please see this example usage <https://github.com/pytorch/ao/blob/main/torchao/quantization/README.md#quantization-flow-example>_.
torchao also provides this table <https://github.com/pytorch/ao#inference>_ listing performance and accuracy results for llama2 and llama3.
For Llama models, you can run generation directly in torchao on the quantized model using their generate.py script as
discussed in this readme <https://github.com/pytorch/ao/tree/main/torchao/_models/llama>_. This way you can compare your own results
to those in the previously-linked table.
.. _use_model_in_wild:
Use your model in the wild
Let's say we're happy with how our model is performing at this point - we want to do something with it! Productionize it for serving, publish on the Hugging Face Hub, etc.
Since we handle checkpoint conversion, you can directly work with standard formats.
Use with Hugging Face from_pretrained()
text
**Case 1: Hugging Face using base model + trained adapters**
Here we load the base model from Hugging Face model hub. Then we load the adapters on top of it using `PeftModel <https://huggingface.co/docs/peft/v0.6.1/en/package_reference/peft_model>`_.
It will look for the files ``adapter_model.safetensors`` for the weights and ``adapter_config.json`` for where to insert them.
.. code-block:: python
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
#TODO: update it to your chosen epoch
trained_model_path = "/tmp/torchtune/llama3_2_3B/lora_single_device/epoch_0"
# Define the model and adapter paths
original_model_name = "meta-llama/Llama-3.2-1B-Instruct"
model = AutoModelForCausalLM.from_pretrained(original_model_name)
# huggingface will look for adapter_model.safetensors and adapter_config.json
peft_model = PeftModel.from_pretrained(model, trained_model_path)
# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(original_model_name)
# Function to generate text
def generate_text(model, tokenizer, prompt, max_length=50):
inputs = tokenizer(prompt, return_tensors="pt")
outputs = model.generate(**inputs, max_length=max_length)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
prompt = "tell me a joke: '"
print("Base model output:", generate_text(peft_model, tokenizer, prompt))
**Case 2: Hugging Face using merged weights**
In this case, Hugging Face will check in ``model.safetensors.index.json`` for which files it should load.
.. code-block:: python
from transformers import AutoModelForCausalLM, AutoTokenizer
#TODO: update it to your chosen epoch
trained_model_path = "/tmp/torchtune/llama3_2_3B/lora_single_device/epoch_0"
model = AutoModelForCausalLM.from_pretrained(
pretrained_model_name_or_path=trained_model_path,
)
# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(trained_model_path, safetensors=True)
# Function to generate text
def generate_text(model, tokenizer, prompt, max_length=50):
inputs = tokenizer(prompt, return_tensors="pt")
outputs = model.generate(**inputs, max_length=max_length)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
prompt = "Complete the sentence: 'Once upon a time...'"
print("Base model output:", generate_text(model, tokenizer, prompt))
Use with vLLM
~~~~~~~~~~~~~
`vLLM <https://docs.vllm.ai/en/latest/>`_ is a fast and easy-to-use library for LLM inference and serving. They include a lot of awesome features like
state-of-the-art serving throughput, continuous batching of incoming requests, quantization, and speculative decoding.
The library will load any .safetensors file. Since we already merged the full model weights and adapter weights, we can safely delete the
adapter weights (or move them) so that vLLM doesn't get confused by those files.
.. code-block:: python
rm /tmp/torchtune/llama3_2_3B/lora_single_device/base_model/adapter_model.safetensors
Now we can run the following script:
.. code-block:: python
from vllm import LLM, SamplingParams
def print_outputs(outputs):
for output in outputs:
prompt = output.prompt
generated_text = output.outputs[0].text
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
print("-" * 80)
#TODO: update it to your chosen epoch
llm = LLM(
model="/tmp/torchtune/llama3_2_3B/lora_single_device/epoch_0",
load_format="safetensors",
kv_cache_dtype="auto",
)
sampling_params = SamplingParams(max_tokens=16, temperature=0.5)
conversation = [
{"role": "system", "content": "You are a helpful assistant"},
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hello! How can I assist you today?"},
{
"role": "user",
"content": "Write an essay about the importance of higher education.",
},
]
outputs = llm.chat(conversation, sampling_params=sampling_params, use_tqdm=False)
print_outputs(outputs)
Upload 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_hub <https://huggingface.co/docs/huggingface_hub/guides/upload>`_.
.. code-block:: python
import huggingface_hub
api = huggingface_hub.HfApi()
#TODO: update it to your chosen epoch
trained_model_path = "/tmp/torchtune/llama3_2_3B/lora_single_device/epoch_0"
username = huggingface_hub.whoami()["name"]
repo_name = "my-model-trained-with-torchtune"
# if the repo doesn't exist
repo_id = huggingface_hub.create_repo(repo_name).repo_id
# if it already exists
repo_id = f"{username}/{repo_name}"
api.upload_folder(
folder_path=trained_model_path,
repo_id=repo_id,
repo_type="model",
create_pr=False
)
If you prefer, you can also try the cli version `huggingface-cli upload <https://huggingface.co/docs/huggingface_hub/en/guides/cli#huggingface-cli-upload>`_.
|
Hopefully this tutorial gave you some insights into how you can use torchtune for
your own workflows. Happy Tuning!
---
### Source Tutorials First Finetune Tutorial (`docs/source/tutorials/first_finetune_tutorial.rst`)
.. _finetune_llama_label:
========================
Fine-Tune Your First LLM
========================
This guide will walk you through the process of launching your first finetuning
job using torchtune.
.. grid:: 2
.. grid-item-card:: :octicon:`mortar-board;1em;` What you will learn
* How to download a model from the `Hugging Face Hub <https://huggingface.co/docs/hub/en/index>`_
* How to modify a recipe's parameters to suit your needs
* How to run a finetune
.. grid-item-card:: :octicon:`list-unordered;1em;` Prerequisites
* Be familiar with the :ref:`overview of torchtune<overview_label>`
* Make sure to :ref:`install torchtune`
.. _download_llama_label:
Downloading a model
-------------------
The first step in any finetuning job is to download a pretrained base model. torchtune supports an integration
with the `Hugging Face Hub <https://huggingface.co/docs/hub/en/index>`_ - a collection of the latest and greatest model weights.
For this tutorial, you're going to use the `Llama2 7B model from Meta <https://llama.meta.com/>`_. Llama2 is a "gated model",
meaning that you need to be granted access in order to download the weights. Follow `these instructions <https://huggingface.co/meta-llama>`_ on the official Meta page
hosted on Hugging Face to complete this process. This should take less than 5 minutes. To verify that you have the access, go to the `model page <https://huggingface.co/meta-llama/Llama-2-7b-hf/tree/main>`_.
You should be able to see the model files. If not, you may need to accept the agreement to complete the process.
.. note::
Alternatively, you can opt to download the model directly through the Llama2 repository.
See `this page <https://llama.meta.com/get-started#getting-the-models>`_ for more details.
Once you have authorization, you will need to authenticate with Hugging Face Hub. The easiest way to do so is to provide an
access token to the download script. You can find your token `here <https://huggingface.co/settings/tokens>`_.
Then, it's as simple as:
.. code-block:: bash
tune download meta-llama/Llama-2-7b-hf \
--output-dir /tmp/Llama-2-7b-hf \
--hf-token <ACCESS TOKEN>
This command will also download the model tokenizer and some other helpful files such as a Responsible Use guide.
|
Selecting a recipe
------------------
Recipes are the primary entry points for torchtune users.
These can be thought of as **hackable, singularly-focused scripts for interacting with LLMs** including training,
inference, evaluation, and quantization.
Each recipe consists of three components:
* **Configurable parameters**, specified through yaml configs and command-line overrides
* **Recipe script**, entry-point which puts everything together including parsing and validating configs, setting up the environment, and correctly using the recipe class
* **Recipe class**, core logic needed for training, exposed through a set of APIs
.. note::
To learn more about the concept of "recipes", check out our technical deep-dive: :ref:`recipe_deepdive`.
torchtune provides built-in recipes for finetuning on single device, on multiple devices with `FSDP <https://pytorch.org/blog/introducing-pytorch-fully-sharded-data-parallel-api/>`_,
using memory efficient techniques like `LoRA <https://arxiv.org/abs/2106.09685>`_, and more! Check out all our built-in recipes in our :ref:`recipes overview<recipes_overview_label>`. You can also utilize the
:code:`tune ls` command to print out all recipes and corresponding configs.
.. code-block:: bash
$ 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 the purposes of this tutorial, you'll will be using the recipe for finetuning a Llama2 model using `LoRA <https://arxiv.org/abs/2106.09685>`_ on
a single device. For a more in-depth discussion on LoRA in torchtune, you can see the complete ":ref:`lora_finetune_label`" tutorial.
.. note::
**Why have a separate recipe for single device vs. distributed?** This is discussed in
":ref:`recipe_deepdive`" but one of our :ref:`core principles <design_principles_label>` in torchtune is minimal abstraction and boilerplate code.
If you only want to train on a single GPU, our single-device recipe ensures you don't have to worry about additional
features like FSDP that are only required for distributed training.
|
.. _tune_cp_label:
Modifying a config
------------------
YAML configs hold most of the important information needed for running your recipe.
You can set hyperparameters, specify metric loggers like `WandB <wandb.ai>`_, select a new dataset, and more.
For a list of all currently supported datasets, see :ref:`datasets`.
There are two ways to modify an existing config:
**Override existing parameters from the command line**
You can override existing parameters from the command line using a :code:`key=value` format. Let's say
you want to set the number of training epochs to 1.
.. code-block:: bash
tune run <RECIPE> --config <CONFIG> epochs=1
**Copy the config through `tune cp` and modify directly**
If you want to make more substantial changes to the config, you can use the :ref:`tune <cli_label>` CLI to copy it to your local directory.
.. code-block:: bash
$ tune cp llama2/7B_lora_single_device custom_config.yaml
Copied file to custom_config.yaml
Now you can update the custom YAML config any way you like. Try setting the random seed in order to make replication easier,
changing the LoRA rank, update batch size, etc.
.. note::
Check out ":ref:`config_tutorial_label`" for a deeper dive on configs in torchtune.
|
Training a model
----------------
Now that you have a model in the proper format and a config that suits your needs, let's get training!
Just like all the other steps, you will be using the tune CLI tool to launch your finetuning run.
.. code-block:: bash
$ tune run lora_finetune_single_device --config llama2/7B_lora_single_device epochs=1
INFO:torchtune.utils.logging:Running LoRAFinetuneRecipeSingleDevice with resolved config:
Writing logs to /tmp/lora_finetune_output/log_1713194212.txt
INFO:torchtune.utils.logging:Model is initialized with precision torch.bfloat16.
INFO:torchtune.utils.logging:Tokenizer is initialized from file.
INFO:torchtune.utils.logging:Optimizer and loss are initialized.
INFO:torchtune.utils.logging:Loss is initialized.
INFO:torchtune.utils.logging:Dataset and Sampler are initialized.
INFO:torchtune.utils.logging:Learning rate scheduler is initialized.
1|52|Loss: 2.3697006702423096: 0%|▏ | 52/25880 [00:24<3:55:01, 1.83it/s]
You can see that all the modules were successfully initialized and the model has started training.
You can monitor the loss and progress through the `tqdm <https://tqdm.github.io/>`_ bar but torchtune
will also log some more metrics, such as GPU memory usage, at an interval defined in the config.
|
Next steps
----------
Now that you have trained your model and set up your environment, let's take a look at what we can do with our
new model by checking out the ":ref:`E2E Workflow Tutorial<e2e_flow>`".
---
### Source Tutorials Llama3 (`docs/source/tutorials/llama3.rst`)
.. _llama3_label:
========================
Meta Llama3 in torchtune
========================
.. grid:: 2
.. grid-item-card:: :octicon:`mortar-board;1em;` You will learn how to:
* Download the Llama3-8B-Instruct weights and tokenizer
* Fine-tune Llama3-8B-Instruct with LoRA and QLoRA
* Evaluate your fine-tuned Llama3-8B-Instruct model
* Generate text with your fine-tuned model
* Quantize your model to speed up generation
.. grid-item-card:: :octicon:`list-unordered;1em;` Prerequisites
* Be familiar with :ref:`torchtune<overview_label>`
* Make sure to :ref:`install torchtune`
Llama3-8B
---------
`Meta Llama 3 <https://llama.meta.com/llama3>`_ is a new family of models released by Meta AI that improves upon the performance of the Llama2 family
of models across a `range of different benchmarks <https://huggingface.co/meta-llama/Meta-Llama-3-8B#base-pretrained-models>`_.
Currently there are two different sizes of Meta Llama 3: 8B and 70B. In this tutorial we will focus on the 8B size model.
There are a few main changes between Llama2-7B and Llama3-8B models:
- Llama3-8B uses `grouped-query attention <https://arxiv.org/abs/2305.13245>`_ instead of the standard multi-head attention from Llama2-7B
- Llama3-8B has a larger vocab size (128,256 instead of 32,000 from Llama2 models)
- Llama3-8B uses a different tokenizer than Llama2 models (`tiktoken <https://github.com/openai/tiktoken>`_ instead of `sentencepiece <https://github.com/google/sentencepiece>`_)
- Llama3-8B uses a larger intermediate dimension in its MLP layers than Llama2-7B
- Llama3-8B uses a higher base value to calculate theta in its `rotary positional embeddings <https://arxiv.org/abs/2104.09864>`_
|
Getting access to Llama3-8B-Instruct
------------------------------------
For this tutorial, we will be using the instruction-tuned version of Llama3-8B. First, let's download the model from Hugging Face. You will need to follow the instructions
on the `official Meta page <https://github.com/meta-llama/llama3/blob/main/README.md>`_ to gain access to the model.
Next, make sure you grab your Hugging Face token from `here <https://huggingface.co/settings/tokens>`_.
.. code-block:: bash
tune download meta-llama/Meta-Llama-3-8B-Instruct \
--output-dir <checkpoint_dir> \
--hf-token <ACCESS TOKEN>
|
Fine-tuning Llama3-8B-Instruct in torchtune
-------------------------------------------
torchtune provides `LoRA <https://arxiv.org/abs/2106.09685>`_, `QLoRA <https://arxiv.org/abs/2305.14314>`_, and full fine-tuning
recipes for fine-tuning Llama3-8B on one or more GPUs. For more on LoRA in torchtune, see our :ref:`LoRA Tutorial <lora_finetune_label>`.
For more on QLoRA in torchtune, see our :ref:`QLoRA Tutorial <qlora_finetune_label>`.
Let's take a look at how we can fine-tune Llama3-8B-Instruct with LoRA on a single device using torchtune. In this example, we will fine-tune
for one epoch on a common instruct dataset for illustrative purposes. The basic command for a single-device LoRA fine-tune is
.. code-block:: bash
tune run lora_finetune_single_device --config llama3/8B_lora_single_device
.. note::
To see a full list of recipes and their corresponding configs, simply run ``tune ls`` from the command line.
We can also add :ref:`command-line overrides <cli_override>` as needed, e.g.
.. code-block:: bash
tune run lora_finetune_single_device --config llama3/8B_lora_single_device \
checkpointer.checkpoint_dir=<checkpoint_dir> \
tokenizer.path=<checkpoint_dir>/tokenizer.model \
checkpointer.output_dir=<checkpoint_dir>
This will load the Llama3-8B-Instruct checkpoint and tokenizer from ``<checkpoint_dir>`` used in the :ref:`tune download <tune_download_label>` command above,
then save a final checkpoint in the same directory following the original format. For more details on the
checkpoint formats supported in torchtune, see our :ref:`checkpointing deep-dive `.
.. note::
To see the full set of configurable parameters for this (and other) configs we can use :ref:`tune cp <tune_cp_cli_label>` to copy (and modify)
the default config. :ref:`tune cp <tune_cp_cli_label>` can be used with recipe scripts too, in case you want to make more custom changes
that cannot be achieved by directly modifying existing configurable parameters. For more on :ref:`tune cp <tune_cp_cli_label>` see the section on
:ref:`modifying configs <tune_cp_label>` in our ":ref:`finetune_llama_label`" tutorial.
Once training is complete, the model checkpoints will be saved and their locations will be logged. For
LoRA fine-tuning, the final checkpoint will contain the merged weights, and a copy of just the (much smaller) LoRA weights
will be saved separately.
In our experiments, we observed a peak memory usage of 18.5 GB. The default config can be trained on a consumer GPU with 24 GB VRAM.
If you have multiple GPUs available, you can run the distributed version of the recipe.
torchtune makes use of the `FSDP <https://pytorch.org/tutorials/intermediate/FSDP_tutorial.html>`_ APIs from PyTorch Distributed
to shard the model, optimizer states, and gradients. This should enable you to increase your batch size, resulting in faster overall training.
For example, on two devices:
.. code-block:: bash
tune run --nproc_per_node 2 lora_finetune_distributed --config llama3/8B_lora
Finally, if we want to use even less memory, we can leverage torchtune's QLoRA recipe via:
.. TODO (SalmanMohammadi) ref qlora recipe page
.. code-block:: bash
tune run lora_finetune_single_device --config llama3/8B_qlora_single_device
Since our default configs enable full bfloat16 training, all of the above commands can be run with
devices having at least 24 GB of VRAM, and in fact the QLoRA recipe should have peak allocated memory
below 10 GB. You can also experiment with different configurations of LoRA and QLoRA, or even run a full fine-tune.
Try it out!
|
Evaluating fine-tuned Llama3-8B models with EleutherAI's Eval Harness
---------------------------------------------------------------------
Now that we've fine-tuned our model, what's next? Let's take our LoRA-finetuned model from the
preceding section and look at a couple different ways we can evaluate its performance on the tasks we care about.
First, torchtune provides an integration with
`EleutherAI's evaluation harness <https://github.com/EleutherAI/lm-evaluation-harness>`_
for model evaluation on common benchmark tasks.
.. note::
Make sure you've first installed the evaluation harness via :code:`pip install "lm_eval==0.4.*"`.
For this tutorial we'll use the `truthfulqa_mc2 <https://github.com/sylinrl/TruthfulQA>`_ 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. First, let's copy the config so we can point the YAML
file to our fine-tuned checkpoint files.
.. code-block:: bash
tune cp eleuther_evaluation ./custom_eval_config.yaml
Next, we modify ``custom_eval_config.yaml`` to include the fine-tuned checkpoints.
.. code-block:: yaml
model:
_component_: torchtune.models.llama3.llama3_8b
checkpointer:
_component_: torchtune.training.FullModelMetaCheckpointer
# directory with the checkpoint files
# this should match the output_dir specified during
# fine-tuning
checkpoint_dir: <checkpoint_dir>
# checkpoint files for the fine-tuned model. These will be logged
# at the end of your fine-tune
checkpoint_files: [
meta_model_0.pt
]
output_dir: <checkpoint_dir>
model_type: LLAMA3
# Make sure to update the tokenizer path to the right
# checkpoint directory as well
tokenizer:
_component_: torchtune.models.llama3.llama3_tokenizer
path: <checkpoint_dir>/tokenizer.model
Finally, we can run evaluation using our modified config.
.. code-block:: bash
tune run eleuther_eval --config ./custom_eval_config.yaml
Try it for yourself and see what accuracy your model gets!
|
Generating text with our fine-tuned Llama3 model
------------------------------------------------
.. TODO (SalmanMohammadi) ref generate recipe page
Next, let's look at one other way we can evaluate our model: generating text! torchtune provides a
`recipe for generation <https://github.com/pytorch/torchtune/blob/main/recipes/generate.py>`_ as well.
Similar to what we did, let's copy and modify the default generation config.
.. code-block:: bash
tune cp generation ./custom_generation_config.yaml
Now we modify ``custom_generation_config.yaml`` to point to our checkpoint and tokenizer.
.. code-block:: yaml
model:
_component_: torchtune.models.llama3.llama3_8b
checkpointer:
_component_: torchtune.training.FullModelMetaCheckpointer
# directory with the checkpoint files
# this should match the output_dir specified during
# fine-tuning
checkpoint_dir: <checkpoint_dir>
# checkpoint files for the fine-tuned model. These will be logged
# at the end of your fine-tune
checkpoint_files: [
meta_model_0.pt
]
output_dir: <checkpoint_dir>
model_type: LLAMA3
# Make sure to update the tokenizer path to the right
# checkpoint directory as well
tokenizer:
_component_: torchtune.models.llama3.llama3_tokenizer
path: <checkpoint_dir>/tokenizer.model
Running generation with our LoRA-finetuned model, we see the following output:
.. code-block:: bash
tune run generate --config ./custom_generation_config.yaml \
prompt.user="Hello, my name is"
[generate.py:122] Hello, my name is Sarah and I am a busy working mum of two young children, living in the North East of England.
...
[generate.py:135] Time for inference: 10.88 sec total, 18.94 tokens/sec
[generate.py:138] Bandwidth achieved: 346.09 GB/s
[generate.py:139] Memory used: 18.31 GB
Faster generation via quantization
----------------------------------
We rely on `torchao <https://github.com/pytorch/ao>`_ for `post-training quantization <https://github.com/pytorch/ao/tree/main/torchao/quantization#quantization>`_.
To quantize the fine-tuned model after installing torchao we can run the following command::
# we also support `int8_weight_only()` and `int8_dynamic_activation_int8_weight()`, see
# https://github.com/pytorch/ao/tree/main/torchao/quantization#other-available-quantization-techniques
# for a full list of techniques that we support
from torchao.quantization.quant_api import quantize_, int4_weight_only
quantize_(model, int4_weight_only())
After quantization, we rely on torch.compile for speedups. For more details, please see `this example usage <https://github.com/pytorch/ao/blob/main/torchao/quantization/README.md#quantization-flow-example>`_.
torchao also provides `this table <https://github.com/pytorch/ao#inference>`_ listing performance and accuracy results for ``llama2`` and ``llama3``.
For Llama models, you can run generation directly in torchao on the quantized model using their ``generate.py`` script as
discussed in `this readme <https://github.com/pytorch/ao/tree/main/torchao/_models/llama>`_. This way you can compare your own results
to those in the previously-linked table.
This is just the beginning of what you can do with Meta Llama3 using torchtune and the broader ecosystem.
We look forward to seeing what you build!
---
### Source Tutorials Llama Kd Tutorial (`docs/source/tutorials/llama_kd_tutorial.rst`)
.. _llama_kd_label:
====================================================================
Distilling Llama3.1 8B into Llama3.2 1B using Knowledge Distillation
====================================================================
This guide will teach you about knowledge distillation (KD) and show you how you can use torchtune to distill a Llama3.1 8B model into Llama3.2 1B.
If you already know what knowledge distillation is and want to get straight to running your own distillation in torchtune,
you can jump to the `KD recipe in torchtune`_ tutorial.
.. grid:: 2
.. grid-item-card:: :octicon:`mortar-board;1em;` What you will learn
* What KD is and how it can help improve model performance
* An overview of KD components in torchtune
* How to distill from a teacher to student model using torchtune
* How to experiment with different KD configurations
.. grid-item-card:: :octicon:`list-unordered;1em;` Prerequisites
* Be familiar with :ref:`torchtune<overview_label>`
* Make sure to :ref:`install torchtune`
* Make sure you have downloaded the :ref:`Llama3 model weights<download_llama_label>`
* Be familiar with :ref:`LoRA<lora_finetune_label>`
What is Knowledge Distillation?
-------------------------------
`Knowledge Distillation <https://arxiv.org/pdf/1503.02531>`_ is a widely used compression technique
that transfers knowledge from a larger (teacher) model to a smaller (student) model. Larger models have
more parameters and capacity for knowledge, however, this larger capacity is also more computationally
expensive to deploy. Knowledge distillation can be used to compress the knowledge of a larger model into
a smaller model. The idea is that performance of smaller models can be improved by learning from larger
model's outputs.
How does Knowledge Distillation work?
-------------------------------------
Knowledge is transferred from the teacher to student model by training it on a transfer set where the
student is trained to imitate the token-level probability distributions of the teacher. The diagram below
is a simplified representation of how KD works.
.. image:: /_static/img/kd-simplified.png
The total loss can be configured in many ways. The default KD config in torchtune combines the cross-entropy (CE) loss with the
forward `Kullback-Leibler (KL) divergence <https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence>`_ loss,
which is used in standard KD approaches. Forward KL divergence aims to minimize the difference by forcing the student's
distribution to align with all of the teacher's distributions. However, aligning the student distribution to the whole
teacher distribution may not be effective and there are multiple papers, such as `MiniLLM <https://arxiv.org/pdf/2306.08543>`_,
`DistiLLM <https://arxiv.org/pdf/2402.03898>`_, and `Generalized KD <https://arxiv.org/pdf/2306.13649>`_,
that introduce new KD losses to address the limitations. For this tutorial, let's take a look at the implementation of
the forward KL divergence loss.
.. code-block:: python
import torch
import torch.nn.functional as F
class ForwardKLLoss(torch.nn.Module):
def __init__(self, ignore_index: int = -100)
super().__init__()
self.ignore_index = ignore_index
def forward(self, student_logits, teacher_logits, labels) -> torch.Tensor:
# Implementation from https://github.com/jongwooko/distillm
# Computes the softmax of the teacher logits
teacher_prob = F.softmax(teacher_logits, dim=-1, dtype=torch.float32)
# Computes the student log softmax probabilities
student_logprob = F.log_softmax(student_logits, dim=-1, dtype=torch.float32)
# Computes the forward KL divergence
prod_probs = teacher_prob * student_logprob
# Compute the sum
x = torch.sum(prod_probs, dim=-1).view(-1)
# We don't want to include the ignore labels in the average
mask = (labels != self.ignore_index).int()
# Loss is averaged over non-ignored targets
return -torch.sum(x * mask.view(-1), dim=0) / torch.sum(mask.view(-1), dim=0)
There are some details omitted to simplify the computation, but if you'd like to know more,
you can see the implementation in :class:`~torchtune.modules.loss.ForwardKLLoss`.
By default, the KD configs use :class:`~torchtune.modules.loss.ForwardKLWithChunkedOutputLoss` to reduce memory.
The current implementation only supports student and teacher models that have the same output
logit shape and same tokenizer.
KD recipe in torchtune
----------------------
With torchtune, we can easily apply knowledge distillation to Llama3, as well as other LLM model families.
Let's take a look at how you could distill a model using torchtune's `KD recipe <https://github.com/pytorch/torchtune/blob/4234b78b914af23384ce0348f564e2119d107a96/recipes/knowledge_distillation_single_device.py>`_.
First, make sure that you have downloaded all the model weights. For this example, we'll use the Llama3.1-8B as teacher and Llama3.2-1B as student.
.. code-block:: bash
tune download meta-llama/Meta-Llama-3.1-8B-Instruct --output-dir /tmp/Meta-Llama-3.1-8B-Instruct --ignore-patterns "original/consolidated.00.pth" --hf-token <HF_TOKEN>
tune download meta-llama/Llama-3.2-1B-Instruct --output-dir /tmp/Llama-3.2-1B-Instruct --ignore-patterns "original/consolidated.00.pth" --hf-token <HF_TOKEN>
Then, we will fine-tune the teacher model using LoRA. Based on our experiments and previous work,
we've found that KD performs better when the teacher model is already fine-tuned on the target dataset.
.. code-block:: bash
tune run lora_finetune_single_device --config llama3_1/8B_lora_single_device
Finally, we can run the following command to distill the fine-tuned 8B model into the 1B model on a single GPU.
.. code-block:: bash
tune run knowledge_distillation_single_device --config llama3_2/8B_to_1B_KD_lora_single_device
Ablation studies
----------------
In the previous example, we used the LoRA fine-tuned 8B teacher model and baseline 1B student model,
but we may want to experiment a bit with different configurations and hyperparameters.
For this tutorial, we are going to fine-tune on the :class:`~torchtune.datasets.alpaca_cleaned_dataset`
and evaluate the models on `truthfulqa_mc2 <https://github.com/EleutherAI/lm-evaluation-harness/tree/feff1b55c57993c4d42c8f913a22eeec395cd690/lm_eval/tasks/truthfulqa>`_,
`hellaswag <https://github.com/EleutherAI/lm-evaluation-harness/tree/517aadc/lm_eval/tasks/hellaswagd>`_
and `commonsense_qa <https://github.com/EleutherAI/lm-evaluation-harness/tree/b62b9bd/lm_eval/tasks/commonsense_qa>`_ tasks
through the EleutherAI `LM evaluation harness <https://github.com/EleutherAI/lm-evaluation-harness/tree/main>`_.
Let's take a look at the effects of:
#. Using a fine-tuned teacher model
#. Using a fine-tuned student model
#. Hyperparameter tuning of kd_ratio and learning rate
#. Teacher and student models with closer number of parameters
Using a fine-tuned teacher model
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
The default settings in the config uses the fine-tuned teacher model. Now, let's take a look at the
effects of not fine-tuning the teacher model first. To change the teacher model, you can modify the
``teacher_checkpointer`` in the config:
.. code-block:: yaml
teacher_checkpointer:
_component_: torchtune.training.FullModelHFCheckpointer
checkpoint_dir: /tmp/Meta-Llama-3.1-8B-Instruct/
checkpoint_files: [
model-00001-of-00004.safetensors,
model-00002-of-00004.safetensors,
model-00003-of-00004.safetensors,
model-00004-of-00004.safetensors
]
In the table below, we can see that standard fine-tuning of the 1B model achieves better accuracy
than the baseline 1B model. By using the fine-tuned 8B teacher model, we see comparable results
for truthfulqa and improvement for hellaswag and commonsense. When using the baseline 8B as a
teacher, we see improvement across all metrics, but lower than the other configurations.
.. image:: /_static/img/kd-finetune-teacher.png
Taking a look at the losses, using the baseline 8B as teacher results in a higher loss than
using the fine-tuned teacher model. The KD loss also remains relatively constant, suggesting
that the teacher model should have the same distributions as the transfer dataset.
Using a fine-tuned student model
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
For these experiments, let's take a look at the effects of KD when the student model is already
fine-tuned. In these experiments, we look at different combinations of baseline and fine-tuned 8B
and 1B models. To change the student model, you can first fine-tune the 1B model then modify the
student model checkpointer in the config:
.. code-block:: yaml
checkpointer:
_component_: torchtune.training.FullModelHFCheckpointer
checkpoint_dir: /tmp/Llama-3.2-1B-Instruct/
checkpoint_files: [
hf_model_0001_0.pt
]
Using the fine-tuned student model boosts accuracy even further for truthfulqa, but the accuracy
drops for hellaswag and commonsense. Using a fine-tuned teacher model and baseline student
model achieved the best results on hellaswag and commonsense dataset. Based on these findings,
the best configuration will change depending on which evaluation dataset and metric you are optimizing for.
.. image:: /_static/img/kd-finetune-student.png
Based on the loss graphs, using a fine-tuned teacher model results in a lower loss irrespective of
whether the student model is fine-tuned or not. It's also interesting to note that the class loss
starts to increase when using a fine-tuned student model.
Hyperparameter tuning: learning rate
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
By default, the config has the learning rate as :math:`3e^{-4}`, which is the same as the LoRA configs. For these experiments,
we changed the learning rate from as high as :math:`1e^{-3}` to as low as :math:`1e^{-5}`. To change the learning rate,
you can simply override the learning rate parameter using:
.. code-block:: bash
tune run knowledge_distillation_single_device --config llama3_2/8B_to_1B_KD_lora_single_device optimizer.lr=1e-3
Based on the results, the optimal learning rate changes depending on which metric you are optimizing for.
.. image:: /_static/img/kd-hyperparam-lr.png
Based on the loss graphs, all learning rates result in similar losses except for :math:`1e^{-5}`, which has a higher KD and class loss.
Hyperparameter tuning: KD ratio
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
In the config, we have the ``kd_ratio`` as 0.5, which gives even weightings to both the class and KD loss. In these experiments,
we look at the effects of different KD ratios, where 0 only uses the class loss and 1 only uses the KD loss.
Similar to changing the learning rate, the KD ratio can be adjusted using:
.. code-block:: bash
tune run knowledge_distillation_single_device --config llama3_2/8B_to_1B_KD_lora_single_device kd_ratio=0.25
Overall, the evaluation results are slightly better for higher KD ratios.
.. image:: /_static/img/kd-hyperparam-kd-ratio.png
Qwen2 1.5B to 0.5B
^^^^^^^^^^^^^^^^^^
The KD recipe can also be applied to different model families. Here we look at the effect of KD when the number of
parameters between the teacher and student models are closer. For this experiment, we used Qwen2 1.5B and Qwen2 0.5B, the configs for which can be found in
`qwen2/1.5_to_0.5B_KD_lora_single_device <https://github.com/pytorch/torchtune/blob/113807c2ce44791ac3354a40c93b87720c5bfc45/recipes/configs/qwen2/1.5_to_0.5B_KD_lora_single_device.yaml>`_
config. Here we see that training on the alpaca cleaned dataset only improves truthful_qa performance and drops the metrics for the other evaluation tasks.
For truthful_qa, KD improves the student model performance by 5.8% whereas fine-tuning improves performance by 1.3%.
.. image:: /_static/img/kd-qwen2-res.png
---
### Source Tutorials Lora Finetune (`docs/source/tutorials/lora_finetune.rst`)
.. _lora_finetune_label:
============================
Fine-Tuning Llama2 with LoRA
============================
This guide will teach you about `LoRA <https://arxiv.org/abs/2106.09685>`_, a parameter-efficient finetuning technique,
and show you how you can use torchtune to finetune a Llama2 model with LoRA.
If you already know what LoRA is and want to get straight to running
your own LoRA finetune in torchtune, you can jump to :ref:`LoRA finetuning recipe in torchtune<lora_recipe_label>`.
.. grid:: 2
.. grid-item-card:: :octicon:`mortar-board;1em;` What you will learn
* What LoRA is and how it saves memory during finetuning
* An overview of LoRA components in torchtune
* How to run a LoRA finetune using torchtune
* How to experiment with different LoRA configurations
.. grid-item-card:: :octicon:`list-unordered;1em;` Prerequisites
* Be familiar with :ref:`torchtune<overview_label>`
* Make sure to :ref:`install torchtune`
* Make sure you have downloaded the :ref:`Llama2-7B model weights<download_llama_label>`
What is LoRA?
-------------
`LoRA <https://arxiv.org/abs/2106.09685>`_ is an adapter-based method for
parameter-efficient finetuning that adds trainable low-rank decomposition matrices to different layers of a neural network,
then freezes the network's remaining parameters. LoRA is most commonly applied to
transformer models, in which case it is common to add the low-rank matrices
to some of the linear projections in each transformer layer's self-attention.
.. note::
If you're unfamiliar, check out these references for the `definition of rank <https://en.wikipedia.org/wiki/Rank_(linear_algebra)>`_
and discussion of `low-rank approximations <https://en.wikipedia.org/wiki/Low-rank_approximation>`_.
By finetuning with LoRA (as opposed to finetuning all model parameters),
you can expect to see memory savings due to a substantial reduction in the
number of parameters with gradients. When using an optimizer with momentum,
like `AdamW <https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html>`_,
you can expect to see further memory savings from the optimizer state.
.. note::
LoRA memory savings come primarily from gradient and optimizer states,
so if your model's peak memory comes in its :code:`forward()` method, then LoRA
may not reduce peak memory.
How does LoRA work?
-------------------
LoRA replaces weight update matrices with a low-rank approximation. In general, weight updates
for an arbitrary :code:`nn.Linear(in_dim,out_dim)` layer could have rank as high as
:code:`min(in_dim,out_dim)`. LoRA (and other related papers such as `Aghajanyan et al. <https://arxiv.org/abs/2012.13255>`_)
hypothesize that the `intrinsic dimension <https://en.wikipedia.org/wiki/Intrinsic_dimension>`_
of these updates during LLM fine-tuning can in fact be much lower.
To take advantage of this property, LoRA finetuning will freeze the original model,
then add a trainable weight update from a low-rank projection. More explicitly, LoRA trains two
matrices :code:`A` and :code:`B`. :code:`A` projects the inputs down to a much smaller rank (often four or eight in practice), and
:code:`B` projects back up to the dimension output by the original linear layer.
The image below gives a simplified representation of a single weight update step from a full finetune
(on the left) compared to a weight update step with LoRA (on the right). The LoRA matrices :code:`A` and :code:`B`
serve as an approximation to the full rank weight update in blue.
.. image:: /_static/img/lora_diagram.png
Although LoRA introduces a few extra parameters in the model :code:`forward()`, only the :code:`A` and :code:`B` matrices are trainable.
This means that with a rank :code:`r` LoRA decomposition, the number of gradients we need to store reduces
from :code:`in_dim*out_dim` to :code:`r*(in_dim+out_dim)`. (Remember that in general :code:`r`
is much smaller than :code:`in_dim` and :code:`out_dim`.)
For example, in the 7B Llama2's self-attention, :code:`in_dim=out_dim=4096` for the Q, K,
and V projections. This means a LoRA decomposition of rank :code:`r=8` will reduce the number of trainable
parameters for a given projection from :math:`4096 * 4096 \approx 15M` to :math:`8 * 8192 \approx 65K`, a
reduction of over 99%.
Let's take a look at a minimal implementation of LoRA in native PyTorch.
.. code-block:: python
import torch
from torch import nn
class LoRALinear(nn.Module):
def __init__(
self,
in_dim: int,
out_dim: int,
rank: int,
alpha: float,
dropout: float
):
# These are the weights from the original pretrained model
self.linear = nn.Linear(in_dim, out_dim, bias=False)
# These are the new LoRA params. In general rank << in_dim, out_dim
self.lora_a = nn.Linear(in_dim, rank, bias=False)
self.lora_b = nn.Linear(rank, out_dim, bias=False)
# Rank and alpha are commonly-tuned hyperparameters
self.rank = rank
self.alpha = alpha
# Most implementations also include some dropout
self.dropout = nn.Dropout(p=dropout)
# The original params are frozen, and only LoRA params are trainable.
self.linear.weight.requires_grad = False
self.lora_a.weight.requires_grad = True
self.lora_b.weight.requires_grad = True
def forward(self, x: torch.Tensor) -> torch.Tensor:
# This would be the output of the original model
frozen_out = self.linear(x)
# lora_a projects inputs down to the much smaller self.rank,
# then lora_b projects back up to the output dimension
lora_out = self.lora_b(self.lora_a(self.dropout(x)))
# Finally, scale by the alpha parameter (normalized by rank)
# and add to the original model's outputs
return frozen_out + (self.alpha / self.rank) * lora_out
There are some other details around initialization which we omit here, but if you'd like to know more
you can see our implementation in :class:`~torchtune.modules.peft.LoRALinear`.
Now that we understand what LoRA is doing, let's look at how we can apply it to our favorite models.
Applying LoRA to Llama2 models
------------------------------
With torchtune, we can easily apply LoRA to Llama2 with a variety of different configurations.
Let's take a look at how to construct Llama2 models in torchtune with and without LoRA.
.. code-block:: python
from torchtune.models.llama2 import llama2_7b, lora_llama2_7b
# Build Llama2 without any LoRA layers
base_model = llama2_7b()
# The default settings for lora_llama2_7b will match those for llama2_7b
# We just need to define which layers we want LoRA applied to.
# Within each self-attention, we can choose from ["q_proj", "k_proj", "v_proj", and "output_proj"].
# We can also set apply_lora_to_mlp=True or apply_lora_to_output=True to apply LoRA to other linear
# layers outside of the self-attention.
lora_model = lora_llama2_7b(lora_attn_modules=["q_proj", "v_proj"])
.. note::
Calling :func:`lora_llama_2_7b <torchtune.models.llama2.lora_llama2_7b>` alone will not handle the definition of which parameters are trainable.
See :ref:`below<setting_trainable_params>` for how to do this.
Let's inspect each of these models a bit more closely.
.. code-block:: bash
# Print the first layer's self-attention in the usual Llama2 model
>>> print(base_model.layers[0].attn)
MultiHeadAttention(
(q_proj): Linear(in_features=4096, out_features=4096, bias=False)
(k_proj): Linear(in_features=4096, out_features=4096, bias=False)
(v_proj): Linear(in_features=4096, out_features=4096, bias=False)
(output_proj): Linear(in_features=4096, out_features=4096, bias=False)
(pos_embeddings): RotaryPositionalEmbeddings()
)
# Print the same for Llama2 with LoRA weights
>>> print(lora_model.layers[0].attn)
MultiHeadAttention(
(q_proj): LoRALinear(
(dropout): Dropout(p=0.0, inplace=False)
(lora_a): Linear(in_features=4096, out_features=8, bias=False)
(lora_b): Linear(in_features=8, out_features=4096, bias=False)
)
(k_proj): Linear(in_features=4096, out_features=4096, bias=False)
(v_proj): LoRALinear(
(dropout): Dropout(p=0.0, inplace=False)
(lora_a): Linear(in_features=4096, out_features=8, bias=False)
(lora_b): Linear(in_features=8, out_features=4096, bias=False)
)
(output_proj): Linear(in_features=4096, out_features=4096, bias=False)
(pos_embeddings): RotaryPositionalEmbeddings()
)
Notice that our LoRA model's layer contains additional weights in the Q and V projections,
as expected. Additionally, inspecting the type of :code:`lora_model` and
:code:`base_model`, would show that they are both instances of the same :class:`~torchtune.modules.TransformerDecoder`.
(Feel free to verify this for yourself.)
Why does this matter? torchtune makes it easy to load checkpoints for LoRA directly from our Llama2
model without any wrappers or custom checkpoint conversion logic.
.. code-block:: python
# Assuming that base_model already has the pretrained Llama2 weights,
# this will directly load them into your LoRA model without any conversion necessary.
lora_model.load_state_dict(base_model.state_dict(), strict=False)
.. note::
Whenever loading weights with :code:`strict=False`, you should verify that any missing or extra keys in
the loaded :code:`state_dict` are as expected. torchtune's LoRA recipes do this by default via
:func:`validate_missing_and_unexpected_for_lora() <torchtune.modules.peft.validate_missing_and_unexpected_for_lora>`.
Once we've loaded the base model weights, we also want to set only LoRA parameters to trainable.
.. _setting_trainable_params:
.. code-block:: python
from torchtune.modules.peft.peft_utils import get_adapter_params, set_trainable_params
# Fetch all params from the model that are associated with LoRA.
lora_params = get_adapter_params(lora_model)
# Set requires_grad=True on lora_params, and requires_grad=False on all others.
set_trainable_params(lora_model, lora_params)
# Print the total number of parameters
total_params = sum([p.numel() for p in lora_model.parameters()])
trainable_params = sum([p.numel() for p in lora_model.parameters() if p.requires_grad])
print(
f"""
{total_params} total params,
{trainable_params}" trainable params,
{(100.0 * trainable_params / total_params):.2f}% of all params are trainable.
"""
)
6742609920 total params,
4194304 trainable params,
0.06% of all params are trainable.
.. note::
If you are directly using the LoRA recipe (as detailed :ref:`here<lora_recipe_label>`), you need only pass the
relevant checkpoint path. Loading model weights and setting trainable parameters will be taken care
of in the recipe.
.. _lora_recipe_label:
LoRA finetuning recipe in torchtune
-----------------------------------
Finally, we can put it all together and finetune a model using torchtune's `LoRA recipe <https://github.com/pytorch/torchtune/blob/48626d19d2108f92c749411fbd5f0ff140023a25/recipes/lora_finetune.py>`_.
Make sure that you have first downloaded the Llama2 weights and tokenizer by following :ref:`these instructions<download_llama_label>`.
You can then run the following command to perform a LoRA finetune of Llama2-7B with two GPUs (each having VRAM of at least 16GB):
.. code-block:: bash
tune run --nnodes 1 --nproc_per_node 2 lora_finetune_distributed --config llama2/7B_lora
.. note::
Make sure to point to the location of your Llama2 weights and tokenizer. This can be done
either by adding :code:`checkpointer.checkpoint_files=[my_model_checkpoint_path] tokenizer_checkpoint=my_tokenizer_checkpoint_path`
or by directly modifying the :code:`7B_lora.yaml` file. See our "":ref:`config_tutorial_label`" recipe
for more details on how you can easily clone and modify torchtune configs.
.. note::
You can modify the value of :code:`nproc_per_node` depending on (a) the number of GPUs you have available,
and (b) the memory constraints of your hardware.
The preceding command will run a LoRA finetune with torchtune's factory settings, but we may want to experiment a bit.
Let's take a closer look at some of the :code:`lora_finetune_distributed` config.
.. code-block:: yaml
# Model Arguments
model:
_component_: lora_llama2_7b
lora_attn_modules: ['q_proj', 'v_proj']
lora_rank: 8
lora_alpha: 16
...
We see that the default is to apply LoRA to Q and V projections with a rank of 8.
Some experiments with LoRA have found that it can be beneficial to apply LoRA to all linear layers in
the self-attention, and to increase the rank to 16 or 32. Note that this is likely to increase our max memory,
but as long as we keep :code:`rank<`_
floating-point format. This can be done via the command:
.. code-block:: bash
tune run lora_finetune_single_device --config llama2/7B_lora_single_device
On a single device, we may need to be more cognizant of our peak memory. Let's run a few experiments
to see our peak memory during a finetune. We will experiment along two axes:
first, which model layers have LoRA applied, and second, the rank of each LoRA layer. (We will scale
alpha in parallel to LoRA rank, as discussed above.)
To compare the results of our experiments, we can evaluate our models on `truthfulqa_mc2 <https://github.com/sylinrl/TruthfulQA>`_, a task from
the `TruthfulQA <https://arxiv.org/abs/2109.07958>`_ benchmark for language models. For more details on how to run this and other evaluation tasks
with torchtune's EleutherAI evaluation harness integration, see our :ref:`End-to-End Workflow Tutorial <eval_harness_label>`.
Previously, we only enabled LoRA for the linear layers in each self-attention module, but in fact there are other linear
layers we can apply LoRA to: MLP layers and our model's final output projection. Note that for Llama-2-7B the final output
projection maps to the vocabulary dimension (32000 instead of 4096 as in the other linear layers), so enabling LoRA for this layer will increase
our peak memory a bit more than the other layers. We can make the following changes to our config:
.. code-block:: yaml
# Model Arguments
model:
_component_: lora_llama2_7b
lora_attn_modules: ['q_proj', 'k_proj', 'v_proj', 'output_proj']
apply_lora_to_mlp: True
apply_lora_to_output: True
...
.. note::
All the finetuning runs below use the `llama2/7B_lora_single_device <https://github.com/pytorch/torchtune/blob/main/recipes/configs/llama2/7B_lora_single_device.yaml>`_
config, which has a default batch size of 2. Modifying the batch size (or other hyperparameters, e.g. the optimizer) will impact both peak memory
and final evaluation results.
.. list-table::
:widths: 25 25 25 25 25
:header-rows: 1
* - LoRA Layers
- Rank
- Alpha
- Peak Memory
- Accuracy (truthfulqa_mc2)
* - Q and V only
- 8
- 16
- **15.57 GB**
- 0.475
* - all layers
- 8
- 16
- 15.87 GB
- 0.508
* - Q and V only
- 64
- 128
- 15.86 GB
- 0.504
* - all layers
- 64
- 128
- 17.04 GB
- **0.514**
We can see that our baseline settings give the lowest peak memory, but our evaluation performance is relatively lower.
By enabling LoRA for all linear layers and increasing the rank to 64, we see almost a 4% absolute improvement
in our accuracy on this task, but our peak memory also increases by about 1.4GB. These are just a couple simple
experiments; we encourage you to run your own finetunes to find the right tradeoff for your particular setup.
Additionally, if you want to decrease your model's peak memory even further (and still potentially achieve similar
model quality results), you can check out our :ref:`QLoRA tutorial<qlora_finetune_label>`.
---
### Source Tutorials Memory Optimizations (`docs/source/tutorials/memory_optimizations.rst`)
.. _memory_optimization_overview_label:
============================
Memory Optimization Overview
============================
**Author**: `Salman Mohammadi <https://github.com/SalmanMohammadi>`_
torchtune comes with a host of plug-and-play memory optimization components which give you lots of flexibility
to ``tune`` our recipes to your hardware. This page provides a brief glossary of these components and how you might use them.
To make things easy, we've summarized these components in the following table:
.. csv-table:: Memory optimization components
:header: "Component", "When to use?"
:widths: auto
":ref:`glossary_precision`", "You'll usually want to leave this as its default ``bfloat16``. It uses 2 bytes per model parameter instead of 4 bytes when using ``float32``."
":ref:`glossary_act_ckpt`", "Use when you're memory constrained and want to use a larger model, batch size or context length. Be aware that it will slow down training speed."
":ref:`glossary_act_off`", "Similar to activation checkpointing, this can be used when memory constrained, but may decrease training speed. This **should** be used alongside activation checkpointing."
":ref:`glossary_grad_accm`", "Helpful when memory-constrained to simulate larger batch sizes. Not compatible with optimizer in backward. Use it when you can already fit at least one sample without OOMing, but not enough of them."
":ref:`glossary_low_precision_opt`", "Use when you want to reduce the size of the optimizer state. This is relevant when training large models and using optimizers with momentum, like Adam. Note that lower precision optimizers may reduce training stability/accuracy."
":ref:`glossary_opt_in_bwd`", "Use it when you have large gradients and can fit a large enough batch size, since this is not compatible with ``gradient_accumulation_steps``."
":ref:`glossary_cpu_offload`", "Offloads optimizer states and (optionally) gradients to CPU, and performs optimizer steps on CPU. This can be used to significantly reduce GPU memory usage at the cost of CPU RAM and training speed. Prioritize using it only if the other techniques are not enough."
":ref:`glossary_lora`", "When you want to significantly reduce the number of trainable parameters, saving gradient and optimizer memory during training, and significantly speeding up training. This may reduce training accuracy"
":ref:`glossary_qlora`", "When you are training a large model, since quantization will save 1.5 bytes * (# of model parameters), at the potential cost of some training speed and accuracy."
":ref:`glossary_dora`", "a variant of LoRA that may improve model performance at the cost of slightly more memory."
.. note::
In its current state, this tutorial is focused on single-device optimizations. Check in soon as we update this page
for the latest memory optimization features for distributed fine-tuning.
.. _glossary_precision:
Model Precision
---------------
*What's going on here?*
We use the term "precision" to refer to the underlying data type used to represent the model and optimizer parameters.
We support two data types in torchtune:
.. note::
We recommend diving into Sebastian Raschka's `blogpost on mixed-precision techniques <https://sebastianraschka.com/blog/2023/llm-mixed-precision-copy.html>`_
for a deeper understanding of concepts around precision and data formats.
* ``fp32``, commonly referred to as "full-precision", uses 4 bytes per model and optimizer parameter.
* ``bfloat16``, referred to as "half-precision", uses 2 bytes per model and optimizer parameter - effectively half
the memory of ``fp32``, and also improves training speed. Generally, if your hardware supports training with ``bfloat16``,
we recommend using it - this is the default setting for our recipes.
.. note::
Another common paradigm is "mixed-precision" training: where model weights are in ``bfloat16`` (or ``fp16``), and optimizer
states are in ``fp32``. Currently, we don't support mixed-precision training in torchtune.
*Sounds great! How do I use it?*
Simply use the ``dtype`` flag or config entry in all our recipes! For example, to use half-precision training in ``bf16``,
set ``dtype=bf16``.
.. _glossary_act_ckpt:
Activation Checkpointing
------------------------
*What's going on here?*
The relevant section in the `PyTorch documentation <https://pytorch.org/docs/stable/checkpoint.html>`_ explains this concept well.
To quote:
Activation checkpointing is a technique that trades compute for memory.
Instead of keeping tensors needed for backward alive until they are used in
gradient computation during backward, forward computation in checkpointed
regions omits saving tensors for backward and recomputes them during the backward pass.
This setting is helpful for when you're memory-constrained, especially due to larger batch sizes or longer context lengths.
However, these savings in memory come at the cost of training speed (i.e. tokens-per-second),
and in most cases training can slow down quite a bit as a result of this activation recomputation.
*Sounds great! How do I use it?*
To enable activation checkpointing, use ``enable_activation_checkpointing=True``.
.. _glossary_act_off:
Activation Offloading
---------------------
*What's going on here?*
You may have just read about activation checkpointing! Similar to checkpointing, offloading is a memory
efficiency technique that allows saving GPU VRAM by temporarily moving activations to CPU and bringing
them back when needed in the backward pass.
See `PyTorch autograd hook tutorial <https://pytorch.org/tutorials/intermediate/autograd_saved_tensors_hooks_tutorial.html#saving-tensors-to-cpu>`_
for more details about how this is implemented through :func:`torch.autograd.graph.saved_tensors_hooks`.
This setting is especially helpful for larger batch sizes, or longer context lengths when you're memory constrained.
While of course it takes runtime and resources to move Tensors from GPU to CPU and back, the implementation in
torchtune uses multiple CUDA streams (when available) in order to overlap the extra communication with the computation
to hide the extra runtime. As the communication workload is variable depending on the number and size of tensors being
offloaded, we do not recommend using it unless :ref:`glossary_act_ckpt` is also enabled, in which case only the checkpointed
tensors will be offloaded.
*Sounds great! How do I use it?*
To enable activation offloading, use the ``enable_activation_offloading`` config entry or flag
in our lora finetuning single device recipe, e.g. ``enable_activation_offloading=True``. To allow
usage of streams, make sure you are on a torch version equal to or later than PyTorch.
.. _glossary_grad_accm:
Gradient Accumulation
---------------------
*What's going on here?*
Gradient accumulation allows you to simulate large batch sizes by *accumulating* gradients over several
batches before updating model parameters using the optimizer. Concretely, the total number of samples used
for a gradient update is when using gradient accumulation is:
``total_batch_size = batch_size * gradient_accumulation_steps``
For example: with ``batch_size=1`` and ``gradient_accumulation_steps=32`` we get a total batch size of 32.
.. note::
For other components in torchtune which use "steps", such as :ref:`metric logging <metric_logging_label>`, or
:func:`learning rate schedulers <torchtune.training.lr_schedulers.get_cosine_schedule_with_warmup>`, a "step" is counted as a
single update to model parameters, rather than a single model forward pass with the data.
Suppose ``gradient_accumulation_steps = 4`` and ``log_every_n_steps = 10``.
Metrics would be logged every 10 global steps, which translates to every 40 model forward passes.
For this reason, metric logging will appear less frequently when training with gradient accumulation,
and progress bars may update more slowly.
If you're using one of our distributed recipes, simply multiply by the number of devices:
``total_batch_size = batch_size * gradient_accumulation_steps * num_devices``
Gradient accumulation is especially useful when you can fit at least one sample in your GPU. In this case, artificially increasing the batch by
accumulating gradients might give you faster training speeds than using other memory optimization techniques that trade-off memory for speed, like :ref:`activation checkpointing <glossary_act_ckpt>`.
*Sounds great! How do I use it?*
All of our finetuning recipes support simulating larger batch sizes by accumulating gradients. Just set the
``gradient_accumulation_steps`` flag or config entry.
.. note::
Gradient accumulation should always be set to 1 when :ref:`fusing the optimizer step into the backward pass <glossary_opt_in_bwd>`.
Optimizers
----------
.. _glossary_low_precision_opt:
Lower Precision Optimizers
^^^^^^^^^^^^^^^^^^^^^^^^^^
*What's going on here?*
In addition to :ref:`reducing model and optimizer precision <glossary_precision>` during training, we can further reduce precision in our optimizer states.
All of our recipes support lower-precision optimizers from the `torchao <https://github.com/pytorch/ao/tree/main/torchao/prototype/low_bit_optim>`_ library.
For single device recipes, we also support `bitsandbytes <https://huggingface.co/docs/bitsandbytes/main/en/index>`_.
A good place to start might be the :class:`torchao.prototype.low_bit_optim.AdamW8bit` and :class:`bitsandbytes.optim.PagedAdamW8bit` optimizers.
Both reduce memory by quantizing the optimizer state dict. Paged optimizers will also offload to CPU if there isn't enough GPU memory available. In practice,
you can expect higher memory savings from bnb's PagedAdamW8bit but higher training speed from torchao's AdamW8bit.
*Sounds great! How do I use it?*
To use this in your recipes, make sure you have installed torchao (``pip install torchao``) or bitsandbytes (``pip install bitsandbytes``). Then, enable
a low precision optimizer using the :ref:`cli_label`:
.. code-block:: bash
tune run <RECIPE> --config <CONFIG> \
optimizer=torchao.prototype.low_bit_optim.AdamW8bit
.. code-block:: bash
tune run <RECIPE> --config <CONFIG> \
optimizer=bitsandbytes.optim.PagedAdamW8bit
or by directly :ref:`modifying a config file<config_tutorial_label>`:
.. code-block:: yaml
optimizer:
_component_: bitsandbytes.optim.PagedAdamW8bit
lr: 2e-5
.. _glossary_opt_in_bwd:
Fusing Optimizer Step into Backward Pass
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
*What's going on here?*
Stateful optimizers (e.g. optimizers which use momentum) are the default in modern deep learning due to their stable convergence properties.
However, maintaining a state of gradient statistics comes at the cost of additional memory usage. An immediate alternative might be to
turn to stateless optimizers such as `stochastic gradient descent <https://pytorch.org/docs/stable/generated/torch.optim.SGD.html>`_
without momentum, which don't require any additional memory usage, but will likely result in worse convergence during training.
Can we find a middle ground here? Let's consider a technique which enables the use of "stateful" optimizers such as `AdamW <https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html>`_
without the memory overhead of gradient statistics, and without sacrificing their desirable convergence properties.
How is this possible, you might ask? By *completely removing the buffer of gradients* which are stored by the optimizer during its ``step()``.
To understand how this works, we encourage you to read through the relevant PyTorch tutorial on this concept:
`How to save memory by fusing the optimizer step into the backward pass <https://pytorch.org/tutorials/intermediate/optimizer_step_in_backward_tutorial.html>`_.
*Sounds great! How do I use it?*
In torchtune, you can enable this feature using :class:`OptimizerInBackward`, which wraps around a normal PyTorch optimizer like so and can be used
as a drop-in replacement for your normal Optimizer:
.. code-block:: python
OptimizerInBackward(
model.parameters(),
torch.optim.AdamW,
lr=0.01,
)
In order to use this OOTB with a torchtune recipe and config, simply specify your chosen optimizer and then add the key ``optimizer_in_bwd: True``, which will
construct a :class:`OptimizerInBackward` from the base optimizer class.
.. code-block:: yaml
optimizer:
_component_: torch.optim.AdamW
lr: 0.01
optimizer_in_bwd: True
This feature works best when using a stateful optimizer with a model with a lot of parameters, and when you don't need to use :ref:`gradient accumulation <glossary_grad_accm>`.
You won't see meaningful impact when finetuning LoRA recipes, since in this case the number of parameters being updated are small.
.. _glossary_cpu_offload:
Offloading Optimizer/Gradient states to CPU
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
*What's going on here?*
We've mentioned above the concept of optimizer states - memory used by the stateful optimizers to maintain a state of gradient statistics, and
model gradients - tensors used to store gradients when we perform model backwards passes. We support using CPU offloading in our single-device recipes
through the `CPUOffloadOptimizer <https://github.com/pytorch/ao/tree/main/torchao/prototype/low_bit_optim#optimizer-cpu-offload>`_ from ``torchao``.
This optimizer can wrap any base optimizer and works by keeping the optimizer states and performing the optimizer step on CPU, thus reducing
GPU memory usage by the size of the optimizer states. Additionally, we can also offload gradients to the CPU by using `offload_gradients=True`.
If finetuning on a single-device, another option is to use the ``PagedAdamW8bit`` from bitsandbytes, mentioned :ref:`above <glossary_low_precision_opt>`, which will *only* offload to CPU
when there is not enough GPU available.
*Sounds great! How do I use it?*
To use this optimizer in your recipes, set the ``optimizer`` key in your config to :class:`torchao.prototype.low_bit_optim.CPUOffloadOptimizer`, which
will use the :class:`torch.optim.AdamW` optimizer with ``fused=True`` as the base optimizer. For example, to use this optimizer to offload
both optimizer states and gradients to CPU:
.. code-block:: bash
tune run <RECIPE> --config <CONFIG> \
optimizer=optimizer=torchao.prototype.low_bit_optim.CPUOffloadOptimizer \
optimizer.offload_gradients=True \
lr=4e-5
or by directly :ref:`modifying a config file<config_tutorial_label>`:
.. code-block:: yaml
optimizer:
_component_: torchao.prototype.low_bit_optim.CPUOffloadOptimizer
offload_gradients: True
# additional key-word arguments can be passed to torch.optim.AdamW
lr: 4e-5
or using it directly in your code, which allows you to change the base optimizer:
.. code-block:: python
from torchao.prototype.low_bit_optim import CPUOffloadOptimizer
from torch.optim import Adam
optimizer = CPUOffloadOptimizer(
model.parameters(), # your model here
Adam,
lr=1e-5,
fused=True
)
Some helpful hints from the ``torchao`` `CPUOffloadOptimizer page <https://github.com/pytorch/ao/tree/main/torchao/prototype/low_bit_optim#optimizer-cpu-offload>`_:
* The CPU optimizer step is often the bottleneck when optimizer CPU offload is used. To minimize the slowdown, it is recommended to (1) use full ``bf16`` training so that parameters, gradients, and optimizer states are in ``bf16``; and (2) give GPU more work per optimizer step to amortize the offloading time (e.g. larger batch size with activation checkpointing, gradient accumulation).
* Gradient accumulation should always be set to 1 when ``offload_gradients=True``, as gradients are cleared on GPU every backward pass.
* This optimizer works by keeping a copy of parameters and pre-allocating gradient memory on CPU. Therefore, expect your RAM usage to increase by 4x model size.
* This optimizer is only supported for single-device recipes. To use CPU-offloading in distributed recipes, use ``fsdp_cpu_offload=True`` instead. See :class:`torch.distributed.fsdp.FullyShardedDataParallel` for more details and `FSDP1 vs FSDP2 <https://github.com/pytorch/torchtitan/blob/main/docs/fsdp.md>`_ to see how they differ.
.. _glossary_peft:
Parameter Efficient Fine-Tuning (PEFT)
--------------------------------------
.. _glossary_lora:
Low Rank Adaptation (LoRA)
^^^^^^^^^^^^^^^^^^^^^^^^^^
*What's going on here?*
You can read our tutorial on :ref:`finetuning Llama2 with LoRA<lora_finetune_label>` to understand how LoRA works, and how to use it.
Simply stated, LoRA greatly reduces the number of trainable parameters, thus saving significant gradient and optimizer
memory during training.
*Sounds great! How do I use it?*
You can finetune using any of our recipes with the ``lora_`` prefix, e.g. :ref:`lora_finetune_single_device<lora_finetune_recipe_label>`. These recipes utilize
LoRA-enabled model builders, which we support for all our models, and also use the ``lora_`` prefix, e.g.
the :func:`torchtune.models.llama3.llama3` model has a corresponding :func:`torchtune.models.llama3.lora_llama3`.
We aim to provide a comprehensive set of configurations to allow you to get started with training with LoRA quickly,
just specify any config with ``_lora`` in its name, e.g:
.. code-block:: bash
tune run lora_finetune_single_device --config llama3/8B_lora_single_device
There are two sets of parameters to customize LoRA to suit your needs. Firstly, the parameters which control
which linear layers LoRA should be applied to in the model:
* ``lora_attn_modules: list[str]`` accepts a list of strings specifying which layers of the model to apply
LoRA to:
* ``q_proj`` applies LoRA to the query projection layer.
* ``k_proj`` applies LoRA to the key projection layer.
* ``v_proj`` applies LoRA to the value projection layer.
* ``output_proj`` applies LoRA to the attention output projection layer.
Whilst adding more layers to be fine-tuned may improve model accuracy,
this will come at the cost of increased memory usage and reduced training speed.
* ``apply_lora_to_mlp: Bool`` applies LoRA to the MLP in each transformer layer.
* ``apply_lora_to_output: Bool`` applies LoRA to the model's final output projection.
This is usually a projection to vocabulary space (e.g. in language models), but
other modelling tasks may have different projections - classifier models will project
to the number of classes, for example
.. note::
Models which use tied embeddings (such as Gemma and Qwen2 1.5B and 0.5B) for the
final output projection do not support ``apply_lora_to_output``.
These are all specified under the ``model`` flag or config entry, i.e:
.. code-block:: bash
tune run lora_finetune_single_device --config llama3/8B_lora_single_device \
model.apply_lora_to_mlp=True \
model.lora_attn_modules=["q_proj","k_proj","v_proj","output_proj"]
.. code-block:: yaml
model:
_component_: torchtune.models.llama3.lora_llama3_8b
apply_lora_to_mlp: True
model.lora_attn_modules: ["q_proj", "k_proj", "v_proj","output_proj"]
Secondly, parameters which control the scale of the impact of LoRA on the model:
* ``lora_rank: int`` affects the scale of the LoRA decomposition, where ``lora_rank << in_dim`` and ``lora_rank << out_dim``
\- the dimensions of an arbitrary linear layer in the model. Concretely, ``lora_rank`` reduces the number of gradients stored
in a linear fashion from ``in_dim * out_dim`` to ``lora_rank * (in_dim + out_dim)``. Typically, we have ``lora_rank in [8, 256]``.
* ``lora_alpha: float`` affects the magnitude of the LoRA updates. A larger alpha results in larger updates to the base model weights
, potentially at the cost of training stability, conversely, smaller alpha can stabilize training at the cost of slower learning.
We provide default settings for these parameters which we've tested with all of our models, but we encourage you to adjust them
to your specific use case. Typically, one jointly changes ``lora_rank`` and ``lora_alpha`` together, where ``lora_alpha ~= 2*lora_rank``.
* ``lora_dropout`` introduces dropout in the LoRA layers to help regularize training. We default to 0.0 for all of our models.
As above, these parameters are also specified under the ``model`` flag or config entry:
.. code-block:: bash
tune run lora_finetune_single_device --config llama3/8B_lora_single_device \
model.apply_lora_to_mlp=True \
model.lora_attn_modules=["q_proj","k_proj","v_proj","output_proj"] \
model.lora_rank=32 \
model.lora_alpha=64
.. code-block:: yaml
model:
_component_: torchtune.models.llama3.lora_llama3_8b
apply_lora_to_mlp: True
lora_attn_modules: ["q_proj", "k_proj", "v_proj","output_proj"]
lora_rank: 32
lora_alpha: 64
.. note::
To get a deeper sense of how LoRA parameters affect memory usage during training,
see the :ref:`relevant section in our Llama2 LoRA tutorial<lora_tutorial_memory_tradeoff_label>`.
.. _glossary_qlora:
Quantized Low Rank Adaptation (QLoRA)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
*What's going on here?*
`QLoRA <https://arxiv.org/abs/2305.14314>`_ is a memory enhancement on top of `LoRA <https://arxiv.org/abs/2106.09685>`_
that maintains the frozen model parameters from LoRA in 4-bit quantized precision, thereby reducing memory usage.
This is enabled through a novel 4-bit NormalFloat (NF4) data type proposed by the authors, which allows for 4-8x less
parameter memory usage whilst retaining model accuracy. You can read our tutorial on :ref:`finetuning Llama2 with QLoRA<qlora_finetune_label>`
for a deeper understanding of how it works.
When considering using QLoRA to reduce memory usage, it's worth noting that QLoRA is slower than LoRA and may not be worth it if
the model you are finetuning is small. In numbers, QLoRA saves roughly 1.5 bytes * (# of model parameters). Also, although QLoRA quantizes the model,
it minimizes accuracy degradation by up-casting quantized parameters to the original higher precision datatype during model forward passes - this up-casting may incur penalties to training speed.
The :ref:`relevant section <qlora_compile_label>` in our QLoRA tutorial demonstrates the usage of ``torch.compile`` to address this by speeding up training.
*Sounds great! How do I use it?*
You can finetune using QLoRA with any of our LoRA recipes, i.e. recipes with the ``lora_`` prefix, e.g. :ref:`lora_finetune_single_device<lora_finetune_recipe_label>`. These recipes utilize
QLoRA-enabled model builders, which we support for all our models, and also use the ``qlora_`` prefix, e.g.
the :func:`torchtune.models.llama3.llama3_8b` model has a corresponding :func:`torchtune.models.llama3.qlora_llama3_8b`.
We aim to provide a comprehensive set of configurations to allow you to get started with training with QLoRA quickly,
just specify any config with ``_qlora`` in its name.
All the rest of the LoRA parameters remain the same for QLoRA - check out the section above on :ref:`LoRA <glossary_lora>`
to see how to configure these parameters.
To configure from the command line:
.. code-block:: bash
tune run lora_finetune_single_device --config llama3/8B_qlora_single_device \
model.apply_lora_to_mlp=True \
model.lora_attn_modules=["q_proj","k_proj","v_proj"] \
model.lora_rank=32 \
model.lora_alpha=64
or, by modifying a config:
.. code-block:: yaml
model:
_component_: torchtune.models.qlora_llama3_8b
apply_lora_to_mlp: True
lora_attn_modules: ["q_proj", "k_proj", "v_proj"]
lora_rank: 32
lora_alpha: 64
.. _glossary_dora:
Weight-Decomposed Low-Rank Adaptation (DoRA)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
*What's going on here?*
`DoRA <https://arxiv.org/abs/2402.09353>`_ is another PEFT technique which builds on-top of LoRA by
further decomposing the pre-trained weights into two components: magnitude and direction. The magnitude component
is a scalar vector that adjusts the scale, while the direction component corresponds to the original LoRA decomposition and
updates the orientation of weights.
DoRA adds a small overhead to LoRA training due to the addition of the magnitude parameter, but it has been shown to
improve the performance of LoRA, particularly at low ranks.
*Sounds great! How do I use it?*
Much like LoRA and QLoRA, you can finetune using DoRA with any of our LoRA recipes. We use the same model builders for LoRA
as we do for DoRA, so you can use the ``lora_`` version of any model builder with ``use_dora=True``. For example, to finetune
:func:`torchtune.models.llama3.llama3_8b` with DoRA, you would use :func:`torchtune.models.llama3.lora_llama3_8b` with ``use_dora=True``:
.. code-block:: bash
tune run lora_finetune_single_device --config llama3/8B_lora_single_device \
model.use_dora=True
.. code-block:: yaml
model:
_component_: torchtune.models.lora_llama3_8b
use_dora: True
Since DoRA extends LoRA, the parameters for :ref:`customizing LoRA <glossary_lora>` are identical. You can also quantize the base model weights like in :ref:`glossary_qlora` by using ``quantize=True`` to reap
even more memory savings!
.. code-block:: bash
tune run lora_finetune_single_device --config llama3/8B_lora_single_device \
model.apply_lora_to_mlp=True \
model.lora_attn_modules=["q_proj","k_proj","v_proj"] \
model.lora_rank=16 \
model.lora_alpha=32 \
model.use_dora=True \
model.quantize_base=True
.. code-block:: yaml
model:
_component_: torchtune.models.lora_llama3_8b
apply_lora_to_mlp: True
lora_attn_modules: ["q_proj", "k_proj", "v_proj"]
lora_rank: 16
lora_alpha: 32
use_dora: True
quantize_base: True
.. note::
Under the hood, we've enabled DoRA by adding the :class:`~torchtune.modules.peft.DoRALinear` module, which we swap
out for :class:`~torchtune.modules.peft.LoRALinear` when ``use_dora=True``.
.. _glossary_distrib:
.. TODO
.. Distributed
.. -----------
.. .. _glossary_fsdp:
.. Fully Sharded Data Parallel (FSDP)
.. ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.. All our ``_distributed`` recipes use `FSDP <https://pytorch.org/docs/stable/fsdp.html>`.
.. .. _glossary_fsdp2:
---
### Source Tutorials Multinode (`docs/source/tutorials/multinode.rst`)
.. _multinode_tutorial:
=====================
Multi-node finetuning
=====================
Congratulations! You've finally escaped the struggles of being "GPU poor" and now have access to a multi-node setup.
You can bid farewell to the days of sweating over memory-efficient optimizations, but get ready for new challenges as you navigate the complexities of distributed computing.
.. grid:: 2
.. grid-item-card:: :octicon:`mortar-board;1em;` You will learn:
* Why multi-node training is useful
* How to set up the torchtune package on a SLURM cluster
* How to fine-tune a Llama3.3 70B model w/ full parameter updates (not LoRA)
.. grid-item-card:: :octicon:`list-unordered;1em;` Prerequisites
* Be familiar with distributed training in torchtune
* Already know basic SLURM commands
.. _advantages_multi_node_label:
Advantages of multi-node training
---------------------------------
More machines means more memory! This is cool for several reasons:
1. **Bigger models**: With more memory, you can train larger models such as `Llama3.1 405B <https://ai.meta.com/blog/meta-llama-3-1/>`_, `Deepseek-V3 <https://www.deepseek.com/>`_, and more.
2. **Longer data**: For many fine-tuning tasks like writing code, it's helpful to have long context lengths; however longer context length means more memory needed for activations.
3. **Higher quality**: With more memory, you can do full parameter updates (not LoRA) and use optimizers like `AdamW <https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html>`_ (not low-precision optimizers), both of which can potentially improve the quality of your training.
4. **Faster training**: With the ability to fit more data in memory, you can use higher batch sizes *and* turn off memory optimizations like :ref:`activation checkpointing<glossary_act_ckpt>` thereby decreasing the time it takes for training to complete.
.. note::
**Low inter-node bandwidth & FSDP** We utilize PyTorch's **Fully Sharded Data Parallel** to distribute models over multiple devices. In order to distribute training, FSDP runs an `all-gather <https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html#allgather>`_ operation
for each forward pass and an all-gather (usually) plus a `reduce-scatter <https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html#reducescatter>`_ operation for each backwards pass. These operations (usually) block training from continuing until completed and with a slow
inter-node connection, training speed may be reduced. For more on this, please refer to `this Github Issue <https://github.com/pytorch/pytorch/issues/102434>`_.
Training Llama3.3 70B on 2 nodes
--------------------------------
Let's get training! We'll be utilizing a common cluster workflow manager called `SLURM <https://slurm.schedmd.com/documentation.html>`_ and assume you have a decent working knowledge of SLURM for this tutorial.
First, we need to install torchtune. Although pretty much as straightforward as the :ref:`normal install instructions`,
it's recommended that you install the package into a virtual environment that is accessible from all nodes in your cluster like a shared filesystem.
Next, we need to download the `Llama3.3 70B <https://huggingface.co/meta-llama/Llama-3.3-70B-Instruct>`_ model to your shared filesystem. You'll need to make sure you have the correct credentials following the steps
outlined :ref:`here<tune_download_label>`.
.. code-block:: bash
$ tune download meta-llama/Llama-3.3-70B-Instruct --ignore-patterns "consolidated/*.pth" --output-dir SHARED_FS/Llama-3.3-70B-Instruct
Now that we have a downloaded model, let's check out our example SLURM bash script.
.. literalinclude:: ../../../recipes/full_finetune_multinode.slurm
**There's a lot of information in this script but here are the high-level parts:**
* We utilize SLURM specific commands like number of nodes, tasks, CPUs available, etc.
* We are using `torchrun <https://pytorch.org/docs/stable/elastic/run.html>`_ and the `full_finetune_distributed <https://github.com/pytorch/torchtune/blob/main/recipes/full_finetune_distributed.py>`_ recipe to train just like on single node
* You can consider several cluster-specific environment variables (``NCCL_BUFFSIZE``, ``NCCL_DEBUG``, ``FI_PROVIDER``, etc.) in order to maximize GPU utilization, debug, and more.
.. note::
We may need to explicitly set the network interface for distributed backends. You can read more about `PyTorch distributed backends here <https://pytorch.org/docs/stable/distributed.html#common-environment-variables>`_
but it's also helpful to know that you can find your network interface by running `ipconfig <https://en.wikipedia.org/wiki/Ipconfig#:~:text=ipconfig%20(standing%20for%20%22Internet%20Protocol,ipconfig>`_ from a specific node.
After we update the shared filesystem in the bash script, we can launch using `sbatch <https://slurm.schedmd.com/sbatch.html>`_.
.. code-block:: bash
sbatch full_finetune_multinode.slurm
And the output of `squeue <https://slurm.schedmd.com/squeue.html>`_ should show our job running:
.. code-block:: bash
$ squeue
JOBID PARTITION NAME USER ST TIME NODES NODELIST(REASON)
1 train torchtun slurm R 0:03 2 slurm-worker-[1-2]
Once training has completed, which should take roughly seven minutes in total (880 tok/s) with the default config, we can follow the :ref:`instructions here` in order to upload our beautiful new model to the Hugging Face Hub!
Future development
------------------
We've covered the basics of how to launch a fine-tuning job with SLURM on two nodes with FSDP. There's still more things we're cooking up,
including...
**2D parallelism**: Utilizing both FSDP *and* tensor parallelism in what is commonly referred to as `2D parallelism <https://pytorch.org/tutorials/intermediate/TP_tutorial.html>`_ will decrease memory requirements even further, allowing us to lean even harder
into the advantages listed :ref:`above<advantages_multi_node_label>`.
**Longer context (ring attention, etc)**: More memory and more machines means we can train on longer sequences and tag advantage of neat tricks like ring attention, where tokens are split across
GPUs. You can read more about our plans for torchtune in `this Github RFC <https://github.com/pytorch/torchtune/issues/1244>`_.
**Want other optimizations?** Feel free to let us know by `opening up a Github Issue <https://github.com/pytorch/torchtune/issues/new?q=sort%3Aupdated-desc+is%3Aissue+is%3Aopen&template=Blank+issue>`_ on our repo or `dropping us a line in Discord <https://discord.gg/Zsf8xgT7>`_!
---
### Source Tutorials Qat Finetune (`docs/source/tutorials/qat_finetune.rst`)
.. _qat_finetune_label:
===========================
Fine-Tuning Llama3 with QAT
===========================
Quantization-Aware Training (QAT) is a common technique for users to quantize their
models without incurring significant degradations in accuracy or perplexity. In this
tutorial, we’ll walk through how to apply QAT during fine-tuning, quantize the
resulting model, and evaluate your quantized model using torchtune.
.. grid:: 2
.. grid-item-card:: :octicon:`mortar-board;1em;` What you will learn
* What QAT is and how it helps reduce quantization degradation
* How to run QAT during fine-tuning in torchtune
* End-to-end example of connecting QAT, quantization, and evaluation recipes
.. grid-item-card:: :octicon:`list-unordered;1em;` Prerequisites
* Be familiar with :ref:`torchtune<overview_label>`
* Make sure to :ref:`install torchtune`
* Make sure you have downloaded the :ref:`Llama3-8B model weights<download_llama_label>`
.. _what_is_qat_label:
What is QAT?
------------
`Quantization-Aware Training <https://pytorch.org/blog/introduction-to-quantization-on-pytorch/#quantization-aware-training>`_ (QAT) refers to simulating quantization numerics during
training or fine-tuning, with the end goal of ultimately producing a higher quality
quantized model compared to simple post-training quantization (PTQ). During QAT,
the weights and/or activations are “fake quantized”, meaning they are transformed
as if they were being quantized, but kept in the original data type (e.g. bfloat16)
without being actually cast to lower bit-widths. Thus, fake quantization allows the
model to adjust for quantization noise when updating the weights, hence the training
process is “aware” that the model will ultimately be quantized after training.
.. code-block:: python
# PTQ: x_q is quantized and cast to int8
# scale and zero point (zp) refer to parameters used to quantize x_float
# qmin and qmax refer to the range of quantized values
x_q = (x_float / scale + zp).round().clamp(qmin, qmax).cast(int8)
# QAT: x_fq is still in float
# Fake quantize simulates the numerics of quantize + dequantize
x_fq = (x_float / scale + zp).round().clamp(qmin, qmax)
x_fq = (x_fq - zp) * scale
QAT typically involves applying a transformation to your model before and after training.
For example, in the `torchao QAT implementation <https://github.com/pytorch/ao/blob/v0.2.0/torchao/quantization/prototype/qat.py>`_,
these are represented as the ``prepare()`` and ``convert()`` steps: (1) ``prepare()`` inserts fake quantize
operations into linear layers, and (2) ``convert()`` transforms the fake quantize operations
to actual quantize and dequantize operations after training, thereby producing a quantized
model (dequantize operations are typically fused with linear after lowering).
Between these two steps, training can proceed exactly as before.
.. image:: /_static/img/qat_diagram.png
.. _apply_qat_label:
Applying QAT to Llama3 models
-----------------------------
We can easily apply the above QAT transformations to Llama3 for fine-tuning,
leveraging the APIs in torchao as follows:
.. code-block:: python
import copy
import torch
from torchao.quantization import quantize_
from torchao.quantization.qat import (
FakeQuantizeConfig,
IntXQuantizationAwareTrainingConfig,
)
from torchtune.models.llama3 import llama3_8b
model = llama3_8b()
original_model = copy.deepcopy(model)
# Config for int8 dynamic asymmetric per token activations +
# int4 symmetric per group weights, only for linear layers
activation_config = FakeQuantizeConfig(torch.int8, "per_token", is_symmetric=False)
weight_config = FakeQuantizeConfig(torch.int4, group_size=32)
qat_config = IntXQuantizationAwareTrainingConfig(activation_config, weight_config)
# Prepare the model for quantization-aware fine-tuning.
#
# This step inserts "fake quantize" ops that simulate
# quantization numerics during fine-tuning without
# actually casting the activations/weights to lower-bit
# dtypes like in "real" quantization.
quantize_(model, qat_config)
prepared_model = model
The model is now ready for QAT fine-tuning! If we print the model we’ll see that
all linear layers have been swapped with :code:`FakeQuantizedLinear`, which simulates
the numerics of int8 dynamic asymmetric per token activations + int4 symmetric
per group weights:
.. code-block:: bash
>>> original_model.layers[0].attn
MultiHeadAttention(
(q_proj): Linear(in_features=4096, out_features=4096, bias=False)
(k_proj): Linear(in_features=4096, out_features=1024, bias=False)
(v_proj): Linear(in_features=4096, out_features=1024, bias=False)
(output_proj): Linear(in_features=4096, out_features=4096, bias=False)
(pos_embeddings): RotaryPositionalEmbeddings()
)
.. code-block:: bash
>>> prepared_model.layers[0].attn
MultiHeadAttention(
(q_proj): FakeQuantizedLinear(
in_features=4096, out_features=4096, bias=False
(activation_fake_quantizer): FakeQuantizer(FakeQuantizeConfig(dtype=torch.int8, granularity=PerToken(), mapping_type=<MappingType.ASYMMETRIC: 3>, scale_precision=torch.float32, zero_point_precision=torch.int32, zero_point_domain=<ZeroPointDomain.INT: 1>, is_dynamic=True, range_learning=False))
(weight_fake_quantizer): FakeQuantizer(FakeQuantizeConfig(dtype=torch.int4, granularity=PerGroup(group_size=32), mapping_type=<MappingType.SYMMETRIC: 1>, scale_precision=torch.float32, zero_point_precision=torch.int32, zero_point_domain=<ZeroPointDomain.INT: 1>, is_dynamic=True, range_learning=False))
)
(k_proj): FakeQuantizedLinear(
in_features=4096, out_features=1024, bias=False
(activation_fake_quantizer): FakeQuantizer(FakeQuantizeConfig(dtype=torch.int8, granularity=PerToken(), mapping_type=<MappingType.ASYMMETRIC: 3>, scale_precision=torch.float32, zero_point_precision=torch.int32, zero_point_domain=<ZeroPointDomain.INT: 1>, is_dynamic=True, range_learning=False))
(weight_fake_quantizer): FakeQuantizer(FakeQuantizeConfig(dtype=torch.int4, granularity=PerGroup(group_size=32), mapping_type=<MappingType.SYMMETRIC: 1>, scale_precision=torch.float32, zero_point_precision=torch.int32, zero_point_domain=<ZeroPointDomain.INT: 1>, is_dynamic=True, range_learning=False))
)
(v_proj): FakeQuantizedLinear(
in_features=4096, out_features=1024, bias=False
(activation_fake_quantizer): FakeQuantizer(FakeQuantizeConfig(dtype=torch.int8, granularity=PerToken(), mapping_type=<MappingType.ASYMMETRIC: 3>, scale_precision=torch.float32, zero_point_precision=torch.int32, zero_point_domain=<ZeroPointDomain.INT: 1>, is_dynamic=True, range_learning=False))
(weight_fake_quantizer): FakeQuantizer(FakeQuantizeConfig(dtype=torch.int4, granularity=PerGroup(group_size=32), mapping_type=<MappingType.SYMMETRIC: 1>, scale_precision=torch.float32, zero_point_precision=torch.int32, zero_point_domain=<ZeroPointDomain.INT: 1>, is_dynamic=True, range_learning=False))
)
(output_proj): FakeQuantizedLinear(
in_features=4096, out_features=4096, bias=False
(activation_fake_quantizer): FakeQuantizer(FakeQuantizeConfig(dtype=torch.int8, granularity=PerToken(), mapping_type=<MappingType.ASYMMETRIC: 3>, scale_precision=torch.float32, zero_point_precision=torch.int32, zero_point_domain=<ZeroPointDomain.INT: 1>, is_dynamic=True, range_learning=False))
(weight_fake_quantizer): FakeQuantizer(FakeQuantizeConfig(dtype=torch.int4, granularity=PerGroup(group_size=32), mapping_type=<MappingType.SYMMETRIC: 1>, scale_precision=torch.float32, zero_point_precision=torch.int32, zero_point_domain=<ZeroPointDomain.INT: 1>, is_dynamic=True, range_learning=False))
)
(pos_embeddings): RotaryPositionalEmbeddings()
)
After fine-tuning, we can convert the model to get an actual quantized model:
.. code-block:: python
from torchao.quantization.qat import (
FromIntXQuantizationAwareTrainingConfig,
)
from torchao.quantization import (
Int8DynamicActivationIntxWeightConfig,
)
from torchao.quantization.granularity import PerGroup
# Fine-tune as before
train_loop(prepared_model)
# Convert the fake quantized model into an actual quantized model
#
# First, we swap `FakeQuantizedLinear` back to `torch.nn.Linear`
# while keeping the QAT fine-tuned weights. Then, we perform standard
# post-training quantization (PTQ), which inserts quantized activation
# and weight tensor subclasses
quantize_(prepared_model, FromIntXQuantizationAwareTrainingConfig())
quantize_(prepared_model, Int8DynamicActivationIntxWeightConfig(weight_dtype=torch.int4, weight_granularity=PerGroup(32)))
converted_model = prepared_model
The model is now fully quantized to int8 and int4 and ready for inference
or generation. If we print the model now, we will see the linear layers
are now swapped back to :code:`torch.nn.Linear`, but with quantized tensor
activations and weights:
.. code-block:: bash
>>> converted_model.layers[0].attn
MultiHeadAttention(
(q_proj): Linear(in_features=4096, out_features=4096, weight=LinearActivationQuantizedTensor(activation=<function _int8_asymm_per_token_quant at 0x7f801ce08790>, weight=AffineQuantizedTensor(shape=torch.Size([4096, 4096]), block_size=(1, 32), device=cpu, _layout=PlainLayout(), tensor_impl_dtype=torch.int8, quant_min=-8, quant_max=7)))
(k_proj): Linear(in_features=4096, out_features=1024, weight=LinearActivationQuantizedTensor(activation=<function _int8_asymm_per_token_quant at 0x7f801ce08790>, weight=AffineQuantizedTensor(shape=torch.Size([1024, 4096]), block_size=(1, 32), device=cpu, _layout=PlainLayout(), tensor_impl_dtype=torch.int8, quant_min=-8, quant_max=7)))
(v_proj): Linear(in_features=4096, out_features=1024, weight=LinearActivationQuantizedTensor(activation=<function _int8_asymm_per_token_quant at 0x7f801ce08790>, weight=AffineQuantizedTensor(shape=torch.Size([1024, 4096]), block_size=(1, 32), device=cpu, _layout=PlainLayout(), tensor_impl_dtype=torch.int8, quant_min=-8, quant_max=7)))
(output_proj): Linear(in_features=4096, out_features=4096, weight=LinearActivationQuantizedTensor(activation=<function _int8_asymm_per_token_quant at 0x7f801ce08790>, weight=AffineQuantizedTensor(shape=torch.Size([4096, 4096]), block_size=(1, 32), device=cpu, _layout=PlainLayout(), tensor_impl_dtype=torch.int8, quant_min=-8, quant_max=7)))
(pos_embeddings): RotaryPositionalEmbeddings()
)
QAT finetuning recipe in torchtune
----------------------------------
Putting it all together, we can now fine-tune a model using torchtune’s :ref:`QAT recipe<qat_distributed_recipe_label>`.
Make sure that you have first downloaded the Llama3 weights and tokenizer by
following :ref:`these instructions<download_llama_label>`. In this tutorial,
we use the following settings to demonstrate QAT’s effectiveness in recovering
quantization degradation compared to directly quantizing a model fine-tuned
without QAT. You can copy the default QAT config and make the following
modifications accordingly:
.. code-block:: bash
tune cp llama3/8B_qat_full custom_8B_qat_full.yaml
.. code-block:: yaml
dataset:
_component_: torchtune.datasets.text_completion_dataset
source: allenai/c4
column: text
name: en
split: train
...
epochs: 1
max_steps_per_epoch: 2000
fake_quant_after_n_steps: 1000
By default, this uses the :code:`torchtune.training.quantization.Int8DynActInt4WeightQATQuantizer`,
which uses the same fake quantization configurations as the example above.
Empirically, we observed that disabling fake quantization for the first N steps
led to better results, presumably because doing so allows the weights to stabilize
before we start introducing quantization noise to the fine-tuning process.
For this reason, here we disable fake quantization for the first 1000 steps.
You can then use the following command to run fine-tuning with QAT using the above
config. This workload requires at least 6 GPUs, each with VRAM of at least 80GB.
By default, this uses the int8 dynamic per token activations + int4 grouped per
channel weights quantization configuration as shown above:
.. code-block:: bash
tune run --nnodes 1 --nproc_per_node 6 qat_distributed --config custom_8B_qat_full.yaml
.. note::
Make sure to point to the location of your Llama3 weights and tokenizer. This can be done
either by adding :code:`checkpointer.checkpoint_files=[my_model_checkpoint_path] tokenizer_checkpoint=my_tokenizer_checkpoint_path`
or by directly modifying the :code:`8B_qat_full.yaml` file. See our :ref:`config_tutorial_label`
for more details on how you can easily clone and modify torchtune configs.
.. note::
QAT introduces memory and computation overheads compared to regular fine-tuning,
since fake quantization fundamentally involves extra ops and requires cloning
the weights to avoid mutating them when computing the fake quantized values.
In general, we expect around 30% decrease in fine-tuning speed for models like
Llama3-8B. With activation checkpointing, the increase in memory footprint per
GPU is minimal (< 5GB per GPU).
Quantizing the QAT model
------------------------
Note that the QAT recipe above produces an unquantized bfloat16 model. The model
structure is exactly the same as the model produced with regular full fine-tuning
without QAT, just with different weights. To actually get a quantized model,
copy and make the following modifications to the quantization config:
.. code-block:: bash
tune cp quantization custom_quantization.yaml
.. code-block:: yaml
model:
_component_: torchtune.models.llama3.llama3_8b
checkpointer:
_component_: torchtune.training.FullModelMetaCheckpointer
checkpoint_dir: <your QAT checkpoint dir>
checkpoint_files: [ft-model-00001-of-00001.bin]
output_dir: <your QAT checkpoint dir>
model_type: LLAMA3
...
quantizer:
_component_: torchtune.training.quantization.Int8DynActInt4WeightQATQuantizer
groupsize: 256
The following command performs the convert step in the QAT flow, which actually
quantizes the float model to a model with quantized weights:
.. code-block:: bash
tune run quantize --config custom_quantization.yaml
.. note::
Make sure to use the same QAT quantizer you used to fine-tune your model,
otherwise the numerics will be off and the quantized model will perform poorly.
.. _qat_eval_label:
Evaluating the quantized model
------------------------------
Now that we have a quantized model, we can run some evaluations on it and compare the
results against regular fine-tuning without QAT (i.e. post-training quantization).
To achieve this, we use `EleutherAI’s evaluation harness <https://github.com/EleutherAI/lm-evaluation-harness>`_
integrated in torchtune. First, copy the evaluation config and make the following changes:
.. code-block:: bash
tune cp eleuther_evaluation custom_eleuther_evaluation.yaml
.. code-block:: yaml
model:
_component_: torchtune.models.llama3.llama3_8b
checkpointer:
_component_: torchtune.training.FullModelTorchTuneCheckpointer
checkpoint_dir: <your quantized model checkpoint dir>
checkpoint_files: [ft-model-00001-of-00001-8da4w.bin]
output_dir: <your quantized model checkpoint dir>
model_type: LLAMA3
...
tasks: ["hellaswag", "wikitext"]
quantizer:
_component_: torchtune.training.quantization.Int8DynActInt4WeightQuantizer
groupsize: 256
.. note::
Since we are passing in a quantized model, be sure to use the corresponding
post-training quantizer instead of the QAT quantizer. For example, if you
used the :code:`Int8DynActInt4WeightQATQuantizer` during fine-tuning, you
should specify :code:`Int8DynActInt4WeightQuantizer` in this step. See the
`quantization recipe <https://github.com/pytorch/torchtune/blob/main/recipes/quantize.py>`_
for a full list of supported quantizers.
Now run the evaluation recipe:
.. code-block:: bash
tune run eleuther_eval --config my_eleuther_evaluation.yaml
The results should look something like this:
.. code-block:: bash
# QAT quantized model evaluation results (int8 activations + int4 weights)
| Tasks |Version|Filter|n-shot| Metric |Value | |Stderr|
|---------|------:|------|-----:|---------------|-----:|---|------|
|wikitext | 2|none | 0|word_perplexity|9.9148|± |N/A |
| | |none | 0|byte_perplexity|1.5357|± |N/A |
| | |none | 0|bits_per_byte |0.6189|± |N/A |
|hellaswag| 1|none | 0|acc |0.5687|± |0.0049|
| | |none | 0|acc_norm |0.7536|± |0.0043|
Comparing these results to the model fine-tuned without QAT, we can see that
QAT was able to recover a significant portion of the quantization degradations
from the original unquantized model compared to PTQ. For example, normalized
accuracy in the hellaswag task dropped by 2.20% with PTQ but only 0.74% with
QAT when compared to the original unquantized model. Similarly, word perplexity
in the wikitext task increased by 2.048 with PTQ but only 1.190 with QAT (lower
is better).
.. code-block:: bash
# PTQ quantized model evaluation results (int8 activations + int4 weights)
| Tasks |Version|Filter|n-shot| Metric | Value | |Stderr|
|---------|------:|------|-----:|---------------|------:|---|------|
|wikitext | 2|none | 0|word_perplexity|10.7735|± |N/A |
| | |none | 0|byte_perplexity| 1.5598|± |N/A |
| | |none | 0|bits_per_byte | 0.6413|± |N/A |
|hellaswag| 1|none | 0|acc | 0.5481|± |0.0050|
| | |none | 0|acc_norm | 0.7390|± |0.0044|
.. code-block:: bash
# Float model evaluation results (bfloat16)
| Tasks |Version|Filter|n-shot| Metric |Value | |Stderr|
|---------|------:|------|-----:|---------------|-----:|---|------|
|wikitext | 2|none | 0|word_perplexity|8.7251|± |N/A |
| | |none | 0|byte_perplexity|1.4994|± |N/A |
| | |none | 0|bits_per_byte |0.5844|± |N/A |
|hellaswag| 1|none | 0|acc |0.5740|± |0.0049|
| | |none | 0|acc_norm |0.7610|± |0.0043|
Thus, the QAT flow produced a quantized model that outperforms the post-training
quantized model. Importantly, the quantized model structure is identical in both
flows, and so the model size, memory usage, and all other performance
characteristics are also the same.
Note that although the weights are quantized to int4, the quantized model size
for both the QAT and the PTQ flows are 8.187 GB, while the original float model
is 14.958 GB. This is because this quantizer uses int8 to represent the weights
as PyTorch does not have native int4 dtype support. A more efficient representation
is to pack the int4 weights, which will halve the quantized model size. This is
what the Int4WeightOnlyQuantizer does, and the corresponding QAT quantizer will
be added in the future.
Lowering QAT model to device (optional)
---------------------------------------
One important motivation for quantizing a model is to be able to run it in resource
constrained environments. You can further lower your QAT Llama3 model to edge devices
such as smartphones using `executorch <https://github.com/pytorch/executorch/>`_ by
following `these instructions <https://github.com/pytorch/executorch/tree/main/examples/models/llama2>`_.
For example, the following command lowers the model to the XNNPACK backend:
.. code-block:: bash
python -m examples.models.llama2.export_llama --checkpoint <your QAT checkpoint> -p -kv --use_sdpa_with_kv_cache -X -qmode 8da4w --group_size 256 -d fp32 --metadata '{"get_bos_id":128000, "get_eos_id":128001}' --embedding-quantize 4,32 --output_name="llama3_8da4w.pte"
This results in a much smaller quantized model of size 3.881 GB. When benchmarked on a OnePlus 12 smartphone, this model also achieved the same inference and generation speeds as the post-training quantized model. This is because the model structures are the same across the two flows:
.. list-table::
:widths: 25 25 25
:header-rows: 1
* -
- QAT
- PTQ
* - Quantized model size
- 3.881 GB
- 3.881 GB
* - Inference speed
- 9.709 tok/s
- 9.815 tok/s
* - Generation speed
- 11.316 tok/s
- 11.364 tok/s
---
### Source Tutorials Qlora Finetune (`docs/source/tutorials/qlora_finetune.rst`)
.. _qlora_finetune_label:
=============================
Fine-Tuning Llama2 with QLoRA
=============================
In this tutorial, we'll learn about `QLoRA <https://arxiv.org/abs/2305.14314>`_, an enhancement on top of
`LoRA <https://arxiv.org/abs/2106.09685>`_ that maintains frozen model parameters in 4-bit quantized precision, thereby reducing memory usage. We'll
walk through how QLoRA can be utilized within torchtune to finetune a Llama2-7b model in <10 GB of memory.
It is highly recommended to first develop an understanding of :ref:`LoRA finetuning in torchtune<lora_finetune_label>`.
.. grid:: 2
.. grid-item-card:: :octicon:`mortar-board;1em;` What you will learn
* How QLoRA saves memory over LoRA finetuning
* An overview of QLoRA in torchtune
* How to run a QLoRA finetune in torchtune
.. grid-item-card:: :octicon:`list-unordered;1em;` Prerequisites
* Be familiar with :ref:`torchtune<overview_label>`
* Make sure to :ref:`install torchtune`
* Make sure you have downloaded the :ref:`Llama2-7B model weights<download_llama_label>`
* Be familiar with :ref:`LoRA in torchtune<lora_finetune_label>`
What is QLoRA?
---------------
QLoRA builds on top of LoRA to enable further
memory savings. In LoRA, model parameters can be thought of as existing in two partitions: adapters, which are
low-rank matrices added to different layers of a neural network, and base model parameters, which are parameters that are part of
the original model. In vanilla LoRA-style training, both these parameters are held in the same precision (typically fp32 or `bf16 <https://en.wikipedia.org/wiki/Bfloat16_floating-point_format#bfloat16_floating-point_format>`_.), and
therefore activations and intermediate gradients computed are in fp32/bf16.
QLoRA further quantizes the base model parameters into a bespoke 4-bit NormalFloat (`NF4 <https://www.youtube.com/watch?v=TPcXVJ1VSRI&t=563s>`_) data type, resulting in 4-8x less parameter memory usage while
largely retaining model accuracy. As a result, the vast majority of parameters only take up 4 bits (as opposed to 16 or 32 bits by bf16/fp32 dtypes). This
quantization is done through the method highlighted in the original `QLoRA paper <https://arxiv.org/abs/2305.14314>`_. Adapter
parameters are still held in the original precision, and activations, gradients, and optimizer states still exist in the higher precision to preserve
accuracy.
The QLoRA authors introduce two key abstractions to decrease memory usage and avoid accuracy degradation: the bespoke 4-bit NormatFloat
type, and a double quantization method that quantizes the quantization parameters themselves to save even more memory. torchtune uses
the `NF4Tensor <https://github.com/pytorch/ao/blob/b9beaf351e27133d189b57d6fa725b1a7824a457/torchao/dtypes/nf4tensor.py#L153>`_ abstraction from the `torchao library <https://github.com/pytorch/ao>`_ to build QLoRA components as specified in the paper.
torchao is a PyTorch-native library that allows you to quantize and prune your models.
.. _qlora_core_highlevel:
Using QLoRA to save memory
----------------------------------------
In this section, we'll overview how to apply QLoRA to a :class:`~torchtune.modules.peft.LoRALinear` layer in torchtune. For a deep dive into details on QLoRA in torchtune and underlying abstractions,
please see the :ref:`QLoRA in torchtune deepdive <qlora_deepdive_label>` section of this tutorial.
A core idea of QLoRA is the distinction between compute and storage datatypes (dtypes). Specifically, QLoRA stores base model parameters in 4-bit precision (i.e. the storage dtype), and runs
computation in an original higher precision (the compute dtype), generally either fp32 or bf16. As a first step, QLoRA needs to quantize these base model parameters to 4-bit precision
and store them.
To quantize a :class:`~torchtune.modules.peft.LoRALinear` layer in the QLoRA style, simply pass in the ``quantize_base`` flag as ``True`` into :class:`~torchtune.modules.peft.LoRALinear`. This flag
will result in base model weights being quantized and backed by the ``NF4Tensor`` dtype. Forward passes will also be automatically handled to work with the ``NF4Tensor`` dtype,
specifically, the ``NF4`` base weight will be de-quantized to the compute precision, activation will be computed, and only the 4-bit parameter will be stored for gradient computation
in the backward pass, avoiding extra memory usage that would be incurred by storing the higher precision compute dtype.
Here's an example of creating a quantized ``LoRALinear`` layer in comparison to an unquantized ``LoRALinear`` layer. As we can see, the quantized layer consumes
~8x less memory than the unquantized counterpart.
.. code-block:: python
import torch
from torchtune.modules.peft import LoRALinear
torch.set_default_device("cuda")
qlora_linear = LoRALinear(512, 512, rank=8, alpha=0.1, quantize_base=True)
print(torch.cuda.memory_allocated()) # 177,152 bytes
del qlora_linear
torch.cuda.empty_cache()
lora_linear = LoRALinear(512, 512, rank=8, alpha=0.1, quantize_base=False)
print(torch.cuda.memory_allocated()) # 1,081,344 bytes
Using QLoRA in torchtune
----------------------------
We'll now cover how you can initialize a QLoRA-enabled Llama2-7b model as well as some details around
checkpointing with QLoRA.
With torchtune, you can use a simple builder similar to the LoRA builder (:func:`lora_llama_2_7b <torchtune.models.llama2.lora_llama2_7b>`) to apply QLoRA to Llama2 models. Here's a simple example of
initializing a Llama2-7b model with QLoRA enabled:
.. code-block:: python
from torchtune.models.llama2 import qlora_llama2_7b
qlora_model = qlora_llama2_7b(lora_attn_modules=["q_proj", "v_proj"])
Under the hood, this will apply LoRA to the ``q_proj`` and ``v_proj`` matrices in all attention layers, and further quantize the base parameters
in these matrices to the ``NF4`` dtype. Note that quantization of base model parameters is only applied to layers that are configured to have
LoRA adapters added. For example, in this case, ``k_proj`` and ``output_proj`` in the attention layers don't have LoRA applied to them, so their
base model parameters are not quantized. We can see this by printing the base model parameter dtypes for a particular attention layer:
.. code-block:: python
attn = qlora_model.layers[0].attn
print(type(attn.q_proj.weight)) # <class 'torchao.quantization.quantize_.workflows.nf4.nf4_tensor.NF4Tensor'>
print(type(attn.k_proj.weight)) # <class 'torch.nn.parameter.Parameter'>
Next, there are a couple of details essential to checkpointing (i.e. ``state_dict``) of QLoRA-enabled models.
To integrate well with torchtune's :ref:`checkpointing <checkpointing_label>`, we need to convert ``NF4Tensors`` back to their
original precision (generally fp32/bf16). This allows QLoRA-trained checkpoints to interoperate well with the rest of the ecosystem, within
torchtune and beyond (e.g. post-training quantization, evaluation, inference). This conversion process also allows LoRA adapter weights to be merged back into the base model as done
in a typical LoRA training flow.
To achieve this, when using torchtune's :func:`lora_llama_2_7b <torchtune.models.llama2.lora_llama2_7b>` builder, we automatically register a hook,
:func:`reparametrize_as_dtype_state_dict_post_hook <torchtune.modules.common_utils.reparametrize_as_dtype_state_dict_post_hook>`,
that runs after calling ``.state_dict()`` on the top level model. This hook converts ``NF4Tensors`` back to their original precision, while also offloading these
converted tensors to the CPU. This offloading is to avoid peaking memory; if we did not, we would have to maintain an entire bf16/fp32 copy of the ``state_dict``
on GPU.
.. _qlora_compile_label:
Putting it all together: QLoRA finetune
-----------------------------------------
Putting it all together, we can now finetune a model using torchtune's :ref:`LoRA single-device finetuning <lora_finetune_recipe_label>` recipe,
with a `QLoRA configuration <https://github.com/pytorch/torchtune/blob/main/recipes/configs/llama2/7B_qlora_single_device.yaml>`_.
Make sure that you have first downloaded the Llama2 weights and tokenizer by following :ref:`these instructions<download_llama_label>`.
You can then run the following command to perform a QLoRA finetune of Llama2-7B on a single GPU.
.. code-block:: bash
tune run lora_finetune_single_device --config llama2/7B_qlora_single_device
.. note::
Make sure to correctly point to the location of your Llama2 weights and tokenizer. This can be done
either by adding :code:`checkpointer.checkpoint_files=[my_model_checkpoint_path] tokenizer_checkpoint=my_tokenizer_checkpoint_path`
or by directly modifying the :code:`7B_qlora_single_device.yaml` file. See our ":ref:`config_tutorial_label`" recipe
for more details on how you can easily clone and modify torchtune configs.
By default, this run should log peak memory stats at model initialization time and every 100
iterations during training. Let's understand the memory savings enabled by QLoRA on top of LoRA training. LoRA training
can be run as follows:
.. code-block:: bash
tune run lora_finetune_single_device --config llama2/7B_lora_single_device
You should see the memory usage printed out during model initialization and training. An example log for LoRA model initialization is as follows:
.. code-block:: python
Memory Stats after model init::
GPU peak memory allocation: 13.96 GB
GPU peak memory reserved: 13.98 GB
GPU peak memory active: 13.96 GB
The following table compares the QLoRA's memory reserved during model initialization and training against vanilla LoRA's.
We can see that QLoRA reduces peak memory by about 35% during model initialization, and about 40% during model training:
================== ================================== ================================
Finetuning method Peak memory reserved, model init Peak memory reserved, training
================== ================================== ================================
LoRA 13.98 GB 15.57 GB
QLoRA 9.13 GB 9.29 GB
================== ================================== ================================
From the logs, one can see that the out-of-the-box training performance is quite slow, slower than 1 iteration per
second:
.. code-block:: python
1|149|Loss: 0.9157477021217346: 1%| | 149/25880 [02:08<6:14:19, 1.15it/s
To speed things up, we can leverage `torch.compile <https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html>`_ to compile our model and run the compiled result. To work with
QLoRA training, a nightly build of PyTorch must be used. To update PyTorch to the latest nightly,
please see `the installation instructions <https://pytorch.org/get-started/locally/>`_. Once updated,
you can specify the compile flag as ``True`` via a config override:
.. code-block:: bash
tune run lora_finetune_single_device --config llama2/7B_qlora_single_device compile=True
From the logs, we can see about a 200% speed up (after a few hundred iterations once the training has stabilized):
.. code-block:: python
1|228|Loss: 0.8158286809921265: 1%| | 228/25880 [11:59<1:48:16, 3.95it/s
A comparison of the smoothed loss curves between QLoRA and LoRA can be seen below.
.. image:: /_static/img/qlora_exp.png
.. note::
The above figure was generated with W&B. You can use torchtune's :class:`~torchtune.training.metric_logging.WandBLogger`
to generate similar loss curves, but you will need to install W&B and setup an account separately. For more details on
using W&B in torchtune, see our ":ref:`wandb_logging`" recipe.
As an exercise, you can also try running some evaluation tasks or manually inspecting generations
output by your saved checkpoints (which can be found in :code:`output_dir`).
In the final section, we'll go over a deep dive on how a QLoRA component can be built from a LoRA component.
.. _qlora_deepdive_label:
Deep-dive: Building QLoRA from LoRA
-----------------------------------------
This deep-dive section resumes from the :ref:`Using QLoRA to save memory<qlora_core_highlevel>` portion of this tutorial and dives into how quantization is done with ``NF4Tensor`` and handled appropriately in the forward pass.
First, we'll begin with
a vanilla minimal LoRA layer, taken from :ref:`the LoRA tutorial <lora_finetune_label>` and augmented to support quantization:
.. code-block:: python
:emphasize-lines: 3, 13, 19, 20, 39, 40, 41
import torch
from torch import nn
import torch.nn.functional as F
from torchao.quantization import to_nf4
from torchao.quantization.quantize_.workflows.nf4.nf4_tensor import linear_nf4
class LoRALinear(nn.Module):
def __init__(
self,
in_dim: int,
out_dim: int,
rank: int,
alpha: float,
dropout: float,
quantize_base: bool
):
# These are the weights from the original pretrained model
self.linear = nn.Linear(in_dim, out_dim, bias=False)
self.linear_weight = self.linear.weight
# Use torchao's to_nf4 API to quantize the base weight if needed.
if quantize_base:
self.linear_weight = to_nf4(self.linear_weight)
# These are the new LoRA params. In general rank << in_dim, out_dim
self.lora_a = nn.Linear(in_dim, rank, bias=False)
self.lora_b = nn.Linear(rank, out_dim, bias=False)
# Rank and alpha are commonly-tuned hyperparameters
self.rank = rank
self.alpha = alpha
# Most implementations also include some dropout
self.dropout = nn.Dropout(p=dropout)
# The original params are frozen, and only LoRA params are trainable.
self.linear.weight.requires_grad = False
self.lora_a.weight.requires_grad = True
self.lora_b.weight.requires_grad = True
def forward(self, x: torch.Tensor) -> torch.Tensor:
# frozen_out would be the output of the original model
if quantize_base:
# Call into torchao's linear_nf4 to run linear forward pass w/quantized weight.
frozen_out = linear_nf4(x, self.weight)
else:
frozen_out = F.linear(x, self.weight)
# lora_a projects inputs down to the much smaller self.rank,
# then lora_b projects back up to the output dimension
lora_out = self.lora_b(self.lora_a(self.dropout(x)))
# Finally, scale by the alpha parameter (normalized by rank)
# and add to the original model's outputs
return frozen_out + (self.alpha / self.rank) * lora_out
As mentioned above, torchtune takes a dependency on torchao for some of the core components required for QLoRA. This includes the
``NF4Tensor``, as well as helpful utilities including ``to_nf4`` and ``linear_nf4``.
The key changes on top of the LoRA layer are the usage of the ``to_nf4`` and ``linear_nf4`` APIs.
``to_nf4`` accepts an unquantized (bf16 or fp32) tensor and produces an ``NF4`` representation of the weight. See the `implementation <https://github.com/pytorch/ao/blob/c40358072f99b50cd7e58ec11e0e8d90440e3e25/torchao/dtypes/nf4tensor.py#L587>`_ of ``to_nf4`` for more details.
``linear_nf4`` handles the forward pass and autograd when running with quantized base model weights. It computes the forward pass as a regular
``F.linear`` with the incoming activation and unquantized weight. The quantized weight is saved for backward, as opposed to the unquantized version of the weight, to avoid extra
memory usage due to storing higher precision variables to compute gradients in the backward pass. See `linear_nf4 <https://github.com/pytorch/ao/blob/main/torchao/dtypes/nf4tensor.py#L577>`_ for more details.
---
### Source Templates Autosummary Class (`docs/source/_templates/autosummary/class.rst`)
.. role:: hidden
:class: hidden-section
.. currentmodule:: {{ module }}
{{ name | underline}}
.. autoclass:: {{ name }}
:members:
---
### Source Templates Autosummary Function (`docs/source/_templates/autosummary/function.rst`)
.. role:: hidden
:class: hidden-section
.. currentmodule:: {{ module }}
{{ name | underline}}
.. autofunction:: {{ name }}
---
--- METRICS ---
- Files Extracted: 56
- Estimated Token Budget: ~96156 tokens
- Canonical Reference: https://codewiki.google/github.com/meta-pytorch/torchtune