Combine GPTHuggingfaceDatasetConfig input sources into `source_schema` #255

nitsanluke · 2025-05-07T19:13:52Z

✨ 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.

…nfig_for_multi_source

tscholak

one minor suggestion, otherwise LGTM!

fast_llm/data/preparator/gpt_memmap/config.py

tscholak · 2025-05-08T18:26:14Z

fast_llm/data/preparator/gpt_memmap/config.py

+class SourceSchemaConfig(Config):
+    pass
+
+class TextColumnConfig(SourceSchemaConfig):


Suggested change

class TextColumnConfig(SourceSchemaConfig):

class TextColumnConfig(SourceSchemaConfig):

type: ClassVar[str] = "text_column"

@jlamypoirier, what's the endorsed idiom for this?

Nothing until we merge @245. Then we can add to registry. This one won't work unless it's marked as a classvar

@nitsanluke's PR is assuming #245 is getting merged. It's branching off from the generalize_dynamic_classes branch of #245.

Oh, then it needs to be registered after the class definition, ex. https://github.com/ServiceNow/Fast-LLM/pull/245/files#diff-7001a0c0d12d5adf8e07c112218448985de3f25255f166f01d001768c7130788R148. Here something like SourceSchemaConfig.register_subclass("text_column", TextColumnConfig).

Co-authored-by: Torsten Scholak <torsten.scholak@googlemail.com>

…com:ServiceNow/Fast-LLM into restructure_dataset_config_for_multi_source

Co-authored-by: Torsten Scholak <torsten.scholak@googlemail.com>

…com:ServiceNow/Fast-LLM into restructure_dataset_config_for_multi_source

sohamparikh · 2025-05-08T22:58:19Z

fast_llm/data/preparator/gpt_memmap/prepare.py

+    _data_column: str
+    _loss_masking_spans_column: str | None


qq: do we need to make these class variables? With more schemas coming in, we would have to define class variables for fields in each of them. Wouldn't it be simpler to directly access them from the schema config?

the type class variable is a discriminator field. It tells fast-llm to instantiate the right class when deserializing the config from yaml

We only need one such discriminator field (type: ClassVar[str]) per source schema config class

Wait a sec. This isn’t the class I thought I was looking at. I think you’re right @sohamparikh, I don’t think we need these fields here. @nitsanluke?

@sohamparikh The idea was to separate the set of vars used by prepare from the config; since we can't change config vars after defining. So it become complex to use the same var with multiple data sources (need to set during validation). Eg: For the concat case the new col that would be generated during the execution will take the place of _data_column and rest of the prepare code doesn't have to change. The separation allows for different source schema to use the class vars to setup for the specific use-case they are dealing with. We can discuss further to see if the alternative has more benefits.

Seems reasonable to use a typing.ClassVar[str] in the config class. Or for more flexibility it could be a @property or @cachedproperty

Thanks @nitsanluke, makes sense to separate them from the config
With more modalities coming in (image, audio), could we rename _data_column to _text_column?

Renamed to _text_column

jlamypoirier · 2025-05-09T15:43:53Z

fast_llm/data/preparator/gpt_memmap/config.py

-        default="text",
-        desc="Field of the dataset to use.",
+    source_schema: SourceSchemaConfig = Field(
+        default_factory=TextColumnConfig,


default_factory doesn't actually work and is removed for config classes in #245. I'm working on a way to set a default subclass https://github.com/ServiceNow/Fast-LLM/pull/245/files#diff-7001a0c0d12d5adf8e07c112218448985de3f25255f166f01d001768c7130788R44.

Will remove it from here and keep a TODO.

nitsanluke added 2 commits May 7, 2025 19:02

combine data source inputs to data_source

94e56e1

Merge branch 'generalize_dynamic_classes' into restructure_dataset_co…

6a30d76

…nfig_for_multi_source

tscholak approved these changes May 8, 2025

View reviewed changes

fast_llm/data/preparator/gpt_memmap/config.py Outdated Show resolved Hide resolved

tscholak reviewed May 8, 2025

View reviewed changes

nitsanluke and others added 9 commits May 8, 2025 15:27

Update fast_llm/data/preparator/gpt_memmap/config.py

1dab7de

Co-authored-by: Torsten Scholak <torsten.scholak@googlemail.com>

Update fast_llm/data/preparator/gpt_memmap/config.py

36b42b9

Co-authored-by: Torsten Scholak <torsten.scholak@googlemail.com>

merge

c6876ac

Merge branch 'restructure_dataset_config_for_multi_source' of github.…

eadd49a

…com:ServiceNow/Fast-LLM into restructure_dataset_config_for_multi_source

Update fast_llm/data/preparator/gpt_memmap/config.py

a5b06d8

Co-authored-by: Torsten Scholak <torsten.scholak@googlemail.com>

Merge branch 'restructure_dataset_config_for_multi_source' of github.…

272c63f

…com:ServiceNow/Fast-LLM into restructure_dataset_config_for_multi_source

remove duplicate

cbcde98

name change

694181f

adding ClassVar type

fdf44d3

nitsanluke mentioned this pull request May 8, 2025

Concat prompt and completion cols for tokenizing #257

Draft

14 tasks

sohamparikh reviewed May 8, 2025

View reviewed changes

nitsanluke changed the title ~~Combine GPTHuggingfaceDatasetConfig input sources into data_source~~ Combine GPTHuggingfaceDatasetConfig input sources into source_schema May 9, 2025

jlamypoirier reviewed May 9, 2025

View reviewed changes

tobyzl2 mentioned this pull request May 12, 2025

DPO #223

Merged

25 tasks

sohamparikh mentioned this pull request May 12, 2025

Support chat template in prepare #262

Open

4 tasks

nitsanluke added 3 commits May 14, 2025 15:00

rename to _text_column

0909768

remove default_factory for source_schema

1a6b78b

minor comment

662f318

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Combine GPTHuggingfaceDatasetConfig input sources into `source_schema` #255

Combine GPTHuggingfaceDatasetConfig input sources into `source_schema` #255

nitsanluke commented May 7, 2025 •

edited

Loading

tscholak left a comment

tscholak May 8, 2025 •

edited

Loading

tscholak May 8, 2025

jlamypoirier May 8, 2025

tscholak May 8, 2025

jlamypoirier May 9, 2025

sohamparikh May 8, 2025

tscholak May 8, 2025

tscholak May 8, 2025

tscholak May 8, 2025

nitsanluke May 9, 2025 •

edited

Loading

jlamypoirier May 9, 2025

sohamparikh May 9, 2025

nitsanluke May 14, 2025

jlamypoirier May 9, 2025

nitsanluke May 14, 2025

	class TextColumnConfig(SourceSchemaConfig):
	class TextColumnConfig(SourceSchemaConfig):
	type: ClassVar[str] = "text_column"

Combine GPTHuggingfaceDatasetConfig input sources into source_schema #255

Are you sure you want to change the base?

Combine GPTHuggingfaceDatasetConfig input sources into source_schema #255

Conversation

nitsanluke commented May 7, 2025 • edited Loading

✨ Description

🔍 Type of change

📝 Changes

✅ Checklist

General

Dependencies and Configuration

Testing

tscholak left a comment

Choose a reason for hiding this comment

tscholak May 8, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

nitsanluke May 9, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Combine GPTHuggingfaceDatasetConfig input sources into `source_schema` #255

Combine GPTHuggingfaceDatasetConfig input sources into `source_schema` #255

nitsanluke commented May 7, 2025 •

edited

Loading

tscholak May 8, 2025 •

edited

Loading

nitsanluke May 9, 2025 •

edited

Loading