construct dtensor with dtensor_from_local api (#60206)

pangengzheng · web-flow · commit c9d7f8b257fb · 2023-12-26T11:32:47.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/api.py b/python/paddle/distributed/auto_parallel/api.py
@@ -1309,15 +1309,13 @@ def build_distributed_tensor(local_tensor, dist_attr):
                     )
                 else:
                     raise ValueError(f"dim {dim} is not supported.")
-            # TODO(pangengzheng): construct dist_tensor with _dtensor_from_local api when it is ready.
-            global_tensor = paddle.zeros(global_shape, dtype=local_tensor.dtype)
             mesh = dist.ProcessMesh(
                 np.array(dist_attr["process_group"]).reshape(
                     dist_attr["process_shape"]
                 )
             )
             placements = to_placements(dist_attr["dims_mapping"], mesh)
-            dist_tensor = dist.shard_tensor(global_tensor, mesh, placements)
+            dist_tensor = dtensor_from_local(local_tensor, mesh, placements)
             assert (
                 dist_tensor._local_value().shape == local_tensor.shape
             ), f"local tensor shape {dist_tensor._local_value().shape} not equal to local_tensor.shape:{local_tensor.shape}"

Original file line number	Diff line number	Diff line change
`@@ -1309,15 +1309,13 @@ def build_distributed_tensor(local_tensor, dist_attr):`
`1309`	`1309`	`)`
`1310`	`1310`	`else:`
`1311`	`1311`	`raise ValueError(f"dim {dim} is not supported.")`
`1312`		`- # TODO(pangengzheng): construct dist_tensor with _dtensor_from_local api when it is ready.`
`1313`		`- global_tensor = paddle.zeros(global_shape, dtype=local_tensor.dtype)`
`1314`	`1312`	`mesh = dist.ProcessMesh(`
`1315`	`1313`	`np.array(dist_attr["process_group"]).reshape(`
`1316`	`1314`	`dist_attr["process_shape"]`
`1317`	`1315`	`)`
`1318`	`1316`	`)`
`1319`	`1317`	`placements = to_placements(dist_attr["dims_mapping"], mesh)`
`1320`		`- dist_tensor = dist.shard_tensor(global_tensor, mesh, placements)`
	`1318`	`+ dist_tensor = dtensor_from_local(local_tensor, mesh, placements)`
`1321`	`1319`	`assert (`
`1322`	`1320`	`dist_tensor._local_value().shape == local_tensor.shape`
`1323`	`1321`	`), f"local tensor shape {dist_tensor._local_value().shape} not equal to local_tensor.shape:{local_tensor.shape}"`