| 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) |