Simple NCCL example with TCP/IP backend
Simple NCCL with TCP/IP example
What is your question?
I can't find out how to "add all_reduce" so lightning would be usable in a very simple NCCL TCP/IP environment. I have a very simple pytorch program that does exactly that; I can't see how I could integrate this with lightning (even though it might be that there is a really simple solution). I like the concepts of lightning very much but what I am missing is something like a "custom all_reduce" that's being called after each sample / batch / epoch (I think it might depend which one makes most sense). Here's a RUNNING example of what I would like to integrate, and here's just the basic structure:
init_method = 'tcp://192.168.1.154:23456'
backend = 'nccl'
world_size = 2 # n
manual_seed = 12345

class Partition(object):
    """ Dataset partitioning helper """
    ………
class DataPartitioner(object):
    """ Dataset partitioning """
    ………    

class Net(nn.Module):
        ………
    
def average_gradients(model):
    """ Distributed Gradient averaging. """
    size = float(dist.get_world_size())
    for param in model.parameters():
        dist.all_reduce(param.grad.data, op=torch.distributed.ReduceOp.SUM)
        param.grad.data /= size

def run(rank, size):
    torch.manual_seed(1234)
    ………
    for epoch in range(10):
        epoch_loss = 0.0
        for data, target in train_set:
            ………
        print('Rank ', dist.get_rank(), ', epoch ',
              epoch, ': ', epoch_loss / num_batches)
        average_gradients(model)
        
def start(rank, world_size):
    dist.init_process_group(
        backend=backend,
        init_method=init_method,
        rank=rank,
        world_size=world_size
    )
    run(rank, world_size)

rank = int(sys.argv[1])
size = int(sys.argv[2])
start(rank, size)
What have you tried?
This can easily be run (Node 0: #python3 whatever.py 0 2; Node 0: #python3 whatever.py 1 2 ) and given nccl is available and both nodes "see" each other on a TCP/IP level this is working.
I can see how I can just make the model a pl.LightningModule, data loading is pretty much how Lightning likes it, too. I simply have no clue how to get the average_gradients method into Lightning so it would use dist.all_reduce(param.grad.data, op=torch.distributed.ReduceOp.SUM)... The example is mostly from PyTorch distributed docs.
What's your environment?

OS: Linux
Packaging pip, conda, nvidia-docker-runtime
Version: cloned 2020/05/27 from GitHub; pip install .

Any help would be very appreciated.