How to Configure Tensor Parallel Mode (All-Reduce vs Reduce-Scatter) in Nanotron

Set the tp_mode field in ParallelismConfig to TensorParallelLinearMode.ALL_REDUCE or TensorParallelLinearMode.REDUCE_SCATTER to switch between all-reduce and reduce-scatter collectives for tensor-parallel linear layers.

Nanotron, Hugging Face's lightweight library for large-scale transformer training, exposes granular control over tensor-parallel communication through the TensorParallelLinearMode enum. Configuring this tensor parallel mode determines whether row-parallel linear layers aggregate gradients via an all-reduce or a reduce-scatter operation, directly impacting inter-GPU bandwidth utilization.

Tensor Parallel Modes Explained

Nanotron defines the communication strategy in src/nanotron/parallel/tensor_parallel/enum.py through the TensorParallelLinearMode enum:

from enum import Enum, auto

class TensorParallelLinearMode(Enum):
    ALL_REDUCE = auto()        # Row-linear layers use all-reduce on the output

    REDUCE_SCATTER = auto()    # Row-linear layers use reduce-scatter; column-linear layers use all-gather

ALL_REDUCE Mode

In ALL_REDUCE mode, TensorParallelRowLinear shards the input dimension and performs an all_reduce_sum operation on the output tensor. TensorParallelColumnLinear shards the output dimension and uses all_gather on the input. This pattern mirrors the original Megatron-LM implementation and is generally simpler but may incur higher latency when many GPUs share an interconnect.

REDUCE_SCATTER Mode

In REDUCE_SCATTER mode, TensorParallelRowLinear executes a reduce_scatter_sum operation, splitting the reduced result across ranks so each receives only its slice. TensorParallelColumnLinear still uses all_gather on the input. This reduces the total data volume exchanged during the row-parallel step, making it preferable when inter-GPU bandwidth is the primary bottleneck.

Implementation Details

Both TensorParallelColumnLinear and TensorParallelRowLinear forward calls pass the configured mode to functional kernels in src/nanotron/parallel/tensor_parallel/nn.py:

# From nn.py

return column_linear(..., tp_mode=self.mode, ...)
return row_linear(..., tp_mode=self.mode, ...)

The functional kernels in src/nanotron/parallel/tensor_parallel/functional.py read the tp_mode parameter and dispatch to the appropriate collective operations (all_reduce_sum, reduce_scatter_sum, or all_gather) based on the enum value.

How to Configure Tensor Parallel Mode

You can specify the mode through three equivalent interfaces, all of which ultimately set the tp_mode attribute inside ParallelismConfig (defined in src/nanotron/config/parallelism_config.py).

Python Configuration

Instantiate NanotronConfig directly with the enum value:

from nanotron.parallel.tensor_parallel.enum import TensorParallelLinearMode
from nanotron.config import NanotronConfig, ParallelismConfig

config = NanotronConfig(
    parallelism=ParallelismConfig(
        tp=4,
        tp_mode=TensorParallelLinearMode.REDUCE_SCATTER,
        dp=1,
        pp=1,
    )
)

YAML or JSON Configuration

When loading from a dictionary, provide the string representation; the loader casts it to the enum automatically:

parallelism:
  tp: 4
  tp_mode: REDUCE_SCATTER
  dp: 1
  pp: 1

Command-Line Interface

The training and generation entry points expose a --tp-mode flag that accepts all-reduce or reduce-scatter (case-insensitive):

python -m nanotron.run_train \
    --tp 4 \
    --tp-mode reduce-scatter \
    --config-file config.yaml

Performance Trade-offs

Choose your tensor parallel mode based on your hardware topology and tensor-parallel size:

• ALL_REDUCE: Best for small to medium TP sizes (≤ 8 GPUs) where latency dominates and the interconnect is high-speed. The simpler communication pattern is well-tested and stable.

• REDUCE_SCATTER: Optimal for larger TP sizes (≥ 8 GPUs) or when the inter-GPU bandwidth is constrained. By reducing the data volume sent to each rank during the row-linear forward pass, this mode improves bandwidth efficiency at the cost of slightly more complex communication scheduling.

Both modes produce mathematically identical results; the difference is purely in the collective-communication implementation.

Summary

• TensorParallelLinearMode controls the communication pattern for tensor-parallel linear layers and is defined in src/nanotron/parallel/tensor_parallel/enum.py.
• ALL_REDUCE uses all_reduce_sum for row-linear layers; REDUCE_SCATTER uses reduce_scatter_sum to minimize bandwidth.
• Set the mode via ParallelismConfig.tp_mode in Python, YAML, or through the --tp-mode CLI flag.
• Use REDUCE_SCATTER for large tensor-parallel degrees or bandwidth-constrained environments; use ALL_REDUCE for simplicity with smaller topologies.

Frequently Asked Questions

What is the default tensor parallel mode in Nanotron?

According to the source code, the default is typically ALL_REDUCE unless explicitly overridden in the configuration. You can verify the active mode by inspecting trainer.config.parallelism.tp_mode after initialization.

Does switching between ALL_REDUCE and REDUCE_SCATTER affect numerical accuracy?

No. Both modes are mathematically equivalent and produce identical loss curves; they differ only in the distributed communication primitive used to aggregate partial results across tensor-parallel ranks.

Which mode should I use for a 64-GPU tensor parallel setup?

For large tensor-parallel sizes (≥ 8 GPUs), REDUCE_SCATTER is recommended because it reduces the amount of data transferred over the interconnect during the row-parallel linear layer computation, alleviating bandwidth bottlenecks that become pronounced at scale.

Where does the actual communication logic reside?

The collective operation selection logic lives in src/nanotron/parallel/tensor_parallel/functional.py, while the enum definition is in src/nanotron/parallel/tensor_parallel/enum.py. The linear layer wrappers in src/nanotron/parallel/tensor_parallel/nn.py bridge the configuration to the functional kernels.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}