fix prompt incrementing

add formatting exception for llama2 "".res

Author: lessw2020

dist_run.py

@@ -219,7 +219,6 @@ def _update_padded_sequence(
     prompt_lengths: List[int],
 ) -> None:
     for i in range(len(prompt_lengths)):
-        prompt_lengths[i] += 1
         padded_sequence[i, prompt_lengths[i] - 1] = new_token[i, 0]
         # logger.info(f"updated prompt {i} with new token {new_token[i, 0]}")
 
@@ -427,10 +426,6 @@ def main(args):
                         [decode_results[i][0]], device=device
                     )  # decode_results[i][0]
 
-                # increment prompt lengths for next token
-                for i in range(len(prompt_lengths)):
-                    prompt_lengths[i] += 1
-
             # sendrecv between last and first ranks, only if:
             # first_pp_rank != last_pp_rank.
             if pp_rank == last_pp_rank and pp_rank != first_pp_rank:
@@ -446,6 +441,10 @@ def main(args):
                     group=pp_group,
                 )
 
+            # increment prompt lengths for next token
+            for i in range(len(prompt_lengths)):
+                prompt_lengths[i] += 1
+
             # Update input sequence with new token
             if pp_rank == first_pp_rank:
                 _update_padded_sequence(padded_sequence, new_token, prompt_lengths)