# How to Set Up Distributed Training with Multiple GPUs Using NCCL Backend in rwkv-clip

> Learn how to set up distributed training with multiple GPUs using the NCCL backend in rwkv-clip. Easily configure efficient multi-GPU training with torchrun and DistributedDataParallel.

- Repository: [DeepGlint/rwkv-clip](https://github.com/deepglint/rwkv-clip)
- Tags: how-to-guide
- Published: 2026-02-28

---

**Launch [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py) with `torchrun` to spawn one process per GPU, which automatically initializes the NCCL process group and wraps the model in `DistributedDataParallel` for efficient multi-GPU training.**

The rwkv-clip repository implements cross-modal contrastive learning using PyTorch's native Distributed Data Parallel (DDP) framework with the **NCCL** communication backend. Setting up distributed training requires proper initialization of the process group, wrapping the model for synchronized training, and configuring a shard-aware data loader to partition batches across devices. This guide covers the exact implementation details found in [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py) and [`dali.py`](https://github.com/deepglint/rwkv-clip/blob/main/dali.py) to help you scale training across any number of NCCL-compatible GPUs.

## Initializing the NCCL Process Group

The distributed training setup begins in the first 20 lines of [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py), where the script reads environment variables set by the launch utility and initializes the communication backend.

```python
import torch.distributed as distributed

# Lines 16-19 read environment variables set by torchrun

local_rank = int(os.environ["LOCAL_RANK"])
rank = int(os.environ["RANK"])
world_size = int(os.environ["WORLD_SIZE"])

# Initialize NCCL backend (line ~20)

distributed.init_process_group(backend="nccl")
torch.cuda.set_device(local_rank)

```

The **NCCL** (NVIDIA Collective Communications Library) backend provides optimized peer-to-peer GPU communication for NVIDIA hardware. The `torch.cuda.set_device(local_rank)` call ensures each process exclusively targets one GPU, preventing device contention.

## Wrapping the Model for Distributed Training

After loading the RWKV-CLIP architecture, the code converts batch normalization layers to synchronized versions and wraps the model with **DistributedDataParallel** (DDP). This happens in the `main()` function at lines 35-42 of [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py):

```python

# Convert to synchronized batch norm for consistent statistics across GPUs

RWKV_CLIP_model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(RWKV_CLIP_model)

# Wrap with DDP

RWKV_CLIP_model = torch.nn.parallel.DistributedDataParallel(
    module=RWKV_CLIP_model,
    bucket_cap_mb=32,
    find_unused_parameters=True,
    static_graph=True)

```

The **SyncBatchNorm** conversion ensures that batch statistics are aggregated across all GPUs during the forward pass, which is critical for consistent normalization behavior in distributed settings. The DDP wrapper parameters include:
- **`bucket_cap_mb=32`**: Sets the gradient bucket size to 32 MB for efficient overlapping of communication and computation
- **`find_unused_parameters=True`**: Handles potential unused parameters in the RWKV architecture
- **`static_graph=True`**: Optimizes performance when the model graph remains constant across iterations

## Sharding Data Across GPUs with DALI

The data loader in [`dali.py`](https://github.com/deepglint/rwkv-clip/blob/main/dali.py) implements automatic data sharding to ensure each GPU processes a distinct subset of the training data. This occurs at lines 56-60 using NVIDIA DALI's MXNet reader:

```python

# Inside dali_dataloader() function

pipe = fn.readers.mxnet(
    path=data_path,
    index_path=index_path,
    num_shards=world_size,  # Total number of GPUs

    shard_id=rank,          # Current GPU rank

    ...
)

```

The **`num_shards=world_size`** and **`shard_id=rank`** parameters split the MXNet record files (`.rec` and `.idx` format) across all participating processes. Each GPU loads only its assigned shard, eliminating data duplication and ensuring the effective batch size scales linearly with the number of GPUs.

## Launching Distributed Jobs

### Using torchrun (Recommended)

For single-node multi-GPU training, use **torchrun** (PyTorch's elastic launch utility) which automatically sets the required environment variables (`RANK`, `LOCAL_RANK`, `WORLD_SIZE`):

```bash
torchrun \
    --nnodes=1 \
    --nproc_per_node=4 \
    --rdzv_id=rwkv_clip \
    --rdzv_backend=c10d \
    train.py \
      --output ./output_dir \
      --train-data /data/my_dataset \
      --train-num-samples 1280000 \
      --batch-size 256 \
      --epochs 32 \
      --lr 0.1 \
      --precision bf16

```

This command spawns four processes, each assigned to a different GPU via `LOCAL_RANK`. The NCCL backend establishes peer-to-peer communication channels automatically using these environment variables.

### Programmatic Launch with torch.multiprocessing

For custom launch scenarios, use `torch.multiprocessing.spawn` to programmatically set the environment variables:

```python
import os
import torch
import torch.distributed as dist
from train import main, get_args

def run(rank, world_size):
    os.environ["RANK"] = str(rank)
    os.environ["LOCAL_RANK"] = str(rank)
    os.environ["WORLD_SIZE"] = str(world_size)
    args = get_args()  # Parses CLI flags from train.py

    main(args)

if __name__ == "__main__":
    world_sz = torch.cuda.device_count()
    torch.multiprocessing.spawn(run, args=(world_sz,), nprocs=world_sz)

```

Both methods trigger the same initialization path in [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py) because the script reads the environment variables at startup to determine the process group configuration.

## Multi-Node Cluster Configuration

To scale beyond a single node, add the `--nnodes` and `--node_rank` parameters along with a rendezvous endpoint:

```bash
torchrun \
    --nnodes=2 \
    --node_rank=0 \
    --nproc_per_node=8 \
    --rdzv_id=rwkv_clip \
    --rdzv_backend=c10d \
    --rdzv_endpoint=192.168.1.100:29500 \
    train.py \
      --train-data /shared/data \
      --batch-size 256

```

NCCL can span multiple nodes when the network fabric (e.g., InfiniBand or high-speed Ethernet) supports GPU Direct RDMA. Ensure all nodes use the same NCCL-enabled PyTorch build and can resolve the rendezvous endpoint.

## Verifying NCCL Communication

To troubleshoot connectivity issues or verify GPU topology detection, set the `NCCL_DEBUG` environment variable before launching:

```bash
NCCL_DEBUG=INFO torchrun --nproc_per_node=4 train.py ...

```

This outputs the NCCL topology detection, ring formation, and communication statistics to the console, helping diagnose mismatched CUDA versions or network fabric issues.

## Summary

- **Launch with `torchrun`** to automatically configure `RANK`, `LOCAL_RANK`, and `WORLD_SIZE` for each process
- **Initialize NCCL** via `distributed.init_process_group(backend="nccl")` in [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py) (lines 1-20)
- **Convert to SyncBatchNorm** before wrapping with `DistributedDataParallel` using `bucket_cap_mb=32` and `static_graph=True` (lines 35-42)
- **Shard data automatically** via DALI's `num_shards=world_size` and `shard_id=rank` parameters in [`dali.py`](https://github.com/deepglint/rwkv-clip/blob/main/dali.py) (lines 56-60)
- **Use MXNet record format** (`.rec`/`.idx` files) required by the DALI-based loader
- **Enable debug logging** with `NCCL_DEBUG=INFO` to verify multi-GPU communication

## Frequently Asked Questions

### How do environment variables configure the NCCL process group in rwkv-clip?

The [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py) script reads `RANK`, `LOCAL_RANK`, and `WORLD_SIZE` from the environment at lines 16-19. These variables are automatically set by `torchrun` or manually when using `torch.multiprocessing.spawn`. The `LOCAL_RANK` determines which GPU the process uses via `torch.cuda.set_device()`, while `RANK` and `WORLD_SIZE` define the process's global identity and total participant count for the NCCL backend.

### Why does rwkv-clip require MXNet record files instead of standard image folders?

The data pipeline in [`dali.py`](https://github.com/deepglint/rwkv-clip/blob/main/dali.py) uses NVIDIA DALI's MXNet reader (`fn.readers.mxnet`) which expects `.rec` and `.idx` files. This format enables efficient sharding across GPUs using `num_shards` and `shard_id` parameters at lines 56-60. Convert your image-text dataset to MXNet record format before training to ensure compatibility with the distributed loader.

### What is the purpose of the SyncBatchNorm conversion in distributed training?

The `torch.nn.SyncBatchNorm.convert_sync_batchnorm()` call at line 36 in [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py) converts standard batch normalization layers to synchronized versions. This ensures that mean and variance statistics are aggregated across all GPUs during the forward pass, maintaining consistent normalization behavior regardless of batch size per GPU. Without this conversion, batch statistics would differ across devices, destabilizing training.

### Can I use mixed precision training with DDP in rwkv-clip?

Yes. The training script supports mixed precision through `torch.cuda.amp` or native bfloat16, controlled by the `--precision` argument. The DDP wrapper created at lines 35-42 handles gradient synchronization for all precision types automatically. Set `--precision bf16` or `--precision fp16` when launching to enable mixed precision training without additional configuration changes.