# How PyTorch Lightning DDP Strategy is Configured for Multi-GPU Training in LeanAgent

> Learn how LeanAgent configures PyTorch Lightning's DDP strategy leveraging NCCL timeout, four GPUs, and bf16-mixed precision for stable multi-GPU training of retrieval models.

- Repository: [LeanDojo/leanagent](https://github.com/lean-dojo/leanagent)
- Tags: how-to-guide
- Published: 2026-03-05

---

**LeanAgent configures PyTorch Lightning's DDP strategy with a one-year NCCL timeout, four GPUs, and bf16-mixed precision to enable stable distributed training for retrieval models and Fisher information computation.**

The LeanAgent repository implements large-scale neural theorem proving using distributed training across multiple GPUs. According to the LeanAgent source code, the project specifically utilizes PyTorch Lightning's Distributed Data Parallel (DDP) strategy with custom NCCL environment variables to prevent communication timeouts during long-running training sessions on four GPU devices.

## DDP Strategy Configuration Overview

LeanAgent runs both retrieval training and Fisher-information computation using PyTorch Lightning's `DDPStrategy` with identical configurations. The setup spans two primary files: [`leanagent.py`](https://github.com/lean-dojo/leanagent/blob/main/leanagent.py) handles the main retrieval model training, while [`compute_fisher.py`](https://github.com/lean-dojo/leanagent/blob/main/compute_fisher.py) manages the Fisher Information Matrix calculation. Both implementations configure a `DDPStrategy` object with an extended timeout, set critical NCCL environment variables, and explicitly request four GPU devices with mixed-precision training.

## Main Training Loop Configuration (leanagent.py)

In [`leanagent.py`](https://github.com/lean-dojo/leanagent/blob/main/leanagent.py) (lines 1185-1194), the primary trainer initialization establishes the DDP strategy with specific environment tweaks for stability. The configuration sets a one-year NCCL timeout to prevent hangs during extended training runs on large clusters.

### Environment Setup for NCCL Stability

Before creating the trainer, the code configures NCCL environment variables to handle asynchronous errors and extend the communication timeout:

```python
VERY_LONG_TIMEOUT = 7 * 24 * 60 * 60 * 52          # 1 year (seconds)

os.environ['TORCH_NCCL_ASYNC_ERROR_HANDLING'] = '1'
os.environ['NCCL_TIMEOUT'] = str(VERY_LONG_TIMEOUT * 1000)  # ms

```

### Trainer Initialization with DDP Strategy

The `pl.Trainer` is instantiated with the `DDPStrategy` object passed to the `strategy` parameter, alongside gradient clipping and mixed-precision settings:

```python
ddp_strategy = DDPStrategy(timeout=timedelta(seconds=VERY_LONG_TIMEOUT))

trainer = pl.Trainer(
    accelerator="gpu",
    gradient_clip_val=1.0,
    precision="bf16-mixed",
    strategy=ddp_strategy,
    devices=4,                     # use 4 GPUs

    accumulate_grad_batches=4,
    callbacks=[lr_monitor, checkpoint_callback, early_stop_callback],
    max_epochs=current_epoch + epochs_per_repo,
    log_every_n_steps=1,
    num_sanity_val_steps=0,
    default_root_dir=custom_log_dir,
)

# Synchronisation points

trainer.strategy.barrier()   # before loading data

trainer.fit(model, datamodule=data_module, ckpt_path=model_checkpoint_path)
trainer.strategy.barrier()   # after training

```

## Fisher Information Computation (compute_fisher.py)

The Fisher Information Matrix calculation in [`compute_fisher.py`](https://github.com/lean-dojo/leanagent/blob/main/compute_fisher.py) (lines 53-66) reuses the identical DDP configuration to ensure consistency across distributed processes. This trainer also runs on four GPUs with the same NCCL timeout settings:

```python
VERY_LONG_TIMEOUT = 7 * 24 * 60 * 60 * 52  # 1 year

os.environ['TORCH_NCCL_ASYNC_ERROR_HANDLING'] = '1'
os.environ['NCCL_TIMEOUT'] = str(VERY_LONG_TIMEOUT * 1000)

ddp_strategy = DDPStrategy(timeout=timedelta(seconds=VERY_LONG_TIMEOUT))

fisher_trainer = pl.Trainer(
    accelerator="gpu",
    precision="bf16-mixed",
    strategy=ddp_strategy,
    devices=4,               # same 4-GPU setup

    max_epochs=1,
    log_every_n_steps=1,
    num_sanity_val_steps=0,
)

fisher_trainer.strategy.barrier()
fisher_trainer.fit(fisher_module, datamodule=data_module)
fisher_trainer.strategy.barrier()

```

## Synchronization and Barrier Usage

LeanAgent explicitly calls `trainer.strategy.barrier()` before and after critical sections to guarantee that all four GPU processes reach synchronization points simultaneously. These barriers appear before data module loading, after training completion, and surrounding Fisher information computation. This pattern prevents race conditions and ensures consistent checkpoint saving across distributed processes.

## Summary

- **DDPStrategy configuration**: Uses `DDPStrategy(timeout=timedelta(seconds=VERY_LONG_TIMEOUT))` with a one-year timeout to prevent NCCL communication hangs.
- **GPU allocation**: Explicitly sets `devices=4` and `accelerator="gpu"` for CUDA-based distributed training.
- **Precision settings**: Configures `precision="bf16-mixed"` for BFloat16 mixed-precision training on compatible hardware.
- **Environment variables**: Sets `TORCH_NCCL_ASYNC_ERROR_HANDLING` to `'1'` and `NCCL_TIMEOUT` to approximately one year in milliseconds.
- **Synchronization**: Implements `trainer.strategy.barrier()` calls to coordinate the four GPU processes during data loading and training phases.

## Frequently Asked Questions

### Why does LeanAgent use a one-year NCCL timeout?

LeanAgent sets the NCCL timeout to approximately one year (52 weeks) using `VERY_LONG_TIMEOUT = 7 * 24 * 60 * 60 * 52` to prevent communication hangs during long-running training sessions on large clusters. This extreme timeout value ensures that distributed processes remain connected during extended Fisher information computations and retrieval model training without dropping due to default NCCL timeout limits.

### How many GPUs does LeanAgent use for distributed training?

According to the source code in both [`leanagent.py`](https://github.com/lean-dojo/leanagent/blob/main/leanagent.py) and [`compute_fisher.py`](https://github.com/lean-dojo/leanagent/blob/main/compute_fisher.py), LeanAgent explicitly configures `devices=4` in the PyTorch Lightning Trainer, indicating that the DDP strategy is designed to run across four GPU devices simultaneously for both retrieval training and Fisher information calculation.

### What precision mode does LeanAgent use for multi-GPU training?

LeanAgent configures the PyTorch Lightning Trainer with `precision="bf16-mixed"`, enabling BFloat16 mixed-precision training. This setting reduces memory usage and increases training speed on modern NVIDIA GPUs while maintaining numerical stability for the retrieval models and Fisher information matrices.

### Where is the DDP strategy configured in the LeanAgent codebase?

The DDP strategy is configured in two primary locations: [`leanagent.py`](https://github.com/lean-dojo/leanagent/blob/main/leanagent.py) (lines 1185-1194) for the main retrieval model training loop, and [`compute_fisher.py`](https://github.com/lean-dojo/leanagent/blob/main/compute_fisher.py) (lines 53-66) for the Fisher Information Matrix computation. Both files create a `DDPStrategy` object with identical timeout settings and pass it to the `strategy` parameter of the `pl.Trainer` constructor.