lancedb/docs/web/training/distributed-training.mdx

---
title: "Distributed Training"
sidebarTitle: "Distributed Training"
description: "Shard StreamingDataset across ranks and workers, with elastic determinism across topologies."
icon: server
availability:
  embedded: available
  enterprise: available
---

To understand the wide range of scenarios handled by [`StreamingDataset`](/training/data-loading) 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](/training/data-loading).

## Parallelism & Elasticity {#parallelism-and-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.

<CodeGroup>
```py Python icon=Python
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)
```
</CodeGroup>

## Next steps {#next}

<CardGroup cols={2}>
  <Card title="Checkpointing" icon="floppy-disk" href="/training/checkpointing">
    Save progress so a failed run can pick back up.
  </Card>
  <Card title="Resuming" icon="rotate-left" href="/training/resuming">
    Load a saved checkpoint back, on the same or a different topology.
  </Card>
</CardGroup>