[Auto Parallel] Enable to_static unittest (#69326)

deepllz · web-flow · commit 1dbe1253c52f · 2024-11-13T14:53:17.000+08:00
* [Auto Parallel] Enable to_static unittest

* [Auto Parallel] Enable to_static unittest
diff --git a/python/paddle/distributed/auto_parallel/intermediate/parallel_base.py b/python/paddle/distributed/auto_parallel/intermediate/parallel_base.py
@@ -12,12 +12,29 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import paddle
 import paddle.distributed as dist
+from paddle import pir
+from paddle.base.framework import (
+    in_dygraph_mode,
+    in_pir_mode,
+)
 from paddle.distributed import fleet
 from paddle.nn import Layer
 from paddle.optimizer import Optimizer
 
 
+def is_tensor(tensor):
+    if in_dygraph_mode():
+        return isinstance(tensor, paddle.Tensor)
+    elif in_pir_mode():
+        return isinstance(tensor, pir.Value)
+    else:
+        raise RuntimeError(
+            "PipelineParallel are only supported in dynamic or pir mode."
+        )
+
+
 class ParallelOptimizer:
     def __init__(self, optimizer, level=None):
         self.level = None
diff --git a/python/paddle/distributed/auto_parallel/intermediate/pipeline_parallel.py b/python/paddle/distributed/auto_parallel/intermediate/pipeline_parallel.py
@@ -15,11 +15,10 @@
 from collections import OrderedDict
 from enum import Enum
 
-import paddle
 import paddle.distributed as dist
 from paddle.distributed import fleet
 
-from .parallel_base import ParallelModel, ParallelOptimizer
+from .parallel_base import ParallelModel, ParallelOptimizer, is_tensor
 
 
 class SplitPoint(Enum):
@@ -55,19 +54,21 @@ def forward_post_hook(layer, input, output):
             # reshard to next pipeline stage
             if isinstance(output, (dict, OrderedDict)):
                 for key, tensor in output.items():
+                    assert is_tensor(tensor)
                     output[key] = dist.reshard(
                         tensor,
                         self.get_mesh(pipeline_stage_index + 1),
                         tensor.placements,
                     )
             elif isinstance(output, (list, tuple)):
                 for i in range(len(output)):
+                    assert is_tensor(output[i])
                     output[i] = dist.reshard(
                         output[i],
                         self.get_mesh(pipeline_stage_index + 1),
                         output[i].placements,
                     )
-            elif isinstance(output, paddle.Tensor):
+            elif is_tensor(output):
                 output = dist.reshard(
                     output,
                     self.get_mesh(pipeline_stage_index + 1),
diff --git a/python/paddle/distributed/auto_parallel/intermediate/tensor_parallel.py b/python/paddle/distributed/auto_parallel/intermediate/tensor_parallel.py
@@ -17,7 +17,7 @@
 import paddle
 import paddle.distributed as dist
 
-from .parallel_base import ParallelModel, ParallelOptimizer
+from .parallel_base import ParallelModel, ParallelOptimizer, is_tensor
 
 
 class PlanBase:
@@ -162,7 +162,7 @@ def sp_split(x, process_mesh, need_transpose):
         target_x = x[0]
     else:
         target_x = x
-    assert isinstance(target_x, paddle.Tensor)
+    assert is_tensor(target_x)
     assert len(target_x.shape) == 3
     if need_transpose:
         target_x = paddle.transpose(target_x, perm=[1, 0, 2])
@@ -187,7 +187,7 @@ def sp_reduce_scatter(x, process_mesh, need_transpose):
         target_x = x[0]
     else:
         target_x = x
-    assert isinstance(target_x, paddle.Tensor)
+    assert is_tensor(target_x)
     assert len(target_x.shape) == 3
     placements = target_x.placements
     if placements is None:
diff --git a/test/auto_parallel/hybrid_strategy/parallel_api.py b/test/auto_parallel/hybrid_strategy/parallel_api.py
@@ -420,14 +420,13 @@ def run_llama(self, to_static=0):
             for step, inputs in enumerate(dist_loader()):
                 input_ids, labels = inputs
                 loss = dist_model(input_ids, labels)
-                logging.info(step, loss)
-
+                logging.info(f"step: {step}  loss: {loss}")
                 if step >= 10:
                     break
 
     def run_test_cases(self):
         self.run_llama(to_static=0)
-        # self.run_llama(to_static=1)
+        self.run_llama(to_static=1)
 
 
 if __name__ == '__main__':
diff --git a/test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama.py b/test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama.py
@@ -12,14 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
 import tempfile
 import unittest
 
 import collective.test_communication_api_base as test_base
 
-os.environ['FLAGS_enable_pir_api'] = '0'
-
 
 class TestShardingParallelAPI(test_base.CommunicationTestDistBase):
     def setUp(self):
diff --git a/test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama_mp.py b/test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama_mp.py
@@ -12,14 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
 import tempfile
 import unittest
 
 import collective.test_communication_api_base as test_base
 
-os.environ['FLAGS_enable_pir_api'] = '0'
-
 
 class TestTensorParallelAPI(test_base.CommunicationTestDistBase):
     def setUp(self):
