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

Code before:
        device = device.value

        self.empty_partition_cache()

        assert self.optimizer.__class__ == deepspeed.ops.adam.fused_adam.FusedAdam, f"Offloading is supported only for DeepSpeed FusedAdam."

        def needs_offload(target):
            # return True
            return target not in self.offloaded_states and (include == None or target in include)

        # HP param
        if needs_offload(OffloadStateTypeEnum.hp_params):
            if pin_memory:
                if not hasattr(self, "hp_params_pin_buffers"):
                    self.hp_params_pin_buffers = [
                        get_accelerator().pin_memory(torch.empty_like(t, device=device))
                        for t in self.fp32_partitioned_groups_flat
                    ]

                for src_tensor, dest_buf in zip(self.fp32_partitioned_groups_flat, self.hp_params_pin_buffers):
                    dest_buf.copy_(src_tensor, non_blocking=non_blocking)
                    src_tensor.data = dest_buf
            else:
                for buf in self.fp32_partitioned_groups_flat:
                    buf.data = buf.data.to(device, non_blocking=non_blocking)
            self.offloaded_states.add(OffloadStateTypeEnum.hp_params)

        # LP param
        if needs_offload(OffloadStateTypeEnum.lp_params):
            if pin_memory:
                if not hasattr(self, "lp_param_contiguous_pin_buffer"):
                    self.lp_param_contiguous_pin_buffer = get_accelerator().pin_memory(
                        torch.empty_like(self.lp_param_buffer, device=device))
                    self.lp_params_pin_buffers = [
                        get_accelerator().pin_memory(torch.empty_like(p.ds_tensor, device=device))
                        for p in self.module.parameters()
                    ]
                self.lp_param_contiguous_pin_buffer.copy_(self.lp_param_buffer, non_blocking=non_blocking)
                self.lp_param_buffer.data = self.lp_param_contiguous_pin_buffer

                for p, buf in zip(self.module.parameters(), self.lp_params_pin_buffers):
                    buf.copy_(p.ds_tensor.data, non_blocking=non_blocking)
                    p.ds_tensor.data = buf
            else:
                self.lp_param_buffer.data = self.lp_param_buffer.to(device, non_blocking=non_blocking)
                for p in self.module.parameters():
                    p.ds_tensor.data = p.ds_tensor.data.to(device, non_blocking=non_blocking)

            self.fp16_partitioned_groups_flat.clear()
            self.offloaded_states.add(OffloadStateTypeEnum.lp_params)


Review comment (suggestion):
I don't understand the need for two pinned buffers `self.lp_param_contiguous_pin_buffer` and `self.lp_params_pin_buffers`. Let's chat about this. 

Code after:
        def needs_offload(target):
            # return True
            return target not in self.offloaded_states and (include == None or target in include)

        # HP param
        if needs_offload(OffloadStateTypeEnum.hp_params):
            if pin_memory:
                if not hasattr(self, "hp_params_pin_buffers"):
                    self.hp_params_pin_buffers = [
                        get_accelerator().pin_memory(torch.empty_like(t, device=device))
                        for t in self.fp32_partitioned_groups_flat
                    ]

                for src_tensor, dest_buf in zip(self.fp32_partitioned_groups_flat, self.hp_params_pin_buffers):
                    dest_buf.copy_(src_tensor, non_blocking=non_blocking)
                    src_tensor.data = dest_buf
            else:
                for buf in self.fp32_partitioned_groups_flat:
                    buf.data = buf.data.to(device, non_blocking=non_blocking)
            self.offloaded_states.add(OffloadStateTypeEnum.hp_params)

        # LP param
        if needs_offload(OffloadStateTypeEnum.lp_params):
            if pin_memory:
                if not hasattr(self, "lp_param_contiguous_pin_buffer"):
                    self.lp_param_contiguous_pin_buffer = get_accelerator().pin_memory(
                        torch.empty_like(self.lp_param_buffer, device=device))
                self.lp_param_contiguous_pin_buffer.copy_(self.lp_param_buffer, non_blocking=non_blocking)
                cpu_buffer = self.lp_param_contiguous_pin_buffer
            else:
                cpu_buffer = self.lp_param_buffer.to(device, non_blocking=non_blocking)

            self.lp_param_buffer.data = cpu_buffer
            for tensor, offset, tensor_numel in get_mapping_to_flat_buffer(
                [p.ds_tensor for p in self.module.parameters()]):
                tensor.data = cpu_buffer.narrow(0, offset, tensor_numel)

            self.fp16_partitioned_groups_flat.clear()
            self.offloaded_states.add(OffloadStateTypeEnum.lp_params)

        # LP grad
        if needs_offload(OffloadStateTypeEnum.lp_grads):
            if pin_memory:
                if not hasattr(self, "lp_grad_partitions_flat_pin_buffers"):
                    self.lp_grad_partitions_flat_pin_buffers = get_accelerator().pin_memory(
                        torch.empty_like(self.grad_partitions_flat_buffer, device=device))
                self.lp_grad_partitions_flat_pin_buffers.copy_(self.grad_partitions_flat_buffer,
                                                               non_blocking=non_blocking)
                self.grad_partitions_flat_buffer.data = self.lp_grad_partitions_flat_pin_buffers
            else:
                self.grad_partitions_flat_buffer.data = self.grad_partitions_flat_buffer.data.to(device)