[Recompute] Update recompute for hybrid parallel interface. (#3211)

Co-authored-by: Zhong Hui <[email protected]>

Author: wuhuachaocoding

examples/language_model/gpt-3/dygraph/modeling.py

@@ -1178,4 +1178,10 @@ def _logits_helper(embedding, output):
                          loss_fn=GPTPretrainingCriterionPipe(),
                          topology=topology,
                          seg_method="layer:TransformerDecoderLayer",
-                         recompute_interval=1 if use_recompute else 0)
+                         recompute_interval=1 if use_recompute else 0,
+                         recompute_ctx={
+                             "mp_group":
+                             fleet.fleet._hcg.get_model_parallel_group(),
+                             "offload": False,
+                             "partition": False
+                         })

examples/language_model/moe/dygraph/modeling.py

@@ -1165,5 +1165,9 @@ def _logits_helper(embedding, output):
                          topology=topology,
                          seg_method="layer:TransformerDecoderLayer",
                          recompute_interval=recompute_interval,
-                         recompute_partition=False,
-                         recompute_offload=False)
+                         recompute_ctx={
+                             "mp_group":
+                             fleet.fleet._hcg.get_model_parallel_group(),
+                             "offload": False,
+                             "partition": False
+                         })