Skip to content

Commit c83f9c8

Browse files
committed
merge: update branch to latest main
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
2 parents 0e292d9 + 0c2443e commit c83f9c8

7 files changed

Lines changed: 629 additions & 12 deletions

File tree

‎docs/fern/versions/nightly/pages/training/data-preparation.mdx‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -191,6 +191,28 @@ dataset.preprocessing = PromptCompletionSFTPreprocessingConfig(
191191

192192
Structured multi-turn rows require chat preprocessing; Bridge does not silently flatten them into prompt-completion text.
193193

194+
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 and round up to a whole row. Omitting the weights gives every source one pass.
195+
196+
```python
197+
from megatron.bridge.data.builders import DirectHFSFTDatasetConfig, HFDatasetSourceConfig
198+
199+
dataset = DirectHFSFTDatasetConfig(
200+
seq_length=4096,
201+
source=[
202+
HFDatasetSourceConfig(dataset_name="squad"),
203+
HFDatasetSourceConfig(dataset_name="gsm8k"),
204+
],
205+
source_weights=[1.0, 2.0],
206+
blend_seed=1234,
207+
validation_source=HFDatasetSourceConfig(dataset_name="squad", split="validation"),
208+
do_test=False,
209+
)
210+
```
211+
212+
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.
213+
214+
Weights count rows, not tokens. SFT rows vary in length, so sources with equal row counts can still contribute very different token counts.
215+
194216
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.
195217

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

‎docs/training/data-preparation.md‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,28 @@ dataset.preprocessing = PromptCompletionSFTPreprocessingConfig(
189189

190190
Structured multi-turn rows require chat preprocessing; Bridge does not silently flatten them into prompt-completion text.
191191

192+
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 and round up to a whole row. Omitting the weights gives every source one pass.
193+
194+
```python
195+
from megatron.bridge.data.builders import DirectHFSFTDatasetConfig, HFDatasetSourceConfig
196+
197+
dataset = DirectHFSFTDatasetConfig(
198+
seq_length=4096,
199+
source=[
200+
HFDatasetSourceConfig(dataset_name="squad"),
201+
HFDatasetSourceConfig(dataset_name="gsm8k"),
202+
],
203+
source_weights=[1.0, 2.0],
204+
blend_seed=1234,
205+
validation_source=HFDatasetSourceConfig(dataset_name="squad", split="validation"),
206+
do_test=False,
207+
)
208+
```
209+
210+
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.
211+
212+
Weights count rows, not tokens. SFT rows vary in length, so sources with equal row counts can still contribute very different token counts.
213+
192214
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.
193215

194216
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‎

Lines changed: 90 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
"""Serializable config and runtime builder for direct Hugging Face SFT."""
1616

1717
import logging
18-
from collections.abc import Callable
18+
from collections.abc import Callable, Mapping
1919
from dataclasses import dataclass, field
2020
from functools import partial
2121
from typing import Any, Literal
@@ -40,7 +40,9 @@
4040
HFDatasetSourceConfig,
4141
hf_dataset_supports_split,
4242
load_and_adapt_hf_dataset,
43+
load_and_blend_hf_datasets,
4344
prepare_hf_dataset_sources,
45+
resolve_blend_weights,
4446
)
4547
from megatron.bridge.models.hf_pretrained.utils import is_safe_repo
4648
from megatron.bridge.training.tokenizers.tokenizer import MegatronTokenizer
@@ -51,17 +53,29 @@
5153
CollateFunction = Callable[..., dict[str, torch.Tensor]]
5254

5355

56+
def _as_source(entry: "HFDatasetSourceConfig | Mapping[str, Any]") -> HFDatasetSourceConfig:
57+
"""Rebuild a source that an override round trip flattened into a mapping."""
58+
if isinstance(entry, HFDatasetSourceConfig):
59+
return entry
60+
return HFDatasetSourceConfig(**dict(entry))
61+
62+
5463
@dataclass(kw_only=True)
5564
class DirectHFSFTDatasetConfig(DataloaderConfig):
5665
"""Serializable configuration for direct Hugging Face SFT datasets.
5766
5867
Chat preprocessing is the compatibility default for multimodal and
5968
conversation sources. New text recipes should select chat or paired-text
6069
preprocessing explicitly.
70+
71+
``source`` takes one source or a list of them. A list blends the training
72+
rows across the sources, weighted by ``source_weights``.
6173
"""
6274

6375
seq_length: int
64-
source: HFDatasetSourceConfig
76+
source: HFDatasetSourceConfig | list[HFDatasetSourceConfig]
77+
source_weights: list[float] | None = None
78+
blend_seed: int = 1234
6579
validation_source: HFDatasetSourceConfig | None = None
6680
test_source: HFDatasetSourceConfig | None = None
6781
preprocessing: SFTPreprocessingConfig = field(default_factory=ChatSFTPreprocessingConfig)
@@ -82,7 +96,11 @@ def validate(self) -> None:
8296
if self.seq_length <= 0:
8397
raise ValueError("seq_length must be greater than 0.")
8498
validate_sft_preprocessing_config(self.preprocessing)
85-
self.source.validate()
99+
self._validate_sources()
100+
if self.validation_source is not None:
101+
self.validation_source = _as_source(self.validation_source)
102+
if self.test_source is not None:
103+
self.test_source = _as_source(self.test_source)
86104
if self.do_validation and self.validation_source is not None:
87105
self._inherit_source_adapter_kwargs(self.validation_source)
88106
self.validation_source.validate()
@@ -101,12 +119,54 @@ def validate(self) -> None:
101119
if self.in_batch_packing_pad_to_multiple_of <= 0:
102120
raise ValueError("in_batch_packing_pad_to_multiple_of must be greater than 0.")
103121

122+
@property
123+
def training_sources(self) -> list[HFDatasetSourceConfig]:
124+
"""The sources the training rows are drawn from, blended or not."""
125+
return list(self.source) if isinstance(self.source, list) else [self.source]
126+
127+
@property
128+
def split_source(self) -> HFDatasetSourceConfig | None:
129+
"""The source that derives validation and test splits, if one is unambiguous."""
130+
sources = self.training_sources
131+
return sources[0] if len(sources) == 1 else None
132+
133+
def _validate_sources(self) -> None:
134+
"""Validate the training sources and any blend weights."""
135+
# Any override re-serializes this config through OmegaConf, which returns
136+
# dataclasses as plain mappings even when the override targeted another field.
137+
if isinstance(self.source, list):
138+
self.source = [_as_source(entry) for entry in self.source]
139+
else:
140+
self.source = _as_source(self.source)
141+
142+
sources = self.training_sources
143+
if isinstance(self.source, list) or self.source_weights is not None:
144+
resolve_blend_weights(sources, self.source_weights)
145+
# random.Random(None) reseeds from system entropy, so an unset seed
146+
# would give every rank and every rebuild a different blend order.
147+
if not isinstance(self.blend_seed, int) or isinstance(self.blend_seed, bool):
148+
raise ValueError(f"blend_seed must be an integer, got {self.blend_seed!r}.")
149+
for entry in sources:
150+
entry.validate()
151+
152+
if len(sources) > 1:
153+
for name, split_source in (
154+
("validation_source", self.validation_source),
155+
("test_source", self.test_source),
156+
):
157+
enabled = self.do_validation if name == "validation_source" else self.do_test
158+
if enabled and split_source is None:
159+
raise ValueError(
160+
f"A blended source has no single split to derive from; set {name} or disable that split."
161+
)
162+
104163
def _inherit_source_adapter_kwargs(self, split_source: HFDatasetSourceConfig) -> None:
105164
"""Fill unset adapter arguments on another split of the training source."""
106-
if split_source.dataset_name != self.source.dataset_name or not self.source.adapter_kwargs:
165+
source = self.split_source
166+
if source is None or split_source.dataset_name != source.dataset_name or not source.adapter_kwargs:
107167
return
108168
split_adapter_kwargs = dict(split_source.adapter_kwargs or {})
109-
for key, value in self.source.adapter_kwargs.items():
169+
for key, value in source.adapter_kwargs.items():
110170
if split_adapter_kwargs.get(key) is None:
111171
split_adapter_kwargs[key] = value
112172
split_source.adapter_kwargs = split_adapter_kwargs
@@ -197,20 +257,31 @@ def select_direct_hf_sft_collate(
197257
raise ValueError("Prompt-completion preprocessing supports text-only examples.")
198258

199259

260+
def load_direct_hf_sft_train_examples(config: DirectHFSFTDatasetConfig) -> list[dict[str, Any]]:
261+
"""Load the training rows, blending them when the config names several sources."""
262+
sources = config.training_sources
263+
if len(sources) == 1 and config.source_weights is None:
264+
return load_direct_hf_sft_examples(sources[0], config.preprocessing)
265+
rows = load_and_blend_hf_datasets(sources, config.source_weights, seed=config.blend_seed)
266+
return normalize_sft_examples(rows, config.preprocessing)
267+
268+
200269
def build_direct_hf_sft_split(
201270
config: DirectHFSFTDatasetConfig,
202271
source: HFDatasetSourceConfig,
203272
target_length: int,
204273
processor: Any,
205274
*,
206275
collate_impl: CollateFunction | None = None,
276+
examples: list[dict[str, Any]] | None = None,
207277
) -> DirectSFTDataset | None:
208278
"""Build one requested direct-HF SFT split."""
209279
if target_length <= 0:
210280
return None
211281
from megatron.bridge.data.collators.registry import model_collate_required_for_all_examples
212282

213-
examples = load_direct_hf_sft_examples(source, config.preprocessing)
283+
if examples is None:
284+
examples = load_direct_hf_sft_examples(source, config.preprocessing)
214285
if collate_impl is None and model_collate_required_for_all_examples(type(processor).__name__):
215286
if not isinstance(config.preprocessing, ChatSFTPreprocessingConfig):
216287
raise ValueError(
@@ -256,7 +327,7 @@ def build(
256327
self.config.do_validation
257328
and context.valid_samples > 0
258329
and self.config.validation_source is None
259-
and not hf_dataset_supports_split(self.config.source, "validation")
330+
and not hf_dataset_supports_split(self.config.split_source, "validation")
260331
):
261332
raise ValueError(
262333
"The selected Hugging Face source has no validation split; disable validation or set one."
@@ -265,14 +336,19 @@ def build(
265336
self.config.do_test
266337
and context.test_samples > 0
267338
and self.config.test_source is None
268-
and not hf_dataset_supports_split(self.config.source, "test")
339+
and not hf_dataset_supports_split(self.config.split_source, "test")
269340
):
270341
raise ValueError("The selected Hugging Face source has no test split; disable test or set one.")
271-
validation_source = self.config.validation_source or self.config.source.with_split("validation")
272-
test_source = self.config.test_source or self.config.source.with_split("test")
342+
split_source = self.config.split_source
343+
validation_source = self.config.validation_source or (
344+
split_source.with_split("validation") if split_source is not None else None
345+
)
346+
test_source = self.config.test_source or (
347+
split_source.with_split("test") if split_source is not None else None
348+
)
273349
requested_sources = []
274350
if context.train_samples > 0:
275-
requested_sources.append(self.config.source)
351+
requested_sources.extend(self.config.training_sources)
276352
if self.config.do_validation and context.valid_samples > 0:
277353
requested_sources.append(validation_source)
278354
if self.config.do_test and context.test_samples > 0:
@@ -287,12 +363,14 @@ def build(
287363
and getattr(context.tokenizer, "library", None) in {"sentencepiece", "tiktoken"}
288364
):
289365
processor = normalize_direct_hf_sft_processor(context.tokenizer)
366+
train_examples = load_direct_hf_sft_train_examples(self.config) if context.train_samples > 0 else None
290367
train_dataset = build_direct_hf_sft_split(
291368
self.config,
292-
self.config.source,
369+
self.config.training_sources[0],
293370
context.train_samples,
294371
processor,
295372
collate_impl=self._collate_impl,
373+
examples=train_examples,
296374
)
297375
valid_dataset = (
298376
build_direct_hf_sft_split(

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

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@
1414

1515
"""Declarative Hugging Face sources, named presets, and shared normalization."""
1616

17+
import math
18+
import random
1719
from collections.abc import Sequence
1820
from dataclasses import dataclass, replace
1921
from pathlib import Path
@@ -324,3 +326,107 @@ def load_and_adapt_hf_dataset(source: HFDatasetSourceConfig) -> list[dict[str, A
324326
adapter_name=resolved.schema_adapter,
325327
adapter_kwargs=resolved.adapter_kwargs,
326328
)
329+
330+
331+
def resolve_blend_weights(
332+
sources: Sequence[HFDatasetSourceConfig],
333+
weights: Sequence[float] | None,
334+
) -> list[float]:
335+
"""Validate blend weights against their sources and fill in the uniform default.
336+
337+
Args:
338+
sources: The sources taking part in the blend.
339+
weights: One positive weight per source, or None for equal weights.
340+
341+
Returns:
342+
One weight per source.
343+
344+
Raises:
345+
ValueError: If no source is given, if the counts disagree, or if a weight
346+
is not a positive finite number.
347+
"""
348+
if not sources:
349+
raise ValueError("A blend must name at least one source.")
350+
if weights is None:
351+
return [1.0] * len(sources)
352+
if len(weights) != len(sources):
353+
raise ValueError(f"Blend sources and weights must be equal in number, got {len(sources)} and {len(weights)}.")
354+
355+
resolved = []
356+
for weight in weights:
357+
# NaN and infinity pass every comparison below, so they are rejected by
358+
# kind rather than by range.
359+
value = float(weight)
360+
if not math.isfinite(value) or value <= 0:
361+
raise ValueError(f"Blend weights must be positive finite numbers, got {weight}.")
362+
resolved.append(value)
363+
return resolved
364+
365+
366+
def blend_sft_rows(
367+
rows_per_source: Sequence[Sequence[dict[str, Any]]],
368+
weights: Sequence[float],
369+
*,
370+
seed: int,
371+
) -> list[dict[str, Any]]:
372+
"""Draw rows from several sources under per-source epoch counts.
373+
374+
A weight is how many passes the blend takes over a source: ``1.0`` keeps
375+
every row once, ``2.5`` keeps every row twice plus half of them again, and
376+
``0.5`` keeps half of them. Fractional passes draw without replacement and
377+
round up to a whole row, so a positive weight always keeps at least one.
378+
379+
Weights count rows, not tokens. SFT rows vary in length, so equal row counts
380+
do not make equal token counts.
381+
382+
Args:
383+
rows_per_source: Canonical SFT rows for each source of the blend.
384+
weights: One positive weight per source.
385+
seed: Seeds the fractional draws and the shuffle that interleaves sources.
386+
387+
Returns:
388+
The blended rows, holding references to the rows that were passed in.
389+
390+
Raises:
391+
ValueError: If a source contributed no rows.
392+
"""
393+
rng = random.Random(seed)
394+
blended: list[dict[str, Any]] = []
395+
for index, (rows, weight) in enumerate(zip(rows_per_source, weights)):
396+
if not rows:
397+
raise ValueError(f"Blend source {index} produced no rows.")
398+
full_passes = int(weight)
399+
for _ in range(full_passes):
400+
blended.extend(rows)
401+
# Rounding up rather than to nearest, so a positive weight always draws
402+
# something. A share below half a row would otherwise drop the source.
403+
remainder = math.ceil(len(rows) * (weight - full_passes))
404+
if remainder:
405+
blended.extend(rng.sample(list(rows), remainder))
406+
407+
rng.shuffle(blended)
408+
return blended
409+
410+
411+
def load_and_blend_hf_datasets(
412+
sources: Sequence[HFDatasetSourceConfig],
413+
weights: Sequence[float] | None,
414+
*,
415+
seed: int,
416+
) -> list[dict[str, Any]]:
417+
"""Load several Hugging Face splits and mix them under per-source epoch counts.
418+
419+
Each source is adapted to canonical SFT rows first, so sources with unrelated
420+
column layouts can be blended as long as they share a schema adapter target.
421+
422+
Args:
423+
sources: The sources taking part in the blend.
424+
weights: One positive weight per source, or None for one pass each.
425+
seed: Seeds the fractional draws and the shuffle that interleaves sources.
426+
427+
Returns:
428+
The blended canonical SFT rows.
429+
"""
430+
resolved_weights = resolve_blend_weights(sources, weights)
431+
rows_per_source = [load_and_adapt_hf_dataset(source) for source in sources]
432+
return blend_sft_rows(rows_per_source, resolved_weights, seed=seed)

0 commit comments

Comments
 (0)