use Ke's variable names for send/rcv

Author: lessw2020

dist_run.py

@@ -373,8 +373,9 @@ def main(args):
     first_pp_rank = 0
     last_pp_rank = pp_group_size - 1
 
-    send_destination = dist.get_global_rank(pp_group, first_pp_rank)
-    recv_source = dist.get_global_rank(pp_group, last_pp_rank)
+    # 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)
 
     # encode the prompt
     input_ids = _encode_strings(
@@ -435,13 +436,13 @@ def main(args):
             if pp_rank == last_pp_rank and pp_rank != first_pp_rank:
                 dist.send(
                     new_token,
-                    dst=send_destination,
+                    dst=first_pp_rank_global_id,
                     group=pp_group,
                 )
             elif pp_rank == first_pp_rank and pp_rank != last_pp_rank:
                 dist.recv(
                     new_token,
-                    src=recv_source,
+                    src=last_pp_rank_global_id,
                     group=pp_group,
                 )