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

Code before:
            self._release_sub_group(sub_group_id, timer_names)

    def checkpoint_event_prologue(self):
        self._partition_all_parameters()

    def checkpoint_event_epilogue(self):
        if len(self.persistent_parameters) > 0:
            self.persistent_parameters[0].all_gather(self.persistent_parameters)

    def empty_partition_cache(self):
        self.parameter_offload.empty_partition_cache()

    def offload_states(self,
                       include: Container[OffloadStateTypeEnum] = None,
                       device: OffloadDeviceEnum = OffloadDeviceEnum.cpu,
                       pin_memory: bool = True,
                       non_blocking: bool = False):
        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 = [

Review comment (suggestion):
I think it would be faster and reduce internal fragmentation to pin one large buffer instead of many smaller buffers. We could then use `narrow` to partition the large buffer. 

Code after:
        if len(self.persistent_parameters) > 0:
            self.persistent_parameters[0].all_gather(self.persistent_parameters)

    def empty_partition_cache(self):
        self.parameter_offload.empty_partition_cache()

    def offload_states(self,
                       include: Container[OffloadStateTypeEnum] = None,
                       device: OffloadDeviceEnum = OffloadDeviceEnum.cpu,
                       pin_memory: bool = True,
                       non_blocking: bool = False):
        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_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(