Method full

deepspeed/runtime/utils.py:715–732

(self, device=None)

        return tensor_part, partition

    def full(self, device=None):
        if device is None:
            device = self.orig_device

        # Allocate the full tensor as a flat buffer.
        full_numel = prod(self.full_size())
        flat_tensor = torch.zeros([full_numel], dtype=self.local_data.dtype, device=device)
        if self.even_split:
            # Collect the full tensor
            dist.all_gather_into_tensor(flat_tensor, self.local_data, group=self.group)
        else:
            for part_id in range(self.num_parts):
                part_size = self.partition[part_id + 1] - self.partition[part_id]
                buf = flat_tensor.narrow(0, start=self.partition[part_id], length=part_size)
                if part_id == self.rank:
                    buf.copy_(self.local_data)
                dist.broadcast(buf, part_id, self.group)
        return flat_tensor.view(self.full_size()).clone().detach()

    def to_meta(self):
        """Returns a torch.LongTensor that encodes partitioning information.