How to Set Up Sequence Parallelism with Context Parallel Size in Nanotron

In Hugging Face Nanotron, enable sequence parallelism by setting tp_mode="reduce_scatter" and configure context parallelism via context_parallel_size (cp) in the ParallelismConfig, ensuring that tp × pp × dp × cp equals your total world size.

Nanotron is a scalable framework for training large language models developed by Hugging Face. Setting up sequence parallelism with context parallel size allows you to distribute long sequences across multiple GPUs while minimizing communication overhead. This guide walks through the configuration based on the actual implementation in the huggingface/nanotron repository.

Understanding Sequence and Context Parallelism

Nanotron implements sequence parallelism through the reduce_scatter tensor-parallel mode. When activated, each tensor-parallel rank processes a distinct slice of the sequence dimension, eliminating the need for all-reduce operations across tensor-parallel groups.

Context parallelism complements this by splitting the batch dimension across a new parallel dimension labeled cp (context parallel). According to the source code in nanotron/parallel/context.py, the implementation enforces a strict constraint on the total world size:

tp * pp * dp * cp == WORLD_SIZE

This validation occurs in nanotron/parallel/context.py at lines 27-32, ensuring the product of tensor-parallel (tp), pipeline-parallel (pp), data-parallel (dp), and context-parallel (cp) sizes matches the total number of available GPUs.

Configuration Parameters

The parallelism settings are defined in nanotron/config/parallelism_config.py. To enable both features, modify these specific fields in your configuration:

• tp_mode: Set to "reduce_scatter" to activate sequence parallelism (default is "all_reduce"). Defined at line 24.
• context_parallel_size: Integer specifying the context-parallel degree (default 1). Defined at line 39.
• tensor_parallel_size: Number of tensor-parallel replicas.
• pipeline_parallel_size: Number of pipeline stages.
• data_parallel_size: Standard data parallelism degree.

When the trainer initializes, it reads these values from ParallelismConfig and constructs the distributed environment. In nanotron/trainer.py at line 175, the trainer passes context_parallel_size to the ParallelContext constructor, which then creates the context-parallel process group self.cp_pg at lines 77-88 in nanotron/parallel/context.py.

Step-by-Step Implementation

Configuring the Parallelism Settings

Create or modify your Nanotron configuration file to specify both the tensor-parallel mode and context-parallel size:

from nanotron.config import Config, ParallelismConfig

config = Config(
    parallelism=ParallelismConfig(
        tp=4,                     # tensor-parallel degree

        pp=2,                     # pipeline-parallel degree

        dp=1,                     # data-parallel degree

        cp=2,                     # context-parallel degree

        tp_mode="reduce_scatter"  # enables sequence parallelism

    ),
    # ... additional model and training configuration ...

)

For this example, the required world size is 4 × 2 × 1 × 2 = 16 GPUs.

Launching with the Correct World Size

Ensure your job launch matches the calculated world size. Using the SLURM launcher provided in the repository, you can pass the context-parallel size directly via the --cp flag (implemented in slurm_launcher.py at lines 331-353):

sbatch run_train.sh \
    --tp 4 --pp 2 --dp 1 --cp 2 \
    --tp-mode reduce_scatter

Alternatively, launch directly with torchrun:

torchrun --nproc_per_node=16 your_training_script.py

Verifying the Context Parallel Group

After instantiation, verify that the context-parallel process group exists by inspecting the trainer's parallel context:

# Inside your training script

from nanotron.trainer import Trainer

trainer = Trainer(config)
print(trainer.parallel_context.cp_pg)  # Should show the process group info

The data collator utilizes this group to determine rank and size information for batch handling, as seen in nanotron/data/clm_collator.py at lines 71 and 100. Additionally, nanotron/helpers.py logs all parallel dimensions including cp at line 649, allowing you to confirm the configuration in your training logs.

Communication Optimization with Reduce-Scatter

When sequence parallelism is active via tp_mode="reduce_scatter", Nanotron eliminates redundant gradient synchronization. Because each tensor-parallel rank handles a unique sequence slice, the framework skips the typical all-reduce operation across TP ranks.

This optimization appears in run_generate.py at line 125, where the code explicitly avoids TP-level synchronization when sequence parallelism is enabled. The result is reduced communication overhead during both forward and backward passes, particularly beneficial for long-context training.

Summary

• Sequence parallelism is enabled by setting tp_mode="reduce_scatter" in ParallelismConfig (defined in nanotron/config/parallelism_config.py).
• Context parallelism uses the context_parallel_size (cp) parameter to distribute batches across additional GPUs.
• The product of all parallel dimensions (tp × pp × dp × cp) must equal WORLD_SIZE, enforced in nanotron/parallel/context.py.
• The ParallelContext constructor creates the cp_pg process group for context-parallel communication.
• Using reduce_scatter mode eliminates unnecessary all-reduce operations, as referenced in run_generate.py.

Frequently Asked Questions

Can I use sequence parallelism without context parallelism?

Yes. Set tp_mode="reduce_scatter" while keeping context_parallel_size=1 (the default). This activates sequence parallelism across tensor-parallel ranks without splitting the batch dimension further. The cp dimension remains inactive but the communication optimizations for sequence parallelism still apply.

What happens if my world size does not match the parallel configuration?

Nanotron raises an error during ParallelContext initialization. The code at lines 27-32 in nanotron/parallel/context.py explicitly validates that tp * pp * dp * cp == WORLD_SIZE. If the product does not match the number of available GPUs, the training job will fail immediately with a clear assertion message.

How does context parallelism affect batch processing?

Context parallelism splits the global batch across the cp dimension, effectively increasing the micro-batch size processed by each GPU while maintaining the same global batch size. The CLMCollator in nanotron/data/clm_collator.py uses the cp_pg process group to coordinate data loading and ensure each rank receives the correct slice of the sequence, as seen in the rank and size lookups at lines 71 and 100.

Where is the context parallel process group initialized?

The context-parallel process group self.cp_pg is created in the ParallelContext class constructor at lines 77-88 of nanotron/parallel/context.py. This occurs when the trainer instantiates ParallelContext at line 175 of nanotron/trainer.py, passing the context_parallel_size value from your configuration.

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