[WIP] Support Continual Pretraining Multi Dataset using Streaming #2355
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2355
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
| # Dataset | ||
| dataloader: | ||
| shuffle: True | ||
| collate_fn: torchtune.data.padded_collate_sft |
There was a problem hiding this comment.
I had to add this to get it to work. Not sure if it is considered a hack
| if "data_dir" in load_dataset_kwargs: | ||
| load_dataset_kwargs[ | ||
| "data_files" | ||
| ] = f"{load_dataset_kwargs.pop('data_dir')}/*.jsonl" |
There was a problem hiding this comment.
I did this to support providing a local data directory that contains .jsonl files
| - source: cerebras/SlimPajama-627B | ||
| split: train | ||
| column: text | ||
| weight: 0.9 | ||
| - source: bigcode/the-stack-dedup | ||
| split: train | ||
| column: content | ||
| weight: 0.1 |
There was a problem hiding this comment.
Here I am supporting passing large generic HuggingFace datasets, and loading from them in streaming mode without downloading the full dataset.
I also got passing local directory with .jsonl file but I forgot how I set it in the .jsonl file.
| @@ -0,0 +1,948 @@ | |||
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |||
There was a problem hiding this comment.
All of the contents of this file except the _setup_data() function was copied from recipes/full_finetune_distributed.py
Here is the diff comparing the 2 files:
https://www.diffchecker.com/CjtUG31A/
| utils.log_rank_zero(log, "Optimizer is initialized.") | ||
| return optimizer | ||
|
|
||
| def _setup_data( |
There was a problem hiding this comment.
This _setup_data(..) function was copied from recipes/lora_finetune_distributed_multi_dataset.py and modified.
You can view the modifications in this online Diff: https://www.diffchecker.com/Cg4AiMni/
Context
What is the purpose of this PR? Is it to
Please link to any issues this PR addresses.
Changelog
What are the changes made in this PR?
I had to make some workarounds to the multidata approach implemented in #1929 . I am not sure if this is because I didn't understand the codebase or if some classes related to data loading and collating need to be refactored.
Cc @andrewkho, @RdoubleA, @ebsmothers
Test plan
Please make sure to do each of the following if applicable to your PR. If you're unsure about any one of these just ask and we will happily help. We also have a contributing page for some guidance on contributing.
pre-commit install)pytest testspytest tests -m integration_testUX
If your function changed a public API, please add a dummy example of what the user experience will look like when calling it.
Here is a docstring example
and a tutorial example