| 666 | |
| 667 | |
| 668 | class PartitionedTensor: |
| 669 | |
| 670 | def __init__(self, tensor, group, partition_meta=None): |
| 671 | super().__init__() |
| 672 | |
| 673 | self.group = group |
| 674 | self.num_parts = dist.get_world_size(group=self.group) |
| 675 | self.rank = dist.get_rank(group=self.group) |
| 676 | self.orig_size = list(tensor.size()) |
| 677 | self.orig_device = tensor.device |
| 678 | self.local_data, self.partition = self._partition_tensor(tensor) |
| 679 | self.even_split = tensor.numel() % self.num_parts == 0 |
| 680 | |
| 681 | @classmethod |
| 682 | def from_meta(cls, meta, local_part, group, device=get_accelerator().device_name()): |
| 683 | assert meta.dtype == torch.long |
| 684 | dummy = torch.ones(dist.get_world_size(group=group)) |
| 685 | part_obj = cls(tensor=dummy, group=group) |
| 686 | |
| 687 | meta = meta.tolist() |
| 688 | |
| 689 | # [N, list0, ..., listN-1] |
| 690 | part_obj.orig_size = meta[1:(1 + meta[0])] |
| 691 | meta = meta[1 + meta[0]:] |
| 692 | |
| 693 | part_obj.orig_device = device |
| 694 | part_obj.local_data = local_part.detach() |
| 695 | |
| 696 | part_obj.group = group |
| 697 | |
| 698 | # Partition is encoded like the rowptr of a CSR matrix: |
| 699 | # [num_parts, rank, 0, part_1, ..., part_num_parts] |
| 700 | # TODO: support shuffle between different partition granularities |
| 701 | assert part_obj.num_parts == meta[0] |
| 702 | assert part_obj.rank == meta[1] |
| 703 | part_obj.partition = meta[2:] # length num_parts+1 |
| 704 | |
| 705 | return part_obj |
| 706 | |
| 707 | def _partition_tensor(self, tensor): |
| 708 | partition = partition_uniform(num_items=tensor.numel(), num_parts=self.num_parts) |
| 709 | start = partition[self.rank] |
| 710 | length = partition[self.rank + 1] - start |
| 711 | tensor_part = tensor.detach().contiguous().view(-1).narrow(0, start=start, length=length).clone() |
| 712 | |
| 713 | return tensor_part, partition |
| 714 | |
| 715 | def full(self, device=None): |
| 716 | if device is None: |
| 717 | device = self.orig_device |
| 718 | |
| 719 | # Allocate the full tensor as a flat buffer. |
| 720 | full_numel = prod(self.full_size()) |
| 721 | flat_tensor = torch.zeros([full_numel], dtype=self.local_data.dtype, device=device) |
| 722 | if self.even_split: |
| 723 | # Collect the full tensor |
| 724 | dist.all_gather_into_tensor(flat_tensor, self.local_data, group=self.group) |
| 725 | else: |
no outgoing calls