Skip to content

feat(data): add indexed storage for packed SFT - #5226

Closed
yaoyu-33 wants to merge 3 commits into
mainfrom
yuya/packed-sft-indexed
Closed

yaoyu-33 wants to merge 3 commits into
mainfrom
yuya/packed-sft-indexed

Conversation

@yaoyu-33

@yaoyu-33 yaoyu-33 commented Aug 1, 2026 •

Copy link
Copy Markdown
Contributor

What does this PR do?

Use Megatron Core .bin/.idx IndexedDataset pairs as the default storage for offline-packed text SFT and PEFT data, while retaining explicit Parquet and deprecated NumPy compatibility paths.

Changelog

  • Add a versioned packed-SFT IndexedDataset schema with one pack per item, token IDs in the lower 31 bits, a target-aligned loss-mask bit, and stored sequence boundaries.
  • Route offline preparation, runtime loading, path validation, multi-shard resolution, builder-managed caching, and LoRA FLOP statistics through indexed data.
  • Change builder-managed outputs from training_<length>.idx.parquet to the logical prefix training_<length>.sft, producing .sft.bin and .sft.idx.
  • Keep explicit .parquet output and loading for migration and A/B checks; keep deprecated .npy loading behavior.
  • Add bounded stale-NFS retries, per-shard schema checks, sample-cap mapping fixes, and parity checks that detect row-count mismatches.
  • Publish local pairs under a reader/writer filesystem lock, restore the previous pair if either publication step fails, and preserve recovery backups if rollback itself fails.
  • Support prebuilt msc:// pairs with automatic MCore feature enablement, shared index-cache configuration, and remote range reads. Direct object-storage writes are rejected: prepare locally and upload .bin before .idx.
  • Add compare_packed_sft_formats.py for row-level semantic parity and sequential-read microbenchmarks.
  • Add a complete packed SFT indexed-dataset tutorial and update the text-SFT, data-preparation, and packed-sequence guides.

The .bin/.idx container is now shared with pretraining, but the packed-SFT payload remains a distinct versioned schema because SFT also needs loss masks and sequence boundaries.

Validation

  • 184 targeted unit/integration executions passed on CW against the original implementation, including end-to-end JSONL -> offline packing -> .bin/.idx -> builder load -> packed collate.
  • After multiple independent review/fix rounds, focused compatibility-isolated local runs passed: 26 indexed storage tests, 40 packing algorithm/statistics tests, 3 Parquet/indexed comparison tests, 1 end-to-end builder/sample-cap test, and 10 tutorial tests.
  • A locked multi-storage-client==0.51.0 smoke test passed for local pair generation followed by msc://default auto-enable, remote index caching/range reading, and row decode.
  • Full repository pre-commit passed after the final fixes; git diff --check, compile checks, and secret/path scans passed.
  • CW/Lustre synthetic warm-read comparison: 512 rows x 4096 tokens (2,097,152 tokens), all input_ids, loss_mask, and seq_start_id values matched in three runs. Median sequential decode was about 2.32M tokens/s for Parquet and 174.9M tokens/s for bin/idx. Storage was 5.16 MiB and 8.03 MiB respectively. This is a decoder microbenchmark, not an end-to-end training-throughput claim.

The full functional builder module previously stalled while downloading legacy Megatron-LM release assets before executing a test. Fresh post-review CW attempts did not enter the Python payload because Pyxis/container initialization stalled on multiple nodes; no test failure occurred. The final changes are therefore covered by focused local tests and still require normal NVIDIA CI validation.

GitHub Actions CI

This remains a draft PR. NVIDIA unit/L0 workflows have not run because copy-pr-bot requires additional validation before NVIDIA runners can start.

Before your PR is "Ready for review"

Pre checks:

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests?
  • Did you add or update any necessary documentation?
  • Does the PR affect components that are optional to install? No new dependency is added. Parquet remains lazily imported on Parquet-only paths; MSC must already be installed and configured for msc:// inputs.

Additional Information

@yaoyu-33 yaoyu-33 added feature New capabilities, enhancements, or enablement work area:data Dataset builders, preprocessing, and samplers labels Aug 1, 2026
@copy-pr-bot

copy-pr-bot Bot commented Aug 1, 2026

Copy link
Copy Markdown

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.

@yaoyu-33
yaoyu-33 force-pushed the yuya/packed-sft-indexed branch from 251ff14 to f0fac0b Compare August 2, 2026 16:23
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
@yaoyu-33
yaoyu-33 force-pushed the yuya/packed-sft-indexed branch from f0fac0b to 3cb860d Compare August 3, 2026 05:02
@yaoyu-33

yaoyu-33 commented Aug 4, 2026

Copy link
Copy Markdown
Contributor Author

Closing in favor of #5289, which adds weighted multi-source blending for the existing packed Parquet path.

The real NVIDIA Nemotron 3 Nano 30B-A3B SFT A/B on 8 H100s (sequence length 2048, GBS 128, four interleaved cold optimizer steps) did not show a bin/idx benefit: indexed averaged 104.679 s/step versus 103.810 s/step for Parquet. Parquet was nominally 0.83% faster, while same-format run-to-run variation was 1.1-2.6%, so the formats are effectively tied at this model scale. Batch generation took only about 3-4 ms of a roughly 104 s step.

The benchmark was limited to first complete optimizer steps because subsequent DeepEP steps hit the existing timeout; all four measured steps completed successfully with no skipped or NaN iterations. Given the lack of a material training-speed benefit and Parquet being about 51% smaller in the test dataset, the additional indexed storage format is not worth carrying.

@yaoyu-33 yaoyu-33 closed this Aug 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:data Dataset builders, preprocessing, and samplers feature New capabilities, enhancements, or enablement work

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant