feat(data): blend packed Parquet SFT sources

[yaoyu-33]

Summary

• add packed_train_data_blend=([sources], [weights]) for offline-packed GPT SFT Parquet data
• treat each file, directory, or glob as one logical source while preserving shard support within a source
• build deterministic low-discrepancy weighted mappings with seeded source-local shuffling and optional over/under-sampling through max_train_samples
• allow a packed blend to be the standalone training source when validation and test are disabled, or pair it with a raw-data source for validation and test splits
• validate paths, positive finite weights, mutually exclusive single/blended paths, and reject pad_cu_seqlens because each source requires its own packing metadata
• document the supported JSONL-to-packed-Parquet-to-blend workflow and remove stale packed-storage guidance

Motivation

A real NVIDIA Nemotron 3 Nano 30B-A3B SFT comparison on 8 H100s at sequence length 2048 and GBS 128 found bin/idx and Parquet statistically indistinguishable: indexed averaged 104.679 s per cold optimizer step and Parquet averaged 103.810 s. Since storage-format overhead is negligible at this model scale, weighted multi-source Parquet support is more useful than adding another packed storage format.

Tests

Validated on the current head in a Megatron Bridge Linux container:

• python -m pytest -q tests/unit_tests/data/packing/test_parquet.py tests/unit_tests/data/builders/test_gpt_sft_config.py tests/unit_tests/data/builders/test_gpt_sft_builder.py — 101 passed
• python -m pytest -q tests/unit_tests/training/test_config.py -k "packed or parquet or packing" — 36 passed, 243 deselected
• python -m pytest -q tests/unit_tests/doc_consistency/test_readme_consistency.py tests/unit_tests/tutorials/test_text_dataset_tutorials.py — 26 passed
• end-to-end smoke: independently packed two JSONL sources to Parquet, then built a 3:1 blend through GPTSFTDatasetBuilder; 8 requested samples produced deterministic 6/2 source counts and collated successfully
• pre-commit run --all-files — passed

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[claude]

Light review — LGTM with one edge-case note.

This adds packed_train_data_blend for weighted sampling across independently packed Parquet SFT sources. Validation, standalone-source semantics, and the low-discrepancy blend scheduler are all reasonable and well covered by the new unit tests. Docs/tutorials are updated consistently and the error-message wording in text-only-sft/README.md matches the actual validate() strings.

Correctness (verified):

• PackedSequenceSpecs.post_init rejects setting both packed_train_data_path and packed_train_data_blend; blend weights are validated for finiteness/positivity/finite-sum.
• GPTSFTDatasetConfig.validate() correctly relaxes the mutual-exclusion rule and requires do_validation=False/do_test=False for a standalone blend.
• Negative-index padding convention in the blend getitem (zeroing loss_mask) matches the base packed datasets.
• train_path_packed now raises when a blend is configured, preventing a stale single-path lookup.

Observation (non-blocking): _build_packed_parquet_blend passes the single builder-managed pack_metadata_path (self.pack_metadata) to every source, and forwards pad_cu_seqlens=self._pad_cu_seqlens. If a user enables pad_cu_seqlens with a standalone blend, all sources read the same metadata file, which (a) will not exist for externally pre-packed sources, tripping the assert in GPTSFTPackedDataset, and (b) would apply one source's global max_seqlen to all sources even if it did exist. Consider validating that pad_cu_seqlens is incompatible with packed_train_data_blend, or deriving per-source metadata.

Suggested test cases:

• No perf tests impacted.
• Unit: test_packed_train_data_blend_rejects_pad_cu_seqlens (or equivalent) covering pad_cu_seqlens=True with a standalone blend, to pin the intended behavior of the shared-metadata edge case above.

[yaoyu-33]

/ok to test f651638

[yaoyu-33]

Addressed the packed-blend metadata edge case in 542cd9874: packed_train_data_blend now rejects pad_cu_seqlens, the blend builder no longer forwards one shared metadata file to every source, and a regression test pins the behavior. Current-head EOS validation passed: 100 data tests, 36 packing-related training-config tests, 26 docs/tutorial consistency tests, and all pre-commit hooks.

[yaoyu-33]

/ok to test 542cd98

[claude]

Light review - looks good overall. This adds weighted blending across multiple pre-packed Parquet SFT sources (packed_train_data_blend), with layered validation in PackedSequenceSpecs, GPTSFTDatasetConfig.validate(), and the new GPTSFTPackedParquetBlendDataset. The logic is sound and test coverage is solid.

Verified:

• Source-selection validation splits the old has_local==has_hf check into distinct both-set and none-set errors, and admits a standalone blend only when do_validation=False and do_test=False.
• PackedSequenceSpecs rejects packed_train_data_path+blend together, non-Parquet sources, fewer than 2 sources, mismatched weights, non-positive/non-finite weights, non-finite weight sums, and pad_cu_seqlens.
• train_path_packed raises when a blend is configured, so the pre-packing path is not silently mixed with a blend.
• Weight normalization is consistent (floats in config, sum-to-1 in the dataset), matching the weights-need-not-sum-to-one docs.
• Low-discrepancy source scheduling matches Megatron blended-dataset semantics; the int16 source cap and the count message agree.
• Doc import examples resolve: GPTSFTDatasetConfig and PackedSequenceSpecs are exported from their documented modules.

Minor (non-blocking):

• tutorials/data/text-only-sft/README.md troubleshooting bullet is keyed on the string 'Exactly one text-only SFT source' but describes the neither-set case. With this PR the neither-set case raises a different message ('A text-only SFT source must be set: dataset_root, hf_dataset, or packed_train_data_blend.'), so a user searching for that message will not match this bullet. Consider updating the entry.

Suggested test cases:

• test_packed_train_data_blend_round_trip
• test_standalone_packed_train_data_blend_rejects_additional_splits
• test_packed_train_data_blend_rejects_single_path_and_blend
• test_packed_train_data_blend_validates_weights
• test_packed_train_data_blend_rejects_legacy_numpy_source
• test_packed_train_data_blend_rejects_padded_cu_seqlens
• test_config_requires_a_non_competing_source
• TestPackedParquetBlendDataset::test_weighted_source_counts_and_samples
• TestPackedParquetBlendDataset::test_default_size_and_deterministic_mapping
• TestPackedParquetBlendDataset::test_oversampling_wraps_each_source
• TestPackedParquetBlendDataset::test_negative_index_zeroes_loss_mask
• TestPackedParquetBlendDataset::test_collate_delegates_to_packed_sft_collator
• test_gpt_sft_builder_builds_weighted_parquet_blend
• Consider adding: a blend paired with dataset_root where do_validation=True builds the train split from the blend while validation comes from the local source (the documented combined-source path is untested).

No perf tests impacted.

[yaoyu-33]

Addressed the latest review suggestions in 185947b65: the tutorial now distinguishes the competing-source and missing-source validation messages, and a new integration-style unit test verifies the documented combined path where a weighted packed-Parquet blend supplies training while local JSONL supplies and prepares validation. EOS rerun: 67 relevant data/docs tests passed and all pre-commit hooks passed.

[copy-pr-bot]

/ok to test 185947b65f7b53b2d704c1d3c3f39efcd95dbcbe

@yaoyu-33, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

[yaoyu-33]

/ok to test 185947b

[yaoyu-33]

CI attribution update: run 33600044355 built the container successfully; diffusion unit tests, import check, and core shard 2 passed. Core shards 0/1 failed only in tests/unit_tests/doc_consistency/test_model_verification_catalog.py (test_generated_outputs_and_navigation_are_current and test_models_map_to_canonical_guides), so L0 jobs were skipped by dependency. This is a current-base regression: the main run for base commit 08fc1c60e (33591645933) fails the exact same two tests after the Muse Glimmer card was added without a canonical model guide/generated catalog refresh. PR #5289 does not touch the model-verification card, generator, or model-guide files. I will sync/re-run after the base is repaired rather than mixing an unrelated model-doc fix into this data PR.

[yaoyu-33]

/ok to test 227a2e7

[yaoyu-33]

/ok to test bb5e3ce

[yaoyu-33]

/ok to test 1aa87db

[yaoyu-33]

/ok to test 8f0cb6e

[github-actions]

Automatic Claude reviews have been retired. To request a pull-request review, post a comment containing:

/review

Add model=claude to use a Claude reviewer (the default is model=codex). mode=light|strict selects the review depth; for example, /review model=claude mode=strict. Comment /review help for all options.

[github-actions]

🌿 Preview your docs: https://nvidia-preview-pr-5289-8f0cb6e2e7ab.docs.buildwithfern.com/nemo/megatron-bridge