Support alternative distributed communication

[rblake-llnl]

Is your feature request related to a problem? Please describe.
I am working on a cluster that forbids opening sockets on a compute node that anyone can connect to, unauthenticated. This means I can't use the 'env://' initialization method for torch.distributed. Furthermore, I may want to train on systems that do not have CUDA, and therefore would like the option of using an mpi backend for torch.distributed.

Describe the solution you'd like

• Options to change the distributed backend for torch.dist
• Options to change the default bootstrapping to include things like shared filesystem support.

Describe alternatives you've considered

• I've debated coding up everything myself using horovod, but I don't like how it locks me down into MPI and won't let me use nccl afterwards. Using horovod also breaks pytorch-lightning's abstraction of handling such parallelism for you, and conflicts with the pytorch default of using the 'dp' backend when possible.

[williamFalcon]

@rblake-llnl

That's great. Would love a PR for this.

Is the best UX here to add a:

# default
Trainer(distributed_backend_config = {...})

Or use another flag called:

Trainer(distributed_process_group='nccl')

[williamFalcon]

so, your PR is great, but i wonder if we should make it a flag instead to access your implementation of that. the reason is that it will allow others to access it who may not know how to implement it.

second, perhaps a hook in the init_ddp method can make it even more customizable?

@neggert any thoughts?

[rblake-llnl]

Sure! I'm open to absolutely anything, I just took a first stab at the problem based on your UX above. I agree that it's a bit clunky.

You all have done a great job at designing a great UX for this library. Let me know what you want and I can code it up.

For my purposes, as long as I can call this method with dist_backend_config(['nccl'], {'init_method' : 'file:///path/to/file'}) I'll be happy in the short term.

Adding support for things like horovod would also be nice, but to do this you'd have to use horovod's DistributedSampler and horovod's DistributedDataParallel model decorator.

[neggert]

Would allowing the user to specify backend and init_method args that get passed to init_process_group be sufficient to accomplish everything we need? The trouble with passing through *args, **kwargs is that now users have to dig through multiple layers of documentation to figure out what values are allowed.

I'd propose adding a couple new args to Trainer that default to the current values.

Trainer(
    ...
    torch_dist_backend="nccl",
    torch_dist_init_method="env://",
)

Note that I don't love these names, so feel free to chime in with suggestions. distributed_backend is already taken.

I think there's another issue open asking for distributed CPU training. There are some other changes needed for that, but being able to set torch_dist_backend="gloo" is, I think, a necessary step.

[williamFalcon]

I like this suggestion, it might help to enumerate possible future support cases:

1. nccl
2. gloo
3. whatever new one comes up?

Non-slurm cluster or other ddp implementations wouldn't fall under this new flag (i think). This flag in particular is about configuring ddp?

Maybe it makes sense to do it as a hook then? In fact, the hook is how we might support arbitrary clusters?

def setup_distributed_backend():
   port = ...
   ...
   init_process_group(...)

[capdevc]

As far as I can tell, the only way to use DDP right now without SLURM is to manually construct and set the relevant SLURM environment variables (SLURM_NODEID and SLURM_NODELIST.. maybe more?). Otherwise, it looks like the current code overwrites whatever is set in the MASTER_ADDR environment variable with 127.0.0.2 and assumes that the current node is node 0.

While I'd love to have "set one flag and it works" support for things like Horovod, I think that #357 at least makes it possible to use DDP in the same enviornment as one would without Lightning, which would be great for kubernetes/sagemaker/databricks/etc. users.

I'd also like to see "set one flag and it works" support for things like Horovod. I'll probably make it work for my use case, but would be happy to turn that into a PR if given a little guidance on API

[williamFalcon]

@capdevc lightning does not overwrite MASTER_ADDR. If it is not passed it uses the default, otherwise it uses the user provided version.

Lightning always assumes the current node is 0 unless an environment var is present to change that.

Non SLURM systems would need a way to set this.

@capdevc if you want to go for it, i think horovod support would be great.

@rblake-llnl , @capdevc can you guys propose an API alternative which does the following:

1. support variations on DDP backend (nccl, etc...)
2. supports other ddp implementations.
3. final case is where the user still wants to do something on their own.

Currently i'm thinking:
Option 1

# default
Trainer(distributed_backend = 'ddp_nccl')
Trainer(distributed_backend = 'ddp_something_else')

Option 2

# default
Trainer(distributed_backend_config = {...}, distributed_backend='ddp')
Trainer(distributed_backend_config = {...}, distributed_backend='horovod')

And to support custom ddp add a hook:

def setup_distributed_backend():
   port = ...
   ...
   init_process_group(...)

[rblake-llnl]

Adding in something like horovod is more invasive-- it requires a wrapper around the optimizers and requires separate DistributedSampler calls that are horovod-aware.

Maybe we should think about having some sort of interface that all these models inherit from? Users could implement that interface to use their own distributed interface for things like Horovod. We could allow distributed_backend to remain as it is, but also allow it to be an object that implements the prescribed interface.

That way, people who want to use more exotic distributed_backend options could create and initialize objects themselves.

The resulting changes could also clean up the use_dp, use_ddp, and so on flags that are scattered throughout the source code and bleed out into userspace. Code like this in userspace:

        if self.use_ddp:
            train_sampler = DistributedSampler(dataset)

        should_shuffle = train_sampler is None
        loader = DataLoader(
            dataset=dataset,
            batch_size=batch_size,
            shuffle=should_shuffle,
            sampler=train_sampler,
            num_workers=0
        )

could become

        train_sampler = self.sampler(dataset)
        should_shuffle = train_sampler is None
        loader = DataLoader(
            dataset=dataset,
            batch_size=batch_size,
            shuffle=should_shuffle,
            sampler=train_sampler,
            num_workers=0
        )

[williamFalcon]

Great suggestion. I like the idea of a wrapper users can implement for more esoteric versions.
Could you draft a proposal for something like that? (quick, simple and at a high-level so we can modify it quickly before making it super robust)?

@jeffling (this might be relevant to you guys)

So, to repeat what I understood your vision is, have something like the Logger abstraction where people can do their own implementation and submit it if it's useful for the rest of the community?

This code is like this because users need the option to use DDP WITHOUT using distributedSampler. That came up a while back so we added support for it.

        if self.use_ddp:
            train_sampler = DistributedSampler(dataset)

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.