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

Code before:
            instrument_w_nvtx(self.quantization.quant_handle.wait)()
            flat_tensor = self.quantization.backend.dequantize(
                self.quantization.quantized_param, self.quantization.scale_buffer).to(self.params[0].device)

            self.partitions: List[Parameter] = []
            for i in range(self.world_size):
                self.partitions.append(
                    flat_tensor.narrow(0, self.quantization.partition_sz * i, self.quantization.partition_sz))

        # split the single tensor out into individual tensors
        param_offset = 0
        for param in self.params:
            assert param.ds_status == ZeroParamStatus.INFLIGHT, f"expected param {param.ds_summary()} to be inflight"
            partitions: List[Tensor] = []
            ds_tensor_numel = param.ds_tensor.ds_numel
            if self.use_secondary_tensor:
                ds_tensor_numel *= param.ds_secondary_tensor_num_of_groups
            for rank in range(self.world_size):
                param_start = rank * ds_tensor_numel
                if param_start < param.ds_numel:
                    part_to_copy = self.partitions[rank].narrow(0, param_offset,
                                                                min(param.ds_numel - param_start, ds_tensor_numel))
                    partitions.append(part_to_copy)
            param.data = instrument_w_nvtx(torch.cat)(partitions).view(param.ds_shape)
            param.ds_status = ZeroParamStatus.AVAILABLE
            if handle_dependency:
                for part_to_copy in partitions:
                    if not get_accelerator().is_synchronized_device():
                        part_to_copy.record_stream(get_accelerator().current_stream())

            param_offset += ds_tensor_numel

        self.complete = True
        if not handle_dependency:
            AllGatherCoalescedHandle.data_buffer.append(partitions)

    @staticmethod
    def free_buffer():
        AllGatherCoalescedHandle.data_buffer = []


class MultipleAllGatherHandles:

    def __init__(self, handles: List[AllGatherCoalescedHandle]):
        self.handles = handles

    def wait(self, handle_dependency=True) -> None:
        for handle in self.handles:
            handle.wait(handle_dependency)



Review comment (suggestion):
Would the following not work?

```suggestion
            if handle_dependency and not get_accelerator().is_synchronized_device() :
```

Code after:
            instrument_w_nvtx(self.quantization.quant_handle.wait)()
            flat_tensor = self.quantization.backend.dequantize(
                self.quantization.quantized_param, self.quantization.scale_buffer).to(self.params[0].device)

            self.partitions: List[Parameter] = []
            for i in range(self.world_size):
                self.partitions.append(
                    flat_tensor.narrow(0, self.quantization.partition_sz * i, self.quantization.partition_sz))

        # split the single tensor out into individual tensors
        param_offset = 0
        for param in self.params:
            assert param.ds_status == ZeroParamStatus.INFLIGHT, f"expected param {param.ds_summary()} to be inflight"
            partitions: List[Tensor] = []
            ds_tensor_numel = param.ds_tensor.ds_numel
            if self.use_secondary_tensor:
                ds_tensor_numel *= param.ds_secondary_tensor_num_of_groups
            for rank in range(self.world_size):
                param_start = rank * ds_tensor_numel
                if param_start < param.ds_numel:
                    part_to_copy = self.partitions[rank].narrow(0, param_offset,
                                                                min(param.ds_numel - param_start, ds_tensor_numel))
                    partitions.append(part_to_copy)
            param.data = instrument_w_nvtx(torch.cat)(partitions).view(param.ds_shape)
            param.ds_status = ZeroParamStatus.AVAILABLE
            if not get_accelerator().is_synchronized_device() and handle_dependency:
                for part_to_copy in partitions:
                    part_to_copy.record_stream(get_accelerator().current_stream())

            param_offset += ds_tensor_numel

        self.complete = True
        if not get_accelerator().is_synchronized_device() and not handle_dependency:
            # if the device needs to handle dependencies and opts for explicit processing outside the function.
            AllGatherCoalescedHandle.data_buffer.append(partitions)

    @staticmethod
    def free_buffer():
        AllGatherCoalescedHandle.data_buffer = []


class MultipleAllGatherHandles:

    def __init__(self, handles: List[AllGatherCoalescedHandle]):
        self.handles = handles

    def wait(self, handle_dependency=True) -> None:
        for handle in self.handles:
            handle.wait(handle_dependency)