NCCL backend configuration

import os
import torch
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, TensorDataset
from torch.distributed import init_process_group, destroy_process_group

def configure_nccl():
    """
    Set NCCL environment variables before init_process_group().
    These must be set BEFORE importing torch.distributed initialization.
    """
    os.environ['NCCL_DEBUG'] = 'INFO'
    os.environ['NCCL_SOCKET_NTHREADS'] = '4'
    os.environ['NCCL_NSOCKS_PERTHREAD'] = '2'
    os.environ['NCCL_BLOCKING_WAIT'] = '1'
    os.environ['NCCL_ASYNC_ERROR_HANDLING'] = '1'
    os.environ['NCCL_SOCKET_IFNAME'] = 'eth0'
    os.environ['NCCL_TIMEOUT'] = '1800'
    os.environ['NCCL_ALGO'] = 'Ring'
    print(f"NCCL environment configured: DEBUG={os.environ['NCCL_DEBUG']}, TIMEOUT={os.environ['NCCL_TIMEOUT']}s")

def setup_distributed(rank, world_size, backend='nccl'):
    """
    Initialize distributed training with NCCL backend.
    rank: GPU index (0, 1, 2, ...)
    world_size: total number of GPUs
    """
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '29500'
    
    init_process_group(
        backend=backend,
        rank=rank,
        world_size=world_size,
        timeout=torch.distributed.timedelta(seconds=1800)
    )
    torch.cuda.set_device(rank)
    print(f"Rank {rank}: NCCL backend initialized")

def create_model_and_loader(rank, world_size, batch_size=32):
    """
    Create a simple model wrapped in DDP with NCCL backend.
    DDP automatically synchronizes gradients via NCCL all-reduce.
    """
    model = nn.Sequential(
        nn.Linear(10, 64),
        nn.ReLU(),
        nn.Linear(64, 2)
    ).to(rank)
    
    model = DDP(
        model,
        device_ids=[rank],
        process_group=torch.distributed.GroupMember.WORLD
    )
    
    fake_data = torch.randn(1000, 10)
    fake_labels = torch.randint(0, 2, (1000,))
    dataset = TensorDataset(fake_data, fake_labels)
    sampler = torch.utils.data.distributed.DistributedSampler(
        dataset,
        num_replicas=world_size,
        rank=rank,
        shuffle=True
    )
    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        sampler=sampler
    )
    return model, loader

def train_one_step(rank, world_size):
    """
    Single training step demonstrating NCCL gradient synchronization.
    """
    configure_nccl()
    setup_distributed(rank, world_size)
    
    model, loader = create_model_and_loader(rank, world_size)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    
    batch_x, batch_y = next(iter(loader))
    batch_x, batch_y = batch_x.to(rank), batch_y.to(rank)
    
    optimizer.zero_grad()
    logits = model(batch_x)
    loss = criterion(logits, batch_y)
    loss.backward()
    optimizer.step()
    
    if rank == 0:
        print(f"Training step complete. Loss: {loss.item():.4f}")
        print(f"NCCL synchronized gradients across {world_size} GPUs")
    
    destroy_process_group()

if __name__ == '__main__':
    print("NCCL Configuration Example")
    print("This script demonstrates environment variables needed for multi-GPU training.")
    print("In production, run with: torchrun --nproc_per_node=2 script.py")
    print("Key NCCL variables set:")
    print("  NCCL_DEBUG=INFO → logs all collective operations")
    print("  NCCL_TIMEOUT=1800 → 30 minute timeout for hanging collectives")
    print("  NCCL_SOCKET_NTHREADS=4 → worker threads for socket communication")
    print("  NCCL_ALGO=Ring → use ring reduction topology (alternatives: Tree)")