avoid copy in to_tensor when data is tensor (#67169)

zhiqiu · web-flow · commit 710ed51fd389 · 2024-08-08T11:10:53.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/api.py b/python/paddle/distributed/auto_parallel/api.py
@@ -188,12 +188,7 @@ def shard_tensor(
     if stop_gradient is None:
         stop_gradient = getattr(data, "stop_gradient", True)
 
-    if isinstance(data, EagerParamBase) and not data._is_initialized():
-        assert (
-            data._init_func is not None
-        ), "Get an uninitialized param with an unregistered init_func."
-        tensor = data
-    elif paddle.framework.in_pir_mode():
+    if paddle.framework.in_pir_mode():
         assert isinstance(
             data, (type(None), pir.Value)
         ), "input tensor is not pir value."
@@ -202,10 +197,19 @@ def shard_tensor(
         ), "shard_tensor() input data only supported dense tensor type right."
         tensor = data
     else:
-        # `paddle.to_tensor` supports both dynamic and static mode
-        tensor = paddle.to_tensor(
-            data, dtype=dtype, place=place, stop_gradient=stop_gradient
-        )
+        if isinstance(data, EagerParamBase) and not data._is_initialized():
+            assert (
+                data._init_func is not None
+            ), "Get an uninitialized param with an unregistered init_func."
+            tensor = data
+        elif isinstance(data, paddle.Tensor) and dtype is None:
+            # if place is not equal, it is handled in paddle.Tensor()
+            tensor = data
+        else:
+            # `paddle.to_tensor` supports both dynamic and static mode
+            tensor = paddle.to_tensor(
+                data, dtype=dtype, place=place, stop_gradient=stop_gradient
+            )
 
     if paddle.in_dynamic_mode():
         # here the dist tensor is deep copy constructed
