[Feature] Add --streaming_shard: split the streaming dataset across data-parallel ranks

[Rapisurazurite]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

Adds an opt-in --streaming_shard that splits the streaming training set across data-parallel ranks, so each rank preprocesses only its own share instead of rank 0 preprocessing everything and scattering the batches.

Problem

Under --streaming true, DataLoaderDispatcher has only rank 0 iterate the dataset; it runs the whole cluster's tokenization / image encoding and then scatters one batch per rank through scatter_object_list (swift/dataloader/dispatcher.py). Because the encode is inside the streaming dataset (IterablePackingDataset, or EncodePreprocessor's lazy map), global preprocessing throughput does not depend on world size — adding ranks adds no data-supply capacity.

This is already documented (docs/source_en/Instruction/Command-line-parameters.md, --streaming: "preprocessing is only performed on rank 0 and then distributed to other processes ... can become a bottleneck") and reported in #6758 and #1939.

Two details make it sharper:

• Without packing there is no parallelism at all: num_proc is only passed to map for a non-streaming HfDataset (swift/dataset/preprocessor/core.py), so the entire cluster encodes in one Python process.
• With packing, every rank still spawns dataset_num_proc encoder processes, but only rank 0's are ever fed; the rest idle for the whole run.

Design

The split is by blocks of consecutive raw samples, sized to the amount of raw data the loader consumes per item it emits — the per-device batch size, or packing_interval when packing. This is the same block rule accelerate's IterableDatasetShard and Megatron-LM's MegatronPretrainingSampler use for one-to-one samples. Because the default path already hands rank r the r-th batch of every round, without packing the per-rank sample sequence is identical when preprocessing is deterministic and maps every raw sample to one output (test below). Samples dropped by truncation or bad media, expanded by truncation_strategy=split, or randomly augmented can change the encoded assignment.

With packing the guarantee is deliberately weaker: rank 0 scatters individual packs rather than whole packing rounds, so matching packing_interval only lines the raw packing windows up. Each rank rebuilds its own packs and carries the leftover bin locally; pack grouping, pack counts, and finite-stream tail coverage can differ from dispatch.

Rank, world size and block size are bound by the dataloader, not at dataset construction time, because that is the only layer that knows which group the data is parallelized over: Megatron-SWIFT initializes its parallel state after the dataset is prepared, and under sequence parallel the relevant group is the dp group rather than WORLD. Binding inside DataLoaderDispatcher means DDP, sequence parallel and Megatron-SWIFT all share one code path, via the rank / world_size / group properties SequenceParallelDispatcher and MegatronDataLoaderDispatcher already override — no fallback branch is needed, and ranks in the same TP/PP group naturally get the same shard.

No new class is introduced: the sharded iteration is a second path inside DataLoaderDispatcher, selected by whether a shard state was attached. The default path is unchanged.

Because each rank consumes its own shard, ranks can run out at different times, so a one-byte all_reduce(MAX) per step implements a first-exhausted policy and stops them together. This avoids a distributed hang, but when filtering or packing makes output batch counts uneven it can leave an unbounded tail on other ranks. Dataset shuffling normally reduces systematic imbalance, and step-based runs that reach max_steps before exhaustion are unaffected. For reference, neither of accelerate's non-dispatch paths gives us this: .shard() splits the source but has no cross-rank stop at all, and IterableDatasetShard gets equal batch counts by construction only because every rank materializes every rank's samples, which is exactly the cost this PR avoids. Note also that the current dispatcher gets lockstep termination for free by scattering None, so sharding must not silently drop it.

sequenceDiagram
    participant S as Raw stream
    participant R0 as Rank 0
    participant RN as Other DP ranks
    R0->>S: Read assigned raw blocks
    RN->>S: Read assigned raw blocks
    R0->>R0: Encode and pack locally
    RN->>RN: Encode and pack locally
    R0->>RN: Synchronize exhausted flags
    alt All ranks have a batch
        R0->>R0: Yield local batch
        RN->>RN: Yield local batch
    else Any rank is exhausted
        R0-->>RN: Stop iteration together
    end

Loading

Open design questions — feedback welcome

The current patch intentionally keeps the smallest policy surface, but I would appreciate maintainer guidance on two follow-ups:

1. Exhaustion policy. The patch uses first_exhausted because it does not repeat data and preserves the lockstep termination guarantee of the existing dispatcher. A torchtitan-style infinite policy would remove the per-step collective and avoid truncating longer shards, but shorter shards would be repeatedly sampled and true epoch semantics would disappear. Would you prefer an explicit first_exhausted | infinite option, or keep the current single policy?
2. Packing assignment. block_size=packing_interval aligns raw packing windows, but independently rebuilt packs can still have very different counts. Alternatives are sample-stride assignment for better statistical balance, or initially limiting streaming_shard to non-packing data and handling packing in a follow-up. Is preserving the current packing support useful despite the documented first-exhausted semantics?

Compatibility and limits

• Defaults to False; when --streaming false, it is reset with an info log. The existing dispatch path is unchanged.
• All ranks must iterate the raw stream in the same deterministic order; the built-in loader uses the same shuffle seed on every rank.
• Training set only; the validation set keeps the existing path. GRPO and GKD are rejected because their encoding is deferred past dataset post-processing.
• Rejected together with --deepspeed_autotp_size: the streaming path does not replicate a batch across tensor-parallel ranks. Note this predates this PR and also affects the default path — BatchSamplerShard is TP-aware via tp_size but DataLoaderDispatcher uses WORLD. Happy to fix that separately if you agree it is a bug.
• Cost: every rank reads the raw stream in full; the template encoding and streaming packing stages are split. Expensive custom dataset preprocessing or parquet with inlined image bytes should be evaluated first, especially multi-node.
• With packing, every rank now uses its dataset_num_proc encoders instead of only rank 0, so per-node CPU/memory can grow with (ranks per node) x dataset_num_proc. Without packing, streaming encoding remains single-process per rank.

Tests

The test file was reduced from eight process-spawning tests (about 4 minutes) to two focused tests (49.93 seconds on the project container):

• test_lazy_block_assignment: public filter(with_indices=True) remains unsharded until the dataloader binds rank/world size/block size, then selects the expected blocks.
• test_distributed_streaming_shard: one 2-rank process group covers deterministic one-to-one equivalence (including a partial tail and num_workers=1), local resume skipping, and deliberately imbalanced packing. The packing assertion records the chosen first-exhausted policy explicitly: both ranks emit one pack, while the long-sample rank consumes 1 raw sample and the short-sample rank consumes 128; no bounded-tail guarantee is claimed.

flake8 / yapf / isort are clean with the versions pinned in .pre-commit-config.yaml.

Data-pipeline benchmark (no model)

Multi-node pipeline throughput (2 nodes x Ascend 910B, hccl, 50 ms/sample encode, no packing). This isolates the supply side, with no compute floor hiding it:

ranks	default (dispatch)	--streaming_shard true	speedup
2 x 8 = 16	20.17 samples/s	319.14 samples/s	15.8x
2 x 16 = 32	20.19 samples/s	634.88 samples/s	31.4x

Doubling the ranks leaves the default path exactly flat (20.17 -> 20.19; per-rank rate halves, 1.26 -> 0.63 steps/s) while the sharded path doubles. Both match theory: one process at 50 ms/sample is 20 samples/s, and 16 / 32 x 20 = 320 / 640 vs 319.14 / 634.88 measured.