Support masking of partial dialogue in multi-turn chat datasets

[jiatong-yu]

According to the documentation, it appears that Torchtune currently supports training using either all “assistant” | “user” content or all of “assistant” content in a multi-turn conversation. However, a common use case is training on a specific subset of responses, such as only the most recent “assistant” responses in a conversation.

What is the recommended approach for achieving this with Torchtune?

[felipemello1]

hey @jiatong-yu , you should be able to write your own custom message_transform / dataset.

Here is our wiki: https://pytorch.org/torchtune/main/basics/message_transforms.html

Take a look at how its done in the chat dataset: chatdataset.

torchtune/torchtune/datasets/_chat.py

Line 21 in aa8f365

train_on_input: bool = False,

Then, in your config, you can pass:

tune run <recipe> <config> --config dataset._component_:path.to.my.custom.dataset

[calvinpelletier]

Here's some additional info: #2111 (comment)