Repository: microsoft/DeepSpeed
File: deepspeed/runtime/zero/partition_parameters.py (Python)

Code before:

        # Wait ensures the operation is enqueued, but not necessarily complete.
        launch_handles[-1].wait()
        if quantize:
            for quant_handle in launch_quantize_handles:
                quant_handle.wait()

        # assign to param.data (not copy)
        for i, param in enumerate(param_list):
            gathered_tensor = allgather_params[i]
            if quantize:
                gathered_tensor = self.quantizer_module.dequantize(gathered_tensor, allgather_quantize_scale[i])
            param.data = gathered_tensor.narrow(0, 0, param.ds_numel).view(param.ds_shape).data

        # guarantee the communication to be completed
        if not get_accelerator().resolves_data_dependency():
            get_accelerator().synchronize()

        return None

    @torch.no_grad()
    def _allgather_params(self, param_list, hierarchy=0):
        if len(param_list) == 0:
            return

        if self.allgather_single_param:
            for param in param_list:
                partition_size = param.ds_tensor.ds_numel
                tensor_size = partition_size * self.num_partitions

                flat_tensor = torch.empty(tensor_size, dtype=param.ds_tensor.dtype, device=self.local_device)
                if self.use_all_gather_into_tensor:
                    dist.all_gather_into_tensor(flat_tensor,
                                                param.ds_tensor.to(get_accelerator().device_name()),
                                                group=self.get_partition_dp_group(param),
                                                async_op=False)
                else:
                    partitions = []
                    for i in range(self.num_partitions):
                        partitions.append(flat_tensor.narrow(0, partition_size * i, partition_size))
                        if i == self.get_partition_rank():
                            partitions[i].data.copy_(param.ds_tensor.data, non_blocking=True)
                    dist.all_gather(partitions,
                                    partitions[self.get_partition_rank()],
                                    group=self.get_partition_dp_group(param),
                                    async_op=False)

                if hasattr(param.ds_tensor, 'ds_quant_scale'):
                    scale_size = param.ds_tensor.ds_quant_scale.numel()
                    scale_tensor_size = scale_size * self.num_partitions
                    flat_scale_tensor = torch.empty(scale_tensor_size,

Review comment (refactor):
It seems the same refactoring opportunity exists here:
1. Refactor current logic into a helper that processes list of params
2. Use helper appropriatly for both cases

```python
     if self.allgather_single_param:
         for param in param_list:
              _all_gather_params_helper([param], ...)
     else:
           _all_gather_params_helper(param_list, ...)
```

Code after:
                quant_handle.wait()

        # assign to param.data (not copy)
        for i, param in enumerate(param_list):
            gathered_tensor = allgather_params[i]
            if quantize:
                gathered_tensor = self.quantizer_module.dequantize(gathered_tensor, allgather_quantize_scale[i])
            param.data = gathered_tensor.narrow(0, 0, param.ds_numel).view(param.ds_shape).data

        # guarantee the communication to be completed
        if not get_accelerator().resolves_data_dependency():
            get_accelerator().synchronize()

        return None

    def _allgather_params_sequential(self, param_list, hierarchy=0):
        if len(param_list) == 0:
            return

        for param in param_list:
            partition_size = param.ds_tensor.ds_numel
            tensor_size = partition_size * self.num_partitions

            flat_tensor = torch.empty(tensor_size, dtype=param.ds_tensor.dtype, device=self.local_device)
            flat_tensor.requires_grad = False
            if self.use_all_gather_into_tensor:
                dist.all_gather_into_tensor(flat_tensor,
                                            param.ds_tensor.to(get_accelerator().device_name()),
                                            group=self.get_partition_dp_group(param),
                                            async_op=False)
            else:
                partitions = []
                for i in range(self.num_partitions):
                    partitions.append(flat_tensor.narrow(0, partition_size * i, partition_size))
                    if i == self.get_partition_rank():
                        partitions[i].data.copy_(param.ds_tensor.data, non_blocking=True)
                dist.all_gather(partitions,
                                partitions[self.get_partition_rank()],
                                group=self.get_partition_dp_group(param),
                                async_op=False)

            if hasattr(param.ds_tensor, 'ds_quant_scale'):
                scale_size = param.ds_tensor.ds_quant_scale.numel()
                scale_tensor_size = scale_size * self.num_partitions
                flat_scale_tensor = torch.empty(scale_tensor_size,
                                                dtype=param.ds_tensor.ds_quant_scale.dtype,
                                                device=self.local_device)
                flat_scale_tensor.requires_grad = False
                if self.use_all_gather_into_tensor:
                    dist.all_gather_into_tensor(flat_scale_tensor,
                                                param.ds_tensor.ds_quant_scale.to(get_accelerator().device_name()),