Added MessagesDataloader so we can just use messages in our datasets rather than tokenized inputs

README.md

@@ -46,7 +46,7 @@ Here is an end-to-end workflow of performing data preparation for training along
 
 ### Data preparation
 
-Below is the set of commands to run to prepare datasets into the appropriate Mosaic Data Shard (MDS) format, which is a pre-tokenized version of the data, that we will use for training.
+Below is the set of commands to run to prepare datasets into the appropriate Mosaic Data Shard (MDS) format, which is either a pre-tokenized version of the data or the raw chat-message data, that we will use for training.
 
 Below is the command to prepare preference data -- which can be used for reward model or offline RL (e.g. DPO) training:
 
@@ -60,37 +60,44 @@ python data/unified_tokenize_dataset.py --dataset_name allenai/ultrafeedback_bin
 --split train_prefs
 ```
 
-Below is the command to prepare prompt data -- which can be used for online RL (e.g. PPO) training:
+Below is the command to prepare single-turn message data -- which can be used for online RL (e.g. PPO) training:
 
 <!--pytest.mark.skip-->
 ```bash
 cd scripts
 python data/unified_tokenize_dataset.py --dataset_name allenai/ultrafeedback_binarized_cleaned \
---local_dir prompt_data \
---dataset_type single_prompt \
+--local_dir single_message_data \
+--dataset_type single_message \
 --tokenizer_name meta-llama/Llama-3.1-8B-Instruct \
 --split train_prefs
 ```
 
-To further enable online RL with [verifiable rewards](https://arxiv.org/abs/2411.15124) you can use the following command:
+To further enable online RL with [verifiable rewards](https://arxiv.org/abs/2411.15124) you can use the following command to prepare the chat-message data and their corresponding verifiable answers:
 
 <!--pytest.mark.skip-->
 ```bash
 cd scripts
-python data/unified_tokenize_dataset.py --dataset_name <hf_dataset_name> \
+python data/messages_dataset_to_mds.py --dataset_path <hf_dataset_name> \
 --local_dir verifiable_data \
---dataset_type verifiable_answers \
---tokenizer_name meta-llama/Llama-3.1-8B-Instruct \
 --split train \
 ```
 
-We currently support the following two HuggingFace datasets for verifiable rewards:
+For RLVR, We currently support the following two HuggingFace datasets for verifiable rewards:
 
 - GMS8k: `openai/gsm8k`
 - MATH: `DigitalLearningGmbH/MATH-lighteval`
 
 The data preparation scripts also supports additional arguments for specifying the subset of the HuggingFace dataset `--subset <str>` and max sequence length `--max_length <int>`
 
+For custom datasets, you can create a custom preprocessing function in the `compose-rl/scripts/data/messages_preprocessing_utils.py` file, or you can preprocess your own dataset directly, save it locally as a .jsonl file, and then use the following command to convert it to the MDS format:
+<!--pytest.mark.skip-->
+```bash
+cd scripts
+python data/messages_dataset_to_mds.py --dataset_path <path_to_jsonl_file> \
+--local_dir custom_dataset \
+```
+
+
 ### Model training
 
 Below are the scripts to launch training runs assuming you ran the data preparation scripts above. Additionally, these scripts assume that we are in the root directory where Compose RL and LLM Foundry were cloned. This is because we utilize [LLM Foundry's Registry System](https://github.com/mosaicml/llm-foundry/?tab=readme-ov-file#registry) in order to take advantage of existing features in LLM Foundry.
@@ -129,7 +136,7 @@ Below is the command to run Online PPO training:
 ```bash
 composer llm-foundry/scripts/train/train.py \
 compose-rl/yamls/local_ppo.yaml \
-train_loader.dataset.local=/compose-rl/scripts/prompt_data/ \
+train_loader.dataset.local=/compose-rl/scripts/single_message_data/ \
 train_loader.dataset.split=train_prefs
 ```
 

compose_rl/algorithms/online/callback.py

@@ -649,7 +649,12 @@ def _get_next_iter_prompts(self):
                 # Explode the batch into multiple batches for each generation
                 for _ in range(self.generations_per_prompt):
                     # For keys that do not require additional processing
-                    if key in ['prompt_len', 'verified_answer', 'prompt_id']:
+                    if key in [
+                        'prompt_len',
+                        'verified_answer',
+                        'prompt_id',
+                        'messages',
+                    ]:
                         curr_values.append(batch[key])
                         continue
 
@@ -680,7 +685,6 @@ def _get_next_iter_prompts(self):
                 if key == 'verified_answer':
                     ret_batch[key] = list(flatten(curr_values))
                 else:
-                    # this is an edge case that we will not hit currently, but just handling it as needed
                     ret_batch[key] = curr_values
 
         return ret_batch

compose_rl/algorithms/online/generation_utils/vllm_utils.py

@@ -239,7 +239,7 @@ def create_vllm_engines(
         log.info(f'vllm: {num_gpus=}, {num_engines=}')
 
         vllm_engines.append(
-            LLMRayActor.options(
+            LLMRayActor.options(  # type: ignore
                 num_cpus=num_gpus,
                 num_gpus=num_gpus,
                 scheduling_strategy=scheduling_strategy,

compose_rl/data/__init__.py

@@ -7,9 +7,11 @@
 )
 from compose_rl.data.dataloader import (
     build_finegrained_preference_dataloader,
+    build_messages_dataloader,
     build_pairwise_preference_dataloader,
     build_prompt_dataloader,
 )
+from compose_rl.data.messages_data import messages_dataset_collate_fn
 from compose_rl.data.preference_data import (
     finegrained_preference_dataset_collate_fn,
     pairwise_preference_dataset_collate_fn,
@@ -19,10 +21,12 @@
 __all__ = [
     'build_pairwise_preference_dataloader',
     'build_finegrained_preference_dataloader',
+    'build_messages_dataloader',
     'build_prompt_dataloader',
     'DummyDataset',
     'finegrained_preference_dataset_collate_fn',
     'MinibatchRolloutBuffer',
     'pairwise_preference_dataset_collate_fn',
     'prompt_dataset_collate_fn',
+    'messages_dataset_collate_fn',
 ]

compose_rl/data/dataloader.py

@@ -10,6 +10,10 @@
 from torch.utils.data import DataLoader
 from transformers import PreTrainedTokenizer
 
+from compose_rl.data.messages_data import (
+    MessagesStreamingDataset,
+    messages_dataset_collate_fn,
+)
 from compose_rl.data.preference_data import (
     FinegrainedPreferenceStreamingDataset,
     PairwisePreferenceStreamingDataset,
@@ -25,6 +29,7 @@
     'build_finegrained_preference_dataloader',
     'build_pairwise_preference_dataloader',
     'build_prompt_dataloader',
+    'build_messages_dataloader',
 ]
 
 
@@ -71,10 +76,15 @@ def build_preference_dataloader(
         streams = None
         if streams_dict is not None:
             streams = [Stream(**stream) for stream in streams_dict.values()]
+        if issubclass(
+            dataset_cls,
+            MessagesStreamingDataset,
+        ) and 'tokenizer' not in dataset_cfg:
+            dataset_cfg['tokenizer'] = tokenizer
 
         streaming_dataset = dataset_cls(
-            streams=streams,
-            batch_size=device_batch_size,
+            streams=streams,  # type: ignore
+            batch_size=device_batch_size,  # type: ignore
             **dataset_cfg,
         )
 
@@ -111,3 +121,8 @@ def build_preference_dataloader(
     PromptStreamingDataset,
     prompt_dataset_collate_fn,
 )
+
+build_messages_dataloader = generate_dataloader_builder(
+    MessagesStreamingDataset,
+    messages_dataset_collate_fn,
+)

compose_rl/data/messages_data.py

@@ -0,0 +1,174 @@
+# Copyright 2025 MosaicML ComposeRL authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Build a prompt dataset and dataloader for training."""
+
+import logging
+from typing import Any, TypeAlias
+
+import torch
+from streaming import StreamingDataset
+from transformers import (
+    AutoTokenizer,
+    DataCollatorForLanguageModeling,
+    PreTrainedTokenizerBase,
+)
+
+import compose_rl.utils as utils
+
+log = logging.getLogger(__name__)
+
+Messages: TypeAlias = list[dict[str, str]]
+
+
+def messages_dataset_collate_fn(
+    tokenizer: PreTrainedTokenizerBase,
+    max_seq_len: int,
+    batch: list[dict[str, Any]],
+) -> dict[str, Any]:
+    """Collator for messages data.
+
+    Args:
+        batch (List[Dict[str, Any]]): A list of data samples to collate.
+        tokenizer (PreTrainedTokenizer): The model's tokenizer.
+        max_seq_len (int): The maximum sequence length of the model.
+    """
+    if tokenizer.pad_token_id is None:
+        raise ValueError('Tokenizer must have a PAD token.')
+
+    ref_collate_fn = DataCollatorForLanguageModeling(
+        tokenizer=tokenizer,
+        mlm=False,
+        mlm_probability=0.0,
+    )
+
+    keys = batch[0].keys()
+    collated_batch: dict[str, Any] = {}
+    for key in keys:
+        cur_values = [item[key] for item in batch]
+        if key in ['prompt_len']:
+            collated_batch[key] = torch.stack(cur_values).squeeze(dim=1)
+        elif key == 'prompt_id':
+            collated_batch[key] = torch.tensor(cur_values)
+        elif key in ['verified_answer']:
+            collated_batch[key] = list(  # pyright: ignore[reportGeneralTypeIssues]
+                utils.flatten(cur_values),
+            )
+        elif key == 'messages':
+            collated_batch[key] = cur_values
+        elif key == 'prompt':
+            collated_batch[key] = ref_collate_fn(cur_values)['input_ids']
+        else:
+            raise ValueError(f'Invalid key: {key}')
+
+    collated_batch['prompt_attention_mask'] = torch.logical_not(
+        torch.eq(collated_batch['prompt'],
+                 tokenizer.pad_token_id),  # type: ignore
+    )
+
+    return collated_batch
+
+
+class MessagesStreamingDataset(StreamingDataset):
+    """Dataloader for streaming in messages and converting to prompts."""
+
+    def __init__(
+        self,
+        max_gen_len: int,
+        max_seq_len: int,
+        tokenizer: str | PreTrainedTokenizerBase,
+        **kwargs: dict[str, Any],
+    ):
+        self.max_gen_len = max_gen_len
+        self.max_seq_len = max_seq_len
+        if isinstance(tokenizer, str):
+            self.tokenizer = AutoTokenizer.from_pretrained(tokenizer)
+        else:
+            self.tokenizer = tokenizer
+        super().__init__(**kwargs)
+
+    def _validate_messages(self, messages: Messages) -> bool:
+        """Validates the message.
+
+        A valid message is a list of dictionaries with
+        a 'role' key and a 'content' key.
+
+        Args:
+            messages (Messages): The messages to validate.
+
+        Returns:
+            bool: True if the messages are valid, False otherwise.
+        """
+        if not isinstance(messages, list):
+            return False
+        for message in messages:
+            if not isinstance(message, dict):
+                return False
+            if 'role' not in message:
+                return False
+            if 'content' not in message:
+                return False
+            if not isinstance(message['content'], str):
+                return False
+        return True
+
+    def _tokenize_messages(self, messages: Messages) -> torch.Tensor:
+        if not self._validate_messages(messages):
+            raise ValueError(f'Invalid messages received. Got: {messages=}')
+        return torch.tensor(
+            self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=True,
+                add_generation_prompt=True,
+            ),
+            dtype=torch.int64,
+        )
+
+    # How to process a sample
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        """Gets the messages + (optionally) metadata at the given index.
+
+        Args:
+            idx (int): the index where we fetch the data in the StreamingDataset.
+        """
+        sample = super().__getitem__(idx)
+        messages = sample['messages']
+        metadata = sample.get('metadata', {})
+        prompt: torch.Tensor = self._tokenize_messages(messages)
+
+        # TODO (bcui): Maybe add in an option to truncate a prompt by a given length?
+        if len(prompt) + self.max_gen_len > self.max_seq_len:
+            truncate_len = len(prompt) + self.max_gen_len - self.max_seq_len
+            log.info(f'Truncating prompt by: {truncate_len}')
+            prompt = prompt[:-truncate_len]
+
+        prompt_len = torch.Tensor([len(prompt)]).to(dtype=torch.int64)
+        #  Create the return dictionary
+        item_dict = {
+            'prompt_id': idx,
+            'prompt': prompt,
+            'prompt_len': prompt_len,
+            'messages': messages,
+            # 'metadata': metadata,  # removing metadata for now
+        }
+
+        # extract the verified answer, if there is one
+        verified_answer = metadata.get('verified_answer', None)
+        if verified_answer is None:
+            # for backwards compatibility, we also check the sample directly for the verified answer
+            verified_answer = sample.get('verified_answer', None)
+        if verified_answer:
+            if isinstance(verified_answer, str):
+                _answer = verified_answer
+            else:
+                try:
+                    _answer = verified_answer.decode('utf-8', errors='strict')
+                except UnicodeDecodeError as e:
+                    log.error(
+                        f'Failed to decode verifed_answer with error: {e}',
+                    )
+                    _answer = ''
+
+            item_dict['verified_answer'] = _answer  # type: ignore
+
+        return item_dict

pyproject.toml

@@ -59,6 +59,7 @@ hf_ppo_lm = "compose_rl.algorithms.online:ComposerHFPolicyLM"
 pairwise_preference = "compose_rl.data:build_pairwise_preference_dataloader"
 finegrained_preference = "compose_rl.data:build_finegrained_preference_dataloader"
 prompt = "compose_rl.data:build_prompt_dataloader"
+messages = "compose_rl.data:build_messages_dataloader"
 
 [project.entry-points."llmfoundry_callbacks_with_config"]
 offline_rl = "compose_rl.algorithms.offline:ReferencePolicyCallback"