# How to Implement Parameter Tying Across Pipeline Stages in Nanotron

> Learn to implement parameter tying across pipeline stages in Nanotron. This guide explains how to use tie_parameters and sync_tied_weights_gradients for efficient model parallelism. Reduce memory and boost performance.

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

---

**To implement parameter tying across pipeline stages in Nanotron, identify parameters by their hierarchical names, specify the global pipeline ranks that share them, invoke `tie_parameters()` with a reduction operation like `dist.ReduceOp.SUM`, and synchronize gradients after the backward pass using `sync_tied_weights_gradients()`.**

Parameter tying allows multiple neural network layers to share a single physical weight tensor, reducing memory consumption and enforcing representation consistency. In the Nanotron framework, this capability extends across pipeline-parallel boundaries, enabling layers residing on different pipeline stages to reference the same parameters while maintaining correct gradient flow through explicit cross-stage synchronization.

## Understanding Parameter Tying in Pipeline Parallelism

When training large models with pipeline parallelism (PP), different layers reside on distinct devices. **Parameter tying** in this context means sharing weights between layers that may exist on separate pipeline stages. Nanotron handles this through metadata tracking and explicit gradient synchronization rather than physical tensor sharing across device boundaries.

The implementation in [`src/nanotron/parallel/tied_parameters.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/parallel/tied_parameters.py) manages the full lifecycle of tied parameters: creation, process group initialization, and gradient reduction. The system ensures that while each pipeline stage maintains its own tensor instance, the gradients are aggregated across stages after each backward pass to keep the weights synchronized.

## Core Components for Cross-Stage Tying

### NanotronParameter Class

The `NanotronParameter` class, defined in [`src/nanotron/parallel/parameters.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/parallel/parameters.py), extends `torch.nn.Parameter` to store tying metadata. When a parameter is tied, this subclass tracks the `global_ranks` participating in the tie, the reduction operation, and references to the root module. You can check if a parameter is tied using the `is_tied` attribute and retrieve metadata via `get_tied_info()`.

### tie_parameters Function

Located at lines 31-84 of [`src/nanotron/parallel/tied_parameters.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/parallel/tied_parameters.py), this function performs the actual parameter substitution. It accepts:
- `root_module`: The model containing the parameters
- `ties`: A list of `(parameter_name, tuple_of_global_ranks)` pairs
- `parallel_context`: The `ParallelContext` instance
- `reduce_op`: The reduction operation (typically `dist.ReduceOp.SUM`)

The function verifies that all specified ranks belong to the **same data-parallel replica** before performing in-place substitution of the parameters with shared `NanotronParameter` instances.

### Process Group Creation

The `create_pg_for_tied_weights` function (lines 85-102) automatically builds dedicated process groups for each unique set of global ranks sharing a parameter. These groups are stored in `parallel_context.world_ranks_to_pg`, enabling efficient communication during gradient synchronization without requiring global barriers.

### Gradient Synchronization

After `loss.backward()`, you must call `sync_tied_weights_gradients` (lines 120-168) to aggregate gradients across pipeline stages. This function groups tied parameters by their process groups and reduction operators, then performs all-reduce operations to ensure all stages receive the same gradient update.

## Step-by-Step Implementation

Below is a complete example demonstrating how to tie weights across two pipeline stages (PP=2):

```python

# example_tie_across_pp.py

import torch
from torch import nn
from nanotron import distributed as dist
from nanotron.parallel import ParallelContext
from nanotron.parallel.tied_parameters import tie_parameters, sync_tied_weights_gradients

@dist.rerun_if_address_is_in_use()
def run():
    init = dist.init_distributed(tp=1, dp=1, pp=2)
    init(_main)()

def _main(parallel_context: ParallelContext):
    # Build a model split across pipeline stages

    if dist.get_rank(parallel_context.pp_pg) == 0:
        model = nn.ModuleDict({"dense0": nn.Linear(10, 10, device="cuda")})
    else:
        model = nn.ModuleDict({"dense1": nn.Linear(10, 10, device="cuda")})

    # Define ties: (parameter_path, tuple_of_global_pp_ranks)

    ties = [
        ("dense0.weight", (0,)),   # global rank 0

        ("dense1.weight", (1,)),   # global rank 1

    ]

    # Perform the tie with SUM reduction for gradients

    tie_parameters(
        root_module=model,
        ties=ties,
        parallel_context=parallel_context,
        reduce_op=dist.ReduceOp.SUM,
    )

    # Forward pass

    x = torch.randn(4, 10, device="cuda")
    if dist.get_rank(parallel_context.pp_pg) == 0:
        out = model.dense0(x)
    else:
        out = model.dense1(x)
    
    loss = out.sum()
    loss.backward()

    # Sync tied gradients across PP ranks

    sync_tied_weights_gradients(
        module=model,
        parallel_context=parallel_context,
        grad_accumulator=None,
    )

    # Gradients are now identical on both ranks

    parallel_context.destroy()

if __name__ == "__main__":
    run()

```

Run the example with:

```bash
python example_tie_across_pp.py

```

## Handling Gradient Synchronization

The `sync_tied_weights_gradients` function at [`src/nanotron/parallel/tied_parameters.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/parallel/tied_parameters.py) lines 120-168 is critical for cross-stage tying. After backward propagation completes, this function:

1. Identifies all parameters with `is_tied=True` that require gradient reduction
2. Groups them by their process group and reduction operation
3. Executes all-reduce (or all-reduce-coalesced) to aggregate gradients across the specified global ranks

Without this call, each pipeline stage would update its local copy of the tied weight independently, causing the shared parameters to diverge.

## Important Constraints and Validation

**Data-Parallel Replica Restriction**: The `tie_parameters` function strictly enforces that all tied ranks must belong to a single data-parallel replica. This check (implemented at lines 31-84 of [`tied_parameters.py`](https://github.com/huggingface/nanotron/blob/main/tied_parameters.py) via `assert len(dp_ranks) == 1`) ensures that each physical model replica maintains a consistent view of the tied weight. You cannot tie parameters across different DP replicas, as this would violate the consistency guarantees required for distributed training.

**Same Device vs. Different Devices**: The mechanism works whether tied modules reside on the same device (in-place sharing) or different devices (metadata-only sharing with explicit gradient sync). When on distinct PP ranks, the parameters remain independent `NanotronParameter` objects until `sync_tied_weights_gradients` reconciles their gradients.

## Summary

- **Identify parameters** using hierarchical names like `"dense0.weight"` when calling `tie_parameters()`
- **Specify global ranks** as tuples in the ties list to indicate which pipeline stages share each parameter
- **Use `dist.ReduceOp.SUM`** for gradient aggregation to maintain consistent updates across stages
- **Call `sync_tied_weights_gradients`** after every backward pass to synchronize gradients across pipeline boundaries
- **Respect DP constraints** by ensuring all tied ranks belong to the same data-parallel replica
- **Reference test examples** in [`tests/test_tie_weights.py`](https://github.com/huggingface/nanotron/blob/main/tests/test_tie_weights.py) for production-ready validation patterns

## Frequently Asked Questions

### Can I tie parameters across different data-parallel replicas?

No. Nanotron explicitly prohibits tying parameters across different data-parallel replicas. The `tie_parameters` function asserts that all participating global ranks belong to the same DP replica (`assert len(dp_ranks) == 1`). This restriction ensures that each physical model replica maintains a consistent view of the tied weights throughout training.

### What is the difference between tying on the same device versus different pipeline stages?

When tied parameters reside on the same device, Nanotron can share the underlying physical tensor directly. When tying across pipeline stages (different devices), each stage maintains its own tensor instance, and consistency is achieved through metadata tracking in `NanotronParameter` and explicit gradient synchronization via `sync_tied_weights_gradients`. The API remains identical for both scenarios.

### When should I call `sync_tied_weights_gradients` during training?

You must call `sync_tied_weights_gradients` immediately after `loss.backward()` and before `optimizer.step()`. This timing ensures that gradients from all pipeline stages are aggregated across the tied process groups before the optimizer updates the weights. Calling it after the optimizer step would result in divergent weights across stages.

### How do I determine which global ranks to specify in the ties tuple?

The tuple should contain the global ranks (relative to the world group) of the pipeline stages that own the tied parameter copies. For a configuration with TP=1, DP=1, and PP=2, the global ranks are typically 0 and 1. You can verify the mapping using `parallel_context.pp_pg` and `dist.get_rank()`. The test suite in [`tests/test_tie_weights.py`](https://github.com/huggingface/nanotron/blob/main/tests/test_tie_weights.py) provides concrete examples of rank mapping for various parallel configurations.