Skip to content

Conversation

nitsanluke
Copy link
Contributor

@nitsanluke nitsanluke commented May 7, 2025

✨ Description

This PR creates a common interface for all GPTHuggingfaceDatasetConfig input columns via the new source_schema variable. Beyond the variable filed we require additional keys to preprocess and tokenize different types of datasets. (eg: SFT, combine cols, etc).
Therefore we have created a new variable source_schema which can accommodate these different data sources specific preprocessing and tokenization. Current variables field and loss_masking_spans are moved into TextColumnConfig as a type of input/data source.

Merge after #245

🔍 Type of change

Select all that apply:

  • 🐛 Bug fix (non-breaking change that addresses a specific issue)
  • 🚀 New feature (non-breaking change that adds functionality)
  • ⚠️ Breaking change (a change that could affect existing functionality)
  • 📈 Performance improvement/optimization (improves speed, memory usage, or efficiency)
  • 🛠️ Code refactor (non-functional changes that improve code readability, structure, etc.)
  • 📦 Dependency bump (updates dependencies, including Dockerfile or package changes)
  • 📝 Documentation change (updates documentation, including new content or typo fixes)
  • 🔧 Infrastructure/Build change (affects build process, CI/CD, or dependencies)

📝 Changes

List the key changes introduced in this PR:

✅ Checklist

Make sure the following tasks are completed before submitting the PR:

General

  • 📜 I have read and followed the contributing guidelines.
  • 🏷️ I am using a clear and descriptive PR title that summarizes the key change or feature introduced.
  • 🎉 The functionality is complete, and I have tested the changes.
  • 📝 I have updated the documentation if needed.
  • ⚠️ The change does not introduce any new issues (e.g., runtime warnings, type checker errors, linting problems, unhandled edge cases).
  • 🧩 I have commented my code, especially in hard-to-understand areas.

Dependencies and Configuration

  • 🐋 I have updated the Docker configuration or dependencies, if applicable.
  • 🔄 I have ensured compatibility with the existing setup after dependency changes.

Testing

  • 🧪 I have added or updated tests to cover my changes.
  • ✔️ New and existing tests pass locally with my changes.
  • 🚦 I have tested these changes on GPUs and verified training stability.
  • 🏋️ I have tested the changes on realistic training workloads, if applicable.

Copy link
Collaborator

@tscholak tscholak left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

one minor suggestion, otherwise LGTM!

@nitsanluke nitsanluke changed the title Combine GPTHuggingfaceDatasetConfig input sources into data_source Combine GPTHuggingfaceDatasetConfig input sources into source_schema May 9, 2025
@tobyzl2 tobyzl2 mentioned this pull request May 12, 2025
Merged
25 tasks
@nitsanluke nitsanluke changed the base branch from generalize_dynamic_classes to main June 3, 2025 18:32
@nitsanluke nitsanluke requested a review from jlamypoirier June 3, 2025 20:30
@nitsanluke nitsanluke marked this pull request as ready for review June 3, 2025 20:30

if self._loss_masking_spans_column is not None:
if self._loss_masking_spans_column not in dataset.column_names:
raise ValueError(f"Dataset does not have spans field '{self._loss_masking_spans_column}'.")
tokenize_fn = self._tokenize_batch_with_spans
elif self._config.dataset.chosen_text is not None and self._config.dataset.rejected_text is not None:
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @tobyzl2 can pls make sure the DPO conditions are properly met.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @nitsanluke, sharing a few lines here for checking DPO conditions. Essentially we want to ensure that these 3 are met

  1. If loss masking spans (SFT) are already enabled, preference spans (chosen/rejected) should not also be enabled.
  2. Chosen and rejected spans should either be both be specified or neither should be specified
  3. If both chosen and rejected are specified, make sure that they are within the dataset columns

if self._config.dataset.loss_masking_spans is not None and (
self._config.dataset.chosen_text is not None or self._config.dataset.rejected_text is not None
):
raise ValueError(f"Can not enable both loss masking spans and chosen/rejected loss masking spans.")
if (self._config.dataset.chosen_text is None) != (self._config.dataset.rejected_text is None):
raise ValueError(f"Both chosen and rejected loss masking spans must be specified if one is specified.")

elif self._config.dataset.chosen_text is not None and self._config.dataset.rejected_text is not None:
if self._config.dataset.chosen_text not in dataset.column_names:
raise ValueError(f"Dataset does not have chosen spans field '{self._config.dataset.chosen_text}'.")
if self._config.dataset.rejected_text not in dataset.column_names:
raise ValueError(f"Dataset does not have rejected spans field '{self._config.dataset.rejected_text}'.")

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @tobyzl2! I'm adding the checks back. There is some reduandancy on self._config.dataset.loss_masking_spans is not None but will leave it as is.

@nitsanluke nitsanluke requested a review from sohamparikh June 3, 2025 20:35
@nitsanluke
Copy link
Contributor Author

Sample config for the default text column tokenizing

loading_workers: 1
tokenize_workers: 1
saving_workers: 1
output_path: ./test_output/cpt_test_gsm8k
dataset:
  path: openai/gsm8k
  config_name: main
  split: train
  trust_remote_code: true
  source_schema:
    type: text_column
    input_column: question


tokenizer:
  path: /mnt/checkpoints/upstream/Mistral-Nemo-Base-2407/

@jlamypoirier
Copy link
Collaborator

Is this ready to merge?

@nitsanluke nitsanluke merged commit d9bb084 into main Jun 16, 2025
4 checks passed
@nitsanluke nitsanluke deleted the restructure_dataset_config_for_multi_source branch June 16, 2025 14:47
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants