Support for unstructured text corpus datasets for CPT

[RdoubleA]

Context

What is the purpose of this PR? Is it to

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

Please link to any issues this PR addresses: #845, #809

Continued pre-training involves an identical data processing pipeline as standard pre-training, where the model simply predicts the next token and completes the text. It does not require any templating or prompt formatting. Our existing dataset classes are specifically designed for instruct and chat tuning, but don't support free-form text corpuses.

Here, we add a dataset class TextDataset that simply calls load_dataset and tokenizes the text directly without any further processing. It uses encode which simply adds BOS/EOS tokens if needed. This should be compatible with llama2, llama3, and other model's formatting requirements.

Changelog

What are the changes made in this PR?

• Add TextDataset and appropriate tests
• Add an example dataset builder, cnn_dailymail_articles_dataset and appropriate tests

Test plan

Please make sure to do each of the following if applicable to your PR. (If you're not sure about any one of these just ask and we will happily help.)

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

Ran the full finetune distributed recipe with the cnn dailymail dataset.

2024-04-24:23:12:06,234 INFO     [_utils.py:34] Running FullFinetuneRecipeDistributed with resolved config:

batch_size: 2
checkpointer:
  _component_: torchtune.utils.FullModelMetaCheckpointer
  checkpoint_dir: /tmp/Meta-Llama-3-8B/original/
  checkpoint_files:
  - consolidated.00.pth
  model_type: LLAMA3
  output_dir: /tmp/Meta-Llama-3-8B/
  recipe_checkpoint: null
dataset:
  _component_: torchtune.datasets.cnn_dailymail_articles_dataset
device: cuda
dtype: bf16
enable_activation_checkpointing: true
epochs: 3
gradient_accumulation_steps: 1
log_every_n_steps: null
loss:
  _component_: torch.nn.CrossEntropyLoss
max_steps_per_epoch: null
metric_logger:
  _component_: torchtune.utils.metric_logging.WandBLogger
  log_dir: /tmp/alpaca-llama3-finetune
  name: cnn_dailymail
  project: torchtune
model:
  _component_: torchtune.models.llama3.llama3_8b
optimizer:
  _component_: torch.optim.AdamW
  foreach: false
  lr: 2.0e-05
output_dir: /tmp/alpaca-llama3-finetune
resume_from_checkpoint: false
seed: null
shuffle: true
tokenizer:
  _component_: torchtune.models.llama3.llama3_tokenizer
  path: /tmp/Meta-Llama-3-8B/original/tokenizer.model

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 55a6de3 with merge base 29ae975 ():
💚 Looks good so far! There are no failures yet. 💚

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

[rohan-varma]

This is cool, do we in general plan to support CPT in torchtune and plan for an RFC around this possibly?

[RdoubleA]

do we in general plan to support CPT in torchtune and plan for an RFC around this possibly?

I don't know that CPT will be sufficiently different from our finetune recipes to warrant its own, besides different data handling. @rohan-varma would you know :) or maybe @pbontrager ?

[pbontrager]

I don't know that CPT will be sufficiently different from our finetune recipes to warrant its own, besides different data handling. @rohan-varma would you know :) or maybe @pbontrager ?

I don't think so initially. I think this dataset is good and we could add a recipe in the future if there are additional custom things that people want for advanced use cases.

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 55.20833% with 43 lines in your changes are missing coverage. Please review.

Project coverage is 27.20%. Comparing base (cb8e65a) to head (6ead093).
Report is 2 commits behind head on main.

Files	Patch %	Lines
torchtune/datasets/_text_completion.py	44.00%	14 Missing ⚠️
...torchtune/datasets/test_text_completion_dataset.py	44.44%	10 Missing ⚠️
...s/torchtune/datasets/test_cnn_dailymail_dataset.py	57.89%	8 Missing ⚠️
tests/torchtune/datasets/test_wikitext_dataset.py	57.89%	8 Missing ⚠️
torchtune/data/_utils.py	50.00%	1 Missing ⚠️
torchtune/datasets/_cnn_dailymail.py	80.00%	1 Missing ⚠️
torchtune/datasets/_wikitext.py	80.00%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             main     #868       +/-   ##
===========================================
- Coverage   67.10%   27.20%   -39.91%     
===========================================
  Files         174      180        +6     
  Lines        7423     7518       +95     
===========================================
- Hits         4981     2045     -2936     
- Misses       2442     5473     +3031     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[joecummings]

What is "name"?

[RdoubleA]

let me add a comment, but it's basically specifying the subset of data

[joecummings]

Any update here? I want to train on https://huggingface.co/datasets/allenai/dolma soon.

[joecummings]

Sorry one last nit: Can you add an examples section like this?

torchtune/torchtune/datasets/_chat.py

Line 137 in 29ae975

Examples:

[joecummings]

It'll make the docs and usage much clearer

[joecummings]

🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥