Add masking strategies to message transforms #2284

supreethmanyam · 2025-01-20T23:07:56Z

Context

What is the purpose of this PR? Is it to

add a new feature
fix a bug
update tests and/or documentation
other (please add here)

Please link to any issues this PR addresses.
#2261

Changelog

What are the changes made in this PR?

Added mask_messages utility in torchtune/data/_messages.py to mask each message given a list of messages based on user provided masking_strategy
Fixes in-built message transforms to always mask system message
Added tests for calling message transforms with masking strategies

Test plan

Please make sure to do each of the following if applicable to your PR. If you're unsure about any one of these just ask and we will happily help. We also have a contributing page for some guidance on contributing.

run pre-commit hooks and linters (make sure you've first installed via pre-commit install)
add unit tests for any new functionality
update docstrings for any new or updated methods or classes
run unit tests via pytest tests
run recipe tests via pytest tests -m integration_test
manually run any new or modified recipes with sufficient proof of correctness
include relevant commands and any other artifacts in this summary (pastes of loss curves, eval results, etc.)

UX

If your function changed a public API, please add a dummy example of what the user experience will look like when calling it.
Here is a docstring example
and a tutorial example

I did not change any public API
I have added an example to docs or docstrings

pytorch-bot · 2025-01-20T23:07:59Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2284

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 806f437 with merge base d5d12fe ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

supreethmanyam · 2025-01-21T01:26:31Z

@RdoubleA Added mask_messages utility in torchtune/data/_messages.py along with relevant unit tests.
I will update the documentation once the changes look good. Please let me know.

RdoubleA

This looks great, really like the tests 👌 left a couple of comments

RdoubleA · 2025-01-23T01:19:58Z

torchtune/data/_messages.py

@@ -176,12 +185,24 @@ class InputOutputToMessages(Transform):

    def __init__(
        self,
-        train_on_input: bool = False,
+        masking_strategy: Optional[str] = "train_on_all",


if the default value for train_on_input was False, this would correspond to train_on_assistant. Let's keep that same default behavior.

Updated default strategy

RdoubleA · 2025-01-23T01:20:20Z

torchtune/data/_messages.py

@@ -305,11 +332,23 @@ class ChosenRejectedToMessages(Transform):

    def __init__(
        self,
-        train_on_input: bool = False,
+        masking_strategy: Optional[str] = "train_on_all",


same comment here

Updated default strategy

RdoubleA · 2025-01-23T01:20:45Z

torchtune/data/_messages.py


    Raises:
        ValueError: If ``column_map`` is provided and ``conversations`` not in ``column_map``.
    """

    def __init__(
        self,
-        train_on_input: bool = False,
+        masking_strategy: Optional[str] = "train_on_all",


same comment here

Updated default strategy

RdoubleA · 2025-01-23T01:21:51Z

torchtune/data/_messages.py

    """

    def __init__(
-        self, train_on_input: bool = True, column_map: Optional[Dict[str, str]] = None
+        self,
+        masking_strategy: Optional[str] = "train_on_all",


contrary to the other comment, this default value should remain train_on_all since train_on_input was defaulted to True. so let's keep this

RdoubleA · 2025-01-23T01:25:10Z

torchtune/data/_messages.py

+        ValueError: If the masking strategy is not one of the supported strategies:
+                    `train_on_all`, `train_on_assistant`, `train_on_last`.
+    """
+    if masking_strategy not in ["train_on_all", "train_on_assistant", "train_on_last"]:


One suggestion that I didn't mention in the issue is that we could make masking_strategy an Enum so that it can be type checked more easily. users can still specify as a string in the config and in the message transforms above, but here in mask_messages you could convert to the Enum, and the ValueError will automatically be captured.

Converted masking_strategy to Enum and raises ValueError automatically if doesn't match with mapping

supreethmanyam · 2025-01-28T03:00:47Z

@RdoubleA

In addition to converting masking_strategy to Enum and fixing default strategy in built-in transforms, I made following small changes:

I modified structure of column_map usage in InputOutputToMessages and AlpacaToMessages to make it consistent across transforms
I realized there were no unit tests for AlpacaToMessages, I added a few tests.

ebsmothers

Sorry @supreethmanyam, I think this PR fell through the cracks a bit. I have a couple comments but after that I think it looks good. cc @RdoubleA to confirm

ebsmothers · 2025-02-28T17:45:13Z

torchtune/data/_messages.py

@@ -177,31 +193,42 @@ class InputOutputToMessages(Transform):

    def __init__(
        self,
-        train_on_input: bool = False,
+        masking_strategy: Optional[str] = "train_on_assistant",


Why do we switch the order of these arguments? If someone is passing InputOutputToMessages(False) this will break them, right? (Since we unfortunately don't enforce keyword-only args on these APIs). Personally I would leave train_on_input in the first position and put masking_strategy last

agreed, let's keep the positions of the arguments the same until we deprecate

Ah! my bad!

I rearranged the order of arguments. I moved train_on_input to first as it was before and moved masking_strategy to the last. Updated docstrings accordingly.

ebsmothers · 2025-02-28T17:46:51Z

tests/torchtune/data/test_messages.py

+            ("train_on_all", [[False, False], [False, False]]),
+            ("train_on_assistant", [[True, False], [True, False]]),
+            ("train_on_last", [[True, False], [True, False]]),


Ideally we should explicitly test a case where the results of train_on_assistant and train_on_last differ

I have added test cases using multi turn samples where results train_on_assistant and train_on_last differ where applicable (ChosenRejectedToMessages, ShareGPTToMessages, OpenAIToMessages)

ebsmothers · 2025-02-28T17:47:51Z

tests/torchtune/data/test_messages.py

@@ -479,6 +609,200 @@ def test_call_image_messages(self, mock_load_image):
        mock_load_image.assert_called_once_with("https://example.com")


+class TestAlpacaToMessages:


What's the difference between this and the existing TestAlpacaToMessages here? They look fairly similar

I did not know we had TestAlpacaToMessages in test_alpaca_dataset.py and added one in test_messages.py.

Both are similar and the one in test_messages.py runs more tests. Please let me know if I can remove TestAlpacaToMessages from here

ebsmothers · 2025-02-28T17:50:29Z

torchtune/data/_messages.py

        return {"messages": messages}
+
+
+def validate_messages(


Similar q here: don't we have an existing validate_messages function?

it was just moved down here I believe

I moved validate_messages down to keep all the message transforms at one place.

RdoubleA · 2025-02-28T18:52:42Z

torchtune/data/_messages.py

    ):
-        self.train_on_input = train_on_input
-        self.column_map = column_map
+        if train_on_input is not None:


I thought we had a deprecated param utility, but it looks like it never got reviewed: #2321. should make a note there to also replace this code if this PR merges first @ebsmothers

…iffer

supreethmanyam · 2025-03-12T03:04:08Z

@RdoubleA @ebsmothers
I have made all the changes and updated the documentation. Requesting another review.

Current references to train_on_input in the docs comes from dataset modules, hence I haven't made any changes yet. Please let me know if I should create another issue for dataset module changes. I will be happy to make changes to dataset modules to reflect using masking_strategy instead of train_on_input and update docs accordingly.

ebsmothers · 2025-03-26T04:34:09Z

Thanks @supreethmanyam for your patience on this one. Since we have not been great about timely reviews here, I took the liberty of pushing some changes directly to the PR (hope that's alright). Mainly I changed the deprecation logic slightly: when train_on_input and masking_strategy are both passed, we should still use train_on_input but warn the user (otherwise it's possible for silent behavior changes without the user actually migrating). While doing this, I noticed that our SlimOrca test case seems incorrect -- it appears as though we were never masking the system message (see here). cc @RdoubleA to sanity check me here. After this (and green CI) I think this should be good to land. Thanks again for your patience!

RdoubleA · 2025-03-26T04:40:41Z

That's a tough one... it does look like system is being masked, but if train_on_input is true it will not be masked, same as the user message. So yeah, big oof

ebsmothers

Really appreciate you adding this @supreethmanyam!

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jan 20, 2025

supreethmanyam force-pushed the add-chat-loss-masking-strategies branch from db9a565 to 255bdcd Compare January 20, 2025 23:14

Add masking strategies to message transforms

ac6272f

supreethmanyam force-pushed the add-chat-loss-masking-strategies branch from 255bdcd to ac6272f Compare January 20, 2025 23:22

Add tests for masking strategies in message transforms

4e97a57

RdoubleA mentioned this pull request Jan 21, 2025

v0.6.0 tracker #2232

Open

RdoubleA reviewed Jan 23, 2025

View reviewed changes

supreethmanyam added 2 commits January 27, 2025 16:51

Merge branch 'pytorch:main' into add-chat-loss-masking-strategies

0381a13

Convert MaskingStrategy to Enum and fix default strategy in transforms

8ef107a

Merge branch 'pytorch:main' into add-chat-loss-masking-strategies

98ad8a4

ebsmothers reviewed Feb 28, 2025

View reviewed changes

RdoubleA reviewed Feb 28, 2025

View reviewed changes

RdoubleA mentioned this pull request Mar 6, 2025

Will support multi-turn conversations? #2463

Open

supreethmanyam added 4 commits March 11, 2025 17:25

Merge branch 'main' into add-chat-loss-masking-strategies

b3bbdf7

Rearrange the order of train_on_input and masking_strategy arguments

90c7c22

Add tests where the results of train_on_assistant and train_on_last d…

7d573ef

…iffer

Update documentation

6684dd9

supreethmanyam requested review from ebsmothers and RdoubleA March 12, 2025 03:05

cleanup, fix test case

806f437

ebsmothers approved these changes Mar 26, 2025

View reviewed changes

ebsmothers merged commit 32d195c into pytorch:main Mar 26, 2025
17 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add masking strategies to message transforms #2284

Add masking strategies to message transforms #2284

supreethmanyam commented Jan 20, 2025 •

edited

Loading

pytorch-bot bot commented Jan 20, 2025 •

edited

Loading

supreethmanyam commented Jan 21, 2025 •

edited

Loading

RdoubleA left a comment

RdoubleA Jan 23, 2025

supreethmanyam Jan 28, 2025

RdoubleA Jan 23, 2025

supreethmanyam Jan 28, 2025

RdoubleA Jan 23, 2025

supreethmanyam Jan 28, 2025

RdoubleA Jan 23, 2025

RdoubleA Jan 23, 2025

supreethmanyam Jan 28, 2025 •

edited

Loading

supreethmanyam commented Jan 28, 2025

ebsmothers left a comment

ebsmothers Feb 28, 2025

RdoubleA Feb 28, 2025

supreethmanyam Mar 12, 2025

ebsmothers Feb 28, 2025

supreethmanyam Mar 12, 2025

ebsmothers Feb 28, 2025

supreethmanyam Mar 12, 2025

ebsmothers Feb 28, 2025

RdoubleA Feb 28, 2025

supreethmanyam Mar 12, 2025

RdoubleA Feb 28, 2025

supreethmanyam commented Mar 12, 2025 •

edited

Loading

ebsmothers commented Mar 26, 2025

RdoubleA commented Mar 26, 2025

ebsmothers left a comment

		@@ -479,6 +609,200 @@ def test_call_image_messages(self, mock_load_image):
		mock_load_image.assert_called_once_with("https://example.com")


		class TestAlpacaToMessages:

Add masking strategies to message transforms #2284

Add masking strategies to message transforms #2284

Conversation

supreethmanyam commented Jan 20, 2025 • edited Loading

Context

Changelog

Test plan

UX

pytorch-bot bot commented Jan 20, 2025 • edited Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2284

✅ No Failures

supreethmanyam commented Jan 21, 2025 • edited Loading

RdoubleA left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

supreethmanyam Jan 28, 2025 • edited Loading

Choose a reason for hiding this comment

supreethmanyam commented Jan 28, 2025

ebsmothers left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

supreethmanyam commented Mar 12, 2025 • edited Loading

ebsmothers commented Mar 26, 2025

RdoubleA commented Mar 26, 2025

ebsmothers left a comment

Choose a reason for hiding this comment

supreethmanyam commented Jan 20, 2025 •

edited

Loading

pytorch-bot bot commented Jan 20, 2025 •

edited

Loading

supreethmanyam commented Jan 21, 2025 •

edited

Loading

supreethmanyam Jan 28, 2025 •

edited

Loading

supreethmanyam commented Mar 12, 2025 •

edited

Loading