diff --git a/torchtune/models/mistral/_tokenizer.py b/torchtune/models/mistral/_tokenizer.py
index b541c993d9..c3bbc8a4a7 100644
--- a/torchtune/models/mistral/_tokenizer.py
+++ b/torchtune/models/mistral/_tokenizer.py
@@ -122,7 +122,10 @@ def decode(
        s

diff --git a/torchtune/models/gemma/_tokenizer.py b/torchtune/models/gemma/_tokenizer.py
index d4655c5f48..e5eb89e230 100644
--- a/torchtune/models/gemma/_tokenizer.py
+++ b/torchtune/models/gemma/_tokenizer.py
@@ -99,6 +99,8 @@ def decode(
     def tokenize_messages(
         self,
         messages: List[Message],
+        *,
+        add_eos: bool = True,
     ) -> Tuple[List[int], List[bool]]:
         r"""Tokenize a list of messages one at a time then concatenate them,
         returning a list of tokens and a list of masks.
@@ -125,6 +127,7 @@ def tokenize_messages(
         Args:
             messages (List[Message]): A list of messages, each containing role, content,
                 and masked attributes.
+            add_eos (bool): Whether to append EOS after assistant message, default to True

         Returns:
             Tuple[List[int], List[bool]]: The tokenized messages
@@ -138,7 +141,7 @@ def tokenize_messages(
             tokenizer=self,
             messages=templated_messages,
             bos_id=self.bos_id,
-            eos_id=self.eos_id,
+            eos_id=self.eos_id if add_eos else None,
         )

     def __call__(
diff --git a/torchtune/models/mistral/_tokenizer.py b/torchtune/models/mistral/_tokenizer.py
index b541c993d9..c3bbc8a4a7 100644
--- a/torchtune/models/mistral/_tokenizer.py
+++ b/torchtune/models/mistral/_tokenizer.py
@@ -122,7 +122,10 @@ def decode(
         return self._spm_model.decode(token_ids)

     def tokenize_messages(
-        self, messages: List[Message]
+        self,
+        messages: List[Message],
+        *,
+        add_eos: bool = True,
     ) -> Tuple[List[int], List[bool]]:
         r"""Tokenize a list of messages one at a time then concatenate them,
         returning a list of tokens and a list of masks.
@@ -154,6 +157,7 @@ def tokenize_messages(
         Args:
             messages (List[Message]): A list of messages, each containing role, content,
                 and masked attributes.
+            add_eos (bool): Whether to append EOS after assistant message, default to True

         Returns:
             Tuple[List[int], List[bool]]: The tokenized messages
@@ -167,7 +171,7 @@ def tokenize_messages(
             tokenizer=self,
             messages=templated_messages,
             bos_id=self.bos_id,
-            eos_id=self.eos_id,
+            eos_id=self.eos_id if add_eos else None,
         )

     def __call__(
diff --git a/torchtune/models/phi3/_tokenizer.py b/torchtune/models/phi3/_tokenizer.py
index bd1466497a..b48b1d93a3 100644
--- a/torchtune/models/phi3/_tokenizer.py
+++ b/torchtune/models/phi3/_tokenizer.py
@@ -241,9 +241,9 @@ def tokenize_messages(
         # Finally, truncate if necessary
         if self.max_seq_len and len(tokenized_messages) >= self.max_seq_len:
             tokenized_messages = truncate(
-                tokenized_messages, self.max_seq_len, self.eos_id
+                tokenized_messages, self.max_seq_len, self.eos_id if add_eos else None
             )
-            mask = truncate(mask, self.max_seq_len, message.masked)
+            mask = truncate(mask, self.max_seq_len, message.masked if add_eos else None)

         return tokenized_messages, mask

diff --git a/torchtune/models/qwen2/_tokenizer.py b/torchtune/models/qwen2/_tokenizer.py
index a4f749c39c..952100002e 100644
--- a/torchtune/models/qwen2/_tokenizer.py
+++ b/torchtune/models/qwen2/_tokenizer.py
@@ -381,9 +381,9 @@ def tokenize_messages(
         # Finally, truncate if necessary
         if self.max_seq_len:
             tokenized_messages = truncate(
-                tokenized_messages, self.max_seq_len, self.eos_id
+                tokenized_messages, self.max_seq_len, self.eos_id if add_eos else None
             )
-            mask = truncate(mask, self.max_seq_len, True)
+            mask = truncate(mask, self.max_seq_len, True if add_eos else None)

         return tokenized_messages, mask

Add an optional parameter `add_eos` (defaulting to True) to control whether the EOS token is appended when tokenizing messages.

Prompt:
The following is the text of a github issue and related comments for this repository:

Ensure all tokenizers have add_eos parameter and ensure eos_token is not added if add_eos=False
 - Phi3 tokenizer
 - Gemma tokenizer
 - Mistral tokenizer
 - Qwen2 tokenizer

Resolve the issue described above. Do not modify tests or user code. The issue can be resolved by modifying the actual library/application code. Make the minimal changes to resolve the issue while adhering to the specification and coding style of the codebase.

diff --git a/tests/torchtune/models/gemma/test_gemma_tokenizer.py b/tests/torchtune/models/gemma/test_gemma_tokenizer.py
index 205f0fae1f..41a6ef79fd 100644
--- a/tests/torchtune/models/gemma/test_gemma_tokenizer.py
+++ b/tests/torchtune/models/gemma/test_gemma_tokenizer.py
@@ -242,3 +242,232 @@ def test_tokenize_messages(self, tokenizer):
         expected_mask = [True] * 75 + [False] * 125
         assert expected_tokens == tokens
         assert expected_mask == mask
+
+    def test_tokenize_messages_drop_eos(self, tokenizer):
+        messages = [
+            Message(
+                role="user",
+                content="Below is an instruction that describes a task. Write a response "
+                "that appropriately completes the request.\n\n### Instruction:\nGenerate "
+                "a realistic dating profile bio.\n\n### Response:\n",
+                masked=True,
+            ),
+            Message(
+                role="assistant",
+                content="I'm an outgoing and friendly person who loves spending time with "
+                "friends and family. I'm also a big-time foodie and love trying out new "
+                "restaurants and different cuisines. I'm a big fan of the arts and enjoy "
+                "going to museums and galleries. I'm looking for someone who shares my "
+                "interest in exploring new places, as well as someone who appreciates a "
+                "good conversation over coffee.",
+            ),
+        ]
+        tokens, mask = tokenizer.tokenize_messages(messages, add_eos=False)
+        expected_tokens = [
+            1,
+            323,
+            418,
+            202,
+            31,
+            128,
+            15,
+            120,
+            47,
+            88,
+            584,
+            23,
+            1665,
+            182,
+            9,
+            434,
+            295,
+            85,
+            4,
+            780,
+            47,
+            636,
+            9,
+            1094,
+            213,
+            23,
+            9,
+            69,
+            69,
+            164,
+            1153,
+            299,
+            35,
+            961,
+            132,
+            237,
+            7,
+            5,
+            761,
+            4,
+            12,
+            0,
+            313,
+            120,
+            47,
+            88,
+            584,
+            166,
+            493,
+            171,
+            54,
+            299,
+            9,
+            906,
+            244,
+            19,
+            186,
+            767,
+            303,
+            671,
+            92,
+            209,
+            24,
+            190,
+            52,
+            38,
+            4,
+            12,
+            0,
+            1243,
+            7,
+            69,
+            135,
+            213,
+            166,
+            6,
+            21,
+            45,
+            128,
+            71,
+            58,
+            38,
+            14,
+            10,
+            652,
+            35,
+            462,
+            101,
+            1306,
+            7,
+            341,
+            171,
+            20,
+            14,
+            127,
+            26,
+            652,
+            7,
+            10,
+            1268,
+            4,
+            6,
+            21,
+            45,
+            591,
+            9,
+            566,
+            22,
+            994,
+            913,
+            38,
+            20,
+            52,
+            24,
+            10,
+            1306,
+            734,
+            14,
+            71,
+            365,
+            1382,
+            7,
+            10,
+            801,
+            105,
+            88,
+            244,
+            985,
+            7,
+            4,
+            6,
+            21,
+            45,
+            9,
+            566,
+            126,
+            180,
+            11,
+            5,
+            1137,
+            7,
+            10,
+            1089,
+            151,
+            8,
+            1156,
+            213,
+            342,
+            7,
+            10,
+            384,
+            104,
+            54,
+            470,
+            4,
+            6,
+            21,
+            45,
+            287,
+            14,
+            33,
+            125,
+            135,
+            24,
+            101,
+            512,
+            66,
+            7,
+            28,
+            822,
+            15,
+            542,
+            69,
+            59,
+            110,
+            14,
+            365,
+            229,
+            7,
+            3,
+            36,
+            267,
+            36,
+            125,
+            135,
+            24,
+            101,
+            1503,
+            182,
+            9,
+            222,
+            1661,
+            191,
+            332,
+            92,
+            92,
+            24,
+            24,
+            4,
+            2,
+        ]
+        # Drop eos token.
+        expected_tokens = expected_tokens[:-1]
+        # On 1 less then with eos
+        expected_mask = [True] * 75 + [False] * 124
+        assert expected_tokens == tokens
+        assert expected_mask == mask
diff --git a/tests/torchtune/models/mistral/test_mistral_tokenizer.py b/tests/torchtune/models/mistral/test_mistral_tokenizer.py
index 13fc9cd546..b54cefd03e 100644
--- a/tests/torchtune/models/mistral/test_mistral_tokenizer.py
+++ b/tests/torchtune/models/mistral/test_mistral_tokenizer.py
@@ -472,3 +472,215 @@ def test_tokenize_messages_chat_template(self, messages):
         expected_mask = [True] * 88 + [False] * 125
         assert expected_tokens == tokens
         assert expected_mask == mask
+
+    def test_tokenize_message_drop_eos(self, messages):
+        tokenizer = self.tokenizer(template=False)
+        tokens, mask = tokenizer.tokenize_messages(messages, add_eos=False)
+        expected_tokens = [
+            1,
+            323,
+            418,
+            202,
+            31,
+            128,
+            15,
+            120,
+            47,
+            88,
+            584,
+            23,
+            1665,
+            182,
+            9,
+            434,
+            295,
+            85,
+            4,
+            780,
+            47,
+            636,
+            9,
+            1094,
+            213,
+            23,
+            9,
+            69,
+            69,
+            164,
+            1153,
+            299,
+            35,
+            961,
+            132,
+            237,
+            7,
+            5,
+            761,
+            4,
+            12,
+            0,
+            313,
+            120,
+            47,
+            88,
+            584,
+            166,
+            493,
+            171,
+            54,
+            299,
+            9,
+            906,
+            244,
+            19,
+            186,
+            767,
+            303,
+            671,
+            92,
+            209,
+            24,
+            190,
+            52,
+            38,
+            4,
+            12,
+            0,
+            1243,
+            7,
+            69,
+            135,
+            213,
+            166,
+            6,
+            21,
+            45,
+            128,
+            71,
+            58,
+            38,
+            14,
+            10,
+            652,
+            35,
+            462,
+            101,
+            1306,
+            7,
+            341,
+            171,
+            20,
+            14,
+            127,
+            26,
+            652,
+            7,
+            10,
+            1268,
+            4,
+            6,
+            21,
+            45,
+            591,
+            9,
+            566,
+            22,
+            994,
+            913,
+            38,
+            20,
+            52,
+            24,
+            10,
+            1306,
+            734,
+            14,
+            71,
+            365,
+            1382,
+            7,
+            10,
+            801,
+            105,
+            88,
+            244,
+            985,
+            7,
+            4,
+            6,
+            21,
+            45,
+            9,
+            566,
+            126,
+            180,
+            11,
+            5,
+            1137,
+            7,
+            10,
+            1089,
+            151,
+            8,
+            1156,
+            213,
+            342,
+            7,
+            10,
+            384,
+            104,
+            54,
+            470,
+            4,
+            6,
+            21,
+            45,
+            287,
+            14,
+            33,
+            125,
+            135,
+            24,
+            101,
+            512,
+            66,
+            7,
+            28,
+            822,
+            15,
+            542,
+            69,
+            59,
+            110,
+            14,
+            365,
+            229,
+            7,
+            3,
+            36,
+            267,
+            36,
+            125,
+            135,
+            24,
+            101,
+            1503,
+            182,
+            9,
+            222,
+            1661,
+            191,
+            332,
+            92,
+            92,
+            24,
+            24,
+            4,
+            2,
+        ]
+        # Drop eos token.
+        expected_tokens = expected_tokens[:-1]
+        # On 1 less then with eos
+        expected_mask = [True] * 75 + [False] * 124
+        assert expected_tokens == tokens
+        assert expected_mask == mask
diff --git a/tests/torchtune/models/phi3/test_phi3_tokenizer.py b/tests/torchtune/models/phi3/test_phi3_tokenizer.py
index d28920a480..93fb2c54ce 100644
--- a/tests/torchtune/models/phi3/test_phi3_tokenizer.py
+++ b/tests/torchtune/models/phi3/test_phi3_tokenizer.py
@@ -489,3 +489,246 @@ def test_tokenize_messages_no_system_prompt(self, tokenizer):
         expected_mask = [True] * 77 + [False] * 126
         assert expected_tokens == tokens
         assert expected_mask == mask
+
+    def test_tokenize_messages_drop_eos(self, tokenizer):
+        messages = [
+            Message(role="system", content="You are a helpful assistant", masked=True),
+            Message(
+                role="user",
+                content="Below is an instruction that describes a task. Write a response "
+                "that appropriately completes the request.\n\n### Instruction:\nGenerate "
+                "a realistic dating profile bio.\n\n### Response:\n",
+                masked=True,
+            ),
+            Message(
+                role="assistant",
+                content="I'm an outgoing and friendly person who loves spending time with "
+                "friends and family. I'm also a big-time foodie and love trying out new "
+                "restaurants and different cuisines. I'm a big fan of the arts and enjoy "
+                "going to museums and galleries. I'm looking for someone who shares my "
+                "interest in exploring new places, as well as someone who appreciates a "
+                "good conversation over coffee.",
+            ),
+        ]
+        tokens, mask = tokenizer.tokenize_messages(messages, add_eos=False)
+        expected_tokens = [
+            1,
+            32006,
+            272,
+            84,
+            9,
+            615,
+            454,
+            1974,
+            19,
+            32007,
+            32010,
+            323,
+            418,
+            202,
+            31,
+            128,
+            15,
+            120,
+            47,
+            88,
+            584,
+            23,
+            1665,
+            182,
+            9,
+            434,
+            295,
+            85,
+            4,
+            780,
+            47,
+            636,
+            9,
+            1094,
+            213,
+            23,
+            9,
+            69,
+            69,
+            164,
+            1153,
+            299,
+            35,
+            961,
+            132,
+            237,
+            7,
+            5,
+            761,
+            4,
+            12,
+            0,
+            313,
+            120,
+            47,
+            88,
+            584,
+            166,
+            493,
+            171,
+            54,
+            299,
+            9,
+            906,
+            244,
+            19,
+            186,
+            767,
+            303,
+            671,
+            92,
+            209,
+            24,
+            190,
+            52,
+            38,
+            4,
+            12,
+            0,
+            1243,
+            7,
+            69,
+            135,
+            213,
+            166,
+            32007,
+            32001,
+            6,
+            21,
+            45,
+            128,
+            71,
+            58,
+            38,
+            14,
+            10,
+            652,
+            35,
+            462,
+            101,
+            1306,
+            7,
+            341,
+            171,
+            20,
+            14,
+            127,
+            26,
+            652,
+            7,
+            10,
+            1268,
+            4,
+            6,
+            21,
+            45,
+            591,
+            9,
+            566,
+            22,
+            994,
+            913,
+            38,
+            20,
+            52,
+            24,
+            10,
+            1306,
+            734,
+            14,
+            71,
+            365,
+            1382,
+            7,
+            10,
+            801,
+            105,
+            88,
+            244,
+            985,
+            7,
+            4,
+            6,
+            21,
+            45,
+            9,
+            566,
+            126,
+            180,
+            11,
+            5,
+            1137,
+            7,
+            10,
+            1089,
+            151,
+            8,
+            1156,
+            213,
+            342,
+            7,
+            10,
+            384,
+            104,
+            54,
+            470,
+            4,
+            6,
+            21,
+            45,
+            287,
+            14,
+            33,
+            125,
+            135,
+            24,
+            101,
+            512,
+            66,
+            7,
+            28,
+            822,
+            15,
+            542,
+            69,
+            59,
+            110,
+            14,
+            365,
+            229,
+            7,
+            3,
+            36,
+            267,
+            36,
+            125,
+            135,
+            24,
+            101,
+            1503,
+            182,
+            9,
+            222,
+            1661,
+            191,
+            332,
+            92,
+            92,
+            24,
+            24,
+            4,
+            32007,
+        ]
+
+        # Drop eos token
+        expected_tokens = expected_tokens[:]
+        # On 1 less then with eos
+        expected_mask = [True] * 86 + [False] * 126
+        assert expected_tokens == tokens
+        assert expected_mask == mask
diff --git a/tests/torchtune/models/qwen2/test_qwen2_tokenizer.py b/tests/torchtune/models/qwen2/test_qwen2_tokenizer.py
index 9416ed7f5a..b11d6679e1 100644
--- a/tests/torchtune/models/qwen2/test_qwen2_tokenizer.py
+++ b/tests/torchtune/models/qwen2/test_qwen2_tokenizer.py
@@ -444,3 +444,194 @@ def test_tokenize_messages_gt_max_seq_len(self, messages):
         tokens, mask = tokenizer.tokenize_messages(messages)
         assert len(tokens) == 10
         assert len(mask) == 10
+
+    def test_tokenize_message_drop_eos(self, messages):
+        tokenizer = self.tokenizer(template=False)
+        expected_tokens = [
+            33,
+            214,
+            174,
+            156,
+            194,
+            130,
+            197,
+            184,
+            446,
+            789,
+            113,
+            98,
+            1914,
+            13,
+            346,
+            788,
+            98,
+            706,
+            102,
+            182,
+            184,
+            1916,
+            176,
+            762,
+            83,
+            113,
+            103,
+            874,
+            269,
+            13,
+            94,
+            94,
+            2,
+            2,
+            2,
+            483,
+            197,
+            25,
+            94,
+            885,
+            98,
+            1226,
+            1960,
+            348,
+            114,
+            1123,
+            399,
+            1583,
+            78,
+            13,
+            94,
+            94,
+            2,
+            2,
+            2,
+            360,
+            1733,
+            102,
+            182,
+            25,
+            94,
+            40,
+            1791,
+            194,
+            453,
+            70,
+            78,
+            114,
+            120,
+            967,
+            176,
+            618,
+            628,
+            1275,
+            794,
+            294,
+            1095,
+            445,
+            212,
+            1356,
+            120,
+            1299,
+            13,
+            223,
+            1791,
+            451,
+            98,
+            127,
+            181,
+            1047,
+            375,
+            915,
+            380,
+            120,
+            1448,
+            1732,
+            114,
+            453,
+            447,
+            1219,
+            64,
+            187,
+            921,
+            120,
+            742,
+            107,
+            84,
+            122,
+            893,
+            13,
+            223,
+            1791,
+            98,
+            127,
+            181,
+            123,
+            124,
+            131,
+            103,
+            744,
+            82,
+            120,
+            1506,
+            416,
+            114,
+            128,
+            1429,
+            182,
+            253,
+            82,
+            120,
+            163,
+            330,
+            105,
+            262,
+            13,
+            223,
+            1791,
+            155,
+            1551,
+            171,
+            1951,
+            628,
+            296,
+            64,
+            237,
+            886,
+            1390,
+            130,
+            883,
+            1678,
+            447,
+            306,
+            279,
+            113,
+            11,
+            215,
+            785,
+            215,
+            1951,
+            628,
+            378,
+            101,
+            66,
+            72,
+            593,
+            98,
+            984,
+            208,
+            1580,
+            167,
+            510,
+            737,
+            318,
+            1278,
+            13,
+            2000,
+        ]
+
+        # Remove the EOS token
+        expected_tokens = expected_tokens[:-1]
+        # On 1 less then with eos
+        expected_mask = [True] * 61 + [False] * 115
+
+        tokens, mask = tokenizer.tokenize_messages(messages, add_eos=False)
+        assert tokens == expected_tokens
+        assert mask == expected_mask


diff --git a/torchtune/datasets/_packed.py b/torchtune/datasets/_packed.py
index 8cee72eed9..69127f646e 100644
--- a/torchtune/datasets/_packed.py
+++ b/torchtune/datasets/_packed.py
@@ -150,8 +153,7 @@ def _pack(self) -> None:
            i
@@ -161,6 +163,13 @@ def _pack(self) -> None:
    def _should_stop_packing(

diff --git a/torchtune/datasets/_packed.py b/torchtune/datasets/_packed.py
index 8cee72eed9..69127f646e 100644
--- a/torchtune/datasets/_packed.py
+++ b/torchtune/datasets/_packed.py
@@ -136,12 +136,15 @@ def _pack(self) -> None:
             # Update the current pack
             current_pack["tokens"] += tokens
             current_pack["labels"] += labels
-            current_pack["input_pos"] += list(range(seq_len))
+            current_pack["input_pos"] += [x % self.max_seq_len for x in range(seq_len)]
             current_pack["seq_lens"] += [seq_len]

             # If the current pack is over the max_seq_len, add it to self.packs and
             # retain any truncated or bumped samples for next pack
-            if len(current_pack["tokens"]) > self.max_seq_len:
+            while (
+                len(current_pack["tokens"]) > self.max_seq_len
+                and not self._should_stop_packing()
+            ):
                 current_pack = self._split_and_add_pack(current_pack)

             if rank == 0:
@@ -150,8 +153,7 @@ def _pack(self) -> None:
             # Keep track of previous sample boundary
             self.previous_sample_boundary = len(current_pack["tokens"])

-            # If max packs is set, stop packing when we reach that number
-            if self.max_packs is not None and len(self.packs) == self.max_packs:
+            if self._should_stop_packing():
                 break

         # Handle the last pack if there's leftover and we haven't filled up the max packs
@@ -161,6 +163,13 @@ def _pack(self) -> None:
             # No need to handle splitting at this point so we can just add the current pack
             self._add_pack(current_pack)

+    def _should_stop_packing(self) -> bool:
+        """If max packs is set, stop packing when we reach that number."""
+
+        if self.max_packs is not None and len(self.packs) == self.max_packs:
+            return True
+        return False
+
     def _split_and_add_pack(self, current_pack: PACK_TYPE) -> PACK_TYPE:
         """Splits the current pack at the boundary, processes it, adds it to ``self.packs`` and
         returns the start of the next pack."""

Wrap long sequences to max_seq_len and continuously split packs until either reaching max_packs or having a valid-length pack.

Prompt:
The following is the text of a github issue and related comments for this repository:

PackedDataset cannot handle long sequence whose length is larger than 2*max_seq_len when using split_across_pack=True
As the title, it reports a runtime error in the self._pad_pack()
```
f = {"tokens":list(range(121)),"labels":list(range(121))}
x = PackedDataset([f],max_seq_len=60,split_across_pack=True)
for i in x:
    print(i)
```
```
    num_range = torch.arange(
                ^^^^^^^^^^^^^
RuntimeError: upper bound and larger bound inconsistent with step sign
```

Resolve the issue described above. Do not modify tests or user code. The issue can be resolved by modifying the actual library/application code. Make the minimal changes to resolve the issue while adhering to the specification and coding style of the codebase.

diff --git a/torchtune/datasets/_packed.py b/torchtune/datasets/_packed.py
index 8cee72eed9..69127f646e 100644
--- a/torchtune/datasets/_packed.py
+++ b/torchtune/datasets/_packed.py
@@ -150,8 +153,7 @@ def _pack(self) -> None:
             # Keep track of previous sample boundary
             self.previous_sample_boundary = len(current_pack["tokens"])

-            # If max packs is set, stop packing when we reach that number
-            if self.max_packs is not None and len(self.packs) == self.max_packs:
+            if self._should_stop_packing():
                 break

         # Handle the last pack if there's leftover and we haven't filled up the max packs
@@ -161,6 +163,13 @@ def _pack(self) -> None:
             # No need to handle splitting at this point so we can just add the current pack
             self._add_pack(current_pack)

+    def _should_stop_packing(self) -> bool:
+        """If max packs is set, stop packing when we reach that number."""
+
+        if self.max_packs is not None and len(self.packs) == self.max_packs:
+            return True
+        return False
+
     def _split_and_add_pack(self, current_pack: PACK_TYPE) -> PACK_TYPE:
         """Splits the current pack at the boundary, processes it, adds it to ``self.packs`` and
         returns the start of the next pack."""

diff --git a/tests/torchtune/datasets/test_packed_dataset.py b/tests/torchtune/datasets/test_packed_dataset.py
index 208ac333f5..4d587dc98f 100644
--- a/tests/torchtune/datasets/test_packed_dataset.py
+++ b/tests/torchtune/datasets/test_packed_dataset.py
@@ -147,6 +147,34 @@ def test_packed_dataset(
         torch.testing.assert_close(packed[0]["seq_lens"], expected_seq_lens)
         torch.testing.assert_close(packed[0]["input_pos"], expected_input_pos)

+    @pytest.mark.parametrize("max_seq_len", [13])
+    @pytest.mark.parametrize("sample_size", [14, 27, 40])
+    @pytest.mark.parametrize("max_packs", [5, 200, 3100])
+    @pytest.mark.parametrize("split_across_pack", [True])
+    def test_chunked_case(self, max_seq_len, sample_size, max_packs, split_across_pack):
+        dataset = DummyDataset(sample_size)
+        packed = PackedDataset(
+            dataset,
+            max_seq_len=max_seq_len,
+            max_packs=max_packs,
+            split_across_pack=split_across_pack,
+        )
+
+        # Check we get right number of packs
+        correct_num_packs = self._calculate_num_packs(
+            len(dataset), max_seq_len, sample_size, split_across_pack, max_packs
+        )
+        assert len(packed) == correct_num_packs
+
+        # Check all fields are same length
+        assert all(
+            len(pack["tokens"]) == len(pack["labels"]) == len(pack["input_pos"])
+            for pack in packed
+        )
+
+        # Check that all sum(seq_lens) are equal to max_seq_len
+        assert all(pack["seq_lens"].sum().item() == max_seq_len for pack in packed)
+
     def test_packed_dataset_real_data(self):
         expected_tokenized_prompts = [
             torch.tensor([0, 4, 2, 1, 7, 4, -1, 0, 1, 9]),


diff --git a/docs/source/api_ref_datasets.rst b/docs/source/api_ref_datasets.rst
index cc0f6da466..40def346e4 100644
--- a/docs/source/api_ref_datasets.rst
+++ b/docs/source/api_ref_datasets.rst
@@ -64,8 +64,6 @@ Class representations for the above dataset builders.
     :toctree: generated/
     :nosignatures:

-    InstructDataset
-    ChatDataset
     TextCompletionDataset
     ConcatDataset
     PackedDataset
diff --git a/docs/source/deep_dives/configs.rst b/docs/source/deep_dives/configs.rst
index 54eef7144d..0f86a29a58 100644
--- a/docs/source/deep_dives/configs.rst
+++ b/docs/source/deep_dives/configs.rst
@@ -119,7 +119,7 @@ keyword arguments not specified in the config if we'd like:
         tokenizer: ModelTokenizer,
         train_on_input: bool = True,
         max_seq_len: int = 512,
-    ) -> InstructDataset:
+    ) -> SFTDataset:

     from torchtune import config

diff --git a/docs/source/tutorials/datasets.rst b/docs/source/tutorials/datasets.rst
index a45239a523..781573b89e 100644
--- a/docs/source/tutorials/datasets.rst
+++ b/docs/source/tutorials/datasets.rst
@@ -491,15 +491,10 @@ Fully customized datasets
 -------------------------

 More advanced tasks and dataset formats that don't fit into the templating and processing
-that :class:`~torchtune.datasets.InstructDataset`, :class:`~torchtune.datasets.ChatDataset`,
-and :class:`~torchtune.datasets.TextCompletionDataset` provide may require you to create your own dataset
-class for more flexibility. Let's walk through the :class:`~torchtune.datasets.PreferenceDataset`,
+that :class:`~torchtune.datasets.SFTDataset` and :class:`~torchtune.datasets.TextCompletionDataset` provide may require
+you to create your own dataset class for more flexibility. Let's walk through the :class:`~torchtune.datasets.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 :class:`~torchtune.datasets.PreferenceDataset` class,
-you'll notice it's quite similar to :class:`~torchtune.datasets.InstructDataset` with a few
-adjustments for chosen and rejected samples in preference data.
-
 .. code-block:: python

     chosen_message = [
diff --git a/torchtune/datasets/__init__.py b/torchtune/datasets/__init__.py
index de2e22beda..b0c7c11738 100644
--- a/torchtune/datasets/__init__.py
+++ b/torchtune/datasets/__init__.py
@@ -6,12 +6,12 @@

 from torchtune.datasets import multimodal
 from torchtune.datasets._alpaca import alpaca_cleaned_dataset, alpaca_dataset
-from torchtune.datasets._chat import chat_dataset, ChatDataset
+from torchtune.datasets._chat import chat_dataset
 from torchtune.datasets._cnn_dailymail import cnn_dailymail_articles_dataset
 from torchtune.datasets._concat import ConcatDataset
 from torchtune.datasets._grammar import grammar_dataset
 from torchtune.datasets._hh_rlhf_helpful import hh_rlhf_helpful_dataset
-from torchtune.datasets._instruct import instruct_dataset, InstructDataset
+from torchtune.datasets._instruct import instruct_dataset
 from torchtune.datasets._packed import PackedDataset
 from torchtune.datasets._preference import preference_dataset, PreferenceDataset
 from torchtune.datasets._samsum import samsum_dataset
@@ -30,9 +30,7 @@
     "grammar_dataset",
     "samsum_dataset",
     "stack_exchange_paired_dataset",
-    "InstructDataset",
     "slimorca_dataset",
-    "ChatDataset",
     "instruct_dataset",
     "preference_dataset",
     "chat_dataset",
diff --git a/torchtune/datasets/_chat.py b/torchtune/datasets/_chat.py
index f18961d36f..b9f5639706 100644
--- a/torchtune/datasets/_chat.py
+++ b/torchtune/datasets/_chat.py
@@ -4,111 +4,12 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.

-from typing import Any, Callable, Dict, List, Mapping, Optional, Union
+from typing import Any, Dict, Optional, Union

-import numpy as np
-
-from datasets import load_dataset
-from torch.utils.data import Dataset
-from torchtune.data._chat_formats import ChatFormat
-from torchtune.data._common import CROSS_ENTROPY_IGNORE_IDX
-from torchtune.data._messages import (
-    Message,
-    OpenAIToMessages,
-    ShareGPTToMessages,
-    validate_messages,
-)
+from torchtune.data._messages import OpenAIToMessages, ShareGPTToMessages
 from torchtune.datasets._packed import PackedDataset
 from torchtune.datasets._sft import SFTDataset
 from torchtune.modules.tokenizers import ModelTokenizer
-from torchtune.utils._logging import deprecated
-
-
-@deprecated(msg="Please use `torchtune.datasets.SFTDataset` for custom chat data.")
-class ChatDataset(Dataset):
-    """
-    Note:
-        This class is deprecated and will be removed in a future release. Please use
-        :class:`~torchtune.datasets.SFTDataset` or :func:`~torchtune.datasets.chat_dataset`
-        for custom chat data.
-
-    Class that supports any custom dataset with multiturn conversations.
-
-    The general flow from loading a sample to tokenized prompt is:
-    load sample -> apply transform -> foreach turn{format into template -> tokenize}
-
-    Use ``convert_to_messages`` to prepare your dataset into the Llama2 chat format
-    and roles::
-
-        [
-            Message(
-                role=<system|user|assistant>,
-                content=<message>,
-            ),
-            ...
-        ]
-
-    This class supports multi-turn conversations. If a tokenizer sample with multiple
-    turns does not fit within ``max_seq_len`` then it is truncated.
-
-    Args:
-        tokenizer (ModelTokenizer): Tokenizer used by the model that implements the ``tokenize_messages`` method.
-        source (str): path to dataset repository on Hugging Face. For local datasets,
-            define source as the data file type (e.g. "json", "csv", "text") and pass
-            in the filepath in ``data_files``. See Hugging Face's ``load_dataset``
-            (https://huggingface.co/docs/datasets/en/package_reference/loading_methods#datasets.load_dataset.path)
-            for more details.
-        convert_to_messages (Callable[[Mapping[str, Any]], List[Message]]): function that keys into the desired field in the sample
-            and converts to a list of :class:`~torchtune.data.Message` that follows the Llama format with the expected keys
-        chat_format (Optional[ChatFormat]): template used to format the chat. This is used to add structured text around the actual
-            messages, such as the [INST] tags in Llama2 and in Mistral. The extra text will still get tokenized as normal text, not
-            as special tokens. In models like Llama3 where the tokenizer adds tags as special tokens, ``chat_format`` is not needed,
-            unless you want to structure messages in a particular way for inference.
-        max_seq_len (int): Maximum number of tokens in the returned input and label token id lists.
-        train_on_input (bool): Whether the model is trained on the prompt or not. Default is False.
-        **load_dataset_kwargs (Dict[str, Any]): additional keyword arguments to pass to ``load_dataset``,
-            such as ``data_files`` or ``split``.
-    """
-
-    def __init__(
-        self,
-        *,
-        tokenizer: ModelTokenizer,
-        source: str,
-        convert_to_messages: Callable[[Mapping[str, Any]], List[Message]],
-        chat_format: Optional[ChatFormat] = None,
-        max_seq_len: int,
-        train_on_input: bool = False,
-        **load_dataset_kwargs: Dict[str, Any],
-    ) -> None:
-
-        self._tokenizer = tokenizer
-        self._data = load_dataset(source, **load_dataset_kwargs)
-        self._convert_to_messages = convert_to_messages
-        self.chat_format = chat_format
-        self.max_seq_len = max_seq_len
-        self.train_on_input = train_on_input
-
-    def __len__(self):
-        return len(self._data)
-
-    def __getitem__(self, index: int) -> Dict[str, List[int]]:
-        sample = self._data[index]
-        return self._prepare_sample(sample)
-
-    def _prepare_sample(self, sample: Mapping[str, Any]) -> Dict[str, List[int]]:
-        messages = self._convert_to_messages(sample, self.train_on_input)
-        if self.chat_format is not None:
-            messages = self.chat_format.format(messages)
-        validate_messages(messages)
-        tokens, mask = self._tokenizer.tokenize_messages(
-            messages,
-        )
-        # Wherever mask == True, set to CROSS_ENTROPY_IGNORE_IDX. Otherwise keep as tokens
-        labels = list(np.where(mask, CROSS_ENTROPY_IGNORE_IDX, tokens))
-        assert len(tokens) == len(labels)
-
-        return {"tokens": tokens, "labels": labels}

 def chat_dataset(
diff --git a/torchtune/datasets/_instruct.py b/torchtune/datasets/_instruct.py
index e291feadb3..82d9bb267e 100644
--- a/torchtune/datasets/_instruct.py
+++ b/torchtune/datasets/_instruct.py
@@ -4,130 +4,12 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.

-from typing import Any, Callable, Dict, List, Mapping, Optional, Union
+from typing import Any, Dict, Optional, Union

-import numpy as np
-from datasets import load_dataset
-from torch.utils.data import Dataset
-from torchtune.data import (
-    CROSS_ENTROPY_IGNORE_IDX,
-    InputOutputToMessages,
-    InstructTemplate,
-    Message,
-    validate_messages,
-)
+from torchtune.data import InputOutputToMessages
 from torchtune.datasets._packed import PackedDataset
 from torchtune.datasets._sft import SFTDataset
 from torchtune.modules.tokenizers import ModelTokenizer
-from torchtune.utils._logging import deprecated
-
-
-@deprecated(
-    msg="Please use `torchtune.datasets.SFTDataset` or :func:`~torchtune.datasets.instruct_dataset` for custom instruct data."
-)
-class InstructDataset(Dataset):
-    """
-    Note:
-        This class is deprecated and will be removed in a future release. Please use
-        :class:`~torchtune.datasets.SFTDataset` or :func:`~torchtune.datasets.instruct_dataset`
-        for custom instruct data.
-
-    Class that supports any custom dataset with instruction-based prompts and a
-    configurable template.
-
-    The general flow from loading a sample to tokenized prompt is:
-    load sample -> apply transform -> format into template -> tokenize
-
-    If the column/key names differ from the expected names in the :class:`~torchtune.data.InstructTemplate`,
-    then the ``column_map`` argument can be used to provide this mapping.
-
-    Masking of the prompt during training is controlled by the ``train_on_input`` flag, which is
-    set to ``False`` by default.
-    - If ``train_on_input`` is True, the prompt is used during training and
-    contributes to the loss.
-    - If ``train_on_input`` is False, the prompt is masked out (tokens replaced with -100)
-
-    Args:
-        tokenizer (ModelTokenizer): Tokenizer used by the model that implements the ``tokenize_messages`` method.
-        source (str): path to dataset repository on Hugging Face. For local datasets,
-            define source as the data file type (e.g. "json", "csv", "text") and pass
-            in the filepath in ``data_files``. See Hugging Face's ``load_dataset``
-            (https://huggingface.co/docs/datasets/en/package_reference/loading_methods#datasets.load_dataset.path)
-            for more details.
-        template (InstructTemplate): template used to format the prompt. If the placeholder variable
-            names in the template do not match the column/key names in the dataset, use ``column_map`` to map them.
-        transform (Optional[Callable]): transform to apply to the sample before formatting to the template.
-            Default is None.
-        column_map (Optional[Dict[str, str]]): a mapping from the expected placeholder names in the template
-            to the column/key names in the sample. If None, assume these are identical.
-            The output column can be indicated using the ``output`` key mapping.
-            If no placeholder for the ``output`` column is provided in ``column_map`` it is assumed to be ``output``.
-        train_on_input (bool): Whether the model is trained on the prompt or not. Default is False.
-        max_seq_len (Optional[int]): Maximum number of tokens in the returned input and label token id lists.
-            Default is None, disabling truncation. We recommend setting this to the highest you can fit in memory
-            and is supported by the model. For example, llama2-7B supports up to 4096 for sequence length.
-        **load_dataset_kwargs (Dict[str, Any]): additional keyword arguments to pass to ``load_dataset``,
-            such as ``data_files`` or ``split``.
-    Raises:
-        ValueError: If ``template`` is not an instance of :class:`torchtune.data.InstructTemplate`
-    """
-
-    def __init__(
-        self,
-        tokenizer: ModelTokenizer,
-        source: str,
-        template: InstructTemplate,
-        transform: Optional[Callable] = None,
-        column_map: Optional[Dict[str, str]] = None,
-        train_on_input: bool = False,
-        max_seq_len: Optional[int] = None,
-        **load_dataset_kwargs: Dict[str, Any],
-    ) -> None:
-        if not isinstance(template(), InstructTemplate):
-            raise ValueError(
-                f"template must be an InstructTemplate class, not {type(template())}"
-            )
-
-        self._tokenizer = tokenizer
-        self._data = load_dataset(source, **load_dataset_kwargs)
-        self.template = template
-        self._transform = transform
-        self._column_map = column_map
-        self.train_on_input = train_on_input
-        self.max_seq_len = max_seq_len
-
-    def __len__(self):
-        return len(self._data)
-
-    def __getitem__(self, index: int) -> Dict[str, List[int]]:
-        sample = self._data[index]
-        return self._prepare_sample(sample)
-
-    def _prepare_sample(self, sample: Mapping[str, Any]) -> Dict[str, List[int]]:
-        transformed_sample = self._transform(sample) if self._transform else sample
-
-        prompt = self.template.format(transformed_sample, self._column_map)
-        key_output = (
-            self._column_map["output"]
-            if self._column_map and "output" in self._column_map
-            else "output"
-        )
-        messages = [
-            Message(role="user", content=prompt, masked=(not self.train_on_input)),
-            Message(role="assistant", content=transformed_sample[key_output]),
-        ]
-
-        validate_messages(messages)
-
-        tokens, mask = self._tokenizer.tokenize_messages(
-            messages,
-        )
-
-        # Wherever mask == True, set to CROSS_ENTROPY_IGNORE_IDX. Otherwise keep as tokens
-        labels = list(np.where(mask, CROSS_ENTROPY_IGNORE_IDX, tokens))
-        assert len(tokens) == len(labels)
-
-        return {"tokens": tokens, "labels": labels}

 def instruct_dataset(

Prompt:
The following is the text of a github issue and related comments for this repository:

Delete `ChatDataset` and `InstructDataset`
These have been deprecated for a cycle and can now be removed.

Resolve the issue described above. Do not modify tests or user code. The issue can be resolved by modifying the actual library/application code. Make the minimal changes to resolve the issue while adhering to the specification and coding style of the codebase.

diff --git a/tests/torchtune/datasets/test_chat_dataset.py b/tests/torchtune/datasets/test_chat_dataset.py
index 23db9569fe..8e99ad92ae 100644
--- a/tests/torchtune/datasets/test_chat_dataset.py
+++ b/tests/torchtune/datasets/test_chat_dataset.py
@@ -7,9 +7,8 @@
 import pytest
 from tests.common import ASSETS
 from tests.test_utils import DummyChatFormat, DummyTokenizer
-from torchtune.data import get_sharegpt_messages
 from torchtune.data._common import CROSS_ENTROPY_IGNORE_IDX
-from torchtune.datasets import chat_dataset, ChatDataset
+from torchtune.datasets import chat_dataset

 class TestChatDataset:
@@ -18,62 +17,6 @@ def chat_format(self):
         return DummyChatFormat

     def test_get_item(self, chat_format):
-        expected_tokenized_prompts = [
-            [
-                0,
-                7,
-                3,
-                3,
-                2,
-                2,
-                10,
-                5,
-                4,
-                2,
-                3,
-                7,
-                2,
-                5,
-                10,
-                3,
-                7,
-                2,
-                4,
-                2,
-                3,
-                -1,
-                0,
-                5,
-                6,
-                11,
-                10,
-                1,
-                6,
-                -1,
-            ]
-        ]
-        prompt_lengths = (15, 5)
-        expected_labels = [
-            [CROSS_ENTROPY_IGNORE_IDX] * prompt_lengths[0]
-            + [3, 7, 2, 4, 2, 3, -1]
-            + [CROSS_ENTROPY_IGNORE_IDX] * prompt_lengths[1]
-            + [1, 6, -1]
-        ]
-        ds = ChatDataset(
-            tokenizer=DummyTokenizer(),
-            source="json",
-            convert_to_messages=get_sharegpt_messages,
-            chat_format=chat_format,
-            max_seq_len=100,
-            train_on_input=False,
-            data_files=str(ASSETS / "chat_tiny.json"),
-            split="train",
-        )
-        assert len(ds) == 1
-        prompt, label = ds[0]["tokens"], ds[0]["labels"]
-        assert prompt == expected_tokenized_prompts[0]
-        assert label == expected_labels[0]
-
         expected_tokenized_prompts = [
             [
                 0,
diff --git a/tests/torchtune/datasets/test_instruct_dataset.py b/tests/torchtune/datasets/test_instruct_dataset.py
index c734bec885..04f3c2f49c 100644
--- a/tests/torchtune/datasets/test_instruct_dataset.py
+++ b/tests/torchtune/datasets/test_instruct_dataset.py
@@ -9,7 +9,7 @@
 from tests.test_utils import DummyTokenizer
 from torchtune.data import InstructTemplate
 from torchtune.data._common import CROSS_ENTROPY_IGNORE_IDX
-from torchtune.datasets import instruct_dataset, InstructDataset
+from torchtune.datasets import instruct_dataset

 def dummy_transform(sample):
@@ -29,38 +29,6 @@ def format(cls, sample, column_map):
 class TestInstructDataset:
     @pytest.mark.parametrize("train_on_input", [True, False])
     def test_get_item(self, train_on_input):
-        template = DummyTemplate
-        expected_tokenized_prompts = [
-            [0, 12, 4, 4, 2, 2, 2, 7, 10, 9, 2, 2, 5, 2, 2, 6, 10, -1],
-            [0, 12, 2, 2, 8, 2, 15, 10, 9, 8, 3, 15, 3, 4, 9, 3, 15, 10, -1],
-        ]
-        prompt_lengths = (10, 9)
-        expected_labels = [
-            [CROSS_ENTROPY_IGNORE_IDX] * prompt_lengths[0] + [2, 2, 5, 2, 2, 6, 10, -1],
-            [CROSS_ENTROPY_IGNORE_IDX] * prompt_lengths[1]
-            + [8, 3, 15, 3, 4, 9, 3, 15, 10, -1],
-        ]
-
-        dataset = InstructDataset(
-            tokenizer=DummyTokenizer(),
-            source="json",
-            template=template,
-            transform=dummy_transform,
-            train_on_input=train_on_input,
-            data_files=str(ASSETS / "instruct_tiny.json"),
-            column_map={"output": "response"},
-            split="train",
-        )
-        assert len(dataset) == 2
-
-        for i in range(len(dataset)):
-            prompt, label = dataset[i]["tokens"], dataset[i]["labels"]
-            assert prompt == expected_tokenized_prompts[i]
-            if train_on_input:
-                assert label == expected_tokenized_prompts[i]
-            else:
-                assert label == expected_labels[i]
-
         expected_tokenized_prompts = [
             [0, 6, 4, 6, 4, 4, 2, 2, 2, 7, 2, 2, 5, 2, 2, 6, -1],
             [0, 6, 4, 6, 2, 2, 8, 2, 15, 8, 3, 15, 3, 4, 9, 3, 15, -1],


diff --git a/docs/source/api_ref_modules.rst b/docs/source/api_ref_modules.rst
index d1fb60144a..76dae8600d 100644
--- a/docs/source/api_ref_modules.rst
+++ b/docs/source/api_ref_modules.rst
@@ -113,20 +113,3 @@ Functions used for preprocessing images.

     transforms.Transform
     transforms.VisionCrossAttentionMask
-
-Reinforcement Learning From Human Feedback (RLHF)
---------------------------------------------------
-Components and losses for RLHF algorithms like PPO and DPO:
-
-.. autosummary::
-   :toctree: generated/
-   :nosignatures:
-
-    rlhf.estimate_advantages
-    rlhf.get_rewards_ppo
-    rlhf.truncate_sequence_at_first_stop_token
-    rlhf.loss.PPOLoss
-    rlhf.loss.DPOLoss
-    rlhf.loss.RSOLoss
-    rlhf.loss.IPOLoss
-    rlhf.loss.SimPOLoss
diff --git a/docs/source/api_ref_rlhf.rst b/docs/source/api_ref_rlhf.rst
new file mode 100644
index 0000000000..e68ca8aed1
--- /dev/null
+++ b/docs/source/api_ref_rlhf.rst
@@ -0,0 +1,20 @@
+===============
+torchtune.rlhf
+===============
+
+.. currentmodule:: torchtune.rlhf
+
+Components and losses for RLHF algorithms like PPO and DPO:
+
+.. autosummary::
+   :toctree: generated/
+   :nosignatures:
+
+    estimate_advantages
+    get_rewards_ppo
+    truncate_sequence_at_first_stop_token
+    loss.PPOLoss
+    loss.DPOLoss
+    loss.RSOLoss
+    loss.IPOLoss
+    loss.SimPOLoss
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 98676e4ce9..f3a95281fe 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -155,5 +155,6 @@ torchtune tutorials.
    api_ref_generation
    api_ref_models
    api_ref_modules
+   api_ref_rlhf
    api_ref_training
    api_ref_utilities
diff --git a/recipes/configs/llama2/7B_lora_dpo.yaml b/recipes/configs/llama2/7B_lora_dpo.yaml
index 78b16c6e04..f6acfcb76e 100644
--- a/recipes/configs/llama2/7B_lora_dpo.yaml
+++ b/recipes/configs/llama2/7B_lora_dpo.yaml
@@ -62,7 +62,7 @@ lr_scheduler:
   num_warmup_steps: 100

 loss:
-  _component_: torchtune.modules.rlhf.loss.DPOLoss
+  _component_: torchtune.rlhf.loss.DPOLoss
   beta: 0.1
   label_smoothing: 0

diff --git a/recipes/configs/llama2/7B_lora_dpo_single_device.yaml b/recipes/configs/llama2/7B_lora_dpo_single_device.yaml
index 5fb50ff6e8..d51847d1e7 100644
--- a/recipes/configs/llama2/7B_lora_dpo_single_device.yaml
+++ b/recipes/configs/llama2/7B_lora_dpo_single_device.yaml
@@ -61,7 +61,7 @@ lr_scheduler:
   num_warmup_steps: 100

 loss:
-  _component_: torchtune.modules.rlhf.loss.DPOLoss
+  _component_: torchtune.rlhf.loss.DPOLoss
   beta: 0.1
   label_smoothing: 0

diff --git a/recipes/configs/mistral/7B_full_ppo_low_memory.yaml b/recipes/configs/mistral/7B_full_ppo_low_memory.yaml
index adf11ce486..1cf7dd974a 100644
--- a/recipes/configs/mistral/7B_full_ppo_low_memory.yaml
+++ b/recipes/configs/mistral/7B_full_ppo_low_memory.yaml
@@ -167,7 +167,7 @@ lmbda: 0.95

 # PPO hyperparameters
 loss:
-  _component_: torchtune.modules.rlhf.loss.PPOLoss
+  _component_: torchtune.rlhf.loss.PPOLoss
   epsilon: 0.2
   value_coeff: 0.1
   value_clip_range: 0.2
diff --git a/recipes/lora_dpo_distributed.py b/recipes/lora_dpo_distributed.py
index 1d928e4fc8..d381dc9832 100644
--- a/recipes/lora_dpo_distributed.py
+++ b/recipes/lora_dpo_distributed.py
@@ -18,10 +18,9 @@
 from torch.distributed import destroy_process_group, init_process_group
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader, DistributedSampler
-from torchtune import config, modules, training, utils
+from torchtune import config, modules, rlhf, training, utils
 from torchtune.data import CROSS_ENTROPY_IGNORE_IDX, padded_collate_dpo
 from torchtune.datasets import ConcatDataset
-from torchtune.modules import rlhf
 from torchtune.modules.peft import (
     disable_adapter,
     DoRALinear,
@@ -32,8 +31,8 @@
     set_trainable_params,
     validate_missing_and_unexpected_for_lora,
 )
-from torchtune.modules.rlhf.loss import SimPOLoss
 from torchtune.recipe_interfaces import FTRecipeInterface
+from torchtune.rlhf.loss import SimPOLoss
 from tqdm import tqdm

 log = utils.get_logger("DEBUG")
@@ -95,10 +94,10 @@ class LoRADPORecipeDistributed(FTRecipeInterface):
         - Logging. Terminal, Disk, WandB and TensorBoard are all supported.

     The following losses are supported in this recipe:
-        - :class:`~torchtune.modules.rlhf.loss.DPOLoss`: Direct Preference Optimization (DPO).
-        - :class:`~torchtune.modules.rlhf.loss.RSOPLoss`: Rejection Sampling Optimization (RSO).
-        - :class:`~torchtune.modules.rlhf.loss.IPO`: Identity Preference Optimization (IPO).
-        - :class:`~torchtune.modules.rlhf.loss.SimPOLoss`: Simple Preference Optimization (SimPO).
+        - :class:`~torchtune.rlhf.loss.DPOLoss`: Direct Preference Optimization (DPO).
+        - :class:`~torchtune.rlhf.loss.RSOPLoss`: Rejection Sampling Optimization (RSO).
+        - :class:`~torchtune.rlhf.loss.IPO`: Identity Preference Optimization (IPO).
+        - :class:`~torchtune.rlhf.loss.SimPOLoss`: Simple Preference Optimization (SimPO).

     For a full list of example configs for this recipe, run ``tune ls`` on the command line. Each config
     has example commands for how to kick-off training.
@@ -582,7 +581,7 @@ def concatenated_forward(
         all_log_probs = rlhf.get_batch_log_probs(
             all_logits,
             concatenated_labels,
-            # see :class:`~torchtune.modules.rlhf.loss.dpo.SimPOLoss`
+            # see :class:`~torchtune.rlhf.loss.dpo.SimPOLoss`
             return_average_logprobs=isinstance(self._loss_fn, SimPOLoss),
         )

diff --git a/recipes/lora_dpo_single_device.py b/recipes/lora_dpo_single_device.py
index 8411d8088d..4cb53b2afb 100644
--- a/recipes/lora_dpo_single_device.py
+++ b/recipes/lora_dpo_single_device.py
@@ -18,10 +18,9 @@
 from torch import nn
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader, DistributedSampler
-from torchtune import config, modules, training, utils
+from torchtune import config, modules, rlhf, training, utils
 from torchtune.data import CROSS_ENTROPY_IGNORE_IDX, padded_collate_dpo
 from torchtune.datasets import ConcatDataset
-from torchtune.modules import rlhf
 from torchtune.modules.peft import (
     disable_adapter,
     get_adapter_params,
@@ -30,9 +29,9 @@
     validate_missing_and_unexpected_for_lora,
     validate_state_dict_for_lora,
 )
-
-from torchtune.modules.rlhf.loss import SimPOLoss
 from torchtune.recipe_interfaces import FTRecipeInterface
+
+from torchtune.rlhf.loss import SimPOLoss
 from tqdm import tqdm

 log = utils.get_logger("DEBUG")
@@ -56,10 +55,10 @@ class LoRADPORecipeSingleDevice(FTRecipeInterface):

     The following losses are supported in this recipe:
-        - :class:`~torchtune.modules.rlhf.loss.DPOLoss`: Direct Preference Optimization (DPO).
-        - :class:`~torchtune.modules.rlhf.loss.RSOPLoss`: Rejection Sampling Optimization (RSO).
-        - :class:`~torchtune.modules.rlhf.loss.IPOLoss`: Identity Preference Optimization (IPO).
-        - :class:`~torchtune.modules.rlhf.loss.SimPOLoss`: Simple Preference Optimization (SimPO).
+        - :class:`~torchtune.rlhf.loss.DPOLoss`: Direct Preference Optimization (DPO).
+        - :class:`~torchtune.rlhf.loss.RSOPLoss`: Rejection Sampling Optimization (RSO).
+        - :class:`~torchtune.rlhf.loss.IPOLoss`: Identity Preference Optimization (IPO).
+        - :class:`~torchtune.rlhf.loss.SimPOLoss`: Simple Preference Optimization (SimPO).

     Assumptions:
         - Checkpoints are ONLY saved at epoch boundaries. In case of failure, work done
@@ -471,7 +470,7 @@ def concatenated_forward(
         all_log_probs = rlhf.get_batch_log_probs(
             all_logits,
             concatenated_labels,
-            # see :class:`~torchtune.modules.rlhf.loss.dpo.SimPOLoss`
+            # see :class:`~torchtune.rlhf.loss.dpo.SimPOLoss`
             return_average_logprobs=isinstance(self._loss_fn, SimPOLoss),
         )

diff --git a/recipes/ppo_full_finetune_single_device.py b/recipes/ppo_full_finetune_single_device.py
index bdd63e8cdc..b9840fc067 100644
--- a/recipes/ppo_full_finetune_single_device.py
+++ b/recipes/ppo_full_finetune_single_device.py
@@ -17,12 +17,11 @@
 from torch import nn
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader, DistributedSampler
-from torchtune import config, modules, training, utils
+from torchtune import config, modules, rlhf, training, utils
 from torchtune.data import padded_collate
 from torchtune.datasets import ConcatDataset
-from torchtune.modules import rlhf
-from torchtune.modules.rlhf import PPOStats, Trajectory
 from torchtune.recipe_interfaces import FTRecipeInterface
+from torchtune.rlhf import PPOStats, Trajectory
 from tqdm import tqdm

@@ -680,7 +679,7 @@ def generate_trajectory(self, input_ids: torch.Tensor) -> Trajectory:
             input_ids (torch.Tensor): tensor of input token IDs with shape [b, seq_length]

         Returns:
-            Trajectory: An instance of :class:`~torchtune.modules.rlhf.Trajectory` comprising
+            Trajectory: An instance of :class:`~torchtune.rlhf.Trajectory` comprising
                 the current trajectory.
         """
         batch_size, context_length = input_ids.shape
@@ -799,7 +798,7 @@ def generate_trajectory_batched(self, input_ids: torch.Tensor) -> Trajectory:
             input_ids (torch.Tensor): tensor of input token IDs with shape [b, seq_length]

         Returns:
-            Trajectory: An instance of :class:`~torchtune.modules.rlhf.Trajectory`, comprising
+            Trajectory: An instance of :class:`~torchtune.rlhf.Trajectory`, comprising
                 the current trajectory.
         """
         trajectories: List[Trajectory] = []
@@ -947,7 +946,7 @@ def _ppo_step(
             context_length (int): input ids sequence length

         Returns:
-            PPOStats: An instance of :class:`~torchtune.modules.rlhf.PPOStats`, a NamedTuple containing:
+            PPOStats: An instance of :class:`~torchtune.rlhf.PPOStats`, a NamedTuple containing:
                - loss (torch.Tensor): The total PPO loss.
                - policy_loss (torch.Tensor): The policy function loss.
                - value_loss (torch.Tensor): The value function loss.
diff --git a/torchtune/modules/rlhf/__init__.py b/torchtune/rlhf/__init__.py
similarity index 100%
rename from torchtune/modules/rlhf/__init__.py
rename to torchtune/rlhf/__init__.py
diff --git a/torchtune/modules/rlhf/_generation.py b/torchtune/rlhf/_generation.py
similarity index 100%
rename from torchtune/modules/rlhf/_generation.py
rename to torchtune/rlhf/_generation.py
diff --git a/torchtune/modules/rlhf/_types.py b/torchtune/rlhf/_types.py
similarity index 100%
rename from torchtune/modules/rlhf/_types.py
rename to torchtune/rlhf/_types.py
diff --git a/torchtune/modules/rlhf/loss/__init__.py b/torchtune/rlhf/loss/__init__.py
similarity index 100%
rename from torchtune/modules/rlhf/loss/__init__.py
rename to torchtune/rlhf/loss/__init__.py
diff --git a/torchtune/modules/rlhf/loss/dpo.py b/torchtune/rlhf/loss/dpo.py
similarity index 99%
rename from torchtune/modules/rlhf/loss/dpo.py
rename to torchtune/rlhf/loss/dpo.py
index d1ca35cd11..d49dbad676 100644
--- a/torchtune/modules/rlhf/loss/dpo.py
+++ b/torchtune/rlhf/loss/dpo.py
@@ -255,7 +255,7 @@ class SimPOLoss(nn.Module):

     SimPO is pretty much identitcal to DPO but uses average logprobs to eliminate the need for a reference model to regularize
     the policy during training. It also uses a target reward margin to guide the policy towards better responses.
-    This is kind of the same intuition as in :class:`~torchtune.modules.rlhf.loss.IPOLoss`, but instead of optimizing against
+    This is kind of the same intuition as in :class:`~torchtune.rlhf.loss.IPOLoss`, but instead of optimizing against
     a margin between the reference policy and policy models, we're optimizing against a margin between the chosen and
     rejected responses.

diff --git a/torchtune/modules/rlhf/loss/ppo.py b/torchtune/rlhf/loss/ppo.py
similarity index 99%
rename from torchtune/modules/rlhf/loss/ppo.py
rename to torchtune/rlhf/loss/ppo.py
index 0cef4a5301..d4770802f7 100644
--- a/torchtune/modules/rlhf/loss/ppo.py
+++ b/torchtune/rlhf/loss/ppo.py
@@ -8,7 +8,7 @@

 import torch
 import torch.nn as nn
-from torchtune.modules import rlhf
+from torchtune import rlhf

 class PPOLoss(nn.Module):
diff --git a/torchtune/modules/rlhf/rewards.py b/torchtune/rlhf/rewards.py
similarity index 98%
rename from torchtune/modules/rlhf/rewards.py
rename to torchtune/rlhf/rewards.py
index dd9f970376..f0e42ca58c 100644
--- a/torchtune/modules/rlhf/rewards.py
+++ b/torchtune/rlhf/rewards.py
@@ -19,7 +19,7 @@ def get_reward_penalty_mask(
     Calculates a mask to penalise scores corresponding to sequences generated during PPO, where True indicates the score
     at the corresponding position should be penalised.
     This function assumes sequences have already been truncated at an EOS, if present, and padded to length,
-    e.g. by :func:`torchtune.modules.rlhf.sequence_processing.truncate_sequence_at_first_stop_token`.
+    e.g. by :func:`torchtune.rlhf.sequence_processing.truncate_sequence_at_first_stop_token`.

     Scores are penalised such that:
     - If ``min_response_length`` is set, scores for sequences with ``length < min_response_length`` are penalised.
diff --git a/torchtune/modules/rlhf/sequence_processing.py b/torchtune/rlhf/sequence_processing.py
similarity index 99%
rename from torchtune/modules/rlhf/sequence_processing.py
rename to torchtune/rlhf/sequence_processing.py
index adbce7cd6c..9844dd001c 100644
--- a/torchtune/modules/rlhf/sequence_processing.py
+++ b/torchtune/rlhf/sequence_processing.py
@@ -8,8 +8,8 @@

 import torch
 import torch.nn.functional as F
+from torchtune import rlhf
 from torchtune.data import CROSS_ENTROPY_IGNORE_IDX
-from torchtune.modules import rlhf

 def truncate_sequence_at_first_stop_token(
diff --git a/torchtune/modules/rlhf/utils/__init__.py b/torchtune/rlhf/utils/__init__.py
similarity index 100%
rename from torchtune/modules/rlhf/utils/__init__.py
rename to torchtune/rlhf/utils/__init__.py
diff --git a/torchtune/modules/rlhf/utils/_convert_weights.py b/torchtune/rlhf/utils/_convert_weights.py
similarity index 100%
rename from torchtune/modules/rlhf/utils/_convert_weights.py
rename to torchtune/rlhf/utils/_convert_weights.py
diff --git a/torchtune/training/checkpointing/_checkpointer.py b/torchtune/training/checkpointing/_checkpointer.py
index fc4a3d2816..ad681db5a5 100644
--- a/torchtune/training/checkpointing/_checkpointer.py
+++ b/torchtune/training/checkpointing/_checkpointer.py
@@ -18,7 +18,7 @@
 from torchtune.models import convert_weights
 from torchtune.models.phi3._convert_weights import phi3_hf_to_tune, phi3_tune_to_hf
 from torchtune.models.qwen2._convert_weights import qwen2_hf_to_tune, qwen2_tune_to_hf
-from torchtune.modules.rlhf.utils import reward_hf_to_tune, reward_tune_to_hf
+from torchtune.rlhf.utils import reward_hf_to_tune, reward_tune_to_hf
 from torchtune.training.checkpointing._utils import (
     get_path,
     ModelType,

Prompt:
The following is the text of a github issue and related comments for this repository:

Move `torchtune.modules.rlhf` to `torchtune.rlhf`
Make sure to update all imports and usages.

Resolve the issue described above. Do not modify tests or user code. The issue can be resolved by modifying the actual library/application code. Make the minimal changes to resolve the issue while adhering to the specification and coding style of the codebase.

diff --git a/tests/torchtune/modules/rlhf/loss/test_dpo_loss.py b/tests/torchtune/modules/rlhf/loss/test_dpo_loss.py
index ef82aad873..e56621f3ed 100644
--- a/tests/torchtune/modules/rlhf/loss/test_dpo_loss.py
+++ b/tests/torchtune/modules/rlhf/loss/test_dpo_loss.py
@@ -6,7 +6,7 @@

 import pytest
 import torch
-from torchtune.modules.rlhf.loss import DPOLoss, IPOLoss, RSOLoss, SimPOLoss
+from torchtune.rlhf.loss import DPOLoss, IPOLoss, RSOLoss, SimPOLoss

 @pytest.fixture(autouse=True)
diff --git a/tests/torchtune/modules/rlhf/loss/test_ppo_loss.py b/tests/torchtune/modules/rlhf/loss/test_ppo_loss.py
index 97a45dace2..01a2388770 100644
--- a/tests/torchtune/modules/rlhf/loss/test_ppo_loss.py
+++ b/tests/torchtune/modules/rlhf/loss/test_ppo_loss.py
@@ -6,7 +6,7 @@

 import pytest
 import torch
-from torchtune.modules.rlhf.loss import PPOLoss
+from torchtune.rlhf.loss import PPOLoss

 @pytest.fixture(autouse=True)
diff --git a/tests/torchtune/modules/rlhf/test_generation.py b/tests/torchtune/modules/rlhf/test_generation.py
index bfe8237c0c..511ecfcdc4 100644
--- a/tests/torchtune/modules/rlhf/test_generation.py
+++ b/tests/torchtune/modules/rlhf/test_generation.py
@@ -8,9 +8,9 @@

 import torch
 from tests.test_utils import fixed_init_model
+from torchtune import rlhf
 from torchtune.generation._generation import sample
 from torchtune.models.llama2 import llama2
-from torchtune.modules import rlhf

 class TestGenerateNextTokenWithLogits:
@@ -61,7 +61,7 @@ def test_generate_next_token_with_logits(self, generation_model):

 class TestGenerate:
     """
-    Test class for text generation functionality in :func:`~torchtune.modules.rlhf.generate`.
+    Test class for text generation functionality in :func:`~torchtune.rlhf.generate`.
     See `torchtune.tests.utils.test_generation` for context.
     """

diff --git a/tests/torchtune/modules/rlhf/test_rewards.py b/tests/torchtune/modules/rlhf/test_rewards.py
index 0e8ec998fa..4284d1d63d 100644
--- a/tests/torchtune/modules/rlhf/test_rewards.py
+++ b/tests/torchtune/modules/rlhf/test_rewards.py
@@ -5,7 +5,7 @@
 # LICENSE file in the root directory of this source tree.

 import torch
-from torchtune.modules import rlhf
+from torchtune import rlhf

 class TestGetRewards:
@@ -182,7 +182,7 @@ def test_estimate_advantages_with_whitening(self):
             ]
         )

-        # see `torchtune.modules.rlhf.estimate_advantages`
+        # see `torchtune.rlhf.estimate_advantages`
         expected_advantages = returns - values
         expected_whitened_advantages = rlhf.whiten(expected_advantages, shift_mean=True)
         advantages, _ = rlhf.estimate_advantages(values, rewards, gamma, lmbda)
@@ -209,7 +209,7 @@ def test_estimate_advantages_with_masks(self):
             ]
         )

-        # see `torchtune.modules.rlhf.estimate_advantages`
+        # see `torchtune.rlhf.estimate_advantages`
         expected_advantages = returns - values
         expected_advantages = rlhf.whiten(expected_advantages, mask=masks)
         expected_advantages[..., -1] = 0.0
diff --git a/tests/torchtune/modules/rlhf/test_sequence_processing.py b/tests/torchtune/modules/rlhf/test_sequence_processing.py
index 43accdf80c..ae53e6494f 100644
--- a/tests/torchtune/modules/rlhf/test_sequence_processing.py
+++ b/tests/torchtune/modules/rlhf/test_sequence_processing.py
@@ -5,7 +5,7 @@
 # LICENSE file in the root directory of this source tree.

 import torch
-from torchtune.modules import rlhf
+from torchtune import rlhf

 class TestTruncateSequenceAtFirstStopToken:


diff --git a/torchtune/data/_messages.py b/torchtune/data/_messages.py
index a6b356b0ca..bbd3ae5981 100644
--- a/torchtune/data/_messages.py
+++ b/torchtune/data/_messages.py
@@ -170,6 +170,7 @@ class InputOutputToMessages(Transform):
     Raises:
         ValueError: If ``column_map`` is provided and ``input`` not in ``column_map``, or
             ``output`` not in ``column_map``.
+        ValueError: If ``image_dir`` is provided but ``image`` not in ``column_map``.
     """

     def __init__(
@@ -196,6 +197,14 @@ def __init__(
         else:
             self.column_map = {"input": "input", "output": "output", "image": "image"}

+        # Ensure that if a user seems to want to construct a multimodal transform, they provide
+        # a proper column_mapping
+        if "image" not in self.column_map.keys() and image_dir is not None:
+            raise ValueError(
+                f"image_dir is specified as {image_dir} but 'image' is not in column_map. "
+                "Please specify an 'image' key in column_map."
+            )
+
         self.image_dir = image_dir

     def __call__(self, sample: Mapping[str, Any]) -> Mapping[str, Any]:
@@ -206,8 +215,13 @@ def __call__(self, sample: Mapping[str, Any]) -> Mapping[str, Any]:
         if is_multimodal:
             image_path = sample[self.column_map["image"]]
             if isinstance(image_path, str):
+                # Convert image_path to Path obj
+                image_path = Path(image_path)
+
+                # If image_dir is not None, prepend image_dir to image_path
                 if self.image_dir is not None:
                     image_path = self.image_dir / image_path
+
                 # Load if not loaded
                 pil_image = load_image(image_path)
             else:

Add validation to raise a ValueError with message "image_dir is specified as {image_dir} but 'image' is not in column_map. Please specify an 'image' key in column_map" when image_dir is provided but 'image' is not in column_map. Also, convert the image path to a Path object before combining with image_dir.

Prompt:
The following is the text of a github issue and related comments for this repository:

TypeError: unsupported operand type(s) for /: 'str' and 'str' in torchtune/data/_messages.py
Hello, I am trying to run torchtune on a vqa multimodal dataset. When trying to run the torchtune script I ran into this error:
"
File "/home/usr/ml-project-2-a-3/python_venv/lib64/python3.9/site-packages/torchtune/data/_messages.py", line 210, in __call__
    image_path = self.image_dir / image_path
TypeError: unsupported operand type(s) for /: 'str' and 'str'
"

We should convert image_dir to a Path object here. We also must ensure that if a user seems to want to construct a multimodal transform, they provide
a proper column_mapping. If not, we should raise a ValueError with the message "image_dir is specified as {image_dir} but 'image' is not in column_map. Please specify an 'image' key in column_map."

Resolve the issue described above. Do not modify tests or user code. The issue can be resolved by modifying the actual library/application code. Make the minimal changes to resolve the issue while adhering to the specification and coding style of the codebase.

diff --git a/tests/torchtune/data/test_messages.py b/tests/torchtune/data/test_messages.py
index a46cfd9349..86b7d7319f 100644
--- a/tests/torchtune/data/test_messages.py
+++ b/tests/torchtune/data/test_messages.py
@@ -9,6 +9,7 @@
 import pytest

 from PIL import Image
+from tests.common import ASSETS
 from tests.test_utils import (
     assert_dialogue_equal,
     CHAT_SAMPLE,
@@ -24,6 +25,8 @@
     validate_messages,
 )

+PYTORCH_RGB_IMAGE_AS_PIL = Image.open(ASSETS / "rgb_pytorch.png")
+

 class TestMessage:
     @pytest.fixture
@@ -106,6 +109,60 @@ def sample(self):
             "maybe_output": "hello world",
         }

+    @pytest.mark.parametrize(
+        "input_image, expected_image",
+        [
+            ("rgb_pytorch.png", PYTORCH_RGB_IMAGE_AS_PIL),
+            (ASSETS / "rgb_pytorch.png", PYTORCH_RGB_IMAGE_AS_PIL),
+            (PYTORCH_RGB_IMAGE_AS_PIL, PYTORCH_RGB_IMAGE_AS_PIL),
+        ],
+    )
+    def test_call_with_image(self, sample, input_image, expected_image):
+        # Add the image to the sample
+        sample["image"] = input_image
+
+        # Create the transform
+        transform = InputOutputToMessages(
+            column_map={
+                "input": "maybe_input",
+                "output": "maybe_output",
+                "image": "image",
+            },
+            # Need to test if the image_dir is properly joined w/ image
+            image_dir=ASSETS if isinstance(input_image, str) else None,
+        )
+        actual = transform(sample)
+        expected = [
+            Message(
+                role="user",
+                content=[
+                    {"type": "text", "content": "hello world"},
+                    {"type": "image", "content": expected_image},
+                ],
+                masked=True,
+                eot=True,
+            ),
+            Message(role="assistant", content="hello world", masked=False, eot=True),
+        ]
+        assert_dialogue_equal(actual["messages"], expected)
+
+    def test_call_with_image_fails_when_bad_image_inputs_are_passed(self, sample):
+        # Construct a bad column_map without an 'image' key
+        column_map = {
+            "input": "maybe_input",
+            "output": "maybe_output",
+        }
+
+        # Create a transform that expects an image column
+        with pytest.raises(
+            ValueError,
+            match="Please specify an 'image' key in column_map",
+        ):
+            transform = InputOutputToMessages(
+                column_map=column_map,
+                image_dir=ASSETS,
+            )
+
     def test_call(self, sample):
         transform = InputOutputToMessages(
             column_map={"input": "maybe_input", "output": "maybe_output"}


diff --git a/recipes/eleuther_eval.py b/recipes/eleuther_eval.py
index ce07497899..ad6ba41e74 100644
--- a/recipes/eleuther_eval.py
+++ b/recipes/eleuther_eval.py
@@ -28,6 +28,7 @@
 from torchtune.modules.tokenizers import ModelTokenizer
 from torchtune.modules.transforms import Transform
 from torchtune.recipe_interfaces import EvalRecipeInterface
+from torchtune.training import FullModelTorchTuneCheckpointer

 try:
     import lm_eval
@@ -475,13 +476,6 @@ def setup(self, cfg: DictConfig) -> None:

         # Load checkpoint
         checkpointer = config.instantiate(cfg.checkpointer)
-        if quantization_mode is None:
-            ckpt_dict = checkpointer.load_checkpoint()
-        else:
-            # weights_only needs to be False when loading a quantized model
-            # currently loading a quantized model is only supported with the
-            # FullModelTorchTuneCheckpointer
-            ckpt_dict = checkpointer.load_checkpoint(weights_only=False)

         # Initialize model
         with training.set_default_dtype(self.dtype), self.device:
@@ -489,14 +483,32 @@ def setup(self, cfg: DictConfig) -> None:

         # Quantize model if requested
         if quantization_mode is not None:
+            if not isinstance(checkpointer, FullModelTorchTuneCheckpointer):
+                raise ValueError(
+                    "Quantization is only supported for models quantized and saved with the "
+                    "FullModelTorchTuneCheckpointer - please ensure you have quantized your "
+                    "model and are using the quantized weights!"
+                )
+            if "qat" in quantization_mode:
+                raise ValueError(
+                    "You have specified a quantizer with 'QAT' - "
+                    "QAT quantizers should only be used during quantization aware training "
+                    "and when quantizing models. Please use the corresponding post-training "
+                    "quantizer e.g. Int8DynActInt4WeightQuantizer for Int8DynActInt4WeightQATQuantizer."
+                )
             model = quantizer.quantize(model)
             model = model.to(device=self.device, dtype=self.dtype)
-            for k, v in model_state_dict.items():
-                model_state_dict[k] = v.to(self._device)
-            model.load_state_dict(model_state_dict, assign=True)
+            ckpt_dict = checkpointer.load_checkpoint(weights_only=False)[
+                training.MODEL_KEY
+            ]
+            for k, v in ckpt_dict.items():
+                ckpt_dict[k] = v.to(self.device)
+            model.load_state_dict(ckpt_dict, assign=True)
+        else:
+            ckpt_dict = checkpointer.load_checkpoint()[training.MODEL_KEY]
+            model.load_state_dict(ckpt_dict)

         # Load model weights into initialized model
-        model.load_state_dict(ckpt_dict[training.MODEL_KEY])
         self.logger.info(f"Model is initialized with precision {self.dtype}.")

         # Put model in eval mode.

Prompt:
The following is the text of a github issue and related comments for this repository:

model eval error: NameError: name 'model_state_dict' is not defined
i use this commend to run evaluation
```
tune run eleuther_eval --config eleuther_evaluation \
>     tasks="[hellaswag, wikitext]" \
>     model._component_=torchtune.models.llama3.llama3_8b \
>     quantizer._component_=torchtune.training.quantization.Int8DynActInt4WeightQuantizer\
>     quantizer.groupsize=128 \
>     checkpointer._component_=torchtune.training.FullModelTorchTuneCheckpointer \
>     checkpointer.checkpoint_dir="/QAT/output/llama3-8B" \
>     checkpointer.output_dir="/QAT/output/llama3-8B" \
>     checkpointer.checkpoint_files=[meta_model_2-8da4w.pt] \
>     checkpointer.model_type=LLAMA3 \
>     tokenizer._component_=torchtune.models.llama3.llama3_tokenizer \
>     tokenizer.path=/QAT/Meta-Llama-3-8B/original/tokenizer.model
```
But i get this.
```
2024-10-09:08:30:57,790 INFO     [_logging.py:101] Running EleutherEvalRecipe with resolved config:

batch_size: 8
checkpointer:
  _component_: torchtune.training.FullModelTorchTuneCheckpointer
  checkpoint_dir: /QAT/output/llama3-8B
  checkpoint_files:
  - meta_model_2-8da4w.pt
  model_type: LLAMA3
  output_dir: /QAT/output/llama3-8B
device: cuda
dtype: bf16
enable_kv_cache: true
limit: null
max_seq_length: 4096
model:
  _component_: torchtune.models.llama3.llama3_8b
quantizer:
  _component_: torchtune.training.quantization.Int8DynActInt4WeightQuantizer
  groupsize: 128
seed: 1234
tasks:
- hellaswag
- wikitext
tokenizer:
  _component_: torchtune.models.llama3.llama3_tokenizer
  max_seq_len: null
  path: /QAT/Meta-Llama-3-8B/original/tokenizer.model

Traceback (most recent call last):
  File "/usr/local/bin/tune", line 8, in <module>
    sys.exit(main())
  File "/usr/local/lib/python3.10/dist-packages/torchtune/_cli/tune.py", line 49, in main
    parser.run(args)
  File "/usr/local/lib/python3.10/dist-packages/torchtune/_cli/tune.py", line 43, in run
    args.func(args)
  File "/usr/local/lib/python3.10/dist-packages/torchtune/_cli/run.py", line 196, in _run_cmd
    self._run_single_device(args, is_builtin=is_builtin)
  File "/usr/local/lib/python3.10/dist-packages/torchtune/_cli/run.py", line 102, in _run_single_device
    runpy.run_path(str(args.recipe), run_name="__main__")
  File "/usr/lib/python3.10/runpy.py", line 289, in run_path
    return _run_module_code(code, init_globals, run_name,
  File "/usr/lib/python3.10/runpy.py", line 96, in _run_module_code
    _run_code(code, mod_globals, init_globals,
  File "/usr/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/usr/local/lib/python3.10/dist-packages/recipes/eleuther_eval.py", line 576, in <module>
    sys.exit(recipe_main())
  File "/usr/local/lib/python3.10/dist-packages/torchtune/config/_parse.py", line 99, in wrapper
    sys.exit(recipe_main(conf))
  File "/usr/local/lib/python3.10/dist-packages/recipes/eleuther_eval.py", line 571, in recipe_main
    recipe.setup(cfg=cfg)
  File "/usr/local/lib/python3.10/dist-packages/recipes/eleuther_eval.py", line 494, in setup
    for k, v in model_state_dict.items():
NameError: name 'model_state_dict' is not defined
```
i read the code https://github.com/pytorch/torchtune/blob/main/recipes/eleuther_eval.py, i can not find where model_state_dict is defined.  A bug ???
I have used this config file, but get same error
```
model:
  _component_: torchtune.models.llama3.llama3_8b

checkpointer:
  _component_: torchtune.training.FullModelTorchTuneCheckpointer
  checkpoint_dir: /QAT/output/llama3-8B/
  checkpoint_files: [
    meta_model_2-8da4w.pt
  ]
  output_dir: /QAT/output/llama3-8B/
  model_type: LLAMA3

# Tokenizer
tokenizer:
  _component_: torchtune.models.llama3.llama3_tokenizer
  path: /QAT/Meta-Llama-3-8B/original/tokenizer.model
  max_seq_len: null

# Environment
device: cuda
dtype: bf16
seed: 42 # It is not recommended to change this seed, b/c it matches EleutherAI's default seed

# EleutherAI specific eval args
tasks: ["hellaswag"]
limit: null
max_seq_length: 8192
batch_size: 8

# Quantization specific args
quantizer:
  _component_: torchtune.training.quantization.Int8DynActInt4WeightQuantizer
  groupsize: 256
```
Anyone can helps, thanks very much !!!

As part of fixing this issue, you should add checks that
 - if quantization_mode is not none, the checkpointer must be a FullModelTorchTuneCheckpointer
 (raise a ValueError with the message "Quantization is only supported for models quantized and saved with the FullModelTorchTuneCheckpointer - please ensure you have quantized your model and are using the quantized weights!")
 - if quantization_mode is not none, the quantization_mode must not contain "qat"
 (raise a ValueError with the message "You have specified a quantizer with 'QAT' - QAT quantizers should only be used during quantization aware training and when quantizing models. Please use the corresponding post-training quantizer e.g. Int8DynActInt4WeightQuantizer for Int8DynActInt4WeightQATQuantizer.")

Resolve the issue described above. Do not modify tests or user code. The issue can be resolved by modifying the actual library/application code. Make the minimal changes to resolve the issue while adhering to the specification and coding style of the codebase.

diff --git a/tests/recipes/test_eleuther_eval.py b/tests/recipes/test_eleuther_eval.py
index 32eaee4b1b..f09daf2309 100644
--- a/tests/recipes/test_eleuther_eval.py
+++ b/tests/recipes/test_eleuther_eval.py
@@ -14,7 +14,7 @@
 import pytest

 from tests.common import TUNE_PATH
-from tests.recipes.utils import llama2_test_config
+from tests.recipes.utils import llama2_test_config, write_hf_ckpt_config
 from tests.test_utils import CKPT_MODEL_PATHS

@@ -126,6 +126,80 @@ def test_eval_recipe_errors_without_lm_eval(self, capsys, monkeypatch, tmpdir):
             in printed_err
         )

+    @pytest.mark.integration_test
+    def test_eval_recipe_errors_with_quantization_hf_checkpointer(
+        self, capsys, monkeypatch, tmpdir
+    ):
+        ckpt = "llama2_hf"
+        ckpt_path = Path(CKPT_MODEL_PATHS[ckpt])
+        ckpt_dir = ckpt_path.parent
+
+        # Config file needed for model conversion.
+        write_hf_ckpt_config(ckpt_dir)
+
+        cmd = f"""
+        tune run eleuther_eval \
+            --config eleuther_evaluation \
+            output_dir={tmpdir} \
+            checkpointer=torchtune.training.FullModelHFCheckpointer \
+            checkpointer.checkpoint_dir='{ckpt_dir}' \
+            checkpointer.checkpoint_files=[{ckpt_path}]\
+            checkpointer.output_dir={tmpdir} \
+            checkpointer.model_type=LLAMA2 \
+            tokenizer.path=/tmp/test-artifacts/tokenizer.model \
+            tokenizer.prompt_template=null \
+            limit=1 \
+            dtype=fp32 \
+            device=cpu \
+            quantizer._component_=torchtune.training.quantization.Int8DynActInt4WeightQuantizer \
+            quantizer.groupsize=256 \
+        """.split()
+
+        model_config = llama2_test_config()
+        cmd = cmd + model_config
+
+        monkeypatch.setattr(sys, "argv", cmd)
+        with pytest.raises(
+            ValueError,
+            match="Quantization is only supported for models quantized and saved with the "
+            "FullModelTorchTuneCheckpointer",
+        ):
+            runpy.run_path(TUNE_PATH, run_name="__main__")
+
+    @pytest.mark.integration_test
+    def test_eval_recipe_errors_with_qat_quantizer(self, capsys, monkeypatch, tmpdir):
+        ckpt = "llama2_tune"
+        ckpt_path = Path(CKPT_MODEL_PATHS[ckpt])
+        ckpt_dir = ckpt_path.parent
+
+        cmd = f"""
+        tune run eleuther_eval \
+            --config eleuther_evaluation \
+            output_dir={tmpdir} \
+            checkpointer=torchtune.training.FullModelTorchTuneCheckpointer \
+            checkpointer.checkpoint_dir='{ckpt_dir}' \
+            checkpointer.checkpoint_files=[{ckpt_path}]\
+            checkpointer.output_dir={tmpdir} \
+            checkpointer.model_type=LLAMA2 \
+            tokenizer.path=/tmp/test-artifacts/tokenizer.model \
+            tokenizer.prompt_template=null \
+            limit=1 \
+            dtype=fp32 \
+            device=cpu \
+            quantizer._component_=torchtune.training.quantization.Int8DynActInt4WeightQATQuantizer \
+            quantizer.groupsize=32\
+        """.split()
+
+        model_config = llama2_test_config()
+        cmd = cmd + model_config
+
+        monkeypatch.setattr(sys, "argv", cmd)
+        with pytest.raises(
+            ValueError,
+            match="QAT quantizers should only be used during quantization aware training",
+        ):
+            runpy.run_path(TUNE_PATH, run_name="__main__")
+
     @pytest.mark.integration_test
     def test_eval_recipe_errors_with_generate_until_and_mc_tasks(
         self, caplog, capsys, monkeypatch, tmpdir