WebDataset

WebDataset uses .tar files as shards. The format follows two main conventions:

1. Sample Grouping: Files belonging to the same training sample share the same basename when extensions are stripped (e.g., sample1.json and sample1.png form one sample).
2. Shard Numbering: Shards are typically numbered sequentially (e.g., data-000000.tar to data-000010.tar), often referenced using brace notation like data-{000000..000010}.tar.

This format allows for purely sequential I/O, which is highly efficient for large-scale deep learning on local disks or cloud object stores.

To detect when a tarball has been fully consumed (especially for remote files using io.BytesIO), you can add a metadata field like __index_in_shard__ to your samples. When this index is 0, it indicates the start of a new shard, which can be used to monitor consumption or handle single-sample tars.

# Example of adding metadata to track consumption
sample['__index_in_shard__'] = index_in_shard
if sample['__index_in_shard__'] == 0:
    print("Last shard fully consumed.")

To process data in parallel for large-scale WebDataset construction, use Apache Beam's distributed processing. The recommended pattern is to handle each shard in parallel while processing the contents of each shard sequentially. You can implement a custom beam.DoFn to write processed data to a WebDataset tar file using ShardWriter.

import apache_beam as beam

class ProcessAndWriteToWebDataset(beam.DoFn):
    def process(self, element):
        # Process your data
        processed_sample = ...  # Your processing logic here
        # Write to WebDataset
        with ShardWriter('output.tar') as writer:
            writer.write(processed_sample)

with beam.Pipeline() as pipeline:
    (pipeline
     | 'ReadData' >> beam.io.ReadFromSource(...)
     | 'ProcessAndWrite' >> beam.ParDo(ProcessAndWriteToWebDataset()))

When using WebDataset to store small embeddings or output classes, you may encounter significant space and performance overhead due to individual serialization. To improve efficiency:

• Use efficient formats: Store data in .npy or .ten formats.
• Reduce precision: Use 8-bit integers or float16 instead of full precision.
• Batch data: Group multiple samples into a single file to reduce per-record overhead.
• Leverage caching: Use in-memory or on-disk caching to speed up reads.
• Minimize IPC: Reduce DataLoader Inter-Process Communication (IPC) overhead, which is particularly impactful for small records.

When using WebDataset in Distributed Data Parallel (DDP) training, you can improve data randomness and manage epoch lengths by applying transformations to the WebLoader rather than the WebDataset object.

To implement cross-worker shuffling and define a specific epoch length, use the following pattern on your WebLoader:

1. Use .unbatched() to operate on individual samples.
2. Use .shuffle(buffer_size) to shuffle samples in memory.
3. Use .batched(batch_size) to regroup samples into batches.
4. Use .with_epoch(num_batches) to define the epoch length in terms of the number of batches.

When using Fully Sharded Data Parallel (FSDP), you must provide an explicit nodesplitter to the WebDataset class to manage shard distribution across nodes. Use wds.split_by_node to ensure each node receives a distinct subset of the data.

import webdataset as wds

dataset = (
    wds.WebDataset(shard_urls, resampled=True, cache_dir=data_args.local_cache_path, nodesplitter=wds.split_by_node)
    .shuffle(training_args.seed)
    .map(decode_text)
    .map(TokenizeDataset(tokenizer, max_seq_len=data_args.max_seq_length))
)

While wds.MultiDataset is an experimental alternative to torch.utils.data.DataLoader, it is not a direct replacement.

MultiDataset offers advantages in containerized environments and provides simpler shard assignment and sample splitting among workers. However, when using frameworks like PyTorch Lightning, using MultiDataset as a dataloader may cause training to stall; in those cases, use torch.utils.data.DataLoader instead.

If you prefer traditional epoch-based training in a distributed environment, use indexable datasets combined with wids.DistributedChunkedSampler. This ensures balanced data distribution across nodes and prevents duplicates by ensuring each node processes a unique subset of data.

# Example for indexable dataset with DistributedChunkedSampler
sampler = wids.DistributedChunkedSampler(dataset, num_replicas=world_size, rank=rank)
dataloader = DataLoader(dataset, sampler=sampler, batch_size=batch_size)