Table of Contents
  • Docs >
  • Configuring Datasets for Fine-Tuning
Shortcuts

Configuring Datasets for Fine-Tuning

This tutorial will guide you through how to set up a dataset to fine-tune on.

What you will learn

  • How to quickly get started with built-in datasets

  • How to use any dataset from Hugging Face Hub

  • How to use instruct, chat, or text completion datasets

  • How to configure datasets from code, config, or command-line

  • How to fully customize your own dataset

Prerequisites

Datasets are a core component of fine-tuning workflows that serve as a “steering wheel” to guide LLM generation for a particular use case. Many publicly shared open-source datasets have become popular for fine-tuning LLMs and serve as a great starting point to train your model. torchtune gives you the tools to download external community datasets, load in custom local datasets, or create your own datasets.

Built-in datasets

To use one of the built-in datasets in the library, simply import and call the dataset builder function. You can see a list of all supported datasets here.

from torchtune.datasets import alpaca_dataset

# Load in tokenizer
tokenizer = ...
dataset = alpaca_dataset(tokenizer)

# YAML config
dataset:
  _component_: torchtune.datasets.alpaca_dataset

# Command line
tune run full_finetune_single_device --config llama3/8B_full_single_device \
dataset=torchtune.datasets.alpaca_dataset

Hugging Face datasets

We provide first class support for datasets on the Hugging Face hub. Under the hood, all of our built-in datasets and dataset builders are using Hugging Face’s load_dataset() to load in your data, whether local or on the hub.

You can pass in a Hugging Face dataset path to the source parameter in any of our builders to specify which dataset on the hub to download. Additionally, all builders accept any keyword-arguments that load_dataset() supports. You can see a full list on Hugging Face’s documentation.

from torchtune.datasets import text_completion_dataset

# Load in tokenizer
tokenizer = ...
dataset = text_completion_dataset(
    tokenizer,
    source="allenai/c4",
    # Keyword-arguments that are passed into load_dataset
    split="train",
    data_dir="realnewslike",
)

# YAML config
dataset:
  _component_: torchtune.datasets.text_completion_dataset
  source: allenai/c4
  split: train
  data_dir: realnewslike

# Command line
tune run full_finetune_single_device --config llama3/8B_full_single_device \
dataset=torchtune.datasets.text_completion_dataset dataset.source=allenai/c4 \
dataset.split=train dataset.data_dir=realnewslike

Setting max sequence length

The default collator padded_collate() used in all our training recipes will pad samples to the max sequence length within the batch, not globally. If you wish to set an upper limit on the max sequence length globally, you can specify it in the dataset builder with max_seq_len. Any sample in the dataset that is longer than max_seq_len will be truncated in truncate(). The tokenizer’s EOS ids are ensured to be the last token, except in TextCompletionDataset.

Generally, you want the max sequence length returned in each data sample to match the context window size of your model. You can also decrease this value to reduce memory usage depending on your hardware constraints.

from torchtune.datasets import alpaca_dataset

# Load in tokenizer
tokenizer = ...
dataset = alpaca_dataset(
    tokenizer=tokenizer,
    max_seq_len=4096,
)

# YAML config
dataset:
  _component_: torchtune.datasets.alpaca_dataset
  max_seq_len: 4096

# Command line
tune run full_finetune_single_device --config llama3/8B_full_single_device \
dataset.max_seq_len=4096

Sample packing

You can use sample packing with any of the single dataset builders by passing in packed=True. 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.

from torchtune.datasets import alpaca_dataset, PackedDataset

# Load in tokenizer
tokenizer = ...
dataset = alpaca_dataset(
    tokenizer=tokenizer,
    packed=True,
)
print(isinstance(dataset, PackedDataset))  # True

# YAML config
dataset:
  _component_: torchtune.datasets.alpaca_dataset
  packed: True

# Command line
tune run full_finetune_single_device --config llama3/8B_full_single_device \
dataset.packed=True

Custom unstructured text corpus

For continued pre-training, typically a similar data setup to pre-training is used for a simple text completion task. This means no instruct templates, chat formats, and minimal special tokens (only BOS and EOS). To specify an unstructured text corpus, you can use the text_completion_dataset() builder with a Hugging Face dataset or a custom local corpus. Here is how to specify it for local files:

from torchtune.datasets import text_completion_dataset

# Load in tokenizer
tokenizer = ...
dataset = text_completion_dataset(
    tokenizer,
    source="text",
    data_files="path/to/my_data.txt",
    split="train",
)

# YAML config
dataset:
  _component_: torchtune.datasets.text_completion_dataset
  source: text
  data_files: path/to/my_data.txt
  split: train

# Command line
tune run --nproc_per_node 4 full_finetune_distributed --config llama3/8B_full \
dataset=torchtune.datasets.text_completion_dataset dataset.source=text \
dataset.data_files=path/to/my_data.txt dataset.split=train

Custom instruct dataset and instruct templates

If you have a custom instruct dataset that’s not already provided in the library, you can use the instruct_dataset() builder and specify the source path. Instruct datasets typically have multiple columns with text that are formatted into a prompt template.

To fine-tune an LLM on a particular task, a common approach is to create a fixed instruct template that guides the model to generate output with a specific goal. Instruct templates are simply flavor text that structures your inputs for the model. It is model agnostic and is tokenized normally just like any other text, but it can help condition the model to respond better to an expected format. For example, the AlpacaInstructTemplate structures the data in the following way:

"Below is an instruction that describes a task, paired with an input that provides further context. "
"Write a response that appropriately completes the request.\n\n"
"### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n"

Here is an example of a sample that is formatted with AlpacaInstructTemplate:

from torchtune.data import AlpacaInstructTemplate

sample = {
    "instruction": "Classify the following into animals, plants, and minerals",
    "input": "Oak tree, copper ore, elephant",
}
prompt = AlpacaInstructTemplate.format(sample)
print(prompt)
# Below is an instruction that describes a task, paired with an input that provides further context.
# Write a response that appropriately completes the request.
#
# ### Instruction:
# Classify the following into animals, plants, and minerals
#
# ### Input:
# Oak tree, copper ore, elephant
#
# ### Response:
#

We provide other instruct templates <data> for common tasks such summarization and grammar correction. If you need to create your own instruct template for a custom task, you can inherit from InstructTemplate and create your own class.

from torchtune.datasets import instruct_dataset
from torchtune.data import InstructTemplate

class CustomTemplate(InstructTemplate):
    # Define the template as string with {} as placeholders for data columns
    template = ...

    # Implement this method
    @classmethod
    def format(
        cls, sample: Mapping[str, Any], column_map: Optional[Dict[str, str]] = None
    ) -> str:
        ...

# Load in tokenizer
tokenizer = ...
dataset = instruct_dataset(
    tokenizer=tokenizer,
    source="my/dataset/path",
    template="import.path.to.CustomTemplate",
)

# YAML config
dataset:
  _component_: torchtune.datasets.instruct_dataset
  source: my/dataset/path
  template: import.path.to.CustomTemplate

# Command line
tune run full_finetune_single_device --config llama3/8B_full_single_device \
dataset=torchtune.datasets.instruct_dataset dataset.source=my/dataset/path \
dataset.template=import.path.to.CustomTemplate

Custom chat dataset and chat formats

If you have a custom chat/conversational dataset that’s not already provided in the library, you can use the chat_dataset() builder and specify the source path. Chat datasets typically have a single column with multiple back and forth messages between the user and assistant.

Chat formats are similar to instruct templates, except that they format system, user, and assistant messages into a list of messages (see ChatFormat) for a conversational dataset. These can be configured quite similarly to instruct datasets.

Here is how messages would be formatted using the Llama2ChatFormat:

from torchtune.data import Llama2ChatFormat, Message

messages = [
    Message(
        role="system",
        content="You are a helpful, respectful, and honest assistant.",
    ),
    Message(
        role="user",
        content="I am going to Paris, what should I see?",
    ),
    Message(
        role="assistant",
        content="Paris, the capital of France, is known for its stunning architecture..."
    ),
]
formatted_messages = Llama2ChatFormat.format(messages)
print(formatted_messages)
# [
#     Message(
#         role="user",
#         content="[INST] <<SYS>>\nYou are a helpful, respectful and honest assistant.\n<</SYS>>\n\n"
#         "I am going to Paris, what should I see? [/INST] ",
#     ),
#     Message(
#         role="assistant",
#         content="Paris, the capital of France, is known for its stunning architecture..."
#     ),
# ]

Note that the system message is now incorporated in the user message. If you create custom ChatFormats you can also add more advanced behavior.

from torchtune.datasets import chat_dataset
from torchtune.data import ChatFormat

class CustomChatFormat(ChatFormat):
    # Define templates for system, user, assistant messages
    # as strings with {} as placeholders for message content
    system = ...
    user = ...
    assistant = ...

    # Implement this method
    @classmethod
    def format(
        cls,
        sample: List[Message],
    ) -> List[Message]:
        ...

# Load in tokenizer
tokenizer = ...
dataset = chat_dataset(
    tokenizer=tokenizer,
    source="my/dataset/path",
    split="train",
    conversation_style="openai",
    chat_format="import.path.to.CustomChatFormat",
)

# YAML config
dataset:
  _component_: torchtune.datasets.chat_dataset
  source: my/dataset/path
  conversation_style: openai
  chat_format: import.path.to.CustomChatFormat

# Command line
tune run full_finetune_single_device --config llama3/8B_full_single_device \
dataset=torchtune.datasets.chat_dataset dataset.source=my/dataset/path \
dataset.conversation_style=openai dataset.chat_format=import.path.to.CustomChatFormat

Multiple in-memory datasets

It is also possible to train on multiple datasets and configure them individually using our ConcatDataset interface. You can even mix instruct and chat datasets or other custom datasets.

# YAML config
dataset:
  - _component_: torchtune.datasets.instruct_dataset
    source: vicgalle/alpaca-gpt4
    template: torchtune.data.AlpacaInstructTemplate
    split: train
    train_on_input: True
  - _component_: torchtune.datasets.instruct_dataset
    source: samsum
    template: torchtune.data.SummarizeTemplate
    column_map:
      output: summary
    split: train
    train_on_input: False
  - _component_: torchtune.datasets.chat_dataset
    ...

Local and remote datasets

To use a dataset saved on your local hard drive, simply specify the file type for source and pass in the data_files argument using any of the dataset builder functions. We support all file types supported by Hugging Face’s load_dataset, including csv, json, txt, and more.

from torchtune.datasets import instruct_dataset

# Load in tokenizer
tokenizer = ...
# Local files
dataset = instruct_dataset(
    tokenizer=tokenizer,
    source="csv",
    split="train",
    template="import.path.to.CustomTemplate"
    data_files="path/to/my/data.csv",
)
# Remote files
dataset = instruct_dataset(
    tokenizer=tokenizer,
    source="json",
    split="train",
    template="import.path.to.CustomTemplate"
    data_files="https://rajpurkar.github.io/SQuAD-explorer/dataset/train-v2.0.json",
    # You can also pass in any kwarg that load_dataset accepts
    field="data",
)

# YAML config - local files
dataset:
  _component_: torchtune.datasets.instruct_dataset
  source: csv
  template: import.path.to.CustomTemplate
  data_files: path/to/my/data.csv

# YAML config - remote files
dataset:
  _component_: torchtune.datasets.instruct_dataset
  source: json
  template: import.path.to.CustomTemplate
  data_files: https://rajpurkar.github.io/SQuAD-explorer/dataset/train-v2.0.json
  field: data

# Command line - local files
tune run full_finetune_single_device --config llama3/8B_full_single_device \
dataset=torchtune.datasets.chat_dataset dataset.source=csv \
dataset.template=import.path.to.CustomTemplate dataset.data_files=path/to/my/data.csv

Fully customized datasets

More advanced tasks and dataset formats that don’t fit into the templating and processing that InstructDataset, ChatDataset, and TextCompletionDataset provide may require you to create your own dataset class for more flexibility. Let’s walk through the PreferenceDataset, which has custom functionality for RLHF preference data, as an example to understand what you’ll need to do.

If you take a look at the code for the PreferenceDataset class, you’ll notice it’s quite similar to InstructDataset with a few adjustments for chosen and rejected samples in preference data.

chosen_message = [
    Message(role="user", content=prompt, masked=True),
    Message(role="assistant", content=transformed_sample[key_chosen]),
]
rejected_message = [
    Message(role="user", content=prompt, masked=True),
    Message(role="assistant", content=transformed_sample[key_rejected]),
]

chosen_input_ids, c_masks = self._tokenizer.tokenize_messages(
    chosen_message, self.max_seq_len
)
chosen_labels = list(
    np.where(c_masks, CROSS_ENTROPY_IGNORE_IDX, chosen_input_ids)
)

rejected_input_ids, r_masks = self._tokenizer.tokenize_messages(
    rejected_message, self.max_seq_len
)
rejected_labels = list(
    np.where(r_masks, CROSS_ENTROPY_IGNORE_IDX, rejected_input_ids)
)

For a specific dataset that’s easy to customize from the config, you can create a builder function. This is the builder function for the stack_exchanged_paired_dataset(), which creates a PreferenceDataset configured to use a paired dataset from Hugging Face. Notice that we’ve also had to add a custom instruct template as well.

def stack_exchanged_paired_dataset(
    tokenizer: ModelTokenizer,
    max_seq_len: int = 1024,
) -> PreferenceDataset:
    return PreferenceDataset(
        tokenizer=tokenizer,
        source="lvwerra/stack-exchange-paired",
        template=StackExchangedPairedTemplate(),
        column_map={
            "prompt": "question",
            "chosen": "response_j",
            "rejected": "response_k",
        },
        max_seq_len=max_seq_len,
        split="train",
        data_dir="data/rl",
    )

Now we can easily specify our custom dataset from the config, or from command-line.

# This is how you would configure the Alpaca dataset using the builder
dataset:
  _component_: torchtune.datasets.stack_exchanged_paired_dataset
  max_seq_len: 512

# Command line - local files
tune run full_finetune_single_device --config llama3/8B_full_single_device \
dataset=torchtune.datasets.stack_exchanged_paired_dataset dataset.max_seq_len=512

Docs

Access comprehensive developer documentation for PyTorch

View Docs

Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

Resources

Find development resources and get your questions answered

View Resources

To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.