fix pir llama vpp test

Author: waliwali777

test/auto_parallel/hybrid_strategy/semi_auto_llama_pp_gradmerge.py

@@ -120,12 +120,7 @@ def __init__(self):
         self.init_dist_env()
 
     def init_dist_env(self):
-        order = ["dp", "pp", "mp"]
-        dp_degree = self.dp
-        mp_degree = self.mp
-        pp_degree = self.pp
-        degree = [dp_degree, pp_degree, mp_degree]
-        mesh_dims = list(filter(lambda x: x[1] > 1, list(zip(order, degree))))
+        mesh_dims = [("pp", self.pp), ("dp", self.dp), ("mp", self.mp)]
         if not mesh_dims:
             mesh_dims = [("dp", 1)]
         dim_names = [mesh_dim[0] for mesh_dim in mesh_dims]
@@ -212,6 +207,7 @@ def run_llama(self, to_static=0):
                 strategy.pipeline.accumulate_steps = (
                     self.gradient_accumulation_steps
                 )
+                strategy.pipeline.pp_degree = self.pp
                 strategy.pipeline.micro_batch_size = micro_bsz
                 strategy.pipeline.schedule_mode = self.schedule_mode
                 strategy.pipeline.vpp_degree = self.config.virtual_pp_degree

test/auto_parallel/hybrid_strategy/test_semi_auto_parallel_llama_model_vpp.py

@@ -20,7 +20,7 @@
 
 import collective.test_communication_api_base as test_base
 
-os.environ['FLAGS_enable_pir_api'] = '0'
+os.environ['FLAGS_enable_pir_api'] = '1'
 
 
 class TestSemiAutoParallelLlama3DVPP(test_base.CommunicationTestDistBase):