# How to Configure 3D Parallelism (DP+TP+PP) in Hugging Face Nanotron

> Learn to configure 3D parallelism DP+TP+PP in Hugging Face Nanotron. Set YAML config and use torchrun for efficient distributed training with simple steps.

- Repository: [Hugging Face/nanotron](https://github.com/huggingface/nanotron)
- Tags: how-to-guide
- Published: 2026-03-03

---

**Configure 3D parallelism in Nanotron by setting the `dp`, `tp`, and `pp` fields in your YAML config and ensuring `dp * tp * pp equals WORLD_SIZE`, then launch with `torchrun` to automatically initialize the `ParallelContext` process groups.**

Nanotron is a lightweight, open-source framework for pretraining and fine-tuning large language models. It implements **3-dimensional (3D) parallelism**—combining **data parallelism (DP)**, **tensor parallelism (TP)**, and **pipeline parallelism (PP)**—through a hierarchical process group structure managed by the `ParallelContext` class.

## Understanding 3D Parallelism in Nanotron

Nanotron partitions training across three orthogonal dimensions:

- **Data Parallelism (DP)**: Replicates the model across independent data shards; gradients are averaged across `dp` ranks after each backward pass.
- **Tensor Parallelism (TP)**: Splits individual layers (typically linear projections) across `tp` ranks; activations are synchronized via all-reduce or reduce-scatter operations.
- **Pipeline Parallelism (PP)**: Distributes sequential layers across `pp` stages; each stage processes micro-batches and passes activations (via `TensorPointer` objects) to the next stage.

The product of these dimensions must equal the total number of processes: **`dp × tp × pp × cp × ep = WORLD_SIZE`**, where `cp` (context parallelism) and `ep` (expert parallelism) are optional additional dimensions.

## Core Configuration Components

### ParallelismArgs: The Configuration Interface

User-facing parallelism settings are defined in [`src/nanotron/config/parallelism_config.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/config/parallelism_config.py) within the `ParallelismArgs` dataclass. Key fields include:

- `dp`: Integer size of the data-parallel dimension.
- `tp`: Integer size of the tensor-parallel dimension.
- `pp`: Integer size of the pipeline-parallel dimension.
- `tp_mode`: Either `"ALL_REDUCE"` (default) or `"REDUCE_SCATTER"` (enables sequence parallelism).
- `tp_linear_async_communication`: Boolean flag to enable asynchronous communication for tensor-parallel linear layers.

### ParallelContext: Runtime Process Groups

The `ParallelContext` class in [`src/nanotron/parallel/context.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/parallel/context.py) validates the parallelism configuration and constructs the process group hierarchy. During initialization (lines 60–71), it builds a 5-dimensional rank matrix of shape `(ep, pp, dp, cp, tp)` and slices it along each axis to create:

- `tp_pg`: Tensor-parallel group (ranks sharing the same EP, PP, DP, CP indices).
- `pp_pg`: Pipeline-parallel group (ranks sharing the same EP, DP, CP, TP indices).
- `dp_pg`: Data-parallel group (ranks sharing the same EP, PP, CP, TP indices).

These groups are consumed by `TensorParallelRowLinear` and `TensorParallelColumnLinear` (in [`src/nanotron/parallel/tensor_parallel/nn.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/parallel/tensor_parallel/nn.py)) and by the pipeline engine ([`src/nanotron/parallel/pipeline_parallel/engine.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/parallel/pipeline_parallel/engine.py)).

## Step-by-Step Configuration Guide

### Step 1: Define Parallel Dimensions in YAML

Create a configuration file (e.g., [`config_3d.yaml`](https://github.com/huggingface/nanotron/blob/main/config_3d.yaml)) specifying the three dimensions:

```yaml
parallelism:
  dp: 4          # 4 data-parallel replicas

  pp: 2          # 2 pipeline stages

  tp: 2          # 2 tensor-parallel shards per stage

  pp_engine: 1f1b   # 1F1B pipeline schedule

  tp_mode: ALL_REDUCE   # Standard tensor parallelism

  tp_linear_async_communication: false

```

Ensure the product (4 × 2 × 2 = 16) matches your total GPU count.

### Step 2: Set Environment Variables

Before launching, export the standard distributed training variables:

```bash
export WORLD_SIZE=16
export MASTER_ADDR=127.0.0.1
export MASTER_PORT=29500

# RANK and LOCAL_RANK are set automatically by torchrun

```

Nanotron validates that `WORLD_SIZE` equals the product of your parallelism dimensions in `ParallelContext.__init__`.

### Step 3: Launch with torchrun

Execute the training script using `torchrun` or `python -m torch.distributed.run`:

```bash
torchrun --nproc_per_node=16 run_train.py \
  --config-file config_3d.yaml

```

The [`run_train.py`](https://github.com/huggingface/nanotron/blob/main/run_train.py) entry point (located in the repository root) parses the `parallelism` section, instantiates `ParallelContext`, and passes it to `DistributedTrainer`.

## Validating Your 3D Parallelism Setup

To verify that process groups are constructed correctly, inspect the `ParallelContext` after initialization:

```python
from nanotron.parallel.context import ParallelContext

# Initialize with your config values

pc = ParallelContext(
    tensor_parallel_size=2,
    pipeline_parallel_size=2,
    data_parallel_size=4,
)

print(f"TP group ranks: {pc.tp_pg.ranks()}")
print(f"PP group ranks: {pc.pp_pg.ranks()}")
print(f"DP group ranks: {pc.dp_pg.ranks()}")

```

Each print statement lists the global ranks participating in that specific parallel dimension, confirming that `dp * tp * pp` equals your world size.

## Summary

- **3D parallelism** in Nanotron combines data, tensor, and pipeline parallelism through the `ParallelContext` class.
- **Configuration** occurs via `ParallelismArgs` in [`src/nanotron/config/parallelism_config.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/config/parallelism_config.py), typically set through a YAML file with `dp`, `tp`, and `pp` fields.
- **Constraint**: The product of all parallelism dimensions must equal `WORLD_SIZE` (total number of GPUs).
- **Launch**: Use `torchrun` with `--nproc_per_node` matching the total GPU count; Nanotron automatically builds process groups (`tp_pg`, `pp_pg`, `dp_pg`) in [`src/nanotron/parallel/context.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/parallel/context.py).

## Frequently Asked Questions

### What is the mathematical relationship between dp, tp, pp, and world_size?

Nanotron requires that **`dp × tp × pp × cp × ep = WORLD_SIZE`**, where `cp` (context parallelism) and `ep` (expert parallelism) default to 1 if unused. The `ParallelContext` constructor in [`src/nanotron/parallel/context.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/parallel/context.py) validates this equality and raises an assertion error if the dimensions are inconsistent with the number of processes launched by `torchrun`.

### How do I enable asynchronous communication for tensor parallelism?

Set `tp_linear_async_communication: true` in your YAML configuration under the `parallelism` section. This flag is defined in [`src/nanotron/config/parallelism_config.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/config/parallelism_config.py) and consumed by `TensorParallelRowLinear` and `TensorParallelColumnLinear` layers in [`src/nanotron/parallel/tensor_parallel/nn.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/parallel/tensor_parallel/nn.py) to overlap communication with computation during the backward pass.

### Can I use context parallelism alongside DP, TP, and PP?

Yes. Nanotron supports a 5D parallel configuration including **context parallelism (CP)** and **expert parallelism (EP)**. Add `cp: <size>` to your YAML config; the `ParallelContext` will include this dimension in its rank matrix `(ep, pp, dp, cp, tp)`. Ensure your `WORLD_SIZE` accounts for the additional factor (e.g., `dp * tp * pp * cp`).

### Where does the pipeline engine schedule forward and backward passes?

The pipeline engine is implemented in [`src/nanotron/parallel/pipeline_parallel/engine.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/parallel/pipeline_parallel/engine.py). It provides schedulers such as `AllForwardAllBackwardPipelineEngine` and `1f1b` (one-forward-one-backward), which orchestrate micro-batch execution across the `pp_pg` process group. The engine uses `TensorPointer` objects to manage activation communication between pipeline stages defined in [`src/nanotron/parallel/pipeline_parallel/p2p.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/parallel/pipeline_parallel/p2p.py).