[synthetic-data] Add on-the-fly synthetic length-distribution dataset

[jinsooihm]

Stack from ghstack (oldest at bottom):

• -> [synthetic-data] Add on-the-fly synthetic length-distribution dataset #4275
• [data] TextCollator: pad positions within the context window #4274

Streaming data source for DP load-balancing and throughput experiments,
generating random-token sequences whose lengths follow a configurable
distribution. No dataset is written to disk.

Example usage:

from torchtitan.components.data import (
    BucketLengthSpec,
    FirstFitPackingConfig,
    GrainDataLoader,
    LengthBucket,
    RandomTokenProcessor,
    SingleDatasetConfig,
    SyntheticLengthSource,
)

synthetic_ds = SingleDatasetConfig(
    source=SyntheticLengthSource.Config(
        length_spec=BucketLengthSpec(
            buckets=(
                # keep max_len < training.max_context_length; longer draws are dropped
                LengthBucket(min_len=1, max_len=128, weight=3.0),
                LengthBucket(min_len=1024, max_len=2000, weight=1.0),
            ),
        ),
        seed=0,
    ),
    processor=RandomTokenProcessor.Config(),  # vocab from tokenizer
    post_filters=(lambda sample: sample is not None,),
)

# Pack whole sequences into rows, as SFT does. FirstFit preserves each sampled
# length; ConcatThenSplit would re-chunk into uniform blocks and erase the
# distribution you configured.
synthetic_packed_ds = FirstFitPackingConfig(dataset=synthetic_ds)

config.dataloader = GrainDataLoader.Config(dataset=synthetic_packed_ds, shuffle=False)

Another option is to use ConstantTokenProcessor producing same token for perhaps non-MoE models in which token routing does not matter.

scripts/preview_synthetic_lengths.py: inspect a distribution (percentiles, histogram, DP-balance simulation) for a seed without running training.

Tests: tests/unit_tests/components/data/test_synthetic_data.py and
tests/unit_tests/test_preview_synthetic_lengths.py (30 tests).

[tianyu-l]

For MoE load balance, it is

[jinsooihm]

Noob question, but is testing MoE on random tokens a fair way to test if token routing matters?

[tianyu-l]

No, I'm just saying that the note is not accurate

[jinsooihm]

Yes, I understand that. I am just curious about what kind of dataset one should use when doing perf comparisons on MoE.

[tianyu-l]

must be real dataset

[tianyu-l]

Do we run this E2E? What if loss / grad norm explode?

[jinsooihm]

Hmm, my guess would be that if the tokens are random, then the loss / grad would converge to whatever random prediction gives, which should be a finite value?

[tianyu-l]

I think loss would be. Grad norm, I'm not sure

[jinsooihm]

This is what claude tells:

The gradient of cross-entropy w.r.t. the logits is always

$$g = \text{softmax}(z) - e_y = p - e_y.$$

Both $p$ and $e_y$ are probability vectors (entries in [0,1] summing to 1), so their difference is bounded no matter what:

$$|g|_1 \le |p|_1 + |e_y|_1 = 2.$$

That bound is independent of the token/label and of vocab size — a random label gives the same $p - e_y$ form as a real one, so synthetic data can't make it bigger.

Everything downstream is just this bounded signal times the network's Jacobian. Normalization (RMSNorm) caps activation sizes and residual connections keep each layer's Jacobian ~O(1), so a bounded input stays bounded through backprop. And clip_grad_norm_(max_norm=1.0) caps the applied update at 1.0 regardless.

And test running llama3b debugmodel for 1000 steps: