How to Implement Parameter Tying Across Pipeline Stages in Nanotron

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 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, 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, 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):

# 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:

python example_tie_across_pp.py

Handling Gradient Synchronization

The sync_tied_weights_gradients function at 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 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 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 provides concrete examples of rank mapping for various parallel configurations.

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