| import functools |
| import torch |
| import torch.distributed as dist |
| import torch.nn as nn |
| |
| class PythonDDP(nn.Module): |
| """ |
| Python only implementation for DistributedDataParallel module. |
| |
| Unlike the production DistributedDataParallel which relies on many C++ core |
| utils to manage gradient distribution and reduction. This class implement |
| all functions in pure Python such as param bucketing, gradient |
| synchronization and reduction. The only C++ dependency is the common utils: |
| ``dist.all_reduce`` |
| |
| The idea: parallelize gradient calculation and reduction, the same algo as |
| https://pytorch.org/docs/stable/notes/ddp.html, main steps: |
| 1. Distribute params into list of buckets. |
| 2. Register per-param hook to be invoked when grad is ready during backward |
| 3. In the hook, copy grad to corresponding bucket. If bucket is full, kick |
| off an async all_reduce operation to calculate average grad. |
| 4. After backward wait for all async ops to be done. Copy reduced grads back |
| to original places. |
| |
| Two modes are supported, asynchronous reduction (async_reduction=True) and |
| synchronous reduction (async_reduction=False) which shares the same algo as |
| LegacyDistributedDataParallel. |
| |
| Same as DistributedDataParallel to use this class , a process group needs to |
| be initiated. |
| |
| Example:: |
| |
| >>> torch.distributed.init_process_group( |
| >>> backend='gloo', world_size=N, init_method='...' |
| >>> ) |
| >>> pg = dist.distributed_c10d._get_default_group() |
| >>> async_reduction = True |
| >>> module = ToyModel() |
| >>> ddp_model = PythonDDP(module, pg, async_reduction) |
| >>> loss_fn = nn.MSELoss() |
| >>> optimizer = optim.SGD(ddp_model.parameters(), lr=0.001) |
| >>> outputs = ddp_model(torch.randn(20, 10).to(rank)) |
| >>> labels = torch.randn(20, 10).to(rank) |
| >>> loss_fn(outputs, labels).backward() |
| >>> |
| >>> # Reduce param grads |
| >>> ddp_model.all_reduce_grads() |
| >>> optimizer.step() |
| >>> |
| |
| """ |
| |
| class Bucket: |
| """Bucket is a container for list of params. """ |
| |
| def __init__(self, max_buffer_size): |
| self.param_to_offset = {} |
| self.buffer = None |
| self.ready_param_grad_count = 0 |
| self.total_elements = 0 |
| self._MAX_BUFFER_SIZE = max_buffer_size |
| |
| def __str__(self): |
| return "Bucket: num_params={}, total_elements={}, ready_param_grad_count={}".format( |
| len(self.param_to_offset), |
| self.total_elements, |
| self.ready_param_grad_count) |
| |
| def is_full(self): |
| """ |
| Returns whether grad for all the params in current bucket are ready |
| and copied to self.buffer. |
| """ |
| assert self.ready_param_grad_count >= 0 |
| assert self.ready_param_grad_count <= len(self.param_to_offset) |
| return len(self.param_to_offset) == self.ready_param_grad_count |
| |
| def empty(self): |
| self.ready_param_grad_count = 0 |
| |
| def try_hold_param(self, param): |
| """ |
| Checks whether current bucket has enough buffer to hold the incoming |
| param. If there is enough space, distribute param into current |
| bucket and Returns true. Otherwise, returns False. |
| """ |
| if self.total_elements + param.numel() <= self._MAX_BUFFER_SIZE : |
| self.param_to_offset[param] = self.total_elements |
| self.total_elements += param.numel() |
| return True |
| else: |
| return False |
| |
| def __init__(self, module, process_group, async_reduction=True, buffer_size=2 ** 22): |
| super(PythonDDP, self).__init__() |
| |
| self.module = module |
| self.process_group = process_group |
| self.world_size = dist.get_world_size(group=self.process_group) |
| self.async_reduction = async_reduction |
| |
| # Holds all_reduce handles, used when async_reduction is True |
| self.async_handles = set() |
| |
| # Ensure buffer_size is large enough to hold largest param. |
| max_numel = max(p.numel() for p in module.parameters()) |
| assert buffer_size > max_numel, "buffer_size: {} should be larger than largest param: {}".format(buffer_size, max_numel) |
| |
| # Build buckets for params |
| self.param_to_bucket, self.buckets = self._build_buckets_for_params(buffer_size) |
| |
| # Register per-parameter hook to be invoked when grad is ready. |
| for p in self.module.parameters(): |
| assert p.requires_grad |
| p.register_hook(functools.partial(self._on_param_grad_ready, p)) |
| |
| def _build_buckets_for_params(self, max_buffer_size): |
| """ |
| Distributes params into list of buckets. Maintains param -> bucket |
| mapping. Returns tuple of (param_to_buckets, buckets). |
| """ |
| print("_build_buckets_for_params called") |
| params_to_buckets = {} |
| buckets = set() |
| cur_bucket = self.Bucket(max_buffer_size) |
| total_param = 0 |
| for param in self.module.parameters(): |
| total_param += 1 |
| assert param.requires_grad, "param.requires_grad must be True" |
| if cur_bucket.try_hold_param(param): |
| params_to_buckets[param] = cur_bucket |
| buckets.add(cur_bucket) |
| else: |
| new_bucket = self.Bucket(max_buffer_size) |
| assert new_bucket.try_hold_param(param), "param must be holded in a empty bucket" |
| params_to_buckets[param] = new_bucket |
| buckets.add(new_bucket) |
| cur_bucket = new_bucket |
| |
| first_param = next(self.module.parameters()) |
| for bucket in buckets: |
| bucket.buffer = first_param.new(bucket.total_elements) |
| assert bucket.buffer is not None, 'bucket.buffer should not be None' |
| print("len(param_to_bucket)={}, len(buckets)={}".format( |
| len(params_to_buckets), len(buckets))) |
| |
| # Sanity check to ensure all params are distributed correctly into buckets |
| total_params_in_buckets = 0 |
| for bucket in buckets: |
| total_params_in_buckets += len(bucket.param_to_offset) |
| assert total_param == total_params_in_buckets |
| |
| return params_to_buckets, buckets |
| |
| # Callback when param.grad is ready. Note during callback, param.grad won't |
| # be ready yet, we MUST use the given ''grad'' which would be passed upon |
| # callback. |
| def _on_param_grad_ready(self, param, grad): |
| """ |
| Callback when grad for param is ready. Copy grad to its corresponding |
| bucket. When the bucket is full, kickoff an async all_reduce if |
| async_reduction is set, and adds the resultant handle to |
| self.async_handles. |
| |
| .. warning:: |
| Note param.grad isn't set yet. Use the passed grad instead. |
| """ |
| # Validate bucket and offset are set. |
| bucket = self.param_to_bucket.get(param) |
| assert bucket is not None, "Failed to find bucket for param" |
| offset = bucket.param_to_offset.get(param) |
| assert offset is not None, "offset must be set for param" |
| assert bucket.buffer is not None, "buffer must be allocated" |
| |
| # Copy grad to bucket, note param.grad isn't ready yet. |
| sz = param.numel() |
| assert grad is not None |
| assert param.requires_grad |
| assert param.numel() == grad.numel() |
| bucket.buffer[offset : offset + sz].copy_(grad.detach().view(-1)) |
| bucket.ready_param_grad_count += 1 |
| |
| # Kickoff grad reduction async when bucket is full. This ensures grad |
| # reduction and other grad calculation runs in parallel. |
| if self.async_reduction and bucket.is_full(): |
| bucket.buffer.div_(self.world_size) |
| handle = dist.all_reduce( |
| bucket.buffer, dist.ReduceOp.SUM, self.process_group, True) |
| self.async_handles.add(handle) |
| |
| def forward(self, *inputs, **kwargs): |
| return self.module(*inputs, **kwargs) |
| |
| def all_reduce_grads(self): |
| """ |
| Reduces all gradients across worker and updates param gradients. The |
| client should call this func post backward. |
| |
| If async_reduction is True, waits for all async handles (of all_reduce), |
| otherwise, kicks off synchrous all_reduce for all buckets. |
| |
| Once all all buckets are reduced, copy the reduced grads back to their |
| original parameters. After that, reset all buckets in prep for the next |
| iteration. |
| """ |
| if self.async_reduction: |
| for handle in self.async_handles: |
| handle.wait() |
| self.async_handles.clear() |
| else: |
| for bucket in self.buckets: |
| assert bucket.is_full() |
| bucket.buffer.div_(self.world_size) |
| dist.all_reduce(bucket.buffer, dist.ReduceOp.SUM, self.process_group) |
| |
| # Copy reduced-grad back into original place |
| for bucket in self.buckets: |
| assert bucket.is_full() |
| for cur_p, cur_offset in bucket.param_to_offset.items(): |
| sz = cur_p.numel() |
| if cur_p.grad is not None: |
| with torch.no_grad(): |
| cur_p.grad.copy_(bucket.buffer[cur_offset : cur_offset + sz].view_as(cur_p)) |
| else: |
| cur_p.grad = bucket.buffer[cur_offset : cur_offset + sz].view_as(cur_p).clone() |
| |
| # Empty bucket for next epoch |
| for bucket in self.buckets: |
| bucket.empty() |
