Cache DataParallel replicas

[apaszke]

This commits makes DataParallel cache the wrapped module replicas, and helps bring down overhead of replicate on ResNet1001 from ~140ms to 42ms. Remaining time is spent mostly in broadcast_coalesced, so moving it to C++ is the next step.

This commit really starts to push Python to the limit, which can be seen in two places:

• There's this ugly odict hack, because Python inheritance does weird things, and selects a much slower implementation of __getitem__ for OrderedDict subclasses than it could. I'm going to post to Python's mailing lists and clarify why is this happening, but I'm not aware of any reason why this hack would not work. This might seem like a silly thing, but removing this hack costs us 30ms at each forward.
• I was forced to change the implementation of torch.jit.compile for modules, and implement this poor man's inheritance-like thing (including the __instancecheck__ hack so these object still appear to belong to subclasses)... I've tried a few other things, but I can't come up with anything else that wouldn't break. The problem is that having one superclass with __slots__, and another one in C++ confuses Python, and it complains that it can't figure out how to lay them out in memory. I'm happy to discuss alternative solutions.

[apaszke]

I force pushed, but there are no significant code changes. Only the comments @ezyang requested and lint fixes

[ssnl]

Not about this PR, but in broadcast_coalesced, could there be a way to only broadcast modified parameters? A valid use case is finetuning only the last layer.

[apaszke]

To do this you'd need a mechanism to verify which parameters have changed, and it's not that simple. It's not enough to account for optimizers - someone might have loaded a different set of weights, and you need to propagate this change

[ssnl]

Good point. How about when doing inference with a model? Is there a way to avoid broadcasting tensors at all in PyTorch? Could this be an option of DataParallel?

[ezyang]

In general, I'd like a comment on top of DataParallel describing how the general replica strategy works. Otherwise LGTM.

[apaszke]

There is a comment. In general it's all invisible to the user and doesn't need to be. We're taking care of maintaining them in sync. The only thing that changed (and I forgot to update it) is that now __dict__s are shared among all replicas (previously they were shallow copies).

EDIT: I thought you meant the docs, but I think I misunderstood that and you just wanted a comment in the code. Will add this

[apaszke]

@ssnl I don't think so. If someone changes parameters you'd still need to re-broadcast them even in inference mode. In general, you need this "tensor modification hook" that we don't provide, and I'm not sure if we want to leak this logic that deep.

[colesbury]

I really like that we're speeding up DataParallel. I'm a bit concerned about the caching strategy. We're adding state to DataParallel, but still trying to pretend it's stateless.

Is there a way to speed up replicate without adding hidden state (the caching)? For example, we're already special casing _parameters, _buffers and _modules. Can we put a parent class of nn.Module in C++ that handles these three attributes? That would give use quick access to parameters and buffers for broadcasting. Replicate could then be light-weight: create a sort of proxy module that has its own _parameters, _buffers, and _modules but shares __dict__. Sub-modules could be lazily created on access to minimize the overhead at the start of a forward pass.

[ezyang]

I believe the new answer is to just DistributedDataParallel