merge ipu_shard & ipu_stage -> ipu_shard, rename unittest filename (P…

…addlePaddle#71)

python/paddle/fluid/framework.py

@@ -47,7 +47,6 @@
     'program_guard',
     'name_scope',
     'ipu_shard',
-    'ipu_stage',
     'cuda_places',
     'cpu_places',
     'xpu_places',
@@ -81,24 +80,24 @@
 
 
 @signature_safe_contextmanager
-def ipu_shard(ipu_index):
+def ipu_shard(ipu_index=None, ipu_stage=None):
+    """
+    Set model sharding id and pipeline stage.
+
+    Args:
+        ipu_index: set device index of subgraph
+        ipu_stage: set pipeline stage of subgraph
+    """
     global global_ipu_index
+    global global_ipu_stage
     prev_ipu_index = global_ipu_index
+    prev_ipu_stage = global_ipu_stage
     global_ipu_index = ipu_index
+    global_ipu_stage = ipu_stage
     try:
         yield
     finally:
         global_ipu_index = prev_ipu_index
-
-
-@signature_safe_contextmanager
-def ipu_stage(stage):
-    global global_ipu_stage
-    prev_ipu_stage = global_ipu_stage
-    global_ipu_stage = stage
-    try:
-        yield
-    finally:
         global_ipu_stage = prev_ipu_stage
 
 

python/paddle/fluid/tests/unittests/ipu/ernie_training.py

@@ -496,11 +496,11 @@ def __init__(self, src_ids, sentence_ids, config):
         self.src_ids = src_ids
         self.sentence_ids = sentence_ids
 
-        with fluid.ipu_shard(0):
+        with fluid.ipu_shard(ipu_index=0):
             self.position_ids = self._build_position_ids()  # position_ids
             self.input_mask = self._build_input_mask()  # input mask
 
-        with fluid.ipu_shard(1):
+        with fluid.ipu_shard(ipu_index=1):
             self._build_model()
 
     def _build_model(self, emb=None):
@@ -761,7 +761,7 @@ def get_task_output(self, task, task_labels):
     ernie = ErnieModel(src_ids, sent_ids, ernie_config)
     fetch_node = ernie.get_sequence_output()
     if args.is_training:
-        with fluid.ipu_shard(2):
+        with fluid.ipu_shard(ipu_index=2):
             _, mean_mask_lm_loss = ernie.get_lm_output(mask_label, mask_pos)
             fetch_node = mean_mask_lm_loss
             adam = paddle.optimizer.Adam(learning_rate=1e-2)

python/paddle/fluid/tests/unittests/ipu/ipu_shard_test.py → python/paddle/fluid/tests/unittests/ipu/test_ipu_shard.py

@@ -31,21 +31,21 @@ def _test(self):
         a = paddle.static.data(name='data', shape=[None, 1], dtype='int32')
         b = a + 2  # scale : scale * x + bias, ipu_index : no
 
-        with paddle.fluid.ipu_shard(1):
+        with paddle.fluid.ipu_shard(ipu_index=1):
             c = b + 1  # scale, ipu_index : 1
-            with paddle.fluid.ipu_shard(2):
+            with paddle.fluid.ipu_shard(ipu_index=2):
                 d = c * 2  # scale, ipu_index : 2
-            with paddle.fluid.ipu_shard(3):
+            with paddle.fluid.ipu_shard(ipu_index=3):
                 e = d + 3  # scale, ipu_index : 3
-                with paddle.fluid.ipu_shard(1):
+                with paddle.fluid.ipu_shard(ipu_index=1):
                     e = e + 3  # scale, ipu_index : 1
-                    with paddle.fluid.ipu_shard(2):
+                    with paddle.fluid.ipu_shard(ipu_index=2):
                         e = e + 3  # scale, ipu_index : 2
 
-        with paddle.fluid.ipu_shard(1):
+        with paddle.fluid.ipu_shard(ipu_index=1):
             f = paddle.tensor.pow(e, 2.0)  # pow, ipu_index : 1
 
-        with paddle.fluid.ipu_shard(2):
+        with paddle.fluid.ipu_shard(ipu_index=2):
             g = f - 1  # scale, ipu_index : 2
 
         h = g + 1  # scale, ipu_index : no

python/paddle/fluid/tests/unittests/ipu/ipu_pipeline_test.py → python/paddle/fluid/tests/unittests/ipu/test_pipeline.py

@@ -31,21 +31,21 @@ def _test(self):
         a = paddle.static.data(name='data', shape=[None, 1], dtype='int32')
         b = a + 2  # scale : scale * x + bias, ipu_stage : no
 
-        with paddle.fluid.ipu_stage(1):
+        with paddle.fluid.ipu_stage(ipu_stage=1):
             c = b + 1  # scale, ipu_stage : 1
-            with paddle.fluid.ipu_stage(2):
+            with paddle.fluid.ipu_stage(ipu_stage=2):
                 d = c * 2  # scale, ipu_stage : 2
-            with paddle.fluid.ipu_stage(3):
+            with paddle.fluid.ipu_stage(ipu_stage=3):
                 e = d + 3  # scale, ipu_stage : 3
-                with paddle.fluid.ipu_stage(1):
+                with paddle.fluid.ipu_stage(ipu_stage=1):
                     e = e + 3  # scale, ipu_stage : 1
-                    with paddle.fluid.ipu_stage(2):
+                    with paddle.fluid.ipu_stage(ipu_stage=2):
                         e = e + 3  # scale, ipu_stage : 2
 
-        with paddle.fluid.ipu_stage(1):
+        with paddle.fluid.ipu_stage(ipu_stage=1):
             f = paddle.tensor.pow(e, 2.0)  # pow, ipu_stage : 1
 
-        with paddle.fluid.ipu_stage(2):
+        with paddle.fluid.ipu_stage(ipu_stage=2):
             g = f - 1  # scale, ipu_stage : 2
 
         h = g + 1  # scale, ipu_stage : no