Converts a distributed tensor to a dense tensor. ``unshard_dtensor`` first make the ``dist_tensor`` be ``Replicate`` state on all processes and then converts it to a dense ``paddle.Tensor``. It can be treated as a reverse operation of ``shard_tensor``. Args: dist_tensor
(dist_tensor: Tensor)
| 3833 | |
| 3834 | |
| 3835 | def unshard_dtensor(dist_tensor: Tensor) -> Tensor: |
| 3836 | """ |
| 3837 | Converts a distributed tensor to a dense tensor. ``unshard_dtensor`` |
| 3838 | first make the ``dist_tensor`` be ``Replicate`` state on all processes and |
| 3839 | then converts it to a dense ``paddle.Tensor``. It can be treated as a |
| 3840 | reverse operation of ``shard_tensor``. |
| 3841 | |
| 3842 | Args: |
| 3843 | dist_tensor (paddle.Tensor): The distributed tensor which is constructed |
| 3844 | from a dense tensor with ``shard_tensor``. |
| 3845 | |
| 3846 | Returns: |
| 3847 | paddle.Tensor: The original dense tensor of the input ``dist_tensor``. |
| 3848 | |
| 3849 | Examples: |
| 3850 | .. code-block:: pycon |
| 3851 | |
| 3852 | >>> import paddle |
| 3853 | >>> import paddle.distributed as dist |
| 3854 | >>> from paddle.distributed import Replicate, Shard |
| 3855 | |
| 3856 | >>> # doctest: +REQUIRES(env:DISTRIBUTED) |
| 3857 | >>> mesh = dist.ProcessMesh([0, 1], dim_names=["x"]) |
| 3858 | >>> original_tensor = paddle.rand([4, 1024, 512]) |
| 3859 | >>> dist_tensor = dist.shard_tensor(original_tensor, mesh, [Shard(0)]) |
| 3860 | >>> # dense_tensor's shape is the same as original_tensor |
| 3861 | >>> dense_tensor = dist.unshard_dtensor(dist_tensor) |
| 3862 | """ |
| 3863 | if paddle.in_dynamic_mode(): |
| 3864 | # if the input is not a distributed |
| 3865 | # tensor, return it directly |
| 3866 | if dist_tensor.is_dist() is False: |
| 3867 | raise ValueError("The input should be a distributed tensor.") |
| 3868 | |
| 3869 | mesh = dist_tensor.process_mesh |
| 3870 | placements = dist_tensor.placements |
| 3871 | replicate_placements = [dist.Replicate()] * len(placements) |
| 3872 | r_dist_tensor = reshard(dist_tensor, mesh, replicate_placements) |
| 3873 | |
| 3874 | if isinstance(dist_tensor, EagerParamBase): |
| 3875 | return EagerParamBase.from_tensor( |
| 3876 | r_dist_tensor._local_value(), |
| 3877 | **dist_tensor.__dict__, |
| 3878 | ) |
| 3879 | else: |
| 3880 | return paddle.Tensor(r_dist_tensor._local_value()) |
| 3881 | |
| 3882 | elif paddle.framework.in_pir_mode(): |
| 3883 | # in pir mode, we define the logic of unshard_tensor as dist_tensor_type --> dense_tensor_type with global shape. |
| 3884 | dense_tensor_type = paddle.pir.create_shaped_type( |
| 3885 | dist_tensor.type(), dist_tensor.shape |
| 3886 | ) |
| 3887 | dist_tensor.set_type(dense_tensor_type) |
| 3888 | |
| 3889 | return dist_tensor |
| 3890 | |
| 3891 | else: |
| 3892 | raise NotImplementedError( |
nothing calls this directly
no test coverage detected