Data parallelism

Overview

Data parallelism is the simplest way to train a neural network on many GPUs. Every device holds a full copy of the model. Each device reads a different slice of the batch, computes gradients on its slice, and then all devices average their gradients so that every copy takes the same update step. The copies stay identical, and the training runs faster because more examples are processed per step.

It is the first tool most teams reach for because it needs almost no change to the model code. Its main limit is memory: every device must hold the whole model, its gradients, and its optimizer state. When a model no longer fits on one GPU, data parallelism is combined with, or replaced by, sharding and model-splitting methods such as tensor parallelism and pipeline parallelism. For tabular foundation models, which train on huge numbers of small tables, data parallelism is often the main engine of scale.

What it is

Data parallelism is a way to split the work of training along the data axis. The model is replicated. The data is divided. The results are merged.

Picture a team of proofreaders marking up a long book. Each proofreader has the same style guide, which plays the role of the model. The book is cut into sections, and each person works on one section. At the end of the day they meet, compare notes, and agree on a single updated style guide. Everyone starts the next day with the same guide, but each reads new sections.

The core ingredients

• Replicas. Identical copies of the model, one per device (often called workers or ranks).
• Sharded input. The global batch is divided into smaller local batches, one per replica.
• Gradient synchronization. Gradients from all replicas are combined, usually by averaging.
• Synchronized update. Every replica applies the same averaged gradient, so the copies never drift apart.

Synchronous and asynchronous variants

Almost all modern large-scale training uses the synchronous form, because its behavior is easy to reason about and it reproduces the result of a single large batch.

Where it came from

The idea is as old as parallel computing. Split the data, run the same program on every piece, and combine the answers. Early large neural network systems used a central parameter server that held the weights while many workers pulled the weights and pushed gradients. These systems showed that scale helped, but the central server was a bottleneck.

Two things changed the picture. First, research on large-batch training showed that with a suitable learning rate schedule, such as scaling the rate with the batch size and warming it up, models could be trained with very large batches without losing accuracy. A well-known 2017 report from a large research group trained an image classifier in about an hour using this approach. Second, fast collective communication, especially the ring all-reduce algorithm and libraries such as NCCL, removed the central server. Workers talk to their neighbors directly, and the cost of averaging gradients stays nearly constant as workers are added.

Framework support followed. Libraries in PyTorch, TensorFlow, and JAX now offer data parallel wrappers, and higher-level tools such as PyTorch Lightning hide most details. Later work, including the ZeRO family of methods, kept the data parallel programming model but split the model state across workers to save memory. See DeepSpeed ZeRO and FSDP.

How it works

A single training step in synchronous data parallelism has six stages.

Step 1: Start with identical replicas

Before training, one copy of the weights is broadcast to all workers, or every worker is initialized with the same random seed. From this moment, all replicas hold exactly the same numbers.

Step 2: Split the batch

The data loader gives each worker a different, non-overlapping slice. A distributed sampler is used so that each example is seen by one worker per epoch. If the global batch is 1024 examples and there are 8 workers, each worker handles 128.

Step 3: Forward pass on the local slice

Each worker runs the model forward on its own slice and computes a local loss. No communication is needed yet. The work is fully parallel.

Step 4: Backward pass on the local slice

Each worker computes gradients of its local loss. These are local gradients, and they differ from worker to worker because the data differs.

Step 5: All-reduce the gradients

Now the workers must agree. An all-reduce operation sums the gradients across all workers and gives every worker the result. Dividing by the number of workers gives the average.

The most common algorithm is the ring all-reduce. Workers are arranged in a ring. Each gradient tensor is cut into chunks. In one phase, chunks are passed around the ring and added up, so each worker ends up owning the full sum for one chunk. In a second phase, the finished chunks are passed around again so everyone gets all of them. Each worker sends and receives roughly twice the size of the gradient, no matter how many workers there are. This is why the method scales so well.

A daily-life analogy: a group of friends wants to split a restaurant bill. Instead of sending everything to one person, each friend passes a partial tally to the next friend around the table. After one lap everyone knows the total.

Step 6: Apply the same update

Each worker uses the averaged gradient to update its own copy of the weights with the optimizer. Because the gradients, the optimizer state, and the starting weights are identical, the new weights are identical too.

Overlapping communication with computation

The backward pass produces gradients layer by layer, from the last layer to the first. Good implementations start the all-reduce for the last layers while the earlier layers are still computing. Gradients are grouped into buckets so that each message is large enough to use the network well. This overlap hides much of the communication time.

Why the math matches a single big batch

Averaging local gradients over equal-sized slices gives exactly the gradient of the average loss over the whole batch. So data parallel training with 8 workers and a local batch of 128 is mathematically the same as one worker with a batch of 1024. This is the reason the method is so widely trusted.

What stays the same and what changes

• Same: the model, the loss, the optimizer, and the final accuracy target.
• Changes: the global batch size grows with the number of workers, the learning rate usually needs adjusting, and some layers, such as batch normalization, see only the local slice.

Key concepts and terms

A worked example

A numeric walk-through

Suppose a model has a single weight w = 2.0. There are four examples in the global batch and two workers. The loss for an example is (w - t)^2, where t is the target, so its gradient with respect to w is 2(w - t).

The targets are 1, 3, 5, and 7.

• Worker A gets targets 1 and 3. Its gradients are 2(2 - 1) = 2 and 2(2 - 3) = -2. The local average is 0.
• Worker B gets targets 5 and 7. Its gradients are 2(2 - 5) = -6 and 2(2 - 7) = -10. The local average is -8.

All-reduce averages the two local averages: (0 + -8) / 2 = -4.

A single worker with all four examples would compute (2 - 2 - 6 - 10) / 4 = -4. The numbers match. With a learning rate of 0.1, both workers update w to 2.0 - 0.1 * (-4) = 2.4. Both copies agree.

A PyTorch example

The following sketch uses the standard distributed data parallel wrapper. It is launched with one process per GPU.

import os
import torch
import torch.distributed as dist
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, TensorDataset
from torch.utils.data.distributed import DistributedSampler

def main():
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    model = nn.Sequential(nn.Linear(64, 256), nn.ReLU(), nn.Linear(256, 1))
    model = model.cuda(local_rank)
    model = DDP(model, device_ids=[local_rank])

    x = torch.randn(10_000, 64)
    y = torch.randn(10_000, 1)
    ds = TensorDataset(x, y)
    sampler = DistributedSampler(ds, shuffle=True)
    loader = DataLoader(ds, batch_size=128, sampler=sampler)

    opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
    loss_fn = nn.MSELoss()

    for epoch in range(3):
        sampler.set_epoch(epoch)
        for xb, yb in loader:
            xb = xb.cuda(local_rank, non_blocking=True)
            yb = yb.cuda(local_rank, non_blocking=True)
            opt.zero_grad(set_to_none=True)
            loss = loss_fn(model(xb), yb)
            loss.backward()
            opt.step()

    dist.destroy_process_group()

if __name__ == "__main__":
    main()

Launch it with a command such as torchrun --nproc_per_node=8 train.py. The wrapper hooks into the backward pass and runs the all-reduce for you.

Points to notice

• DistributedSampler gives each rank a distinct slice, and set_epoch changes the shuffle every epoch.
• The batch size in the loader is the local batch. The global batch is this value times the world size.
• Saving a checkpoint is normally done by rank zero only, since all replicas are identical.

Gradient accumulation to enlarge the batch

When memory limits the local batch, gradients from several micro-batches can be added before one update. With DDP, communication should be skipped on the intermediate micro-batches to avoid wasted work.

accum = 4
for i, (xb, yb) in enumerate(loader):
    sync = (i + 1) % accum == 0
    ctx = nullcontext() if sync else model.no_sync()
    with ctx:
        loss = loss_fn(model(xb), yb) / accum
        loss.backward()
    if sync:
        opt.step()
        opt.zero_grad(set_to_none=True)

The effective global batch is then local batch times world size times the accumulation count.

Estimating the cost of one step

Total step time in synchronous data parallelism can be pictured as the sum of three parts.

• Compute time. Forward and backward work on the local slice.
• Communication time. The all-reduce of all gradients.
• Waiting time. Idle time caused by uneven workers.

If communication is fully overlapped with the backward pass, only the part that sticks out past the end of computation adds to the step. A larger local batch raises compute time without raising communication time, because the gradient size depends on the model, not on the batch. This is why scaling efficiency is higher with larger local batches and lower with tiny ones.

A simple rule follows from this. To keep efficiency high, make the local batch as large as memory allows, then add workers. Adding workers while shrinking the local batch makes communication a bigger share of each step.

Memory arithmetic for replication

Every replica carries the same four items, so the memory cost does not shrink when workers are added.

This table explains why plain data parallelism hits a wall. Only activations respond to the number of workers, and only if you shrink the local batch. Sharded variants change the first three rows so that they shrink with the number of workers. See FSDP.

Large-batch training and the learning rate

When the global batch grows, one update uses more examples, so the gradient is less noisy. That allows a larger step. Common practice includes the following.

• Scale the learning rate with the global batch for moderate increases.
• Warm up the learning rate over the first steps so the model is not shocked by a large step.
• Watch for a point of diminishing returns, beyond which a larger batch no longer reduces the number of steps needed.
• Consider optimizers designed for very large batches if the batch is extreme.

The critical batch size depends on the task. Noisy tasks with many different cases, such as meta-training over many varied tables, tend to tolerate larger batches than narrow, clean tasks.

A launch configuration example

Cluster jobs are usually described in a short file. The following YAML is a generic sketch of what such a file states.

job:
  name: tabular-pretrain
  nodes: 2
  gpus_per_node: 8
  launcher: torchrun
  entrypoint: train.py
training:
  local_batch_tables: 16
  grad_accum_steps: 2
  precision: bf16
  base_lr: 0.0003
  warmup_steps: 2000
data:
  source: synthetic
  seed_per_rank: true
checkpoint:
  every_steps: 1000
  save_from_rank: 0

With these numbers the world size is 16, the local batch is 16 tables, and the global batch is 16 times 16 times 2, which is 512 tables per update.

Fault tolerance and elastic training

Large jobs run for days, and hardware fails. Three ideas keep them safe.

1. Frequent checkpoints. Save the model, optimizer, scheduler, and random number generator state.
2. Elastic launchers. These restart the group when a worker dies, and can continue with a different number of workers if needed.
3. Health checks. Detect slow or faulty devices early so they can be replaced.

When the number of workers changes, the global batch changes too. Adjust the accumulation steps or the learning rate so that the training behavior stays close to the original plan.

Determinism and debugging

Distributed runs are harder to reproduce than single-device runs. The order of floating-point additions in an all-reduce can differ, which causes tiny numeric differences. Useful habits:

• Fix all random seeds, and give them clear per-rank rules.
• Log the loss on each rank for the first few steps to catch a rank that has diverged.
• Compare the weights across ranks every so often. They should be identical.
• Test with two workers on one machine before launching on a large cluster.
• Use a tiny model and a tiny dataset to check that the loss decreases identically to a single-device run.

Evaluating a distributed run

Evaluation also runs on many workers, and mistakes here are subtle.

• Shard the validation set across workers to save time.
• Gather the raw counts, such as correct predictions and total examples, from all workers.
• Sum the counts and divide once, rather than averaging per-worker ratios.
• Make sure padding or duplicated samples added by the sampler are removed, or the metric will be slightly off.
• Report the metric from a single rank to avoid duplicate logs.

For ranking and probabilistic metrics, gather the predictions to one rank and compute the metric there, since these metrics do not decompose into simple sums.

When data parallelism is not enough

Watch for these signs that it is time to add another kind of parallelism.

• The model plus optimizer state no longer fits on one device, even with a small local batch.
• Activations do not fit, even with checkpointing and a local batch of one.
• Communication time exceeds compute time and cannot be hidden.
• A single layer is too large for a device.

Each sign points to a different fix: sharding for the first, checkpointing or sequence splitting for the second, bigger buckets or faster links for the third, and tensor parallelism for the fourth.

Strengths and limitations

Strengths

• Simple to understand and to implement.
• Needs almost no changes to model code.
• Gives close to linear speedup when communication is hidden behind computation.
• Mathematically equivalent to one large batch, so results are easy to reason about.
• Works with any model architecture and any optimizer.
• Mature, well-tested tooling exists in every major framework.
• Combines well with mixed precision.

Limitations

• Every device must hold the full model, gradients, and optimizer state, so memory is the main wall.
• Communication cost grows with model size, since the whole gradient is synchronized each step.
• Very large global batches can hurt generalization unless the learning rate schedule is tuned.
• The slowest worker sets the pace for everyone.
• A single failed worker can stop the whole job unless elastic training is set up.
• Batch-dependent layers see only the local slice unless they are synchronized.
• Network speed between nodes can limit scaling, especially without fast interconnects.

How it compares with alternatives

In practice the methods are layered. A common large-scale layout uses tensor parallelism inside a node, pipeline parallelism across a few nodes, and data parallelism across groups of those. Libraries such as Megatron-LM implement this combined layout. The product of the three degrees equals the total number of devices.

Choosing between replicate and shard

If the model, its gradients, and its optimizer state fit on one GPU with room for activations, plain data parallelism is simplest and fastest. If not, move to sharded data parallelism before reaching for harder model-splitting methods. Add tensor or pipeline parallelism only when sharding alone cannot fit the model or when communication becomes the limit.

Role in tabular foundation models

Tabular foundation models learn from raw tables: rows and columns, mixed types, missing values, and often no useful column names. They are pre-trained on large numbers of synthetic or real tables and then applied through in-context learning, where labeled rows are shown at inference time and the model predicts new rows with no weight update. Data parallelism fits this workload unusually well. See /model/architecture for how such models are built.

Many small tables make a natural batch

A typical pre-training example for a model like TabPFN is one whole table: a set of labeled training rows plus query rows. Tables are small compared with the text documents used to train language models. A single table fits easily on one GPU.

That means each worker can process its own tables independently. The global batch is simply a collection of different tables, and the gradient average runs over them all. There is no need to split a single table across devices. This is the sweet spot of data parallelism.

Fresh synthetic data on every worker

When pre-training on synthetic tables, the data is generated on the fly. Each worker can create its own tables from its own random seed. There is no shared dataset to read from disk and no sampler needed. As long as the seeds differ per rank, the workers see different tables and the average gradient is an unbiased estimate.

Practical points:

• Give each rank a distinct seed for table generation, and keep the model initialization seed identical.
• Keep the mix of table sizes similar across workers, so no worker is a straggler.
• Log the generator settings so a run can be reproduced.

Variable table shapes and stragglers

Tables differ in the number of rows and columns. A worker handed a large table takes longer than one handed a small table, and synchronous training waits for the slowest. Good practice is to bucket tables by size, or to sample tables so that each worker's total number of cells per step is roughly equal. The same idea appears in sequence models that group similar lengths together.

Mixed types, missing values, and no column names

Each worker encodes its tables with the same rules: numbers are scaled, categories are mapped to learned or hashed embeddings, missing cells receive a mask or a special token, and column names are either ignored or treated as optional hints. Because the encoders are part of the replicated model, all workers apply identical logic, and the gradient average updates a single shared encoder. The model learns to treat columns by their statistics and relationships, not by their names.

Row and column attention across devices

Models that attend over rows and over columns, as in axial attention, keep a full table on one device and run both attention passes locally. Data parallelism then scales the number of tables per step, not the size of any table. When tables grow so large that one cannot fit, teams add sharding or sequence splitting on top, but for most tabular pre-training this is not needed.

Masked cell modeling and imputation

In masked cell modeling, cells are hidden and the model predicts them. This objective is naturally parallel: each worker masks cells in its own tables, and the loss is averaged. The same recipe trains imputation networks. Different workers can use different masking rates and patterns, which adds variety to each global batch and teaches the model to handle missing data of many kinds.

In-context learning and meta-training

Pre-training for in-context learning is a form of meta-learning: each table is a small task, and the model learns to solve a new task from examples shown in context. Data parallelism treats each task as one sample in the batch. A larger global batch means more tasks per update, which gives a smoother gradient across task types. The cost is that very large batches may need a learning rate adjustment. See meta-learning and few-shot learning.

Continual learning

When a tabular model is updated on new tables without forgetting old ones, data parallelism still applies, but care is needed. Replay buffers, as in experience replay, should be sampled consistently across ranks, or each rank should hold its own buffer. Penalty-based methods such as elastic weight consolidation keep an importance estimate for each weight. This estimate must be identical on every replica, so it is computed once and broadcast, or averaged with an all-reduce.

Anomaly detection and sector classification

Training heads for anomaly detection or for semantic and sector classification, such as semantic sector heads, follows the same pattern. Each worker computes the loss on its tables, and gradients are averaged. For evaluation, per-worker metrics should be gathered and combined correctly, for example by summing counts and then dividing, not by averaging per-worker averages when worker sizes differ.

Benchmarking

Reports of speed and quality should state the world size and the global batch. See /model/benchmarks and /model/research.

Practical guidance and best practices

Plan the batch and the learning rate

• Decide the global batch first, then divide by the world size to get the local batch.
• Scale the learning rate with the global batch, commonly in proportion for small increases and more gently for large ones.
• Use a warm-up period so the first steps with a large batch stay stable.
• Test the schedule on a small run before scaling to many devices.

Keep the replicas in sync

• Use the same initial weights, or broadcast from rank zero.
• Use the same optimizer settings on every rank.
• Use a distributed sampler for real datasets.
• Seed the data generation differently per rank, and the model identically.

Make communication cheap

• Use a fast collective library such as NCCL with the right interconnect.
• Group gradients into buckets of suitable size.
• Overlap the all-reduce with the backward pass.
• Send gradients in 16-bit when it is safe. See mixed precision.
• Place workers on the same node when you can, since links inside a node are faster than links between nodes.

Reduce memory pressure

• Use gradient checkpointing to trade compute for activation memory.
• Move to FSDP or DeepSpeed ZeRO if the model state is too large to replicate.
• Use gradient accumulation to reach a large global batch with small local batches.

Handle checkpoints and failures

• Save a checkpoint from a single rank, and load it on all ranks.
• Save the optimizer state and the random number generator state for exact resumes.
• Use elastic launchers that can restart a job after a worker failure.
• Save often enough that a failure costs little.

Measure scaling efficiency

• Record throughput in examples (or cells) per second for 1, 2, 4, and 8 workers.
• Divide the measured speedup by the ideal speedup to get efficiency.
• If efficiency drops, profile to see whether the cause is communication, stragglers, or data loading.

Common mistakes

• Forgetting to set the epoch on the sampler. Every epoch then uses the same shuffle order, which hurts training.
• Using the same data seed on every rank. All workers then see identical data, and the extra devices add no information.
• Not scaling the learning rate. A larger global batch with an unchanged rate often trains more slowly per example.
• Averaging metrics wrongly. Averaging per-worker averages is wrong when workers have different counts.
• Saving checkpoints from every rank. This wastes disk and risks write conflicts.
• Ignoring stragglers. One slow worker, caused by a large table or a busy disk, slows the whole group.
• Running batch normalization on tiny local batches. Statistics become noisy. Use synchronized batch normalization or a batch-independent normalization such as layer normalization.
• Assuming linear scaling forever. Past a point, communication and large-batch effects reduce the gain.
• Skipping a single-device test. A model that does not train on one device will not train on eight.
• Leaving unused parameters in the graph. Wrappers may wait forever for gradients that never arrive, or need a flag to handle them.

Frequently asked questions

What is data parallelism in deep learning?

Data parallelism is a method of distributed training in which every device holds a full copy of the model and processes a different slice of each batch. After computing local gradients, the devices average them so every copy takes the same update. The result matches training on one large batch, but it finishes faster because more examples are handled at once.

How is data parallelism different from model parallelism?

Data parallelism splits the data and replicates the model, while model parallelism splits the model itself across devices. Data parallelism is simpler and works whenever the model fits on one device. Model parallelism, including tensor parallelism and pipeline parallelism, is needed when the model is too large for one device.

What is an all-reduce?

An all-reduce is a communication step in which every worker contributes a value, the values are combined (usually summed), and every worker receives the combined result. In data parallelism it is used to sum gradients. The ring version passes chunks around a ring of workers, and its cost per worker stays roughly constant as workers are added.

Does data parallelism change the final accuracy?

In principle it does not, because averaging gradients over equal slices gives the same gradient as one big batch. In practice the global batch grows, which can change optimization behavior. Adjusting the learning rate and using warm-up usually restores the same quality as small-batch training.

Why does my multi-GPU run not scale linearly?

Several things limit scaling: time spent in communication that is not hidden by computation, slow workers that hold up the group, slow data loading, and small per-device batches that leave the GPU underused. Profiling a run at different world sizes shows which of these dominates. Faster interconnects, larger buckets, and balanced work usually help.

When should I use FSDP or ZeRO instead of plain data parallelism?

Use them when the model, gradients, and optimizer state do not fit on one device, or when memory limits the batch size. They keep the data parallel programming model but split the model state across workers, trading extra communication for much lower memory. See FSDP and DeepSpeed ZeRO.

How does data parallelism help tabular foundation models?

Tabular foundation models are pre-trained on many small tables, and one table fits on a single GPU. Each worker can process different tables, often generated on the fly, and the gradients are averaged. This gives more tasks per update and faster pre-training without any need to split a table across devices.

Do I need the same data on every GPU?

No. Each GPU should see different data, otherwise the extra devices add no new information. The model weights, however, must be identical at the start and are kept identical by applying the same averaged gradient at every step.

What is the global batch size?

The global batch size is the total number of examples (or tables) processed in one optimizer step across all workers. It equals the local batch size times the number of workers times the number of accumulation steps. It is the number that matters for the optimization behavior and the learning rate.

Can data parallelism be combined with mixed precision?

Yes, and it usually is. Lower precision cuts memory and speeds up computation, and 16-bit gradients can reduce communication volume. Keep master weights and optimizer state in 32-bit. See mixed precision for details.

Glossary

• All-gather: A collective that collects pieces from every worker and gives all workers the full set.
• All-reduce: A collective that combines values across workers and shares the result.
• Bucket: A group of gradients communicated together.
• DDP: A common name for distributed data parallel training wrappers.
• Global batch: The total number of examples in one optimizer step across all workers.
• Gradient accumulation: Summing gradients over several micro-batches before one update.
• Local batch: The slice of data one worker processes.
• Rank: The index of a worker within the group.
• Replica: A full copy of the model on one device.
• Scaling efficiency: Measured speedup divided by ideal speedup.
• Straggler: A worker that finishes later than the rest.
• World size: The total number of workers in the group.

Summary

Data parallelism copies the model onto every device, gives each device a different slice of the batch, and averages gradients so that all copies stay identical. It is easy to adopt, mathematically equal to large-batch training, and scales well when communication is overlapped with computation. Its main limit is memory, since each device holds the full model state, which is why sharded variants and model-splitting methods exist.

For tabular foundation models, data parallelism is a natural fit. A pre-training batch is a set of small tables or tasks, each worker handles its own tables, often generated from its own seed, and the averaged gradient trains one shared encoder for mixed types, missing values, and unnamed columns. Pay attention to balanced table sizes, correct metric aggregation, consistent replay and importance estimates in continual learning, and a learning rate matched to the global batch. For related topics, see mixed precision, FSDP, DeepSpeed ZeRO, Megatron-LM, and /model/architecture.

Related

It belongs to Foundation models from scratch in my stack, under Foundation models and applied ML. Others in the same group: