feat(data): blend several sources on the direct-HF SFT path

Let DirectHFSFTDatasetConfig.source take a list of sources so a run can draw
its training rows from several Hugging Face sources at once, weighted by
source_weights. Each weight is an epoch count over its source.

A single source keeps deriving its own validation and test splits. A list has
no single split to derive from, so those splits need explicit sources or have
to be disabled.

Signed-off-by: ryan.u <ryumin93@gmail.com>

Author: bzantium

docs/training/data-preparation.md

@@ -189,6 +189,28 @@ dataset.preprocessing = PromptCompletionSFTPreprocessingConfig(
 
 Structured multi-turn rows require chat preprocessing; Bridge does not silently flatten them into prompt-completion text.
 
+To train on several sources at once, pass a list to `source`. Each entry in `source_weights` is an epoch count for its source: `1.0` keeps every row once, `2.5` keeps every row twice plus half of them again, and `0.5` keeps half of them. Fractional passes draw without replacement. Omitting the weights gives every source one pass.
+
+```python
+from megatron.bridge.data.builders import DirectHFSFTDatasetConfig, HFDatasetSourceConfig
+
+dataset = DirectHFSFTDatasetConfig(
+    seq_length=4096,
+    source=[
+        HFDatasetSourceConfig(dataset_name="squad"),
+        HFDatasetSourceConfig(dataset_name="gsm8k"),
+    ],
+    source_weights=[1.0, 2.0],
+    blend_seed=1234,
+    validation_source=HFDatasetSourceConfig(dataset_name="squad", split="validation"),
+    do_test=False,
+)
+```
+
+Sources are adapted to canonical SFT rows before mixing, so their raw column layouts do not have to match. A single source still derives its own validation and test splits; a list has no single split to derive from, so set `validation_source` and `test_source` explicitly or disable those splits.
+
+Weights count rows, not tokens. SFT rows vary in length, so sources with equal row counts can still contribute very different token counts.
+
 Known semantic datasets should use their preset name, for example `squad`, `gsm8k`, `openmathinstruct2`, `cord_v2`, `raven`, `rdr`, `medpix`, `cv17`, or `llava_video_178k`. Do not combine `dataset_name` with `path_or_dataset`, `subset`, or `schema_adapter`; a preset owns those coupled fields. `split`, `load_kwargs`, and `adapter_kwargs` remain available for split selection and declarative runtime options such as a video root. Presets validate published split support: `raven`, `rdr`, and `llava_video_178k` are train-only; `medpix` and `squad` have no test split; `gsm8k` has no validation split; and OpenMathInstruct-2 exposes training variants only. Disable unsupported derived validation/test splits or supply explicit compatible sources.
 
 For text chat, `hf_processor_path=None` reuses the training tokenizer only when that tokenizer already defines the intended chat template. Otherwise select a vocabulary-compatible instruction processor explicitly, as above.

src/megatron/bridge/data/builders/direct_hf_sft.py

@@ -15,7 +15,7 @@
 """Serializable config and runtime builder for direct Hugging Face SFT."""
 
 import logging
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 from dataclasses import dataclass, field
 from functools import partial
 from typing import Any, Literal
@@ -40,7 +40,9 @@
     HFDatasetSourceConfig,
     hf_dataset_supports_split,
     load_and_adapt_hf_dataset,
+    load_and_blend_hf_datasets,
     prepare_hf_dataset_sources,
+    resolve_blend_weights,
 )
 from megatron.bridge.models.hf_pretrained.utils import is_safe_repo
 from megatron.bridge.training.tokenizers.tokenizer import MegatronTokenizer
@@ -51,17 +53,29 @@
 CollateFunction = Callable[..., dict[str, torch.Tensor]]
 
 
+def _as_source(entry: "HFDatasetSourceConfig | Mapping[str, Any]") -> HFDatasetSourceConfig:
+    """Rebuild a source that an override round trip flattened into a mapping."""
+    if isinstance(entry, HFDatasetSourceConfig):
+        return entry
+    return HFDatasetSourceConfig(**dict(entry))
+
+
 @dataclass(kw_only=True)
 class DirectHFSFTDatasetConfig(DataloaderConfig):
     """Serializable configuration for direct Hugging Face SFT datasets.
 
     Chat preprocessing is the compatibility default for multimodal and
     conversation sources. New text recipes should select chat or paired-text
     preprocessing explicitly.
+
+    ``source`` takes one source or a list of them. A list blends the training
+    rows across the sources, weighted by ``source_weights``.
     """
 
     seq_length: int
-    source: HFDatasetSourceConfig
+    source: HFDatasetSourceConfig | list[HFDatasetSourceConfig]
+    source_weights: list[float] | None = None
+    blend_seed: int = 1234
     validation_source: HFDatasetSourceConfig | None = None
     test_source: HFDatasetSourceConfig | None = None
     preprocessing: SFTPreprocessingConfig = field(default_factory=ChatSFTPreprocessingConfig)
@@ -82,7 +96,7 @@ def validate(self) -> None:
         if self.seq_length <= 0:
             raise ValueError("seq_length must be greater than 0.")
         validate_sft_preprocessing_config(self.preprocessing)
-        self.source.validate()
+        self._validate_sources()
         if self.do_validation and self.validation_source is not None:
             self._inherit_source_adapter_kwargs(self.validation_source)
             self.validation_source.validate()
@@ -101,12 +115,50 @@ def validate(self) -> None:
         if self.in_batch_packing_pad_to_multiple_of <= 0:
             raise ValueError("in_batch_packing_pad_to_multiple_of must be greater than 0.")
 
+    @property
+    def training_sources(self) -> list[HFDatasetSourceConfig]:
+        """The sources the training rows are drawn from, blended or not."""
+        return list(self.source) if isinstance(self.source, list) else [self.source]
+
+    @property
+    def split_source(self) -> HFDatasetSourceConfig | None:
+        """The source that derives validation and test splits, if one is unambiguous."""
+        sources = self.training_sources
+        return sources[0] if len(sources) == 1 else None
+
+    def _validate_sources(self) -> None:
+        """Validate the training sources and any blend weights."""
+        # Any override re-serializes this config through OmegaConf, which returns
+        # dataclasses as plain mappings even when the override targeted another field.
+        if isinstance(self.source, list):
+            self.source = [_as_source(entry) for entry in self.source]
+        else:
+            self.source = _as_source(self.source)
+
+        sources = self.training_sources
+        if isinstance(self.source, list) or self.source_weights is not None:
+            resolve_blend_weights(sources, self.source_weights)
+        for entry in sources:
+            entry.validate()
+
+        if len(sources) > 1:
+            for name, split_source in (
+                ("validation_source", self.validation_source),
+                ("test_source", self.test_source),
+            ):
+                enabled = self.do_validation if name == "validation_source" else self.do_test
+                if enabled and split_source is None:
+                    raise ValueError(
+                        f"A blended source has no single split to derive from; set {name} or disable that split."
+                    )
+
     def _inherit_source_adapter_kwargs(self, split_source: HFDatasetSourceConfig) -> None:
         """Fill unset adapter arguments on another split of the training source."""
-        if split_source.dataset_name != self.source.dataset_name or not self.source.adapter_kwargs:
+        source = self.split_source
+        if source is None or split_source.dataset_name != source.dataset_name or not source.adapter_kwargs:
             return
         split_adapter_kwargs = dict(split_source.adapter_kwargs or {})
-        for key, value in self.source.adapter_kwargs.items():
+        for key, value in source.adapter_kwargs.items():
             if split_adapter_kwargs.get(key) is None:
                 split_adapter_kwargs[key] = value
         split_source.adapter_kwargs = split_adapter_kwargs
@@ -197,20 +249,31 @@ def select_direct_hf_sft_collate(
     raise ValueError("Prompt-completion preprocessing supports text-only examples.")
 
 
+def load_direct_hf_sft_train_examples(config: DirectHFSFTDatasetConfig) -> list[dict[str, Any]]:
+    """Load the training rows, blending them when the config names several sources."""
+    sources = config.training_sources
+    if len(sources) == 1 and config.source_weights is None:
+        return load_direct_hf_sft_examples(sources[0], config.preprocessing)
+    rows = load_and_blend_hf_datasets(sources, config.source_weights, seed=config.blend_seed)
+    return normalize_sft_examples(rows, config.preprocessing)
+
+
 def build_direct_hf_sft_split(
     config: DirectHFSFTDatasetConfig,
     source: HFDatasetSourceConfig,
     target_length: int,
     processor: Any,
     *,
     collate_impl: CollateFunction | None = None,
+    examples: list[dict[str, Any]] | None = None,
 ) -> DirectSFTDataset | None:
     """Build one requested direct-HF SFT split."""
     if target_length <= 0:
         return None
     from megatron.bridge.data.collators.registry import model_collate_required_for_all_examples
 
-    examples = load_direct_hf_sft_examples(source, config.preprocessing)
+    if examples is None:
+        examples = load_direct_hf_sft_examples(source, config.preprocessing)
     if collate_impl is None and model_collate_required_for_all_examples(type(processor).__name__):
         if not isinstance(config.preprocessing, ChatSFTPreprocessingConfig):
             raise ValueError(
@@ -256,7 +319,7 @@ def build(
             self.config.do_validation
             and context.valid_samples > 0
             and self.config.validation_source is None
-            and not hf_dataset_supports_split(self.config.source, "validation")
+            and not hf_dataset_supports_split(self.config.split_source, "validation")
         ):
             raise ValueError(
                 "The selected Hugging Face source has no validation split; disable validation or set one."
@@ -265,14 +328,19 @@ def build(
             self.config.do_test
             and context.test_samples > 0
             and self.config.test_source is None
-            and not hf_dataset_supports_split(self.config.source, "test")
+            and not hf_dataset_supports_split(self.config.split_source, "test")
         ):
             raise ValueError("The selected Hugging Face source has no test split; disable test or set one.")
-        validation_source = self.config.validation_source or self.config.source.with_split("validation")
-        test_source = self.config.test_source or self.config.source.with_split("test")
+        split_source = self.config.split_source
+        validation_source = self.config.validation_source or (
+            split_source.with_split("validation") if split_source is not None else None
+        )
+        test_source = self.config.test_source or (
+            split_source.with_split("test") if split_source is not None else None
+        )
         requested_sources = []
         if context.train_samples > 0:
-            requested_sources.append(self.config.source)
+            requested_sources.extend(self.config.training_sources)
         if self.config.do_validation and context.valid_samples > 0:
             requested_sources.append(validation_source)
         if self.config.do_test and context.test_samples > 0:
@@ -287,12 +355,14 @@ def build(
             and getattr(context.tokenizer, "library", None) in {"sentencepiece", "tiktoken"}
         ):
             processor = normalize_direct_hf_sft_processor(context.tokenizer)
+        train_examples = load_direct_hf_sft_train_examples(self.config) if context.train_samples > 0 else None
         train_dataset = build_direct_hf_sft_split(
             self.config,
-            self.config.source,
+            self.config.training_sources[0],
             context.train_samples,
             processor,
             collate_impl=self._collate_impl,
+            examples=train_examples,
         )
         valid_dataset = (
             build_direct_hf_sft_split(

src/megatron/bridge/data/sources/hf.py

@@ -14,6 +14,7 @@
 
 """Declarative Hugging Face sources, named presets, and shared normalization."""
 
+import random
 from collections.abc import Sequence
 from dataclasses import dataclass, replace
 from pathlib import Path
@@ -324,3 +325,98 @@ def load_and_adapt_hf_dataset(source: HFDatasetSourceConfig) -> list[dict[str, A
         adapter_name=resolved.schema_adapter,
         adapter_kwargs=resolved.adapter_kwargs,
     )
+
+
+def resolve_blend_weights(
+    sources: Sequence[HFDatasetSourceConfig],
+    weights: Sequence[float] | None,
+) -> list[float]:
+    """Validate blend weights against their sources and fill in the uniform default.
+
+    Args:
+        sources: The sources taking part in the blend.
+        weights: One positive weight per source, or None for equal weights.
+
+    Returns:
+        One weight per source.
+
+    Raises:
+        ValueError: If no source is given, if the counts disagree, or if a weight
+            is not positive.
+    """
+    if not sources:
+        raise ValueError("A blend must name at least one source.")
+    if weights is None:
+        return [1.0] * len(sources)
+    if len(weights) != len(sources):
+        raise ValueError(f"Blend sources and weights must be equal in number, got {len(sources)} and {len(weights)}.")
+    for weight in weights:
+        if weight <= 0:
+            raise ValueError(f"Blend weights must be positive, got {weight}.")
+    return [float(weight) for weight in weights]
+
+
+def blend_sft_rows(
+    rows_per_source: Sequence[Sequence[dict[str, Any]]],
+    weights: Sequence[float],
+    *,
+    seed: int,
+) -> list[dict[str, Any]]:
+    """Draw rows from several sources under per-source epoch counts.
+
+    A weight is how many passes the blend takes over a source: ``1.0`` keeps
+    every row once, ``2.5`` keeps every row twice plus half of them again, and
+    ``0.5`` keeps half of them. Fractional passes draw without replacement.
+
+    Weights count rows, not tokens. SFT rows vary in length, so equal row counts
+    do not make equal token counts.
+
+    Args:
+        rows_per_source: Canonical SFT rows for each source of the blend.
+        weights: One positive weight per source.
+        seed: Seeds the fractional draws and the shuffle that interleaves sources.
+
+    Returns:
+        The blended rows, holding references to the rows that were passed in.
+
+    Raises:
+        ValueError: If a source contributed no rows.
+    """
+    rng = random.Random(seed)
+    blended: list[dict[str, Any]] = []
+    for index, (rows, weight) in enumerate(zip(rows_per_source, weights)):
+        if not rows:
+            raise ValueError(f"Blend source {index} produced no rows.")
+        full_passes = int(weight)
+        for _ in range(full_passes):
+            blended.extend(rows)
+        remainder = round(len(rows) * (weight - full_passes))
+        if remainder:
+            blended.extend(rng.sample(list(rows), remainder))
+
+    rng.shuffle(blended)
+    return blended
+
+
+def load_and_blend_hf_datasets(
+    sources: Sequence[HFDatasetSourceConfig],
+    weights: Sequence[float] | None,
+    *,
+    seed: int,
+) -> list[dict[str, Any]]:
+    """Load several Hugging Face splits and mix them under per-source epoch counts.
+
+    Each source is adapted to canonical SFT rows first, so sources with unrelated
+    column layouts can be blended as long as they share a schema adapter target.
+
+    Args:
+        sources: The sources taking part in the blend.
+        weights: One positive weight per source, or None for one pass each.
+        seed: Seeds the fractional draws and the shuffle that interleaves sources.
+
+    Returns:
+        The blended canonical SFT rows.
+    """
+    resolved_weights = resolve_blend_weights(sources, weights)
+    rows_per_source = [load_and_adapt_hf_dataset(source) for source in sources]
+    return blend_sft_rows(rows_per_source, resolved_weights, seed=seed)

tests/functional_tests/test_groups/training/test_sft.py

@@ -180,6 +180,67 @@ def test_direct_hf_prompt_completion_training_step(self, tmp_path):
         finally:
             clear_directories(shared_base_dir)
 
+    @pytest.mark.run_only_on("GPU")
+    def test_direct_hf_blended_sources_training_step(self, tmp_path):
+        """Run one GPT training step over a weighted blend of two direct-HF sources."""
+        initialize_distributed()
+        shared_base_dir = broadcast_path(tmp_path)
+        checkpoint_dir, tensorboard_dir, _, _ = self._setup_directories(shared_base_dir)
+        first_path = os.path.join(shared_base_dir, "blend-a.jsonl")
+        second_path = os.path.join(shared_base_dir, "blend-b.jsonl")
+
+        if torch.distributed.get_rank() == 0:
+            # NullTokenizer reads whitespace-separated integers, so the sources are
+            # told apart by disjoint number ranges rather than by a text label.
+            for path, offset, rows in ((first_path, 0, 8), (second_path, 100, 4)):
+                with open(path, "w", encoding="utf-8") as output_file:
+                    for index in range(rows):
+                        value = offset + index
+                        record = {"input": f"{value} {value}", "output": str(value + 1)}
+                        output_file.write(json.dumps(record) + "\n")
+        torch.distributed.barrier()
+
+        try:
+            config = self._create_config(1, checkpoint_dir, tensorboard_dir, seq_length=16)
+            config.model.vocab_size = 1024
+            config.train.global_batch_size = 4
+            config.train.micro_batch_size = 2
+            config.tokenizer = TokenizerConfig(tokenizer_type="NullTokenizer", vocab_size=1024)
+            config.dataset = DirectHFSFTDatasetConfig(
+                seq_length=16,
+                source=[
+                    HFDatasetSourceConfig(
+                        path_or_dataset="json",
+                        load_kwargs={"data_files": {"train": first_path}},
+                    ),
+                    HFDatasetSourceConfig(
+                        path_or_dataset="json",
+                        load_kwargs={"data_files": {"train": second_path}},
+                    ),
+                ],
+                source_weights=[1.0, 3.0],
+                preprocessing=PromptCompletionSFTPreprocessingConfig(
+                    prompt_column="input",
+                    completion_column="output",
+                ),
+                do_validation=False,
+                do_test=False,
+                enable_in_batch_packing=True,
+                pad_to_multiple_of=1,
+                num_workers=0,
+                persistent_workers=False,
+            )
+
+            pretrain(config, forward_step)
+            verify_checkpoint_files(
+                checkpoint_dir,
+                1,
+                ckpt_format=config.checkpoint.ckpt_format,
+                storage_writers_per_rank=config.checkpoint.storage_writers_per_rank,
+            )
+        finally:
+            clear_directories(shared_base_dir)
+
     @pytest.mark.run_only_on("GPU")
     def test_direct_hf_non_packed_context_parallel_training_step(self, tmp_path):
         """Run direct-HF dynamic padding through a two-rank context-parallel step."""