Merge branch 'main' into lessw2020/batch_decoding

Author: lessw2020

dist_run.py

@@ -390,6 +390,11 @@ def main(args):
     schedule = ScheduleGPipe(stage, mbs)
 
     # with CUDATrackTime() as timer:
+    first_pp_rank = 0
+    last_pp_rank = pp_group_size - 1
+    # Need these global ids due to the API definition of dist.send and recv
+    first_pp_rank_global_id = dist.get_global_rank(pp_group, first_pp_rank)
+    last_pp_rank_global_id = dist.get_global_rank(pp_group, last_pp_rank)
 
     # New token generated each iteration
     total_prompts = len(prompt_lengths)