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

[SeanKski]

This adds the MessagesStreamingDataset, which allows us to use chat-ml messages as our main data interface. This can essentially be thought of as the same as the PromptStreamingDataset but with the tokenizing of the prompts happening on the fly, rather than requiring the samples to be pre-tokenized.

What this allows us to do:

1. Re-use the same dataset across models, rather than having to create a different pre-tokenized MDS dataset for different tokenizers
2. Have the raw messages in our batches so we don't have to untokenize + unchat-template to use chat interfaces such as bpt
3. More easily support tool-use and multi-turn

I've updated the README and example yamls to use messages rather than tokenized-prompts

Couple things to note:

• This is part 1 of a closely-linked PR that moves to using vllm.chat rather than vllm.generate, but I'm keeping the PRs separate for house-keeping reasons
• I couldn't come up with a clean way of caching the tokenizations for samples, so currently, [dataset.__getitem__(0) for _ in range(10)] will freshly tokenize the same sample 10 times. It's a minor thing, but if anyone has ideas on how to properly do that caching that'd be dope!

MCLI runs:

• the control run: (recreation of @jdchang1's run from his reorg PR, runs from compose-rl main)mcli logs grpo-reorg-tVrNI0
• a run with PromptDataset: mcli logs grpo-reorg-prompt-Az2nwg
• a run with MessagesDataset: mcli logs grpo-reorg-messages-HMdrix (note: since this uses @abaheti95's open-r1 filtered dataset which he only has tokenized versions of, I had to untokenize + unchat-template these to convert them to messages, so there might be some differences caused from that)

[abaheti95]

Discussed offline with/ @SeanKski. There is a limitation in the current unified_tokenize_dataset.py, where it only supports converting single prompts to messages. We would like to have an additional preprocessing script that directly ingests a JSONL file for data preprocessing, where each row contains the following keys: messages and verified_answer. This can be directly converted into MDS that the dataloader can handle.

[bcui-db]

Does this flatten this correctly?

curr_values is just a raw list of messages, so I think this would be unpacking all of the messages for multi-turn. Let me know if I'm interpreting the code incorrectly.

[SeanKski]

hmmm good point! I originally thought this was [num_batches_per_update, batch_size, num_turns] (especially since there's concatting and flattening happening for the other keys), but you're right this is already in [num_batches_per_update * batch_size, num_turns]. So I just removed this part and now we should be good!

[bcui-db]

Is this file supposed to be a part of the PR? Or is it still okay to be using unified_tokenize_dataset.py?

[SeanKski]

Discussed offline with/ @SeanKski. There is a limitation in the current unified_tokenize_dataset.py, where it only supports converting single prompts to messages. We would like to have an additional preprocessing script that directly ingests a JSONL file for data preprocessing, where each row contains the following keys: messages and verified_answer. This can be directly converted into MDS that the dataloader can handle.

done! check the new readme and scripts/data/messages_dataset_to_mds.py. Each messages sample gets passed to MDS with messages and metadata and so the verified_answer is located on inside the metadata json