Distributed Training

Applies to: Embedded ✅, Enterprise ✅ To understand the wide range of scenarios handled by StreamingDataset across a distributed job, consider a model trained with stochastic gradient descent (SGD) and distributed data parallelism (DDP). In this example, we need to load batches onto multiple GPUs across multiple servers. After each batch is processed, the GPUs exchange weights and the next batch is loaded. We fix the following terminology:

World size - The number of GPUs being loaded simultaneously. For example, if we have $2$ servers with $4$ GPUs each, the world size is $8$ .
Rank - The identifier of an individual GPU. It is an integer in the range [0, world_size). Each rank is allocated a disjoint portion of the source data.
Batch size - The number of rows processed by one GPU in a single training step. For example, $128$ .
Global batch size - The total number of rows processed across all GPUs in a training step (sometimes called a “minibatch” in machine learning literature). For example, if we have $8$ GPUs with batch size $128$ , the global batch size is $1024$ .

Other concepts, such as read batch size and num_workers, are introduced in PyTorch Data Loader.

Parallelism & Elasticity

When loading data, StreamingDataset partitions the source dataset among disjoint parallel threads called splits. Each rank is allocated a contiguous group of splits, and each DataLoader worker associated to a rank gets a contiguous subgroup of its rank’s splits. If using shuffle=False, data is allocated to splits naively in ordered contiguous blocks of equal size. With shuffle=True (the default), the rows of the source dataset are shuffled and distributed among splits such that each split receives the same number of rows; with shuffle_clump_size set, they are shuffled in clumps of that many contiguous rows. The number of splits defaults to world_size if num_splits is omitted during initialization. This is not recommended, as a later run using a different world size will then construct a different layout and structure of global batches. To get elastic determinism, choose a fixed num_splits that is compatible with every topology you plan to use. For example, a run with num_splits=8 and

4

GPUs assigns

2

splits to each rank. The ranks pull from their splits in round-robin order, producing the same global batches as a run with

8

GPUs and the same num_splits, shuffle_seed, and epoch. The following constraints apply:

num_splits must be divisible by world_size.
With DataLoader worker processes, num_splits must be divisible by world_size * num_workers.
For the same samples to form each global training step across topologies, global_batch_size must be a multiple of num_splits.
The filtered row count must be at least num_splits. If it is not evenly divisible by num_splits, up to num_splits - 1 surplus rows are dropped so that every split has the same length.

For example, to switch between

8

and

6

GPUs, use a common multiple such as

24

for num_splits. Highly composite values can support many layouts; for example, num_splits = 48 is compatible with

10

distinct possible GPU counts. Remember to include num_workers when checking divisibility.

import torch.distributed as dist
from lancedb.streaming import StreamingDataset

dist.init_process_group("nccl")
rank = dist.get_rank()
world_size = dist.get_world_size()

# num_splits=48 is divisible by many integers,
# so this dataset works unchanged as you add or remove GPUs.
ds = StreamingDataset(
    table,
    num_splits=48,
    shuffle_seed=42,
    epoch=current_epoch,
    rank=rank,
    world_size=world_size,
)

for sample in ds:
    train_step(sample)

Introduction

Getting Started

Concepts

Data Ingestion

Search

Data Curation

Feature Engineering

Training

Management

Maintenance

Security

Deployment

SQL Reference

Distributed Training

Parallelism & Elasticity

Next steps

Checkpointing

Resuming

​Parallelism & Elasticity

​Next steps

Checkpointing

Resuming

Parallelism & Elasticity

Next steps