part-5 cherry from: [Distributed]Add unbalance batch for virtual pp (PaddlePaddle#58383)

* add unbalanced batch for vpp

* add unbalanced batch for vpp

* add unbalanced batch for vpp

Author: ForFishes

python/paddle/distributed/fleet/meta_parallel/pipeline_parallel.py

@@ -912,6 +912,7 @@ def __init__(self, layers, hcg, strategy):
         assert layers.get_num_virtual_stages() > 1
 
         # setup for interleave scheduler
+        self._check_sanity()
         self.num_model_chunks = layers.get_num_virtual_stages()
         self.model_chunks = layers.get_model_chunks()
         assert self.model_chunks is not None
@@ -920,7 +921,7 @@ def __init__(self, layers, hcg, strategy):
         self._virtual_pp_rank = 0
         self._reset_counter()
 
-        self._check_sanity()
+        self._assign_vpp_info(self.model_chunks)
 
     def _check_sanity(self):
         assert (