### Source/ Templates/Autosummary/Class .. role:: hidden :class: hidden-section .. currentmodule:: {{ module }} {{ name | underline}} .. autoclass:: {{ name }} :members: --- ### Source/ Templates/Autosummary/Function .. role:: hidden :class: hidden-section .. currentmodule:: {{ module }} {{ name | underline}} .. autofunction:: {{ name }} --- ### Source/Basics/Chat Datasets .. _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 `_. 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 `_ 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": , }, ... ] .. _sharegpt: ``"sharegpt"`` ^^^^^^^^^^^^^^ The associated message transform is :class:`~torchtune.data.ShareGPTToMessages`. The expected format is: .. code-block:: python { "conversations": [ { "from": "system" | "human" | "gpt", "value": , }, ... ] } 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": , }, ... ] } 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 .. _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 .. _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 .. _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 .. _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 `_. 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 `_ 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 ``{"": ""}``. 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 .. _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 .. _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 .. _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 .. _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 `_. 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 `_ 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 `, 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 .. _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 `_, 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 .. _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 `_. 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 `_ 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 .. _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 .. _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 .. _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 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) # .. 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 .. _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 `_, 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 `_ 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 `_ whilst some other models like Phi3 have their own `conversion functions `_ which can be found within their model folders. | Handling different Checkpoint Formats ------------------------------------- torchtune supports three different :ref:`checkpointers`, each of which supports a different checkpoint format. :class:`HFCheckpointer ` ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 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: # 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 ` ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 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 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: # 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 ` ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 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 ` ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 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 ` 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 `. .. 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 `. We hope this deep-dive provided a deeper insight into the checkpointer and associated utilities in torchtune. Happy tuning! --- ### Source/Deep Dives/Comet Logging .. _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 `_. 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 `_. The config used to train the models can be found `here `_. --- ### Source/Deep Dives/Configs .. _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` * Make sure to :ref:`install torchtune` * Understand the :ref:`fundamentals of recipes` 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` ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 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`, 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` 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 ` 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 `_ optimizer, you may need to delete parameters like ``foreach`` which are specific to PyTorch optimizers. Note that this example requires that you have `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 .. _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` and :ref:`Checkpointing `). | **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 `_ in torchtune and the associated `config `_. .. 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 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 ` 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 ` command with any CLI overrides: .. code-block:: bash tune run --config k1=v1 k2=v2 ... --- ### Source/Deep Dives/Wandb Logging .. _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 `_. 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 `_. The config used to train the models can be found `here `_. 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 .. _dpo_recipe_label: ==================================== Direct Preference Optimization ==================================== This recipe supports several `Direct Preference Optimization `_ (DPO)-style fine-tuning techniques. These techniques aim to steer (or `align `_) 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 `_. 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 ` 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 ` 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 # 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 ` (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 ` (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`. .. 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 .. _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 ` 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`. 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` * :ref:`Finetuning Llama2 with QLoRA` * :ref:`Fine-tuning Llama3 with Chat Data` * :ref:`Meta Llama3 in torchtune` * :ref:`Fine-Tune Your First LLM` --- ### Source/Recipes/Qat Distributed .. _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 `_ to implement QAT. This works by :ref:`simulating quantization numerics during fine-tuning `. 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 `_ 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 `_ 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 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 ` 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 `_ 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` or `run inference `_ 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 ` to learn more. Many of our other memory optimization features can be used in this recipe, too: * Adjust :ref:`model precision `. * Use :ref:`activation checkpointing `. * Enable :ref:`gradient accumulation `. * Use :ref:`lower precision optimizers `. You can learn more about all of our memory optimization features in our :ref:`memory optimization overview`. 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 .. _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 `. * Single-device full fine-tuning * Distributed full fine-tuning * Distributed LoRA fine-tuning * :ref:`Direct Preference Optimization (DPO) ` * Proximal Policy Optimization (PPO) * :ref:`Distributed Quantization-Aware Training (QAT)`. 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 `_ (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 ` and :ref:`multi-device ` LoRA finetuning. .. note:: Our recipe documentation is currently in construction. Please feel free to follow the progress in our tracker issue `here `_. --- ### Source/Tutorials/Chat .. _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` * Know how to :ref:`download Llama3 Instruct weights ` 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 `_ for the Llama2 chat model, we can see that special tags are added: .. code-block:: text [INST] <> You are a helpful, respectful, and honest assistant. <> Hi! I am a human. [/INST] Hello there! Nice to meet you! I'm Meta AI, your friendly AI assistant Llama3 Instruct `overhauled `_ 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 `_ 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] <>\nYou are a helpful, respectful, and honest assistant.\n<>\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:`` and :code:`` 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("")) # 1 print(tokenizer._spm_model.spm_model.piece_to_id("")) # 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:`<>` 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])) # '<>' 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("", 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 `_ 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 `_. Now we're ready to start fine-tuning! We'll use the built-in LoRA single device recipe. Use the :ref:`tune cp ` 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 .. _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` * Make sure to :ref:`install torchtune` * Concepts such as :ref:`configs ` 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 `_ 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() `_ 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`. 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 `. 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 ` and use the standard settings from the `default config `_. 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 with `EleutherAI's evaluation harness `_. An example of this is available through the `eleuther_eval `_ recipe. In this tutorial, we're going to directly use this recipe by modifying its associated config `eleuther_evaluation.yaml `_. .. note:: For this section of the tutorial, you should first run :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 $ tune cp eleuther_evaluation ./custom_eval_config.yaml Copied file to custom_eval_config.yaml Notice that we are using the merged weights, and not the LoRA adapters. .. code-block:: yaml # 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: null For this tutorial we'll use the `truthfulqa_mc2 `_ task from the harness. This task measures a model's propensity to be truthful when answering questions and measures the model's zero-shot accuracy on a question followed by one or more true responses and one or more false responses. .. code-block:: text $ tune run eleuther_eval --config ./custom_eval_config.yaml [evaluator.py:324] Running loglikelihood requests ... Generate some output ~~~~~~~~~~~~~~~~~~~~ We've run some evaluations and the model seems to be doing well. But does it really generate meaningful text for the prompts you care about? Let's find out! For this, we'll use the `generate recipe `_ and the associated `config `_. Let's first copy over the config to our local working directory so we can make changes. .. 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 quantization ~~~~~~~~~~~~~~~~~~~~~~~~~~~ We rely on `torchao `_ for `post-training 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 `_. torchao also provides `this table `_ 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 `_. 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()`` ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ **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 `_. 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 `_ 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 `_. .. 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 `_. | Hopefully this tutorial gave you some insights into how you can use torchtune for your own workflows. Happy Tuning! --- ### Source/Tutorials/First Finetune Tutorial .. _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 `_ * 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` * 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 `_ - a collection of the latest and greatest model weights. For this tutorial, you're going to use the `Llama2 7B model from Meta `_. Llama2 is a "gated model", meaning that you need to be granted access in order to download the weights. Follow `these instructions `_ 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 `_. 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 `_ 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 `_. 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 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 `_, using memory efficient techniques like `LoRA `_, and more! Check out all our built-in recipes in our :ref:`recipes overview`. 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 `_ 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 ` 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 `_, 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 --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 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 `_ 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`". --- ### Source/Tutorials/Llama Kd Tutorial .. _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` * Make sure to :ref:`install torchtune` * Make sure you have downloaded the :ref:`Llama3 model weights` * Be familiar with :ref:`LoRA` What is Knowledge Distillation? ------------------------------- `Knowledge Distillation `_ 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 `_ 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 `_, `DistiLLM `_, and `Generalized KD `_, 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 `_. 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 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 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 `_, `hellaswag `_ and `commonsense_qa `_ tasks through the EleutherAI `LM evaluation harness `_. 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 `_ 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 ---