This version of DistributedDataParallel is designed to be used in conjunction with the multiproc.py launcher included with this example. It assumes that your run is using multiprocess with 1 GPU/process, that the model is on the correct device, and that torch.set_device has been use
(module)
| 81 | |
| 82 | |
| 83 | def apply_gradient_allreduce(module): |
| 84 | """ |
| 85 | This version of DistributedDataParallel is designed to be used in conjunction with the multiproc.py |
| 86 | launcher included with this example. It assumes that your run is using multiprocess with 1 |
| 87 | GPU/process, that the model is on the correct device, and that torch.set_device has been |
| 88 | used to set the device. |
| 89 | Parameters are broadcasted to the other processes on initialization of DistributedDataParallel, |
| 90 | and will be all-reduced at the finish of the backward pass. |
| 91 | |
| 92 | Note by Rick: Compared to the simple pytorch's DDP wrapper on modules, which synchronizes at both forward and backward, |
| 93 | wrapping the method around a module only synchronize the gradient (not the output)! This is more efficient when |
| 94 | multiple networks are concatenated together, and if we don't care about the intermediate inputs from all threads. |
| 95 | Also, the method does NOT synchronize buffers (so batchnorm's statistics are on a smaller dateset). |
| 96 | """ |
| 97 | if not hasattr(dist, "_backend"): |
| 98 | module.warn_on_half = True |
| 99 | else: |
| 100 | module.warn_on_half = True if dist._backend == dist.dist_backend.GLOO else False |
| 101 | |
| 102 | for p in module.state_dict().values(): |
| 103 | if not torch.is_tensor(p): |
| 104 | continue |
| 105 | dist.broadcast(p, 0) |
| 106 | |
| 107 | def allreduce_params(): |
| 108 | if module.needs_reduction: |
| 109 | module.needs_reduction = False |
| 110 | buckets = {} |
| 111 | for param in module.parameters(): |
| 112 | if param.requires_grad and param.grad is not None: |
| 113 | tp = param.data.dtype |
| 114 | if tp not in buckets: |
| 115 | buckets[tp] = [] |
| 116 | buckets[tp].append(param) |
| 117 | if module.warn_on_half: |
| 118 | if torch.cuda.HalfTensor in buckets: |
| 119 | print( |
| 120 | "WARNING: gloo dist backend for half parameters may be extremely slow." |
| 121 | + " It is recommended to use the NCCL backend in this case. This currently requires" |
| 122 | + "PyTorch built from top of tree master." |
| 123 | ) |
| 124 | module.warn_on_half = False |
| 125 | |
| 126 | for tp in buckets: |
| 127 | bucket = buckets[tp] |
| 128 | grads = [param.grad.data for param in bucket] |
| 129 | coalesced = _flatten_dense_tensors(grads) |
| 130 | dist.all_reduce(coalesced) |
| 131 | coalesced /= dist.get_world_size() |
| 132 | for buf, synced in zip(grads, _unflatten_dense_tensors(coalesced, grads)): |
| 133 | buf.copy_(synced) |
| 134 | |
| 135 | for param in list(module.parameters()): |
| 136 | |
| 137 | def allreduce_hook(*unused): |
| 138 | Variable._execution_engine.queue_callback(allreduce_params) |
| 139 | |
| 140 | if param.requires_grad: |
no test coverage detected