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

Code before:

            # Move CPU flat tensor to the accelerator memory.
            self.bit16_groups_flat.append(flattened_buffer.to(get_accelerator().current_device_name()))
            del flattened_buffer

            see_memory_usage(f"After flattening and moving param group {i} to GPU", force=False)

            if dist.get_rank(group=self.real_dp_process_group[i]) == 0:
                see_memory_usage(f"After Flattening and after emptying param group {i} cache", force=False)

            # set model bit16 weight to slices of flattened buffer
            self._update_model_bit16_weights(i)

            # divide the flat weights into near equal partition equal to the data parallel degree
            # each process will compute on a different part of the partition
            data_parallel_partitions = self.get_data_parallel_partitions(self.bit16_groups_flat[i], i)
            self.parallel_partitioned_bit16_groups.append(data_parallel_partitions)

            # Record padding required for alignment
            left_boundary = sum([t.numel() for t in data_parallel_partitions[:partition_id]])
            curr_partition_size = data_parallel_partitions[partition_id].numel()

            if orig_group_numel <= left_boundary:
                padding = curr_partition_size
            elif orig_group_numel < left_boundary + curr_partition_size:
                padding = orig_group_numel - left_boundary
            else:
                padding = 0
            self.groups_padding.append(padding)

            # verify that data partition start locations are 4-byte aligned
            for partitioned_data in data_parallel_partitions:
                assert (partitioned_data.data_ptr() % (2 * self.nccl_start_alignment_factor) == 0)

            # A partition of the fp32 master weights that will be updated by this process.
            # Note that the params in single_partition_of_fp32_groups is cloned and detached
            # from the origin params of the model.
            if not fp16_master_weights_and_gradients:
                weights_partition = self.parallel_partitioned_bit16_groups[i][partition_id].to(
                    self.device).clone().float().detach()
            else:
                weights_partition = self.parallel_partitioned_bit16_groups[i][partition_id].to(
                    self.device).clone().half().detach()

            if self.cpu_offload:
                weights_partition = get_accelerator().pin_memory(weights_partition)

            self.single_partition_of_fp32_groups.append(weights_partition)

            # Set local optimizer to have flat params of its own partition.
            # After this, the local optimizer will only contain its own partition of params.

Review comment (suggestion):
Please double check that it shouldn't be 
```suggestion
                padding = (left_boundary + curr_partition_size) - orig_group_numel
```

Code after:

            # Move CPU flat tensor to the accelerator memory.
            self.bit16_groups_flat.append(flattened_buffer.to(get_accelerator().current_device_name()))
            del flattened_buffer

            see_memory_usage(f"After flattening and moving param group {i} to GPU", force=False)

            if dist.get_rank(group=self.real_dp_process_group[i]) == 0:
                see_memory_usage(f"After Flattening and after emptying param group {i} cache", force=False)

            # set model bit16 weight to slices of flattened buffer
            self._update_model_bit16_weights(i)

            # divide the flat weights into near equal partition equal to the data parallel degree
            # each process will compute on a different part of the partition
            data_parallel_partitions = self.get_data_parallel_partitions(self.bit16_groups_flat[i], i)
            self.parallel_partitioned_bit16_groups.append(data_parallel_partitions)

            # Record padding required for alignment
            left_boundary = sum([t.numel() for t in data_parallel_partitions[:partition_id]])
            curr_partition_size = data_parallel_partitions[partition_id].numel()

            if orig_group_numel <= left_boundary:
                padding = curr_partition_size
            elif orig_group_numel < left_boundary + curr_partition_size:
                padding = left_boundary + curr_partition_size - orig_group_numel
            else:
                padding = 0
            self.groups_padding.append(padding)

            # verify that data partition start locations are 4-byte aligned
            for partitioned_data in data_parallel_partitions:
                assert (partitioned_data.data_ptr() % (2 * self.nccl_start_alignment_factor) == 0)

            # A partition of the fp32 master weights that will be updated by this process.
            # Note that the params in single_partition_of_fp32_groups is cloned and detached
            # from the origin params of the model.
            if not fp16_master_weights_and_gradients:
                weights_partition = self.parallel_partitioned_bit16_groups[i][partition_id].to(
                    self.device).clone().float().detach()
            else:
                weights_partition = self.parallel_partitioned_bit16_groups[i][partition_id].to(
                    self.device).clone().half().detach()

            if self.cpu_offload:
                weights_partition = get_accelerator().pin_memory(weights_partition)

            self.single_partition_of_fp32_groups.append(weights_partition)

            # Set local optimizer to have flat params of its own partition.
            # After this, the local optimizer will only contain its own partition of params.